diff --git a/.claude/skills/exec-local-compile/SKILL.md b/.claude/skills/exec-local-compile/SKILL.md index 6a2f3184895d..5004e324254d 100644 --- a/.claude/skills/exec-local-compile/SKILL.md +++ b/.claude/skills/exec-local-compile/SKILL.md @@ -45,7 +45,7 @@ git checkout main && git pull Run the build command (**incremental by default** — omit `-c`/`--clean` unless explicitly requested or the incremental build fails): ```bash -./scripts/build_wheel.py --trt_root /usr/local/tensorrt --benchmarks --use_ccache -a "" -f --nvtx +./scripts/build_wheel.py --trt_root /usr/local/tensorrt --use_ccache -a "" -f --nvtx ``` Replace `` with the target GPU architecture (see Architecture Reference below). If not specified by the user, auto-detect from `nvidia-smi`. @@ -67,7 +67,6 @@ python3 -c "import tensorrt_llm; print(tensorrt_llm.__version__)" | Flag | Description | |------|-------------| | `--trt_root /usr/local/tensorrt` | TensorRT installation path (standard in NVIDIA containers) | -| `--benchmarks` | Build the C++ benchmarks | | `-a ""` | Target GPU architecture(s) | | `--nvtx` | Enable NVTX markers for profiling | | `--use_ccache` | Use ccache for faster recompilation | diff --git a/.claude/skills/exec-slurm-compile/SKILL.md b/.claude/skills/exec-slurm-compile/SKILL.md index 44c60ed374bd..0f05a0ab4db9 100644 --- a/.claude/skills/exec-slurm-compile/SKILL.md +++ b/.claude/skills/exec-slurm-compile/SKILL.md @@ -205,7 +205,6 @@ A successful build ends with a message like `Successfully built tensorrt_llm` or | Flag | Description | |------|-------------| | `--trt_root /usr/local/tensorrt` | TensorRT installation path (standard in NVIDIA containers) | -| `--benchmarks` | Build the C++ benchmarks | | `-a "100-real"` | Target architecture — `100` for Blackwell, `90` for Hopper, etc. | | `--nvtx` | Enable NVTX markers for profiling | | `--no-venv` | Skip virtual environment creation | diff --git a/.claude/skills/exec-slurm-compile/scripts/compile.sh b/.claude/skills/exec-slurm-compile/scripts/compile.sh index 22b7882d55ff..750b10cd8584 100755 --- a/.claude/skills/exec-slurm-compile/scripts/compile.sh +++ b/.claude/skills/exec-slurm-compile/scripts/compile.sh @@ -19,7 +19,7 @@ # Usage: compile.sh [build_wheel_args...] # # Default build_wheel.py flags: -# --trt_root /usr/local/tensorrt --benchmarks -a "100-real" --nvtx --no-venv +# --trt_root /usr/local/tensorrt -a "100-real" --nvtx --no-venv # Any extra arguments after repo_dir are forwarded to build_wheel.py, # overriding the defaults above. @@ -37,7 +37,6 @@ else echo "[compile.sh] Running default build command" python3 ./scripts/build_wheel.py \ --trt_root /usr/local/tensorrt \ - --benchmarks \ -a "100-real" \ --nvtx fi diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 83e0275ad77c..a9a47331064f 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -37,8 +37,6 @@ /tests/unittest @NVIDIA/trt-llm-devs # ===== TensorRT backend (will be deprecated soon) — also on the trt-llm-devs fallback ===== -/cpp/include/tensorrt_llm/plugins @NVIDIA/trt-llm-devs -/cpp/tensorrt_llm/plugins @NVIDIA/trt-llm-devs /tensorrt_llm/builder.py @NVIDIA/trt-llm-devs /tensorrt_llm/commands/build.py @NVIDIA/trt-llm-devs /tensorrt_llm/commands/prune.py @NVIDIA/trt-llm-devs @@ -105,7 +103,6 @@ /cpp/tensorrt_llm/batch_manager @NVIDIA/trt-llm-runtime-devs /cpp/tensorrt_llm/common @NVIDIA/trt-llm-runtime-devs /cpp/tensorrt_llm/executor @NVIDIA/trt-llm-runtime-devs -/cpp/tensorrt_llm/executor_worker @NVIDIA/trt-llm-runtime-devs /cpp/tensorrt_llm/layers @NVIDIA/trt-llm-runtime-devs /cpp/tensorrt_llm/nanobind @NVIDIA/trt-llm-runtime-devs /cpp/tensorrt_llm/runtime @NVIDIA/trt-llm-runtime-devs @@ -169,7 +166,6 @@ /examples/llm-api/quickstart_multimodal.py @NVIDIA/trt-llm-models-devs @NVIDIA/trt-llm-doc-owners /examples/models @NVIDIA/trt-llm-models-devs @NVIDIA/trt-llm-doc-owners /examples/serve/*multimodal* @NVIDIA/trt-llm-models-devs @NVIDIA/trt-llm-doc-owners -/scripts/build_cpp_examples.py @NVIDIA/trt-llm-models-devs /scripts/generate_config_database_tests.py @NVIDIA/trt-llm-models-devs @NVIDIA/trt-llm-doc-owners /scripts/generate_config_table.py @NVIDIA/trt-llm-models-devs @NVIDIA/trt-llm-doc-owners /tensorrt_llm/_torch/models @NVIDIA/trt-llm-models-devs diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 452c3be57ea1..7dccd521319c 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -9,14 +9,14 @@ common-files: &common_files | .devcontainer/make_env.py | .github/scripts/label_community_user.py | .github/scripts/pr_checklist_check.py | - benchmarks/cpp/__init__.py | - benchmarks/cpp/prepare_dataset.py | - benchmarks/cpp/utils/__init__.py | - benchmarks/cpp/utils/convert_nemo_dataset.py | - benchmarks/cpp/utils/generate_rand_loras.py | - benchmarks/cpp/utils/prepare_real_data.py | - benchmarks/cpp/utils/prepare_synthetic_data.py | - benchmarks/cpp/utils/utils.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 | @@ -41,38 +41,11 @@ common-files: &common_files | cpp/tensorrt_llm/deep_ep/strip_nvshmem_helper.py | cpp/tensorrt_llm/kernels/cutlass_kernels/python/generate_kernels.py | cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/copy_cu.py | - cpp/tests/resources/scripts/build_chatglm_engines.py | - cpp/tests/resources/scripts/build_eagle_engines.py | - cpp/tests/resources/scripts/build_enc_dec_engines.py | - cpp/tests/resources/scripts/build_engines_utils.py | - cpp/tests/resources/scripts/build_gpt_engines.py | - cpp/tests/resources/scripts/build_gptj_engines.py | - cpp/tests/resources/scripts/build_llama_engines.py | - cpp/tests/resources/scripts/build_mamba_engines.py | - cpp/tests/resources/scripts/build_medusa_engines.py | - cpp/tests/resources/scripts/build_recurrentgemma_engines.py | - cpp/tests/resources/scripts/build_redrafter_engines.py | - cpp/tests/resources/scripts/generate_expected_chatglm_output.py | - cpp/tests/resources/scripts/generate_expected_eagle_output.py | - cpp/tests/resources/scripts/generate_expected_enc_dec_output.py | - cpp/tests/resources/scripts/generate_expected_gpt_output.py | - cpp/tests/resources/scripts/generate_expected_gptj_output.py | - cpp/tests/resources/scripts/generate_expected_llama_output.py | - cpp/tests/resources/scripts/generate_expected_mamba_output.py | - cpp/tests/resources/scripts/generate_expected_medusa_output.py | - cpp/tests/resources/scripts/generate_expected_recurrentgemma_output.py | - cpp/tests/resources/scripts/generate_expected_redrafter_output.py | - cpp/tests/resources/scripts/generate_hf_gpt_output.py | cpp/tests/resources/scripts/generate_test_lora_weights.py | - cpp/tests/resources/scripts/io_converter.py | docs/source/conf.py | docs/source/helper.py | examples/apps/chat.py | examples/apps/fastapi_server.py | - examples/bindings/executor/example_advanced.py | - examples/bindings/executor/example_basic.py | - examples/bindings/executor/example_debug.py | - examples/bindings/executor/example_logits_processor.py | examples/disaggregated/clients/disagg_client.py | examples/disaggregated/slurm/benchmark/submit.py | examples/dora/normalize_weights.py | @@ -246,7 +219,6 @@ common-files: &common_files | jenkins/scripts/mergeWaiveList.py | jenkins/scripts/open_search_db.py | jenkins/scripts/test_rerun.py | - scripts/build_cpp_examples.py | scripts/build_wheel.py | scripts/check_test_list.py | scripts/dco_check.py | @@ -851,11 +823,8 @@ common-files: &common_files | tests/integration/defs/conftest.py | tests/integration/defs/cpp/conftest.py | tests/integration/defs/cpp/cpp_common.py | - tests/integration/defs/cpp/test_e2e.py | tests/integration/defs/cpp/test_multi_gpu.py | tests/integration/defs/cpp/test_unit_tests.py | - tests/integration/defs/deterministic/mixtral_deterministic.py | - tests/integration/defs/deterministic/test_mixtral_deterministic.py | tests/integration/defs/disaggregated/test_auto_scaling.py | tests/integration/defs/disaggregated/test_disaggregated.py | tests/integration/defs/disaggregated/test_disaggregated_etcd.py | @@ -928,7 +897,6 @@ common-files: &common_files | tests/integration/defs/test_fmha.py | tests/integration/defs/test_list_parser.py | tests/integration/defs/test_list_validation.py | - tests/integration/defs/test_mlpf_results.py | tests/integration/defs/test_sanity.py | tests/integration/defs/test_unittests.py | tests/integration/defs/triton_server/__init__.py | @@ -1081,7 +1049,6 @@ common-files: &common_files | tests/unittest/_torch/thop/serial/test_moe_alltoall.py | tests/unittest/api_stability/api_stability_core.py | tests/unittest/api_stability/test_llm_api.py | - tests/unittest/bindings/binding_test_utils.py | tests/unittest/bindings/test_bindings_moe.py | tests/unittest/bindings/test_bindings_ut.py | tests/unittest/bindings/test_executor_bindings.py | @@ -1169,7 +1136,6 @@ common-files: &common_files | 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_leak.py | tests/unittest/others/test_mapping.py | tests/unittest/others/test_model_dtype.py | tests/unittest/others/test_module.py | @@ -1350,14 +1316,14 @@ legacy-files: &legacy_files | .devcontainer/make_env.py | .github/scripts/label_community_user.py | .github/scripts/pr_checklist_check.py | - benchmarks/cpp/__init__.py | - benchmarks/cpp/prepare_dataset.py | - benchmarks/cpp/utils/__init__.py | - benchmarks/cpp/utils/convert_nemo_dataset.py | - benchmarks/cpp/utils/generate_rand_loras.py | - benchmarks/cpp/utils/prepare_real_data.py | - benchmarks/cpp/utils/prepare_synthetic_data.py | - benchmarks/cpp/utils/utils.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 | @@ -1382,38 +1348,11 @@ legacy-files: &legacy_files | cpp/tensorrt_llm/deep_ep/strip_nvshmem_helper.py | cpp/tensorrt_llm/kernels/cutlass_kernels/python/generate_kernels.py | cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/copy_cu.py | - cpp/tests/resources/scripts/build_chatglm_engines.py | - cpp/tests/resources/scripts/build_eagle_engines.py | - cpp/tests/resources/scripts/build_enc_dec_engines.py | - cpp/tests/resources/scripts/build_engines_utils.py | - cpp/tests/resources/scripts/build_gpt_engines.py | - cpp/tests/resources/scripts/build_gptj_engines.py | - cpp/tests/resources/scripts/build_llama_engines.py | - cpp/tests/resources/scripts/build_mamba_engines.py | - cpp/tests/resources/scripts/build_medusa_engines.py | - cpp/tests/resources/scripts/build_recurrentgemma_engines.py | - cpp/tests/resources/scripts/build_redrafter_engines.py | - cpp/tests/resources/scripts/generate_expected_chatglm_output.py | - cpp/tests/resources/scripts/generate_expected_eagle_output.py | - cpp/tests/resources/scripts/generate_expected_enc_dec_output.py | - cpp/tests/resources/scripts/generate_expected_gpt_output.py | - cpp/tests/resources/scripts/generate_expected_gptj_output.py | - cpp/tests/resources/scripts/generate_expected_llama_output.py | - cpp/tests/resources/scripts/generate_expected_mamba_output.py | - cpp/tests/resources/scripts/generate_expected_medusa_output.py | - cpp/tests/resources/scripts/generate_expected_recurrentgemma_output.py | - cpp/tests/resources/scripts/generate_expected_redrafter_output.py | - cpp/tests/resources/scripts/generate_hf_gpt_output.py | cpp/tests/resources/scripts/generate_test_lora_weights.py | - cpp/tests/resources/scripts/io_converter.py | docs/source/conf.py | docs/source/helper.py | examples/apps/chat.py | examples/apps/fastapi_server.py | - examples/bindings/executor/example_advanced.py | - examples/bindings/executor/example_basic.py | - examples/bindings/executor/example_debug.py | - examples/bindings/executor/example_logits_processor.py | examples/disaggregated/clients/disagg_client.py | examples/disaggregated/slurm/benchmark/submit.py | examples/dora/normalize_weights.py | @@ -1587,7 +1526,6 @@ legacy-files: &legacy_files | jenkins/scripts/mergeWaiveList.py | jenkins/scripts/open_search_db.py | jenkins/scripts/test_rerun.py | - scripts/build_cpp_examples.py | scripts/build_wheel.py | scripts/check_test_list.py | scripts/dco_check.py | @@ -2192,11 +2130,8 @@ legacy-files: &legacy_files | tests/integration/defs/conftest.py | tests/integration/defs/cpp/conftest.py | tests/integration/defs/cpp/cpp_common.py | - tests/integration/defs/cpp/test_e2e.py | tests/integration/defs/cpp/test_multi_gpu.py | tests/integration/defs/cpp/test_unit_tests.py | - tests/integration/defs/deterministic/mixtral_deterministic.py | - tests/integration/defs/deterministic/test_mixtral_deterministic.py | tests/integration/defs/disaggregated/test_auto_scaling.py | tests/integration/defs/disaggregated/test_disaggregated.py | tests/integration/defs/disaggregated/test_disaggregated_etcd.py | @@ -2269,7 +2204,6 @@ legacy-files: &legacy_files | tests/integration/defs/test_fmha.py | tests/integration/defs/test_list_parser.py | tests/integration/defs/test_list_validation.py | - tests/integration/defs/test_mlpf_results.py | tests/integration/defs/test_sanity.py | tests/integration/defs/test_unittests.py | tests/integration/defs/triton_server/__init__.py | @@ -2422,7 +2356,6 @@ legacy-files: &legacy_files | tests/unittest/_torch/thop/serial/test_moe_alltoall.py | tests/unittest/api_stability/api_stability_core.py | tests/unittest/api_stability/test_llm_api.py | - tests/unittest/bindings/binding_test_utils.py | tests/unittest/bindings/test_bindings_moe.py | tests/unittest/bindings/test_bindings_ut.py | tests/unittest/bindings/test_executor_bindings.py | @@ -2510,7 +2443,6 @@ legacy-files: &legacy_files | 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_leak.py | tests/unittest/others/test_mapping.py | tests/unittest/others/test_model_dtype.py | tests/unittest/others/test_module.py | diff --git a/benchmarks/README.md b/benchmarks/README.md index 5d89f412ac5c..87e2c06e6d5d 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -1,10 +1,34 @@ -# TensorRT-LLM Benchmarks +# TensorRT-LLM Benchmarking -## Overview +For benchmarking TensorRT-LLM, use +[`trtllm-bench`](../docs/source/developer-guide/perf-benchmarking.md) and +`trtllm-serve`. -There are currently two workflows to benchmark TensorRT-LLM: -* [`trtllm-bench`](../docs/source/developer-guide/perf-benchmarking.md) - - `trtllm-bench` is native to TensorRT-LLM and is a Python benchmarker for reproducing and testing the performance of TensorRT-LLM. - - _NOTE_: This benchmarking suite is a current work in progress and is prone to large changes. -* [C++ benchmarks](./cpp) - - The recommended workflow that uses TensorRT-LLM C++ API and can take advantage of the latest features of TensorRT-LLM. +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/cpp/__init__.py b/benchmarks/__init__.py similarity index 100% rename from benchmarks/cpp/__init__.py rename to benchmarks/__init__.py diff --git a/benchmarks/cpp/CMakeLists.txt b/benchmarks/cpp/CMakeLists.txt deleted file mode 100644 index cb5ef1ee928b..000000000000 --- a/benchmarks/cpp/CMakeLists.txt +++ /dev/null @@ -1,42 +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. - -include_directories(${PROJECT_SOURCE_DIR}/include) - -set(TOP_LEVEL_DIR "${PROJECT_SOURCE_DIR}/..") - -add_custom_target(benchmarks) - -if(NOT TARGET cxxopts::cxxopts) - add_subdirectory(${CMAKE_BINARY_DIR}/_deps/cxxopts-src - ${CMAKE_CURRENT_BINARY_DIR}/cxxopts) -endif() - -function(add_benchmark test_name test_src) - add_executable(${test_name} ${test_src} utils/utils.cpp) - - target_link_libraries( - ${test_name} PUBLIC ${SHARED_TARGET} nvinfer_plugin_tensorrt_llm - cxxopts::cxxopts) - - target_compile_features(${test_name} PRIVATE cxx_std_17) - target_compile_definitions(${test_name} - PUBLIC TOP_LEVEL_DIR="${TOP_LEVEL_DIR}") - add_dependencies(benchmarks ${test_name}) -endfunction() - -add_benchmark(bertBenchmark bertBenchmark.cpp) -add_benchmark(gptManagerBenchmark gptManagerBenchmark.cpp) -add_benchmark(disaggServerBenchmark disaggServerBenchmark.cpp) diff --git a/benchmarks/cpp/README.md b/benchmarks/cpp/README.md deleted file mode 100644 index ae3287faf06c..000000000000 --- a/benchmarks/cpp/README.md +++ /dev/null @@ -1,367 +0,0 @@ -# Benchmark C++ Runtime - -This document explains how to benchmark the models supported by TensorRT-LLM on a single GPU, a single node with -multiple GPUs or multiple nodes with multiple GPUs using the C++ runtime. - -## Usage - -### 1. Build TensorRT-LLM and benchmarking source code - -Please follow the [`installation document`](../../README.md#installation) to build TensorRT-LLM. - -Note that the benchmarking source code for C++ runtime is not built by default, you can use the argument `--benchmarks` in [`build_wheel.py`](source:scripts/build_wheel.py) to build the corresponding executable. - -### 2. Launch C++ benchmarking (Inflight/V1 batching) - -#### Prepare dataset - -Run a preprocessing script to prepare/generate dataset into a json that `gptManagerBenchmark` can consume later. The processed output json has *input tokens length, input token ids and output tokens length*. - -For `tokenizer`, specifying the path to the local tokenizer that have already been downloaded, or simply the name of the tokenizer from HuggingFace like `meta-llama/Llama-2-7b` will both work. The tokenizer will be downloaded automatically for the latter case. - -This tool can be used in 3 different modes of traffic generation: `dataset`, `token-norm-dist` and `token-unif-dist`. - -##### 1 – Dataset - -The tool will tokenize the words and instruct the model to generate a specified number of output tokens for a request. - -``` -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] -``` - -For datasets that don't have prompt key, set --dataset-prompt instead. -Take [cnn_dailymail dataset](https://huggingface.co/datasets/abisee/cnn_dailymail) for example: -``` -python3 prepare_dataset.py \ - --tokenizer \ - --output cnn_dailymail.json - dataset - --dataset-name cnn_dailymail \ - --dataset-split validation \ - --dataset-config-name 3.0.0 \ - --dataset-input-key article \ - --dataset-prompt "Summarize the following article:" \ - --dataset-output-key "highlights" \ - [--num-requests 100] \ - [--max-input-len 1000] \ - [--output-len-dist 100,10] -``` - -##### 2 – Normal token length distribution - -This mode allows the user to generate normally distributed token lengths with a mean and std deviation specified. -For example, setting `mean=100` and `stdev=10` would generate requests where 95.4% of values are in <80,120> range following the normal probability distribution. Setting `stdev=0` will generate all requests with the same mean number of tokens. - -``` -python prepare_dataset.py \ - --output token-norm-dist.json \ - --tokenizer \ - token-norm-dist \ - --num-requests 100 \ - --input-mean 100 --input-stdev 10 \ - --output-mean 15 --output-stdev 0 -``` - -##### 2 – Uniform token length distribution - -This mode allows the user to generate uniformly distributed token lengths with min and max lengths specified. -For example, setting `min=50` and `max=100` would generate requests where lengths are in the range `[50, 100]` following the uniform probability distribution. Setting `min=x` and `max=x` will generate all requests with the same mean number of tokens `x`. - -``` -python prepare_dataset.py \ - --output token-norm-dist.json \ - --tokenizer \ - token-unif-dist \ - --num-requests 100 \ - --input-min 50 --input-max 100 \ - --output-min 10 --output-max 15 -``` - - -#### Prepare TensorRT-LLM engines - -Before you launch C++ benchmarking, please make sure that you have already built engine(s) using `trtllm-build` command. For more details on building engine(s), please refer to the [Quick Start Guide](../../docs/source/quick-start-guide.md). - -#### Launch benchmarking - -For detailed usage, you can do the following -``` -cd cpp/build - -# You can directly execute the binary for help information -./benchmarks/gptManagerBenchmark --help -``` - -`gptManagerBenchmark` now supports decoder-only models and encoder-decoder models. - -1. Decoder-only Models - - To benchmark decoder-only models, pass in the engine path with `--engine_dir` as executable input argument. - - Take GPT-350M as an example for 2-GPU inflight batching - ``` - mpirun -n 2 ./benchmarks/gptManagerBenchmark \ - --engine_dir ../../examples/models/core/gpt/trt_engine/gpt2-ib/fp16/2-gpu/ \ - --request_rate 10 \ - --dataset ../../benchmarks/cpp/preprocessed_dataset.json \ - --max_num_samples 500 - ``` - - `gptManagerBenchmark` by default uses the high-level C++ API defined by the `executor::Executor` class (see `cpp/include/tensorrt_llm/executor/executor.h`). - -2. Encoder-Decoder Models - To benchmark encoder-decoder models, pass in the encoder engine path with `--encoder_engine_dir` and the decoder engine path with `--decoder_engine_dir` as executable input arguments. `--decoder_engine_dir` is an alias of `--engine_dir`. - - Currently encoder-decoder engines only support `--api executor`, `--type IFB`, `--enable_kv_cache_reuse false`, which are all default values so no specific settings required. - - Prepare t5-small engine from [examples/models/core/enc_dec](/examples/models/core/enc_dec/README.md#convert-and-split-weights) for the encoder-decoder 4-GPU inflight batching example. - - Prepare the dataset suitable for engine input lengths. - ``` - python prepare_dataset.py \ - --tokenizer \ - --output cnn_dailymail.json \ - dataset \ - --dataset-name cnn_dailymail \ - --dataset-split validation \ - --dataset-config-name 3.0.0 \ - --dataset-input-key article \ - --dataset-prompt "Summarize the following article:" \ - --dataset-output-key "highlights" \ - --num-requests 100 \ - --max-input-len 512 \ - --output-len-dist 128,20 - ``` - - Run the benchmark - ``` - mpirun --allow-run-as-root -np 4 ./benchmarks/gptManagerBenchmark \ - --encoder_engine_dir ../../examples/models/core/enc_dec/tmp/trt_engines/t5-small-4gpu/bfloat16/encoder \ - --decoder_engine_dir ../../examples/models/core/enc_dec/tmp/trt_engines/t5-small-4gpu/bfloat16/decoder \ - --dataset cnn_dailymail.json - ``` - - -#### Emulated static batching - -To emulate the deprecated `gptSessionBenchmark` static batching, you can use `gptManagerBenchmark` with the `--static_emulated_batch_size` and `--static_emulated-timeout` arguments. - -Given a `static_emulated_batch_size` of `n` the server will wait for `n` requests to arrive before submitting them to the batch manager at once. If the `static_emulated_timeout` (in ms) is reached before `n` requests are collected, the batch will be submitted prematurely with the current request count. New batches will only be submitted once the previous batch has been processed comepletely. - -Datasets with fixed input/output lengths for benchmarking can be generated with the preprocessing script, e.g. -``` - python prepare_dataset.py \ - --output tokens-fixed-lengths.json \ - --tokenizer \ - token-norm-dist \ - --num-requests 128 \ - --input-mean 60 --input-stdev 0 \ - --output-mean 20 --output-stdev 0 -``` - -Take GPT-350M as an example for single GPU with static batching -``` -./benchmarks/gptManagerBenchmark \ - --engine_dir ../../examples/models/core/gpt/trt_engine/gpt2/fp16/1-gpu/ \ - --request_rate -1 \ - --static_emulated_batch_size 32 \ - --static_emulated_timeout 100 \ - --dataset ../../benchmarks/cpp/tokens-fixed-lengths.json -``` - -#### Benchmarking LoRA - -Using either of the `prepare_dataset.py` methods above, add `--rand-task-id ` to the command. This will add a random `task_id` from `` to `` inclusive. -You can then use `utils/generate_rand_loras.py` to generate random LoRA weights for benchmarking purposes. `utils/generate_rand_loras.py` takes an example LoRA for the model you are benchmarking. -Then you can run `gptManagerBenchmark` with `--type IFB` and `--lora_dir /path/to/utils/generate_rand_loras/output` - -End-to-end LoRA benchmarking script - -``` -git-lfs clone https://huggingface.co/meta-llama/Llama-2-13b-hf -git-lfs clone https://huggingface.co/hfl/chinese-llama-2-lora-13b - -MODEL_CHECKPOINT=Llama-2-13b-hf -CONVERTED_CHECKPOINT=Llama-2-13b-hf-ckpt -TOKENIZER=Llama-2-13b-hf -LORA_ENGINE=Llama-2-13b-hf-engine - -DTYPE=float16 -TP=2 -PP=1 -MAX_LEN=1024 -MAX_BATCH=32 -NUM_LAYERS=40 -MAX_LORA_RANK=64 -NUM_LORA_MODS=7 -EOS_ID=2 - -SOURCE_LORA=chinese-llama-2-lora-13b -CPP_LORA=chinese-llama-2-lora-13b-cpp - -EG_DIR=/tmp/lora-eg - -# Build lora enabled engine -python examples/models/core/llama/convert_checkpoint.py --model_dir ${MODEL_CHECKPOINT} \ - --output_dir ${CONVERTED_CHECKPOINT} \ - --dtype ${DTYPE} \ - --tp_size ${TP} \ - --pp_size 1 - -${HOME}/.local/bin/trtllm-build \ - --checkpoint_dir ${CONVERTED_CHECKPOINT} \ - --output_dir ${LORA_ENGINE} \ - --max_batch_size ${MAX_BATCH} \ - --max_input_len $MAX_LEN \ - --max_seq_len $((2*${MAX_LEN})) \ - --gemm_plugin float16 \ - --lora_plugin float16 \ - --use_paged_context_fmha enable \ - --lora_target_modules attn_q attn_k attn_v attn_dense mlp_h_to_4h mlp_4h_to_h mlp_gate \ - --max_lora_rank ${MAX_LORA_RANK} - -NUM_LORAS=(8 16) -NUM_REQUESTS=1024 - -# Convert LoRA to cpp format -python examples/hf_lora_convert.py \ - -i $SOURCE_LORA \ - --storage-type $DTYPE \ - -o $CPP_LORA - -# Prepare datasets -mkdir -p $EG_DIR/data - -# Prepare dataset without lora_task_id -python benchmarks/cpp/prepare_dataset.py \ - --output "${EG_DIR}/data/token-norm-dist.json" \ - --tokenizer $TOKENIZER \ - token-norm-dist \ - --num-requests $NUM_REQUESTS \ - --input-mean 256 --input-stdev 16 --output-mean 128 --output-stdev 24 - -# Prepare dataset with lora_task_ids from 0 - $nloras -for nloras in ${NUM_LORAS[@]}; do - python benchmarks/cpp/prepare_dataset.py \ - --output "${EG_DIR}/data/token-norm-dist-lora-${nloras}.json" \ - --rand-task-id 0 $(( $nloras - 1 )) \ - --tokenizer $TOKENIZER \ - token-norm-dist \ - --num-requests $NUM_REQUESTS \ - --input-mean 256 --input-stdev 16 --output-mean 128 --output-stdev 24 -done - -# Generate random lora weights for 16 adapters -python benchmarks/cpp/utils/generate_rand_loras.py ${CPP_LORA} ${EG_DIR}/loras 16 - -# Perform benchmarking - -# First run inference without LoRAs -mkdir -p ${EG_DIR}/log-base-lora -mpirun -n ${TP} --output-filename ${EG_DIR}/log-base-lora \ - cpp/build/benchmarks/gptManagerBenchmark \ - --engine_dir $LORA_ENGINE \ - --type IFB \ - --dataset "${EG_DIR}/data/token-norm-dist.json" \ - --lora_host_cache_bytes 8589934592 \ - --lora_num_device_mod_layers $(( 32 * $NUM_LAYERS * $NUM_LORA_MODS * $MAX_LORA_RANK )) \ - --kv_cache_free_gpu_mem_fraction 0.70 \ - --log_level info \ - --eos_id ${EOS_ID} - -# Now run inference with various numbers or loras -# The host cache is set large enough to hold all the LoRAs in lora_dir -# GPU cache is set to hold 16 LoRAs -# This benchmark will preload all the LoRAs into the host cache -# We run inference on a range of active LoRAs exercising different cache miss rates. -for nloras in ${NUM_LORAS[@]}; do - mkdir -p ${EG_DIR}/log-lora-${nloras} - mpirun -n ${TP} --output-filename "${EG_DIR}/log-lora-${nloras}" \ - cpp/build/benchmarks/gptManagerBenchmark \ - --engine_dir $LORA_ENGINE \ - --type IFB \ - --dataset "${EG_DIR}/data/token-norm-dist-lora-${nloras}.json" \ - --lora_host_cache_bytes 8589934592 \ - --lora_num_device_mod_layers $(( 16 * $NUM_LAYERS * $NUM_LORA_MODS * $MAX_LORA_RANK )) \ - --kv_cache_free_gpu_mem_fraction 0.70 \ - --log_level info \ - --eos_id ${EOS_ID} \ - --lora_dir ${EG_DIR}/loras -done -``` - -### 3. [DEPRECATED] Launch C++ static batching benchmarking (Fixed BatchSize/InputLen/OutputLen) - -#### Prepare TensorRT-LLM engine(s) - -Before you launch C++ benchmarking, please make sure that you have already built engine(s) using TensorRT-LLM API, C++ benchmarking code cannot generate engine(s) for you. - -Use `trtllm-build` to build the TRT-LLM engine. Alternatively, if you have already benchmarked Python Runtime, you can reuse the engine(s) built previously, please see that [`document`](../python/README.md). - -#### Launch benchmarking - -For detailed usage, you can do the following -``` -cd cpp/build - -# You can directly execute the binary for help information -./benchmarks/bertBenchmark --help -``` - -*Please note that the expected outputs in that document are only for reference, specific performance numbers depend on the GPU you're using.* - - -### 4.launch C++ disaggServerBenchmark -Currently ,TensorRT-LLM has limited support for disaggregated inference, where context and generation phases of a request can run on different executors. `disaggServerBenchmark` is a tool to benchmark disaggregated inference. - -#### Usage -For detailed usage, you can do the following -``` -cd cpp/build - -# You can directly execute the binary for help information -./benchmarks/disaggServerBenchmark --help -``` -`disaggServerBenchmark` only supports `decoder-only` models. -Here is the basic usage: -``` -export TRTLLM_USE_UCX_KVCACHE=1 -mpirun -n ${proc} benchmarks/disaggServerBenchmark --context_engine_dirs ${context_engine_0},${context_engine_1}...,${context_engine_{m-1}} \ ---generation_engine_dirs ${generation_engine_0},${generation_engine_1}...,${generation_engine_{n-1}} --dataset ${dataset_path} -``` -This command will launch m context engines and n generation engines. You need to ensure `proc` is equal to the sum of the number of processes required for each engine plus 1. Since we use orchestrator mode for `disaggServerBenchmark` we need an additional process as the orchestrator. For example, if there are two context engines (one is TP2_PP1,another is TP1_PP1) and two generation engines(one is TP2_PP1,another is TP1_PP1), then the `proc` value should be set to 7. - -for example: -``` -export TRTLLM_USE_UCX_KVCACHE=1 -mpirun -n 7 benchmarks/disaggServerBenchmark --context_engine_dirs ${llama_7b_tp2_pp1_dir},${llama_7b_tp1_pp1_dir} --generation_engine_dirs ${llama_7b_tp1_pp1_dir},${llama_7b_tp2_pp1_dir} --dataset ${dataset_path} - -# need 6 gpus and 7 processes to launch the benchmark. -``` - -#### Known Issues - -##### 1. error `All available sequence slots are used` - -If generation_engine's pp_size >1, the error "All available sequence slots are used" may occur, setting and adjusting the parameter `--request_rate` may help alleviate the problem. - -##### 2.KVCache transfers are by default via PCIE on single node. -Currently, because of the dependency libraries,KVCache transfers are by default via PCIE on single node. - -If you want to use NVLink, please check the UCX version in the container by running: -``` -ucx_info -v -``` -If the UCX version is less than or equal to 1.17, set `UCX_RNDV_FRAG_MEM_TYPE=cuda` to enable KvCache transfers using NVLink. -If the UCX version is 1.18, please set `UCX_CUDA_COPY_ASYNC_MEM_TYPE=cuda` to enable KvCache transfers using NVLink. diff --git a/benchmarks/cpp/bertBenchmark.cpp b/benchmarks/cpp/bertBenchmark.cpp deleted file mode 100644 index cc10a5b49eee..000000000000 --- a/benchmarks/cpp/bertBenchmark.cpp +++ /dev/null @@ -1,260 +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. - */ - -#include "tensorrt_llm/plugins/api/tllmPlugin.h" -#include "tensorrt_llm/runtime/iTensor.h" -#include "tensorrt_llm/runtime/rawEngine.h" -#include "tensorrt_llm/runtime/tllmLogger.h" -#include "tensorrt_llm/runtime/tllmRuntime.h" -#include "tensorrt_llm/runtime/worldConfig.h" - -#include -#include -#include - -#include -#include -#include -#include -#include -#include - -using namespace tensorrt_llm::runtime; - -namespace trt = nvinfer1; - -namespace -{ - -std::string engineFilename( - std::filesystem::path const& dataPath, WorldConfig const& worldConfig, std::string const& model) -{ - auto constexpr allowExceptions = true; - auto constexpr ignoreComments = true; - auto const jsonFilePath = dataPath / "config.json"; - TLLM_CHECK_WITH_INFO( - std::filesystem::exists(jsonFilePath), std::string("File does not exist: ") + jsonFilePath.string()); - std::ifstream jsonStream(jsonFilePath); - auto const json = nlohmann::json::parse(jsonStream, nullptr, allowExceptions, ignoreComments); - auto const& builderConfig = json.at("builder_config"); - auto const precision = builderConfig.at("precision").template get(); - auto const worldSize = builderConfig.at("tensor_parallel").template get(); - - TLLM_CHECK_WITH_INFO(worldSize == worldConfig.getSize(), "world size mismatch"); - return model + "_" + precision + "_tp" + std::to_string(worldConfig.getSize()) + "_rank" - + std::to_string(worldConfig.getRank()) + ".engine"; -} - -void benchmarkBert(std::string const& modelName, std::filesystem::path const& dataPath, - std::vector const& batchSizes, std::vector const& inLens, bool useGpuDirectStorage, - std::vector const& gpuWeightsPercents, std::shared_ptr const& logger, int warmUp, - int numRuns, int duration) -{ - auto const worldConfig = WorldConfig::mpi(); - auto const enginePath = dataPath / engineFilename(dataPath, worldConfig, modelName); - - for (float gpuWeightsPercent : gpuWeightsPercents) - { - auto rt = std::make_shared( - RawEngine(enginePath), logger.get(), useGpuDirectStorage, gpuWeightsPercent); - rt->addContext(0); - for (auto inLen : inLens) - { - for (auto const batchSize : batchSizes) - { - auto& allocator = rt->getBufferManager(); - TllmRuntime::TensorMap tensorMap{}; - - // input_ids - std::vector inputIdsHost(batchSize * inLen, inLen); - auto inputIdsBuffer = std::shared_ptr{ - allocator.copyFrom(inputIdsHost, ITensor::makeShape({batchSize, inLen}), MemoryType::kGPU)}; - allocator.setZero(*inputIdsBuffer); - tensorMap.insert(std::make_pair("input_ids", inputIdsBuffer)); - // input_lengths - std::vector inputLengthsHost(batchSize); - auto inLensBuffer = std::shared_ptr{ - allocator.copyFrom(inputLengthsHost, ITensor::makeShape({batchSize}), MemoryType::kGPU)}; - allocator.setZero(*inLensBuffer); - tensorMap.insert(std::make_pair("input_lengths", inLensBuffer)); - - rt->setInputTensors(0, tensorMap); - rt->setOutputTensors(0, tensorMap); - cudaDeviceSynchronize(); - - for (auto r = 0; r < warmUp; ++r) - { - rt->executeContext(0); - rt->getStream().synchronize(); - } - cudaDeviceSynchronize(); - - int iterIdx = 0; - float curDuration = 0; - while (iterIdx < numRuns || curDuration / 1000 < duration) - { - auto const start = std::chrono::steady_clock::now(); - rt->executeContext(0); - rt->getStream().synchronize(); - auto const end = std::chrono::steady_clock::now(); - - iterIdx += 1; - curDuration += (static_cast( - std::chrono::duration_cast(end - start).count()) - / 1000); - } - printf("Benchmarking done. Iteration: %d, duration: %.2f sec.\n", iterIdx, curDuration / 1000); - - auto averageLatency = curDuration / iterIdx; - - if (worldConfig.getRank() == 0) - { - printf("[BENCHMARK] batch_size %d input_length %d latency(ms) %.2f\n", batchSize, inLen, - averageLatency); - } - } - } - } -} - -} // namespace - -int main(int argc, char* argv[]) -{ - cxxopts::Options options("TensorRT LLM C++ Runtime Benchmark", "TensorRT LLM C++ Runtime Benchmark for BERT."); - options.add_options()("h,help", "Print usage"); - options.add_options()( - "m,model", "Model name specified for engines.", cxxopts::value()->default_value("bert_base")); - options.add_options()("engine_dir", "Directory that store the engines.", cxxopts::value()); - options.add_options()("batch_size", - "Specify batch size(s) you want to benchmark. Multiple batch sizes can be separated by \";\", example: " - "\"1;8;64\".", - cxxopts::value()->default_value("8")); - options.add_options()("input_len", - "Specify input length(s) you want to benchmark. Multiple input lengths can be " - "separated by \";\", example: \"60;128\".", - cxxopts::value()->default_value("128")); - - options.add_options()("log_level", "Choose log level between verbose/info/warning/error/internal_error.", - cxxopts::value()->default_value("error")); - options.add_options()( - "warm_up", "Specify warm up iterations before benchmark starts.", cxxopts::value()->default_value("2")); - options.add_options()("num_runs", "Minimal number of iterations to run during benchmarking.", - cxxopts::value()->default_value("10")); - options.add_options()("duration", "Minimal duration of iterations to measure in seconds.", - cxxopts::value()->default_value("60")); - options.add_options()("gpu_weights_percent", - "Specify the percentage of weights that reside on GPU (from 0.0 to 1.0). Multiple percentages can be separated " - "by \";\", " - "example: \"0.0;0.5;1.0\".", - cxxopts::value()->default_value("1.0")); - options.add_options()("use_gpu_direct_storage", "Enable GPUDirect Storage (GDS) for loading engine.", - cxxopts::value()->default_value("false")); - - auto result = options.parse(argc, argv); - - if (result.count("help")) - { - std::cout << options.help() << std::endl; - exit(0); - } - - // Argument: Engine directory - if (!result.count("engine_dir")) - { - std::cout << options.help() << std::endl; - TLLM_LOG_ERROR("Please specify engine directory."); - return 1; - } - - // Argument: Batch sizes - std::istringstream ssBatchSizesArg; - ssBatchSizesArg.str(result["batch_size"].as()); - std::vector batchSizes; - for (std::string token; std::getline(ssBatchSizesArg, token, ';');) - { - batchSizes.push_back(std::stoi(token)); - } - - // Argument : Input lengths - std::istringstream ssInLenArg; - ssInLenArg.str(result["input_len"].as()); - std::vector inLens; - for (std::string token; std::getline(ssInLenArg, token, ';');) - { - inLens.push_back(std::stoi(token)); - } - - // Argument: GPU weights percentage - std::istringstream ssGpuPercentArg; - ssGpuPercentArg.str(result["gpu_weights_percent"].as()); - std::vector gpuWeightsPercents; - for (std::string token; std::getline(ssGpuPercentArg, token, ';');) - { - auto gpuWeightsPercent = std::stof(token); - if (gpuWeightsPercent < 0 || gpuWeightsPercent > 1) - { - TLLM_LOG_ERROR( - "--gpu_weights_percent must have percents between 0.0 and 1.0 but got: %f", gpuWeightsPercent); - return 1; - } - gpuWeightsPercents.push_back(gpuWeightsPercent); - } - - // Argument: Log level - auto logger = std::make_shared(); - auto const logLevel = result["log_level"].as(); - if (logLevel == "verbose") - { - logger->setLevel(trt::ILogger::Severity::kVERBOSE); - } - else if (logLevel == "info") - { - logger->setLevel(trt::ILogger::Severity::kINFO); - } - else if (logLevel == "warning") - { - logger->setLevel(trt::ILogger::Severity::kWARNING); - } - else if (logLevel == "error") - { - logger->setLevel(trt::ILogger::Severity::kERROR); - } - else if (logLevel == "internal_error") - { - logger->setLevel(trt::ILogger::Severity::kINTERNAL_ERROR); - } - else - { - TLLM_LOG_ERROR("Unexpected log level: " + logLevel); - return 1; - } - initTrtLlmPlugins(logger.get()); - - try - { - benchmarkBert(result["model"].as(), result["engine_dir"].as(), batchSizes, inLens, - result["use_gpu_direct_storage"].as(), gpuWeightsPercents, logger, result["warm_up"].as(), - result["num_runs"].as(), result["duration"].as()); - } - catch (std::exception const& e) - { - TLLM_LOG_ERROR(e.what()); - return 1; - } - return 0; -} diff --git a/benchmarks/cpp/disaggServerBenchmark.cpp b/benchmarks/cpp/disaggServerBenchmark.cpp deleted file mode 100644 index bc3a7a2659fd..000000000000 --- a/benchmarks/cpp/disaggServerBenchmark.cpp +++ /dev/null @@ -1,1582 +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. - */ - -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/logger.h" -#include "tensorrt_llm/executor/disaggServerUtil.h" -#include "tensorrt_llm/executor/executor.h" -#include "tensorrt_llm/executor/types.h" -#include "tensorrt_llm/plugins/api/tllmPlugin.h" -#include "tensorrt_llm/runtime/common.h" -#include "tensorrt_llm/runtime/gptJsonConfig.h" -#include "tensorrt_llm/runtime/tllmLogger.h" -#include "tensorrt_llm/runtime/utils/mpiUtils.h" -#include "utils/utils.h" - -#include "cxxopts.hpp" -#include - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -using namespace tensorrt_llm::batch_manager; -using namespace tensorrt_llm::runtime; -using namespace tensorrt_llm::benchmark; -using namespace tensorrt_llm::executor::disagg_executor; -namespace texec = tensorrt_llm::executor; -namespace trt = nvinfer1; - -namespace -{ - -class Recorder -{ - -public: - explicit Recorder(std::string opCsvFile, bool streaming = false, int beamWidth = 1, - bool calculateKvCacheTransferTime = true, bool calculateQueueTime = true, std::string responsesJsonFile = "", - bool excludeInputInOutput = false) - : mOpCsvFile(std::move(opCsvFile)) - , mStreaming(streaming) - , mBeamWidth(beamWidth) - , mRespJsonFile(std::move(responsesJsonFile)) - , mOutputHasInput(!excludeInputInOutput) - , mCalculateKVCacheTransferTime(calculateKvCacheTransferTime) - , mCalculateQueueTime(calculateQueueTime) - { - } - - void initialize() - { - mStart = std::chrono::steady_clock::now(); - mSeqLatency.mDataTimes.clear(); - mFtLatency.mDataTimes.clear(); - mGenLatency.mDataTimes.clear(); - mGenFirstTokenLatency.mDataTimes.clear(); - mGenT2TLatency.mDataTimes.clear(); - mGenExcludeFirstIterT2TLatency.mDataTimes.clear(); - mContextReqQueuingLatency.mDataTimes.clear(); - mGenReqQueuingLatency.mDataTimes.clear(); - mGenReqKvCacheTransferLatency.mDataTimes.clear(); - mKvCacheThroughput.mDataTps.clear(); - } - - void finalize() - { - mEnd = std::chrono::steady_clock::now(); - } - - void recordContextQueueLatency(std::vector const& latencies) - { - mContextReqQueuingLatency.mDataTimes.insert( - mContextReqQueuingLatency.mDataTimes.end(), latencies.begin(), latencies.end()); - } - - void recordGenQueueLatency(std::vector const& latencies) - { - mGenReqQueuingLatency.mDataTimes.insert( - mGenReqQueuingLatency.mDataTimes.end(), latencies.begin(), latencies.end()); - } - - void recordKvCacheTransferLatency(std::vector const& latencies) - { - mGenReqKvCacheTransferLatency.mDataTimes.insert( - mGenReqKvCacheTransferLatency.mDataTimes.end(), latencies.begin(), latencies.end()); - } - - void recordKvCacheThroughput(std::vector const& throughputs) - { - mKvCacheThroughput.mDataTps.insert(mKvCacheThroughput.mDataTps.end(), throughputs.begin(), throughputs.end()); - } - - void recordContextStart(SizeType32 inputLength, SizeType32 maxNewTokens, uint64_t requestId, - std::chrono::time_point const& start) - { - mRequestBenchInfos[requestId] = BenchInfo(inputLength, start); - } - - void recordContextEnd(tensorrt_llm::executor::IdType requestId, bool hasError) - { - TLLM_CHECK(mRequestBenchInfos.find(requestId) != mRequestBenchInfos.end()); - mRequestBenchInfos.at(requestId).contextEnd = std::chrono::steady_clock::now(); - mRequestBenchInfos.at(requestId).contextHasError = hasError; - mRequestBenchInfos.at(requestId).decodingIter += 1; - } - - void recordToken(tensorrt_llm::executor::IdType requestId) - { - TLLM_CHECK(mStreaming); - TLLM_CHECK_WITH_INFO(mBeamWidth == 1, "gptManagerBenchmark streaming mode does not support beam > 1"); - TLLM_CHECK(mRequestBenchInfos.find(requestId) != mRequestBenchInfos.end()); - - if (!mRequestBenchInfos.at(requestId).genFirstTokenSeen) - { - mRequestBenchInfos.at(requestId).genFirstTokenTs = std::chrono::steady_clock::now(); - mRequestBenchInfos.at(requestId).genFirstTokenSeen = true; - } - mRequestBenchInfos.at(requestId).decodingIter += 1; - } - - void recordToken(tensorrt_llm::executor::IdType requestId, texec::Response const& response) - { - - TLLM_CHECK(mRequestBenchInfos.find(requestId) != mRequestBenchInfos.end()); - - auto outputTokenIds = response.getResult().outputTokenIds; - - int32_t outputLength = 1; - for (auto const& beam : outputTokenIds) - { - outputLength = std::max(static_cast(beam.size()), outputLength); - } - - mRequestBenchInfos[requestId].outputLength += outputLength; - this->recordToken(requestId); - } - - void recordGenStart( - tensorrt_llm::executor::IdType requestId, std::chrono::time_point const& start) - { - - TLLM_CHECK(mRequestBenchInfos.find(requestId) != mRequestBenchInfos.end()); - mRequestBenchInfos.at(requestId).genStart = start; - } - - void recordGenEnd(tensorrt_llm::executor::IdType requestId, bool hasError) - { - TLLM_CHECK(mRequestBenchInfos.find(requestId) != mRequestBenchInfos.end()); - mRequestBenchInfos.at(requestId).genEnd = std::chrono::steady_clock::now(); - mRequestBenchInfos.at(requestId).genHasError = hasError; - } - - void recordGenEnd(tensorrt_llm::executor::IdType requestId, texec::Response const& response) - { - recordGenEnd(requestId, response.hasError()); - if (!response.hasError()) - { - if (!mStreaming) - { - TLLM_LOG_DEBUG("response.getResult().outputTokenIds"); - auto outputTokenIds = response.getResult().outputTokenIds; - - int32_t outSeqLen = 0; - for (auto const& beam : outputTokenIds) - { - outSeqLen = std::max(static_cast(beam.size()), outSeqLen); - } - if (mOutputHasInput) - { - int inputSeqLen = mRequestBenchInfos[requestId].inputLength; - outSeqLen -= inputSeqLen; - } - mRequestBenchInfos[requestId].outputLength = outSeqLen; - mRequestBenchInfos[requestId].decodingIter = response.getResult().decodingIter; - } - else - { - recordToken(requestId, response); - } - } - } - - void reserve(size_t size) - { - mRequestBenchInfos.reserve(size); - } - - void calculateLatencies() - { - for (auto& reqInfo : mRequestBenchInfos) - { - - reqInfo.second.latency - = std::chrono::duration(reqInfo.second.genEnd - reqInfo.second.contextStart).count(); - reqInfo.second.firstTokenLatency - = std::chrono::duration(reqInfo.second.contextEnd - reqInfo.second.contextStart) - .count(); - reqInfo.second.genLatency - = std::chrono::duration(reqInfo.second.genEnd - reqInfo.second.genStart).count(); - if (mStreaming) - { - reqInfo.second.genFirstTokenLatency - = std::chrono::duration(reqInfo.second.genFirstTokenTs - reqInfo.second.genStart) - .count(); - // include the latency of the second token+ kv Cache transfer latency - - if (reqInfo.second.outputLength > 1) - { - reqInfo.second.avgGenT2TLatency - = std::chrono::duration(reqInfo.second.genEnd - reqInfo.second.genStart) - .count() - / static_cast(reqInfo.second.outputLength - 1); - } - if (reqInfo.second.outputLength > 2) - { - reqInfo.second.avgGenExcludeFirstIterT2TLatency - = std::chrono::duration( - reqInfo.second.genEnd - reqInfo.second.genFirstTokenTs) - .count() - / static_cast(reqInfo.second.outputLength - 2); - } - } - } - } - - void calculateMetrics() - { - - calculateLatencies(); - - int totalOutputTokens{0}; - int totalDecodingIter{0}; - mNumContextErrorSamples = 0; - mNumGenErrorSamples = 0; - mNumSamples = 0; - for (auto const& reqInfo : mRequestBenchInfos) - { - - if (!reqInfo.second.contextHasError && !reqInfo.second.genHasError) - { - mSeqLatency.mDataTimes.push_back(reqInfo.second.latency); - mNumSamples++; - } - if (!reqInfo.second.contextHasError) - { - mFtLatency.mDataTimes.push_back(reqInfo.second.firstTokenLatency); - } - else - { - mNumContextErrorSamples++; - } - if (!reqInfo.second.genHasError) - { - mGenLatency.mDataTimes.push_back(reqInfo.second.genLatency); - totalOutputTokens += reqInfo.second.outputLength; - totalDecodingIter += reqInfo.second.decodingIter; - if (mStreaming) - { - mGenFirstTokenLatency.mDataTimes.push_back(reqInfo.second.genFirstTokenLatency); - - if (reqInfo.second.avgGenT2TLatency.has_value()) - { - mGenT2TLatency.mDataTimes.push_back(reqInfo.second.avgGenT2TLatency.value()); - } - if (reqInfo.second.avgGenExcludeFirstIterT2TLatency.has_value()) - { - mGenExcludeFirstIterT2TLatency.mDataTimes.push_back( - reqInfo.second.avgGenExcludeFirstIterT2TLatency.value()); - } - } - } - else - { - mNumGenErrorSamples++; - } - } - mTotalLatency = std::chrono::duration(mEnd - mStart).count(); - mSeqThroughput = mNumSamples / (mTotalLatency / 1000); - mTokenThroughput = totalOutputTokens / (mTotalLatency / 1000); - mAcceptanceRate = totalDecodingIter - ? (static_cast(totalOutputTokens) / static_cast(totalDecodingIter)) - : 0.0F; - - mSeqLatency.calculate(); - mFtLatency.calculate(); - mGenLatency.calculate(); - if (mStreaming) - { - - mGenFirstTokenLatency.calculate(); - - if (!mGenT2TLatency.mDataTimes.empty()) - { - mGenT2TLatency.calculate(); - std::vector userTokensPerSecond; - userTokensPerSecond.reserve(mGenT2TLatency.mDataTimes.size()); - for (auto const& latency : mGenT2TLatency.mDataTimes) - { - userTokensPerSecond.push_back(1000.F / latency); - } - mAvgUserTokensPerSecond = std::accumulate(userTokensPerSecond.begin(), userTokensPerSecond.end(), 0.F) - / userTokensPerSecond.size(); - } - if (!mGenExcludeFirstIterT2TLatency.mDataTimes.empty()) - { - - mGenExcludeFirstIterT2TLatency.calculate(); - } - } - if (mCalculateQueueTime) - { - - mContextReqQueuingLatency.calculate(); - mGenReqQueuingLatency.calculate(); - } - if (mCalculateKVCacheTransferTime) - { - mGenReqKvCacheTransferLatency.calculate(); - mKvCacheThroughput.calculate(); - } - } - - void report() - { - printf("[BENCHMARK] num_samples %d\n", mNumSamples); - printf("[BENCHMARK] num_context_error_samples %d\n", mNumContextErrorSamples); - printf("[BENCHMARK] num_gen_error_samples %d\n", mNumGenErrorSamples); - printf("\n[BENCHMARK] num_samples %d\n", mNumSamples); - printf("[BENCHMARK] total_latency(ms) %.2f\n", mTotalLatency); - printf("[BENCHMARK] seq_throughput(seq/sec) %.2f\n", mSeqThroughput); - printf("[BENCHMARK] token_throughput(token/sec) %.2f\n", mTokenThroughput); - if (mStreaming) - { - printf("[BENCHMARK] user_tokens_per_second(tokens/sec/user) %.2f\n", mAvgUserTokensPerSecond); - } - printf("[BENCHMARK] avg_acceptance_rate(tokens/decoding steps) %.2f\n\n", mAcceptanceRate); - - mSeqLatency.report(); - mFtLatency.report(); - mGenLatency.report(); - if (mStreaming) - { - mGenFirstTokenLatency.report(); - mGenT2TLatency.report(); - mGenExcludeFirstIterT2TLatency.report(); - } - if (mCalculateQueueTime) - { - mContextReqQueuingLatency.report(); - mGenReqQueuingLatency.report(); - } - if (mCalculateKVCacheTransferTime) - { - mGenReqKvCacheTransferLatency.report(); - mKvCacheThroughput.report(); - } - } - - void writeOpMetricsToCsv() - { - if (!mOpCsvFile.empty()) - { - std::vector headers{"num_samples", "num_context_error_samples", "num_gen_error_samples", - "total_latency(ms)", "seq_throughput(seq/sec)", "token_throughput(token/sec)"}; - auto seqLatencyHeader = mSeqLatency.genHeaders(); - headers.insert(headers.end(), std::make_move_iterator(seqLatencyHeader.begin()), - std::make_move_iterator(seqLatencyHeader.end())); - auto contextLatencyHeader = mFtLatency.genHeaders(); - headers.insert(headers.end(), std::make_move_iterator(contextLatencyHeader.begin()), - std::make_move_iterator(contextLatencyHeader.end())); - auto genLatencyHeader = mGenLatency.genHeaders(); - headers.insert(headers.end(), std::make_move_iterator(genLatencyHeader.begin()), - std::make_move_iterator(genLatencyHeader.end())); - if (mStreaming) - { - auto genFirstTokenHeader = mGenFirstTokenLatency.genHeaders(); - headers.insert(headers.end(), std::make_move_iterator(genFirstTokenHeader.begin()), - std::make_move_iterator(genFirstTokenHeader.end())); - auto genIngterHeader = mGenT2TLatency.genHeaders(); - headers.insert(headers.end(), std::make_move_iterator(genIngterHeader.begin()), - std::make_move_iterator(genIngterHeader.end())); - auto excludeFirstIterIngterHeader = mGenExcludeFirstIterT2TLatency.genHeaders(); - headers.insert(headers.end(), std::make_move_iterator(excludeFirstIterIngterHeader.begin()), - std::make_move_iterator(excludeFirstIterIngterHeader.end())); - headers.push_back("avg_user_tokens_per_second(tokens/sec/user)"); - } - if (mCalculateKVCacheTransferTime) - { - auto genReqKVCacheTransferHeader = mGenReqKvCacheTransferLatency.genHeaders(); - headers.insert(headers.end(), std::make_move_iterator(genReqKVCacheTransferHeader.begin()), - std::make_move_iterator(genReqKVCacheTransferHeader.end())); - auto kvCacheTpHeader = mKvCacheThroughput.genHeaders(); - headers.insert(headers.end(), std::make_move_iterator(kvCacheTpHeader.begin()), - std::make_move_iterator(kvCacheTpHeader.end())); - } - - std::ofstream outputFile(mOpCsvFile); - - if (outputFile.is_open()) - { - for (auto const& header : headers) - { - outputFile << header << ","; - } - outputFile << "\n"; - - outputFile << mNumSamples << "," << mNumContextErrorSamples << "," << mNumGenErrorSamples << "," - << mTotalLatency << "," << mSeqThroughput << "," << mTokenThroughput << "," << mSeqLatency - << "," << mFtLatency << "," << mGenLatency; - if (mStreaming) - { - - outputFile << "," << mGenFirstTokenLatency << "," << mGenT2TLatency << "," - << mGenExcludeFirstIterT2TLatency << "," << mAvgUserTokensPerSecond; - } - if (mCalculateKVCacheTransferTime) - { - outputFile << "," << mGenReqKvCacheTransferLatency << "," << mKvCacheThroughput; - } - - outputFile << "\n"; - } - else - { - std::cerr << "Error opening file '" << mOpCsvFile << "' for writing.\n"; - } - } - } - -private: - struct BenchInfo - { - BenchInfo() = default; - - BenchInfo(int inputLength, std::chrono::time_point start) - : inputLength(inputLength) - , contextStart(start) - { - } - - int inputLength{}; - int outputLength{}; - std::chrono::time_point contextStart; - std::chrono::time_point contextEnd; - std::chrono::time_point genFirstTokenTs; - std::chrono::time_point genStart; - std::chrono::time_point genEnd; - float latency{}; // millisecond - float genLatency{}; - bool contextHasError{false}; - bool genHasError{false}; - float firstTokenLatency{}; - float genFirstTokenLatency{}; - std::optional avgGenT2TLatency; - std::optional avgGenExcludeFirstIterT2TLatency; - bool genFirstTokenSeen{false}; - SizeType32 decodingIter{0}; - }; - - std::unordered_map mRequestBenchInfos; - - std::chrono::time_point mStart; - std::chrono::time_point mEnd; - int mNumSamples{}; - int mNumContextErrorSamples{}; - int mNumGenErrorSamples{}; - float mTotalLatency{}; - float mSeqThroughput{}; - RecordTimeMetric mSeqLatency{"sequence_latency"}; - RecordTimeMetric mFtLatency{"context_latency"}; - RecordTimeMetric mGenLatency{"gen_latency"}; - - RecordTimeMetric mGenFirstTokenLatency{"time_to_gen_first_token"}; - RecordTimeMetric mGenT2TLatency{"inter_token_latency"}; - RecordTimeMetric mGenExcludeFirstIterT2TLatency{"exclude_first_iter_inter_token_latency"}; - RecordTimeMetric mContextReqQueuingLatency{"context_req_queueing_latency"}; - - RecordTimeMetric mGenReqQueuingLatency{"gen_req_queueing_latency"}; - RecordTimeMetric mGenReqKvCacheTransferLatency{"gen_req_kv_cache_transfer_latency"}; - - RecordBwMetric mKvCacheThroughput{"gen_req_kv_cache_transfer_throughput"}; - - float mTokenThroughput{}; - float mAcceptanceRate{}; - - std::string mOpCsvFile; - bool mStreaming; - int mBeamWidth; - std::string mRespJsonFile; - std::unordered_map mResponseTensors; - bool mOutputHasInput; - bool mCalculateKVCacheTransferTime; - bool mCalculateQueueTime; - float mAvgUserTokensPerSecond{}; -}; - -texec::Request makeExecutorContextRequest(Sample const& sample, SizeType32 const& beamWidth, - std::optional const& eosId, std::optional const& padId, bool streaming = false, - bool const& returnContextLogits = false, bool const& returnGenerationLogits = false, - std::optional const& loraConfig = std::nullopt, - std::optional const& lookaheadConfig = std::nullopt, - std::optional const& encoderInputTokenIds = std::nullopt) -{ - auto samplingConfig = texec::SamplingConfig{beamWidth}; - auto outputConfig = texec::OutputConfig{false, returnContextLogits, returnGenerationLogits, false}; - auto request - = texec::Request(sample.inputIds, sample.outputLen, streaming, samplingConfig, outputConfig, eosId, padId, - std::nullopt, // positionIds - std::nullopt, // badWords - std::nullopt, // stopWords - std::nullopt, // embeddingBias - std::nullopt, // speculativeDecoding - std::nullopt, // pTuning - std::nullopt, // multimodalInput - std::nullopt, // multimodalEmbedding - std::nullopt, // mRopeConfig - loraConfig, // loraConfig - lookaheadConfig, // lookaheadConfig - std::nullopt, // kvCacheRetentionConfig - std::nullopt, // logitsPostProcessorName - std::nullopt, // logitsPostProcessor - encoderInputTokenIds.has_value() ? encoderInputTokenIds : std::nullopt, - std::nullopt); // cacheSalt - request.setRequestType(tensorrt_llm::executor::RequestType::REQUEST_TYPE_CONTEXT_ONLY); - return request; -} - -class DisaggExecutorServer -{ - -public: - DisaggExecutorServer(std::vector const& contextEnginePaths, - std::vector const& genEnginePaths, - std::optional>> const& deviceIdsForInstance, int32_t maxBeamWidth, - texec::CapacitySchedulerPolicy capacitySchedulerPolicy, BenchmarkParams const& benchmarkParams, - std::shared_ptr recorder, std::chrono::milliseconds waitSleep, bool logIterationData, - bool hasContextAwaitThreads, bool hasGenAwaitThreads) - : mRecorder(std::move(recorder)) - , mWaitSleep(waitSleep) - , mConcurrency(benchmarkParams.concurrency) - , mShutdown(false) - , mLogIterationData(logIterationData) - , mEnableCollectKvCacheTransferTime(benchmarkParams.enableCollectkvCacheTransferTime) - , mEnableCollectIterStats(benchmarkParams.enableCollectIterStats) - { - - int worldRank = tensorrt_llm::mpi::MpiComm::world().getRank(); - int worldSize = tensorrt_llm::mpi::MpiComm::world().getSize(); - mIsOrchestrator = (worldRank == 0); - auto contextNum = contextEnginePaths.size(); - auto genNum = genEnginePaths.size(); - int deviceCount = -1; - TLLM_CUDA_CHECK(cudaGetDeviceCount(&deviceCount)); - - std::vector> instances; - auto instanceNum = genNum + contextNum; - if (worldRank == 0) - { - TLLM_LOG_INFO("context enigne num :%d gen enigne num:%d", contextNum, genNum); - } - - int startRank = 0; - std::vector ctxExecutorConfigs; - std::vector genExecutorConfigs; - for (auto in = 0; in < instanceNum; in++) - { - auto&& enginePath = in < contextNum ? contextEnginePaths.at(in) : genEnginePaths.at(in - contextNum); - auto decoderJsonConfig = tensorrt_llm::runtime::GptJsonConfig::parse(enginePath / "config.json"); - size_t instanceRanks = decoderJsonConfig.getWorldSize(); - std::vector participateRank(instanceRanks); - std::vector deviceIds; - if (deviceIdsForInstance.has_value()) - { - deviceIds = deviceIdsForInstance.value().at(in); - } - for (int i = 0; i < instanceRanks; i++) - { - startRank++; - participateRank.at(i) = startRank; - if (!deviceIdsForInstance.has_value()) - { - deviceIds.push_back((startRank - 1) % deviceCount); - } - } - texec::DynamicBatchConfig dynamicBatchConfig(benchmarkParams.enableBatchSizeTuning); - texec::SchedulerConfig schedulerConfig(capacitySchedulerPolicy, std::nullopt, dynamicBatchConfig); - texec::KvCacheConfig kvCacheConfig(benchmarkParams.enableBlockReuse, - benchmarkParams.maxTokensInPagedKvCache, benchmarkParams.maxAttentionWindowVec, - benchmarkParams.sinkTokenLength, benchmarkParams.freeGpuMemoryFractions.at(in), - benchmarkParams.kvHostCacheSize); - texec::ExtendedRuntimePerfKnobConfig extendedRuntimePerfKnobConfig(benchmarkParams.multiBlockMode, - benchmarkParams.enableContextFMHAFP32Acc, benchmarkParams.cudaGraphMode, - benchmarkParams.cudaGraphCacheSize); - texec::ExecutorConfig executorConfig(maxBeamWidth, schedulerConfig, kvCacheConfig, - benchmarkParams.enableChunekedContextVec.at(in).value_or(false)); - executorConfig.setGpuWeightsPercent(benchmarkParams.gpuWeightsPercent); - texec::OrchestratorConfig orchestratorConfig{mIsOrchestrator, "", nullptr, false}; - texec::ParallelConfig parallelConfig{tensorrt_llm::executor::CommunicationType::kMPI, - tensorrt_llm::executor::CommunicationMode::kORCHESTRATOR, deviceIds, participateRank, - orchestratorConfig}; - executorConfig.setParallelConfig(parallelConfig); - if (benchmarkParams.maxBatchSizes.at(in)) - { - executorConfig.setMaxBatchSize(benchmarkParams.maxBatchSizes.at(in).value()); - } - if (benchmarkParams.maxNumTokensVec.at(in)) - { - executorConfig.setMaxNumTokens(benchmarkParams.maxNumTokensVec.at(in).value()); - } - - executorConfig.setDecodingConfig( - texec::DecodingConfig(benchmarkParams.medusaChoices.has_value() ? texec::DecodingMode::Medusa() - : benchmarkParams.executorLookaheadConfig.has_value() ? texec::DecodingMode::Lookahead() - : texec::DecodingMode::Auto(), - benchmarkParams.executorLookaheadConfig, benchmarkParams.medusaChoices)); - executorConfig.setExtendedRuntimePerfKnobConfig(extendedRuntimePerfKnobConfig); - executorConfig.setCacheTransceiverConfig( - texec::CacheTransceiverConfig(texec::CacheTransceiverConfig::BackendType::DEFAULT)); - constexpr int maxIterationsForRequestStats = 1000; - if (mEnableCollectKvCacheTransferTime) - { - executorConfig.setRequestStatsMaxIterations(maxIterationsForRequestStats); - } - if (!benchmarkParams.enableCollectIterStats) - { - executorConfig.setIterStatsMaxIterations(0); - } - - if (in < contextNum) - { - ctxExecutorConfigs.push_back(executorConfig); - } - else - { - genExecutorConfigs.push_back(executorConfig); - } - } - - mDisaggExecutor = std::make_unique(contextEnginePaths, genEnginePaths, - ctxExecutorConfigs, genExecutorConfigs, hasContextAwaitThreads, hasGenAwaitThreads); - - if (mIsOrchestrator) - { - - if (mEnableCollectIterStats || mEnableCollectKvCacheTransferTime) - { - mCollectStatsThread = std::thread(&DisaggExecutorServer::collectStats, this); - } - } - tensorrt_llm::mpi::MpiComm::world().barrier(); - } - - std::vector enqueueContext(std::vector const& requests, - std::optional selectContextId = std::nullopt, bool warmup = false, bool batch = false) - { - std::vector inputLengths; - std::vector maxNewTokens; - if (!warmup) - { - for (auto const& request : requests) - { - inputLengths.push_back(static_cast(request.getInputTokenIds().size())); - maxNewTokens.push_back(request.getMaxTokens()); - } - } - auto const start = std::chrono::steady_clock::now(); - std::vector globalReqIds - = mDisaggExecutor->enqueueContext(requests, selectContextId, batch); - if (!warmup) - { - for (size_t i = 0; i < requests.size(); ++i) - { - mRecorder->recordContextStart(inputLengths.at(i), maxNewTokens.at(i), globalReqIds.at(i), start); - } - } - mNumContextActive += requests.size(); - return globalReqIds; - } - - void enqueueGeneration(std::vector const& requests, - std::vector const& globalRequestIds, - std::optional selectGenIdx = std::nullopt, bool warmup = false, bool batch = false) - { - TLLM_CHECK(globalRequestIds.size() == requests.size()); - auto const start = std::chrono::steady_clock::now(); - mDisaggExecutor->enqueueGeneration(requests, globalRequestIds, selectGenIdx, batch); - if (!warmup) - { - for (int i = 0; i < requests.size(); i++) - { - - mRecorder->recordGenStart(globalRequestIds.at(i), start); - } - } - mNumGenActive += requests.size(); - } - - std::vector waitForContextResponse(SizeType32 numRequests, bool warmup = false) - { - std::vector ret; - ret.reserve(numRequests); - while ((mNumContextActive != 0) || (mNumContextFinished < numRequests)) - { - auto responses = mDisaggExecutor->awaitContextResponses(mWaitSleep); - for (auto&& response : responses) - { - TLLM_CHECK(response.response.getResult().isFinal); - if (response.response.getResult().isFinal) - { - mNumContextActive--; - mNumContextFinished++; - } - if (!warmup) - { - mRecorder->recordContextEnd(response.gid, response.response.hasError()); - } - ret.emplace_back(std::move(response)); - } - } - return ret; - } - - void waitForGenResponse(SizeType32 numRequests, bool warmup = false) - { - while (mNumGenActive > 0 || (mNumGenFinished < numRequests)) - { - auto responses = mDisaggExecutor->awaitGenerationResponses(mWaitSleep); - for (auto&& response : responses) - { - if (response.response.getResult().isFinal) - { - mNumGenActive--; - mNumGenFinished++; - - if (!warmup) - { - mRecorder->recordGenEnd(response.gid, response.response); - } - } - else - { - // streaming - if (!warmup && !response.response.hasError()) - { - mRecorder->recordToken(response.gid, response.response); - } - } - } - } - } - - bool canEnqueue(int numSentRequests) const - { - return mIsOrchestrator && (!mConcurrency || (numSentRequests - mNumGenFinished < mConcurrency)); - } - - ~DisaggExecutorServer() - { - mShutdown = true; - if (mCollectStatsThread.joinable()) - { - mCollectStatsThread.join(); - } - } - - void resetNumFinished() - { - mNumContextFinished = 0; - mNumGenFinished = 0; - } - - void resetNumActive() - { - mNumContextActive = 0; - mNumGenActive = 0; - } - - void collectStats() const - { - while (!mShutdown) - { - std::vector> contextStats; - std::vector> generationStats; - std::vector> - generationRequestStatsPerIteration; - contextStats.reserve(mDisaggExecutor->getContextExecutors().size()); - for (auto&& executor : mDisaggExecutor->getContextExecutors()) - { - if (executor->canEnqueueRequests()) - { - contextStats.emplace_back(executor->getLatestIterationStats()); - } - } - generationStats.reserve(mDisaggExecutor->getGenExecutors().size()); - for (auto&& executor : mDisaggExecutor->getGenExecutors()) - { - if (executor->canEnqueueRequests()) - { - if (mEnableCollectIterStats) - { - generationStats.emplace_back(executor->getLatestIterationStats()); - } - if (mEnableCollectKvCacheTransferTime) - { - - generationRequestStatsPerIteration.emplace_back(executor->getLatestRequestStats()); - } - } - } - if (mEnableCollectIterStats) - { - for (std::size_t i = 0; i < contextStats.size(); i++) - { - auto const& iterStats = contextStats.at(i); - for (auto const& stat : iterStats) - { - SizeType32 numNewActiveRequests = stat.numNewActiveRequests; - if (numNewActiveRequests > 0) - { - auto avgQueueingTime - = static_cast(stat.newActiveRequestsQueueLatencyMS / numNewActiveRequests); - std::vector requestsQueueLatencyMS(numNewActiveRequests, avgQueueingTime); - mRecorder->recordContextQueueLatency(requestsQueueLatencyMS); - } - if (mLogIterationData) - { - TLLM_LOG_INFO( - "ctx_id %d, ctx_stat: %s", i, texec::JsonSerialization::toJsonStr(stat).c_str()); - } - } - } - - for (std::size_t i = 0; i < generationStats.size(); i++) - { - auto const& iterStats = generationStats.at(i); - for (auto const& stat : iterStats) - { - SizeType32 numNewActiveRequests = stat.numNewActiveRequests; - if (numNewActiveRequests > 0) - { - float avgQueueingTime - = static_cast(stat.newActiveRequestsQueueLatencyMS / numNewActiveRequests); - std::vector requestsQueueLatencyMS(numNewActiveRequests, avgQueueingTime); - mRecorder->recordGenQueueLatency(requestsQueueLatencyMS); - } - if (mLogIterationData) - { - TLLM_LOG_INFO( - "gen_id %d, gen_stat: %s", i, texec::JsonSerialization::toJsonStr(stat).c_str()); - } - } - } - } - - if (mEnableCollectKvCacheTransferTime) - { - for (std::size_t i = 0; i < generationRequestStatsPerIteration.size(); i++) - { - auto const& stats = generationRequestStatsPerIteration.at(i); - for (auto const& stat : stats) - { - std::vector kvCacheTransferMs; - std::vector kvCacheThroughput; - for (auto const& requestStat : stat.requestStats) - { - if (requestStat.stage == tensorrt_llm::executor::RequestStage::kGENERATION_COMPLETE) - { - kvCacheTransferMs.push_back( - static_cast(requestStat.disServingStats->kvCacheTransferMS)); - kvCacheThroughput.push_back(static_cast(requestStat.disServingStats->kvCacheSize) - * 8 / (static_cast(requestStat.disServingStats->kvCacheTransferMS) / 1000) - / 1e9f); - } - } - if (kvCacheTransferMs.size() > 0) - { - mRecorder->recordKvCacheTransferLatency(kvCacheTransferMs); - } - if (kvCacheThroughput.size() > 0) - { - mRecorder->recordKvCacheThroughput(kvCacheThroughput); - } - if (mLogIterationData) - { - TLLM_LOG_INFO( - "gen_id %d, gen_req_stat: %s", i, texec::JsonSerialization::toJsonStr(stat).c_str()); - } - } - } - } - auto const waitSleep = std::chrono::milliseconds(50); - std::this_thread::sleep_for(waitSleep); - } - } - - std::unique_ptr const& getDisaggExecutor() const noexcept - { - return mDisaggExecutor; - } - -private: - std::unique_ptr mDisaggExecutor; - - std::atomic mShutdown{false}; - bool mIsOrchestrator{false}; - - std::shared_ptr mRecorder; - std::chrono::milliseconds mWaitSleep; - std::optional mConcurrency; - bool mLogIterationData{false}; - bool const mEnableCollectKvCacheTransferTime; - bool const mEnableCollectIterStats; - std::thread mCollectStatsThread; - std::atomic mNumGenFinished{0}; - std::atomic mNumContextFinished{0}; - std::atomic mNumGenActive{0}; - std::atomic mNumContextActive{0}; -}; - -} // namespace - -void benchmark(std::vector const& contextEngineDirs, - std::vector const& generationEngineDirs, - std::optional>> const& deviceIdsForInstances, std::string const& datasetPath, - std::string const& opCsvFile, int maxNumSamples, int beamWidth, int warmUp, std::optional const& eosId, - std::optional const& padId, BenchmarkParams const& benchmarkParams, - texec::CapacitySchedulerPolicy capacitySchedulerPolicy, std::chrono::milliseconds waitSleep, - bool returnContextLogits, bool returnGenerationLogits, std::optional const staticEmulatedBatchSize, - bool logIterationData, std::optional const maxPromptLen, bool hasContextAwait, bool hasGenAwait) -{ - - auto const& world = tensorrt_llm::mpi::MpiComm::world(); - auto worldRank = world.getRank(); - - // Load dataset - auto const samples = parseWorkloadJson(datasetPath, maxNumSamples, maxPromptLen); - auto const numSamples = samples.size(); - auto recorder = std::make_shared(opCsvFile, benchmarkParams.streaming, beamWidth, - benchmarkParams.enableCollectkvCacheTransferTime, benchmarkParams.enableCollectIterStats); - auto disaggExecutor = std::make_shared(contextEngineDirs, generationEngineDirs, - deviceIdsForInstances, beamWidth, capacitySchedulerPolicy, benchmarkParams, recorder, waitSleep, - logIterationData, hasContextAwait, hasGenAwait); - constexpr size_t numMap = 8; - std::vector> gidToRequestMaps( - numMap); - std::vector mtxForMaps(numMap); - - auto fillRequestMap = [&](std::vector const& reqIds, - std::vector&& requests) - { - TLLM_CHECK(reqIds.size() == requests.size()); - for (size_t i = 0; i < reqIds.size(); i++) - { - - size_t mapIdx = reqIds[i] % numMap; - std::scoped_lock lock(mtxForMaps[mapIdx]); - gidToRequestMaps.at(mapIdx).emplace(reqIds[i], std::move(requests[i])); - } - }; - - auto makeGenRequest = [&](std::vector&& contextResponse) - { - std::vector gids; - gids.reserve(contextResponse.size()); - std::vector genRequest; - genRequest.reserve(contextResponse.size()); - for (auto&& ctxResponse : contextResponse) - { - gids.emplace_back(ctxResponse.gid); - size_t mapIdx = ctxResponse.gid % numMap; - - std::unique_lock lock(mtxForMaps[mapIdx]); - TLLM_CHECK(gidToRequestMaps.at(mapIdx).find(ctxResponse.gid) != gidToRequestMaps.at(mapIdx).end()); - auto ctxRequest = std::move(gidToRequestMaps.at(mapIdx).at(ctxResponse.gid)); - gidToRequestMaps.at(mapIdx).erase(ctxResponse.gid); - lock.unlock(); - ctxRequest.setRequestType(tensorrt_llm::executor::RequestType::REQUEST_TYPE_GENERATION_ONLY); - ctxRequest.setContextPhaseParams(ctxResponse.response.getResult().contextPhaseParams.value()); - genRequest.emplace_back(std::move(ctxRequest)); - } - return std::make_pair(genRequest, gids); - }; - if (worldRank == 0) - { - { // warmup - TLLM_LOG_INFO("Warmup start"); - - size_t contextNum = contextEngineDirs.size(); - size_t generationNum = generationEngineDirs.size(); - for (auto con = 0; con < contextNum; con++) - { - for (auto gen = 0; gen < generationNum; gen++) - { - std::vector contextRequests; - contextRequests.reserve(warmUp); - for (int i = 0; i < warmUp; ++i) - { - contextRequests.emplace_back(makeExecutorContextRequest(samples[0], beamWidth, eosId, padId, - benchmarkParams.streaming, returnContextLogits, returnGenerationLogits, std::nullopt, - benchmarkParams.requestLookaheadConfig)); - } - auto reqIds = disaggExecutor->enqueueContext(contextRequests, con, true); - fillRequestMap(reqIds, std::move(contextRequests)); - auto contextResponse = disaggExecutor->waitForContextResponse(warmUp, true); - auto&& [genRequests, gids] = makeGenRequest(std::move(contextResponse)); - disaggExecutor->enqueueGeneration(genRequests, gids, gen, true); - disaggExecutor->waitForGenResponse(warmUp, true); - disaggExecutor->resetNumFinished(); - disaggExecutor->resetNumActive(); - } - } - - auto const warmUpWaitSleep = std::chrono::milliseconds(50); - std::this_thread::sleep_for(warmUpWaitSleep); - TLLM_LOG_INFO("Warmup done"); - } - - { - - auto timeDelays = computeTimeDelays(benchmarkParams, numSamples - 1); - - std::vector contextRequests; - - for (std::size_t i = 0; i < numSamples; ++i) - { - std::optional loraConfig = std::nullopt; - contextRequests.emplace_back(makeExecutorContextRequest(samples[i], beamWidth, eosId, padId, - benchmarkParams.streaming, returnContextLogits, returnGenerationLogits, loraConfig, - benchmarkParams.requestLookaheadConfig)); - } - - bool const hasDelay - = std::any_of(timeDelays.begin(), timeDelays.end(), [](auto const& delay) { return delay > 0.0; }); - disaggExecutor->resetNumFinished(); - disaggExecutor->resetNumActive(); - - recorder->reserve(numSamples); - recorder->initialize(); - if (!staticEmulatedBatchSize) - { - - std::thread waitContextResponseAndEnqueGenThread{[&]() - { - auto numRequest = numSamples; - while (numRequest > 0) - { - auto contextResponseWithIds - = disaggExecutor->getDisaggExecutor()->awaitContextResponses(waitSleep); - if (contextResponseWithIds.empty()) - { - continue; - } - for (auto&& contextResponseWithId : contextResponseWithIds) - { - recorder->recordContextEnd( - contextResponseWithId.gid, contextResponseWithId.response.hasError()); - } - numRequest -= contextResponseWithIds.size(); - auto&& [genReqeust, genGids] = makeGenRequest(std::move(contextResponseWithIds)); - disaggExecutor->enqueueGeneration(genReqeust, genGids); - } - }}; - - std::thread waitGenResponseThread{[&]() { disaggExecutor->waitForGenResponse(numSamples); }}; - int numSentRequests = 0; - while (numSentRequests < numSamples) - { - - if (disaggExecutor->canEnqueue(numSentRequests)) - { - auto gids = disaggExecutor->enqueueContext({contextRequests.at(numSentRequests)}); - fillRequestMap(gids, {contextRequests.at(numSentRequests)}); - - if (hasDelay && numSentRequests < numSamples - 1) - { - std::this_thread::sleep_for( - std::chrono::milliseconds(static_cast(timeDelays.at(numSentRequests) * 1000))); - } - numSentRequests += 1; - } - } - waitContextResponseAndEnqueGenThread.join(); - waitGenResponseThread.join(); - } - else - { - TLLM_CHECK_WITH_INFO( - !hasDelay, "Executor benchmark doesn't support delays with emulated static batch sizes"); - auto numRequests = contextRequests.size(); - int maxBatchSize = staticEmulatedBatchSize.value(); - for (int req = 0; req < numRequests; req += maxBatchSize) - { - auto batchSize = std::min(static_cast(maxBatchSize), numRequests - req); - - std::vector requestsBatch(std::make_move_iterator(contextRequests.begin() + req), - std::make_move_iterator(contextRequests.begin() + req + static_cast(batchSize))); - // Enqueue in batches - - auto reqIds = disaggExecutor->enqueueContext(requestsBatch); - fillRequestMap(reqIds, std::move(requestsBatch)); - auto contextResponse = disaggExecutor->waitForContextResponse(static_cast(batchSize)); - auto&& [genRequests, genReqIds] = makeGenRequest(std::move(contextResponse)); - disaggExecutor->enqueueGeneration(genRequests, genReqIds); - disaggExecutor->waitForGenResponse(static_cast(batchSize)); - - // Wait for current batch to be done - } - } - } - recorder->finalize(); - // sleep for collect stats - if (benchmarkParams.enableCollectIterStats || benchmarkParams.enableCollectkvCacheTransferTime) - { - auto const collectWaitSleep = std::chrono::milliseconds(50); - std::this_thread::sleep_for(collectWaitSleep); - } - recorder->calculateMetrics(); - recorder->report(); - recorder->writeOpMetricsToCsv(); - } -} - -int main(int argc, char* argv[]) - -{ - cxxopts::Options options("TensorRT LLM DisaggServer Benchmark"); - options.add_options()("h,help", "Print usage"); - options.add_options()("context_engine_dirs", "Directories that store context engines,separator is a ,", - cxxopts::value>()); - options.add_options()("generation_engine_dirs", "Directories that store generation engines,separator is a , ", - cxxopts::value>()); - options.add_options()("device_ids_for_instances", - "device ids for each instances , example: \"[[0,1],[2,3],[4,5,6,7]]\" ", cxxopts::value()); - options.add_options()("dataset", "Dataset that is used for benchmarking BatchManager.", - cxxopts::value()->default_value("")); - options.add_options()( - "output_csv", "Write output metrics to CSV", cxxopts::value()->default_value("")); - options.add_options()("max_num_samples", "maximum number of samples to use from dataset/generate", - cxxopts::value()->default_value("100000")); - options.add_options()( - "beam_width", "Specify beam width you want to benchmark.", cxxopts::value()->default_value("1")); - options.add_options()( - "warm_up", "Specify warm up iterations before benchmark starts.", cxxopts::value()->default_value("2")); - options.add_options()( - "eos_id", "Specify the end-of-sequence token id.", cxxopts::value()->default_value("-1")); - options.add_options()("pad_id", "Specify the padding token id.", cxxopts::value()); - options.add_options()("max_tokens_in_paged_kvcache", "Max tokens in paged K-V Cache.", cxxopts::value()); - options.add_options()( - "max_attention_window", "Max KV cache length per sequence", cxxopts::value>()); - options.add_options()("sink_token_len", "Sink token length in kv cache per sequence.", cxxopts::value()); - options.add_options()( - "random_seed", "integer random seed for exponential time delays.", cxxopts::value()->default_value("420")); - options.add_options()("kv_cache_free_gpu_mem_fractions", "K-V Cache Free Gpu Mem Fraction,each for per instance", - cxxopts::value>()); - options.add_options()("request_rate", - "request rate in reqs/sec. Skipping this arg or negative value will trigger offline/0-delay.", - cxxopts::value()); - options.add_options()("concurrency", "Concurrent number of connections with the server.", cxxopts::value()); - options.add_options()("max_batch_sizes", "The max runtime batch size when benchmarking, each for per instance", - cxxopts::value>()); - options.add_options()("max_num_tokens_per_instance", - "The max runtime number of tokens per batch when benchmarking, each for per instance", - cxxopts::value>()); - options.add_options()( - "enable_batch_size_tuning", "Dynamic tuning of batch size", cxxopts::value()->default_value("false")); - options.add_options()("enable_exp_delays", "Enables exponential delay distr to mimic real world request arrival", - cxxopts::value()->default_value("false")); - options.add_options()("streaming", "Operate in streaming mode", cxxopts::value()->default_value("false")); - options.add_options()( - "enable_kv_cache_reuse", "Enables the KV cache reuse.", cxxopts::value()->default_value("false")); - options.add_options()("enable_chunked_context_per_instance", "Whether to enable context chunking for per instance", - cxxopts::value>()->default_value("false")); - options.add_options()( - "return_context_logits", "Whether to return context logits.", cxxopts::value()->default_value("false")); - options.add_options()("return_generation_logits", "Whether to return generation logits.", - cxxopts::value()->default_value("false")); - - options.add_options()("scheduler_policy", - "Choose scheduler policy between max_utilization/guaranteed_no_evict/static_batch.", - cxxopts::value()->default_value("guaranteed_no_evict")); - - options.add_options()("static_emulated_batch_size", - "Emulate static batching performance with the provided batch size.", cxxopts::value()); - options.add_options()("log_level", "Choose log level between verbose/info/warning/error/internal_error.", - cxxopts::value()->default_value("error")); - options.add_options()("log_iteration_data", "On each decoder iteration, print batch state metadata.", - cxxopts::value()->default_value("false")); - options.add_options()("wait_sleep", "Specify how many milliseconds to sleep each iteration of waitForEmpty loop.", - cxxopts::value()->default_value("25")); - options.add_options()("kv_host_cache_bytes", - "Size of secondary memory pool used for offloading kv cache blocks (in bytes).", - cxxopts::value()->default_value("0")); - options.add_options()( - "max_prompt_len", "Truncate all prompts from dataset to the length specified.", cxxopts::value()); - options.add_options()("gpu_weights_percent", - "Specify the percentage of weights that reside on GPU (from 0.0 to 1.0).", - cxxopts::value()->default_value("1.0")); - options.add_options()( - "medusa_choices", "Medusa choices in the format of [[0], [0, 1], [0, 0, 1]]", cxxopts::value()); - options.add_options()("multi_block_mode", - "Distribute the work across multiple CUDA thread-blocks on the GPU for masked MHA kernel", - cxxopts::value()->default_value("true")); - options.add_options()("cuda_graph_mode", "When enabled, inference is executed with cuda graph.", - cxxopts::value()->default_value("false")); - options.add_options()("cuda_graph_cache_size", - "Specify how many cuda graphs are cached in the runtime. Larger cache gives better perf, but consumes more GPU " - "memory.", - cxxopts::value()->default_value("0")); - options.add_options()("enable_context_fmha_fp32_acc", "Enable FMHA runner FP32 accumulation", - cxxopts::value()->default_value("false")); - options.add_options()("executor_lookahead_config", - "lookahead config in the format of [max_window_size, max_ngram_size, max_verification_set_size]", - cxxopts::value()); - options.add_options()("request_lookahead_config", - "lookahead config in the format of [max_window_size, max_ngram_size, max_verification_set_size], and each <= " - "executor lookahead config", - cxxopts::value()); - options.add_options()("context_await", "When enabled, will has a thread to await context response.", - cxxopts::value()->default_value("true")); - options.add_options()("gen_await", "When enabled,will has a thread to await gen response.", - cxxopts::value()->default_value("true")); - options.add_options()("enable_collect_kvcache_transfer_time", "When enabled, will collect kvcache transfer time.", - cxxopts::value()->default_value("false")); - options.add_options()("enable_collect_iter_stats", "When enabled, will collect iteration stats.", - cxxopts::value()->default_value("false")); - - auto result = options.parse(argc, argv); - - if ((result.count("context_engine_dirs") == 0) || (result.count("generation_engine_dirs") == 0)) - { - std::cout << options.help() << std::endl; - TLLM_LOG_ERROR("Please specify context engine and generation engine directory."); - return 1; - } - // Argument: Log level - auto logger = std::make_shared(); - auto const logLevel = result["log_level"].as(); - if (logLevel == "verbose") - { - logger->setLevel(trt::ILogger::Severity::kVERBOSE); - } - else if (logLevel == "info") - { - logger->setLevel(trt::ILogger::Severity::kINFO); - } - else if (logLevel == "warning") - { - logger->setLevel(trt::ILogger::Severity::kWARNING); - } - else if (logLevel == "error") - { - logger->setLevel(trt::ILogger::Severity::kERROR); - } - else if (logLevel == "internal_error") - { - logger->setLevel(trt::ILogger::Severity::kINTERNAL_ERROR); - } - else - { - TLLM_LOG_ERROR("Unexpected log level: " + logLevel); - return 1; - } - - initTrtLlmPlugins(logger.get()); - - // Argument: Dataset - auto const datasetPath = result["dataset"].as(); - auto const maxNumSamples = result["max_num_samples"].as(); - - // Argument: Output metrics CSV - auto const opCsvFile = result["output_csv"].as(); - - // Argument: beam width - auto const beamWidth = result["beam_width"].as(); - TLLM_CHECK_WITH_INFO(beamWidth == 1, "Currently only support beamWidth=1"); - // Argument: wait_sleep - auto const waitSleep = std::chrono::milliseconds(result["wait_sleep"].as()); - auto const hasContextAwait = result["context_await"].as(); - auto const hasGenAwait = result["gen_await"].as(); - BenchmarkParams benchmarkParams; - benchmarkParams.enableCollectkvCacheTransferTime = result["enable_collect_kvcache_transfer_time"].as(); - benchmarkParams.enableCollectIterStats = result["enable_collect_iter_stats"].as(); - - std::vector contextEngineDirs = result["context_engine_dirs"].as>(); - std::vector generationEngineDirs = result["generation_engine_dirs"].as>(); - if (tensorrt_llm::mpi::MpiComm::world().getRank() == 0) - { - std::string contextEngineStrings; - for (auto&& contextEngineDir : contextEngineDirs) - { - contextEngineStrings += contextEngineDir + ","; - } - std::string generationEnginesStrings; - for (auto&& genEngineDir : generationEngineDirs) - { - generationEnginesStrings += genEngineDir + ","; - } - TLLM_LOG_INFO( - "Will Launch benchmark with %d context engines and %d generation engines. Context Engines:%s ; Generation " - "Engines:%s ;", - contextEngineDirs.size(), generationEngineDirs.size(), contextEngineStrings.c_str(), - generationEnginesStrings.c_str()); - } - std::vector contextEnigePaths; - std::vector generationEnginePaths; - - contextEnigePaths.reserve(contextEngineDirs.size()); - - for (auto& contextEngineDir : contextEngineDirs) - { - - contextEnigePaths.emplace_back(contextEngineDir); - } - generationEnginePaths.reserve(generationEngineDirs.size()); - for (auto& genEngineDir : generationEngineDirs) - { - - generationEnginePaths.emplace_back(genEngineDir); - } - - int const instanceNum = contextEngineDirs.size() + generationEngineDirs.size(); - // Argument: Max tokens in paged K-V Cache - if (result.count("max_tokens_in_paged_kvcache")) - { - benchmarkParams.maxTokensInPagedKvCache = result["max_tokens_in_paged_kvcache"].as(); - } - - // Argument: Max KV cache length - if (result.count("max_attention_window")) - { - benchmarkParams.maxAttentionWindowVec = result["max_attention_window"].as>(); - } - - // Argument: Sink token length - if (result.count("sink_token_len")) - { - benchmarkParams.sinkTokenLength = result["sink_token_len"].as(); - } - - if (result.count("random_seed")) - { - benchmarkParams.randomSeed = result["random_seed"].as(); - } - - // Argument: K-V Cache Free Gpu Mem Fraction - benchmarkParams.freeGpuMemoryFractions.resize(instanceNum); - if (result.count("kv_cache_free_gpu_mem_fractions")) - { - auto fractions = result["kv_cache_free_gpu_mem_fractions"].as>(); - TLLM_CHECK_WITH_INFO(fractions.size() == instanceNum || fractions.size() == 1, - "the number of fraction should be equal to the number of instances or equal to 1"); - for (int i = 0; i < instanceNum; i++) - { - benchmarkParams.freeGpuMemoryFractions.at(i) = fractions.size() == 1 ? fractions[0] : fractions[i]; - } - } - - // Argument: Enable dynamic tuning of batch size - benchmarkParams.enableBatchSizeTuning = result["enable_batch_size_tuning"].as(); - - // Argument: Enable KV cache reuse - benchmarkParams.enableBlockReuse = result["enable_kv_cache_reuse"].as(); - - // Argument: streaming - benchmarkParams.streaming = result["streaming"].as(); - - TLLM_CHECK_WITH_INFO(!(result.count("request_rate") && result.count("concurrency")), - "request_rate and concurrency cannot be specified at the same time."); - - // Argument: request rate - if (result.count("request_rate")) - { - benchmarkParams.requestRate = result["request_rate"].as(); - } - - // Argument: concurrency - if (result.count("concurrency")) - { - benchmarkParams.concurrency = result["concurrency"].as(); - } - - // Argument: max_batch_sizes - benchmarkParams.maxBatchSizes.resize(instanceNum); - if (result.count("max_batch_sizes")) - { - auto batchSizes = result["max_batch_sizes"].as>(); - TLLM_CHECK_WITH_INFO(batchSizes.size() == instanceNum || batchSizes.size() == 1, - "the number of batch size should be equal to the number of instances or equal to 1"); - for (int i = 0; i < instanceNum; i++) - { - benchmarkParams.maxBatchSizes.at(i) = batchSizes.size() == 1 ? batchSizes[0] : batchSizes[i]; - } - } - - // Argument: max_num_tokens_per_instance - benchmarkParams.maxNumTokensVec.resize(instanceNum); - if (result.count("max_num_tokens_per_instance")) - { - auto maxNumTokensVec = result["max_num_tokens_per_instance"].as>(); - TLLM_CHECK_WITH_INFO(maxNumTokensVec.size() == instanceNum || maxNumTokensVec.size() == 1, - "the number of max_num_tokens should be equal to the number of instances or equal to 1"); - for (int i = 0; i < instanceNum; i++) - { - benchmarkParams.maxNumTokensVec.at(i) - = maxNumTokensVec.size() == 1 ? maxNumTokensVec[0] : maxNumTokensVec[i]; - } - } - - benchmarkParams.enableExpDelays = result["enable_exp_delays"].as(); - - // Argument: Enable batch stats output - bool logIterationData = result["log_iteration_data"].as(); - - // Argument: Enable chunked context - benchmarkParams.enableChunekedContextVec.resize(instanceNum); - if (result.count("enable_chunked_context_per_instance")) - { - auto enableChunkedContextVec = result["enable_chunked_context_per_instance"].as>(); - - TLLM_CHECK_WITH_INFO(enableChunkedContextVec.size() == instanceNum || enableChunkedContextVec.size() == 1, - "the number of enable_chunked_context_per_instance should be equal to the number of instances or equal to " - "1"); - for (int i = 0; i < instanceNum; i++) - { - benchmarkParams.enableChunekedContextVec.at(i) - = enableChunkedContextVec.size() == 1 ? enableChunkedContextVec[0] : enableChunkedContextVec[i]; - } - } - // Argument: Enable return context logits - bool returnContextLogits = result["return_context_logits"].as(); - TLLM_CHECK_WITH_INFO(returnContextLogits == false, "Currently disaggServer don't support returnContextLogits!"); - // Argument: Enable return context logits - bool returnGenerationLogits = result["return_generation_logits"].as(); - TLLM_CHECK_WITH_INFO( - returnGenerationLogits == false, "Currently disaggServer don't support returnGenerationLogits!"); - - if (result.count("lora_dir")) - { - TLLM_CHECK_WITH_INFO(false, "Currently disaggServer don't support lora!"); - benchmarkParams.loraDir = result["lora_dir"].as(); - } - if (result.count("lora_host_cache_bytes")) - { - TLLM_CHECK_WITH_INFO(false, "Currently disaggServer don't support lora!"); - - benchmarkParams.loraHostCacheSize = result["lora_host_cache_bytes"].as(); - } - if (result.count("lora_num_device_mod_layers")) - { - TLLM_CHECK_WITH_INFO(false, "Currently disaggServer don't support lora!"); - - benchmarkParams.loraDeviceNumModLayers = result["lora_num_device_mod_layers"].as(); - } - - // Argument: How many KV cache blocks (as fraction of number of GPU kv cache blocks). - benchmarkParams.kvHostCacheSize = result["kv_host_cache_bytes"].as(); - TLLM_CHECK_WITH_INFO( - benchmarkParams.kvHostCacheSize == false, "Currently disaggServer don't support kv_host_cache!"); - - // Argument: Medusa choices for the Medusa speculative decoding. - if (result.count("medusa_choices")) - { - TLLM_CHECK_WITH_INFO(false, "Currently disaggServer don't support medusa!"); - - benchmarkParams.medusaChoices = parseVectorOfVectors(result["medusa_choices"].as()); - } - if (result.count("executor_lookahead_config")) - { - TLLM_CHECK_WITH_INFO(false, "Currently disaggServer don't support lookhead!"); - - benchmarkParams.executorLookaheadConfig - = parseLookaheadConfig(result["executor_lookahead_config"].as()); - } - if (result.count("request_lookahead_config")) - { - TLLM_CHECK_WITH_INFO(false, "Currently disaggServer don't support lookhead!"); - - benchmarkParams.requestLookaheadConfig - = parseLookaheadConfig(result["request_lookahead_config"].as()); - } - - // Argument: multi_block_mode - benchmarkParams.multiBlockMode = result["multi_block_mode"].as(); - - // Argument: enable_context_fmha_fp32_acc - benchmarkParams.enableContextFMHAFP32Acc = result["enable_context_fmha_fp32_acc"].as(); - - // Argument: cuda_graph_mode - benchmarkParams.cudaGraphMode = result["cuda_graph_mode"].as(); - - // Argument: cuda_graph_cache_size - benchmarkParams.cudaGraphCacheSize = result["cuda_graph_cache_size"].as(); - - std::optional padId; - // Argument: Padding token id - if (result.count("pad_id")) - { - padId = result["pad_id"].as(); - } - - // Argument: End-of-sentence token id - std::optional eosId = result["eos_id"].as(); - - std::optional batchTimeout; - - std::optional staticEmulatedBatchSize; - // Argument: Static emulated batch size - if (result.count("static_emulated_batch_size")) - { - staticEmulatedBatchSize = result["static_emulated_batch_size"].as(); - } - - // Argument: Scheduler policy - texec::CapacitySchedulerPolicy capacitySchedulerPolicy; - auto const capacitySchedulerPolicyArg = result["scheduler_policy"].as(); - if (capacitySchedulerPolicyArg == "max_utilization") - { - capacitySchedulerPolicy = texec::CapacitySchedulerPolicy::kMAX_UTILIZATION; - } - else if (capacitySchedulerPolicyArg == "guaranteed_no_evict") - { - capacitySchedulerPolicy = texec::CapacitySchedulerPolicy::kGUARANTEED_NO_EVICT; - } - else if (capacitySchedulerPolicyArg == "static_batch") - { - capacitySchedulerPolicy = texec::CapacitySchedulerPolicy::kSTATIC_BATCH; - } - else - { - TLLM_LOG_ERROR("Unexpected scheduler policy: " + capacitySchedulerPolicyArg); - return 1; - } - - // Argument: max_prompt_len - std::optional maxPromptLen; - if (result.count("max_prompt_len")) - { - maxPromptLen = result["max_prompt_len"].as(); - } - - // Argument: GPU weights percentage - auto gpuWeightsPercent = result["gpu_weights_percent"].as(); - if (gpuWeightsPercent < 0 || gpuWeightsPercent > 1) - { - TLLM_LOG_ERROR("--gpu_weights_percent must be between 0.0 and 1.0 but got: %f", gpuWeightsPercent); - return 1; - } - benchmarkParams.gpuWeightsPercent = gpuWeightsPercent; - - std::optional>> deviceIdsForInstance = std::nullopt; - if (result.count("device_ids_for_instances")) - { - deviceIdsForInstance = parseVectorOfVectors(result["device_ids_for_instances"].as()); - } - benchmark(contextEnigePaths, generationEnginePaths, deviceIdsForInstance, datasetPath, opCsvFile, maxNumSamples, - beamWidth, result["warm_up"].as(), eosId, padId, benchmarkParams, capacitySchedulerPolicy, waitSleep, - returnContextLogits, returnContextLogits, staticEmulatedBatchSize, logIterationData, maxPromptLen, - hasContextAwait, hasGenAwait); -} diff --git a/benchmarks/cpp/gptManagerBenchmark.cpp b/benchmarks/cpp/gptManagerBenchmark.cpp deleted file mode 100644 index 287cbba343ce..000000000000 --- a/benchmarks/cpp/gptManagerBenchmark.cpp +++ /dev/null @@ -1,1557 +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. - */ - -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/logger.h" -#include "tensorrt_llm/executor/executor.h" -#include "tensorrt_llm/executor/tensor.h" -#include "tensorrt_llm/executor/types.h" -#include "tensorrt_llm/plugins/api/tllmPlugin.h" -#include "tensorrt_llm/runtime/common.h" -#include "tensorrt_llm/runtime/gptJsonConfig.h" -#include "tensorrt_llm/runtime/tllmLogger.h" -#include "tensorrt_llm/runtime/utils/mpiUtils.h" -#include "tensorrt_llm/runtime/utils/numpyUtils.h" -#include "tensorrt_llm/runtime/worldConfig.h" -#include "utils/utils.h" - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -using namespace tensorrt_llm::batch_manager; -using namespace tensorrt_llm::runtime; -using namespace tensorrt_llm::benchmark; -namespace texec = tensorrt_llm::executor; -namespace trt = nvinfer1; -namespace fs = std::filesystem; - -namespace -{ - -using TensorPtr = ITensor::SharedPtr; - -class LoraLib -{ -public: - LoraLib(std::string const& loraDir) - : mLoraDir(loraDir) - , mBufferManager(std::make_shared()) - , mTaskPaths(parseDirPaths(mLoraDir)) - , mLoras(readLoras(mTaskPaths)) - { - } - - TensorPtr getLoraWeights(uint64_t taskId) const - { - return mLoras.at(taskId).first; - } - - TensorPtr getLoraConfig(uint64_t taskId) const - { - return mLoras.at(taskId).second; - } - - void clear() - { - mLoras.clear(); - } - - std::map> const& getLoras() - { - return mLoras; - } - -private: - std::string const mLoraDir; - BufferManager mBufferManager; - std::map mTaskPaths; - std::map> mLoras; - - std::map> readLoras(std::map taskPaths) - { - std::map> loras; - for (auto const& [id, p] : taskPaths) - { - TensorPtr loraWeights - = utils::loadNpy(mBufferManager, (p / "model.lora_weights.npy").string(), MemoryType::kCPU); - TensorPtr loraConfig - = utils::loadNpy(mBufferManager, (p / "model.lora_config.npy").string(), MemoryType::kCPU); - loras.insert_or_assign(id, std::make_pair(loraWeights, loraConfig)); - } - return loras; - } - - std::map parseDirPaths(std::string const& loraDir) - { - std::map taskPaths; - if (loraDir == "") - { - return taskPaths; - } - for (auto const& entry : fs::recursive_directory_iterator(loraDir)) - { - if (entry.is_directory()) - { - auto taskId = parseId(entry.path()); - taskPaths.insert_or_assign(taskId, entry.path()); - } - } - return taskPaths; - } - - uint64_t parseId(fs::path p) - { - auto fn = p.filename().string(); - auto dashPos = fn.find_first_of("-"); - std::string idStr = fn; - if (dashPos != std::string::npos) - { - auto idStr = fn.substr(0, dashPos); - } - uint64_t id = static_cast(std::stoi(idStr)); - return id; - } -}; - -} // namespace - -struct BenchInfo -{ - BenchInfo() = default; - - BenchInfo(int inputLength, std::chrono::time_point start) - : inputLength(inputLength) - , start(start) - { - } - - int inputLength; - int outputLength{0}; - std::chrono::time_point start; - std::chrono::time_point end; - std::chrono::time_point firstTokenTs; - float latency{}; // millisecond - bool hasError{false}; - float firstTokenLatency{}; - std::optional avgGenT2TLatency{}; - bool firstTokenSeen{false}; - SizeType32 decodingIter{0}; -}; - -class Recorder -{ - using TensorPtr = ITensor::SharedPtr; - -public: - explicit Recorder(std::string opCsvFile, bool streaming = false, int beamWidth = 1, - std::string responsesJsonFile = "", bool excludeInputInOutput = false) - : mOpCsvFile(std::move(opCsvFile)) - , mStreaming(streaming) - , mBeamWidth(beamWidth) - , mRespJsonFile(std::move(responsesJsonFile)) - , mOutputHasInput(!excludeInputInOutput) - { - } - - void initialize() - { - mStart = std::chrono::steady_clock::now(); - mRequestsQueueingLatencies.clear(); - } - - void finalize() - { - mEnd = std::chrono::steady_clock::now(); - } - - void recordQueueLatency(std::vector const& latencies) - { - mRequestsQueueingLatencies.insert(mRequestsQueueingLatencies.end(), latencies.begin(), latencies.end()); - } - - // number of output tokens not calculated from output sequence here, instead set to max_output_len - // - if eos_id == -1 (default behavior), this is correct since output seq will have max permissible length. - // - However, if eos_id != -1, the token size of output sequence may be less than max_output_len, and token - // throughput may be inaccurate - void recordStart( - SizeType32 inputLength, uint64_t requestId, std::chrono::time_point const& start) - { - TLLM_CHECK_WITH_INFO(mRequestBenchInfos.find(requestId) == mRequestBenchInfos.end(), - "Request %lu already exists in record before start, please report a bug to developers.", requestId); - std::lock_guard const lock(mRequestBenchInfosMutex); - mRequestBenchInfos[requestId] = BenchInfo(inputLength, start); - } - - void recordToken( - texec::Response const& response, std::chrono::time_point const& tokenTime) - { - auto const requestId = response.getRequestId(); - auto outputTokenIds = response.getResult().outputTokenIds; - - int32_t outputLength = 1; - for (auto const& beam : outputTokenIds) - { - outputLength = std::max(static_cast(beam.size()), outputLength); - } - - std::lock_guard const lock(mRequestBenchInfosMutex); - mRequestBenchInfos[requestId].outputLength += outputLength; - - if (!mRequestBenchInfos[requestId].firstTokenSeen) - { - mRequestBenchInfos[requestId].firstTokenTs = tokenTime; - mRequestBenchInfos[requestId].firstTokenSeen = true; - } - - mRequestBenchInfos[requestId].decodingIter += 1; - } - - void recordEnd(texec::Response const& response, std::chrono::time_point const& end) - { - auto const requestId = response.getRequestId(); - // Get the actual output length - if (!response.hasError()) - { - if (!mStreaming) - { - TLLM_LOG_DEBUG("response.getResult().outputTokenIds"); - auto outputTokenIds = response.getResult().outputTokenIds; - - int32_t outSeqLen = 0; - for (auto const& beam : outputTokenIds) - { - outSeqLen = std::max(static_cast(beam.size()), outSeqLen); - } - if (mOutputHasInput) - { - int inputSeqLen = mRequestBenchInfos[requestId].inputLength; - outSeqLen -= inputSeqLen; - } - std::lock_guard const lock(mRequestBenchInfosMutex); - mRequestBenchInfos[requestId].outputLength = outSeqLen; - mRequestBenchInfos[requestId].decodingIter = response.getResult().decodingIter; - - // We record the first beam for the response file - mResponseTensors[requestId] = outputTokenIds[0]; - } - else - { - TLLM_CHECK_WITH_INFO(mBeamWidth == 1, "gptManagerBenchmark streaming mode does not support beam > 1"); - this->recordToken(response, end); - } - } - - std::lock_guard const lock(mRequestBenchInfosMutex); - mRequestBenchInfos[requestId].end = end; - mRequestBenchInfos[requestId].hasError = response.hasError(); - } - - float calcPercentile(std::vector const& latencies, int percentile) - { - int const index = static_cast(std::ceil((percentile / 100.0) * latencies.size())) - 1; - return latencies[index]; - } - - void calculateLatencies() - { - for (auto& reqInfo : mRequestBenchInfos) - { - reqInfo.second.latency - = std::chrono::duration(reqInfo.second.end - reqInfo.second.start).count(); - if (mStreaming) - { - reqInfo.second.firstTokenLatency - = std::chrono::duration(reqInfo.second.firstTokenTs - reqInfo.second.start) - .count(); - if (reqInfo.second.outputLength > 1) - { - reqInfo.second.avgGenT2TLatency - = std::chrono::duration(reqInfo.second.end - reqInfo.second.firstTokenTs) - .count() - / static_cast(reqInfo.second.outputLength - 1); - } - } - } - } - - void calculateMetrics() - { - calculateLatencies(); - - std::vector reqLatencies; - std::vector ftLatencies; - std::vector genT2TLatencies; - std::vector userTokensPerSecond; - - int totalOutputTokens{0}; - int totalDecodingIter{0}; - mNumErrorSamples = 0; - mNumSamples = 0; - for (auto reqInfo : mRequestBenchInfos) - { - if (!reqInfo.second.hasError) - { - reqLatencies.push_back(reqInfo.second.latency); - totalOutputTokens += reqInfo.second.outputLength; - totalDecodingIter += reqInfo.second.decodingIter; - - if (mStreaming) - { - ftLatencies.push_back(reqInfo.second.firstTokenLatency); - - if (reqInfo.second.avgGenT2TLatency) - { - genT2TLatencies.push_back(reqInfo.second.avgGenT2TLatency.value()); - } - if (reqInfo.second.avgGenT2TLatency.value() > 0) - { - userTokensPerSecond.push_back(1000.F / reqInfo.second.avgGenT2TLatency.value()); - } - } - ++mNumSamples; - } - else - { - ++mNumErrorSamples; - } - } - - mTotalLatency = std::chrono::duration(mEnd - mStart).count(); - mSeqThroughput = mNumSamples / (mTotalLatency / 1000); - mTokenThroughput = totalOutputTokens / (mTotalLatency / 1000); - mAcceptanceRate = totalDecodingIter - ? (static_cast(totalOutputTokens) / static_cast(totalDecodingIter)) - : 0.0f; - - mAvgSeqLatency = std::accumulate(reqLatencies.begin(), reqLatencies.end(), 0.F) / reqLatencies.size(); - - std::sort(reqLatencies.begin(), reqLatencies.end()); - - mP99SeqLatency = calcPercentile(reqLatencies, 99); - mP90SeqLatency = calcPercentile(reqLatencies, 90); - mP50SeqLatency = calcPercentile(reqLatencies, 50); - mMaxSeqLatency = reqLatencies.back(); - mMinSeqLatency = reqLatencies.front(); - - if (mStreaming) - { - mAvgFtLatency = std::accumulate(ftLatencies.begin(), ftLatencies.end(), 0.F) / ftLatencies.size(); - - std::sort(ftLatencies.begin(), ftLatencies.end()); - - mP99FtLatency = calcPercentile(ftLatencies, 99); - mP90FtLatency = calcPercentile(ftLatencies, 90); - mP50FtLatency = calcPercentile(ftLatencies, 50); - mMaxFtLatency = ftLatencies.back(); - mMinFtLatency = ftLatencies.front(); - - if (!genT2TLatencies.empty()) - { - mAvgGenT2TLatency - = std::accumulate(genT2TLatencies.begin(), genT2TLatencies.end(), 0.F) / genT2TLatencies.size(); - - std::sort(genT2TLatencies.begin(), genT2TLatencies.end()); - - mP99GenT2TLatency = calcPercentile(genT2TLatencies, 99); - mP90GenT2TLatency = calcPercentile(genT2TLatencies, 90); - mP50GenT2TLatency = calcPercentile(genT2TLatencies, 50); - mMaxGenT2TLatency = genT2TLatencies.back(); - mMinGenT2TLatency = genT2TLatencies.front(); - } - - if (!userTokensPerSecond.empty()) - { - mAvgUserTokensPerSecond = std::accumulate(userTokensPerSecond.begin(), userTokensPerSecond.end(), 0.F) - / userTokensPerSecond.size(); - std::sort(userTokensPerSecond.begin(), userTokensPerSecond.end()); - mP99UserTokensPerSecond = calcPercentile(userTokensPerSecond, 99); - mP90UserTokensPerSecond = calcPercentile(userTokensPerSecond, 90); - mP50UserTokensPerSecond = calcPercentile(userTokensPerSecond, 50); - mMaxUserTokensPerSecond = userTokensPerSecond.back(); - mMinUserTokensPerSecond = userTokensPerSecond.front(); - } - - mAvgReqQueueingLatency - = std::accumulate(mRequestsQueueingLatencies.begin(), mRequestsQueueingLatencies.end(), 0.F) - / mRequestsQueueingLatencies.size(); - std::sort(mRequestsQueueingLatencies.begin(), mRequestsQueueingLatencies.end()); - mP99ReqQueueingLatency = calcPercentile(mRequestsQueueingLatencies, 99); - mP90ReqQueueingLatency = calcPercentile(mRequestsQueueingLatencies, 90); - mP50ReqQueueingLatency = calcPercentile(mRequestsQueueingLatencies, 50); - mMaxReqQueueingLatency = mRequestsQueueingLatencies.back(); - mMinReqQueueingLatency = mRequestsQueueingLatencies.front(); - } - } - - void report() - { - - printf("[BENCHMARK] num_samples %d\n", mNumSamples); - printf("[BENCHMARK] num_error_samples %d\n", mNumErrorSamples); - printf("\n[BENCHMARK] num_samples %d\n", mNumSamples); - printf("[BENCHMARK] total_latency(ms) %.2f\n", mTotalLatency); - printf("[BENCHMARK] seq_throughput(seq/sec) %.2f\n", mSeqThroughput); - printf("[BENCHMARK] token_throughput(token/sec) %.2f\n", mTokenThroughput); - printf("[BENCHMARK] avg_acceptance_rate(tokens/decoding steps) %.2f\n\n", mAcceptanceRate); - - printf("[BENCHMARK] avg_sequence_latency(ms) %.2f\n", mAvgSeqLatency); - printf("[BENCHMARK] max_sequence_latency(ms) %.2f\n", mMaxSeqLatency); - printf("[BENCHMARK] min_sequence_latency(ms) %.2f\n", mMinSeqLatency); - printf("[BENCHMARK] p99_sequence_latency(ms) %.2f\n", mP99SeqLatency); - printf("[BENCHMARK] p90_sequence_latency(ms) %.2f\n", mP90SeqLatency); - printf("[BENCHMARK] p50_sequence_latency(ms) %.2f\n\n", mP50SeqLatency); - - if (mStreaming) - { - printf("[BENCHMARK] avg_time_to_first_token(ms) %.2f\n", mAvgFtLatency); - printf("[BENCHMARK] max_time_to_first_token(ms) %.2f\n", mMaxFtLatency); - printf("[BENCHMARK] min_time_to_first_token(ms) %.2f\n", mMinFtLatency); - printf("[BENCHMARK] p99_time_to_first_token(ms) %.2f\n", mP99FtLatency); - printf("[BENCHMARK] p90_time_to_first_token(ms) %.2f\n", mP90FtLatency); - printf("[BENCHMARK] p50_time_to_first_token(ms) %.2f\n\n", mP50FtLatency); - - printf("[BENCHMARK] avg_inter_token_latency(ms) %.2f\n", mAvgGenT2TLatency); - printf("[BENCHMARK] max_inter_token_latency(ms) %.2f\n", mMaxGenT2TLatency); - printf("[BENCHMARK] min_inter_token_latency(ms) %.2f\n", mMinGenT2TLatency); - printf("[BENCHMARK] p99_inter_token_latency(ms) %.2f\n", mP99GenT2TLatency); - printf("[BENCHMARK] p90_inter_token_latency(ms) %.2f\n", mP90GenT2TLatency); - printf("[BENCHMARK] p50_inter_token_latency(ms) %.2f\n\n", mP50GenT2TLatency); - - printf("[BENCHMARK] avg_user_tokens_per_second(tokens/sec/user) %.2f\n", mAvgUserTokensPerSecond); - printf("[BENCHMARK] max_user_tokens_per_second(tokens/sec/user) %.2f\n", mMaxUserTokensPerSecond); - printf("[BENCHMARK] min_user_tokens_per_second(tokens/sec/user) %.2f\n", mMinUserTokensPerSecond); - printf("[BENCHMARK] p99_user_tokens_per_second(tokens/sec/user) %.2f\n", mP99UserTokensPerSecond); - printf("[BENCHMARK] p90_user_tokens_per_second(tokens/sec/user) %.2f\n", mP90UserTokensPerSecond); - printf("[BENCHMARK] p50_user_tokens_per_second(tokens/sec/user) %.2f\n\n", mP50UserTokensPerSecond); - - printf("[BENCHMARK] avg_request_queueing_latency(ms) %.2f\n", mAvgReqQueueingLatency); - printf("[BENCHMARK] max_request_queueing_latency(ms) %.2f\n", mMaxReqQueueingLatency); - printf("[BENCHMARK] min_request_queueing_latency(ms) %.2f\n", mMinReqQueueingLatency); - printf("[BENCHMARK] p99_request_queueing_latency(ms) %.2f\n", mP99ReqQueueingLatency); - printf("[BENCHMARK] p90_request_queueing_latency(ms) %.2f\n", mP90ReqQueueingLatency); - printf("[BENCHMARK] p50_request_queueing_latency(ms) %.2f\n\n", mP50ReqQueueingLatency); - } - } - - void writeOpMetricsToCsv() - { - if (!mOpCsvFile.empty()) - { - std::vector headers = {"num_samples", "num_error_samples", "total_latency(ms)", - "seq_throughput(seq/sec)", "token_throughput(token/sec)", "avg_sequence_latency(ms)", - "max_sequence_latency(ms)", "min_sequence_latency(ms)", "p99_sequence_latency(ms)", - "p90_sequence_latency(ms)", "p50_sequence_latency(ms)", "avg_acceptance_rate(tokens/decoding steps)"}; - - if (mStreaming) - { - std::vector streamingHeaders = { - "avg_time_to_first_token(ms)", - "max_time_to_first_token(ms)", - "min_time_to_first_token(ms)", - "p99_time_to_first_token(ms)", - "p90_time_to_first_token(ms)", - "p50_time_to_first_token(ms)", - "avg_inter_token_latency(ms)", - "max_inter_token_latency(ms)", - "min_inter_token_latency(ms)", - "p99_inter_token_latency(ms)", - "p90_inter_token_latency(ms)", - "p50_inter_token_latency(ms)", - "avg_user_tokens_per_second(tokens/sec/user)", - "max_user_tokens_per_second(tokens/sec/user)", - "min_user_tokens_per_second(tokens/sec/user)", - "p99_user_tokens_per_second(tokens/sec/user)", - "p90_user_tokens_per_second(tokens/sec/user)", - "p50_user_tokens_per_second(tokens/sec/user)", - }; - - headers.insert(headers.end(), streamingHeaders.begin(), streamingHeaders.end()); - } - - std::ofstream outputFile(mOpCsvFile); - - if (outputFile.is_open()) - { - for (auto const& header : headers) - { - outputFile << header << ","; - } - outputFile << "\n"; - outputFile << mNumSamples << "," << mNumErrorSamples << "," << mTotalLatency << "," << mSeqThroughput - << "," << mTokenThroughput << "," << mAvgSeqLatency << "," << mMaxSeqLatency << "," - << mMinSeqLatency << "," << mP99SeqLatency << "," << mP90SeqLatency << "," << mP50SeqLatency - << "," << mAcceptanceRate; - if (mStreaming) - { - outputFile << "," << mAvgFtLatency << "," << mMaxFtLatency << "," << mMinFtLatency << "," - << mP99FtLatency << "," << mP90FtLatency << "," << mP50FtLatency << "," - << mAvgGenT2TLatency << "," << mMaxGenT2TLatency << "," << mMinGenT2TLatency << "," - << mP99GenT2TLatency << "," << mP90GenT2TLatency << "," << mP50GenT2TLatency << "," - << mAvgUserTokensPerSecond << "," << mMaxUserTokensPerSecond << "," - << mMinUserTokensPerSecond << "," << mP99UserTokensPerSecond << "," - << mP90UserTokensPerSecond << "," << mP50UserTokensPerSecond << ","; - } - - outputFile << "\n"; - } - else - { - std::cerr << "Error opening file '" << mOpCsvFile << "' for writing.\n"; - } - } - } - - void dumpResponseSeqs() - { - if (mRespJsonFile.empty()) - return; - nlohmann::json jsonResponses = nlohmann::json::array(); - for (auto const& [respId, respTokensTensor] : mResponseTensors) - { - auto respTokens = mResponseTensors[respId]; - int respLength = respTokens.size(); - int* respBufferPtr = respTokens.data(); - - if (mOutputHasInput) - { - int inputSeqLen = mRequestBenchInfos[respId].inputLength; - respBufferPtr += inputSeqLen; - respLength -= inputSeqLen; - } - - std::vector outputTokens(respLength); - std::copy(respBufferPtr, respBufferPtr + respLength, outputTokens.begin()); - - nlohmann::json currResp; - currResp["response_id"] = respId; - currResp["response_tokens"] = outputTokens; - jsonResponses.push_back(currResp); - } - std::ofstream outFile(mRespJsonFile); - outFile << jsonResponses; - outFile.close(); - } - -private: - std::unordered_map mRequestBenchInfos; - - std::chrono::time_point mStart; - std::chrono::time_point mEnd; - int mNumSamples{}; - int mNumErrorSamples{}; - float mTotalLatency{}; - float mSeqThroughput{}; - float mAvgSeqLatency{}; - float mAvgGenT2TLatency{}; - float mAvgUserTokensPerSecond{}; - float mAvgFtLatency{}; - float mTokenThroughput{}; - float mAcceptanceRate{}; - float mP99SeqLatency{}; - float mP90SeqLatency{}; - float mP50SeqLatency{}; - float mMaxSeqLatency{}; - float mMinSeqLatency{}; - float mP99FtLatency{}; - float mP90FtLatency{}; - float mP50FtLatency{}; - float mMaxFtLatency{}; - float mMinFtLatency{}; - float mP99GenT2TLatency{}; - float mP90GenT2TLatency{}; - float mP50GenT2TLatency{}; - float mMaxGenT2TLatency{}; - float mMinGenT2TLatency{}; - float mP99UserTokensPerSecond{}; - float mP90UserTokensPerSecond{}; - float mP50UserTokensPerSecond{}; - float mMaxUserTokensPerSecond{}; - float mMinUserTokensPerSecond{}; - float mAvgReqQueueingLatency{}; - float mP99ReqQueueingLatency{}; - float mP90ReqQueueingLatency{}; - float mP50ReqQueueingLatency{}; - float mMaxReqQueueingLatency{}; - float mMinReqQueueingLatency{}; - std::vector mRequestsQueueingLatencies{}; - - std::string mOpCsvFile; - bool mStreaming; - int mBeamWidth; - std::string mRespJsonFile; - std::unordered_map mResponseTensors; - bool mOutputHasInput; - std::mutex mRequestBenchInfosMutex; - -}; // class Recorder - -class ExecutorServer -{ -public: - ExecutorServer(std::optional const& decoderTrtEnginePath, - std::optional const& encoderTrtEnginePath, texec::BatchingType batchingType, - int32_t maxBeamWidth, texec::CapacitySchedulerPolicy capacitySchedulerPolicy, - BenchmarkParams const& benchmarkParams, std::shared_ptr recorder, std::chrono::milliseconds waitSleep, - bool logIterationData, texec::ModelType executorModelType) - : mRecorder(std::move(recorder)) - , mWaitSleep(waitSleep) - , mConcurrency(benchmarkParams.concurrency) - , mActiveCount(0) - , mNumFinished(0) - , mShutdown(false) - , mLogIterationData(logIterationData) - { - texec::DynamicBatchConfig dynamicBatchConfig( - benchmarkParams.enableBatchSizeTuning, benchmarkParams.enableMaxNumTokensTuning); - texec::SchedulerConfig schedulerConfig(capacitySchedulerPolicy, std::nullopt, dynamicBatchConfig); - - texec::KvCacheConfig kvCacheConfig(benchmarkParams.enableBlockReuse, benchmarkParams.maxTokensInPagedKvCache, - benchmarkParams.maxAttentionWindowVec, benchmarkParams.sinkTokenLength, - benchmarkParams.freeGpuMemoryFraction, benchmarkParams.kvHostCacheSize, - benchmarkParams.crossKvCacheFraction); - texec::PeftCacheConfig peftCacheConfig(0, benchmarkParams.loraDeviceNumModLayers, 8, 64, 4, 4, 4, 24, 8, - std::nullopt, benchmarkParams.loraHostCacheSize); - texec::ExtendedRuntimePerfKnobConfig extendedRuntimePerfKnobConfig(benchmarkParams.multiBlockMode, - benchmarkParams.enableContextFMHAFP32Acc, benchmarkParams.cudaGraphMode, - benchmarkParams.cudaGraphCacheSize); - texec::ExecutorConfig executorConfig( - maxBeamWidth, schedulerConfig, kvCacheConfig, benchmarkParams.enableChunkedContext, true); - executorConfig.setEnableTrtOverlap(benchmarkParams.enableTrtOverlap); - executorConfig.setGpuWeightsPercent(benchmarkParams.gpuWeightsPercent); - executorConfig.setPeftCacheConfig(peftCacheConfig); - executorConfig.setBatchingType(batchingType); - if (benchmarkParams.maxBatchSize) - { - executorConfig.setMaxBatchSize(benchmarkParams.maxBatchSize.value()); - } - if (benchmarkParams.maxNumTokens) - { - executorConfig.setMaxNumTokens(benchmarkParams.maxNumTokens.value()); - } - - auto decodingMode = texec::DecodingMode::Auto(); - if (benchmarkParams.medusaChoices.has_value()) - { - decodingMode = texec::DecodingMode::Medusa(); - } - else if (benchmarkParams.executorLookaheadConfig.has_value()) - { - decodingMode = texec::DecodingMode::Lookahead(); - } - else if (benchmarkParams.eagleConfig.has_value()) - { - decodingMode = texec::DecodingMode::Eagle(); - } - - executorConfig.setDecodingConfig(texec::DecodingConfig(decodingMode, benchmarkParams.executorLookaheadConfig, - benchmarkParams.medusaChoices, benchmarkParams.eagleConfig)); - executorConfig.setExtendedRuntimePerfKnobConfig(extendedRuntimePerfKnobConfig); - - if (executorModelType == texec::ModelType::kDECODER_ONLY) - { - mExecutor - = std::make_unique(decoderTrtEnginePath.value(), executorModelType, executorConfig); - } - else if (executorModelType == texec::ModelType::kENCODER_DECODER) - { - mExecutor = std::make_unique( - encoderTrtEnginePath.value(), decoderTrtEnginePath.value(), executorModelType, executorConfig); - } - else if (executorModelType == texec::ModelType::kENCODER_ONLY) - { - mExecutor - = std::make_unique(encoderTrtEnginePath.value(), executorModelType, executorConfig); - } - else - { - TLLM_LOG_ERROR("not a supported executor model type in executor server."); - } - - auto const& world = tensorrt_llm::mpi::MpiComm::world(); - auto worldRank = world.getRank(); - if (worldRank == 0) - { - mCollectStatsThread = std::thread(&ExecutorServer::collectStats, this); - } - } - - ~ExecutorServer() - { - mShutdown = true; - if (mCollectStatsThread.joinable()) - { - mCollectStatsThread.join(); - } - } - - void enqueue(std::vector requests, bool warmup = false) - { - try - { - std::vector inputLengths; - for (auto const& request : requests) - { - inputLengths.push_back(request.getInputTokenIds().size()); - } - auto const start = std::chrono::steady_clock::now(); - auto reqIds = mExecutor->enqueueRequests(std::move(requests)); - for (int req = 0; req < reqIds.size(); ++req) - { - if (!warmup) - { - mRecorder->recordStart(inputLengths.at(req), reqIds.at(req), start); - } - mActiveCount++; - } - } - catch (std::exception const& e) - { - TLLM_THROW("%s", e.what()); - } - } - - void resetNumFinished() - { - mNumFinished = 0; - } - - bool canEnqueue(int numSentRequests) const - { - return !mConcurrency || (numSentRequests - mNumFinished < mConcurrency); - } - - void waitForResponses(SizeType32 numRequests, bool warmup = false) - { - while (mActiveCount || (mNumFinished < numRequests)) - { - auto responses = mExecutor->awaitResponses(mWaitSleep); - auto const tokenTime = std::chrono::steady_clock::now(); - for (auto const& response : responses) - { - if (response.getResult().isFinal) - { - mActiveCount--; - mNumFinished++; - if (!warmup) - { - mRecorder->recordEnd(response, tokenTime); - } - } - else - { - if (!warmup && !response.hasError()) - { - mRecorder->recordToken(response, tokenTime); - } - } - } - } - } - - void collectStats() const - { - while (!mShutdown) - { - auto iterStats = mExecutor->getLatestIterationStats(); - for (auto const& iterStat : iterStats) - { - SizeType32 numNewActiveRequests = iterStat.numNewActiveRequests; - if (numNewActiveRequests > 0) - { - float avgQueueingTime - = static_cast(iterStat.newActiveRequestsQueueLatencyMS / numNewActiveRequests); - std::vector requestsQueueLatencyMS(numNewActiveRequests, avgQueueingTime); - mRecorder->recordQueueLatency(requestsQueueLatencyMS); - } - if (mLogIterationData) - { - TLLM_LOG_INFO(texec::JsonSerialization::toJsonStr(iterStat)); - } - } - auto const waitSleep = std::chrono::milliseconds(50); - std::this_thread::sleep_for(waitSleep); - } - } - -private: - std::unique_ptr mExecutor; - std::thread mCollectStatsThread; - std::shared_ptr mRecorder; - std::chrono::milliseconds mWaitSleep; - std::optional mConcurrency; - std::atomic mActiveCount; - std::atomic mNumFinished; - std::atomic mShutdown; - bool mLogIterationData; -}; // class ExecutorServer - -namespace -{ - -texec::Request makeExecutorRequest(Sample const& sample, SizeType32 const& beamWidth, - std::optional const& eosId, std::optional const& padId, bool streaming = false, - bool const& returnContextLogits = false, bool const& returnGenerationLogits = false, - std::optional const& loraConfig = std::nullopt, - std::optional const& lookaheadConfig = std::nullopt, - std::optional encoderInputTokenIds = std::nullopt, - std::optional temperature = std::nullopt) -{ - auto samplingConfig = texec::SamplingConfig{beamWidth}; - samplingConfig.setTemperature(temperature); - auto outputConfig = texec::OutputConfig{false, returnContextLogits, returnGenerationLogits, false}; - return texec::Request(sample.inputIds, sample.outputLen, streaming, samplingConfig, outputConfig, eosId, padId, - std::nullopt, // positionIds - std::nullopt, // badWords - std::nullopt, // stopWords - std::nullopt, // embeddingBias - std::nullopt, // speculativeDecoding - std::nullopt, // pTuning - std::nullopt, // multimodalInput - std::nullopt, // multimodalEmbedding - std::nullopt, // mRopeConfig - loraConfig, // loraConfig - lookaheadConfig, // lookaheadConfig - std::nullopt, // kvCacheRetentionConfig - std::nullopt, // logitsPostProcessorName - std::nullopt, // logitsPostProcessor - encoderInputTokenIds.has_value() ? encoderInputTokenIds : std::nullopt, - std::nullopt); // cacheSalt -} - -void benchmarkExecutor(std::optional const& decoderEngineDir, - std::optional const& encoderEngineDir, texec::BatchingType batchingType, - std::string const& datasetPath, std::string const& opCsvFile, int maxNumSamples, int beamWidth, int warmUp, - std::optional const& eosId, std::optional const& padId, BenchmarkParams const& benchmarkParams, - texec::CapacitySchedulerPolicy capacitySchedulerPolicy, std::chrono::milliseconds waitSleep, - bool returnContextLogits, bool returnGenerationLogits, std::optional const staticEmulatedBatchSize, - bool logIterationData, std::optional const maxPromptLen, texec::ModelType executorModelType, - std::string const& responsesJsonFile) -{ - auto const& world = tensorrt_llm::mpi::MpiComm::world(); - auto worldRank = world.getRank(); - - // Load dataset - auto const samples = parseWorkloadJson(datasetPath, maxNumSamples, maxPromptLen); - auto const numSamples = samples.size(); - - auto recorder = std::make_shared(opCsvFile, benchmarkParams.streaming, beamWidth, responsesJsonFile); - int32_t decoderStartTokenId = 0; - std::shared_ptr executorServer; - - if (executorModelType == texec::ModelType::kDECODER_ONLY) - { - TLLM_CHECK_WITH_INFO( - decoderEngineDir.has_value(), "decoder models require a path to decoder engine in executor benchmark."); - executorServer - = std::make_shared(decoderEngineDir.value(), std::nullopt, batchingType, beamWidth, - capacitySchedulerPolicy, benchmarkParams, recorder, waitSleep, logIterationData, executorModelType); - } - else if (executorModelType == texec::ModelType::kENCODER_DECODER) - { - TLLM_CHECK_WITH_INFO(encoderEngineDir.has_value(), - "encoder-decoder models require a path to encoder engine in executor benchmark."); - executorServer = std::make_shared(decoderEngineDir.value(), encoderEngineDir.value(), - batchingType, beamWidth, capacitySchedulerPolicy, benchmarkParams, recorder, waitSleep, logIterationData, - executorModelType); - try - { - std::ifstream decoderJsonConfigPath(decoderEngineDir.value() / "config.json"); - auto const decoderPretrainedConfig - = nlohmann::json::parse(decoderJsonConfigPath, nullptr, true, true).at("pretrained_config"); - decoderStartTokenId = decoderPretrainedConfig.at("decoder_start_token_id").template get(); - } - catch (nlohmann::json::out_of_range& e) - { - TLLM_LOG_ERROR( - "Parameter %s cannot be read from decoder config.json in pretrained_config. Using default id %d.", - std::string("decoder_start_token_id").c_str(), decoderStartTokenId); - } - catch (nlohmann::json::type_error const& e) - { - TLLM_LOG_ERROR( - "Parameter %s has error type in decoder config.json in pretrained_config. Using default id %d.", - std::string("decoder_start_token_id").c_str(), decoderStartTokenId); - } - } - else if (executorModelType == texec::ModelType::kENCODER_ONLY) - { - TLLM_CHECK_WITH_INFO( - encoderEngineDir.has_value(), "encoder models require a path to encoder engine in executor benchmark."); - executorServer - = std::make_shared(std::nullopt, encoderEngineDir.value(), batchingType, beamWidth, - capacitySchedulerPolicy, benchmarkParams, recorder, waitSleep, logIterationData, executorModelType); - } - else - { - TLLM_LOG_ERROR("not a supported executor model type in executor benchmark."); - return; - } - - if (worldRank == 0) - { - if (benchmarkParams.loraDir) - { - auto startLoraLoad = std::chrono::steady_clock::now(); - LoraLib loras(benchmarkParams.loraDir.value()); - std::vector requests; - for (auto& [taskId, p] : loras.getLoras()) - { - // squeeze lora configs and weights since LoraConfig requires them to be 2D tensors - p.first->squeeze(0); - p.second->squeeze(0); - texec::LoraConfig loraConfig( - taskId, texec::detail::ofITensor(p.first), texec::detail::ofITensor(p.second)); - if (executorModelType == texec::ModelType::kENCODER_DECODER) - { - Sample s{std::vector{decoderStartTokenId}, 1, static_cast(taskId)}; - requests.emplace_back(makeExecutorRequest(s, beamWidth, eosId, padId, false, false, false, - loraConfig, std::nullopt, std::vector{1, 2, 3, 4, 5})); - } - else - { - Sample s{std::vector{1, 2, 3, 4, 5}, 1, static_cast(taskId)}; - requests.emplace_back( - makeExecutorRequest(s, beamWidth, eosId, padId, false, false, false, loraConfig, std::nullopt)); - } - } - executorServer->enqueue(std::move(requests), true); - executorServer->waitForResponses(loras.getLoras().size(), true); - auto endLoraLoad = std::chrono::steady_clock::now(); - printf("[BENCHMARK] time to preload LoRAs(ms) %.2f\n", - std::chrono::duration(endLoraLoad - startLoraLoad).count()); - } - // Warm up - { - std::vector requests; - for (auto i = 0; i < warmUp; ++i) - { - if (executorModelType == texec::ModelType::kENCODER_DECODER) - { - Sample s{std::vector{decoderStartTokenId}, samples[0].outputLen, samples[0].taskId}; - requests.emplace_back(makeExecutorRequest(s, beamWidth, eosId, padId, benchmarkParams.streaming, - returnContextLogits, returnGenerationLogits, std::nullopt, - benchmarkParams.requestLookaheadConfig, samples[0].inputIds)); - } - else - { - requests.emplace_back(makeExecutorRequest(samples[0], beamWidth, eosId, padId, - benchmarkParams.streaming, returnContextLogits, returnGenerationLogits, std::nullopt, - benchmarkParams.requestLookaheadConfig, std::nullopt, benchmarkParams.temperature)); - } - } - executorServer->enqueue(std::move(requests), true); - executorServer->waitForResponses(warmUp, true); - } - - // Benchmark - { - auto timeDelays = computeTimeDelays(benchmarkParams, numSamples - 1); - - // Create requests - recorder->initialize(); - std::vector requests; - - for (std::size_t i = 0; i < numSamples; ++i) - { - std::optional loraConfig; - if (samples[i].taskId >= 0) - { - loraConfig = texec::LoraConfig(samples[i].taskId); - } - if (executorModelType == texec::ModelType::kENCODER_DECODER) - { - Sample s{std::vector{decoderStartTokenId}, samples[i].outputLen, samples[i].taskId}; - requests.emplace_back(makeExecutorRequest(s, beamWidth, eosId, padId, benchmarkParams.streaming, - returnContextLogits, returnGenerationLogits, loraConfig, benchmarkParams.requestLookaheadConfig, - samples[i].inputIds)); - } - else - { - requests.emplace_back(makeExecutorRequest(samples[i], beamWidth, eosId, padId, - benchmarkParams.streaming, returnContextLogits, returnGenerationLogits, loraConfig, - benchmarkParams.requestLookaheadConfig, std::nullopt, benchmarkParams.temperature)); - } - } - - bool const hasDelay - = std::any_of(timeDelays.begin(), timeDelays.end(), [](auto const& delay) { return delay > 0.0; }); - executorServer->resetNumFinished(); - if (!staticEmulatedBatchSize) - { - // Launch a thread that will wait for responses - std::thread waitThread( - [numSamples, executorServer]() { executorServer->waitForResponses(numSamples); }); - - // Enqueue requests one by one - int numSentRequests = 0; - while (numSentRequests < numSamples) - { - if (executorServer->canEnqueue(numSentRequests)) - { - executorServer->enqueue({requests.at(numSentRequests)}); - if (hasDelay && numSentRequests < numSamples - 1) - { - std::this_thread::sleep_for( - std::chrono::milliseconds(static_cast(timeDelays.at(numSentRequests) * 1000))); - } - numSentRequests += 1; - } - } - waitThread.join(); - } - else - { - TLLM_CHECK_WITH_INFO( - !hasDelay, "Executor benchmark doesn't support delays with emulated static batch sizes"); - SizeType32 numRequests = requests.size(); - SizeType32 maxBatchSize = staticEmulatedBatchSize.value(); - for (SizeType32 req = 0; req < numRequests; req += maxBatchSize) - { - auto batchSize = std::min(maxBatchSize, numRequests - req); - - std::vector requestsBatch(std::make_move_iterator(requests.begin() + req), - std::make_move_iterator(requests.begin() + req + batchSize)); - // Enqueue in batches - executorServer->enqueue(std::move(requestsBatch)); - // Wait for current batch to be done - executorServer->waitForResponses(batchSize); - } - } - } - recorder->finalize(); - recorder->calculateMetrics(); - recorder->report(); - recorder->writeOpMetricsToCsv(); - recorder->dumpResponseSeqs(); - // Send terminateReqId to terminate servers on all ranks - // Sever on rank 0 will broadcast the terminate signal to other servers on multi-GPU cases - } -} - -} // namespace - -int main(int argc, char* argv[]) -{ - cxxopts::Options options( - "TensorRT LLM BatchManager Benchmark", "TensorRT LLM BatchManager Benchmark for GPT and GPT-like models."); - options.add_options()("h,help", "Print usage"); - options.add_options()("engine_dir, decoder_engine_dir", "Directory that store the engines of decoder models.", - cxxopts::value()); - options.add_options()( - "encoder_engine_dir", "Directory that store the engines of the encoder models.", cxxopts::value()); - options.add_options()( - "api", "API type: gptManager or executor.", cxxopts::value()->default_value("executor")); - options.add_options()("type", - "Batching type: choose between inflight/static. (IFB/V1 options are going to be deprecated)", - cxxopts::value()->default_value("inflight")); - options.add_options()("dataset", "Dataset that is used for benchmarking BatchManager.", - cxxopts::value()->default_value("")); - options.add_options()( - "output_csv", "Write output metrics to CSV", cxxopts::value()->default_value("")); - options.add_options()("max_num_samples", "maximum number of samples to use from dataset/generate", - cxxopts::value()->default_value("100000")); - options.add_options()( - "beam_width", "Specify beam width you want to benchmark.", cxxopts::value()->default_value("1")); - options.add_options()( - "warm_up", "Specify warm up iterations before benchmark starts.", cxxopts::value()->default_value("2")); - options.add_options()( - "eos_id", "Specify the end-of-sequence token id.", cxxopts::value()->default_value("-1")); - options.add_options()("pad_id", "Specify the padding token id.", cxxopts::value()); - options.add_options()("max_tokens_in_paged_kvcache", "Max tokens in paged K-V Cache.", cxxopts::value()); - options.add_options()( - "max_attention_window", "Max KV cache length per sequence", cxxopts::value>()); - options.add_options()("sink_token_len", "Sink token length in kv cache per sequence.", cxxopts::value()); - options.add_options()( - "random_seed", "integer random seed for exponential time delays.", cxxopts::value()->default_value("420")); - options.add_options()( - "kv_cache_free_gpu_mem_fraction", "K-V Cache Free Gpu Mem Fraction.", cxxopts::value()); - options.add_options()( - "cross_kv_cache_fraction", "Cross K-V Cache Fraction (from 0.0 to 1.0).", cxxopts::value()); - options.add_options()("request_rate", - "request rate in reqs/sec. Skipping this arg or negative value will trigger offline/0-delay.", - cxxopts::value()); - options.add_options()("concurrency", "Concurrent number of connections with the server.", cxxopts::value()); - options.add_options()("max_batch_size", "The max runtime batch size when benchmarking", cxxopts::value()); - options.add_options()( - "max_num_tokens", "The max runtime number of tokens per batch when benchmarking", cxxopts::value()); - options.add_options()( - "enable_batch_size_tuning", "Dynamic tuning of batch size", cxxopts::value()->default_value("false")); - options.add_options()("enable_max_num_tokens_tuning", "Dynamic tuning of max num tokens", - cxxopts::value()->default_value("false")); - options.add_options()("enable_exp_delays", "Enables exponential delay distr to mimic real world request arrival", - cxxopts::value()->default_value("false")); - options.add_options()("streaming", - "Operate in streaming mode. Note: it reflects time-to-first-token and inter-token-latency", - cxxopts::value()->default_value("false")); - options.add_options()( - "enable_kv_cache_reuse", "Enables the KV cache reuse.", cxxopts::value()->default_value("true")); - options.add_options()( - "enable_chunked_context", "Whether to enable context chunking.", cxxopts::value()->default_value("true")); - options.add_options()( - "return_context_logits", "Whether to return context logits.", cxxopts::value()->default_value("false")); - options.add_options()("return_generation_logits", "Whether to return generation logits.", - cxxopts::value()->default_value("false")); - - options.add_options()("scheduler_policy", - "Choose scheduler policy between max_utilization/guaranteed_no_evict/static_batch.", - cxxopts::value()->default_value("guaranteed_no_evict")); - - options.add_options()("static_emulated_batch_size", - "Emulate static batching performance with the provided batch size.", cxxopts::value()); - options.add_options()("log_level", "Choose log level between verbose/info/warning/error/internal_error.", - cxxopts::value()->default_value("warning")); - options.add_options()("log_iteration_data", "On each decoder iteration, print batch state metadata.", - cxxopts::value()->default_value("false")); - options.add_options()("wait_sleep", "Specify how many milliseconds to sleep each iteration of waitForEmpty loop.", - cxxopts::value()->default_value("25")); - options.add_options()("lora_dir", "Directory containing LoRAs", cxxopts::value()->default_value("")); - options.add_options()("lora_host_cache_bytes", "LoRA host cache memory in bytes", cxxopts::value()); - options.add_options()("lora_num_device_mod_layers", "LoRA number 1d cache rows", cxxopts::value()); - options.add_options()("kv_host_cache_bytes", - "Size of secondary memory pool used for offloading kv cache blocks (in bytes).", - cxxopts::value()->default_value("0")); - options.add_options()( - "max_prompt_len", "Truncate all prompts from dataset to the length specified.", cxxopts::value()); - - options.add_options()("gpu_weights_percent", - "Specify the percentage of weights that reside on GPU (from 0.0 to 1.0).", - cxxopts::value()->default_value("1.0")); - options.add_options()( - "medusa_choices", "Medusa choices in the format of [[0], [0, 1], [0, 0, 1]]", cxxopts::value()); - options.add_options()( - "eagle_choices", "Eagle choices in the format of [[0], [0, 1], [0, 0, 1]]", cxxopts::value()); - options.add_options()("eagle_posterior_threshold", - "Minimum token probability threshold for typical acceptance. Enables typical acceptance in Eagle", - cxxopts::value()); - options.add_options()("temperature", "Sampling temperature for each request", cxxopts::value()); - options.add_options()( - "eagle_use_dynamic_tree", "Whether to use Eagle-2", cxxopts::value()->default_value("false")); - options.add_options()("eagle_dynamic_tree_max_top_k", - "The max topK for dynamic tree, also the number of draft tokens that will expand for each node", - cxxopts::value()); - - options.add_options()("multi_block_mode", - "Distribute the work across multiple CUDA thread-blocks on the GPU for masked MHA kernel", - cxxopts::value()->default_value("true")); - options.add_options()("cuda_graph_mode", "When enabled, inference is executed with cuda graph.", - cxxopts::value()->default_value("false")); - options.add_options()("cuda_graph_cache_size", - "Specify how many cuda graphs are cached in the runtime. Larger cache gives better perf, but consumes more GPU " - "memory.", - cxxopts::value()->default_value("0")); - options.add_options()("enable_trt_overlap", "Enable TRT Overlap", cxxopts::value()->default_value("false")); - - options.add_options()("enable_context_fmha_fp32_acc", "Enable FMHA runner FP32 accumulation", - cxxopts::value()->default_value("false")); - options.add_options()("executor_lookahead_config", - "lookahead config in the format of [max_window_size, max_ngram_size, max_verification_set_size]", - cxxopts::value()); - options.add_options()("request_lookahead_config", - "lookahead config in the format of [max_window_size, max_ngram_size, max_verification_set_size], and each <= " - "executor lookahead config", - cxxopts::value()); - options.add_options()("responses_json", "Write output response sequences to a json file", - cxxopts::value()->default_value("")); - - auto result = options.parse(argc, argv); - - if (result.count("help")) - { - std::cout << options.help() << std::endl; - return 0; - } - - // Argument: Engine directory - if (!result.count("engine_dir") && !result.count("encoder_engine_dir")) - { - std::cout << options.help() << std::endl; - TLLM_LOG_ERROR("Please specify engine directory."); - return 1; - } - - // Argument: Batching Type - auto const type = result["type"].as(); - texec::BatchingType batchingType{texec::BatchingType::kINFLIGHT}; - if (type == "V1" || type == "static") - { - if (type == "V1") - { - TLLM_LOG_WARNING("type option \"V1\" is going to be renamed to \"static\"."); - } - bool streaming = result["streaming"].as(); - if (streaming) - { - TLLM_LOG_ERROR("Streaming is not supported in static batching.\n"); - return 1; - } - batchingType = texec::BatchingType::kSTATIC; - } - else if (type == "IFB" || type == "inflight") - { - if (type == "IFB") - { - TLLM_LOG_WARNING("type option \"IFB\" is going to be renamed to \"inflight\"."); - } - batchingType = texec::BatchingType::kINFLIGHT; - } - else - { - TLLM_LOG_ERROR("Unexpected batching type: %s", type.c_str()); - return 1; - } - - // Argument: Dataset - auto const datasetPath = result["dataset"].as(); - auto const maxNumSamples = result["max_num_samples"].as(); - - // Argument: Output metrics CSV - auto const opCsvFile = result["output_csv"].as(); - - // Argument: beam width - auto const beamWidth = result["beam_width"].as(); - - // Argument: wait_sleep - auto const waitSleep = std::chrono::milliseconds(result["wait_sleep"].as()); - BenchmarkParams benchmarkParams; - - // Argument: Max tokens in paged K-V Cache - if (result.count("max_tokens_in_paged_kvcache")) - { - benchmarkParams.maxTokensInPagedKvCache = result["max_tokens_in_paged_kvcache"].as(); - } - - // Argument: Max KV cache length - if (result.count("max_attention_window")) - { - benchmarkParams.maxAttentionWindowVec = result["max_attention_window"].as>(); - } - - // Argument: Sink token length - if (result.count("sink_token_len")) - { - benchmarkParams.sinkTokenLength = result["sink_token_len"].as(); - } - - if (result.count("random_seed")) - { - benchmarkParams.randomSeed = result["random_seed"].as(); - } - - // Argument: K-V Cache Free Gpu Mem Fraction - if (result.count("kv_cache_free_gpu_mem_fraction")) - { - benchmarkParams.freeGpuMemoryFraction = result["kv_cache_free_gpu_mem_fraction"].as(); - } - // Argument: K-V Cache Cross Attention Fraction. Only applicable to enc-dec models. - if (result.count("encoder_engine_dir") && result.count("decoder_engine_dir")) - { - if (result.count("cross_kv_cache_fraction")) - { - benchmarkParams.crossKvCacheFraction = result["cross_kv_cache_fraction"].as(); - } - else - { - benchmarkParams.crossKvCacheFraction - = 0.5f; // default value if not set. but non enc-dec should not even have this param set - } - } - - // Argument: Enable dynamic tuning of batch size - benchmarkParams.enableBatchSizeTuning = result["enable_batch_size_tuning"].as(); - - // Argument: Enable dynamic tuning of max num tokens - benchmarkParams.enableMaxNumTokensTuning = result["enable_max_num_tokens_tuning"].as(); - - // Argument: Enable KV cache reuse - benchmarkParams.enableBlockReuse = result["enable_kv_cache_reuse"].as(); - - // Argument: streaming - benchmarkParams.streaming = result["streaming"].as(); - - TLLM_CHECK_WITH_INFO(!(result.count("request_rate") && result.count("concurrency")), - "request_rate and concurrency cannot be specified at the same time."); - - // Argument: request rate - if (result.count("request_rate")) - { - benchmarkParams.requestRate = result["request_rate"].as(); - } - - // Argument: concurrency - if (result.count("concurrency")) - { - benchmarkParams.concurrency = result["concurrency"].as(); - } - - // Argument: request rate - if (result.count("max_batch_size")) - { - benchmarkParams.maxBatchSize = result["max_batch_size"].as(); - } - - // Argument: request rate - if (result.count("max_num_tokens")) - { - benchmarkParams.maxNumTokens = result["max_num_tokens"].as(); - } - - benchmarkParams.enableExpDelays = result["enable_exp_delays"].as(); - - // Argument: Enable batch stats output - bool logIterationData = result["log_iteration_data"].as(); - - if (logIterationData) - { - TLLM_LOG_WARNING("Setting log_iteration_data to true adds overheads and may result in lower perf"); - } - - // Argument: Enable chunked context - benchmarkParams.enableChunkedContext = result["enable_chunked_context"].as(); - - // Argument: Enable return context logits - bool returnContextLogits = result["return_context_logits"].as(); - - // Argument: Enable return context logits - bool returnGenerationLogits = result["return_generation_logits"].as(); - - if (result.count("lora_dir")) - { - benchmarkParams.loraDir = result["lora_dir"].as(); - } - if (result.count("lora_host_cache_bytes")) - { - benchmarkParams.loraHostCacheSize = result["lora_host_cache_bytes"].as(); - } - if (result.count("lora_num_device_mod_layers")) - { - benchmarkParams.loraDeviceNumModLayers = result["lora_num_device_mod_layers"].as(); - } - - // Argument: How many KV cache blocks (as fraction of number of GPU kv cache blocks). - benchmarkParams.kvHostCacheSize = result["kv_host_cache_bytes"].as(); - - // Argument: Medusa choices for the Medusa speculative decoding. - if (result.count("medusa_choices")) - { - benchmarkParams.medusaChoices = parseVectorOfVectors(result["medusa_choices"].as()); - } - // Argument: Eagle choices for the Eagle speculative decoding. - if (result.count("eagle_choices") || result.count("eagle_posterior_threshold") - || result.count("eagle_use_dynamic_tree") || result.count("eagle_dynamic_tree_max_top_k")) - { - std::optional posteriorThreshold; - if (result.count("eagle_posterior_threshold")) - { - posteriorThreshold = result["eagle_posterior_threshold"].as(); - } - std::optional choices; - if (result.count("eagle_choices")) - { - choices = parseVectorOfVectors(result["eagle_choices"].as()); - } - bool eagleUseDynamicTree = false; - if (result.count("eagle_use_dynamic_tree")) - { - eagleUseDynamicTree = result["eagle_use_dynamic_tree"].as(); - } - std::optional eagleDynamicTreeMaxTopK; - if (result.count("eagle_dynamic_tree_max_top_k")) - { - eagleDynamicTreeMaxTopK = result["eagle_dynamic_tree_max_top_k"].as(); - } - benchmarkParams.eagleConfig = texec::EagleConfig( - choices, !posteriorThreshold.has_value(), posteriorThreshold, eagleUseDynamicTree, eagleDynamicTreeMaxTopK); - } - if (result.count("temperature")) - { - benchmarkParams.temperature = result["temperature"].as(); - } - - if (result.count("executor_lookahead_config")) - { - benchmarkParams.executorLookaheadConfig - = parseLookaheadConfig(result["executor_lookahead_config"].as()); - } - if (result.count("request_lookahead_config")) - { - benchmarkParams.requestLookaheadConfig - = parseLookaheadConfig(result["request_lookahead_config"].as()); - } - - // Argument: multi_block_mode - benchmarkParams.multiBlockMode = result["multi_block_mode"].as(); - - // Argument: enable_context_fmha_fp32_acc - benchmarkParams.enableContextFMHAFP32Acc = result["enable_context_fmha_fp32_acc"].as(); - - // Argument: cuda_graph_mode - benchmarkParams.cudaGraphMode = result["cuda_graph_mode"].as(); - - // Argument: cuda_graph_cache_size - benchmarkParams.cudaGraphCacheSize = result["cuda_graph_cache_size"].as(); - - // Argument: enable_trt_overlap - benchmarkParams.enableTrtOverlap = result["enable_trt_overlap"].as(); - - std::optional padId; - // Argument: Padding token id - if (result.count("pad_id")) - { - padId = result["pad_id"].as(); - } - - // Argument: End-of-sentence token id - std::optional eosId = result["eos_id"].as(); - - std::optional staticEmulatedBatchSize; - // Argument: Static emulated batch size - if (result.count("static_emulated_batch_size")) - { - staticEmulatedBatchSize = result["static_emulated_batch_size"].as(); - } - - // Argument: Scheduler policy - texec::CapacitySchedulerPolicy capacitySchedulerPolicy; - auto const capacitySchedulerPolicyArg = result["scheduler_policy"].as(); - if (capacitySchedulerPolicyArg == "max_utilization") - { - capacitySchedulerPolicy = texec::CapacitySchedulerPolicy::kMAX_UTILIZATION; - } - else if (capacitySchedulerPolicyArg == "guaranteed_no_evict") - { - capacitySchedulerPolicy = texec::CapacitySchedulerPolicy::kGUARANTEED_NO_EVICT; - } - else if (capacitySchedulerPolicyArg == "static_batch") - { - capacitySchedulerPolicy = texec::CapacitySchedulerPolicy::kSTATIC_BATCH; - } - else - { - TLLM_LOG_ERROR("Unexpected scheduler policy: " + capacitySchedulerPolicyArg); - return 1; - } - - // Argument: max_prompt_len - std::optional maxPromptLen; - if (result.count("max_prompt_len")) - { - maxPromptLen = result["max_prompt_len"].as(); - } - - // Argument: GPU weights percentage - auto gpuWeightsPercent = result["gpu_weights_percent"].as(); - if (gpuWeightsPercent < 0 || gpuWeightsPercent > 1) - { - TLLM_LOG_ERROR("--gpu_weights_percent must be between 0.0 and 1.0 but got: %f", gpuWeightsPercent); - return 1; - } - benchmarkParams.gpuWeightsPercent = gpuWeightsPercent; - - // Argument: Log level - auto logger = std::make_shared(); - auto const logLevel = result["log_level"].as(); - if (logLevel == "verbose") - { - logger->setLevel(trt::ILogger::Severity::kVERBOSE); - } - else if (logLevel == "info") - { - logger->setLevel(trt::ILogger::Severity::kINFO); - } - else if (logLevel == "warning") - { - logger->setLevel(trt::ILogger::Severity::kWARNING); - } - else if (logLevel == "error") - { - logger->setLevel(trt::ILogger::Severity::kERROR); - } - else if (logLevel == "internal_error") - { - logger->setLevel(trt::ILogger::Severity::kINTERNAL_ERROR); - } - else - { - TLLM_LOG_ERROR("Unexpected log level: " + logLevel); - return 1; - } - - initTrtLlmPlugins(logger.get()); - - // Argument: output sequences JSON - auto const responsesJsonFile = result["responses_json"].as(); - - // Argument: API - auto const api = result["api"].as(); - if (api == "executor") - { - texec::ModelType executorModelType; - std::optional decoderEngineDir = std::nullopt, encoderEngineDir = std::nullopt; - if (result.count("encoder_engine_dir") && result.count("decoder_engine_dir")) - { - TLLM_CHECK_WITH_INFO(api == "executor", "encoder-decoder only support executor api."); - TLLM_CHECK_WITH_INFO( - batchingType == texec::BatchingType::kINFLIGHT, "encoder-decoder only support inflight batching."); - executorModelType = texec::ModelType::kENCODER_DECODER; - encoderEngineDir = result["encoder_engine_dir"].as(); - decoderEngineDir = result["decoder_engine_dir"].as(); - } - else if (result.count("engine_dir")) - { - executorModelType = texec::ModelType::kDECODER_ONLY; - decoderEngineDir = result["engine_dir"].as(); - } - else - { - executorModelType = texec::ModelType::kENCODER_ONLY; - encoderEngineDir = result["encoder_engine_dir"].as(); - } - try - { - benchmarkExecutor(decoderEngineDir, encoderEngineDir, batchingType, datasetPath, opCsvFile, maxNumSamples, - beamWidth, result["warm_up"].as(), eosId, padId, benchmarkParams, capacitySchedulerPolicy, - waitSleep, returnContextLogits, returnGenerationLogits, staticEmulatedBatchSize, logIterationData, - maxPromptLen, executorModelType, responsesJsonFile); - } - catch (std::exception const& e) - { - TLLM_LOG_ERROR(e.what()); - return 1; - } - } - else if (api == "gptManager") - { - TLLM_LOG_ERROR("gptManager is deprecated, please use the executor API."); - return 1; - } - else - { - TLLM_LOG_ERROR("api parameter must be gptManager or executor"); - return 1; - } - - return 0; -} diff --git a/benchmarks/cpp/utils/utils.cpp b/benchmarks/cpp/utils/utils.cpp deleted file mode 100644 index 0cbcf1c0468d..000000000000 --- a/benchmarks/cpp/utils/utils.cpp +++ /dev/null @@ -1,170 +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. - */ - -#include "utils.h" -#include "tensorrt_llm/common/config.h" -#include "tensorrt_llm/common/logger.h" -#include - -#include -#include - -TRTLLM_NAMESPACE_BEGIN - -namespace benchmark -{ - -std::vector> parseVectorOfVectors(std::string const& input) -{ - std::vector> result; - std::regex outer_regex(R"(\[(.*?)\])"); - std::regex inner_regex(R"(\d+)"); - auto outer_begin = std::sregex_iterator(input.begin(), input.end(), outer_regex); - auto outer_end = std::sregex_iterator(); - - for (std::sregex_iterator i = outer_begin; i != outer_end; ++i) - { - std::smatch match = *i; - std::string inner_str = match.str(1); - std::vector inner_vec; - auto inner_begin = std::sregex_iterator(inner_str.begin(), inner_str.end(), inner_regex); - auto inner_end = std::sregex_iterator(); - - for (std::sregex_iterator j = inner_begin; j != inner_end; ++j) - { - std::smatch inner_match = *j; - inner_vec.push_back(std::stoi(inner_match.str())); - } - result.push_back(inner_vec); - } - return result; -} - -texec::LookaheadDecodingConfig parseLookaheadConfig(std::string const& input) -{ - std::regex regex("\\[ *(\\d+) *, *(\\d+) *, *(\\d+) *\\]"); - std::smatch match; - if (std::regex_match(input, match, regex)) - { - TLLM_CHECK(match.size() == 4); - auto w = std::stoi(match[1]); - auto n = std::stoi(match[2]); - auto g = std::stoi(match[3]); - return texec::LookaheadDecodingConfig(w, n, g); - } - else - { - TLLM_LOG_WARNING("cannot parse lookahead config from '%s'", input.c_str()); - return texec::LookaheadDecodingConfig(); - } -} - -Samples parseWorkloadJson( - std::filesystem::path const& datasetPath, int maxNumSamples, std::optional const maxPromptLen) -{ - auto constexpr allowExceptions = true; - auto constexpr ignoreComments = true; - TLLM_CHECK_WITH_INFO(std::filesystem::exists(datasetPath), "File does not exist: %s", datasetPath.c_str()); - std::ifstream jsonStream(datasetPath); - auto json = nlohmann::json::parse(jsonStream, nullptr, allowExceptions, ignoreComments); - - Samples samples; - - for (auto const& sample : json["samples"]) - { - if (samples.size() >= maxNumSamples) - break; - int32_t taskId = sample.count("task_id") ? sample["task_id"].template get() : -1; - auto input_ids(sample["input_ids"].template get>()); - if (maxPromptLen && (input_ids.size() > maxPromptLen.value())) - { - input_ids.resize(maxPromptLen.value()); - } - samples.emplace_back(Sample{std::move(input_ids), sample["output_len"], taskId}); - } - - if (samples.size() < maxNumSamples) - { - TLLM_LOG_WARNING( - "Dataset size %zu is smaller than given max_num_samples " - "%d, max_num_samples will be ignored.\n", - samples.size(), maxNumSamples); - } - return samples; -} - -std::vector generateRandomExponentialValues(int count, float lambda, int seed) -{ - // Set a constant seed for reproducibility - std::mt19937 gen(seed); - - // Create an exponential distribution object - std::exponential_distribution distribution(lambda); - - // Generate random numbers from the exponential distribution - std::vector randomValues; - for (int i = 0; i < count; ++i) - { - double randomValue = distribution(gen); - randomValues.push_back(randomValue); - } - - return randomValues; -} - -std::vector computeTimeDelays(BenchmarkParams const& benchmarkParams, int numDelays) -{ - std::vector timeDelays; - if (benchmarkParams.requestRate.has_value() && benchmarkParams.requestRate.value() > 0.0) - { - if (benchmarkParams.enableExpDelays) - { - timeDelays = generateRandomExponentialValues( - numDelays, benchmarkParams.requestRate.value(), benchmarkParams.randomSeed); - } - else - { - timeDelays.assign(numDelays, 1.0 / benchmarkParams.requestRate.value()); - } - } - else - { - timeDelays.assign(numDelays, 0.0); - } - - return timeDelays; -} - -std::ostream& operator<<(std::ostream& os, RecordTimeMetric const& metric) -{ - os << metric.mAvg << "," << metric.mMax << "," << metric.mMin << "," << metric.mP99 << "," << metric.mP90 << "," - << metric.mP50; - return os; -} - -std::ostream& operator<<(std::ostream& os, RecordBwMetric const& metric) -{ - os << metric.mAvg << "," << metric.mMax << "," << metric.mMin << "," << metric.mP99 << "," << metric.mP90 << "," - << metric.mP50; - return os; -} - -} // namespace benchmark - -TRTLLM_NAMESPACE_END diff --git a/benchmarks/cpp/utils/utils.h b/benchmarks/cpp/utils/utils.h deleted file mode 100644 index fba30fee69ae..000000000000 --- a/benchmarks/cpp/utils/utils.h +++ /dev/null @@ -1,244 +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. - */ - -#include "tensorrt_llm/common/config.h" -#include "tensorrt_llm/executor/executor.h" - -#include -#include -#include -#include -#include -#include -#include -#include - -#pragma once - -TRTLLM_NAMESPACE_BEGIN - -namespace benchmark -{ - -// using namespace tensorrt_llm::batch_manager; -using namespace tensorrt_llm::runtime; - -namespace texec = tensorrt_llm::executor; - -std::vector> parseVectorOfVectors(std::string const& input); - -texec::LookaheadDecodingConfig parseLookaheadConfig(std::string const& input); - -struct BenchmarkParams -{ - std::optional maxTokensInPagedKvCache{std::nullopt}; - std::optional freeGpuMemoryFraction{std::nullopt}; - std::vector> freeGpuMemoryFractions{std::nullopt}; - - std::optional crossKvCacheFraction{std::nullopt}; - bool enableTrtOverlap{false}; - bool enableBatchSizeTuning{false}; - bool enableMaxNumTokensTuning{false}; - bool enableBlockReuse{false}; - bool enableChunkedContext{true}; - bool streaming{false}; - bool enableExpDelays{false}; - std::vector> enableChunekedContextVec{std::nullopt}; - std::optional requestRate{std::nullopt}; - std::optional concurrency{std::nullopt}; - std::optional maxBatchSize{std::nullopt}; - std::vector> maxBatchSizes{std::nullopt}; - std::optional maxNumTokens{std::nullopt}; - std::vector> maxNumTokensVec{std::nullopt}; - int randomSeed = 430; - std::optional> maxAttentionWindowVec{std::nullopt}; - std::optional sinkTokenLength{std::nullopt}; - bool multiBlockMode{true}; - bool enableContextFMHAFP32Acc{false}; - bool cudaGraphMode{false}; - SizeType32 cudaGraphCacheSize{0}; - - // lora / peft params - std::optional loraDir{std::nullopt}; - SizeType32 loraDeviceNumModLayers{0}; - size_t loraHostCacheSize{1024 * 2024 * 1024}; - - // KV cache block offloading - size_t kvHostCacheSize{0}; - - // Weights offloading - float gpuWeightsPercent{1.0}; - - // Decoding params - std::optional>> medusaChoices; - - std::optional eagleConfig; - std::optional temperature; - - std::optional executorLookaheadConfig; - std::optional requestLookaheadConfig; - - bool enableCollectkvCacheTransferTime = false; - bool enableCollectIterStats = false; -}; - -struct RecordTimeMetric -{ - - RecordTimeMetric(std::string tag) - : mTag(std::move(tag)) - { - } - - std::string mTag; - - std::vector mDataTimes; - - float mAvg; - float mP99; - float mP95; - float mP90; - float mP50; - float mMax; - float mMin; - - static float calcPercentile(std::vector const& latencies, int percentile) - { - int const index = static_cast(std::ceil((percentile / 100.0) * latencies.size())) - 1; - return latencies[index]; - } - - void calculate() - { - TLLM_CHECK_WITH_INFO(mDataTimes.size() > 0, "No data to calculate for tag:%s", mTag.c_str()); - mAvg = std::accumulate(mDataTimes.begin(), mDataTimes.end(), 0.F) / mDataTimes.size(); - - std::sort(mDataTimes.begin(), mDataTimes.end()); - - mP99 = calcPercentile(mDataTimes, 99); - mP90 = calcPercentile(mDataTimes, 90); - mP50 = calcPercentile(mDataTimes, 50); - mMax = mDataTimes.back(); - mMin = mDataTimes.front(); - } - - void report() const - { - - printf("[BENCHMARK] avg_%s(ms) %.2f\n", mTag.c_str(), mAvg); - printf("[BENCHMARK] max_%s(ms) %.2f\n", mTag.c_str(), mMax); - printf("[BENCHMARK] min_%s(ms) %.2f\n", mTag.c_str(), mMin); - - printf("[BENCHMARK] p99_%s(ms) %.2f\n", mTag.c_str(), mP99); - - printf("[BENCHMARK] p90_%s(ms) %.2f\n", mTag.c_str(), mP90); - - printf("[BENCHMARK] p50_%s(ms) %.2f\n\n", mTag.c_str(), mP50); - } - - std::vector genHeaders() const - { - std::string timeTag = mTag + "(ms)"; - return { - "avg_" + timeTag, "max_" + timeTag, "min_" + timeTag, "p99" + timeTag, "p90" + timeTag, "p50" + timeTag}; - } -}; - -struct RecordBwMetric -{ - - RecordBwMetric(std::string tag) - : mTag(std::move(tag)) - { - } - - std::string mTag; - - std::vector mDataTps; - - float mAvg; - float mP99; - float mP95; - float mP90; - float mP50; - float mMax; - float mMin; - - static float calcPercentile(std::vector const& throughputs, int percentile) - { - int const index = static_cast(std::ceil((percentile / 100.0) * throughputs.size())) - 1; - return throughputs[index]; - } - - void calculate() - { - TLLM_CHECK_WITH_INFO(mDataTps.size() > 0, "No data to calculate for tag:%s", mTag.c_str()); - mAvg = std::accumulate(mDataTps.begin(), mDataTps.end(), 0.F) / mDataTps.size(); - - std::sort(mDataTps.begin(), mDataTps.end(), std::greater()); - - mP99 = calcPercentile(mDataTps, 99); - mP90 = calcPercentile(mDataTps, 90); - mP50 = calcPercentile(mDataTps, 50); - mMax = mDataTps.front(); - mMin = mDataTps.back(); - } - - void report() const - { - - printf("[BENCHMARK] avg_%s(Gb/sec) %.8f\n", mTag.c_str(), mAvg); - printf("[BENCHMARK] max_%s(Gb/sec) %.8f\n", mTag.c_str(), mMax); - printf("[BENCHMARK] min_%s(Gb/sec) %.8f\n", mTag.c_str(), mMin); - - printf("[BENCHMARK] p99_%s(Gb/sec) %.8f\n", mTag.c_str(), mP99); - - printf("[BENCHMARK] p90_%s(Gb/sec) %.8f\n", mTag.c_str(), mP90); - - printf("[BENCHMARK] p50_%s(Gb/sec) %.8f\n\n", mTag.c_str(), mP50); - } - - std::vector genHeaders() const - { - std::string tpTag = mTag + "(Gb/sec)"; - return {"avg_" + tpTag, "max_" + tpTag, "min_" + tpTag, "p99" + tpTag, "p90" + tpTag, "p50" + tpTag}; - } -}; - -std::ostream& operator<<(std::ostream& os, RecordTimeMetric const& metric); -std::ostream& operator<<(std::ostream& os, RecordBwMetric const& metric); - -struct Sample -{ - std::vector inputIds; - int32_t outputLen; - int32_t taskId; -}; - -using Samples = std::vector; - -Samples parseWorkloadJson( - std::filesystem::path const& datasetPath, int maxNumSamples, std::optional const maxPromptLen); - -std::vector generateRandomExponentialValues(int count, float lambda, int seed); - -std::vector computeTimeDelays(BenchmarkParams const& benchmarkParams, int numDelays); - -} // namespace benchmark - -TRTLLM_NAMESPACE_END diff --git a/benchmarks/cpp/prepare_dataset.py b/benchmarks/prepare_dataset.py similarity index 95% rename from benchmarks/cpp/prepare_dataset.py rename to benchmarks/prepare_dataset.py index 3b9665fd2902..a1f1358f11ff 100644 --- a/benchmarks/cpp/prepare_dataset.py +++ b/benchmarks/prepare_dataset.py @@ -55,7 +55,7 @@ def validate_tokenizer(self): required=True, type=str, help= - "Tokenizer dir for the model run by gptManagerBenchmark, or the model name from HuggingFace." + "Tokenizer dir for the benchmarked model, or the model name from HuggingFace." ) @click.option("--output", type=str, @@ -92,7 +92,7 @@ def validate_tokenizer(self): help="Trust remote code.") @click.pass_context def cli(ctx, **kwargs): - """This script generates dataset input for gptManagerBenchmark.""" + """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': diff --git a/benchmarks/cpp/utils/__init__.py b/benchmarks/utils/__init__.py similarity index 100% rename from benchmarks/cpp/utils/__init__.py rename to benchmarks/utils/__init__.py diff --git a/benchmarks/cpp/utils/convert_nemo_dataset.py b/benchmarks/utils/convert_nemo_dataset.py similarity index 100% rename from benchmarks/cpp/utils/convert_nemo_dataset.py rename to benchmarks/utils/convert_nemo_dataset.py diff --git a/benchmarks/cpp/utils/generate_rand_loras.py b/benchmarks/utils/generate_rand_loras.py similarity index 100% rename from benchmarks/cpp/utils/generate_rand_loras.py rename to benchmarks/utils/generate_rand_loras.py diff --git a/benchmarks/cpp/utils/prepare_real_data.py b/benchmarks/utils/prepare_real_data.py similarity index 100% rename from benchmarks/cpp/utils/prepare_real_data.py rename to benchmarks/utils/prepare_real_data.py diff --git a/benchmarks/cpp/utils/prepare_synthetic_data.py b/benchmarks/utils/prepare_synthetic_data.py similarity index 100% rename from benchmarks/cpp/utils/prepare_synthetic_data.py rename to benchmarks/utils/prepare_synthetic_data.py diff --git a/benchmarks/cpp/utils/utils.py b/benchmarks/utils/utils.py similarity index 100% rename from benchmarks/cpp/utils/utils.py rename to benchmarks/utils/utils.py diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index a323b32b82b5..5dec92599ae6 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -45,7 +45,6 @@ add_compile_definitions(TRTLLM_ABI_NAMESPACE=${TRTLLM_ABI_NAMESPACE}) # Build options option(BUILD_PYT "Build in PyTorch TorchScript class mode" ON) option(BUILD_TESTS "Build Google tests" ON) -option(BUILD_BENCHMARKS "Build benchmarks" ON) option(BUILD_DEEP_EP "Build the Deep EP module" ON) option(BUILD_DEEP_GEMM "Build the DeepGEMM module" ON) option(BUILD_FLASH_MLA "Build the FlashMLA module" ON) @@ -126,12 +125,6 @@ else() message(STATUS "Not building Google tests") endif() -if(BUILD_BENCHMARKS) - message(STATUS "Building benchmarks") -else() - message(STATUS "Not building benchmarks") -endif() - if(BUILD_MICRO_BENCHMARKS) message(STATUS "Building C++ micro benchmarks") else() @@ -249,8 +242,6 @@ if(ENABLE_MULTI_DEVICE) endif() # TRT dependencies -find_package(TensorRT 10 REQUIRED COMPONENTS OnnxParser) -set(TRT_LIB TensorRT::NvInfer) get_filename_component(TRT_LLM_ROOT_DIR ${CMAKE_CURRENT_SOURCE_DIR} PATH) @@ -291,7 +282,6 @@ include_directories( ${CUDAToolkit_INCLUDE_DIRS} ${CUDAToolkit_INCLUDE_DIRS}/cccl ${CUDNN_ROOT_DIR}/include - $ ${maybe_nvtx_includedir} ${CMAKE_BINARY_DIR}/_deps/cutlass-src/include ${CMAKE_BINARY_DIR}/_deps/cutlass-src/tools/util/include @@ -669,11 +659,6 @@ if(BUILD_TESTS) add_subdirectory(tests) endif() -if(BUILD_BENCHMARKS) - add_subdirectory(${TRT_LLM_ROOT_DIR}/benchmarks/cpp - ${CMAKE_BINARY_DIR}/benchmarks) -endif() - if(BUILD_MICRO_BENCHMARKS) add_subdirectory(${TRT_LLM_ROOT_DIR}/cpp/micro_benchmarks ${CMAKE_BINARY_DIR}/micro_benchmarks) @@ -688,6 +673,6 @@ if(MEASURE_BUILD_TIME) endif() set(BUILD_WHEEL_TARGETS - tensorrt_llm;nvinfer_plugin_tensorrt_llm + tensorrt_llm CACHE STRING "Targets used to build wheel") add_custom_target(build_wheel_targets DEPENDS ${BUILD_WHEEL_TARGETS}) diff --git a/cpp/cmake/modules/FindTensorRT.cmake b/cpp/cmake/modules/FindTensorRT.cmake deleted file mode 100644 index 9e7e35b51bae..000000000000 --- a/cpp/cmake/modules/FindTensorRT.cmake +++ /dev/null @@ -1,190 +0,0 @@ -# -# SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. -# - -# TensorRT install path in docker image -set(TensorRT_WELL_KNOWN_ROOT /usr/local/tensorrt) - -find_path( - TensorRT_INCLUDE_DIR - NAMES NvInfer.h - PATHS ${TensorRT_WELL_KNOWN_ROOT}/include) - -function(_tensorrt_get_version) - unset(TensorRT_VERSION_STRING PARENT_SCOPE) - set(_hdr_file "${TensorRT_INCLUDE_DIR}/NvInferVersion.h") - - if(NOT EXISTS "${_hdr_file}") - return() - endif() - - file(STRINGS "${_hdr_file}" IS_10_11_NEW_MACRO REGEX "TRT_MAJOR_ENTERPRISE") - if(IS_10_11_NEW_MACRO) - file(STRINGS "${_hdr_file}" VERSION_STRINGS - REGEX "#define TRT_.+_ENTERPRISE.*") - foreach(TYPE MAJOR MINOR PATCH BUILD) - string(REGEX MATCH "TRT_${TYPE}_ENTERPRISE [0-9]+" TRT_TYPE_STRING - ${VERSION_STRINGS}) - string(REGEX MATCH "[0-9]+" TensorRT_VERSION_${TYPE} ${TRT_TYPE_STRING}) - endforeach(TYPE) - else() - file(STRINGS "${_hdr_file}" VERSION_STRINGS REGEX "#define NV_TENSORRT_.*") - foreach(TYPE MAJOR MINOR PATCH BUILD) - string(REGEX MATCH "NV_TENSORRT_${TYPE} [0-9]+" TRT_TYPE_STRING - ${VERSION_STRINGS}) - string(REGEX MATCH "[0-9]+" TensorRT_VERSION_${TYPE} ${TRT_TYPE_STRING}) - endforeach(TYPE) - endif() - - set(TensorRT_VERSION_MAJOR - ${TensorRT_VERSION_MAJOR} - PARENT_SCOPE) - set(TensorRT_VERSION_STRING - "${TensorRT_VERSION_MAJOR}.${TensorRT_VERSION_MINOR}.${TensorRT_VERSION_PATCH}.${TensorRT_VERSION_BUILD}" - PARENT_SCOPE) -endfunction(_tensorrt_get_version) - -_tensorrt_get_version() - -macro(_tensorrt_find_dll VAR) - find_file( - ${VAR} - NAMES ${ARGN} - HINTS ${TensorRT_ROOT} - PATH_SUFFIXES bin) -endmacro(_tensorrt_find_dll) - -find_library( - TensorRT_LIBRARY - NAMES "nvinfer_${TensorRT_VERSION_MAJOR}" nvinfer - PATHS ${TensorRT_WELL_KNOWN_ROOT}/lib) - -if(WIN32) - _tensorrt_find_dll(TensorRT_DLL "nvinfer_${TensorRT_VERSION_MAJOR}.dll" - nvinfer.dll) -endif() - -if(TensorRT_LIBRARY) - set(TensorRT_LIBRARIES ${TensorRT_LIBRARIES} ${TensorRT_LIBRARY}) -endif(TensorRT_LIBRARY) - -if(TensorRT_FIND_COMPONENTS) - list(REMOVE_ITEM TensorRT_FIND_COMPONENTS "nvinfer") - - if("OnnxParser" IN_LIST TensorRT_FIND_COMPONENTS) - find_path( - TensorRT_OnnxParser_INCLUDE_DIR - NAMES NvOnnxParser.h - PATHS ${TensorRT_WELL_KNOWN_ROOT}/include) - - find_library( - TensorRT_OnnxParser_LIBRARY - NAMES "nvonnxparser_${TensorRT_VERSION_MAJOR}" nvonnxparser - PATHS ${TensorRT_WELL_KNOWN_ROOT}/lib) - if(TensorRT_OnnxParser_LIBRARY AND TensorRT_LIBRARIES) - set(TensorRT_LIBRARIES ${TensorRT_LIBRARIES} - ${TensorRT_OnnxParser_LIBRARY}) - set(TensorRT_OnnxParser_FOUND TRUE) - endif() - - if(WIN32) - _tensorrt_find_dll( - TensorRT_OnnxParser_DLL "nvonnxparser_${TensorRT_VERSION_MAJOR}.dll" - nvonnxparser.dll) - endif() - endif() - - if("Plugin" IN_LIST TensorRT_FIND_COMPONENTS) - find_path( - TensorRT_Plugin_INCLUDE_DIR - NAMES NvInferPlugin.h - PATHS ${TensorRT_WELL_KNOWN_ROOT}/include) - - find_library( - TensorRT_Plugin_LIBRARY - NAMES "nvinfer_plugin_${TensorRT_VERSION_MAJOR}" nvinfer_plugin - PATHS ${TensorRT_WELL_KNOWN_ROOT}/lib) - - if(TensorRT_Plugin_LIBRARY AND TensorRT_LIBRARIES) - set(TensorRT_LIBRARIES ${TensorRT_LIBRARIES} ${TensorRT_Plugin_LIBRARY}) - set(TensorRT_Plugin_FOUND TRUE) - endif() - - if(WIN32) - _tensorrt_find_dll( - TensorRT_Plugin_DLL "nvinfer_plugin_${TensorRT_VERSION_MAJOR}.dll" - nvinfer_plugin.dll) - endif() - endif() -endif() - -include(FindPackageHandleStandardArgs) -find_package_handle_standard_args( - TensorRT - FOUND_VAR TensorRT_FOUND - REQUIRED_VARS TensorRT_LIBRARY TensorRT_LIBRARIES TensorRT_INCLUDE_DIR - VERSION_VAR TensorRT_VERSION_STRING - HANDLE_COMPONENTS) - -if(NOT TARGET TensorRT::NvInfer) - add_library(TensorRT::NvInfer SHARED IMPORTED) - target_include_directories(TensorRT::NvInfer SYSTEM - INTERFACE "${TensorRT_INCLUDE_DIR}") - if(WIN32) - set_property(TARGET TensorRT::NvInfer PROPERTY IMPORTED_LOCATION - "${TensorRT_DLL}") - set_property(TARGET TensorRT::NvInfer PROPERTY IMPORTED_IMPLIB - "${TensorRT_LIBRARY}") - else() - set_property(TARGET TensorRT::NvInfer PROPERTY IMPORTED_LOCATION - "${TensorRT_LIBRARY}") - endif() -endif() - -if(NOT TARGET TensorRT::OnnxParser AND "OnnxParser" IN_LIST - TensorRT_FIND_COMPONENTS) - add_library(TensorRT::OnnxParser SHARED IMPORTED) - target_include_directories(TensorRT::OnnxParser SYSTEM - INTERFACE "${TensorRT_OnnxParser_INCLUDE_DIR}") - target_link_libraries(TensorRT::OnnxParser INTERFACE TensorRT::NvInfer) - if(WIN32) - set_property(TARGET TensorRT::OnnxParser - PROPERTY IMPORTED_LOCATION "${TensorRT_OnnxParser_DLL}") - set_property(TARGET TensorRT::OnnxParser - PROPERTY IMPORTED_IMPLIB "${TensorRT_OnnxParser_LIBRARY}") - else() - set_property(TARGET TensorRT::OnnxParser - PROPERTY IMPORTED_LOCATION "${TensorRT_OnnxParser_LIBRARY}") - endif() -endif() - -if(NOT TARGET TensorRT::Plugin AND "Plugin" IN_LIST TensorRT_FIND_COMPONENTS) - add_library(TensorRT::Plugin SHARED IMPORTED) - target_include_directories(TensorRT::Plugin SYSTEM - INTERFACE "${TensorRT_Plugin_INCLUDE_DIR}") - target_link_libraries(TensorRT::Plugin INTERFACE TensorRT::NvInfer) - if(WIN32) - set_property(TARGET TensorRT::Plugin PROPERTY IMPORTED_LOCATION - "${TensorRT_Plugin_DLL}") - set_property(TARGET TensorRT::Plugin PROPERTY IMPORTED_IMPLIB - "${TensorRT_Plugin_LIBRARY}") - else() - set_property(TARGET TensorRT::Plugin PROPERTY IMPORTED_LOCATION - "${TensorRT_Plugin_LIBRARY}") - endif() -endif() - -mark_as_advanced(TensorRT_INCLUDE_DIR TensorRT_LIBRARY TensorRT_LIBRARIES) diff --git a/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h b/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h index 0fcee1435005..9d28fa26c4ee 100644 --- a/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h +++ b/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h @@ -22,6 +22,7 @@ #include "tensorrt_llm/batch_manager/llmRequest.h" #include "tensorrt_llm/batch_manager/rnnCacheTransBuffer.h" #include "tensorrt_llm/batch_manager/rnnStateManager.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/executor/cacheCommunicator.h" #include "tensorrt_llm/executor/dataTransceiverState.h" #include "tensorrt_llm/runtime/utils/mpiUtils.h" @@ -238,7 +239,7 @@ class CacheTransceiver : public BaseCacheTransceiver public: CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheManager, executor::kv_cache::CacheState::ModelConfig const& cacheStateModelCfg, runtime::WorldConfig const& worldConfig, - std::vector const& attentionLayerNumPerPP, nvinfer1::DataType dataType, + std::vector const& attentionLayerNumPerPP, tensorrt_llm::DataType dataType, executor::kv_cache::CacheState::AttentionType attentionType = executor::kv_cache::CacheState::AttentionType::kDEFAULT, std::optional cacheTransceiverConfig = std::nullopt, @@ -246,7 +247,7 @@ class CacheTransceiver : public BaseCacheTransceiver CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheManager, std::vector numKvHeadsPerLayer, SizeType32 sizePerHead, SizeType32 tokensPerBlock, runtime::WorldConfig const& worldConfig, - std::vector const& attentionLayerNumPerPP, nvinfer1::DataType dataType, + std::vector const& attentionLayerNumPerPP, tensorrt_llm::DataType dataType, executor::kv_cache::CacheState::AttentionType attentionType = executor::kv_cache::CacheState::AttentionType::kDEFAULT, std::optional cacheTransceiverConfig = std::nullopt, diff --git a/cpp/include/tensorrt_llm/batch_manager/createNewDecoderRequests.h b/cpp/include/tensorrt_llm/batch_manager/createNewDecoderRequests.h index bc619a34bc03..600927af9645 100644 --- a/cpp/include/tensorrt_llm/batch_manager/createNewDecoderRequests.h +++ b/cpp/include/tensorrt_llm/batch_manager/createNewDecoderRequests.h @@ -20,6 +20,7 @@ #include "tensorrt_llm/batch_manager/common.h" #include "tensorrt_llm/common/algorithm.h" #include "tensorrt_llm/common/optionalRef.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/executor/executor.h" #include "tensorrt_llm/runtime/common.h" #include "tensorrt_llm/runtime/iTensor.h" @@ -66,16 +67,17 @@ class CreateNewDecoderRequests : Algorithm std::vector> operator()(runtime::ModelConfig const& modelConfig, runtime::WorldConfig const& worldConfig, executor::DecodingConfig const& decodingConfig, RequestVector const& contextRequests, - nvinfer1::DataType logitsType, DecoderInputBuffers& inputBuffers, runtime::decoder::DecoderState& decoderState, - CudaStream const& runtimeStream, CudaStream const& decoderStream, SizeType32 maxSequenceLength, - SizeType32 beamWidth, OptionalRef medusaBuffers) const; + tensorrt_llm::DataType logitsType, DecoderInputBuffers& inputBuffers, + runtime::decoder::DecoderState& decoderState, CudaStream const& runtimeStream, CudaStream const& decoderStream, + SizeType32 maxSequenceLength, SizeType32 beamWidth, OptionalRef medusaBuffers) const; [[nodiscard]] std::tuple, std::vector> createDecoderRequests(RequestVector const& finishedContextRequests, TensorPtr const& inputIds, executor::DecodingConfig const& decodingConfig, runtime::decoder::DecoderState& decoderState, - nvinfer1::DataType logitsType, runtime::ModelConfig const& modelConfig, runtime::WorldConfig const& worldConfig, - runtime::CudaStream const& runtimeStream, runtime::CudaStream const& decoderStream, - SizeType32 maxSequenceLength, OptionalRef medusaBuffers) const; + tensorrt_llm::DataType logitsType, runtime::ModelConfig const& modelConfig, + runtime::WorldConfig const& worldConfig, runtime::CudaStream const& runtimeStream, + runtime::CudaStream const& decoderStream, SizeType32 maxSequenceLength, + OptionalRef medusaBuffers) const; private: bool mSpeculativeDecodingFastLogits; diff --git a/cpp/include/tensorrt_llm/batch_manager/guidedDecoder.h b/cpp/include/tensorrt_llm/batch_manager/guidedDecoder.h deleted file mode 100644 index 9a577b61ad51..000000000000 --- a/cpp/include/tensorrt_llm/batch_manager/guidedDecoder.h +++ /dev/null @@ -1,67 +0,0 @@ -/* - * Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. - * - * 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. - */ - -#pragma once - -#include "tensorrt_llm/batch_manager/common.h" -#include "tensorrt_llm/executor/executor.h" -#include "tensorrt_llm/runtime/bufferManager.h" -#include "tensorrt_llm/runtime/iTensor.h" - -namespace xgrammar -{ -class GrammarMatcher; -class GrammarCompiler; -} // namespace xgrammar - -namespace tensorrt_llm::batch_manager -{ -class DecoderInputBuffers; - -class GuidedDecoder -{ -public: - using TensorPtr = runtime::ITensor::SharedPtr; - using SizeType32 = tensorrt_llm::runtime::SizeType32; - using BitmaskT = uint32_t; - - GuidedDecoder(executor::GuidedDecodingConfig const& guidedDecodingConfig, SizeType32 maxNumSequences, - SizeType32 vocabSizePadded, nvinfer1::DataType logitsDtype, runtime::BufferManager const& runtimeBufferManager); - void build(ScheduledRequests const& scheduledRequests); - void execute(DecoderInputBuffers const& decoderInputBuffers, runtime::BufferManager const& runtimeBufferManager); - -private: - executor::GuidedDecodingConfig::GuidedDecodingBackend mGuidedDecodingBackend; - std::vector> mXGrammarMatchers; - std::shared_ptr mXGrammarCompiler; - - SizeType32 mMaxNumSequences; - SizeType32 mVocabSizePadded; - SizeType32 mBitmaskSize; // CeilDiv(vocabSizePadded, 32) - nvinfer1::DataType mLogitsDtype; - - TensorPtr mLogitsBitmask; // [mMaxNumRequests, mBitmaskSize] - TensorPtr mLogitsBitmaskHost; // [mMaxNumRequests, mBitmaskSize] - TensorPtr mLogitsBitmaskPtrVec; // [mMaxNumRequests], pointers to the logitsBitmask in a batch - TensorPtr mLogitsBitmaskPtrVecHost; // [mMaxNumRequests] - TensorPtr mLogitsPtrVec; // [mMaxNumRequests], pointers to the logits in a batch - TensorPtr mLogitsPtrVecHost; // [mMaxNumRequests] - - // BufferManager with a dedicated stream for async copy of buffers for guided decoding. - runtime::BufferManager mCopyBufferManager; -}; - -} // namespace tensorrt_llm::batch_manager diff --git a/cpp/include/tensorrt_llm/batch_manager/handleContextLogits.h b/cpp/include/tensorrt_llm/batch_manager/handleContextLogits.h deleted file mode 100644 index cb77545578c8..000000000000 --- a/cpp/include/tensorrt_llm/batch_manager/handleContextLogits.h +++ /dev/null @@ -1,53 +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. - */ - -#pragma once - -#include "tensorrt_llm/batch_manager/common.h" -#include "tensorrt_llm/common/algorithm.h" -#include "tensorrt_llm/common/optionalRef.h" -#include "tensorrt_llm/runtime/modelConfig.h" - -namespace tensorrt_llm::runtime -{ -class BufferManager; -class CudaStream; -} // namespace tensorrt_llm::runtime - -namespace tensorrt_llm::batch_manager -{ - -class DecoderInputBuffers; -class MedusaBuffers; - -class HandleContextLogits : Algorithm -{ -public: - template - using OptionalRef = tensorrt_llm::common::OptionalRef; - - constexpr static auto name{"HandleContextLogits"}; - - HandleContextLogits() = default; - - runtime::SizeType32 operator()(DecoderInputBuffers& inputBuffers, RequestVector const& contextRequests, - runtime::ITensor::SharedPtr const& logits, std::vector const& numContextLogitsVec, - runtime::ModelConfig const& modelConfig, runtime::BufferManager const& manager, - OptionalRef medusaBuffers) const; -}; - -} // namespace tensorrt_llm::batch_manager diff --git a/cpp/include/tensorrt_llm/batch_manager/handleGenerationLogits.h b/cpp/include/tensorrt_llm/batch_manager/handleGenerationLogits.h deleted file mode 100644 index f9fd58800a6f..000000000000 --- a/cpp/include/tensorrt_llm/batch_manager/handleGenerationLogits.h +++ /dev/null @@ -1,53 +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. - */ - -#pragma once - -#include "common.h" -#include "tensorrt_llm/common/algorithm.h" -#include "tensorrt_llm/common/optionalRef.h" -#include "tensorrt_llm/runtime/modelConfig.h" - -namespace tensorrt_llm::runtime -{ -class BufferManager; -} // namespace tensorrt_llm::runtime - -namespace tensorrt_llm::batch_manager -{ - -class DecoderInputBuffers; -class RuntimeBuffers; -class MedusaBuffers; - -class HandleGenerationLogits : Algorithm -{ -public: - template - using OptionalRef = tensorrt_llm::common::OptionalRef; - - constexpr static auto name{"HandleGenerationLogits"}; - - HandleGenerationLogits() = default; - - void operator()(DecoderInputBuffers& inputBuffers, RequestVector const& generationRequests, - runtime::ITensor::SharedPtr const& logits, runtime::SizeType32 logitsIndex, - runtime::ModelConfig const& modelConfig, runtime::BufferManager const& manager, - OptionalRef genRuntimeBuffers, OptionalRef medusaBuffers) const; -}; - -} // namespace tensorrt_llm::batch_manager diff --git a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h index e821cafd5c3f..01adf276f878 100644 --- a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h +++ b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h @@ -23,6 +23,7 @@ #include "tensorrt_llm/batch_manager/llmRequest.h" // TODO forward declare #include "tensorrt_llm/batch_manager/radixBlockTree.h" #include "tensorrt_llm/common/optionalRef.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/executor/executor.h" #include "tensorrt_llm/executor/transferAgent.h" #include "tensorrt_llm/kernels/kvCacheIndex.h" @@ -32,7 +33,6 @@ #include "tensorrt_llm/runtime/iBuffer.h" #include "tensorrt_llm/runtime/iTensor.h" #include "tensorrt_llm/runtime/worldConfig.h" -#include #include #include @@ -140,7 +140,7 @@ struct PoolConfiguration { SizeType32 windowSize; SizeType32 sizePerHead; - nvinfer1::DataType dtype; + tensorrt_llm::DataType dtype; }; struct LinearAttentionMetadata @@ -872,7 +872,7 @@ class WindowBlockManager using BlockMap = std::unordered_multimap; using BlockMapIterRange = std::pair; - explicit WindowBlockManager(nvinfer1::DataType dtype, SizeType32 windowSize, + explicit WindowBlockManager(tensorrt_llm::DataType dtype, SizeType32 windowSize, std::vector const& managedLayers, std::vector const& numKvHeadsPerLayer, SizeType32 sizePerHead, SizeType32 tokensPerBlock, bool isSWA, SizeType32 blocksInPrimaryPool, SizeType32 blocksInSecondaryPool, SizeType32 maxNumSequences, std::shared_ptr stream, @@ -1037,7 +1037,7 @@ class WindowBlockManager //! host pools with mixed precisions when constructed with a per-window //! dtype map. Empty pools or NVFP4-scale pools are routed through the //! per-pool tensor metadata instead. - [[nodiscard]] nvinfer1::DataType getDataType() const noexcept + [[nodiscard]] tensorrt_llm::DataType getDataType() const noexcept { return mDataType; } @@ -1127,7 +1127,7 @@ class WindowBlockManager [[nodiscard]] SizeType32 getNumEltsPerContainer() const { #ifdef ENABLE_FP4 - return mDataType == nvinfer1::DataType::kFP4 ? 2 : 1; + return mDataType == tensorrt_llm::DataType::kFP4 ? 2 : 1; #else return 1; #endif @@ -1192,7 +1192,7 @@ class WindowBlockManager return mLayerToIndexWithinPool.at(layerIdx); } - void setOffsets(kernels::KVCacheIndex* offsetsPtr, nvinfer1::Dims const& offsetsShape, SizeType32 beamIdx, + void setOffsets(kernels::KVCacheIndex* offsetsPtr, tensorrt_llm::Dims const& offsetsShape, SizeType32 beamIdx, SizeType32 blockIdx, KVCacheBlock::IdType blockId) const; //! \brief Bring offloaded block from secondary to primary memory. @@ -1353,7 +1353,7 @@ class WindowBlockManager } private: - nvinfer1::DataType mDataType; + tensorrt_llm::DataType mDataType; SizeType32 mWindowSize; // Number of blocks in pools @@ -1481,7 +1481,7 @@ class BlockManager explicit BlockManager(std::vector const& numKvHeadsPerLayer, SizeType32 sizePerHead, SizeType32 tokensPerBlock, BlocksPerWindow const& blocksPerWindow, SizeType32 maxNumSequences, CudaStreamPtr stream, SizeType32 maxSequenceLength, SizeType32 maxBeamWidth, - std::vector const& maxAttentionWindowVec, nvinfer1::DataType dtype, SizeType32 sinkBubbleLength, + std::vector const& maxAttentionWindowVec, tensorrt_llm::DataType dtype, SizeType32 sinkBubbleLength, SizeType32 chunkSize, CacheType cacheType = CacheType::kSELF, std::optional secondaryOffloadMinPriority = std::nullopt, std::shared_ptr eventManager = nullptr, bool enablePartialReuse = true, @@ -1563,7 +1563,7 @@ class BlockManager void releaseLastBlock(GenerationRequest& sequence, SizeType32 windowSize); - void setOffsets(kernels::KVCacheIndex* offsetsPtr, nvinfer1::Dims const& offsetsShape, SizeType32 beamIdx, + void setOffsets(kernels::KVCacheIndex* offsetsPtr, tensorrt_llm::Dims const& offsetsShape, SizeType32 beamIdx, SizeType32 blockIdx, KVCacheBlock::IdType blockId, SizeType32 windowSize) const; //! \brief Combined prefix reuse analysis — single radix tree walk. @@ -1621,9 +1621,9 @@ class BlockManager //! \brief Convenience: window_size -> dataType, derived from getPoolConfigurations(). //! For one-pool-per-window managers only; multi-pool-per-window will collide. - [[nodiscard]] std::map getDataTypePerWindow() const + [[nodiscard]] std::map getDataTypePerWindow() const { - std::map result; + std::map result; for (auto const& [windowSize, manager] : mWindowBlockManagers) { result[windowSize] = manager.getDataType(); @@ -1636,7 +1636,7 @@ class BlockManager return mWindowBlockManagers.at(windowSize).getSizePerHead(); } - [[nodiscard]] nvinfer1::DataType getDataTypeForWindow(SizeType32 windowSize) const + [[nodiscard]] tensorrt_llm::DataType getDataTypeForWindow(SizeType32 windowSize) const { return mWindowBlockManagers.at(windowSize).getDataType(); } @@ -2189,7 +2189,7 @@ class BaseKVCacheManager /// head_dim=512). Empty vector = uniform @p sizePerHead / @p dtype across all windows. /// @return Map from window size to tuple of (primary blocks, secondary blocks) [[nodiscard]] static BlocksPerWindow calculateMaxNumBlocks(executor::KvCacheConfig const& config, - nvinfer1::DataType dtype, std::vector const& numKvHeadsPerLayer, SizeType32 sizePerHead, + tensorrt_llm::DataType dtype, std::vector const& numKvHeadsPerLayer, SizeType32 sizePerHead, SizeType32 tokensPerBlock, tensorrt_llm::runtime::WorldConfig const& worldConfig, std::map> const& windowSizeToLayers, uint64_t allottedPrimaryMemBytes, uint64_t allottedSecondaryMemBytes, size_t extraCostMemory, SizeType32 kvFactor, SizeType32 maxBatchSize, @@ -2276,7 +2276,7 @@ class KVCacheManager : public BaseKVCacheManager //! and disagg transfer machinery applies natively. Empty vector = uniform. KVCacheManager(std::vector const& numKvHeadsPerLayer, SizeType32 sizePerHead, SizeType32 tokensPerBlock, BlocksPerWindow const& blocksPerWindow, SizeType32 maxNumSequences, SizeType32 maxBeamWidth, - std::vector const& maxAttentionWindowVec, nvinfer1::DataType dtype, SizeType32 sinkTokenLength, + std::vector const& maxAttentionWindowVec, tensorrt_llm::DataType dtype, SizeType32 sinkTokenLength, CudaStreamPtr stream, SizeType32 maxSequenceLength, SizeType32 chunkSize, bool enableBlockReuse = false, CacheType cacheType = CacheType::kSELF, std::optional secondaryOffloadMinPriority = std::nullopt, @@ -2290,7 +2290,7 @@ class KVCacheManager : public BaseKVCacheManager KVCacheManager(std::vector const& numKvHeadsPerLayer, SizeType32 sizePerHead, SizeType32 tokensPerBlock, BlocksPerWindow const& blocksPerWindow, SizeType32 maxNumSequences, SizeType32 maxBeamWidth, - std::vector const& maxAttentionWindowVec, nvinfer1::DataType dtype, SizeType32 sinkTokenLength, + std::vector const& maxAttentionWindowVec, tensorrt_llm::DataType dtype, SizeType32 sinkTokenLength, int64_t stream, SizeType32 maxSequenceLength, SizeType32 chunkSize, bool enableBlockReuse = false, CacheType cacheType = CacheType::kSELF, std::optional secondaryOffloadMinPriority = std::nullopt, @@ -2304,7 +2304,7 @@ class KVCacheManager : public BaseKVCacheManager KVCacheManager(SizeType32 numLayers, SizeType32 numKvHeads, SizeType32 sizePerHead, SizeType32 tokensPerBlock, BlocksPerWindow const& blocksPerWindow, SizeType32 maxNumSequences, SizeType32 maxBeamWidth, - std::vector const& maxAttentionWindowVec, nvinfer1::DataType dtype, SizeType32 sinkTokenLength, + std::vector const& maxAttentionWindowVec, tensorrt_llm::DataType dtype, SizeType32 sinkTokenLength, CudaStreamPtr stream, SizeType32 maxSequenceLength, SizeType32 chunkSize, bool enableBlockReuse = true, CacheType cacheType = CacheType::kSELF, std::optional secondaryOffloadMinPriority = std::nullopt, @@ -2318,7 +2318,7 @@ class KVCacheManager : public BaseKVCacheManager KVCacheManager(SizeType32 numLayers, SizeType32 numKvHeads, SizeType32 sizePerHead, SizeType32 tokensPerBlock, BlocksPerWindow const& blocksPerWindow, SizeType32 maxNumSequences, SizeType32 maxBeamWidth, - std::vector const& maxAttentionWindowVec, nvinfer1::DataType dtype, SizeType32 sinkTokenLength, + std::vector const& maxAttentionWindowVec, tensorrt_llm::DataType dtype, SizeType32 sinkTokenLength, int64_t stream, SizeType32 maxSequenceLength, SizeType32 chunkSize, bool enableBlockReuse = false, CacheType cacheType = CacheType::kSELF, bool enablePartialReuse = true, bool copyOnpartialReuse = true, bool enableIndexerKCache = false, SizeType32 indexerKCacheQuantBlockSize = 128, @@ -2682,7 +2682,7 @@ class KVCacheManager : public BaseKVCacheManager SizeType32 mMaxNumSequences; // Maximum beam width SizeType32 mMaxBeamWidth; - nvinfer1::DataType mDataType; + tensorrt_llm::DataType mDataType; // Maximum kv cache length per sequence SizeType32 mMaxAttentionWindow; // Number of tokens per block diff --git a/cpp/include/tensorrt_llm/batch_manager/llmRequest.h b/cpp/include/tensorrt_llm/batch_manager/llmRequest.h index bc1ca3e6d012..e0d04f7b4ff7 100644 --- a/cpp/include/tensorrt_llm/batch_manager/llmRequest.h +++ b/cpp/include/tensorrt_llm/batch_manager/llmRequest.h @@ -18,6 +18,7 @@ #include "tensorrt_llm/common/assert.h" #include "tensorrt_llm/common/logger.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/executor/executor.h" #include "tensorrt_llm/runtime/bufferManager.h" #include "tensorrt_llm/runtime/iBuffer.h" @@ -1312,7 +1313,7 @@ class GenericLlmRequest mEncoderOutput = std::move(encoderOutput); } - void allocEncoderOutputHost(SizeType32 encoderHiddenSize, nvinfer1::DataType dataType) + void allocEncoderOutputHost(SizeType32 encoderHiddenSize, tensorrt_llm::DataType dataType) { mEncoderOutputHost = runtime::BufferManager::pinned( runtime::ITensor::makeShape({getEncoderOutputLen(), encoderHiddenSize}), dataType); @@ -1328,13 +1329,13 @@ class GenericLlmRequest return mEncoderHiddenStates; } - void allocEncoderOutput(runtime::BufferManager const& manager, nvinfer1::DataType dataType) + void allocEncoderOutput(runtime::BufferManager const& manager, tensorrt_llm::DataType dataType) { // unique_ptr --> shared_ptr ownership move mEncoderOutput = std::move(manager.emptyTensor(runtime::MemoryType::kGPU, dataType)); } - void allocEncoderHiddenStates(runtime::BufferManager const& manager, nvinfer1::DataType dataType) + void allocEncoderHiddenStates(runtime::BufferManager const& manager, tensorrt_llm::DataType dataType) { // unique_ptr --> shared_ptr ownership move mEncoderHiddenStates = std::move(manager.emptyTensor(runtime::MemoryType::kGPU, dataType)); @@ -1452,7 +1453,7 @@ class GenericLlmRequest mContextLogitsHost = std::move(contextLogitsHost); } - void allocContextLogitsHost(SizeType32 vocabSizePadded, nvinfer1::DataType logitsDataType) + void allocContextLogitsHost(SizeType32 vocabSizePadded, tensorrt_llm::DataType logitsDataType) { mContextLogitsHost = runtime::BufferManager::pinnedPool( runtime::ITensor::makeShape({mPromptLen, vocabSizePadded}), logitsDataType); @@ -1471,7 +1472,7 @@ class GenericLlmRequest mGenerationLogitsHost = std::move(generationLogitsHost); } - void allocGenerationLogitsHost(SizeType32 vocabSizePadded, nvinfer1::DataType logitsDataType) + void allocGenerationLogitsHost(SizeType32 vocabSizePadded, tensorrt_llm::DataType logitsDataType) { if (mIsStreaming) { @@ -1490,7 +1491,7 @@ class GenericLlmRequest } } - void allocTargetModelAcceptedTokenLogitsHost(SizeType32 vocabSizePadded, nvinfer1::DataType logitsDataType) + void allocTargetModelAcceptedTokenLogitsHost(SizeType32 vocabSizePadded, tensorrt_llm::DataType logitsDataType) { mGenerationLogitsHost = runtime::BufferManager::pinnedPool( runtime::ITensor::makeShape({1, getNumDraftTokens() + 1, vocabSizePadded}), logitsDataType); @@ -2356,7 +2357,7 @@ class GenericLlmRequest auto const numWords = static_cast(words.size()); auto const shape = runtime::ITensor::makeShape({2, numWords}); - auto tensor = runtime::BufferManager::pinnedPool(shape, nvinfer1::DataType::kINT32); + auto tensor = runtime::BufferManager::pinnedPool(shape, tensorrt_llm::DataType::kINT32); auto* data = runtime::bufferCast(*tensor); std::memcpy(data, words.data(), numWords * sizeof(int32_t)); std::memcpy(data + numWords, offsets.data(), numWords * sizeof(int32_t)); diff --git a/cpp/include/tensorrt_llm/batch_manager/logitsPostProcessor.h b/cpp/include/tensorrt_llm/batch_manager/logitsPostProcessor.h deleted file mode 100644 index 1916a915e337..000000000000 --- a/cpp/include/tensorrt_llm/batch_manager/logitsPostProcessor.h +++ /dev/null @@ -1,53 +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. - */ - -#pragma once - -#include "common.h" -#include "tensorrt_llm/batch_manager/llmRequest.h" -#include "tensorrt_llm/common/algorithm.h" -#include "tensorrt_llm/runtime/worldConfig.h" - -namespace tensorrt_llm::runtime -{ -class CudaStream; -} - -namespace tensorrt_llm::batch_manager -{ -class DecoderInputBuffers; - -class LogitsPostProcessor : Algorithm -{ -public: - using CudaStreamPtr = std::shared_ptr; - - using LogitsPostProcessorBatched = std::function const&, - std::vector&, - std::vector> const&, CudaStreamPtr const&, - std::vector> const&)>; - - constexpr static auto name{"LogitsPostProcessor"}; - - LogitsPostProcessor() = default; - - bool operator()(DecoderInputBuffers& inputBuffers, bool replicateLogitsPostProcessor, - runtime::WorldConfig const& worldConfig, CudaStreamPtr const& stream, - std::optional const& logitsPostProcessorBatched = std::nullopt) const; -}; - -} // namespace tensorrt_llm::batch_manager diff --git a/cpp/include/tensorrt_llm/batch_manager/makeDecodingBatchInputOutput.h b/cpp/include/tensorrt_llm/batch_manager/makeDecodingBatchInputOutput.h deleted file mode 100644 index 245f4b4b5286..000000000000 --- a/cpp/include/tensorrt_llm/batch_manager/makeDecodingBatchInputOutput.h +++ /dev/null @@ -1,56 +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. - */ - -#pragma once - -#include "common.h" -#include "tensorrt_llm/common/algorithm.h" -#include "tensorrt_llm/common/optionalRef.h" -#include "tensorrt_llm/runtime/common.h" -#include "tensorrt_llm/runtime/iGptDecoderBatched.h" -#include "tensorrt_llm/runtime/modelConfig.h" - -namespace tensorrt_llm::runtime::decoder -{ -class DecoderState; -} // namespace tensorrt_llm::runtime::decoder - -namespace tensorrt_llm::batch_manager -{ -class DecoderInputBuffers; -class RuntimeBuffers; - -class MakeDecodingBatchInputOutput : Algorithm -{ -public: - constexpr static auto name{"MakeDecodingBatchInputOutput"}; - - using SizeType32 = tensorrt_llm::runtime::SizeType32; - using TensorPtr = runtime::ITensor::SharedPtr; - template - using OptionalRef = tensorrt_llm::common::OptionalRef; - - MakeDecodingBatchInputOutput() = default; - - void operator()(DecoderInputBuffers& inputBuffers, runtime::decoder::DecoderState& decoderState, - runtime::ModelConfig const& modelConfig, OptionalRef fusedRuntimeBuffers) const; - - static void createDecoderBatchInputs(DecoderInputBuffers& inputBuffers, std::vector const& activeSlots, - runtime::decoder::DecoderState const& decoderState); -}; - -} // namespace tensorrt_llm::batch_manager diff --git a/cpp/include/tensorrt_llm/batch_manager/medusaBuffers.h b/cpp/include/tensorrt_llm/batch_manager/medusaBuffers.h index ba29be6ede81..5342591840a8 100644 --- a/cpp/include/tensorrt_llm/batch_manager/medusaBuffers.h +++ b/cpp/include/tensorrt_llm/batch_manager/medusaBuffers.h @@ -22,7 +22,6 @@ #include "tensorrt_llm/runtime/iTensor.h" #include "tensorrt_llm/runtime/modelConfig.h" #include "tensorrt_llm/runtime/promptTuningParams.h" -#include "tensorrt_llm/runtime/tllmRuntime.h" #include "tensorrt_llm/runtime/worldConfig.h" namespace tensorrt_llm::batch_manager @@ -36,10 +35,6 @@ class MedusaBuffers using TensorPtr = runtime::ITensor::SharedPtr; using TensorMap = runtime::StringPtrMap; - MedusaBuffers(SizeType32 maxBatchSize, SizeType32 maxBeamWidth, runtime::BufferManager const& manager, - runtime::ModelConfig const& modelConfig, runtime::WorldConfig const& worldConfig, - executor::DecodingConfig const& decodingConfig, runtime::TllmRuntime const& runtime); - void reshape(SizeType32 numCtxSequences, SizeType32 numGenSequences, SizeType32 tokensPerStep); void insertInputTensors( diff --git a/cpp/include/tensorrt_llm/batch_manager/peftCacheManager.h b/cpp/include/tensorrt_llm/batch_manager/peftCacheManager.h index cf65753783e8..ed928e96d811 100644 --- a/cpp/include/tensorrt_llm/batch_manager/peftCacheManager.h +++ b/cpp/include/tensorrt_llm/batch_manager/peftCacheManager.h @@ -25,7 +25,7 @@ #include "tensorrt_llm/runtime/workerPool.h" #include "tensorrt_llm/runtime/worldConfig.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include #include @@ -147,7 +147,7 @@ class PeftCacheManager : public BasePeftCacheManager void updateTaskState(uint64_t taskId, uint64_t reqId, bool terminate = false, bool pause = false); static std::pair getMaxNumSlots(PeftCacheManagerConfig const& config, - nvinfer1::DataType dataType, uint64_t pageWidth, uint64_t max1dModSize, + tensorrt_llm::DataType dataType, uint64_t pageWidth, uint64_t max1dModSize, runtime::BufferManager const& bufferManager); static std::pair getPageManagerConfig( diff --git a/cpp/include/tensorrt_llm/batch_manager/promptTuningBuffers.h b/cpp/include/tensorrt_llm/batch_manager/promptTuningBuffers.h deleted file mode 100644 index a1d8849a8811..000000000000 --- a/cpp/include/tensorrt_llm/batch_manager/promptTuningBuffers.h +++ /dev/null @@ -1,106 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2025 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. - */ - -#pragma once - -#include "tensorrt_llm/batch_manager/common.h" -#include "tensorrt_llm/runtime/bufferManager.h" -#include "tensorrt_llm/runtime/iTensor.h" -#include "tensorrt_llm/runtime/modelConfig.h" -#include "tensorrt_llm/runtime/promptTuningParams.h" -#include "tensorrt_llm/runtime/worldConfig.h" - -namespace tensorrt_llm::batch_manager -{ - -class PromptTuningBuffers -{ - -public: - using SizeType32 = tensorrt_llm::runtime::SizeType32; - using ITensor = tensorrt_llm::runtime::ITensor; - using TensorPtr = runtime::ITensor::SharedPtr; - - runtime::PromptTuningParams mPromptTuningParams; - SizeType32 mMaxPromptVocabSize; - - PromptTuningBuffers(SizeType32 maxBatchSize, runtime::BufferManager const& manager, - runtime::ModelConfig const& modelConfig, runtime::WorldConfig const& worldConfig); - - PromptTuningBuffers(SizeType32 maxBatchSize, runtime::BufferManager const& manager, - runtime::ModelConfig const& modelConfig, runtime::WorldConfig const& worldConfig, bool promptTableOffloading); - - void validate(std::optional const& optReqPromptEmbeddingTable, - std::optional const& optReqPromptVocabSize); - - void fill(RequestVector const& contextRequests, RequestVector const& genRequests, - runtime::BufferManager const& manager, bool packed); - - /* - * The below functions are specific for Chunked Prefill mode - * Chunk Ptable with Ping-Pong Buffer Implementation - * ----------------------------------------------- - * - * Overview: - * The chunk ptable (prompt tuning table) system uses a ping-pong buffer mechanism to efficiently - * manage large embedding tables when operating in context Prefill mode. This allows - * for processing of large embedding tables by loading them in chunks from CPU to GPU memory, - * enabling support for tables that exceed available GPU memory. - * - * Key Components: - * 1. Ping-Pong Buffers (mChunkPtableBuffers): - * - Two alternating GPU buffers that store chunks of the embedding table - * - While the current buffer is being processed by the model, - * the next chunk can be asynchronously loaded into the other buffer - * - Managed through mChunkPtableCurrentIndex (toggles between 0 and 1) - * 2. Start Positions Tracking (mChunkPtableBufferStartPositions): - * - Mainly used for multi-batch processing - * - Maintains the starting position of each batch's data within each buffer - * - Maintained separately for each ping-pong buffer - * - * Memory Optimization: - * - Only two GPU buffers are maintained regardless of total embedding table size - * - Each buffer size is limited to contextChunkSize * hiddenSize - * - Efficient memory usage through chunk-based processing - */ - - bool mPromptTableOffloading; - - bool mChunkPtableInitialized{false}; - std::optional> mChunkPtableBuffers; - std::optional>> mChunkPtableBufferStartPositions; - size_t mChunkPtableCurrentIndex{0}; - - void initializeChunkPtableBuffers(runtime::BufferManager const& manager, runtime::ModelConfig const& modelConfig, - SizeType32 contextChunkSize, std::shared_ptr const& llmReq); - - void switchChunkPtableBuffer(); - - size_t getChunkPtableCurrentIndex(); - - [[nodiscard]] TensorPtr& getChunkPtableBuffer(size_t index); - - [[nodiscard]] SizeType32 getChunkPtableBufferSliceSize(size_t index, size_t batchIdx); - - [[nodiscard]] SizeType32 getChunkPtableBufferStartPosition(size_t index, size_t batchIdx); - - void updateBufferStartPosition(size_t index, SizeType32 numRows); - - void clearBufferStartPositions(size_t index); -}; - -} // namespace tensorrt_llm::batch_manager diff --git a/cpp/include/tensorrt_llm/batch_manager/rnnStateManager.h b/cpp/include/tensorrt_llm/batch_manager/rnnStateManager.h index 5c0bfe136de2..c4f97950a6b9 100644 --- a/cpp/include/tensorrt_llm/batch_manager/rnnStateManager.h +++ b/cpp/include/tensorrt_llm/batch_manager/rnnStateManager.h @@ -17,6 +17,7 @@ #pragma once #include "tensorrt_llm/batch_manager/common.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/executor/dataTransceiverState.h" #include "tensorrt_llm/runtime/bufferManager.h" #include "tensorrt_llm/runtime/iTensor.h" @@ -42,8 +43,8 @@ class RnnStateManager runtime::WorldConfig const& worldConfig, tensorrt_llm::runtime::BufferManager const& bufferManager); RnnStateManager(SizeType32 dState, SizeType32 dConv, SizeType32 numHeads, SizeType32 nGroups, SizeType32 headDim, - SizeType32 maxBatchSize, runtime::WorldConfig const& worldConfig, int64_t stream, nvinfer1::DataType dtype, - nvinfer1::DataType ssmCacheDtype, std::vector const& ppLayers, SizeType32 numLayers); + SizeType32 maxBatchSize, runtime::WorldConfig const& worldConfig, int64_t stream, tensorrt_llm::DataType dtype, + tensorrt_llm::DataType ssmCacheDtype, std::vector const& ppLayers, SizeType32 numLayers); void getPtrBuffers(TensorMap& inputBuffers, runtime::ModelConfig const& modelConfig, runtime::WorldConfig const& worldConfig) const; @@ -68,9 +69,9 @@ class RnnStateManager [[nodiscard]] TensorPtr getSsmStates() const; - [[nodiscard]] nvinfer1::DataType getConvStateDataType() const noexcept; + [[nodiscard]] tensorrt_llm::DataType getConvStateDataType() const noexcept; - [[nodiscard]] nvinfer1::DataType getSsmStateDataType() const noexcept; + [[nodiscard]] tensorrt_llm::DataType getSsmStateDataType() const noexcept; [[nodiscard]] executor::kv_cache::CacheState::RnnModelConfig getRnnCacheStateModelConfig() const noexcept; @@ -111,8 +112,8 @@ class RnnStateManager std::vector mFreeBlocks; std::unordered_map mCacheIndex; std::optional mBufferManager; - nvinfer1::DataType mDtype{nvinfer1::DataType::kFLOAT}; - nvinfer1::DataType mSsmCacheDtype{nvinfer1::DataType::kFLOAT}; + tensorrt_llm::DataType mDtype{tensorrt_llm::DataType::kFLOAT}; + tensorrt_llm::DataType mSsmCacheDtype{tensorrt_llm::DataType::kFLOAT}; // RNN model config (global values before TP/PP split) SizeType32 mDState{0}; diff --git a/cpp/include/tensorrt_llm/batch_manager/runtimeBuffers.h b/cpp/include/tensorrt_llm/batch_manager/runtimeBuffers.h deleted file mode 100644 index 97a4ae67acdd..000000000000 --- a/cpp/include/tensorrt_llm/batch_manager/runtimeBuffers.h +++ /dev/null @@ -1,326 +0,0 @@ -/* - * Copyright (c) 2023-2026, NVIDIA CORPORATION. All rights reserved. - * - * 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. - */ - -#pragma once - -#include "tensorrt_llm/batch_manager/common.h" -#include "tensorrt_llm/batch_manager/rnnStateManager.h" -#include "tensorrt_llm/common/optionalRef.h" -#include "tensorrt_llm/runtime/eagleBuffers.h" -#include "tensorrt_llm/runtime/explicitDraftTokensBuffers.h" -#include "tensorrt_llm/runtime/iTensor.h" -#include "tensorrt_llm/runtime/lookaheadBuffers.h" -#include "tensorrt_llm/runtime/loraManager.h" -#include "tensorrt_llm/runtime/modelConfig.h" -#include "tensorrt_llm/runtime/worldConfig.h" - -#include -#include -#include -#include - -namespace tensorrt_llm::runtime -{ -class TllmRuntime; - -namespace decoder -{ -class DecoderState; -} // namespace decoder -} // namespace tensorrt_llm::runtime - -namespace tensorrt_llm::batch_manager -{ - -namespace kv_cache_manager -{ -class BaseKVCacheManager; -} // namespace kv_cache_manager - -class LlmRequest; - -class EncoderBuffers; -class LoraBuffers; -class MedusaBuffers; -class PromptTuningBuffers; -class RnnStateBuffers; -class TransformerBuffers; - -class RuntimeBuffers -{ -public: - static constexpr auto kLogitsTensorName = "logits"; - static constexpr auto kHiddenStatesOutputTensorName = "hidden_states_output"; - static constexpr auto kHiddenStatesInputTensorName = "hidden_states_input"; - static constexpr auto kInputIdsTensorName = "input_ids"; - static constexpr auto kLastTokenIdsTensorName = "last_token_ids"; - static constexpr auto kHostRequestTypesTensorName = "host_request_types"; - static constexpr auto kContextLengthsTensorName = "context_lengths"; - static constexpr auto kHostContextLengthsTensorName = "host_context_lengths"; - static constexpr auto kSequenceLengthsTensorName = "sequence_length"; - static constexpr auto kPromptEmbeddingTableTensorName = "prompt_embedding_table"; - static constexpr auto kTasksTensorName = "tasks"; - static constexpr auto kPromptVocabSizeTensorName = "prompt_vocab_size"; - static constexpr auto kMRopeRotaryCosSinTensorName = "mrope_rotary_cos_sin"; - static constexpr auto kMRopePositionDeltasTensorName = "mrope_position_deltas"; - - using SizeType32 = runtime::SizeType32; - using TensorPtr = runtime::ITensor::SharedPtr; - using TensorMap = runtime::ITensor::TensorMap; - using PeftTable = runtime::LoraManager::PeftTable; - template - using OptionalRef = tensorrt_llm::common::OptionalRef; - - [[nodiscard]] SizeType32 constexpr getContextIndex() const noexcept - { - return contextIndex; - }; - - void constexpr setContextIndex(SizeType32 index) noexcept - { - contextIndex = index; - }; - - [[nodiscard]] SizeType32 constexpr getNumContextTokens() const noexcept - { - return numContextTokens; - }; - - [[nodiscard]] BatchState getBatchState() const noexcept - { - return {numContextRequests, numGenRequests, getNumTokens(), maxKvCacheLengthRounded}; - }; - -private: - [[nodiscard]] SizeType32 constexpr getNumRequests() const noexcept - { - return numContextRequests + numGenRequests; - }; - - [[nodiscard]] SizeType32 constexpr getNumSequences() const noexcept - { - return numContextRequests + numGenSequences; - }; - - [[nodiscard]] SizeType32 constexpr getNumTokens() const noexcept - { - return numContextTokens + numGenTokens; - }; - - //! Sizes - SizeType32 numContextRequests{}; - SizeType32 numGenRequests{}; - SizeType32 numGenSequences{}; - SizeType32 numContextTokens{}; - SizeType32 numGenTokens{}; - SizeType32 numLogits{}; - SizeType32 maxKvCacheLengthRounded{}; - - //! General - TensorPtr inputsIds; - - TensorPtr contextLengthsHost; - TensorPtr contextLengthsDevice; - TensorPtr sequenceLengthsHost; - - //! Index of selected runtime context. - SizeType32 contextIndex{}; - SizeType32 maxContextLength{}; - -public: - TensorPtr sequenceLengthsDevice; - bool promptTableOffloading; - - //! Prompt-Tuning - std::unique_ptr promptTuningBuffers; - -private: - //! Runtime - //! Type of host tensor: 0 for context, 1 for generation - TensorPtr requestTypes; - - TensorPtr lastTokenIdsHost; - TensorPtr lastTokenIdsDevice; - TensorPtr logitsIdsHost; - - //! Pipeline-Parallelism - TensorPtr hiddenStates; - - //! Mrope - TensorPtr mropeRotaryCosSin; - TensorPtr mropePositionDeltas; - - //! LoRA - std::unique_ptr loraBuffers; - -public: - //! Additional buffers depending on model type - std::unique_ptr transformerBuffers; - std::unique_ptr rnnStateBuffers; - - //! Encoder-Decoder - std::unique_ptr encoderBuffers; - - //! Medusa - std::unique_ptr mMedusaBuffers; - //! Lookahead decoding - std::unique_ptr mLookaheadBuffers; - //! Explicit draft tokens decoding - std::unique_ptr mExplicitDraftTokensBuffers; - //! Eagle decoding - std::unique_ptr mEagleBuffers; - - //! Language adapter routing information if language adapter is presented, [numTokens, numLanguages] - TensorPtr languageAdapterRoutings; - - TensorPtr cacheIndirDecoderIOBatchedCopySrcOffsets; - TensorPtr cacheIndirDecoderIOBatchedCopyDstOffsets; - TensorPtr cacheIndirDecoderIOBatchedCopySizes; - - //! Logits - std::vector numContextLogits; - TensorPtr logits; - - //! Helper cache for store generation logits - struct GenerationLogitsCache - { - static constexpr auto kCACHE_LENGTH = 8; - - //! Buffer for logits between steps to prevent from being overwritten - //! [kCACHE_LENGTH, maxBatchSize * maxBeamWidth, vocabSizePadded] - TensorPtr logits; - //! Record the usage offset of the cacheGenerationLogits buffer - SizeType32 offset{0}; - - //! Temporarily store the transposed results of multiple fragment logits, [maxBeamWidth, kCACHE_LENGTH] - TensorPtr transposedLogits; - - //! Temporarily store logits buffer address during the transposing, [maxBatchSize, kCACHE_LENGTH] - //! One row per batch slot (same layout as fragmentPointerHost) so concurrent flushes for - //! different requests in the same batch never clobber each other's pointer arrays. - TensorPtr fragmentPointerDevice; - - //! Temporarily store logits buffer address during the transposing, [maxBatchSize, kCACHE_LENGTH] - TensorPtr fragmentPointerHost; - - //! Cycling index for workspace - size_t workIdx{0}; - - void cycleWorkIdx() - { - workIdx = (workIdx + 1) % (fragmentPointerHost->getShape().d[0]); - } - - //! Returns matching host and device pointer rows for the current workIdx, then advances - //! workIdx. Always call this instead of the individual getters to avoid ordering bugs. - [[nodiscard]] std::pair getFragmentPointerSlot() - { - TensorPtr host = runtime::ITensor::slice(fragmentPointerHost, workIdx, 1); - TensorPtr device = runtime::ITensor::slice(fragmentPointerDevice, workIdx, 1); - cycleWorkIdx(); - return {std::move(host), std::move(device)}; - }; - }; - - GenerationLogitsCache generationLogitsCache; - - //! Mapping from batch idx to slot id - TensorPtr seqSlots; - TensorPtr seqSlotsDevice; - - //! Explicitly device-copy src offsets to reduce warp stalls in copy batch kernel invocation - //! [mMaxNumRequests], on gpu - TensorPtr mCacheIndirDecoderIOBatchedCopySrcOffsetsSliceDevice; - //! Explicitly device-copy dst offsets to reduce warp stalls in copy batch kernel invocation - //! [mMaxNumRequests], on gpu - TensorPtr mCacheIndirDecoderIOBatchedCopyDstOffsetsSliceDevice; - //! Explicitly device-copy size to reduce warp stalls in copy batch kernel invocation - //! [mMaxNumRequests], on gpu - TensorPtr mCacheIndirDecoderIOBatchedCopyCopySizesDevice; - -private: - //! Re-capture cuda graph when max kv cache len of the batch has changed on kKV_CACHE_LEN_CUDA_GRAPH_ROUND_SIZE. - static SizeType32 constexpr kKV_CACHE_LEN_CUDA_GRAPH_ROUND_SIZE{256}; - - TensorMap mAdditionalOutputTensors; // Tensors storing additional output tensors. - - //! Engine I/O - TensorMap inputMap; - TensorMap outputMap; - -public: - RuntimeBuffers(SizeType32 maxBatchSize, SizeType32 maxBeamWidth, - std::vector const& maxAttentionWindowVec, SizeType32 maxAttentionWindow, SizeType32 sinkTokenLen, - runtime::TllmRuntime const& runtime, runtime::ModelConfig const& modelConfig, - runtime::WorldConfig const& worldConfig, executor::DecodingConfig const& decodingConfig, - bool gatherGenerationLogits, std::optional maxNumTokens = std::nullopt, - std::optional> const& additionalModelOutputs = std::nullopt, - bool promptTableOffloading = false); - - RuntimeBuffers(RuntimeBuffers const& other) = delete; - RuntimeBuffers& operator=(RuntimeBuffers const& other) = delete; - RuntimeBuffers(RuntimeBuffers&& other) = delete; - RuntimeBuffers& operator=(RuntimeBuffers&& other) = delete; - - ~RuntimeBuffers(); - - std::tuple prepareStep(RequestVector const& contextRequests, - RequestVector const& genRequests, SizeType32 maxBeamWidth, SizeType32 maxAttentionWindow, - runtime::decoder::DecoderState const& decoderState, kv_cache_manager::BaseKVCacheManager* kvCacheManager, - kv_cache_manager::BaseKVCacheManager* crossKvCacheManager, rnn_state_manager::RnnStateManager* rnnStateManager, - PeftTable const& peftTable, runtime::TllmRuntime const& runtime, runtime::ModelConfig const& modelConfig, - runtime::WorldConfig const& worldConfig, bool gatherGenerationLogits, bool trtOverlap, - OptionalRef newOutputTokens = std::nullopt); - - void prepareBuffersForCudaGraph(SizeType32 maxSequenceLength); - - void prepareExplicitDraftTokenBuffers(runtime::ExplicitDraftTokensBuffers::Inputs const& explicitDraftTokensBuffers, - runtime::TllmRuntime const& runtime, runtime::ModelConfig const& modelConfig, - runtime::WorldConfig const& worldConfig); - - void prepareEagleBuffers(RequestVector const& contextRequests, RequestVector const& genRequests, - runtime::EagleBuffers::Inputs const& eagleBuffers, runtime::TllmRuntime const& runtime, - runtime::ModelConfig const& modelConfig, runtime::WorldConfig const& worldConfig); - -private: - void create(SizeType32 maxBatchSize, SizeType32 maxBeamWidth, std::vector const& maxAttentionWindowVec, - SizeType32 maxAttentionWindow, SizeType32 sinkTokenLen, runtime::TllmRuntime const& runtime, - runtime::ModelConfig const& modelConfig, runtime::WorldConfig const& worldConfig, - executor::DecodingConfig const& decodingConfig, bool gatherGenerationLogits, - std::optional> const& additionalModelOutputs = std::nullopt); - - //! @brief set max sizes for pre-allocation - void setMaxBufferSizes(SizeType32 maxBatchSize, SizeType32 maxBeamWidth, runtime::ModelConfig const& modelConfig, - std::optional maxNumRuntimeTokens); - - //! @brief set sizes depending on scheduled requests - void setBufferSizes(RequestVector const& contextRequests, RequestVector const& genRequests); - - void reshape(runtime::TllmRuntime const& runtime, runtime::ModelConfig const& modelConfig, - runtime::WorldConfig const& worldConfig, bool gatherGenerationLogits); - - void setFromInputs(RequestVector const& contextRequests, RequestVector const& genRequests, SizeType32 maxBeamWidth, - SizeType32 maxAttentionWindow, runtime::decoder::DecoderState const& decoderState, - kv_cache_manager::BaseKVCacheManager* kvCacheManagerPtr, - kv_cache_manager::BaseKVCacheManager* crossKvCacheManagerPtr, - rnn_state_manager::RnnStateManager* rnnStateManagerPtr, PeftTable const& peftTable, - runtime::TllmRuntime const& runtime, runtime::ModelConfig const& modelConfig, - runtime::WorldConfig const& worldConfig, bool trtOverlap, OptionalRef newOutputTokens); - - void fillIOMaps(runtime::ModelConfig const& modelConfig, runtime::WorldConfig const& worldConfig); -}; - -} // namespace tensorrt_llm::batch_manager diff --git a/cpp/include/tensorrt_llm/batch_manager/transformerBuffers.h b/cpp/include/tensorrt_llm/batch_manager/transformerBuffers.h deleted file mode 100644 index b5254c6357b4..000000000000 --- a/cpp/include/tensorrt_llm/batch_manager/transformerBuffers.h +++ /dev/null @@ -1,145 +0,0 @@ -/* - * Copyright (c) 2023-2024, NVIDIA CORPORATION. All rights reserved. - * - * 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. - */ - -#pragma once - -#include "tensorrt_llm/batch_manager/common.h" -#include "tensorrt_llm/batch_manager/kvCacheType.h" -#include "tensorrt_llm/runtime/bufferManager.h" -#include "tensorrt_llm/runtime/iTensor.h" -#include "tensorrt_llm/runtime/modelConfig.h" -#include "tensorrt_llm/runtime/worldConfig.h" - -namespace tensorrt_llm::runtime -{ -class TllmRuntime; -class MulticastTensor; -} // namespace tensorrt_llm::runtime - -namespace tensorrt_llm::batch_manager -{ - -namespace kv_cache_manager -{ -class BaseKVCacheManager; -} - -class TransformerBuffers -{ -public: - using SizeType32 = runtime::SizeType32; - using TensorPtr = runtime::ITensor::SharedPtr; - using TensorMap = runtime::StringPtrMap; - - static constexpr auto kCrossAttentionMaskTensorName = "cross_attention_mask"; - static constexpr auto kCrossAttentionPackedMaskTensorName = "cross_attention_packed_mask"; - static constexpr auto kPositionIdsTensorName = "position_ids"; - static constexpr auto kCacheIndirectionsTensorName = "cache_indirection"; - static constexpr auto kHostPastKeyValueLengthsTensorName = "host_past_key_value_lengths"; - static constexpr auto kHostSinkTokenLengthTensorName = "host_sink_token_length"; - static constexpr auto kHostMaxAttentionWindowSizesTensorName = "host_max_attention_window_sizes"; - static constexpr auto kHostContextProgressTensorName = "host_context_progress"; - static constexpr auto kKvCacheBlockOffsetsTensorName = "kv_cache_block_offsets"; - static constexpr auto kHostKvCacheBlockOffsetsTensorName = "host_kv_cache_block_offsets"; - static constexpr auto kCrossKvCacheBlockOffsetsTensorName = "cross_kv_cache_block_offsets"; - static constexpr auto kHostCrossKvCacheBlockOffsetsTensorName = "host_cross_kv_cache_block_offsets"; - static constexpr auto kHostCrossKvCachePoolPointersTensorName = "host_cross_kv_cache_pool_pointers"; - static constexpr auto kHostCrossKvCachePoolMappingTensorName = "host_cross_kv_cache_pool_mapping"; - static constexpr auto kSkipCrossAttentionBlocksTensorName = "skip_cross_attn_blocks"; - - TensorPtr pastKeyValueLengths; // Host tensor - TensorPtr positionIds; - - // max kv cache lengths. - TensorPtr maxAttentionWindows; - // sink token lengths. - TensorPtr sinkTokenLengths; - TensorPtr cacheIndirection; - TensorPtr kvCacheBlockOffsetsHost; // [numPools, maxBatch * maxBeamWidth, 2, maxBlocksPerSeq] - TensorPtr kvCacheBlockOffsetsDevice; // [numPools, maxBatch * maxBeamWidth, 2, maxBlocksPerSeq] - TensorPtr contextProgressHost; - - // Cross attention buffers - TensorPtr crossKvCacheBlockPoolPointers = nullptr; - TensorPtr crossKvCacheBlockPoolMapping = nullptr; - TensorPtr crossKvCacheBlockOffsetsHost = nullptr; - TensorPtr crossKvCacheBlockOffsetsDevice = nullptr; - TensorPtr crossAttentionMaskCopySrcOffsets = nullptr; // [maxNumRequest] pinned memory. - TensorPtr crossAttentionMaskCopyDstOffsets = nullptr; // [maxNumRequest] pinned memory. - TensorPtr crossAttentionMaskCopySizes = nullptr; // [maxNumRequest] pinned memory. - TensorPtr crossAttentionMaskDevice = nullptr; // [maxNumTokens, maxEncoderOutputLen] - // This is created to allow mixed memory types of crossAttentionMask (i.e. CPU and GPU). - TensorPtr crossAttentionMaskPinnedHost = nullptr; // [maxNumTokens, maxEncoderOutputLen] - // See more details in tensorrt_llm/kernels/contextFusedMultiHeadAttention/fmhaPackedMask.cu. - // The attention packed mask for FMHA where each bit represents one mask. - TensorPtr crossAttentionPackedMaskDevice - = nullptr; // [maxBatchSize, maxInputLengthInBatch, roundUp(maxEncoderOutputLen, 32)] - // The number of cumulative Q sequence lengths in the mask input, which is used to get mask offsets for different - // requests. - TensorPtr crossAttentionCuQSeqLensDevice = nullptr; // [maxBatchSize + 1] - // The number of cumulative Q sequence lengths in the packed mask, which is used to get mask offsets for different - // requests. - TensorPtr crossAttentionPackedMaskCuMaskRowsDevice = nullptr; // [maxBatchSize + 1] - - TensorPtr cacheIndirBatchedCopySrcOffsets; - TensorPtr cacheIndirBatchedCopyDstOffsets; - TensorPtr cacheIndirBatchedCopySizes; - - TensorPtr fillValuesAlt; - TensorPtr fillValuesAltDevice; - TensorPtr seqSlotsAlt; - TensorPtr seqSlotsAltDevice; - TensorPtr skipCrossAttnBlocks; - - std::shared_ptr gemmAllReduceOutput; - - TransformerBuffers(SizeType32 maxBatchSize, SizeType32 maxBeamWidth, - std::vector const& maxAttentionWindowVec, SizeType32 maxAttentionWindow, SizeType32 sinkTokenLen, - runtime::TllmRuntime const& runtime, runtime::ModelConfig const& modelConfig, - runtime::WorldConfig const& worldConfig); - - void reshape(SizeType32 numSequences, SizeType32 numInputTokens); - - void reshapeKvTensors(SizeType32 maxBatchSize, SizeType32 maxBeamWidth, SizeType32 maxBlocksPerSeq, - kv_cache_manager::CacheType kvCacheType, SizeType32 numPools, runtime::BufferManager const& manager); - - void getBuffers(TensorMap& inputBuffers, TensorMap& outputBuffers, runtime::ModelConfig const& modelConfig) const; - - void copyPositionIds(runtime::TllmRuntime const& runtime, std::vector const& positionIdsHost, - bool isChatGlm, TensorPtr const& decoderPositionIds); - - void copyKvBlockOffsets(RequestVector const& contextRequests, RequestVector const& genRequests, - kv_cache_manager::BaseKVCacheManager const* kvCacheManager, - kv_cache_manager::BaseKVCacheManager const* crossKvCacheManager, runtime::BufferManager const& manager); - - // Copy CacheIndirection from `decoderCacheIndirectionOutput` to `this->cacheIndirection` - void copyCacheIndirection(RequestVector const& genRequests, TensorPtr const& decoderCacheIndirectionOutput, - runtime::CudaStream const& stream); - - void copyCrossAttentionMasks(RequestVector const& contextRequests, RequestVector const& genRequests, - TensorPtr const& decoderContextLengthsDevice, TensorPtr const& encoderInputLengths, - SizeType32 maxDecoderContextLength, SizeType32 maxEncoderInputLengthInBatch, - runtime::TllmRuntime const& runtime); - - void copySkipCrossAttnBlocks(bool const& _skipCrossAttnBlocks, runtime::TllmRuntime const& runtime); - -private: - SizeType32 maxInputLen; - SizeType32 maxEncoderOutputLen; - SizeType32 maxNumTokens; -}; - -} // namespace tensorrt_llm::batch_manager diff --git a/cpp/include/tensorrt_llm/batch_manager/updateDecoderBuffers.h b/cpp/include/tensorrt_llm/batch_manager/updateDecoderBuffers.h deleted file mode 100644 index 526a756e5546..000000000000 --- a/cpp/include/tensorrt_llm/batch_manager/updateDecoderBuffers.h +++ /dev/null @@ -1,51 +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. - */ - -#pragma once - -#include "tensorrt_llm/common/algorithm.h" -#include "tensorrt_llm/runtime/modelConfig.h" - -namespace tensorrt_llm::runtime -{ -class BufferManager; -class CudaEvent; - -namespace decoder -{ -class DecoderState; -} // namespace decoder -} // namespace tensorrt_llm::runtime - -namespace tensorrt_llm::batch_manager -{ - -class DecoderOutputBuffers; - -class UpdateDecoderBuffers : Algorithm -{ -public: - constexpr static auto name{"UpdateDecoderBuffers"}; - - UpdateDecoderBuffers() = default; - - runtime::CudaEvent operator()(runtime::ModelConfig const& modelConfig, DecoderOutputBuffers& decoderOutputBuffers, - runtime::BufferManager const& copyBufferManager, runtime::decoder::DecoderState const& decoderState, - bool returnLogProbs, runtime::CudaEvent const& decoderFinishEvent) const; -}; - -} // namespace tensorrt_llm::batch_manager diff --git a/cpp/include/tensorrt_llm/common/dataType.h b/cpp/include/tensorrt_llm/common/dataType.h index 2f19404f9c94..9b3bb5fdf0f0 100644 --- a/cpp/include/tensorrt_llm/common/dataType.h +++ b/cpp/include/tensorrt_llm/common/dataType.h @@ -19,7 +19,7 @@ #include "tensorrt_llm/common/config.h" #include "tensorrt_llm/common/tllmException.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include TRTLLM_NAMESPACE_BEGIN @@ -27,61 +27,61 @@ TRTLLM_NAMESPACE_BEGIN namespace common { -constexpr static size_t getDTypeSize(nvinfer1::DataType type) +constexpr static size_t getDTypeSize(tensorrt_llm::DataType type) { switch (type) { - case nvinfer1::DataType::kINT64: return 8; - case nvinfer1::DataType::kINT32: [[fallthrough]]; - case nvinfer1::DataType::kFLOAT: return 4; - case nvinfer1::DataType::kBF16: [[fallthrough]]; - case nvinfer1::DataType::kHALF: return 2; - case nvinfer1::DataType::kBOOL: [[fallthrough]]; - case nvinfer1::DataType::kUINT8: [[fallthrough]]; - case nvinfer1::DataType::kINT8: [[fallthrough]]; - case nvinfer1::DataType::kFP8: return 1; - case nvinfer1::DataType::kINT4: TLLM_THROW("Cannot determine size of INT4 data type"); - case nvinfer1::DataType::kFP4: TLLM_THROW("Cannot determine size of FP4 data type"); + case tensorrt_llm::DataType::kINT64: return 8; + case tensorrt_llm::DataType::kINT32: [[fallthrough]]; + case tensorrt_llm::DataType::kFLOAT: return 4; + case tensorrt_llm::DataType::kBF16: [[fallthrough]]; + case tensorrt_llm::DataType::kHALF: return 2; + case tensorrt_llm::DataType::kBOOL: [[fallthrough]]; + case tensorrt_llm::DataType::kUINT8: [[fallthrough]]; + case tensorrt_llm::DataType::kINT8: [[fallthrough]]; + case tensorrt_llm::DataType::kFP8: return 1; + case tensorrt_llm::DataType::kINT4: TLLM_THROW("Cannot determine size of INT4 data type"); + case tensorrt_llm::DataType::kFP4: TLLM_THROW("Cannot determine size of FP4 data type"); default: TLLM_THROW("Unknown dtype %d", static_cast(type)); } return 0; } -constexpr static size_t getDTypeSizeInBits(nvinfer1::DataType type) +constexpr static size_t getDTypeSizeInBits(tensorrt_llm::DataType type) { switch (type) { - case nvinfer1::DataType::kINT64: return 64; - case nvinfer1::DataType::kINT32: [[fallthrough]]; - case nvinfer1::DataType::kFLOAT: return 32; - case nvinfer1::DataType::kBF16: [[fallthrough]]; - case nvinfer1::DataType::kHALF: return 16; - case nvinfer1::DataType::kBOOL: [[fallthrough]]; - case nvinfer1::DataType::kUINT8: [[fallthrough]]; - case nvinfer1::DataType::kINT8: [[fallthrough]]; - case nvinfer1::DataType::kFP8: return 8; - case nvinfer1::DataType::kINT4: [[fallthrough]]; - case nvinfer1::DataType::kFP4: return 4; + case tensorrt_llm::DataType::kINT64: return 64; + case tensorrt_llm::DataType::kINT32: [[fallthrough]]; + case tensorrt_llm::DataType::kFLOAT: return 32; + case tensorrt_llm::DataType::kBF16: [[fallthrough]]; + case tensorrt_llm::DataType::kHALF: return 16; + case tensorrt_llm::DataType::kBOOL: [[fallthrough]]; + case tensorrt_llm::DataType::kUINT8: [[fallthrough]]; + case tensorrt_llm::DataType::kINT8: [[fallthrough]]; + case tensorrt_llm::DataType::kFP8: return 8; + case tensorrt_llm::DataType::kINT4: [[fallthrough]]; + case tensorrt_llm::DataType::kFP4: return 4; default: TLLM_THROW("Unknown dtype %d", static_cast(type)); } return 0; } -[[maybe_unused]] static std::string getDtypeString(nvinfer1::DataType type) +[[maybe_unused]] static std::string getDtypeString(tensorrt_llm::DataType type) { switch (type) { - case nvinfer1::DataType::kFLOAT: return "fp32"; break; - case nvinfer1::DataType::kHALF: return "fp16"; break; - case nvinfer1::DataType::kINT8: return "int8"; break; - case nvinfer1::DataType::kINT32: return "int32"; break; - case nvinfer1::DataType::kBOOL: return "bool"; break; - case nvinfer1::DataType::kUINT8: return "uint8"; break; - case nvinfer1::DataType::kFP8: return "fp8"; break; - case nvinfer1::DataType::kBF16: return "bf16"; break; - case nvinfer1::DataType::kINT64: return "int64"; break; - case nvinfer1::DataType::kINT4: return "int4"; break; - case nvinfer1::DataType::kFP4: return "fp4"; break; + case tensorrt_llm::DataType::kFLOAT: return "fp32"; break; + case tensorrt_llm::DataType::kHALF: return "fp16"; break; + case tensorrt_llm::DataType::kINT8: return "int8"; break; + case tensorrt_llm::DataType::kINT32: return "int32"; break; + case tensorrt_llm::DataType::kBOOL: return "bool"; break; + case tensorrt_llm::DataType::kUINT8: return "uint8"; break; + case tensorrt_llm::DataType::kFP8: return "fp8"; break; + case tensorrt_llm::DataType::kBF16: return "bf16"; break; + case tensorrt_llm::DataType::kINT64: return "int64"; break; + case tensorrt_llm::DataType::kINT4: return "int4"; break; + case tensorrt_llm::DataType::kFP4: return "fp4"; break; default: throw std::runtime_error("Unsupported data type"); break; } diff --git a/cpp/include/tensorrt_llm/common/logger.h b/cpp/include/tensorrt_llm/common/logger.h index d14b4c02e992..9073d21f0088 100644 --- a/cpp/include/tensorrt_llm/common/logger.h +++ b/cpp/include/tensorrt_llm/common/logger.h @@ -50,8 +50,6 @@ constexpr std::string_view formatModule(std::string_view module) return "deepgemm"; else if (module == "executor") return "executor"; - else if (module == "executor_worker") - return "exec_wkr"; else if (module == "flash_mla") return "flashmla"; else if (module == "kernels") @@ -60,8 +58,6 @@ constexpr std::string_view formatModule(std::string_view module) return "layers"; else if (module == "nanobind") return "nanobind"; - else if (module == "plugins") - return "plugins"; else if (module == "runtime") return "runtime"; else if (module == "testing") diff --git a/cpp/include/tensorrt_llm/common/tllmDataType.h b/cpp/include/tensorrt_llm/common/tllmDataType.h new file mode 100644 index 000000000000..9e5567ce280e --- /dev/null +++ b/cpp/include/tensorrt_llm/common/tllmDataType.h @@ -0,0 +1,85 @@ +/* + * 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. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" + +#include + +//! \file tllmDataType.h +//! +//! Standalone, TensorRT-free runtime types that replace the \c nvinfer1 +//! types the shared C++ core historically used as common currency: +//! \c tensorrt_llm::common::DataType and \c tensorrt_llm::common::Dims. +//! These are defined here so the retained tree (runtime, batch manager, executor, +//! kernels and the nanobind bridge) compiles and links without the TensorRT +//! library. They are hoisted into the \c tensorrt_llm namespace with +//! using-declarations because they are common currency across the whole tree +//! (\c tensorrt_llm::DataType / \c tensorrt_llm::Dims). +//! +//! The \c DataType enumerator values intentionally mirror the legacy +//! \c nvinfer1::DataType integer values so that previously-serialized executor +//! configs and KV-cache metadata remain byte-compatible. The \c Dims layout +//! mirrors the legacy \c nvinfer1::Dims (\c int32_t \c nbDims followed by +//! \c int64_t \c d[8]) for the same reason. + +TRTLLM_NAMESPACE_BEGIN + +namespace common +{ + +//! \brief Standalone data-type enum. Values mirror the legacy +//! \c nvinfer1::DataType for serialization/format compatibility. +enum class DataType : int32_t +{ + kFLOAT = 0, + kHALF = 1, + kINT8 = 2, + kINT32 = 3, + kBOOL = 4, + kUINT8 = 5, + kFP8 = 6, + kBF16 = 7, + kINT64 = 8, + kINT4 = 9, + kFP4 = 10, + kE8M0 = 11, +}; + +//! \brief Standalone dimensions type. Layout mirrors the legacy +//! \c nvinfer1::Dims (rank plus up to \c MAX_DIMS 64-bit extents) so serialized +//! shapes remain compatible. +class Dims +{ +public: + //! The maximum rank (number of dimensions) supported for a tensor. + static constexpr int32_t MAX_DIMS{8}; + + //! The rank (number of dimensions). + int32_t nbDims; + + //! The extent of each dimension. + int64_t d[MAX_DIMS]; +}; + +} // namespace common + +using common::DataType; +using common::Dims; + +TRTLLM_NAMESPACE_END diff --git a/cpp/include/tensorrt_llm/executor/dataTransceiverState.h b/cpp/include/tensorrt_llm/executor/dataTransceiverState.h index 5067ae61dc83..b00d44d129e7 100644 --- a/cpp/include/tensorrt_llm/executor/dataTransceiverState.h +++ b/cpp/include/tensorrt_llm/executor/dataTransceiverState.h @@ -18,6 +18,7 @@ #include "tensorrt_llm/batch_manager/llmRequest.h" #include "tensorrt_llm/common/assert.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/executor/types.h" #include "tensorrt_llm/runtime/modelConfig.h" #include "tensorrt_llm/runtime/worldConfig.h" @@ -50,7 +51,7 @@ class CacheState final }; CacheState(ModelConfig modelConfig, runtime::WorldConfig const& worldConfig, - std::vector const& attentionLayerNumPerPP, nvinfer1::DataType dataType, + std::vector const& attentionLayerNumPerPP, tensorrt_llm::DataType dataType, AttentionType attentionType = AttentionType::kDEFAULT, int kvFactor = 2, bool enableBlockReuse = false, bool enablePartialReuse = false, bool hasIndexerKCache = false, SizeType32 indexerDimPerHead = 0, SizeType32 indexerKCacheQuantBlockSize = 128, bool indexerKCacheUseFp4 = false) @@ -71,7 +72,7 @@ class CacheState final CacheState(std::vector nbKvHeadPerLayer, SizeType32 sizePerHead, SizeType32 tokensPerBlock, SizeType32 tensorParallelism, SizeType32 pipelineParallelism, SizeType32 contextParallelism, - std::vector const& attentionLayerNumPerPP, nvinfer1::DataType dataType, + std::vector const& attentionLayerNumPerPP, tensorrt_llm::DataType dataType, AttentionType attentionType = AttentionType::kDEFAULT, int kvFactor = 2, bool enableAttentionDP = false, int DPrank = 0, int DPsize = 0, bool enableBlockReuse = false, bool enablePartialReuse = false, bool hasIndexerKCache = false, SizeType32 indexerDimPerHead = 0, SizeType32 indexerKCacheQuantBlockSize = 128, @@ -92,7 +93,7 @@ class CacheState final CacheState(SizeType32 nbAttentionLayers, SizeType32 nbKvHeads, SizeType32 sizePerHead, SizeType32 tokensPerBlock, SizeType32 tensorParallelism, SizeType32 pipelineParallelism, SizeType32 contextParallelism, - std::vector const& attentionLayerNumPerPP, nvinfer1::DataType dataType, + std::vector const& attentionLayerNumPerPP, tensorrt_llm::DataType dataType, AttentionType attentionType = AttentionType::kDEFAULT, int kvFactor = 2, bool enableAttentionDP = false, int DPrank = 0, int DPsize = 0, bool enableBlockReuse = false, bool enablePartialReuse = false, bool hasIndexerKCache = false, SizeType32 indexerDimPerHead = 0, SizeType32 indexerKCacheQuantBlockSize = 128, @@ -238,8 +239,8 @@ class CacheState final RnnModelConfig mModelConfig; /// Number of RNN layers per pipeline parallelism rank. std::vector mLayerNumPerPP; - nvinfer1::DataType mConvStateDataType; - nvinfer1::DataType mSsmStateDataType; + tensorrt_llm::DataType mConvStateDataType; + tensorrt_llm::DataType mSsmStateDataType; [[nodiscard]] bool operator==(RnnCacheState const& other) const noexcept { @@ -263,7 +264,7 @@ class CacheState final return mAttentionConfig; } - [[nodiscard]] nvinfer1::DataType const& getDataType() const + [[nodiscard]] tensorrt_llm::DataType const& getDataType() const { return mDataType; } @@ -308,7 +309,7 @@ class CacheState final } void setRnnConfig(RnnModelConfig rnnModelConfig, std::vector rnnLayerNumPerPP, - nvinfer1::DataType convStateDataType, nvinfer1::DataType ssmStateDataType) + tensorrt_llm::DataType convStateDataType, tensorrt_llm::DataType ssmStateDataType) { mRnnCacheState = RnnCacheState{ std::move(rnnModelConfig), std::move(rnnLayerNumPerPP), convStateDataType, ssmStateDataType}; @@ -325,12 +326,12 @@ class CacheState final return getRnnCacheState().mModelConfig; } - [[nodiscard]] nvinfer1::DataType getConvStateDataType() const + [[nodiscard]] tensorrt_llm::DataType getConvStateDataType() const { return getRnnCacheState().mConvStateDataType; } - [[nodiscard]] nvinfer1::DataType getSsmStateDataType() const + [[nodiscard]] tensorrt_llm::DataType getSsmStateDataType() const { return getRnnCacheState().mSsmStateDataType; } @@ -395,7 +396,7 @@ class CacheState final friend class tensorrt_llm::executor::Serialization; ModelConfig mModelConfig; ParallelConfig mParallelConfig; - nvinfer1::DataType mDataType; + tensorrt_llm::DataType mDataType; AttentionConfig mAttentionConfig; bool mEnableBlockReuse{false}; bool mEnablePartialReuse{false}; diff --git a/cpp/include/tensorrt_llm/executor/disaggServerUtil.h b/cpp/include/tensorrt_llm/executor/disaggServerUtil.h deleted file mode 100644 index b68dce78738a..000000000000 --- a/cpp/include/tensorrt_llm/executor/disaggServerUtil.h +++ /dev/null @@ -1,158 +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. - */ -#pragma once - -#include "tensorrt_llm/executor/executor.h" - -#include -#include -#include -#include -#include - -namespace tensorrt_llm::executor::disagg_executor -{ - -namespace texec = tensorrt_llm::executor; - -struct ResponseWithId -{ - - tensorrt_llm::executor::Response response; - IdType gid; - - ResponseWithId(tensorrt_llm::executor::Response&& response, IdType gid) - : response(std::move(response)) - , gid(gid) - { - } - - ResponseWithId(tensorrt_llm::executor::Response const& response, IdType gid) - : response(response) - , gid(gid) - { - } - - ResponseWithId(ResponseWithId&& other) noexcept - : response(std::move(other.response)) - , gid(other.gid) - { - other.gid = {}; - } - - ResponseWithId(ResponseWithId const& other) = default; - - ResponseWithId& operator=(ResponseWithId&& other) noexcept - { - if (this != &other) - { - response = std::move(other.response); - gid = other.gid; - other.gid = {}; - } - return *this; - } - - ResponseWithId& operator=(ResponseWithId const& other) - { - - if (this != &other) - { - response = other.response; - gid = other.gid; - } - return *this; - } - - ~ResponseWithId() = default; -}; - -class DisaggExecutorOrchestrator -{ -public: - /// @brief Constructs a DisaggExecutorOrchestrator object. - /// - /// @param ctxEnginePaths A vector of file paths to context engine files. - /// @param genEnginePaths A vector of file paths to generation engine files. - /// @param ctxExecutorConfigs A vector of ExecutorConfig for context executors. - /// @param genExecutorConfigs A vector of ExecutorConfig for generation executors. - /// @param hasContextAwaitThreads Whether or not there are threads that receive response for each generation - /// executor. - /// @param hasGenAwaitThreads Whether or not there are threads that receive response for each generation executor. - - DisaggExecutorOrchestrator(std::vector const& ctxEnginePaths, - std::vector const& genEnginePaths, - std::vector const& ctxExecutorConfigs, - std::vector const& genExecutorConfigs, bool hasContextAwaitThreads, - bool hasGenAwaitThreads); - - /// @brief Enqueue context-only requests to context executors. - /// @param requests A vector of context-only requests. - /// @param selectContextId The index of the context executor to use. If `std::nullopt`, the executor that has the - /// smallest number of inflight requests will be used. - /// @param batch If true,enqueue requests in same context executor.If false, will try to use a different executor - /// for each request. - /// @return A vector of global request ids, corresponding to the order of the requests in `requests`, the id - /// returned may be different from the request id in each executor. - [[nodiscard]] std::vector enqueueContext(std::vector const& requests, - std::optional selectContextId = std::nullopt, bool batch = false); - - /// @brief Enqueue generation-only requests to generation executors. - /// @param requests A vector of generation-only requests. - /// @param globalRequestIds A vector of global request ids, corresponding to the order of the requests,and must be - /// the ids returned by the enqueueContext function. - /// @param selectGenIdx The index of the generation executor to use. If `std::nullopt`, the executor that has the - /// smallest number of inflight requests will be used. - /// @param batch If true,enqueue requests in same generation executor.If false, will try to use a different executor - /// for each request. - - void enqueueGeneration(std::vector const& requests, std::vector const& globalRequestIds, - std::optional selectGenIdx = std::nullopt, bool batch = false); - - /// @brief Await for context responses - /// @param timeout The maximum time to wait for new responses - /// @param contextIdx The index of the context executor to use. If `std::nullopt`, return ready responses in all - /// context executors,if `hasContextAwaitThreads` is true, then this parameter must be std::nullopt. - /// @return A vector of responses with corresponding global request ids - - [[nodiscard]] std::vector awaitContextResponses( - std::optional const& timeout, std::optional contextIdx = std::nullopt); - - /// @brief Await for generation responses - /// @param timeout The maximum time to wait for new responses. - /// @param genIdx The index of the generation executor to use. If `std::nullopt`, return ready responses in all - /// generation executors,if `hasGenAwaitThreads` is true, then this parameter must be std::nullopt. - /// @return A vector of responses with corresponding global request ids. - [[nodiscard]] std::vector awaitGenerationResponses( - std::optional const& timeout, std::optional genIdx = std::nullopt); - - /// @brief Indicates if the current process is allowed to enqueueRequests - [[nodiscard]] bool canEnqueue() const; - - /// @brief Get context executors - [[nodiscard]] std::vector> const& getContextExecutors() const; - - /// @brief Get generation executors - [[nodiscard]] std::vector> const& getGenExecutors() const; - - ~DisaggExecutorOrchestrator(); - -private: - class Impl; - std::unique_ptr mImpl; -}; -} // namespace tensorrt_llm::executor::disagg_executor diff --git a/cpp/include/tensorrt_llm/executor/executor.h b/cpp/include/tensorrt_llm/executor/executor.h index 825b8ad75959..acc0efe18966 100644 --- a/cpp/include/tensorrt_llm/executor/executor.h +++ b/cpp/include/tensorrt_llm/executor/executor.h @@ -18,6 +18,7 @@ #include "tensorrt_llm/executor/tensor.h" #include "tensorrt_llm/executor/types.h" +#include "tensorrt_llm/executor/version.h" #include "tensorrt_llm/runtime/common.h" #include "tensorrt_llm/runtime/runtimeDefaults.h" @@ -50,9 +51,11 @@ namespace tensorrt_llm::executor using SizeType32 = tensorrt_llm::runtime::SizeType32; /// @brief Version of TRT-LLM -char const* version() noexcept; +inline char const* version() noexcept +{ + return kTensorRtLlmVersion; +} -class Model; class Serialization; class DataTransceiverState; @@ -1233,6 +1236,11 @@ class DebugConfig SizeType32 mDebugTensorsMaxIterations; }; +/// @brief Configuration for the orchestrator communication mode. +/// @deprecated Orchestrator mode is non-functional: the worker binary it spawned +/// (executorWorker) was removed together with the TensorRT backend. This class is +/// retained only for serialization and Python-binding compatibility and is a +/// candidate for removal in a follow-up (needs API-stability review). class OrchestratorConfig { public: @@ -1835,7 +1843,13 @@ using KVCacheEventData = std::variant attentionDpRank = std::nullopt); + std::optional attentionDpRank = std::nullopt) + : eventId{eventId} + , data{std::move(data)} + , windowSize{windowSize} + , attentionDpRank{attentionDpRank} + { + } /// @brief The unique id of this event IdType eventId; @@ -1863,119 +1877,6 @@ class KVCacheEventManager std::shared_ptr kvCacheManager; }; -/// @brief The executor is responsible for receiving new requests and sending responses, and running the inference -class Executor -{ - -public: - /// @brief - /// @param modelPath Path to the folder that defines the model to run - /// @param modelType The type of model - /// @param executorConfig The configuration for the executor - Executor(std::filesystem::path const& modelPath, ModelType modelType, ExecutorConfig const& executorConfig); - - Executor(std::filesystem::path const& encoderModelPath, std::filesystem::path const& decoderModelPath, - ModelType modelType, ExecutorConfig const& executorConfig); - - Executor(BufferView const& engineBuffer, std::string const& jsonConfigStr, ModelType modelType, - ExecutorConfig const& executorConfig, - std::optional> const& managedWeights = std::nullopt); - - Executor(BufferView const& encoderEngineBuffer, std::string const& encoderJsonConfigStr, - BufferView const& decoderEngineBuffer, std::string const& decoderJsonConfigStr, ModelType modelType, - ExecutorConfig const& executorConfig); - - Executor(std::shared_ptr model, ExecutorConfig const& executorConfig); - - Executor( - std::shared_ptr encoderModel, std::shared_ptr decoderModel, ExecutorConfig const& executorConfig); - - ~Executor(); - Executor(Executor const& executor) = delete; - Executor& operator=(Executor const& executor) = delete; - Executor(Executor&&) = default; - Executor& operator=(Executor&&) = default; - - /// @brief Enqueue a new request - /// @param request The LLM request which contains input tokens and request parameters - /// @return A unique id that identifies the request - [[nodiscard]] IdType enqueueRequest(Request const& request); - - /// @brief Enqueue a batch of request - [[nodiscard]] std::vector enqueueRequests(std::vector const& requests); - - /// @brief Await for ready responses - /// - /// This overload awaits for any ready responses. In particular, if several requests - /// have been enqueued, this method will provide any ready responses without order guarantees. - /// @param timeout The maximum time to wait for new responses - /// @return A vector of responses - [[nodiscard]] std::vector awaitResponses( - std::optional const& timeout = std::nullopt); - - /// @brief Await for ready responses - /// @param id A request id - /// @param timeout The maximum time to wait for new responses - /// @return A vector of responses - [[nodiscard]] std::vector awaitResponses( - IdType const& requestId, std::optional const& timeout = std::nullopt); - - /// @brief Await for multiple ready responses - /// - /// A multiple ID request behaves as if awaitResponses(IdType, timeout) - /// were invoked on all IDs. The returned vector contains - /// a vector of responses per ID in the same order specified by the requestIds. - /// The same behaviour as awaitResponses(IdType, timeout) applies: - /// * Responses may be empty. - /// * If all responses have already been given for one of the requestIds, - /// then this method will hang unless a timeout is specified. - /// @param requestIds Ids requested - /// @param timeout The maximum time to wait for new responses - /// @return A vector of vector of responses - [[nodiscard]] std::vector> awaitResponses( - std::vector const& requestIds, std::optional const& timeout = std::nullopt); - - /// @brief Get the number of ready responses - /// @param requestId An optional request id - /// @return The number of ready responses - [[nodiscard]] SizeType32 getNumResponsesReady(std::optional const& requestId = std::nullopt) const; - - /// @brief Cancel the request with provided request id - /// @param id The request id for which to cancel the response - void cancelRequest(IdType requestId); - - /// @brief Signals the server to shutdown. - /// @details This call is blocking. Only returns when all requests have terminated or timeout has been reached - void shutdown(); - - /// @brief Returns the per-iterations statistics computed since last call to getLatestIterationStats. - /// Contains at most iterStatsMaxIterations iterations, or all iterations when set to -1. - /// @return Iteration stats - std::deque getLatestIterationStats(); - - /// @brief Returns the request stats of each iteration computed since last call to getLatestRequestStats. - /// Contains at most requestStatsMaxIterations iterations, or all iterations when set to -1. - /// @return Request stats grouped by iterations - std::deque getLatestRequestStats(); - - /// @brief Returns the debug tensors of each iteration computed since last call to getLatestDebugTensors. - /// Contains at most debugTensorsMaxIterations iterations. - /// @return Request debug tensors grouped by iterations - std::deque getLatestDebugTensors(); - - /// @brief Indicates if the current process is allowed to enqueueRequests - [[nodiscard]] bool canEnqueueRequests() const; - - /// @brief Indicates if the current process participates in this executor instance - [[nodiscard]] bool isParticipant() const; - - std::optional> getKVCacheEventManager() const; - -private: - class Impl; - std::unique_ptr mImpl; -}; - /// @brief Class with utility functions to serialize statistics to json string class JsonSerialization { diff --git a/cpp/include/tensorrt_llm/plugins/api/tllmPlugin.h b/cpp/include/tensorrt_llm/plugins/api/tllmPlugin.h deleted file mode 100644 index e3d4613e3d00..000000000000 --- a/cpp/include/tensorrt_llm/plugins/api/tllmPlugin.h +++ /dev/null @@ -1,74 +0,0 @@ -/* - * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. - * - * 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. - */ - -#pragma once - -#include -#include - -// Forward declarations -namespace nvinfer1 -{ -class ILoggerFinder; -class ILogger; - -namespace v_1_0 -{ -class IPluginCreator; -class IPluginCreatorV3One; -class IPluginCreatorInterface; -} // namespace v_1_0 - -} // namespace nvinfer1 - -namespace tensorrt_llm::plugins::api -{ - -auto constexpr kDefaultNamespace = "tensorrt_llm"; - -class LoggerManager -{ -public: - //! Set the logger finder. - void setLoggerFinder(nvinfer1::ILoggerFinder* finder); - - //! Get the logger. - [[maybe_unused]] nvinfer1::ILogger* logger(); - - static LoggerManager& getInstance() noexcept; - - static nvinfer1::ILogger* defaultLogger() noexcept; - -private: - LoggerManager() = default; - - nvinfer1::ILoggerFinder* mLoggerFinder{nullptr}; - std::mutex mMutex; -}; -} // namespace tensorrt_llm::plugins::api - -extern "C" -{ - // This function is used for explicitly registering the TRT-LLM plugins and the default logger. - bool initTrtLlmPlugins(void* logger = tensorrt_llm::plugins::api::LoggerManager::defaultLogger(), - char const* libNamespace = tensorrt_llm::plugins::api::kDefaultNamespace); - - // The functions below are used by TensorRT to when loading a shared plugin library with automatic registering. - // see https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#generating-plugin-library - [[maybe_unused]] void setLoggerFinder([[maybe_unused]] nvinfer1::ILoggerFinder* finder); - [[maybe_unused]] nvinfer1::v_1_0::IPluginCreator* const* getPluginCreators(std::int32_t& nbCreators); - [[maybe_unused]] nvinfer1::v_1_0::IPluginCreatorInterface* const* getCreators(std::int32_t& nbCreators); -} diff --git a/cpp/include/tensorrt_llm/runtime/bufferManager.h b/cpp/include/tensorrt_llm/runtime/bufferManager.h index 8357443dc5ea..321a96ba321e 100644 --- a/cpp/include/tensorrt_llm/runtime/bufferManager.h +++ b/cpp/include/tensorrt_llm/runtime/bufferManager.h @@ -17,10 +17,10 @@ #pragma once #include "tensorrt_llm/common/assert.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/runtime/cudaStream.h" #include "tensorrt_llm/runtime/iBuffer.h" #include "tensorrt_llm/runtime/iTensor.h" -#include #include #include @@ -62,63 +62,63 @@ class BufferManager } } - static auto constexpr kBYTE_TYPE = nvinfer1::DataType::kUINT8; + static auto constexpr kBYTE_TYPE = tensorrt_llm::DataType::kUINT8; //! \brief Allocates an `IBuffer` of the given size on the GPU, using cudaMallocAsync. - [[nodiscard]] IBufferPtr gpu(std::size_t size, nvinfer1::DataType type = kBYTE_TYPE) const; + [[nodiscard]] IBufferPtr gpu(std::size_t size, tensorrt_llm::DataType type = kBYTE_TYPE) const; //! \brief Allocates an `ITensor` of the given dimensions on the GPU, using cudaMallocAsync. - [[nodiscard]] ITensorPtr gpu(nvinfer1::Dims dims, nvinfer1::DataType type = kBYTE_TYPE) const; + [[nodiscard]] ITensorPtr gpu(tensorrt_llm::Dims dims, tensorrt_llm::DataType type = kBYTE_TYPE) const; //! \brief Allocates an `IBuffer` of the given size on the GPU, using cudaMalloc. - [[nodiscard]] static IBufferPtr gpuSync(std::size_t size, nvinfer1::DataType type = kBYTE_TYPE); + [[nodiscard]] static IBufferPtr gpuSync(std::size_t size, tensorrt_llm::DataType type = kBYTE_TYPE); //! \brief Allocates an `ITensor` of the given dimensions on the GPU, using cudaMalloc. - [[nodiscard]] static ITensorPtr gpuSync(nvinfer1::Dims dims, nvinfer1::DataType type = kBYTE_TYPE); + [[nodiscard]] static ITensorPtr gpuSync(tensorrt_llm::Dims dims, tensorrt_llm::DataType type = kBYTE_TYPE); //! \brief Allocates an `IBuffer` of the given size on the CPU. - [[nodiscard]] static IBufferPtr cpu(std::size_t size, nvinfer1::DataType type = kBYTE_TYPE); + [[nodiscard]] static IBufferPtr cpu(std::size_t size, tensorrt_llm::DataType type = kBYTE_TYPE); //! \brief Allocates an `ITensor` of the given dimensions on the CPU. - [[nodiscard]] static ITensorPtr cpu(nvinfer1::Dims dims, nvinfer1::DataType type = kBYTE_TYPE); + [[nodiscard]] static ITensorPtr cpu(tensorrt_llm::Dims dims, tensorrt_llm::DataType type = kBYTE_TYPE); //! \brief Allocates a pinned `IBuffer` of the given size on the CPU. - [[nodiscard]] static IBufferPtr pinned(std::size_t size, nvinfer1::DataType type = kBYTE_TYPE); + [[nodiscard]] static IBufferPtr pinned(std::size_t size, tensorrt_llm::DataType type = kBYTE_TYPE); //! \brief Allocates a pinned `ITensor` of the given dimensions on the CPU. - [[nodiscard]] static ITensorPtr pinned(nvinfer1::Dims dims, nvinfer1::DataType type = kBYTE_TYPE); + [[nodiscard]] static ITensorPtr pinned(tensorrt_llm::Dims dims, tensorrt_llm::DataType type = kBYTE_TYPE); //! \brief Allocates a pinned `IBuffer` of the given size on the CPU in the default memory pool. - [[nodiscard]] static IBufferPtr pinnedPool(std::size_t size, nvinfer1::DataType type = kBYTE_TYPE); + [[nodiscard]] static IBufferPtr pinnedPool(std::size_t size, tensorrt_llm::DataType type = kBYTE_TYPE); //! \brief Allocates a pinned `ITensor` of the given dimensions on the CPU in the default memory pool. - [[nodiscard]] static ITensorPtr pinnedPool(nvinfer1::Dims dims, nvinfer1::DataType type = kBYTE_TYPE); + [[nodiscard]] static ITensorPtr pinnedPool(tensorrt_llm::Dims dims, tensorrt_llm::DataType type = kBYTE_TYPE); //! \brief Allocates an `IBuffer` of the given size in UVM. - [[nodiscard]] static IBufferPtr managed(std::size_t size, nvinfer1::DataType type = kBYTE_TYPE); + [[nodiscard]] static IBufferPtr managed(std::size_t size, tensorrt_llm::DataType type = kBYTE_TYPE); //! \brief Allocates an `ITensor` of the given dimensions in UVM. - [[nodiscard]] static ITensorPtr managed(nvinfer1::Dims dims, nvinfer1::DataType type = kBYTE_TYPE); + [[nodiscard]] static ITensorPtr managed(tensorrt_llm::Dims dims, tensorrt_llm::DataType type = kBYTE_TYPE); //! \brief Allocates an `ITensor` of the given dimensions for NVLS - [[nodiscard]] static ITensorPtr ipcNvls(std::set ranks, nvinfer1::Dims dims, nvinfer1::DataType type); + [[nodiscard]] static ITensorPtr ipcNvls(std::set ranks, tensorrt_llm::Dims dims, tensorrt_llm::DataType type); //! \brief Allocates an `IBuffer` of the given size and memory type. [[nodiscard]] IBufferPtr allocate( - MemoryType memoryType, std::size_t size, nvinfer1::DataType type = kBYTE_TYPE) const; + MemoryType memoryType, std::size_t size, tensorrt_llm::DataType type = kBYTE_TYPE) const; //! \brief Allocates an `ITensor` of the given dimensions and memory type. [[nodiscard]] ITensorPtr allocate( - MemoryType memoryType, nvinfer1::Dims dims, nvinfer1::DataType type = kBYTE_TYPE) const; + MemoryType memoryType, tensorrt_llm::Dims dims, tensorrt_llm::DataType type = kBYTE_TYPE) const; //! \brief Create an empty `IBuffer` of the given memory type. It may be resized later. - [[nodiscard]] IBufferPtr emptyBuffer(MemoryType memoryType, nvinfer1::DataType type = kBYTE_TYPE) const + [[nodiscard]] IBufferPtr emptyBuffer(MemoryType memoryType, tensorrt_llm::DataType type = kBYTE_TYPE) const { return allocate(memoryType, 0, type); } //! \brief Create an empty `ITensor` of the given memory type. It may be reshaped later. - [[nodiscard]] ITensorPtr emptyTensor(MemoryType memoryType, nvinfer1::DataType type = kBYTE_TYPE) const + [[nodiscard]] ITensorPtr emptyTensor(MemoryType memoryType, tensorrt_llm::DataType type = kBYTE_TYPE) const { return allocate(memoryType, ITensor::makeShape({}), type); } @@ -167,7 +167,7 @@ class BufferManager //! \brief Copy `src` into a new `ITensor` with a potentially different memory type. template - [[nodiscard]] ITensorPtr copyFrom(T* src, nvinfer1::Dims dims, MemoryType memoryType) const + [[nodiscard]] ITensorPtr copyFrom(T* src, tensorrt_llm::Dims dims, MemoryType memoryType) const { auto buffer = allocate(memoryType, dims, TRTDataType>::value); copy(src, *buffer); @@ -176,7 +176,7 @@ class BufferManager //! \brief Copy `src` into a new `ITensor` with a potentially different memory type. template - [[nodiscard]] ITensorPtr copyFrom(std::vector const& src, nvinfer1::Dims dims, MemoryType memoryType) const + [[nodiscard]] ITensorPtr copyFrom(std::vector const& src, tensorrt_llm::Dims dims, MemoryType memoryType) const { TLLM_CHECK_WITH_INFO(src.size() == ITensor::volumeNonNegative(dims), common::fmtstr("[TensorRT-LLM][ERROR] Incompatible size %lu and dims %s", src.size(), diff --git a/cpp/include/tensorrt_llm/runtime/decoderState.h b/cpp/include/tensorrt_llm/runtime/decoderState.h index 95d7ff0ffac9..ea2c767c0478 100644 --- a/cpp/include/tensorrt_llm/runtime/decoderState.h +++ b/cpp/include/tensorrt_llm/runtime/decoderState.h @@ -18,6 +18,7 @@ #include "decodingInput.h" #include "decodingOutput.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/runtime/bufferManager.h" #include "tensorrt_llm/runtime/iTensor.h" #include "tensorrt_llm/runtime/speculativeDecodingMode.h" @@ -52,7 +53,7 @@ class DecoderState //! @brief Setup buffers for the decoder excluding speculative decoding. void setup(SizeType32 maxNumSequences, SizeType32 maxBeamWidth, SizeType32 maxAttentionWindow, - SizeType32 sinkTokenLength, SizeType32 maxSequenceLength, nvinfer1::DataType dtype, + SizeType32 sinkTokenLength, SizeType32 maxSequenceLength, tensorrt_llm::DataType dtype, ModelConfig const& modelConfig, WorldConfig const& worldConfig, BufferManager const& bufferManager); //! @brief Setup buffers for the cache indirection. @@ -62,7 +63,7 @@ class DecoderState //! @brief Setup buffers for speculative decoding. void setupSpeculativeDecoding(SpeculativeDecodingMode const& speculativeDecodingMode, - SizeType32 maxTokensPerEngineStep, nvinfer1::DataType dtype, ModelConfig const& modelConfig, + SizeType32 maxTokensPerEngineStep, tensorrt_llm::DataType dtype, ModelConfig const& modelConfig, WorldConfig const& worldConfig, BufferManager const& bufferManager); //! @brief Disable lookahead decoding. @@ -199,7 +200,7 @@ class DecoderState [[nodiscard]] DecodingOutput& getJointDecodingOutput() const; private: - void setupBuffers(nvinfer1::DataType dtype, BufferManager const& bufferManager); + void setupBuffers(tensorrt_llm::DataType dtype, BufferManager const& bufferManager); void reshapeBuffers(SizeType32 maxBatchSize, SizeType32 maxBeamWidth, SizeType32 maxAttentionWindow, SizeType32 sinkTokenLength, SizeType32 maxSequenceLength, ModelConfig const& modelConfig, WorldConfig const& worldConfig, BufferManager const& bufferManager); @@ -208,8 +209,8 @@ class DecoderState void reshapeCacheIndirectionBuffers( SizeType32 maxBatchSize, SizeType32 maxBeamWidth, SizeType32 maxAttentionWindow); - void setupSpeculativeDecodingBuffers( - SpeculativeDecodingMode speculativeDecodingMode, nvinfer1::DataType dtype, BufferManager const& bufferManager); + void setupSpeculativeDecodingBuffers(SpeculativeDecodingMode speculativeDecodingMode, tensorrt_llm::DataType dtype, + BufferManager const& bufferManager); void reshapeSpeculativeDecodingBuffers(SpeculativeDecodingMode const& speculativeDecodingMode, SizeType32 maxTokensPerEngineStep, ModelConfig const& modelConfig, WorldConfig const& worldConfig, BufferManager const& bufferManager); diff --git a/cpp/include/tensorrt_llm/runtime/gptDecoder.h b/cpp/include/tensorrt_llm/runtime/gptDecoder.h index 7e0cc1bb56d2..5a785e84fe75 100644 --- a/cpp/include/tensorrt_llm/runtime/gptDecoder.h +++ b/cpp/include/tensorrt_llm/runtime/gptDecoder.h @@ -22,7 +22,7 @@ #include "tensorrt_llm/runtime/decodingOutput.h" #include "tensorrt_llm/runtime/samplingConfig.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include #include @@ -55,7 +55,7 @@ class IGptDecoder /// @param explicitDraftTokensDType is only used by ExplicitDraftTokens model to WAR the lack of bf16 decoder. virtual void setup(SamplingConfig const& samplingConfig, size_t batchSize, TensorConstPtr const& batchSlots, std::optional const& output = std::nullopt, - std::optional explicitDraftTokensDType = std::nullopt, + std::optional explicitDraftTokensDType = std::nullopt, std::optional> const& lookaheadPrompt = std::nullopt, std::optional> const& lookaheadAlgoConfigs = std::nullopt) = 0; @@ -70,7 +70,7 @@ class IGptDecoder std::optional const& samplingConfig, SizeType32 batchSize, TensorConstPtr batchSlots) = 0; - static std::unique_ptr create(executor::DecodingMode const& mode, nvinfer1::DataType dtype, + static std::unique_ptr create(executor::DecodingMode const& mode, tensorrt_llm::DataType dtype, size_t maxNumSequences, size_t maxBeamWidth, size_t vocabSize, size_t vocabSizePadded, BufferManager::CudaStreamPtr const& stream, std::shared_ptr const& speculativeDecodingModule = nullptr); @@ -90,7 +90,7 @@ class GptDecoder : public virtual IGptDecoder void setup(SamplingConfig const& samplingConfig, size_t batchSize, TensorConstPtr const& batchSlots, std::optional const& output = std::nullopt, - std::optional explicitDraftTokensDType = std::nullopt, + std::optional explicitDraftTokensDType = std::nullopt, std::optional> const& lookaheadPrompt = std::nullopt, std::optional> const& lookaheadAlgoConfigs = std::nullopt) override; @@ -121,17 +121,17 @@ class GptDecoder : public virtual IGptDecoder executor::DecodingMode mDecodingMode; }; -inline std::unique_ptr IGptDecoder::create(executor::DecodingMode const& mode, nvinfer1::DataType dtype, - size_t maxNumSequences, size_t maxBeamWidth, size_t vocabSize, size_t vocabSizePadded, +inline std::unique_ptr IGptDecoder::create(executor::DecodingMode const& mode, + tensorrt_llm::DataType dtype, size_t maxNumSequences, size_t maxBeamWidth, size_t vocabSize, size_t vocabSizePadded, BufferManager::CudaStreamPtr const& stream, std::shared_ptr const& speculativeDecodingModule) { switch (dtype) { - case nvinfer1::DataType::kFLOAT: + case tensorrt_llm::DataType::kFLOAT: return std::make_unique>( mode, maxNumSequences, maxBeamWidth, vocabSize, vocabSizePadded, stream, speculativeDecodingModule); - case nvinfer1::DataType::kHALF: + case tensorrt_llm::DataType::kHALF: return std::make_unique>( mode, maxNumSequences, maxBeamWidth, vocabSize, vocabSizePadded, stream, speculativeDecodingModule); default: diff --git a/cpp/include/tensorrt_llm/runtime/gptDecoderBatched.h b/cpp/include/tensorrt_llm/runtime/gptDecoderBatched.h index 9fcd3262c8ca..d5447f441163 100644 --- a/cpp/include/tensorrt_llm/runtime/gptDecoderBatched.h +++ b/cpp/include/tensorrt_llm/runtime/gptDecoderBatched.h @@ -16,6 +16,7 @@ #pragma once +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/runtime/bufferManager.h" #include "tensorrt_llm/runtime/cudaEvent.h" #include "tensorrt_llm/runtime/cudaStream.h" @@ -48,7 +49,7 @@ class GptDecoderBatched : public IGptDecoderBatched explicit GptDecoderBatched(CudaStreamPtr stream); void setup(executor::DecodingMode const& mode, SizeType32 maxNumSequences, SizeType32 maxBeamWidth, - nvinfer1::DataType dtype, ModelConfig const& modelConfig, WorldConfig const& worldConfig) override; + tensorrt_llm::DataType dtype, ModelConfig const& modelConfig, WorldConfig const& worldConfig) override; void disableLookahead(RequestVector const& genRequests, TensorPtr const& batchSlots) override; diff --git a/cpp/include/tensorrt_llm/runtime/iBuffer.h b/cpp/include/tensorrt_llm/runtime/iBuffer.h index 91d5cd739f32..bf63d3a0da7b 100644 --- a/cpp/include/tensorrt_llm/runtime/iBuffer.h +++ b/cpp/include/tensorrt_llm/runtime/iBuffer.h @@ -22,7 +22,7 @@ #include "tensorrt_llm/kernels/kvCacheIndex.h" #include "tensorrt_llm/runtime/common.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include #ifdef ENABLE_FP8 @@ -88,13 +88,13 @@ struct MemoryTypeString }; //! \brief For converting a TensorRT data type to a C++ data type. -template +template struct DataTypeTraits { }; template <> -struct DataTypeTraits +struct DataTypeTraits { using type = float; static char constexpr name[] = "float"; @@ -102,7 +102,7 @@ struct DataTypeTraits }; template <> -struct DataTypeTraits +struct DataTypeTraits { using type = half; static char constexpr name[] = "half"; @@ -110,7 +110,7 @@ struct DataTypeTraits }; template <> -struct DataTypeTraits +struct DataTypeTraits { using type = std::int8_t; static char constexpr name[] = "int8"; @@ -118,7 +118,7 @@ struct DataTypeTraits }; template <> -struct DataTypeTraits +struct DataTypeTraits { using type = std::int32_t; static char constexpr name[] = "int32"; @@ -126,7 +126,7 @@ struct DataTypeTraits }; template <> -struct DataTypeTraits +struct DataTypeTraits { using type = std::int64_t; static char constexpr name[] = "int64"; @@ -134,7 +134,7 @@ struct DataTypeTraits }; template <> -struct DataTypeTraits +struct DataTypeTraits { using type = std::uint32_t; static char constexpr name[] = "uint32"; @@ -142,7 +142,7 @@ struct DataTypeTraits }; template <> -struct DataTypeTraits +struct DataTypeTraits { using type = std::uint64_t; static char constexpr name[] = "uint64"; @@ -150,7 +150,7 @@ struct DataTypeTraits }; template -struct DataTypeTraits +struct DataTypeTraits { using type = bool; static char constexpr name[] = "bool"; @@ -158,7 +158,7 @@ struct DataTypeTraits }; template -struct DataTypeTraits +struct DataTypeTraits { using type = std::uint8_t; static char constexpr name[] = "uint8"; @@ -167,7 +167,7 @@ struct DataTypeTraits #ifdef ENABLE_BF16 template <> -struct DataTypeTraits +struct DataTypeTraits { using type = __nv_bfloat16; static char constexpr name[] = "bfloat16"; @@ -177,7 +177,7 @@ struct DataTypeTraits #ifdef ENABLE_FP8 template <> -struct DataTypeTraits +struct DataTypeTraits { using type = __nv_fp8_e4m3; static char constexpr name[] = "fp8"; @@ -185,7 +185,7 @@ struct DataTypeTraits }; #endif -template +template struct DataTypeTraits { using type = typename DataTypeTraits::type*; @@ -193,26 +193,26 @@ struct DataTypeTraits static auto constexpr size = sizeof(type); }; -//! \brief A wrapper around `nvinfer1::DataType` that provides a support for pointer types. +//! \brief A wrapper around `tensorrt_llm::DataType` that provides a support for pointer types. class BufferDataType { public: constexpr BufferDataType( // NOLINT(*-explicit-constructor) - nvinfer1::DataType dataType, bool _unsigned = false, bool pointer = false) + tensorrt_llm::DataType dataType, bool _unsigned = false, bool pointer = false) : mDataType{dataType} , mUnsigned{_unsigned} , mPointer{pointer} { } - static auto constexpr kTrtPointerType = nvinfer1::DataType::kINT64; + static auto constexpr kTrtPointerType = tensorrt_llm::DataType::kINT64; - constexpr operator nvinfer1::DataType() const noexcept // NOLINT(*-explicit-constructor) + constexpr operator tensorrt_llm::DataType() const noexcept // NOLINT(*-explicit-constructor) { return mPointer ? kTrtPointerType : mDataType; } - [[nodiscard]] constexpr nvinfer1::DataType getDataType() const noexcept + [[nodiscard]] constexpr tensorrt_llm::DataType getDataType() const noexcept { return mDataType; } @@ -226,24 +226,24 @@ class BufferDataType { switch (mDataType) { - case nvinfer1::DataType::kBOOL: [[fallthrough]]; - case nvinfer1::DataType::kUINT8: return true; + case tensorrt_llm::DataType::kBOOL: [[fallthrough]]; + case tensorrt_llm::DataType::kUINT8: return true; default: return mUnsigned; } } [[nodiscard]] constexpr std::size_t getSize() const noexcept { - return tensorrt_llm::common::getDTypeSize(static_cast(*this)); + return tensorrt_llm::common::getDTypeSize(static_cast(*this)); } [[nodiscard]] constexpr std::size_t getSizeInBits() const noexcept { - return tensorrt_llm::common::getDTypeSizeInBits(static_cast(*this)); + return tensorrt_llm::common::getDTypeSizeInBits(static_cast(*this)); } private: - nvinfer1::DataType mDataType; + tensorrt_llm::DataType mDataType; bool mUnsigned; bool mPointer; }; @@ -257,62 +257,62 @@ struct TRTDataType template <> struct TRTDataType { - static constexpr auto value = nvinfer1::DataType::kFLOAT; + static constexpr auto value = tensorrt_llm::DataType::kFLOAT; }; template <> struct TRTDataType { - static constexpr auto value = nvinfer1::DataType::kHALF; + static constexpr auto value = tensorrt_llm::DataType::kHALF; }; template <> struct TRTDataType { - static constexpr auto value = nvinfer1::DataType::kINT8; + static constexpr auto value = tensorrt_llm::DataType::kINT8; }; template <> struct TRTDataType { - static constexpr auto value = nvinfer1::DataType::kINT32; + static constexpr auto value = tensorrt_llm::DataType::kINT32; }; template <> struct TRTDataType { - static constexpr auto value = BufferDataType{nvinfer1::DataType::kINT32, true}; + static constexpr auto value = BufferDataType{tensorrt_llm::DataType::kINT32, true}; }; template <> struct TRTDataType { - static constexpr auto value = nvinfer1::DataType::kINT64; + static constexpr auto value = tensorrt_llm::DataType::kINT64; }; template <> struct TRTDataType { - static constexpr auto value = BufferDataType{nvinfer1::DataType::kINT64, true}; + static constexpr auto value = BufferDataType{tensorrt_llm::DataType::kINT64, true}; }; template <> struct TRTDataType { - static constexpr auto value = nvinfer1::DataType::kBOOL; + static constexpr auto value = tensorrt_llm::DataType::kBOOL; }; template <> struct TRTDataType { - static constexpr auto value = nvinfer1::DataType::kUINT8; + static constexpr auto value = tensorrt_llm::DataType::kUINT8; }; #ifdef ENABLE_BF16 template <> struct TRTDataType<__nv_bfloat16> { - static constexpr auto value = nvinfer1::DataType::kBF16; + static constexpr auto value = tensorrt_llm::DataType::kBF16; }; #endif @@ -320,7 +320,7 @@ struct TRTDataType<__nv_bfloat16> template <> struct TRTDataType<__nv_fp8_e4m3> { - static constexpr auto value = nvinfer1::DataType::kFP8; + static constexpr auto value = tensorrt_llm::DataType::kFP8; }; #endif @@ -380,7 +380,7 @@ class IBuffer using SharedPtr = std::shared_ptr; using UniqueConstPtr = std::unique_ptr; using SharedConstPtr = std::shared_ptr; - using DataType = nvinfer1::DataType; + using DataType = tensorrt_llm::DataType; //! //! \brief Returns a pointer to underlying array. diff --git a/cpp/include/tensorrt_llm/runtime/iGptDecoderBatched.h b/cpp/include/tensorrt_llm/runtime/iGptDecoderBatched.h index ab55b754f9be..b664bc007f0e 100644 --- a/cpp/include/tensorrt_llm/runtime/iGptDecoderBatched.h +++ b/cpp/include/tensorrt_llm/runtime/iGptDecoderBatched.h @@ -16,6 +16,7 @@ #pragma once +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/runtime/cudaEvent.h" #include "tensorrt_llm/runtime/cudaStream.h" #include "tensorrt_llm/runtime/iTensor.h" @@ -51,7 +52,7 @@ class IGptDecoderBatched //! @brief Setup the decoder before calling `forward()` virtual void setup(executor::DecodingMode const& mode, SizeType32 maxNumSequences, SizeType32 maxBeamWidth, - nvinfer1::DataType dtype, ModelConfig const& modelConfig, WorldConfig const& worldConfig) + tensorrt_llm::DataType dtype, ModelConfig const& modelConfig, WorldConfig const& worldConfig) = 0; //! @brief Disable Lookahead decoding. diff --git a/cpp/include/tensorrt_llm/runtime/iTensor.h b/cpp/include/tensorrt_llm/runtime/iTensor.h index eb5c10eeb691..a85291dd8263 100644 --- a/cpp/include/tensorrt_llm/runtime/iTensor.h +++ b/cpp/include/tensorrt_llm/runtime/iTensor.h @@ -20,7 +20,7 @@ #include "tensorrt_llm/runtime/common.h" #include "tensorrt_llm/runtime/iBuffer.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include #include @@ -33,11 +33,6 @@ #include #include -namespace nvinfer1 -{ -class IExecutionContext; -} - namespace tensorrt_llm::runtime { @@ -50,7 +45,7 @@ class ITensor : virtual public IBuffer using SharedPtr = std::shared_ptr; using UniqueConstPtr = std::unique_ptr; using SharedConstPtr = std::shared_ptr; - using Shape = nvinfer1::Dims; + using Shape = tensorrt_llm::Dims; using DimType64 = std::remove_reference_t; using TensorMap = runtime::StringPtrMap; @@ -352,9 +347,9 @@ class ITensor : virtual public IBuffer //! \param shape The shape of the tensor. //! \param capacity The capacity of the buffer. //! \return An `ITensor`. - static UniquePtr wrap(void* data, nvinfer1::DataType type, Shape const& shape, std::size_t capacity); + static UniquePtr wrap(void* data, tensorrt_llm::DataType type, Shape const& shape, std::size_t capacity); - static UniquePtr wrap(void* data, nvinfer1::DataType type, Shape const& shape) + static UniquePtr wrap(void* data, tensorrt_llm::DataType type, Shape const& shape) { return wrap(data, type, shape, volumeNonNegative(shape)); } diff --git a/cpp/include/tensorrt_llm/runtime/lookaheadBuffers.h b/cpp/include/tensorrt_llm/runtime/lookaheadBuffers.h index ecaa439f2d52..26c6e3886be4 100644 --- a/cpp/include/tensorrt_llm/runtime/lookaheadBuffers.h +++ b/cpp/include/tensorrt_llm/runtime/lookaheadBuffers.h @@ -17,9 +17,9 @@ #pragma once #include "tensorrt_llm/executor/executor.h" +#include "tensorrt_llm/runtime/bufferManager.h" #include "tensorrt_llm/runtime/iTensor.h" #include "tensorrt_llm/runtime/modelConfig.h" -#include "tensorrt_llm/runtime/tllmRuntime.h" #include "tensorrt_llm/runtime/worldConfig.h" namespace tensorrt_llm::runtime @@ -37,47 +37,4 @@ class LookaheadDecodingBuffers TensorPtr positionIds; }; -class LookaheadRuntimeBuffers -{ -public: - using TensorPtr = ITensor::SharedPtr; - using TensorMap = StringPtrMap; - - LookaheadRuntimeBuffers(SizeType32 maxBatchSize, SizeType32 maxBeamWidth, BufferManager const& manager, - ModelConfig const& modelConfig, WorldConfig const& worldConfig, executor::DecodingConfig const& decodingConfig, - TllmRuntime const& runtime); - - void setFromInputs(SizeType32 numCtxSequences, SizeType32 numGenSequences, ITensor const& requestTypes, - ITensor const& seqSlots, LookaheadDecodingBuffers const& decoderLookaheadBuffers, TllmRuntime const& runtime, - ModelConfig const& modelConfig, WorldConfig const& worldConfig) const; - - void reshape(SizeType32 numCtxSequences, SizeType32 numGenSequences, SizeType32 tokensPerStep); - - void insertInputTensors(TensorMap& inputBuffers, TensorMap& outputBuffers, WorldConfig const& worldConfig) const; - - void enableLookaheadDecoding(SizeType32 maxBatchSize, SizeType32 tokensPerStep); - - void disableLookaheadDecoding(); - -public: - TensorPtr cumSumLength; // [1] the cumulative sum of generation length, on pinned - TensorPtr packedMasksDevice; // [forwardBatchSize, tokensPerStep, numPackedMasks], on gpu - TensorPtr generationLengthsDevice; // [forwardBatchSize], on gpu - TensorPtr positionOffsetsDevice; // [forwardBatchSize, tokensPerStep], on gpu - TensorPtr positionIdsDevice; // [forwardBatchSize, tokensPerStep], on gpu - - TensorPtr packedMaskHost; - TensorPtr generationLengthsHost; - TensorPtr positionOffsetsHost; - TensorPtr positionIdsHost; - - TensorPtr packedMaskHostCopy; - TensorPtr generationLengthsHostCopy; - TensorPtr positionOffsetsHostCopy; - TensorPtr positionIdsHostCopy; - TensorPtr useSpecDecoding; - - TensorPtr batchSlotsHostCopy; -}; - } // namespace tensorrt_llm::runtime diff --git a/cpp/include/tensorrt_llm/runtime/loraCache.h b/cpp/include/tensorrt_llm/runtime/loraCache.h index 1d242cdc80c5..eb4ef57494ea 100644 --- a/cpp/include/tensorrt_llm/runtime/loraCache.h +++ b/cpp/include/tensorrt_llm/runtime/loraCache.h @@ -25,8 +25,6 @@ #include "tensorrt_llm/runtime/modelConfig.h" #include "tensorrt_llm/runtime/worldConfig.h" -#include - #include #include #include diff --git a/cpp/include/tensorrt_llm/runtime/loraCachePageManagerConfig.h b/cpp/include/tensorrt_llm/runtime/loraCachePageManagerConfig.h index a995304e94ec..cf1b1a6aac18 100644 --- a/cpp/include/tensorrt_llm/runtime/loraCachePageManagerConfig.h +++ b/cpp/include/tensorrt_llm/runtime/loraCachePageManagerConfig.h @@ -20,7 +20,7 @@ #include "tensorrt_llm/runtime/common.h" #include "tensorrt_llm/runtime/iBuffer.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include #include @@ -36,7 +36,7 @@ namespace tensorrt_llm::runtime class LoraCachePageManagerConfig { public: - explicit constexpr LoraCachePageManagerConfig(runtime::MemoryType memType, nvinfer1::DataType dType, + explicit constexpr LoraCachePageManagerConfig(runtime::MemoryType memType, tensorrt_llm::DataType dType, SizeType32 totalNumPages, SizeType32 maxPagesPerBlock, SizeType32 slotsPerPage, SizeType32 pageWidth, SizeType32 numCopyStreams) : mMemoryType(memType) @@ -59,12 +59,12 @@ class LoraCachePageManagerConfig mMemoryType = memoryType; } - [[nodiscard]] nvinfer1::DataType constexpr getDataType() const noexcept + [[nodiscard]] tensorrt_llm::DataType constexpr getDataType() const noexcept { return mDataType; } - void constexpr setDataType(nvinfer1::DataType const& dtype) noexcept + void constexpr setDataType(tensorrt_llm::DataType const& dtype) noexcept { mDataType = dtype; } @@ -131,7 +131,7 @@ class LoraCachePageManagerConfig private: runtime::MemoryType mMemoryType; - nvinfer1::DataType mDataType; + tensorrt_llm::DataType mDataType; /* * Number cache pages in the cache. @@ -154,7 +154,7 @@ inline std::ostream& operator<<(std::ostream& os, LoraCachePageManagerConfig con { os << "{" << "memoryType=" << static_cast::type>(c.getMemoryType()) - << " dataType=" << static_cast::type>(c.getDataType()) + << " dataType=" << static_cast::type>(c.getDataType()) << " totalNumPages=" << c.getTotalNumPages() << " maxPagesPerBlock=" << c.getMaxPagesPerBlock() << " slotsPerPage=" << c.getSlotsPerPage() << " pageWidth=" << c.getPageWidth() << " initToZero=" << c.getInitToZero() << "}"; diff --git a/cpp/include/tensorrt_llm/runtime/modelConfig.h b/cpp/include/tensorrt_llm/runtime/modelConfig.h index 5bfe7bce9d58..b5f18da07f3b 100644 --- a/cpp/include/tensorrt_llm/runtime/modelConfig.h +++ b/cpp/include/tensorrt_llm/runtime/modelConfig.h @@ -23,7 +23,7 @@ #include "tensorrt_llm/runtime/speculativeDecodingMode.h" #include "tensorrt_llm/runtime/speculativeDecodingModule.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include namespace tensorrt_llm::runtime @@ -101,7 +101,7 @@ class ModelConfig }; explicit ModelConfig(SizeType32 vocabSize, SizeType32 nbLayers, SizeType32 nbAttentionLayers, - SizeType32 nbRnnLayers, SizeType32 nbHeads, SizeType32 hiddenSize, nvinfer1::DataType dtype) + SizeType32 nbRnnLayers, SizeType32 nbHeads, SizeType32 hiddenSize, tensorrt_llm::DataType dtype) : mVocabSize(vocabSize) , mNbLayers(nbLayers) , mNbAttentionLayers(nbAttentionLayers) @@ -137,7 +137,7 @@ class ModelConfig , mUsePositionEmbedding(false) , mUseTokenTypeEmbedding(false) , mSpeculativeDecodingMode(SpeculativeDecodingMode::None()) - , mLogitsDtype(nvinfer1::DataType::kFLOAT) + , mLogitsDtype(tensorrt_llm::DataType::kFLOAT) , mUseShapeInference(true) , mManageWeightsType(ManageWeightsType::kDisabled) , mSkipCrossAttnBlocks(false) @@ -331,7 +331,7 @@ class ModelConfig mSizePerHead = sizePerHead; } - [[nodiscard]] nvinfer1::DataType constexpr getDataType() const noexcept + [[nodiscard]] tensorrt_llm::DataType constexpr getDataType() const noexcept { return mDataType; } @@ -735,20 +735,20 @@ class ModelConfig resetSpeculativeDecodingModule(); } - [[nodiscard]] nvinfer1::DataType getKvDataType() const + [[nodiscard]] tensorrt_llm::DataType getKvDataType() const { if (getQuantMode().hasFp8KvCache()) { - return nvinfer1::DataType::kFP8; + return tensorrt_llm::DataType::kFP8; } if (getQuantMode().hasInt8KvCache()) { - return nvinfer1::DataType::kINT8; + return tensorrt_llm::DataType::kINT8; } else if (getQuantMode().hasFp4KvCache()) { #ifdef ENABLE_FP4 - return nvinfer1::DataType::kFP4; + return tensorrt_llm::DataType::kFP4; #else throw std::runtime_error("Model has FP4 KV cache, but TRT-LLM was not compiled with FP4 enabled."); #endif @@ -800,22 +800,22 @@ class ModelConfig return mSpeculativeDecodingMode; } - void setLogitsDtype(nvinfer1::DataType inputDtype) noexcept + void setLogitsDtype(tensorrt_llm::DataType inputDtype) noexcept { mLogitsDtype = inputDtype; } - [[nodiscard]] nvinfer1::DataType constexpr getLogitsDtype() const noexcept + [[nodiscard]] tensorrt_llm::DataType constexpr getLogitsDtype() const noexcept { return mLogitsDtype; } - void setGemmAllReduceDtype(nvinfer1::DataType inputDtype) noexcept + void setGemmAllReduceDtype(tensorrt_llm::DataType inputDtype) noexcept { mGemmAllReduceDtype = inputDtype; } - [[nodiscard]] nvinfer1::DataType constexpr getGemmAllReduceDtype() const noexcept + [[nodiscard]] tensorrt_llm::DataType constexpr getGemmAllReduceDtype() const noexcept { return mGemmAllReduceDtype; } @@ -945,10 +945,10 @@ class ModelConfig SizeType32 mNbHeads; SizeType32 mHiddenSize; SizeType32 mSizePerHead; - nvinfer1::DataType mDataType; + tensorrt_llm::DataType mDataType; bool mUseGptAttentionPlugin; bool mUseGemmAllReducePlugin; - nvinfer1::DataType mGemmAllReduceDtype; + tensorrt_llm::DataType mGemmAllReduceDtype; bool mUseMambaConv1dPlugin; bool mInputPacked; bool mPagedState; @@ -998,7 +998,7 @@ class ModelConfig SpeculativeDecodingMode mSpeculativeDecodingMode; // Logits datatype - nvinfer1::DataType mLogitsDtype; + tensorrt_llm::DataType mLogitsDtype; bool mUseShapeInference; ManageWeightsType mManageWeightsType; std::string mModelName; diff --git a/cpp/include/tensorrt_llm/runtime/rawEngine.h b/cpp/include/tensorrt_llm/runtime/rawEngine.h deleted file mode 100644 index b219cbe03382..000000000000 --- a/cpp/include/tensorrt_llm/runtime/rawEngine.h +++ /dev/null @@ -1,123 +0,0 @@ -/* - * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. - * - * 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. - */ - -#pragma once - -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/executor/tensor.h" - -#include -#include -#include -#include - -namespace tensorrt_llm::runtime -{ - -class RawEngine -{ -public: - enum Type - { - FilePath, - AddressWithSize, - HostMemory - }; - - explicit RawEngine(std::filesystem::path enginePath) noexcept - : mType(FilePath) - , mEnginePath(std::move(enginePath)) - { - } - - explicit RawEngine(void const* engineAddr, std::size_t engineSize) noexcept - : mType(AddressWithSize) - , mEngineAddr(engineAddr) - , mEngineSize(engineSize) - { - } - - explicit RawEngine(nvinfer1::IHostMemory const* engineBuffer) noexcept - : mType(HostMemory) - , mEngineBuffer(engineBuffer) - { - } - - [[nodiscard]] Type getType() const - { - return mType; - } - - [[nodiscard]] std::filesystem::path getPath() const - { - TLLM_CHECK(mEnginePath.has_value()); - return mEnginePath.value(); - } - - [[nodiscard]] std::optional getPathOpt() const - { - return mEnginePath; - } - - void setPath(std::filesystem::path enginePath) - { - mEnginePath = std::move(enginePath); - } - - [[nodiscard]] std::optional> const& - getManagedWeightsMapOpt() const - { - return mManagedWeightsMap; - } - - void setManagedWeightsMap(std::map managedWeightsMap) - { - mManagedWeightsMap = std::move(managedWeightsMap); - } - - [[nodiscard]] void const* getAddress() const - { - TLLM_CHECK(mType == AddressWithSize); - return mEngineAddr; - } - - [[nodiscard]] std::size_t getSize() const - { - TLLM_CHECK(mType == AddressWithSize); - return mEngineSize; - } - - [[nodiscard]] nvinfer1::IHostMemory const* getHostMemory() const - { - TLLM_CHECK(mType == HostMemory); - return mEngineBuffer; - } - -private: - Type mType; - std::optional mEnginePath; - - struct - { - void const* mEngineAddr{}; - std::size_t mEngineSize{}; - }; - - nvinfer1::IHostMemory const* mEngineBuffer{}; - std::optional> mManagedWeightsMap; -}; - -} // namespace tensorrt_llm::runtime diff --git a/cpp/include/tensorrt_llm/runtime/tllmLogger.h b/cpp/include/tensorrt_llm/runtime/tllmLogger.h deleted file mode 100644 index dd3806ec5242..000000000000 --- a/cpp/include/tensorrt_llm/runtime/tllmLogger.h +++ /dev/null @@ -1,34 +0,0 @@ -/* - * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. - * - * 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. - */ - -#pragma once - -#include - -namespace tensorrt_llm::runtime -{ - -class TllmLogger : public nvinfer1::ILogger -{ -public: - void log(Severity severity, nvinfer1::AsciiChar const* msg) noexcept override; - - Severity getLevel(); - - void setLevel(Severity level); -}; - -} // namespace tensorrt_llm::runtime diff --git a/cpp/include/tensorrt_llm/runtime/utils/debugUtils.h b/cpp/include/tensorrt_llm/runtime/utils/debugUtils.h index 68064d74c7e2..3f6c307a3cde 100644 --- a/cpp/include/tensorrt_llm/runtime/utils/debugUtils.h +++ b/cpp/include/tensorrt_llm/runtime/utils/debugUtils.h @@ -15,6 +15,7 @@ */ #pragma once +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/runtime/bufferManager.h" namespace tensorrt_llm::runtime::utils @@ -24,7 +25,7 @@ template bool tensorHasInvalid(ITensor const& tensor, BufferManager const& manager, std::string const& infoStr); bool tensorHasInvalid( - size_t M, size_t K, nvinfer1::DataType type, void const* data, cudaStream_t stream, std::string const& infoStr); + size_t M, size_t K, tensorrt_llm::DataType type, void const* data, cudaStream_t stream, std::string const& infoStr); int stallStream( char const* name, std::optional stream = std::nullopt, std::optional delay = std::nullopt); diff --git a/cpp/include/tensorrt_llm/runtime/worldConfig.h b/cpp/include/tensorrt_llm/runtime/worldConfig.h index 9ff2d0970df7..272b0fec5ada 100644 --- a/cpp/include/tensorrt_llm/runtime/worldConfig.h +++ b/cpp/include/tensorrt_llm/runtime/worldConfig.h @@ -18,7 +18,6 @@ #include "tensorrt_llm/runtime/common.h" -#include #include #include diff --git a/cpp/micro_benchmarks/mixtureOfExpertsBackendBenchmarkFixture.h b/cpp/micro_benchmarks/mixtureOfExpertsBackendBenchmarkFixture.h index f466a65e871f..354e11184f8d 100644 --- a/cpp/micro_benchmarks/mixtureOfExpertsBackendBenchmarkFixture.h +++ b/cpp/micro_benchmarks/mixtureOfExpertsBackendBenchmarkFixture.h @@ -31,6 +31,7 @@ #include "tensorrt_llm/common/cudaUtils.h" #include "tensorrt_llm/common/memoryUtils.h" #include "tensorrt_llm/common/nvtxUtils.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/kernels/cutlass_kernels/cutlass_preprocessors.h" #include "tensorrt_llm/runtime/bufferManager.h" #include "tensorrt_llm/runtime/cudaStream.h" @@ -380,44 +381,44 @@ class MixtureOfExpertsBenchmark : public ::benchmark::Fixture int64_t mNumExpertsPerNode{}; int64_t mK{}; - constexpr static nvinfer1::DataType toDTypeID() + constexpr static tensorrt_llm::DataType toDTypeID() { if (FP8 || WFP4AFP8) - return nvinfer1::DataType::kFP8; + return tensorrt_llm::DataType::kFP8; if (NVFP4) - return nvinfer1::DataType::kFP4; + return tensorrt_llm::DataType::kFP4; if (INT_QUANT && INT4) - return nvinfer1::DataType::kINT4; + return tensorrt_llm::DataType::kINT4; if (INT_QUANT) - return nvinfer1::DataType::kINT8; + return tensorrt_llm::DataType::kINT8; if (std::is_same_v) - return nvinfer1::DataType::kFLOAT; + return tensorrt_llm::DataType::kFLOAT; if (std::is_same_v) - return nvinfer1::DataType::kHALF; + return tensorrt_llm::DataType::kHALF; #ifdef ENABLE_BF16 if (std::is_same_v) - return nvinfer1::DataType::kBF16; + return tensorrt_llm::DataType::kBF16; #endif TLLM_THROW("Unrecognised format"); }; - constexpr static nvinfer1::DataType toWTypeID() + constexpr static tensorrt_llm::DataType toWTypeID() { if (FP8) - return nvinfer1::DataType::kFP8; + return tensorrt_llm::DataType::kFP8; if (NVFP4 || WFP4AFP8) - return nvinfer1::DataType::kFP4; + return tensorrt_llm::DataType::kFP4; if (INT_QUANT && INT4) - return nvinfer1::DataType::kINT4; + return tensorrt_llm::DataType::kINT4; if (INT_QUANT) - return nvinfer1::DataType::kINT8; + return tensorrt_llm::DataType::kINT8; if (std::is_same_v) - return nvinfer1::DataType::kFLOAT; + return tensorrt_llm::DataType::kFLOAT; if (std::is_same_v) - return nvinfer1::DataType::kHALF; + return tensorrt_llm::DataType::kHALF; #ifdef ENABLE_BF16 if (std::is_same_v) - return nvinfer1::DataType::kBF16; + return tensorrt_llm::DataType::kBF16; #endif TLLM_THROW("Unrecognised format"); }; @@ -427,31 +428,31 @@ class MixtureOfExpertsBenchmark : public ::benchmark::Fixture { if constexpr (std::is_same_v) { - return nvinfer1::DataType::kFP8; + return tensorrt_llm::DataType::kFP8; } else if constexpr (std::is_same_v) { - return nvinfer1::DataType::kFP4; + return tensorrt_llm::DataType::kFP4; } else if constexpr (std::is_same_v) { - return nvinfer1::DataType::kINT8; + return tensorrt_llm::DataType::kINT8; } else if constexpr (std::is_same_v) { - return nvinfer1::DataType::kINT4; + return tensorrt_llm::DataType::kINT4; } else if constexpr (std::is_same_v) { - return nvinfer1::DataType::kBF16; + return tensorrt_llm::DataType::kBF16; } else if constexpr (std::is_same_v) { - return nvinfer1::DataType::kHALF; + return tensorrt_llm::DataType::kHALF; } else if constexpr (std::is_same_v) { - return nvinfer1::DataType::kFLOAT; + return tensorrt_llm::DataType::kFLOAT; } else { diff --git a/cpp/tensorrt_llm/CMakeLists.txt b/cpp/tensorrt_llm/CMakeLists.txt index afd2d3a1f415..5f5e37836a05 100644 --- a/cpp/tensorrt_llm/CMakeLists.txt +++ b/cpp/tensorrt_llm/CMakeLists.txt @@ -144,8 +144,6 @@ add_subdirectory(common) add_subdirectory(kernels) add_subdirectory(layers) add_subdirectory(runtime) -add_subdirectory(testing) -add_subdirectory(executor_worker) set(BATCH_MANAGER_TARGET tensorrt_llm_batch_manager_static) set(BATCH_MANAGER_TARGET_ARCH ${TARGET_ARCH}) @@ -176,7 +174,6 @@ set(TRTLLM_LINK_LIBS ${CUBLASLT_LIB} ${CURAND_LIB} ${CMAKE_DL_LIBS} - ${TRT_LIB} common_src kernels_src flash_mla_src @@ -199,7 +196,6 @@ set(TRTLLM_LINK_LIBS cute_dsl_src layers_src runtime_src - testing_src compressorKernels_src mhcKernels_src userbuffers_src @@ -310,5 +306,3 @@ endif() if(BUILD_FLASH_MLA) add_subdirectory(flash_mla) endif() - -add_subdirectory(plugins) diff --git a/cpp/tensorrt_llm/batch_manager/CMakeLists.txt b/cpp/tensorrt_llm/batch_manager/CMakeLists.txt index 88e2484cc6f1..d893d12fa3e3 100644 --- a/cpp/tensorrt_llm/batch_manager/CMakeLists.txt +++ b/cpp/tensorrt_llm/batch_manager/CMakeLists.txt @@ -32,34 +32,20 @@ set(SRCS contextProgress.cpp dataTransceiver.cpp decoderBuffers.cpp - encoderBuffers.cpp - guidedDecoder.cpp - handleContextLogits.cpp - handleGenerationLogits.cpp kvCacheManager.cpp kvCacheEventManager.cpp kvCacheTransferManager.cpp kvCacheManagerV2Utils.cpp kvCacheManagerV2Utils.cu llmRequest.cpp - logitsPostProcessor.cpp - loraBuffers.cpp - makeDecodingBatchInputOutput.cpp medusaBuffers.cpp microBatchScheduler.cpp pauseRequests.cpp peftCacheManager.cpp - promptTuningBuffers.cpp - rnnStateBuffers.cpp rnnStateManager.cpp rnnCacheFormatter.cpp rnnCacheTransBuffer.cpp - runtimeBuffers.cpp sequenceSlotManager.cpp - transformerBuffers.cpp - trtEncoderModel.cpp - trtGptModelInflightBatching.cpp - updateDecoderBuffers.cpp utils/debugUtils.cpp utils/inflightBatchingUtils.cpp utils/logitsThread.cpp diff --git a/cpp/tensorrt_llm/batch_manager/baseTransBuffer.cpp b/cpp/tensorrt_llm/batch_manager/baseTransBuffer.cpp index cbb8159bce44..1fb522afd7be 100644 --- a/cpp/tensorrt_llm/batch_manager/baseTransBuffer.cpp +++ b/cpp/tensorrt_llm/batch_manager/baseTransBuffer.cpp @@ -20,6 +20,7 @@ #include "tensorrt_llm/common/envUtils.h" #include "tensorrt_llm/common/logger.h" #include "tensorrt_llm/common/opUtils.h" +#include "tensorrt_llm/common/tllmDataType.h" #include #include @@ -98,7 +99,7 @@ void BufferIndexHolder::poison() noexcept } BaseTransBufferManager::BaseTransBufferManager( - size_t transferBufferSize, nvinfer1::DataType dataType, std::optional maxNumTokens) + size_t transferBufferSize, tensorrt_llm::DataType dataType, std::optional maxNumTokens) : mDataType{dataType} , mBufferManager{std::make_shared()} , mMaxNumTokens{maxNumTokens} diff --git a/cpp/tensorrt_llm/batch_manager/baseTransBuffer.h b/cpp/tensorrt_llm/batch_manager/baseTransBuffer.h index 8c1fd313f5a6..88585818d4e5 100644 --- a/cpp/tensorrt_llm/batch_manager/baseTransBuffer.h +++ b/cpp/tensorrt_llm/batch_manager/baseTransBuffer.h @@ -17,6 +17,7 @@ #pragma once +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/runtime/bufferManager.h" #include "tensorrt_llm/runtime/iTensor.h" @@ -225,7 +226,7 @@ class BaseTransBufferManager /// @param dataType Data type for the buffers. /// @param maxNumTokens Optional max tokens for sizing. BaseTransBufferManager( - size_t transferBufferSize, nvinfer1::DataType dataType, std::optional maxNumTokens = std::nullopt); + size_t transferBufferSize, tensorrt_llm::DataType dataType, std::optional maxNumTokens = std::nullopt); struct ConcurrenceResource { @@ -256,7 +257,7 @@ class BaseTransBufferManager bool mOnlyUseDynamicBuffer; bool mUseFabricMemory; size_t mNumberOfElements; - nvinfer1::DataType mDataType; + tensorrt_llm::DataType mDataType; ConcurrenceResource mConcurrenceSendResource; ConcurrenceResource mConcurrenceRecvResource; runtime::BufferManager mBufferManager; diff --git a/cpp/tensorrt_llm/batch_manager/cacheFormatter.h b/cpp/tensorrt_llm/batch_manager/cacheFormatter.h index 458cac8d4382..92d1b9d58a27 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheFormatter.h +++ b/cpp/tensorrt_llm/batch_manager/cacheFormatter.h @@ -28,7 +28,6 @@ #include "tensorrt_llm/executor/dataTransceiverState.h" #include "tensorrt_llm/runtime/bufferManager.h" #include "tensorrt_llm/runtime/utils/mpiUtils.h" -#include #include #include #include diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransBuffer.cpp b/cpp/tensorrt_llm/batch_manager/cacheTransBuffer.cpp index 772c9555f0f3..e06198f9ea60 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheTransBuffer.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheTransBuffer.cpp @@ -21,7 +21,7 @@ #include "tensorrt_llm/common/opUtils.h" #include "tensorrt_llm/executor/executor.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include namespace tensorrt_llm::batch_manager::kv_cache_manager @@ -194,7 +194,7 @@ bool FabricMemory::supportFabricMemory() size_t CacheTransBufferManager::computeTransferBufferSize( KVCacheManager::BaseKVCacheManager* cacheManager, std::optional maxNumTokens, bool transferIndexerKCache) { - nvinfer1::DataType dataType; + tensorrt_llm::DataType dataType; if (transferIndexerKCache) { dataType = cacheManager->getIndexerKCachePool()->getDataType(); diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp index aa05cd033892..f730fb2aaf1e 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp @@ -47,6 +47,7 @@ #include "tensorrt_llm/batch_manager/rnnStateManager.h" #include "tensorrt_llm/common/envUtils.h" #include "tensorrt_llm/common/logger.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/executor/cache_transmission/mpi_utils/connection.h" #include "tensorrt_llm/executor/dataTransceiverState.h" #include "tensorrt_llm/executor/serializeUtils.h" @@ -344,7 +345,7 @@ std::unique_ptr CacheTransceiverFactory::createCacheTransc CacheTransceiver::CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheManager, executor::kv_cache::CacheState::ModelConfig const& cacheStateModelCfg, runtime::WorldConfig const& worldConfig, - std::vector const& attentionLayerNumPerPP, nvinfer1::DataType dataType, + std::vector const& attentionLayerNumPerPP, tensorrt_llm::DataType dataType, executor::kv_cache::CacheState::AttentionType attentionType, std::optional cacheTransceiverConfig, std::vector const& rnnLayerNumPerPP) @@ -457,20 +458,20 @@ CacheTransceiver::CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheMa // Pool dtype is UINT8 (raw byte storage), so we cannot use pool->getDataType(). // Only the byte size matters for split/concat kernel stride calculations — the actual // dtype enum is not interpreted numerically, just used for getDTypeSize() dispatch. - auto dtypeFromSize = [](SizeType32 size) -> nvinfer1::DataType + auto dtypeFromSize = [](SizeType32 size) -> tensorrt_llm::DataType { switch (size) { - case 4: return nvinfer1::DataType::kFLOAT; - case 2: return nvinfer1::DataType::kBF16; - case 1: return nvinfer1::DataType::kFP8; + case 4: return tensorrt_llm::DataType::kFLOAT; + case 2: return tensorrt_llm::DataType::kBF16; + case 1: return tensorrt_llm::DataType::kFP8; default: TLLM_THROW("Unsupported RNN state dtype size: %d", size); } }; TLLM_CHECK_WITH_INFO(linearMeta->rnnSsmDtypeSize > 0, "rnnSsmDtypeSize not set in LinearAttentionMetadata"); TLLM_CHECK_WITH_INFO(linearMeta->rnnConvDtypeSize > 0, "rnnConvDtypeSize not set in LinearAttentionMetadata"); - nvinfer1::DataType ssmDtype = dtypeFromSize(linearMeta->rnnSsmDtypeSize); - nvinfer1::DataType convDtype = dtypeFromSize(linearMeta->rnnConvDtypeSize); + tensorrt_llm::DataType ssmDtype = dtypeFromSize(linearMeta->rnnSsmDtypeSize); + tensorrt_llm::DataType convDtype = dtypeFromSize(linearMeta->rnnConvDtypeSize); mCacheState->setRnnConfig(rnnModelCfg, rnnLayerNumPerPP, convDtype, ssmDtype); // Create RnnCacheTransBufferManager for unified pool path. diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransferLayer.cpp b/cpp/tensorrt_llm/batch_manager/cacheTransferLayer.cpp index d0a54dbb7d3c..b2e0a21537d1 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheTransferLayer.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheTransferLayer.cpp @@ -21,6 +21,7 @@ #include "tensorrt_llm/batch_manager/rnnCacheFormatter.h" #include "tensorrt_llm/common/assert.h" #include "tensorrt_llm/common/logger.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/executor/cache_transmission/agent_utils/connection.h" #include "tensorrt_llm/executor/cache_transmission/cacheSplitConcat.h" @@ -121,7 +122,8 @@ void CacheTransferLayer::unformat(TransferSession& session) const } void CacheTransferLayer::setRnnConfig(executor::kv_cache::CacheState::RnnModelConfig rnnModelConfig, - std::vector rnnLayerNumPerPP, nvinfer1::DataType convStateDataType, nvinfer1::DataType ssmStateDataType) + std::vector rnnLayerNumPerPP, tensorrt_llm::DataType convStateDataType, + tensorrt_llm::DataType ssmStateDataType) { mCacheState.setRnnConfig( std::move(rnnModelConfig), std::move(rnnLayerNumPerPP), convStateDataType, ssmStateDataType); diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransferLayer.h b/cpp/tensorrt_llm/batch_manager/cacheTransferLayer.h index 0506e98197e6..48a171a65c9c 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheTransferLayer.h +++ b/cpp/tensorrt_llm/batch_manager/cacheTransferLayer.h @@ -18,6 +18,7 @@ #pragma once #include "tensorrt_llm/batch_manager/rnnCacheFormatter.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/executor/dataTransceiverState.h" #include "tensorrt_llm/runtime/common.h" @@ -79,8 +80,8 @@ class CacheTransferLayer /// @brief Update the RNN config on the internal CacheState. /// Used by CppMambaHybridCacheManager path where RNN config is set after construction. void setRnnConfig(executor::kv_cache::CacheState::RnnModelConfig rnnModelConfig, - std::vector rnnLayerNumPerPP, nvinfer1::DataType convStateDataType, - nvinfer1::DataType ssmStateDataType); + std::vector rnnLayerNumPerPP, tensorrt_llm::DataType convStateDataType, + tensorrt_llm::DataType ssmStateDataType); [[nodiscard]] kv_cache_manager::BaseKVCacheManager* getCacheManager() const noexcept; diff --git a/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp b/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp index 5c5d3e11a01c..04c3760be5c7 100644 --- a/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp +++ b/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp @@ -33,7 +33,7 @@ #include "tensorrt_llm/runtime/utils/mpiUtils.h" #include "tensorrt_llm/runtime/utils/speculativeChoicesUtils.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" using namespace tensorrt_llm::runtime; @@ -93,7 +93,7 @@ void copySequenceLengths(RequestVector const& contextRequests, DecoderInputBuffe /// @brief Retrieve the embedding bias from the request. This potentially makes a copy of the tensor /// to the appropriate type if the input tensor does not match it. -[[nodiscard]] TensorPtr getEmbeddingBias(nvinfer1::DataType logitsType, TensorPtr const& tensor) +[[nodiscard]] TensorPtr getEmbeddingBias(tensorrt_llm::DataType logitsType, TensorPtr const& tensor) { // Check that embedding bias type is same as logits type. If so, we can return the tensor right away if (tensor->getDataType() == logitsType) @@ -102,7 +102,7 @@ void copySequenceLengths(RequestVector const& contextRequests, DecoderInputBuffe } // Support FP32 input for FP16 embedding bias (in the case of FP8 models) - if (tensor->getDataType() == nvinfer1::DataType::kFLOAT && logitsType == nvinfer1::DataType::kHALF) + if (tensor->getDataType() == tensorrt_llm::DataType::kFLOAT && logitsType == tensorrt_llm::DataType::kHALF) { // Do a deep copy of the tensor to the expected type TLLM_LOG_WARNING( @@ -133,10 +133,10 @@ void copySequenceLengths(RequestVector const& contextRequests, DecoderInputBuffe std::tuple, std::vector, std::vector> CreateNewDecoderRequests::operator()(runtime::ModelConfig const& modelConfig, runtime::WorldConfig const& worldConfig, - executor::DecodingConfig const& decodingConfig, RequestVector const& contextRequests, nvinfer1::DataType logitsType, - DecoderInputBuffers& inputBuffers, runtime::decoder::DecoderState& decoderState, CudaStream const& runtimeStream, - CudaStream const& decoderStream, SizeType32 maxSequenceLength, SizeType32 beamWidth, - OptionalRef medusaBuffers) const + executor::DecodingConfig const& decodingConfig, RequestVector const& contextRequests, + tensorrt_llm::DataType logitsType, DecoderInputBuffers& inputBuffers, runtime::decoder::DecoderState& decoderState, + CudaStream const& runtimeStream, CudaStream const& decoderStream, SizeType32 maxSequenceLength, + SizeType32 beamWidth, OptionalRef medusaBuffers) const { TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); NVTX3_SCOPED_RANGE(CreateNewDecoderRequests); @@ -235,7 +235,7 @@ void initializeBeamSearch(DecodingInput& dJointInput, DecodingOutput& dJointOutp } void initializeEmbeddingBias(DecodingInput& dJointInput, SizeType32 batchSlot, - std::optional const& embeddingBias, nvinfer1::DataType logitsType, + std::optional const& embeddingBias, tensorrt_llm::DataType logitsType, runtime::ModelConfig const& modelConfig, BufferManager const& manager) { TensorPtr const embeddingBiasSlice = ITensor::slice(constPointerCast(dJointInput.embeddingBias), batchSlot, 1); @@ -631,7 +631,7 @@ void newRequestSpeculativeDecoding(DecodingInput& jointDecodingInput, DecodingOu std::tuple, std::vector> CreateNewDecoderRequests::createDecoderRequests(RequestVector const& finishedContextRequests, TensorPtr const& inputIds, executor::DecodingConfig const& decodingConfig, runtime::decoder::DecoderState& decoderState, - nvinfer1::DataType logitsType, runtime::ModelConfig const& modelConfig, runtime::WorldConfig const& worldConfig, + tensorrt_llm::DataType logitsType, runtime::ModelConfig const& modelConfig, runtime::WorldConfig const& worldConfig, runtime::CudaStream const& runtimeStream, runtime::CudaStream const& decoderStream, SizeType32 maxSequenceLength, OptionalRef medusaBuffers) const { diff --git a/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp b/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp index 109417965a75..851f0abb04ef 100644 --- a/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp +++ b/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp @@ -20,9 +20,9 @@ #include "tensorrt_llm/batch_manager/cacheFormatter.h" #include "tensorrt_llm/batch_manager/common.h" #include "tensorrt_llm/batch_manager/kvCacheUtils.h" -#include "tensorrt_llm/batch_manager/runtimeBuffers.h" #include "tensorrt_llm/common/envUtils.h" #include "tensorrt_llm/common/logger.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/common/tllmException.h" #include "tensorrt_llm/common/utils.h" #include "tensorrt_llm/executor/cache_transmission/agent_utils/connection.h" @@ -950,8 +950,8 @@ class CacheSender::Impl public: void setRnnConfig(executor::kv_cache::CacheState::RnnModelConfig rnnModelConfig, - std::vector rnnLayerNumPerPP, nvinfer1::DataType convStateDataType, - nvinfer1::DataType ssmStateDataType) + std::vector rnnLayerNumPerPP, tensorrt_llm::DataType convStateDataType, + tensorrt_llm::DataType ssmStateDataType) { mCacheTransferLayer.setRnnConfig(rnnModelConfig, rnnLayerNumPerPP, convStateDataType, ssmStateDataType); mSelfState.setCacheState(mCacheTransferLayer.getCacheState()); @@ -1697,8 +1697,8 @@ class CacheReceiver::Impl public: void setRnnConfig(executor::kv_cache::CacheState::RnnModelConfig rnnModelConfig, - std::vector rnnLayerNumPerPP, nvinfer1::DataType convStateDataType, - nvinfer1::DataType ssmStateDataType) + std::vector rnnLayerNumPerPP, tensorrt_llm::DataType convStateDataType, + tensorrt_llm::DataType ssmStateDataType) { mCacheTransferLayer.setRnnConfig(rnnModelConfig, rnnLayerNumPerPP, convStateDataType, ssmStateDataType); mSelfState.setCacheState(mCacheTransferLayer.getCacheState()); @@ -1778,7 +1778,8 @@ void CacheSender::sendReadySignal(LlmRequest::RequestIdType requestId, bool isRe } void CacheSender::setRnnConfig(executor::kv_cache::CacheState::RnnModelConfig rnnModelConfig, - std::vector rnnLayerNumPerPP, nvinfer1::DataType convStateDataType, nvinfer1::DataType ssmStateDataType) + std::vector rnnLayerNumPerPP, tensorrt_llm::DataType convStateDataType, + tensorrt_llm::DataType ssmStateDataType) { mImpl->setRnnConfig(std::move(rnnModelConfig), std::move(rnnLayerNumPerPP), convStateDataType, ssmStateDataType); } @@ -1817,7 +1818,8 @@ bool CacheReceiver::receiveReadySignal(TransferSession& session) } void CacheReceiver::setRnnConfig(executor::kv_cache::CacheState::RnnModelConfig rnnModelConfig, - std::vector rnnLayerNumPerPP, nvinfer1::DataType convStateDataType, nvinfer1::DataType ssmStateDataType) + std::vector rnnLayerNumPerPP, tensorrt_llm::DataType convStateDataType, + tensorrt_llm::DataType ssmStateDataType) { mImpl->setRnnConfig(std::move(rnnModelConfig), std::move(rnnLayerNumPerPP), convStateDataType, ssmStateDataType); } diff --git a/cpp/tensorrt_llm/batch_manager/dataTransceiver.h b/cpp/tensorrt_llm/batch_manager/dataTransceiver.h index 8e84a71556af..c1da646916dc 100644 --- a/cpp/tensorrt_llm/batch_manager/dataTransceiver.h +++ b/cpp/tensorrt_llm/batch_manager/dataTransceiver.h @@ -29,6 +29,7 @@ #include "tensorrt_llm/common/assert.h" #include "tensorrt_llm/common/envUtils.h" #include "tensorrt_llm/common/logger.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/executor/cacheCommunicator.h" #include "tensorrt_llm/executor/dataTransceiverState.h" #include "tensorrt_llm/executor/serializeUtils.h" @@ -306,8 +307,8 @@ class CacheSender /// @brief Update the RNN config on the internal CacheState copies. /// Used by CppMambaHybridCacheManager path where RNN config is set after construction. void setRnnConfig(executor::kv_cache::CacheState::RnnModelConfig rnnModelConfig, - std::vector rnnLayerNumPerPP, nvinfer1::DataType convStateDataType, - nvinfer1::DataType ssmStateDataType); + std::vector rnnLayerNumPerPP, tensorrt_llm::DataType convStateDataType, + tensorrt_llm::DataType ssmStateDataType); /// @brief Destructor. virtual ~CacheSender(); @@ -357,8 +358,8 @@ class CacheReceiver /// @brief Update the RNN config on the internal CacheState copies. /// Used by CppMambaHybridCacheManager path where RNN config is set after construction. void setRnnConfig(executor::kv_cache::CacheState::RnnModelConfig rnnModelConfig, - std::vector rnnLayerNumPerPP, nvinfer1::DataType convStateDataType, - nvinfer1::DataType ssmStateDataType); + std::vector rnnLayerNumPerPP, tensorrt_llm::DataType convStateDataType, + tensorrt_llm::DataType ssmStateDataType); /// @brief Destructor. virtual ~CacheReceiver(); diff --git a/cpp/tensorrt_llm/batch_manager/decoderBuffers.cpp b/cpp/tensorrt_llm/batch_manager/decoderBuffers.cpp index fd67bb55e89d..fecc0851d361 100644 --- a/cpp/tensorrt_llm/batch_manager/decoderBuffers.cpp +++ b/cpp/tensorrt_llm/batch_manager/decoderBuffers.cpp @@ -18,6 +18,7 @@ #include "tensorrt_llm/batch_manager/decoderBuffers.h" #include "tensorrt_llm/common/logger.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/runtime/bufferManager.h" #include "tensorrt_llm/runtime/common.h" #include "tensorrt_llm/runtime/decoderState.h" @@ -70,21 +71,21 @@ DecoderOutputBuffers::DecoderOutputBuffers(SizeType32 maxNumSequences, SizeType3 auto constexpr TRTTokenIdType = runtime::TRTDataType::value; sequenceLengthsHost - = BufferManager::pinned(ITensor::makeShape({maxNumSequences, maxBeamWidth}), nvinfer1::DataType::kINT32); + = BufferManager::pinned(ITensor::makeShape({maxNumSequences, maxBeamWidth}), tensorrt_llm::DataType::kINT32); - finishedSumHost = BufferManager::pinned(ITensor::makeShape({maxNumSequences}), nvinfer1::DataType::kINT32); + finishedSumHost = BufferManager::pinned(ITensor::makeShape({maxNumSequences}), tensorrt_llm::DataType::kINT32); newOutputTokensHost = BufferManager::pinned(ITensor::makeShape({maxTokensPerStep, maxNumSequences, maxBeamWidth}), TRTTokenIdType); cumLogProbsHost - = BufferManager::pinned(ITensor::makeShape({maxNumSequences, maxBeamWidth}), nvinfer1::DataType::kFLOAT); + = BufferManager::pinned(ITensor::makeShape({maxNumSequences, maxBeamWidth}), tensorrt_llm::DataType::kFLOAT); logProbsHost = BufferManager::pinned( - ITensor::makeShape({maxNumSequences, maxBeamWidth, maxSeqLen}), nvinfer1::DataType::kFLOAT); + ITensor::makeShape({maxNumSequences, maxBeamWidth, maxSeqLen}), tensorrt_llm::DataType::kFLOAT); finishReasonsHost - = BufferManager::pinned(ITensor::makeShape({maxNumSequences, maxBeamWidth}), nvinfer1::DataType::kUINT8); + = BufferManager::pinned(ITensor::makeShape({maxNumSequences, maxBeamWidth}), tensorrt_llm::DataType::kUINT8); } void DecoderOutputBuffers::enableLookaheadDecoding(SizeType32 maxNumSequences, SizeType32 maxTokensPerStep) @@ -115,9 +116,9 @@ void DecoderOutputBuffers::setupSpeculativeDecoding( if (speculativeDecodingMode.variableDraftLength()) { nextDraftTokensLengthsHost - = BufferManager::pinned(ITensor::makeShape({maxNumSequences}), nvinfer1::DataType::kINT32); + = BufferManager::pinned(ITensor::makeShape({maxNumSequences}), tensorrt_llm::DataType::kINT32); prevDraftTokensLengthsHost - = BufferManager::pinned(ITensor::makeShape({maxNumSequences}), nvinfer1::DataType::kINT32); + = BufferManager::pinned(ITensor::makeShape({maxNumSequences}), tensorrt_llm::DataType::kINT32); } } } @@ -307,17 +308,18 @@ DecoderSlotAsyncSend::~DecoderSlotAsyncSend() SlotDecoderBuffers::SlotDecoderBuffers(SizeType32 maxBeamWidth, SizeType32 maxSeqLen, BufferManager const& manager) { - outputIds = manager.gpu(ITensor::makeShape({maxBeamWidth, maxSeqLen}), nvinfer1::DataType::kINT32); - outputIdsHost = BufferManager::pinned(ITensor::makeShape({maxBeamWidth, maxSeqLen}), nvinfer1::DataType::kINT32); + outputIds = manager.gpu(ITensor::makeShape({maxBeamWidth, maxSeqLen}), tensorrt_llm::DataType::kINT32); + outputIdsHost + = BufferManager::pinned(ITensor::makeShape({maxBeamWidth, maxSeqLen}), tensorrt_llm::DataType::kINT32); - sequenceLengths = manager.gpu(ITensor::makeShape({maxBeamWidth}), nvinfer1::DataType::kINT32); - sequenceLengthsHost = BufferManager::pinned(ITensor::makeShape({maxBeamWidth}), nvinfer1::DataType::kINT32); + sequenceLengths = manager.gpu(ITensor::makeShape({maxBeamWidth}), tensorrt_llm::DataType::kINT32); + sequenceLengthsHost = BufferManager::pinned(ITensor::makeShape({maxBeamWidth}), tensorrt_llm::DataType::kINT32); - cumLogProbs = manager.gpu(ITensor::makeShape({maxBeamWidth}), nvinfer1::DataType::kFLOAT); - cumLogProbsHost = BufferManager::pinned(ITensor::makeShape({maxBeamWidth}), nvinfer1::DataType::kFLOAT); + cumLogProbs = manager.gpu(ITensor::makeShape({maxBeamWidth}), tensorrt_llm::DataType::kFLOAT); + cumLogProbsHost = BufferManager::pinned(ITensor::makeShape({maxBeamWidth}), tensorrt_llm::DataType::kFLOAT); - logProbs = manager.gpu(ITensor::makeShape({maxBeamWidth, maxSeqLen}), nvinfer1::DataType::kFLOAT); - logProbsHost = BufferManager::pinned(ITensor::makeShape({maxBeamWidth, maxSeqLen}), nvinfer1::DataType::kFLOAT); + logProbs = manager.gpu(ITensor::makeShape({maxBeamWidth, maxSeqLen}), tensorrt_llm::DataType::kFLOAT); + logProbsHost = BufferManager::pinned(ITensor::makeShape({maxBeamWidth, maxSeqLen}), tensorrt_llm::DataType::kFLOAT); } } // namespace tensorrt_llm::batch_manager diff --git a/cpp/tensorrt_llm/batch_manager/encoderBuffers.cpp b/cpp/tensorrt_llm/batch_manager/encoderBuffers.cpp deleted file mode 100644 index 56fd393c68d7..000000000000 --- a/cpp/tensorrt_llm/batch_manager/encoderBuffers.cpp +++ /dev/null @@ -1,560 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2025 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. - */ - -#include "encoderBuffers.h" - -#include "tensorrt_llm/batch_manager/llmRequest.h" -#include "tensorrt_llm/common/nvtxUtils.h" -#include "tensorrt_llm/runtime/bufferManager.h" -#include "tensorrt_llm/runtime/common.h" -#include "tensorrt_llm/runtime/iBuffer.h" -#include "tensorrt_llm/runtime/iTensor.h" - -#include - -using namespace tensorrt_llm::runtime; - -namespace tensorrt_llm::batch_manager -{ - -EncoderBuffers::EncoderBuffers( - SizeType32 maxBatchSize, ModelConfig const& modelConfig, WorldConfig const& worldConfig, TllmRuntime const& runtime) -{ - // init empty buffers on cpu/gpu/pinned - init(maxBatchSize, modelConfig, worldConfig, runtime); - - // pre-allocate based on max buffer sizes - // Note: pre-allocation can be done directly instead of empty-->reshape, but it is ok extract the common reshape() - // utility because the buffer shapes can be dynamically set during runtime as well - initBufferSizes(maxBatchSize, modelConfig, worldConfig, runtime); -} - -void EncoderBuffers::init( - SizeType32 maxBatchSize, ModelConfig const& modelConfig, WorldConfig const& worldConfig, TllmRuntime const& runtime) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - auto const& manager = runtime.getBufferManager(); - - auto hiddenStatesType = modelConfig.getDataType(); - - inputFeatures = manager.emptyTensor(MemoryType::kGPU, hiddenStatesType); - inputIds = manager.emptyTensor(MemoryType::kGPU, nvinfer1::DataType::kINT32); - - // in PP, only rank 0 needs the following input fields - if (modelConfig.usePositionEmbedding() && worldConfig.isFirstPipelineParallelRank()) - { - positionIds = manager.emptyTensor(MemoryType::kGPU, nvinfer1::DataType::kINT32); - positionIdsReserved.resize(maxBatchSize * modelConfig.getMaxInputLen()); - std::iota(positionIdsReserved.begin(), positionIdsReserved.end(), 0); - } - if (modelConfig.useTokenTypeEmbedding() && worldConfig.isFirstPipelineParallelRank()) - { - tokenTypeIds = manager.emptyTensor(MemoryType::kGPU, nvinfer1::DataType::kINT32); - tokenTypeIdsReserved.resize(maxBatchSize * modelConfig.getMaxInputLen()); - std::fill(tokenTypeIdsReserved.begin(), tokenTypeIdsReserved.end(), 0); - } - - inputLengths = manager.emptyTensor(MemoryType::kGPU, nvinfer1::DataType::kINT32); - maxInputLength = manager.emptyTensor(MemoryType::kGPU, nvinfer1::DataType::kINT32); - - if (worldConfig.isPipelineParallel()) - { - hiddenStates = manager.emptyTensor(MemoryType::kGPU, hiddenStatesType); - } - if (worldConfig.isLastPipelineParallelRank()) - { - encoderOutput = manager.emptyTensor(MemoryType::kGPU, hiddenStatesType); - } - - if (modelConfig.useLanguageAdapter()) - { - languageAdapterRoutings = manager.emptyTensor(MemoryType::kGPU, TRTDataType::value); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void EncoderBuffers::initBufferSizes( - SizeType32 maxBatchSize, ModelConfig const& modelConfig, WorldConfig const& worldConfig, TllmRuntime const& runtime) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - // get buffer shape based on max values - numRequests = maxBatchSize; - encoderInputLen = maxBatchSize * modelConfig.getMaxInputLen(); - encoderOutputLen = maxBatchSize * modelConfig.getMaxInputLen(); // assume output length <= input length - maxInputLengthInBatch = modelConfig.getMaxInputLen(); - - // update buffer shapes - reshape(runtime, modelConfig, worldConfig); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void EncoderBuffers::updateBufferSizes(RequestVector const& requests, ModelConfig const& modelConfig, - WorldConfig const& worldConfig, TllmRuntime const& runtime) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - numRequests = requests.size(); - encoderInputLen = 0; - encoderOutputLen = 0; - maxInputLengthInBatch = 0; - - // get buffer shape based on actual batched requests - for (auto const& req : requests) - { - encoderInputLen += req->getEncoderInputLen(); - encoderOutputLen += req->getEncoderOutputLen(); - maxInputLengthInBatch - = std::max(maxInputLengthInBatch, req->getEncoderInputLen()); // Decoder input is encoder output - } - - // update buffer shapes - reshape(runtime, modelConfig, worldConfig); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void EncoderBuffers::reshape(TllmRuntime const& runtime, ModelConfig const& modelConfig, WorldConfig const& worldConfig) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - if (modelConfig.isMultiModal()) - { - return; // multimodal models do not need to set position id, etc. or any output tensors - } - - inputIds->reshape(ITensor::makeShape({encoderInputLen})); - if (positionIds) - { - if (modelConfig.isWhisper()) - { - positionIds->reshape(ITensor::makeShape({encoderOutputLen})); - } - else - { - positionIds->reshape(ITensor::makeShape({encoderInputLen})); - } - } - if (tokenTypeIds) - { - tokenTypeIds->reshape(ITensor::makeShape({encoderInputLen})); - } - - inputLengths->reshape(ITensor::makeShape({numRequests})); - maxInputLength->reshape(ITensor::makeShape({maxInputLengthInBatch})); - - if (worldConfig.isPipelineParallel()) - { - hiddenStates->reshape( - ITensor::makeShape({encoderOutputLen, modelConfig.getHiddenSize() * worldConfig.getTensorParallelism()})); - } - if (worldConfig.isLastPipelineParallelRank()) - { - encoderOutput->reshape( - ITensor::makeShape({encoderOutputLen, modelConfig.getHiddenSize() * worldConfig.getTensorParallelism()})); - } - if (modelConfig.useLanguageAdapter()) - { - languageAdapterRoutings->reshape(ITensor::makeShape({encoderInputLen, 1})); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void EncoderBuffers::setFromInputs(RequestVector const& requests, ModelConfig const& modelConfig, - WorldConfig const& worldConfig, TllmRuntime const& runtime) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(encoderBuffersSetFromInputs); - - if (!worldConfig.isFirstPipelineParallelRank()) - { - return; - } - - auto const& manager = runtime.getBufferManager(); - - std::vector inputIdsAll; - std::vector positionIdsAll; - std::vector tokenTypeIdsAll; - std::vector inputLengthsAll; - std::vector languageAdapterRoutingAll; - // use shape to indicates max input length, content is not important - // TODO: change to a scalar value for this from engine side - std::vector maxInputLengthAll(maxInputLengthInBatch, 0); - - if (requests.front()->getEncoderInputFeatures()) - { - if (modelConfig.isMultiModal()) - { - auto batchedInputShape = requests.front()->getEncoderInputFeatures()->getShape(); // [1, 3, H, W] - batchedInputShape.d[0] = encoderInputLen; // [batch_size, 3, H, W] - inputFeatures->reshape(batchedInputShape); - } - else - { - SizeType32 const featureDim = requests.front()->getEncoderInputFeatures()->getShape().d[1]; - TLLM_LOG_DEBUG("EncoderBuffers::setFromInputs - featureDim = %d", featureDim); - inputFeatures->reshape(ITensor::makeShape({encoderInputLen, featureDim})); - } - } - - SizeType32 offset = 0; - - for (auto const& llmReq : requests) - { - SizeType32 const inputLength = llmReq->getEncoderInputLen(); - SizeType32 const outputLength = llmReq->getEncoderOutputLen(); - if (llmReq->getEncoderInputFeatures()) - { - auto const& reqFeatures - = llmReq - ->getEncoderInputFeatures(); // whisper: [length, featureDim]; Vision: [batch_size, channel, W, H] - TLLM_LOG_DEBUG("EncoderBuffers::setFromInputs - request id = %d, input features length = %d", - llmReq->mRequestId, inputLength); - manager.copy(*reqFeatures, *ITensor::slice(inputFeatures, offset, inputLength)); - offset += inputLength; - } - else - { - auto const& reqTokens = *llmReq->getEncoderTokens().value(); - inputIdsAll.insert(inputIdsAll.end(), reqTokens.begin(), reqTokens.end()); - if (tokenTypeIds) - { - tokenTypeIdsAll.insert( - tokenTypeIdsAll.end(), tokenTypeIdsReserved.begin(), tokenTypeIdsReserved.begin() + inputLength); - } - } - if (positionIds) - { - SizeType32 const length = modelConfig.isWhisper() ? outputLength : inputLength; - positionIdsAll.insert( - positionIdsAll.end(), positionIdsReserved.begin(), positionIdsReserved.begin() + length); - } - if (modelConfig.useLanguageAdapter()) - { - auto const languageAdapterRouting - = llmReq->getLanguageAdapterRouting(modelConfig.getNumLanguages().value(), inputLength); - languageAdapterRoutingAll.insert( - languageAdapterRoutingAll.end(), std::begin(languageAdapterRouting), std::end(languageAdapterRouting)); - } - inputLengthsAll.push_back(inputLength); - } - - // copy inputs from host to device - { - NVTX3_SCOPED_RANGE(bufferCopies); - if (requests.front()->getEncoderTokens()) - { - manager.copy(inputIdsAll.data(), *inputIds); - if (tokenTypeIds) - { - manager.copy(tokenTypeIdsAll.data(), *tokenTypeIds); - } - manager.copy(maxInputLengthAll.data(), *maxInputLength); - } - if (positionIds) - { - manager.copy(positionIdsAll.data(), *positionIds); - } - manager.copy(inputLengthsAll.data(), *inputLengths); - if (modelConfig.useLanguageAdapter()) - { - manager.copy(languageAdapterRoutingAll.data(), *languageAdapterRoutings); - } - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void EncoderBuffers::fillIOMaps(ModelConfig const& modelConfig, WorldConfig const& worldConfig) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(runtimeBuffersFillIOMaps); - - inputMap.clear(); - outputMap.clear(); - - // inputs - if (modelConfig.isMultiModal()) - { - inputMap.insert_or_assign("input", inputFeatures); - } - else if (modelConfig.isWhisper()) - { - inputMap.insert_or_assign("input_features", inputFeatures); - inputMap.insert_or_assign("input_lengths", inputLengths); - inputMap.insert_or_assign("position_ids", positionIds); - } - else - { - if (worldConfig.isFirstPipelineParallelRank()) - { - inputMap.insert_or_assign("input_ids", inputIds); - if (positionIds) - { - inputMap.insert_or_assign("position_ids", positionIds); - } - if (tokenTypeIds) - { - inputMap.insert_or_assign("token_type_ids", tokenTypeIds); - } - } - else - { - inputMap.insert_or_assign("hidden_states_input", hiddenStates); - } - inputMap.insert_or_assign("input_lengths", inputLengths); - inputMap.insert_or_assign("max_input_length", maxInputLength); - if (modelConfig.useLanguageAdapter()) - { - inputMap.insert_or_assign("language_adapter_routings", languageAdapterRoutings); - } - } - - // outputs - if (worldConfig.isLastPipelineParallelRank()) - { - outputMap.insert_or_assign("encoder_output", encoderOutput); - } - else - { - outputMap.insert_or_assign("hidden_states_output", hiddenStates); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -std::pair EncoderBuffers::prepareIO( - RequestVector const& requests, ModelConfig const& modelConfig, WorldConfig const& worldConfig, - TllmRuntime const& runtime) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - updateBufferSizes(requests, modelConfig, worldConfig, runtime); - - setFromInputs(requests, modelConfig, worldConfig, runtime); - - fillIOMaps(modelConfig, worldConfig); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); - - return {inputMap, outputMap}; -} - -void EncoderBuffers::rearrangeOutputs(RequestVector const& requests, ModelConfig const& modelConfig, - WorldConfig const& worldConfig, TllmRuntime const& runtime) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(encoderBuffersRearrangeOutput); - - auto const& manager = runtime.getBufferManager(); - - SizeType32 offset = 0, size = 0; - - updateReqOutputShape(requests, runtime, worldConfig, modelConfig); - - for (auto const& req : requests) - { - // copy from internal buffer to request-owned external buffers - size = req->getEncoderOutputLen(); - TLLM_LOG_DEBUG("EncoderBuffers::rearrangeOutputs - req: %d, encoderOutput shape = (%d, %d)", req->mClientId, - req->getEncoderOutput()->getShape().d[0], req->getEncoderOutput()->getShape().d[1]); - TLLM_LOG_DEBUG("EncoderBuffers::rearrangeOutputs - req: %d, enc output size = %d", req->mClientId, size); - - if (worldConfig.isPipelineParallel()) - { - manager.copy(*ITensor::slice(hiddenStates, offset, size), *req->getEncoderHiddenStates()); - } - if (worldConfig.isLastPipelineParallelRank()) - { - if (modelConfig.isMultiModal()) - { - manager.copy( - *ITensor::slice(encoderOutput, offset, size), *(req->getPromptEmbeddingTableMutable().value())); - } - else - { - manager.copy(*ITensor::slice(encoderOutput, offset, size), *req->getEncoderOutput()); - } - } - offset += size; - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void EncoderBuffers::updateReqOutputShape(RequestVector const& requests, TllmRuntime const& runtime, - WorldConfig const& worldConfig, ModelConfig const& modelConfig) -{ - auto const& manager = runtime.getBufferManager(); - - for (auto const& req : requests) - { - if (modelConfig.isMultiModal()) - { - auto shape = encoderOutput->getShape(); // [batch_size, prompt_vocab_size, feature_dim] - shape.d[0] = req->getEncoderOutputLen(); - req->getPromptEmbeddingTableMutable() = manager.emptyTensor(MemoryType::kGPU, encoderOutput->getDataType()); - req->getPromptEmbeddingTableMutable().value()->reshape(shape); - req->setPromptVocabSize(shape.d[1]); - // TODO: extra ids for kv cache reuse - } - else - { - auto encOutLen = req->getEncoderOutputLen(); - // update request-owned external buffer for each request - if (worldConfig.isPipelineParallel()) - { - req->getEncoderHiddenStates()->reshape( - ITensor::makeShape({encOutLen, modelConfig.getHiddenSize() * worldConfig.getTensorParallelism()})); - } - if (worldConfig.isLastPipelineParallelRank()) - { - req->getEncoderOutput()->reshape( - ITensor::makeShape({encOutLen, modelConfig.getHiddenSize() * worldConfig.getTensorParallelism()})); - } - } - } -} - -void EncoderBuffers::create(SizeType32 maxBatchSize, ModelConfig const& modelConfig, TllmRuntime const& runtime) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - auto const& manager = runtime.getBufferManager(); - - inputLengths = manager.emptyTensor(MemoryType::kGPU, nvinfer1::DataType::kINT32); - maxInputLength = manager.emptyTensor(MemoryType::kGPU, nvinfer1::DataType::kINT32); - - hiddenSize = modelConfig.getEncoderHiddenSize(); // full hidden size - // assume encoder & decoder use the same data type - encoderOutput = manager.emptyTensor(MemoryType::kGPU, modelConfig.getDataType()); - encoderOutputReserved = manager.gpu(ITensor::makeShape({1, hiddenSize}), modelConfig.getDataType()); - - crossKvCacheGen = manager.gpu(ITensor::makeShape({1}), nvinfer1::DataType::kBOOL); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void EncoderBuffers::setMaxBufferSizes(SizeType32 maxBatchSize, runtime::ModelConfig const& modelConfig) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - numRequests = maxBatchSize; - encoderInputLen = maxBatchSize * modelConfig.getMaxEncoderLen(); - encoderOutputLen = maxBatchSize * modelConfig.getMaxEncoderLen(); - maxInputLengthInBatch = modelConfig.getMaxEncoderLen(); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void EncoderBuffers::setBufferSizes(RequestVector const& contextRequests, RequestVector const& genRequests) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - numRequests = 0; /// total number of requests that need encoder information (context requests + - /// generation requests * beam width) - encoderInputLen = 0; - encoderOutputLen = 0; - maxInputLengthInBatch = 1; /// maximum encoder length in a batch - - for (auto const& llmReq : contextRequests) - { - numRequests += 1; - encoderInputLen += llmReq->getEncoderInputLen(); - encoderOutputLen += llmReq->getEncoderOutputLen(); - maxInputLengthInBatch = std::max(maxInputLengthInBatch, llmReq->getEncoderInputLen()); - } - - for (auto const& llmReq : genRequests) - { - auto const reqBeamWidth = llmReq->getBeamWidthByIter(); - numRequests += reqBeamWidth; // tile by beam width - maxInputLengthInBatch = std::max(maxInputLengthInBatch, llmReq->getEncoderInputLen()); - } - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void EncoderBuffers::reshape() -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - inputLengths->reshape(ITensor::makeShape({numRequests})); - maxInputLength->reshape(ITensor::makeShape({maxInputLengthInBatch})); - encoderOutput->reshape(ITensor::makeShape({encoderOutputLen, hiddenSize})); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void EncoderBuffers::fill( - RequestVector const& ctxRequests, RequestVector const& genRequests, runtime::BufferManager const& manager) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(encoderBufferCopies); - - std::vector inputLengthsAll; - std::vector maxInputLengthAll(maxInputLength->getShape().d[0], 0); - - SizeType32 offset = 0, size = 0; - for (auto const& requests : {ctxRequests, genRequests}) - { - for (auto const& llmReq : requests) - { - // 1. only ctx requests should gather the encoder output - // 2. only gen requests should tile encoder input lengths info by beam width - bool isCtx = llmReq->isContextInitState(); - if (isCtx) - { - size = llmReq->getEncoderOutputLen(); - auto const encoderOutputSlice = runtime::ITensor::slice(encoderOutput, offset, size); - manager.copy(*llmReq->getEncoderOutput(), *encoderOutputSlice); - offset += size; - - inputLengthsAll.emplace_back(size); - } - else - { - auto const reqBeamWidth = llmReq->getBeamWidthByIter(); - std::fill_n(std::back_inserter(inputLengthsAll), reqBeamWidth, - llmReq->getEncoderOutputLen()); // although encoder output is not needed, gen phase still needs the - // encoder length info for cross kv cache. Also tile by beam width - } - } - } - manager.copy(inputLengthsAll.data(), *inputLengths); - manager.copy(maxInputLengthAll.data(), *maxInputLength); - // crossKvCacheGen unused in engine for now, use default tensor - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void EncoderBuffers::insertInputTensors(TensorMap& inputMap) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - inputMap.insert_or_assign("encoder_output", encoderOutput); - inputMap.insert_or_assign("encoder_input_lengths", inputLengths); - inputMap.insert_or_assign("encoder_max_input_length", maxInputLength); - inputMap.insert_or_assign("cross_kv_cache_gen", crossKvCacheGen); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} -} // namespace tensorrt_llm::batch_manager diff --git a/cpp/tensorrt_llm/batch_manager/encoderBuffers.h b/cpp/tensorrt_llm/batch_manager/encoderBuffers.h deleted file mode 100644 index 64d416280f21..000000000000 --- a/cpp/tensorrt_llm/batch_manager/encoderBuffers.h +++ /dev/null @@ -1,140 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2025 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. - */ - -#pragma once - -#include "tensorrt_llm/batch_manager/common.h" -#include "tensorrt_llm/runtime/bufferManager.h" -#include "tensorrt_llm/runtime/iTensor.h" -#include "tensorrt_llm/runtime/modelConfig.h" -#include "tensorrt_llm/runtime/tllmRuntime.h" -#include "tensorrt_llm/runtime/worldConfig.h" - -namespace tensorrt_llm::batch_manager -{ - -class EncoderBuffers -{ -public: - using SizeType32 = tensorrt_llm::runtime::SizeType32; - using ITensor = tensorrt_llm::runtime::ITensor; - using TensorPtr = runtime::ITensor::SharedPtr; - using TensorMap = runtime::StringPtrMap; - using ModelConfig = runtime::ModelConfig; - using WorldConfig = runtime::WorldConfig; - using TllmRuntime = runtime::TllmRuntime; - - TensorPtr inputIds; - TensorPtr positionIds = nullptr; - TensorPtr tokenTypeIds = nullptr; - - TensorPtr inputLengths; // [numEncoderRequests] - TensorPtr maxInputLength; // [maxInputLengthInBatch] - - // intermediate states in pipeline parallelism - TensorPtr hiddenStates; // [numTokens, hiddenSize] - - // features for multimodal encoders (audio, image, etc.) - TensorPtr - inputFeatures; // [totalNumOfFeatures, featureDim] if remove_padding else [batchSize, featureDim, featureLength] - - // language adapter routing information for encoders if language adapter is presented. - TensorPtr languageAdapterRoutings; // [numTokens, numLanguages] - - // encoder output - TensorPtr encoderOutput; // [numEncoderTokens, hiddenSize] - - // output buffer owned by llmRequest, such that it's per-request output buffer - // encoderBuffers class can init and reshape each buffer, without maintaining a list/set of inflight buffers - // TODO in progress: to support BS>1 encoder, need (1) internal scratch space tensors to save the contiguous - // batched output (2) copy from CONTIGUOUS scratch tensor to individual request's DISCRETE output tensor after - // execution To standardize the implementation, for both BS=1 and BS>1, we use internal buffer to store BS=1/BS>1 - // results, and copy to request's external buffers. For BS=1, this introduces a redundancy copy, but ok for now. - - EncoderBuffers() = default; - EncoderBuffers(SizeType32 maxBatchSize, ModelConfig const& modelConfig, WorldConfig const& worldConfig, - TllmRuntime const& runtime); - - std::pair prepareIO(RequestVector const& requests, - ModelConfig const& modelConfig, WorldConfig const& worldConfig, TllmRuntime const& runtime); - - void rearrangeOutputs(RequestVector const& requests, ModelConfig const& modelConfig, WorldConfig const& worldConfig, - TllmRuntime const& runtime); - - //! @brief set shape of individual request's encoder output (Ptuning embedding table if multimodal) - void updateReqOutputShape(RequestVector const& requests, TllmRuntime const& runtime, WorldConfig const& worldConfig, - ModelConfig const& modelConfig); - -private: - SizeType32 numRequests{}; - SizeType32 encoderInputLen{}; - SizeType32 encoderOutputLen{}; - SizeType32 maxInputLengthInBatch{}; // max input length in a batch - - // prefilled with deterministic values to avoid runtime creation - std::vector positionIdsReserved; - std::vector tokenTypeIdsReserved; - - // engine I/O - TensorMap inputMap; - TensorMap outputMap; - - void init(SizeType32 maxBatchSize, ModelConfig const& modelConfig, WorldConfig const& worldConfig, - TllmRuntime const& runtime); - - //! @brief pre-allocate max buffer sizes during init - void initBufferSizes(SizeType32 maxBatchSize, ModelConfig const& modelConfig, WorldConfig const& worldConfig, - TllmRuntime const& runtime); - - //! @brief update actual buffer usage of requests during runtime - void updateBufferSizes(RequestVector const& requests, ModelConfig const& modelConfig, - WorldConfig const& worldConfig, TllmRuntime const& runtime); - - void reshape(TllmRuntime const& runtime, ModelConfig const& modelConfig, WorldConfig const& worldConfig); - - void setFromInputs(RequestVector const& requests, ModelConfig const& modelConfig, WorldConfig const& worldConfig, - TllmRuntime const& runtime); - - void fillIOMaps(ModelConfig const& modelConfig, WorldConfig const& worldConfig); - - // additional members that are Encoder-Decoder specific -private: - TensorPtr encoderOutputReserved; // [1, hiddenSize], dummy tensor for gen phase - TensorPtr crossKvCacheGen; // [1] - SizeType32 hiddenSize; // full hidden size (after multiplying tensor parallelism) - -public: - void create(SizeType32 maxBatchSize, ModelConfig const& modelConfig, TllmRuntime const& runtime); - - SizeType32 getMaxInputLengthInBatch() const - { - return maxInputLengthInBatch; - }; - - void setMaxBufferSizes(SizeType32 maxBatchSize, runtime::ModelConfig const& modelConfig); - - void setBufferSizes(RequestVector const& contextRequests, RequestVector const& genRequests); - - void reshape(); - - void fill( - RequestVector const& ctxRequests, RequestVector const& genRequests, runtime::BufferManager const& manager); - - void insertInputTensors(TensorMap& inputMap); -}; - -} // namespace tensorrt_llm::batch_manager diff --git a/cpp/tensorrt_llm/batch_manager/guidedDecoder.cpp b/cpp/tensorrt_llm/batch_manager/guidedDecoder.cpp deleted file mode 100644 index cb2264ec8003..000000000000 --- a/cpp/tensorrt_llm/batch_manager/guidedDecoder.cpp +++ /dev/null @@ -1,224 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2025 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. - */ - -#include "tensorrt_llm/batch_manager/guidedDecoder.h" -#include "tensorrt_llm/batch_manager/decoderBuffers.h" -#include "tensorrt_llm/batch_manager/llmRequest.h" -#include "tensorrt_llm/common/envUtils.h" -#include "tensorrt_llm/kernels/logitsBitmask.h" - -#include -#include - -using namespace tensorrt_llm::runtime; - -namespace tensorrt_llm::batch_manager -{ - -GuidedDecoder::GuidedDecoder(executor::GuidedDecodingConfig const& guidedDecodingConfig, SizeType32 maxNumSequences, - SizeType32 vocabSizePadded, nvinfer1::DataType logitsDtype, BufferManager const& runtimeBufferManager) - : mGuidedDecodingBackend{guidedDecodingConfig.getBackend()} - , mMaxNumSequences{maxNumSequences} - , mVocabSizePadded{vocabSizePadded} - , mBitmaskSize{common::ceilDiv(mVocabSizePadded, 32)} - , mLogitsDtype{logitsDtype} - , mCopyBufferManager{std::make_shared()} -{ - TLLM_CHECK_WITH_INFO(mGuidedDecodingBackend != executor::GuidedDecodingConfig::GuidedDecodingBackend::kLLGUIDANCE, - "LLGuidance is not supported for guided decoding in C++ runtime."); - if (mGuidedDecodingBackend == executor::GuidedDecodingConfig::GuidedDecodingBackend::kXGRAMMAR) - { - mXGrammarMatchers.resize(mMaxNumSequences); - xgrammar::VocabType vocabType = xgrammar::VocabType::RAW; - bool addPrefixSpace = false; - auto const& tokenizerStr = guidedDecodingConfig.getTokenizerStr(); - if (tokenizerStr) - { - auto const& metadata = xgrammar::TokenizerInfo::DetectMetadataFromHF(tokenizerStr.value()); - auto const& metadataJson = nlohmann::json::parse(metadata); - vocabType = metadataJson.at("vocab_type").template get(); - addPrefixSpace = metadataJson.at("add_prefix_space").template get(); - } - auto const& tokenizerInfo = xgrammar::TokenizerInfo(guidedDecodingConfig.getEncodedVocab().value(), vocabType, - mVocabSizePadded, guidedDecodingConfig.getStopTokenIds(), addPrefixSpace); - - auto const cacheLimitGb = common::getFloatEnv("XGRAMMAR_CACHE_LIMIT_GB"); - mXGrammarCompiler = std::make_shared(tokenizerInfo, /*max_threads=*/8, - /*cache_enabled=*/true, - /*cache_limit_bytes=*/static_cast(cacheLimitGb.value_or(1.0f) * 1024 * 1024 * 1024)); - - auto const logitsPtrDtype = BufferDataType{mLogitsDtype, false, true}; - auto constexpr bitmaskDtype = TRTDataType::value; - auto constexpr bitmaskPtrDtype = TRTDataType::value; - - mLogitsBitmask = runtimeBufferManager.gpu(ITensor::makeShape({mMaxNumSequences, mBitmaskSize}), bitmaskDtype); - mLogitsBitmaskHost = BufferManager::pinned(ITensor::makeShape({mMaxNumSequences, mBitmaskSize}), bitmaskDtype); - mLogitsBitmaskPtrVec = runtimeBufferManager.gpu(ITensor::makeShape({mMaxNumSequences}), bitmaskPtrDtype); - mLogitsBitmaskPtrVecHost = BufferManager::pinned(ITensor::makeShape({mMaxNumSequences}), bitmaskPtrDtype); - mLogitsPtrVec = runtimeBufferManager.gpu(ITensor::makeShape({mMaxNumSequences}), logitsPtrDtype); - mLogitsPtrVecHost = BufferManager::pinned(ITensor::makeShape({mMaxNumSequences}), logitsPtrDtype); - } -} - -void GuidedDecoder::build(ScheduledRequests const& scheduledRequests) -{ - if (mGuidedDecodingBackend == executor::GuidedDecodingConfig::GuidedDecodingBackend::kXGRAMMAR) - { - for (auto const& requests : {scheduledRequests.contextRequests, scheduledRequests.generationRequests}) - { - for (auto const& llmReq : requests) - { - auto const& guidedDecodingParams = llmReq->getGuidedDecodingParams(); - if (!guidedDecodingParams.has_value()) - { - continue; - } - auto const seqSlot = llmReq->mSeqSlot.value(); - if (llmReq->isContextInitState() && llmReq->isFirstContextChunk()) - { - // The request is in the first context forward step (considering kv cache reuse). - auto const& guideType = guidedDecodingParams->getGuideType(); - auto const& guide = guidedDecodingParams->getGuide(); - switch (guideType) - { - case executor::GuidedDecodingParams::GuideType::kJSON: - { - mXGrammarMatchers.at(seqSlot) = std::make_shared( - mXGrammarCompiler->CompileBuiltinJSONGrammar()); - break; - } - case executor::GuidedDecodingParams::GuideType::kJSON_SCHEMA: - { - mXGrammarMatchers.at(seqSlot) = std::make_shared( - mXGrammarCompiler->CompileJSONSchema(guide.value())); - break; - } - case executor::GuidedDecodingParams::GuideType::kREGEX: - { - mXGrammarMatchers.at(seqSlot) = std::make_shared( - mXGrammarCompiler->CompileRegex(guide.value())); - break; - } - case executor::GuidedDecodingParams::GuideType::kEBNF_GRAMMAR: - { - mXGrammarMatchers.at(seqSlot) = std::make_shared( - mXGrammarCompiler->CompileGrammar(guide.value())); - break; - } - case executor::GuidedDecodingParams::GuideType::kSTRUCTURAL_TAG: - { - mXGrammarMatchers.at(seqSlot) = std::make_shared( - mXGrammarCompiler->CompileStructuralTag(guide.value())); - break; - } - default: - { - TLLM_THROW("Unsupported guide type."); - } - } - } - else if (llmReq->isGenerationInProgressState()) - { - // The request is in a generation forward step. - // Currently, guided decoding does not support with beam search. - mXGrammarMatchers.at(seqSlot)->AcceptToken(llmReq->getLastTokens(0)); - } - else - { - continue; - } - - // Fill the bitmask on host and asynchorously copy to device using mCopyBufferManager. - auto const logitsBitmask = ITensor::at(mLogitsBitmask, {seqSlot}); - auto const logitsBitmaskHost = ITensor::at(mLogitsBitmaskHost, {seqSlot}); - - std::array bitmaskShape{mBitmaskSize}; - DLTensor logitsBitmaskDlt{logitsBitmaskHost->data(), DLDevice{kDLCPU, 0}, 1, DLDataType{kDLInt, 32, 1}, - bitmaskShape.data(), nullptr, 0}; - mXGrammarMatchers.at(seqSlot)->FillNextTokenBitmask(&logitsBitmaskDlt); - mCopyBufferManager.copy(*logitsBitmaskHost, *logitsBitmask); - } - } - } -} - -void GuidedDecoder::execute(DecoderInputBuffers const& decoderInputBuffers, BufferManager const& runtimeBufferManager) -{ - auto const& stream = runtimeBufferManager.getStream(); - - // Wait for mCopyBufferManager finishing the H2D copy of logitsBitmask - // TODO(enweiz): Move the H2D copy of logitsBitmaskPtrVec to buildGuidedDecoding. - // This may not bring too much perf gain because of the small size of logitsBitmaskPtrVec. - // TODO(enweiz): For chunked context, we currently build mask cache at the first context chunk, and apply - // the mask at the last context chunk. So, ideally we should sync the stream at the last context chunk. - CudaEvent event{}; - mCopyBufferManager.getStream().record(event); - stream.wait(event); - - if (mGuidedDecodingBackend == executor::GuidedDecodingConfig::GuidedDecodingBackend::kXGRAMMAR - && !decoderInputBuffers.decoderRequests.empty()) - { - SizeType32 batchIdx{0}; - for (size_t requestIdx = 0; requestIdx < decoderInputBuffers.decoderRequests.size(); ++requestIdx) - { - auto const& llmReq = decoderInputBuffers.decoderRequests.at(requestIdx); - - auto const& guidedDecodingParams = llmReq->getGuidedDecodingParams(); - if (guidedDecodingParams.has_value()) - { - auto const seqSlot = llmReq->mSeqSlot.value(); - - auto const& logits = decoderInputBuffers.decoderLogits.at(requestIdx); - auto const logitsBitmask = ITensor::at(mLogitsBitmask, {seqSlot}); - - // Use void* to unify the code for different mLogitsDtype - *reinterpret_cast(ITensor::at(mLogitsPtrVecHost, {batchIdx})->data()) = logits->data(); - *reinterpret_cast(ITensor::at(mLogitsBitmaskPtrVecHost, {batchIdx})->data()) - = logitsBitmask->data(); - - ++batchIdx; - } - } - if (batchIdx > 0) - { - runtimeBufferManager.copy( - *ITensor::slice(mLogitsPtrVecHost, 0, batchIdx), *ITensor::slice(mLogitsPtrVec, 0, batchIdx)); - runtimeBufferManager.copy(*ITensor::slice(mLogitsBitmaskPtrVecHost, 0, batchIdx), - *ITensor::slice(mLogitsBitmaskPtrVec, 0, batchIdx)); - - auto logitsBitmaskPtrVec = bufferCast(*mLogitsBitmaskPtrVec); - if (mLogitsDtype == nvinfer1::DataType::kFLOAT) - { - auto logitsPtrVec = bufferCast(*mLogitsPtrVec); - tensorrt_llm::kernels::invokeLogitsBitmask( - logitsPtrVec, logitsBitmaskPtrVec, batchIdx, mVocabSizePadded, stream.get()); - } - else if (mLogitsDtype == nvinfer1::DataType::kHALF) - { - auto logitsPtrVec = bufferCast(*mLogitsPtrVec); - tensorrt_llm::kernels::invokeLogitsBitmask( - logitsPtrVec, logitsBitmaskPtrVec, batchIdx, mVocabSizePadded, stream.get()); - } - else - { - TLLM_THROW("Unsupported logits data type."); - } - } - } -} - -} // namespace tensorrt_llm::batch_manager diff --git a/cpp/tensorrt_llm/batch_manager/handleContextLogits.cpp b/cpp/tensorrt_llm/batch_manager/handleContextLogits.cpp deleted file mode 100644 index 6f4a541ffcbb..000000000000 --- a/cpp/tensorrt_llm/batch_manager/handleContextLogits.cpp +++ /dev/null @@ -1,176 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2025 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. - */ - -#include "tensorrt_llm/batch_manager/handleContextLogits.h" - -#include "tensorrt_llm/batch_manager/decoderBuffers.h" -#include "tensorrt_llm/batch_manager/llmRequest.h" -#include "tensorrt_llm/batch_manager/medusaBuffers.h" -#include "tensorrt_llm/batch_manager/runtimeBuffers.h" -#include "tensorrt_llm/common/nvtxUtils.h" -#include "tensorrt_llm/runtime/iTensor.h" -#include "tensorrt_llm/runtime/runtimeKernels.h" -#include "tensorrt_llm/runtime/utils/debugUtils.h" - -namespace tr = tensorrt_llm::runtime; -namespace tru = tensorrt_llm::runtime::utils; - -namespace tensorrt_llm::batch_manager -{ - -using BufferManager = tensorrt_llm::runtime::BufferManager; -using TensorPtr = runtime::ITensor::SharedPtr; -using ITensor = runtime::ITensor; -using SizeType32 = tensorrt_llm::runtime::SizeType32; - -namespace -{ - -//! @brief Copy logits from context phase to beginning of generation logits. -//! @details Usually, this concerns logits of 1 token. In speculative decoding this concerns draftLen + 1 tokens. -void copyLastContextLogits(TensorPtr const& contextLogits, LlmRequest& llmReq, BufferManager const& bufferManager) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - auto const numLogits = contextLogits->getShape().d[0]; - for (int beam = 0; beam < llmReq.getBeamWidthByIter(); beam++) - { - // [beamWidth, mMaxNewTokens, vocabSizePadded] -> [numLogits, vocabSizePadded] - auto beamHostTensorPtr = ITensor::slice(llmReq.getGenerationLogitsHost(), {beam, 0}, numLogits); - bufferManager.copy(*contextLogits, *beamHostTensorPtr); - } - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void setupMedusaLogits(std::vector& medusaLogitsHeads, TensorPtr const& medusaLogitsDevice, - SizeType32 medusaHeads, SizeType32 logitsIndex, SizeType32 numLogits) -{ - for (SizeType32 hi = 0; hi < medusaHeads; ++hi) - { - TensorPtr logitsHead = ITensor::slice(medusaLogitsDevice, hi, 1); - logitsHead->squeeze(0); - medusaLogitsHeads[hi] = ITensor::slice(logitsHead, logitsIndex, numLogits); - } -} - -} // namespace - -SizeType32 HandleContextLogits::operator()(DecoderInputBuffers& inputBuffers, RequestVector const& contextRequests, - tr::ITensor::SharedPtr const& logits, std::vector const& numContextLogitsVec, - tr::ModelConfig const& modelConfig, tr::BufferManager const& manager, - OptionalRef medusaBuffers) const -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(HandleContextLogits); - - auto& decoderRequests = inputBuffers.decoderRequests; - decoderRequests.clear(); - decoderRequests.reserve(contextRequests.size()); - auto& allDecoderLogits = inputBuffers.decoderLogits; - allDecoderLogits.clear(); - allDecoderLogits.reserve(contextRequests.size()); - - SizeType32 batchIndex{0}; - SizeType32 logitsIndex{0}; - // Copy logits into decoderBuffers.logits - for (auto const& llmReq : contextRequests) - { - auto const numContextLogits = numContextLogitsVec.at(batchIndex); - auto const draftLength = llmReq->isLastContextChunk() ? llmReq->getNumDraftTokens() : 0; - - TLLM_LOG_DEBUG("logitsIndex: %d", logitsIndex); - TLLM_LOG_DEBUG("numContextLogits %d", numContextLogits); - TLLM_LOG_DEBUG("draftLength: %d", draftLength); - - if (modelConfig.computeContextLogits()) - { - // Since the computational graph has been modified, only the last token is needed. - TLLM_CHECK_WITH_INFO(!modelConfig.getSpeculativeDecodingMode().isMedusa() - && !modelConfig.getSpeculativeDecodingMode().isLookaheadDecoding(), - "Return context logits is not supported with Medusa and Lookahead decoding"); - - if (llmReq->getReturnContextLogits()) - { - if (llmReq->getPrepopulatedPromptLen() > 0) - { - TLLM_LOG_WARNING( - "Because of KV cache reuse, not all context logits could be produced for request %lu.", - llmReq->mRequestId); - } - TensorPtr contextLogitsDeviceView = ITensor::slice(logits, logitsIndex, numContextLogits); - TensorPtr contextLogitsHostView = ITensor::slice( - llmReq->getContextLogitsHost(), llmReq->getContextCurrentPosition(), numContextLogits); - // Copy to host directly - manager.copy(*contextLogitsDeviceView, *contextLogitsHostView); - } - } - logitsIndex += numContextLogits + draftLength; - - // Get the logits from the last context token and draft tokens - auto const numDecoderLogits = 1 + draftLength; - auto const seqSlot = llmReq->mSeqSlot.value(); - TensorPtr logitsView = ITensor::slice(logits, logitsIndex - numDecoderLogits, numDecoderLogits); - - if (modelConfig.getSpeculativeDecodingMode().hasDraftLogits()) - { - auto& medusaLogitsHeads = inputBuffers.predictedDraftLogits.at(seqSlot); - TLLM_CHECK(medusaBuffers); - setupMedusaLogits(medusaLogitsHeads, medusaBuffers->medusaLogitsDevice, - modelConfig.getSpeculativeDecodingModule().getMaxDraftPathLen(), logitsIndex - numDecoderLogits, - numDecoderLogits); - } - - // Save the last token logits of context into generation logits or - // save the accepted token logits from target model - if (llmReq->getReturnGenerationLogits()) - { - copyLastContextLogits(logitsView, *llmReq, manager); - } - - TLLM_CHECK_DEBUG_WITH_INFO(tru::tensorHasInvalid(*logitsView, manager, "logits") == false, - "Found invalid number (NaN or Inf) in logits"); - - if (llmReq->isLastContextChunk()) - { - TensorPtr decoderLogits; - auto const reqBeamWidth = llmReq->getBeamWidthByIter(); - if (reqBeamWidth > 1) - { - // Tile logits of context requests - auto const& logitsShape = logitsView->getShape(); - auto const logitsType = logitsView->getDataType(); - decoderLogits = manager.gpu(ITensor::makeShape({reqBeamWidth, logitsShape.d[1]}), logitsType); - tensorrt_llm::runtime::kernels::tileTensor( - *decoderLogits, *logitsView, reqBeamWidth, manager.getStream()); - decoderLogits->unsqueeze(0); - } - else - { - decoderLogits = logitsView; - decoderLogits->unsqueeze(1); - } - decoderRequests.push_back(llmReq); - allDecoderLogits.emplace_back(std::move(decoderLogits)); - } - - ++batchIndex; - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); - return logitsIndex; -} - -} // namespace tensorrt_llm::batch_manager diff --git a/cpp/tensorrt_llm/batch_manager/handleGenerationLogits.cpp b/cpp/tensorrt_llm/batch_manager/handleGenerationLogits.cpp deleted file mode 100644 index e2a7486b050a..000000000000 --- a/cpp/tensorrt_llm/batch_manager/handleGenerationLogits.cpp +++ /dev/null @@ -1,161 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2025 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. - */ - -#include "tensorrt_llm/batch_manager/handleGenerationLogits.h" - -#include "tensorrt_llm/batch_manager/decoderBuffers.h" -#include "tensorrt_llm/batch_manager/llmRequest.h" -#include "tensorrt_llm/batch_manager/medusaBuffers.h" -#include "tensorrt_llm/batch_manager/runtimeBuffers.h" -#include "tensorrt_llm/batch_manager/utils/inflightBatchingUtils.h" -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/nvtxUtils.h" -#include "tensorrt_llm/runtime/iTensor.h" -#include "tensorrt_llm/runtime/utils/debugUtils.h" - -namespace tr = tensorrt_llm::runtime; -namespace tru = tensorrt_llm::runtime::utils; - -namespace tensorrt_llm::batch_manager -{ - -using BufferManager = tensorrt_llm::runtime::BufferManager; -using TensorPtr = runtime::ITensor::SharedPtr; -using ITensor = runtime::ITensor; -using SizeType32 = tensorrt_llm::runtime::SizeType32; - -namespace -{ - -//! @brief Copy logits from generation phase under streaming mode. -void copyStreamingGenerationLogits(BufferManager const& bufferManager, LlmRequest& llmReq) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - // If llmRequest is streaming, directly copy to host. - // Only one token's logits needs to be copied each time. - TLLM_CHECK(llmReq.getGenerationLogitsFragmentsSize() == 1); - - SizeType32 numGenerationToken = llmReq.getMaxBeamNumTokens() - llmReq.mPromptLen; - TensorPtr const& generationLogitsHost - = llmReq.getGenerationLogitsHost(); // [mMaxNewTokens (or 1), beamWidth, vocabSizePadded] - - TensorPtr hostTensorPtr - = ITensor::slice(generationLogitsHost, numGenerationToken, 1); // [1, beamWidth, vocabSizePadded] - TensorPtr deviceTensorPtr = *(llmReq.getGenerationLogitsFragments().begin()); - - bufferManager.copy(*deviceTensorPtr, *hostTensorPtr); - llmReq.clearGenerationLogitsFragments(); - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void setupMedusaLogits(std::vector& medusaLogitsHeads, TensorPtr const& medusaLogitsDevice, - SizeType32 medusaHeads, SizeType32 logitsIndex, SizeType32 numLogits) -{ - for (SizeType32 hi = 0; hi < medusaHeads; ++hi) - { - TensorPtr logitsHead = ITensor::slice(medusaLogitsDevice, hi, 1); - logitsHead->squeeze(0); - medusaLogitsHeads[hi] = ITensor::slice(logitsHead, logitsIndex, numLogits); - } -} - -} // namespace - -void HandleGenerationLogits::operator()(DecoderInputBuffers& inputBuffers, RequestVector const& generationRequests, - tr::ITensor::SharedPtr const& logits, tr::SizeType32 logitsIndex, tr::ModelConfig const& modelConfig, - tr::BufferManager const& manager, OptionalRef genRuntimeBuffers, - OptionalRef medusaBuffers) const -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(HandleGenerationLogits); - - auto& decoderRequests = inputBuffers.decoderRequests; - decoderRequests.reserve(decoderRequests.size() + generationRequests.size()); - auto& allDecoderLogits = inputBuffers.decoderLogits; - allDecoderLogits.reserve(allDecoderLogits.size() + generationRequests.size()); - - for (auto const& llmReq : generationRequests) - { - auto const reqBeamWidth = llmReq->getBeamWidthByIter(); - auto const seqSlot = llmReq->mSeqSlot.value(); - - auto const draftLength = llmReq->getNumDraftTokens(); - auto const numLogits = draftLength + reqBeamWidth; - - TLLM_CHECK(draftLength == 0 || reqBeamWidth == 1); - - TLLM_LOG_DEBUG("logitsIndex: %d", logitsIndex); - TLLM_LOG_DEBUG("draftLength: %d", draftLength); - TLLM_LOG_DEBUG("reqBeamWidth: %d", reqBeamWidth); - - // genRuntimeBuffers.logits shape: [numGen*reqBeamWidth, vocabSize] - // logitsView shape: [numLogits, vocabSize] - TensorPtr logitsView = ITensor::slice(logits, logitsIndex, numLogits); - TLLM_CHECK_DEBUG_WITH_INFO(tru::tensorHasInvalid(*logitsView, manager, "logits") == false, - "Found invalid number (NaN or Inf) in logits"); - - TLLM_CHECK(llmReq->isGenerationInProgressState()); - TensorPtr decoderLogits; - if (reqBeamWidth > 1) - { - decoderLogits = logitsView; - decoderLogits->unsqueeze(0); - } - else - { - decoderLogits = logitsView; - decoderLogits->unsqueeze(1); - } - decoderRequests.push_back(llmReq); - allDecoderLogits.emplace_back(std::move(decoderLogits)); - - if (llmReq->getReturnGenerationLogits()) - { - TLLM_CHECK_WITH_INFO(modelConfig.getSpeculativeDecodingMode().isNone() - || modelConfig.getSpeculativeDecodingMode().isDraftTokensExternal(), - "Only speculative decoding with external draft tokens supports returning generation logits"); - - // Push into fragments vector - llmReq->addGenerationLogitsFragment(logitsView); - TLLM_CHECK( - llmReq->getGenerationLogitsFragmentsSize() <= RuntimeBuffers::GenerationLogitsCache::kCACHE_LENGTH); - if (llmReq->isStreaming()) - { - copyStreamingGenerationLogits(manager, *llmReq); - } - // Copy back to host for every kCACHE_LENGTH steps to mitigate GPU memory pressure - else if (llmReq->getGenerationLogitsFragmentsSize() == RuntimeBuffers::GenerationLogitsCache::kCACHE_LENGTH) - { - TLLM_CHECK(genRuntimeBuffers); - auto constexpr beforeDecoder = true; - utils::copyGenerationLogits(genRuntimeBuffers->generationLogitsCache, manager, *llmReq, beforeDecoder); - } - } - if (modelConfig.getSpeculativeDecodingMode().hasDraftLogits()) - { - auto& medusaLogitsHeads = inputBuffers.predictedDraftLogits.at(seqSlot); - TLLM_CHECK(medusaBuffers); - setupMedusaLogits(medusaLogitsHeads, medusaBuffers->medusaLogitsDevice, - modelConfig.getSpeculativeDecodingModule().getMaxDraftPathLen(), logitsIndex, draftLength); - } - logitsIndex += numLogits; - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -} // namespace tensorrt_llm::batch_manager diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp index fd5f8d4ddd75..266cdafc8d06 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp @@ -27,6 +27,7 @@ #include "tensorrt_llm/common/envUtils.h" #include "tensorrt_llm/common/logger.h" #include "tensorrt_llm/common/memoryUtils.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/executor/executor.h" #include "tensorrt_llm/kernels/kvCacheIndex.h" #include "tensorrt_llm/runtime/common.h" @@ -586,7 +587,7 @@ std::map BlockManager::calculateWindowSizeToShare( BlockManager::BlockManager(std::vector const& numKvHeadsPerLayer, SizeType32 sizePerHead, SizeType32 tokensPerBlock, BlocksPerWindow const& blocksPerWindow, SizeType32 maxNumSequences, CudaStreamPtr stream, SizeType32 maxSequenceLength, SizeType32 maxBeamWidth, std::vector const& maxAttentionWindowVec, - nvinfer1::DataType dtype, SizeType32 sinkBubbleLength, SizeType32 chunkSize, CacheType cacheType, + tensorrt_llm::DataType dtype, SizeType32 sinkBubbleLength, SizeType32 chunkSize, CacheType cacheType, std::optional secondaryOffloadMinPriority, std::shared_ptr eventManager, bool enablePartialReuse, bool copyOnPartialReuse, std::shared_ptr kvCacheConnectorManager, @@ -740,7 +741,7 @@ BlockManager::BlockManager(std::vector const& numKvHeadsPerLayer, Si "Maybe you tried changing either of them to an std::unordered_map?"); } -WindowBlockManager::WindowBlockManager(nvinfer1::DataType dtype, SizeType32 windowSize, +WindowBlockManager::WindowBlockManager(tensorrt_llm::DataType dtype, SizeType32 windowSize, std::vector const& managedLayers, std::vector const& numKvHeadsPerLayer, SizeType32 sizePerHead, SizeType32 tokensPerBlock, bool isSWA, SizeType32 blocksInPrimaryPool, SizeType32 blocksInSecondaryPool, SizeType32 maxNumSequences, std::shared_ptr stream, @@ -840,7 +841,7 @@ WindowBlockManager::WindowBlockManager(nvinfer1::DataType dtype, SizeType32 wind // to specify FP4 related parameters (scale dtypes, etc)? This can also be passed // in the constructor. constexpr SizeType32 kQuantBlockSizeNVFP4 = 16; - if (dtype == nvinfer1::DataType::kFP4) + if (dtype == tensorrt_llm::DataType::kFP4) { createBlockScalePools(kQuantBlockSizeNVFP4); } @@ -1092,7 +1093,7 @@ void BlockManager::allocatePools(bool useUvm) void WindowBlockManager::allocatePools(bool useUvm) { - constexpr nvinfer1::DataType kScaleDtypeNVFP4 = nvinfer1::DataType::kFP8; + constexpr tensorrt_llm::DataType kScaleDtypeNVFP4 = tensorrt_llm::DataType::kFP8; bool const requestFabricMemory = tc::getEnvKVCachePoolUseFabricMemory(); bool const fabricMemorySupported = FabricMemory::supportFabricMemory(); @@ -1118,21 +1119,21 @@ void WindowBlockManager::allocatePools(bool useUvm) auto blockSize = pool.blockSize; auto poolDtype = pool.containsBlockScales ? kScaleDtypeNVFP4 : mDataType; #ifdef ENABLE_FP4 - auto const poolIsFP4 = poolDtype == nvinfer1::DataType::kFP4; + auto const poolIsFP4 = poolDtype == tensorrt_llm::DataType::kFP4; #else auto const poolIsFP4 = false; #endif if (poolIsFP4) { - poolDtype = nvinfer1::DataType::kINT8; + poolDtype = tensorrt_llm::DataType::kINT8; } if (pool.containsIndexerKCache) { - poolDtype = nvinfer1::DataType::kUINT8; + poolDtype = tensorrt_llm::DataType::kUINT8; } - nvinfer1::Dims cacheShape = isRecurrentState() + tensorrt_llm::Dims cacheShape = isRecurrentState() ? ITensor::makeShape({pool.numLayers, mNumPrimaryBlocks, mKVFactor, blockSize}) : ITensor::makeShape({mNumPrimaryBlocks, pool.numLayers, mKVFactor, blockSize}); pool.layerFirstLayout = isRecurrentState(); @@ -1166,7 +1167,7 @@ void WindowBlockManager::allocatePools(bool useUvm) if (mNumSecondaryBlocks > 0) { - nvinfer1::Dims cacheShapeOffload = isRecurrentState() + tensorrt_llm::Dims cacheShapeOffload = isRecurrentState() ? ITensor::makeShape({pool.numLayers, mNumSecondaryBlocks, mKVFactor, blockSize}) : ITensor::makeShape({mNumSecondaryBlocks, pool.numLayers, mKVFactor, blockSize}); TLLM_LOG_DEBUG("[%s] Allocating secondary pool with %d blocks for %d layers with %d kv heads", @@ -1344,7 +1345,7 @@ BlockPtr WindowBlockManager::getFreeBlock(GenerationRequest& sequence, executor: return block; } -void WindowBlockManager::setOffsets(tk::KVCacheIndex* offsetsPtr, nvinfer1::Dims const& offsetsShape, +void WindowBlockManager::setOffsets(tk::KVCacheIndex* offsetsPtr, tensorrt_llm::Dims const& offsetsShape, SizeType32 beamIdx, SizeType32 blockIdx, KVCacheBlock::IdType blockId) const { auto constexpr kIdx = 0; @@ -1382,7 +1383,7 @@ void WindowBlockManager::setOffsets(tk::KVCacheIndex* offsetsPtr, nvinfer1::Dims } } -void BlockManager::setOffsets(tk::KVCacheIndex* offsetsPtr, nvinfer1::Dims const& offsetsShape, SizeType32 beamIdx, +void BlockManager::setOffsets(tk::KVCacheIndex* offsetsPtr, tensorrt_llm::Dims const& offsetsShape, SizeType32 beamIdx, SizeType32 blockIdx, KVCacheBlock::IdType blockId, SizeType32 windowSize) const { mWindowBlockManagers.at(windowSize).setOffsets(offsetsPtr, offsetsShape, beamIdx, blockIdx, blockId); @@ -3229,7 +3230,7 @@ void WindowBlockManager::schedulingReleaseBlocks(RequestIdType requestId) KVCacheManager::KVCacheManager(SizeType32 numLayers, SizeType32 numKvHeads, SizeType32 sizePerHead, SizeType32 tokensPerBlock, BlocksPerWindow const& blocksPerWindow, SizeType32 maxNumSequences, - SizeType32 maxBeamWidth, std::vector const& maxAttentionWindowVec, nvinfer1::DataType dtype, + SizeType32 maxBeamWidth, std::vector const& maxAttentionWindowVec, tensorrt_llm::DataType dtype, SizeType32 sinkTokenLength, int64_t stream, runtime::SizeType32 maxSequenceLength, SizeType32 chunkSize, bool enableBlockReuse, CacheType cacheType, bool enablePartialReuse, bool copyOnPartialReuse, bool enableIndexerKCache, SizeType32 indexerKCacheQuantBlockSize, SizeType32 indexerKCacheIndexHeadDim, @@ -3246,7 +3247,7 @@ KVCacheManager::KVCacheManager(SizeType32 numLayers, SizeType32 numKvHeads, Size KVCacheManager::KVCacheManager(std::vector const& numKvHeadsPerLayer, SizeType32 sizePerHead, SizeType32 tokensPerBlock, BlocksPerWindow const& blocksPerWindow, SizeType32 maxNumSequences, - SizeType32 maxBeamWidth, std::vector const& maxAttentionWindowVec, nvinfer1::DataType dtype, + SizeType32 maxBeamWidth, std::vector const& maxAttentionWindowVec, tensorrt_llm::DataType dtype, SizeType32 sinkTokenLength, int64_t stream, runtime::SizeType32 maxSequenceLength, SizeType32 chunkSize, bool enableBlockReuse, CacheType cacheType, std::optional secondaryOffloadMinPriority, std::shared_ptr eventManager, bool enablePartialReuse, bool copyOnPartialReuse, @@ -3265,7 +3266,7 @@ KVCacheManager::KVCacheManager(std::vector const& numKvHeadsPerLayer KVCacheManager::KVCacheManager(std::vector const& numKvHeadsPerLayer, SizeType32 sizePerHead, SizeType32 tokensPerBlock, BlocksPerWindow const& blocksPerWindow, SizeType32 maxNumSequences, - SizeType32 maxBeamWidth, std::vector const& maxAttentionWindowVec, nvinfer1::DataType dtype, + SizeType32 maxBeamWidth, std::vector const& maxAttentionWindowVec, tensorrt_llm::DataType dtype, SizeType32 sinkTokenLength, CudaStreamPtr stream, runtime::SizeType32 maxSequenceLength, SizeType32 chunkSize, bool enableBlockReuse, CacheType cacheType, std::optional secondaryOffloadMinPriority, std::shared_ptr eventManager, bool enablePartialReuse, bool copyOnPartialReuse, @@ -3307,7 +3308,7 @@ KVCacheManager::KVCacheManager(std::vector const& numKvHeadsPerLayer KVCacheManager::KVCacheManager(SizeType32 numLayers, SizeType32 numKvHeads, SizeType32 sizePerHead, SizeType32 tokensPerBlock, BlocksPerWindow const& blocksPerWindow, SizeType32 maxNumSequences, - SizeType32 maxBeamWidth, std::vector const& maxAttentionWindowVec, nvinfer1::DataType dtype, + SizeType32 maxBeamWidth, std::vector const& maxAttentionWindowVec, tensorrt_llm::DataType dtype, SizeType32 sinkTokenLength, CudaStreamPtr stream, runtime::SizeType32 maxSequenceLength, SizeType32 chunkSize, bool enableBlockReuse, CacheType cacheType, std::optional secondaryOffloadMinPriority, std::shared_ptr eventManager, bool enablePartialReuse, bool copyOnPartialReuse, @@ -3340,7 +3341,7 @@ void KVCacheManager::allocatePools(bool useUvm) // a future per-window override map can mix precisions inside a single manager. auto const poolDataType = primaryPool->getDataType(); #ifdef ENABLE_FP4 - auto const isFp4 = poolDataType == nvinfer1::DataType::kFP4; + auto const isFp4 = poolDataType == tensorrt_llm::DataType::kFP4; #else auto const isFp4 = false; #endif @@ -4283,7 +4284,7 @@ std::map computeWindowSizeShares( } // namespace BlocksPerWindow BaseKVCacheManager::calculateMaxNumBlocks(executor::KvCacheConfig const& config, - nvinfer1::DataType dtype, std::vector const& numKvHeadsPerLayer, SizeType32 sizePerHead, + tensorrt_llm::DataType dtype, std::vector const& numKvHeadsPerLayer, SizeType32 sizePerHead, SizeType32 tokensPerBlock, WorldConfig const& worldConfig, std::map> const& windowSizeToLayers, uint64_t allottedPrimaryMemBytes, uint64_t allottedSecondaryMemBytes, size_t extraCostMemory, SizeType32 kvFactor, SizeType32 maxBatchSize, diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp index c28d1e476137..33ad5b4a1675 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp @@ -22,6 +22,7 @@ #include "tensorrt_llm/batch_manager/kvCacheEventManager.h" #include "tensorrt_llm/batch_manager/kvCacheManager.h" #include "tensorrt_llm/common/logger.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/executor/executor.h" #include "tensorrt_llm/kernels/kvCachePartialCopy.h" #include "tensorrt_llm/runtime/bufferManager.h" @@ -163,8 +164,8 @@ void KVCacheTransferManager::copyBlock(BlockPtr const& src, BlockPtr const& dst, // If no partial tokens or if the dataType is not supported for partial copy, copy entire block. // Note that nvfp4 kv cache SFs use an interleaved layout, so we need to copy the entire block. - if (numTokensToCopy <= 0 || srcPtr->getDataType() == nvinfer1::DataType::kINT4 - || srcPtr->getDataType() == nvinfer1::DataType::kFP4 || containsBlockScales) + if (numTokensToCopy <= 0 || srcPtr->getDataType() == tensorrt_llm::DataType::kINT4 + || srcPtr->getDataType() == tensorrt_llm::DataType::kFP4 || containsBlockScales) { // For partial copy not implemented with these data types, // just do a full copy. @@ -461,8 +462,8 @@ std::size_t KVCacheTransferManager::computeBlockTransferBytes( // Mirror the logic in copyBlock: a partial copy only happens when numTokensToCopy > 0, // the data type supports it (not kINT4/kFP4), not block scales, and numTokensToCopy < tokensPerBlock. - bool const isPartialCopy = numTokensToCopy > 0 && dataType != nvinfer1::DataType::kINT4 - && dataType != nvinfer1::DataType::kFP4 && !pool.containsBlockScales + bool const isPartialCopy = numTokensToCopy > 0 && dataType != tensorrt_llm::DataType::kINT4 + && dataType != tensorrt_llm::DataType::kFP4 && !pool.containsBlockScales && numTokensToCopy < pool.tokensPerBlock; if (isPartialCopy) diff --git a/cpp/tensorrt_llm/batch_manager/logitsPostProcessor.cpp b/cpp/tensorrt_llm/batch_manager/logitsPostProcessor.cpp deleted file mode 100644 index 95b324f0f2ec..000000000000 --- a/cpp/tensorrt_llm/batch_manager/logitsPostProcessor.cpp +++ /dev/null @@ -1,88 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2025 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. - */ - -#include "tensorrt_llm/batch_manager/logitsPostProcessor.h" - -#include "tensorrt_llm/batch_manager/decoderBuffers.h" -#include "tensorrt_llm/batch_manager/llmRequest.h" -#include "tensorrt_llm/batch_manager/runtimeBuffers.h" -#include "tensorrt_llm/common/nvtxUtils.h" -#include "tensorrt_llm/runtime/iTensor.h" - -namespace tr = tensorrt_llm::runtime; - -namespace tensorrt_llm::batch_manager -{ - -using TensorPtr = runtime::ITensor::SharedPtr; -using ITensor = runtime::ITensor; -using SizeType32 = tensorrt_llm::runtime::SizeType32; - -bool LogitsPostProcessor::operator()(DecoderInputBuffers& inputBuffers, bool replicateLogitsPostProcessor, - tr::WorldConfig const& worldConfig, CudaStreamPtr const& stream, - std::optional const& logitsPostProcessorBatched) const -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(LogitsPostProcessor); - - // Arguments for batched processor - std::vector reqIdsVec; - std::vector logitsVec; - std::vector> beamTokensVec; - std::vector> clientIdsVec; - - bool logitsPostProcessorIsApplied = false; - for (size_t batchIdx = 0; batchIdx < inputBuffers.decoderRequests.size(); ++batchIdx) - { - auto const& llmReq = inputBuffers.decoderRequests.at(batchIdx); - auto& logits = inputBuffers.decoderLogits.at(batchIdx); - - // Invoke non-batched processor or collect arguments for batched processor - if (llmReq->mLogitsPostProcessor) - { - logitsPostProcessorIsApplied = true; - if (replicateLogitsPostProcessor || worldConfig.isFirstTensorParallelRank()) - { - (*llmReq->mLogitsPostProcessor)( - llmReq->mRequestId, logits, llmReq->getTokens(), stream, llmReq->mClientId); - } - } - else if (llmReq->mApplyLogitsPostProcessorBatched) - { - reqIdsVec.push_back(llmReq->mRequestId); - logitsVec.push_back(logits); - beamTokensVec.emplace_back(llmReq->getTokens()); - clientIdsVec.push_back(llmReq->mClientId); - } - } - - // Invoke batched processor - if (!reqIdsVec.empty()) - { - logitsPostProcessorIsApplied = true; - if (replicateLogitsPostProcessor || worldConfig.isFirstTensorParallelRank()) - { - (*logitsPostProcessorBatched)(reqIdsVec, logitsVec, beamTokensVec, stream, clientIdsVec); - } - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); - - return logitsPostProcessorIsApplied; -} - -} // namespace tensorrt_llm::batch_manager diff --git a/cpp/tensorrt_llm/batch_manager/loraBuffers.cpp b/cpp/tensorrt_llm/batch_manager/loraBuffers.cpp deleted file mode 100644 index b67b72f6c49a..000000000000 --- a/cpp/tensorrt_llm/batch_manager/loraBuffers.cpp +++ /dev/null @@ -1,109 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2025 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. - */ - -#include "loraBuffers.h" - -#include "tensorrt_llm/batch_manager/llmRequest.h" -#include "tensorrt_llm/runtime/loraUtils.h" - -namespace tensorrt_llm::batch_manager -{ - -LoraBuffers::LoraBuffers(SizeType32 maxBatchSize, SizeType32 maxBeamWidth, runtime::TllmRuntime const& tllmRuntime, - runtime::ModelConfig const& modelConfig, runtime::WorldConfig const& worldConfig) -{ - auto const localNbLayers - = modelConfig.getNbAttentionLayers(worldConfig.getPipelineParallelism(), worldConfig.getPipelineParallelRank()); - auto const firstLayerId = worldConfig.getPipelineParallelRank() * localNbLayers; - - auto nbModelConfigs = static_cast(modelConfig.getLoraModules().size()); - - // there are 3 pointers: LoRA A, LoRA B, and a DoRA magnitude (null if not DoRA) - auto loraWeightsPtrsShape - = runtime::ITensor::makeShape({nbModelConfigs, localNbLayers, maxBatchSize * maxBeamWidth, 3}); - auto loraAdapterSizesShape - = runtime::ITensor::makeShape({nbModelConfigs, localNbLayers, maxBatchSize * maxBeamWidth}); - - auto firstModuleName = std::string(modelConfig.getLoraModules().front().name()); - auto ptrsFieldName = firstModuleName + "_lora_weights_pointers_" + std::to_string(firstLayerId); - auto rankFieldName = firstModuleName + "_lora_ranks_" + std::to_string(firstLayerId); - auto weightsPtrDtype = tllmRuntime.getEngine().getTensorDataType(ptrsFieldName.c_str()); - auto ranksDtype = tllmRuntime.getEngine().getTensorDataType(rankFieldName.c_str()); - - mLoraManager.create(modelConfig); - - mLoraWeightsPointersHost = runtime::BufferManager::pinned(loraWeightsPtrsShape, weightsPtrDtype); - mLoraAdapterSizesHost = runtime::BufferManager::pinned(loraAdapterSizesShape, ranksDtype); -} - -void LoraBuffers::fill(RequestVector const& contextRequests, RequestVector const& genRequests, - PeftTable const& peftTable, runtime::BufferManager const& manager, runtime::ModelConfig const& modelConfig, - runtime::WorldConfig const& worldConfig) -{ - manager.setZero(*mLoraWeightsPointersHost); - manager.setZero(*mLoraAdapterSizesHost); - - SizeType32 batchIdx{0}; - for (auto const& requests : {contextRequests, genRequests}) - { - for (auto const& llmReq : requests) - { - auto const optReqLoraWeights = llmReq->getLoraWeights(); - auto const optReqLoraConfig = llmReq->getLoraConfig(); - - auto const isContextRequest = llmReq->isContextInitState(); - auto const beamWidth = isContextRequest ? 1 : llmReq->mSamplingConfig.beamWidth; - auto const peftIt = peftTable.find(llmReq->mRequestId); - if (peftIt != peftTable.end()) - { - auto const& peftValues = peftIt->second; - if (!peftValues.empty()) - { - mLoraManager.fillInputTensors(mLoraWeightsPointersHost, mLoraAdapterSizesHost, peftIt->second, - batchIdx, beamWidth, modelConfig, worldConfig); - } - } - ++batchIdx; - } - } -} - -void LoraBuffers::validate(std::optional const& optTaskId, - std::optional const& optReqLoraWeights, std::optional const& optReqLoraConfig, - runtime::ModelConfig const& modelConfig, runtime::WorldConfig const& worldConfig) -{ - runtime::lora::loraValidateRequestTensors(optTaskId, optReqLoraWeights, optReqLoraConfig, modelConfig, worldConfig); -} - -void LoraBuffers::insertInputTensors(TensorMap& inputTensors, TensorPtr weightsPtrs, TensorPtr adapterSizes, - runtime::ModelConfig const& modelConfig, runtime::WorldConfig const& worldConfig) const -{ - mLoraManager.insertInputTensors(inputTensors, weightsPtrs, adapterSizes, modelConfig, worldConfig); -} - -void LoraBuffers::reshape(SizeType32 numSequences) -{ - auto weightsPtrsShape = mLoraWeightsPointersHost->getShape(); - weightsPtrsShape.d[2] = numSequences; - mLoraWeightsPointersHost->reshape(weightsPtrsShape); - - auto adapterSizesShape = mLoraAdapterSizesHost->getShape(); - adapterSizesShape.d[2] = numSequences; - mLoraAdapterSizesHost->reshape(adapterSizesShape); -} - -} // namespace tensorrt_llm::batch_manager diff --git a/cpp/tensorrt_llm/batch_manager/loraBuffers.h b/cpp/tensorrt_llm/batch_manager/loraBuffers.h deleted file mode 100644 index 3ba68995518f..000000000000 --- a/cpp/tensorrt_llm/batch_manager/loraBuffers.h +++ /dev/null @@ -1,61 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2025 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. - */ - -#pragma once - -#include "tensorrt_llm/batch_manager/common.h" -#include "tensorrt_llm/runtime/bufferManager.h" -#include "tensorrt_llm/runtime/common.h" -#include "tensorrt_llm/runtime/iTensor.h" -#include "tensorrt_llm/runtime/loraManager.h" -#include "tensorrt_llm/runtime/modelConfig.h" -#include "tensorrt_llm/runtime/tllmRuntime.h" -#include "tensorrt_llm/runtime/worldConfig.h" - -namespace tensorrt_llm::batch_manager -{ - -class LoraBuffers -{ -public: - using SizeType32 = tensorrt_llm::runtime::SizeType32; - using PeftTable = runtime::LoraManager::PeftTable; - using TensorPtr = runtime::ITensor::SharedPtr; - using TensorMap = runtime::StringPtrMap; - - TensorPtr mLoraWeightsPointersHost; - TensorPtr mLoraAdapterSizesHost; - - runtime::LoraManager mLoraManager; - - LoraBuffers(SizeType32 maxBatchSize, SizeType32 maxBeamWidth, runtime::TllmRuntime const& tllmRuntime, - runtime::ModelConfig const& modelConfig, runtime::WorldConfig const& worldConfig); - - static void validate(std::optional const& optTaskId, - std::optional const& optReqLoraWeights, std::optional const& optReqLoraConfig, - runtime::ModelConfig const& modelConfig, runtime::WorldConfig const& worldConfig); - - void fill(RequestVector const& contextRequests, RequestVector const& genRequests, PeftTable const& peftTable, - runtime::BufferManager const& manager, runtime::ModelConfig const& modelConfig, - runtime::WorldConfig const& worldConfig); - - void insertInputTensors(TensorMap& inputTensors, TensorPtr weightsPtrs, TensorPtr adapterSizes, - runtime::ModelConfig const& modelConfig, runtime::WorldConfig const& worldConfig) const; - - void reshape(SizeType32 numSequences); -}; -} // namespace tensorrt_llm::batch_manager diff --git a/cpp/tensorrt_llm/batch_manager/makeDecodingBatchInputOutput.cpp b/cpp/tensorrt_llm/batch_manager/makeDecodingBatchInputOutput.cpp deleted file mode 100644 index 3e494a6383ec..000000000000 --- a/cpp/tensorrt_llm/batch_manager/makeDecodingBatchInputOutput.cpp +++ /dev/null @@ -1,198 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2025 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. - */ - -#include "tensorrt_llm/batch_manager/makeDecodingBatchInputOutput.h" -#include "tensorrt_llm/batch_manager/decoderBuffers.h" -#include "tensorrt_llm/batch_manager/llmRequest.h" -#include "tensorrt_llm/batch_manager/runtimeBuffers.h" -#include "tensorrt_llm/common/cudaUtils.h" -#include "tensorrt_llm/common/logger.h" -#include "tensorrt_llm/runtime/decoderState.h" -#include "tensorrt_llm/runtime/iGptDecoderBatched.h" - -namespace tr = tensorrt_llm::runtime; - -namespace tensorrt_llm::batch_manager -{ -using SizeType32 = MakeDecodingBatchInputOutput::SizeType32; -using TensorPtr = MakeDecodingBatchInputOutput::TensorPtr; - -void MakeDecodingBatchInputOutput::createDecoderBatchInputs(DecoderInputBuffers& inputBuffers, - std::vector const& activeSlots, runtime::decoder::DecoderState const& decoderState) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - auto const& numDecodingEngineTokens = decoderState.getNumDecodingEngineTokens(); - auto const& maxDecodingEngineTokens = decoderState.getMaxDecodingEngineTokens(); - auto const& maxDecodingDecoderTokens = decoderState.getMaxDecodingDecoderTokens(); - auto const maxDecoderSteps = common::ceilDiv(maxDecodingEngineTokens, maxDecodingDecoderTokens); - - auto& batchSlots = inputBuffers.forwardBatchSlots; - auto& decoderLogits = inputBuffers.decoderLogits; - - for (SizeType32 step = 0; step < maxDecoderSteps; ++step) - { - batchSlots.at(step)->resize(activeSlots.size()); - } - - auto constexpr singleRequest = 1; - - std::vector batchSizes(maxDecoderSteps); - std::vector> batchLogits(maxDecoderSteps); - auto maxActiveDecoderSteps = 1; - for (size_t batchIdx = 0; batchIdx < activeSlots.size(); ++batchIdx) - { - auto const slot = activeSlots.at(batchIdx); - auto const& logits = decoderLogits.at(batchIdx); - - auto const numDecoderSteps = common::ceilDiv(numDecodingEngineTokens.at(slot), maxDecodingDecoderTokens); - maxActiveDecoderSteps = std::max(maxActiveDecoderSteps, numDecoderSteps); - for (SizeType32 step = 0; step < numDecoderSteps; ++step) - { - auto batchSlotsRange = tr::BufferRange(*batchSlots.at(step)); - batchSlotsRange[batchSizes[step]] = slot; - batchSizes[step]++; - auto logitsSlice = tr::ITensor::slice(logits, step, singleRequest); - batchLogits[step].emplace_back(std::move(logitsSlice)); - } - } - - for (SizeType32 step = 0; step < maxDecoderSteps; ++step) - { - batchSlots.at(step)->resize(batchSizes[step]); - } - batchLogits.resize(maxActiveDecoderSteps); - - inputBuffers.maxDecoderSteps = maxActiveDecoderSteps; - inputBuffers.batchLogits = batchLogits; - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -namespace -{ - -std::pair, std::vector> getActiveSlots(RequestVector const& decoderRequests) -{ - std::vector activeSlots; - std::vector generationSteps; - for (auto const& llmReq : decoderRequests) - { - activeSlots.push_back(llmReq->mSeqSlot.value()); - generationSteps.push_back(llmReq->getDecodingIter()); - } - - return {activeSlots, generationSteps}; -} - -//! @brief Sets inputs for explicit draft tokens. -void setExplicitDraftTokensInputs(tr::DecodingInput& dInput, RuntimeBuffers const& fusedRuntimeBuffers) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - TLLM_CHECK(fusedRuntimeBuffers.mExplicitDraftTokensBuffers); - auto const& explicitDraftTokensInputs = fusedRuntimeBuffers.mExplicitDraftTokensBuffers->engineOutputs; - auto const& explicitDraftTokensLastInputs = fusedRuntimeBuffers.mExplicitDraftTokensBuffers->engineInputs; - - dInput.explicitDraftTokensInputs = tr::DecodingInput::ExplicitDraftTokensInputs(); - dInput.explicitDraftTokensInputs->nextDraftTokens = explicitDraftTokensInputs.nextDraftTokens; - dInput.explicitDraftTokensInputs->nextFlatTokens = explicitDraftTokensInputs.nextFlatTokens; - dInput.explicitDraftTokensInputs->nextDraftIndices = explicitDraftTokensInputs.nextDraftIndices; - dInput.explicitDraftTokensInputs->nextDraftProbs = explicitDraftTokensInputs.nextDraftProbs; - dInput.explicitDraftTokensInputs->lastDraftTokens = explicitDraftTokensLastInputs.draftTokens; - dInput.explicitDraftTokensInputs->lastDraftIndices = explicitDraftTokensLastInputs.draftIndices; - dInput.explicitDraftTokensInputs->lastPositionIdsBase = explicitDraftTokensLastInputs.positionIdsBase; - dInput.explicitDraftTokensInputs->masks = explicitDraftTokensInputs.masks; - dInput.explicitDraftTokensInputs->packedPositionIds = explicitDraftTokensInputs.packedPositionIds; - dInput.explicitDraftTokensInputs->bestPathLengths = explicitDraftTokensInputs.bestPathLengths; - dInput.explicitDraftTokensInputs->bestPathIndices = explicitDraftTokensInputs.bestPathIndices; - dInput.explicitDraftTokensInputs->nextGenerationLengths = explicitDraftTokensInputs.nextGenerationLengths; - dInput.explicitDraftTokensInputs->lastGenerationLengths = explicitDraftTokensLastInputs.generationLengths; - dInput.explicitDraftTokensInputs->maxGenLengthDevice = explicitDraftTokensInputs.maxGenToken; - // Slots in request order - dInput.explicitDraftTokensInputs->seqSlots = fusedRuntimeBuffers.seqSlots; - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -//! @brief Sets inputs for eagle decoding. -void setEagleInputs(tr::DecodingInput& dInput, RuntimeBuffers const& fusedRuntimeBuffers) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - TLLM_CHECK(fusedRuntimeBuffers.mEagleBuffers); - auto const& eagleInputs = fusedRuntimeBuffers.mEagleBuffers->engineOutputs; - auto const& eagleLastInputs = fusedRuntimeBuffers.mEagleBuffers->engineInputs; - - dInput.eagleInputs = tr::DecodingInput::EagleInputs(); - dInput.eagleInputs->nextDraftTokens = eagleInputs.nextDraftTokens; - dInput.eagleInputs->nextDraftLens = eagleInputs.nextDraftLens; - dInput.eagleInputs->nextDraftPaths = eagleInputs.nextDraftPaths; - dInput.eagleInputs->lastDraftTokens = eagleLastInputs.draftTokens; - dInput.eagleInputs->lastDraftLens = eagleLastInputs.draftLens; - dInput.eagleInputs->lastDraftPaths = eagleLastInputs.draftPaths; - dInput.eagleInputs->acceptedTokens = eagleInputs.acceptedTokens; - dInput.eagleInputs->acceptedLens = eagleInputs.acceptedLens; - dInput.eagleInputs->acceptedPathIds = eagleInputs.acceptedPaths; - dInput.eagleInputs->chunkedContextNextTokens = eagleInputs.chunkedContextNextTokens; - // Slots in request order - dInput.eagleInputs->seqSlots = fusedRuntimeBuffers.seqSlots; - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -} // namespace - -void MakeDecodingBatchInputOutput::operator()(DecoderInputBuffers& inputBuffers, - runtime::decoder::DecoderState& decoderState, runtime::ModelConfig const& modelConfig, - OptionalRef fusedRuntimeBuffers) const -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - auto [activeSlots, generationSteps] = getActiveSlots(inputBuffers.decoderRequests); - - createDecoderBatchInputs(inputBuffers, activeSlots, decoderState); - - auto const maxBeamWidth = decoderState.getMaxBeamWidth(); - if (maxBeamWidth > 1) - { - // For Variable-Beam-Width-Search - decoderState.getJointDecodingInput().generationSteps = generationSteps; - } - - if (modelConfig.getSpeculativeDecodingMode().hasDraftLogits()) - { - decoderState.getJointDecodingInput().medusaInputs->medusaLogits = inputBuffers.predictedDraftLogits; - } - - if (modelConfig.getSpeculativeDecodingMode().isExplicitDraftTokens()) - { - TLLM_CHECK(fusedRuntimeBuffers); - // requires mCtxGenFusion == true - setExplicitDraftTokensInputs(decoderState.getJointDecodingInput(), *fusedRuntimeBuffers); - } - else if (modelConfig.getSpeculativeDecodingMode().isEagle()) - { - TLLM_CHECK(fusedRuntimeBuffers); - // requires mCtxGenFusion == true - setEagleInputs(decoderState.getJointDecodingInput(), *fusedRuntimeBuffers); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -} // namespace tensorrt_llm::batch_manager diff --git a/cpp/tensorrt_llm/batch_manager/medusaBuffers.cpp b/cpp/tensorrt_llm/batch_manager/medusaBuffers.cpp index eb40208739cf..32935e683b83 100644 --- a/cpp/tensorrt_llm/batch_manager/medusaBuffers.cpp +++ b/cpp/tensorrt_llm/batch_manager/medusaBuffers.cpp @@ -17,99 +17,10 @@ #include "tensorrt_llm/batch_manager/medusaBuffers.h" #include "tensorrt_llm/runtime/bufferManager.h" -#include "tensorrt_llm/runtime/medusaModule.h" -#include "tensorrt_llm/runtime/utils/speculativeChoicesUtils.h" namespace tensorrt_llm::batch_manager { -MedusaBuffers::MedusaBuffers(SizeType32 maxBatchSize, SizeType32 maxBeamWidth, runtime::BufferManager const& manager, - runtime::ModelConfig const& modelConfig, runtime::WorldConfig const& worldConfig, - executor::DecodingConfig const& decodingConfig, runtime::TllmRuntime const& runtime) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - TLLM_CHECK_WITH_INFO(maxBeamWidth == 1, "Medusa does not support beam search"); - - auto const& engine = runtime.getEngine(); - - auto const maxNumSequences = maxBatchSize; - - auto const medusaModule = std::dynamic_pointer_cast( - modelConfig.getSpeculativeDecodingModulePtr()); - - auto const medusaHeads = medusaModule->getMaxDraftPathLen(); - auto const maxPathLen = medusaModule->getMaxPathLen(); // medusaHeads + 1 - auto const maxMedusaTokens = medusaModule->getMaxDecodingDraftTokens(); - auto const maxDecodingTokens = medusaModule->getMaxDecodingTokens(); // maxMedusaTokens + 1 - auto const numPackedMasks = medusaModule->getNumPackedMasks(); - - auto const vocabSizePadded = modelConfig.getVocabSizePadded(worldConfig.getSize()); - - if (worldConfig.isLastPipelineParallelRank()) - { - auto logitsType = engine.getTensorDataType("medusa_logits"); - medusaLogitsDevice = manager.gpu( - ITensor::makeShape({medusaHeads, maxBatchSize, maxDecodingTokens, vocabSizePadded}), logitsType); - } - - // Note: reserved for variable sequence length support. - medusaGenerationLengthsHost - = runtime::BufferManager::pinned(ITensor::makeShape({maxNumSequences}), nvinfer1::DataType::kINT32); - // TODO: pack batch and tokensPerStep into one dim to support variable sequence length without padddings. - attentionPackedMaskHost = runtime::BufferManager::pinned( - ITensor::makeShape({maxNumSequences, maxDecodingTokens, numPackedMasks}), nvinfer1::DataType::kINT32); - medusaPositionOffsetsHost = runtime::BufferManager::pinned( - ITensor::makeShape({maxNumSequences, maxDecodingTokens}), nvinfer1::DataType::kINT32); - medusaTreeIdsHost = runtime::BufferManager::pinned( - ITensor::makeShape({maxNumSequences, maxMedusaTokens}), nvinfer1::DataType::kINT32); - medusaPathsHost = runtime::BufferManager::pinned( - ITensor::makeShape({maxNumSequences, maxDecodingTokens, maxPathLen}), nvinfer1::DataType::kINT32); - - TensorPtr medusaPositionOffsetsHostSlice = ITensor::slice(medusaPositionOffsetsHost, 0, 1); - medusaPositionOffsetsHostSlice->squeeze(0); - TensorPtr medusaTreeIdsHostSlice = ITensor::slice(medusaTreeIdsHost, 0, 1); - medusaTreeIdsHostSlice->squeeze(0); - TensorPtr medusaPathsHostSlice = ITensor::slice(medusaPathsHost, 0, 1); - medusaPathsHostSlice->squeeze(0); - TensorPtr attentionPackedMaskHostSlice = ITensor::slice(attentionPackedMaskHost, 0, 1); - attentionPackedMaskHostSlice->squeeze(0); - - // Init buffers for 1 request - auto const& choices = decodingConfig.getMedusaChoices().value_or(medusaModule->getMedusaChoices()); - runtime::utils::initTensorsFromChoices(*medusaModule, choices, mTopKs, medusaGenerationLengthsHost, - medusaPositionOffsetsHostSlice, medusaTreeIdsHostSlice, medusaPathsHostSlice, attentionPackedMaskHostSlice); - - auto scatterToBatch = [maxBatchSize, &manager](TensorPtr& data) - { - auto srcSlice = ITensor::slice(data, 0, 1); - // Populate data from the 1st request to the other requests in the batch - for (SizeType32 bi = 1; bi < maxBatchSize; ++bi) - { - auto dstSlice = ITensor::slice(data, bi, 1); - manager.copy(*srcSlice, *dstSlice); - } - }; - - scatterToBatch(medusaPositionOffsetsHost); - scatterToBatch(medusaTreeIdsHost); - scatterToBatch(medusaPathsHost); - scatterToBatch(attentionPackedMaskHost); - - // Copy buffers to device - // 1st dimension of packed mask is num_total_generation_tokens now (packed without paddings). - attentionPackedMaskHost->reshape(ITensor::makeShape({maxNumSequences * maxDecodingTokens, numPackedMasks})); - attentionPackedMaskDevice = manager.copyFrom(*attentionPackedMaskHost, runtime::MemoryType::kGPU); - medusaGenerationLengthsDevice = manager.copyFrom(*medusaGenerationLengthsHost, runtime::MemoryType::kGPU); - medusaPositionOffsetsDevice = manager.copyFrom(*medusaPositionOffsetsHost, runtime::MemoryType::kGPU); - medusaTreeIdsDevice = manager.copyFrom(*medusaTreeIdsHost, runtime::MemoryType::kGPU); - medusaPathsDevice = manager.copyFrom(*medusaPathsHost, runtime::MemoryType::kGPU); - - // use speculative decoding buffer - medusaUseSpecDecoding = manager.cpu(ITensor::makeShape({1}), nvinfer1::DataType::kINT32); - runtime::bufferCast(*medusaUseSpecDecoding)[0] = 1; - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - void MedusaBuffers::reshape(SizeType32 /* numCtxSequences */, SizeType32 numGenSequences, SizeType32 tokensPerStep) { TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); diff --git a/cpp/tensorrt_llm/batch_manager/peftCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/peftCacheManager.cpp index 0bf9a989fd65..89cc475b82ee 100644 --- a/cpp/tensorrt_llm/batch_manager/peftCacheManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/peftCacheManager.cpp @@ -30,7 +30,7 @@ #include "tensorrt_llm/runtime/workerPool.h" #include "tensorrt_llm/runtime/worldConfig.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include #include @@ -52,7 +52,8 @@ PeftTaskNotCachedException::PeftTaskNotCachedException(std::string const& msg) PeftTaskNotCachedException::~PeftTaskNotCachedException() noexcept = default; std::pair PeftCacheManager::getMaxNumSlots(PeftCacheManagerConfig const& config, - nvinfer1::DataType dataType, uint64_t pageWidth, uint64_t max1dModSize, runtime::BufferManager const& bufferManager) + tensorrt_llm::DataType dataType, uint64_t pageWidth, uint64_t max1dModSize, + runtime::BufferManager const& bufferManager) { TLLM_LOG_DEBUG("max1dModeSize=%llu", max1dModSize); TLLM_LOG_DEBUG("pageWidth=%llu", pageWidth); diff --git a/cpp/tensorrt_llm/batch_manager/promptTuningBuffers.cpp b/cpp/tensorrt_llm/batch_manager/promptTuningBuffers.cpp deleted file mode 100644 index 1cf73a2c0d21..000000000000 --- a/cpp/tensorrt_llm/batch_manager/promptTuningBuffers.cpp +++ /dev/null @@ -1,323 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2025 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. - */ - -#include "tensorrt_llm/batch_manager/promptTuningBuffers.h" - -#include "tensorrt_llm/batch_manager/llmRequest.h" -#include "tensorrt_llm/common/nvtxUtils.h" - -namespace tensorrt_llm::batch_manager -{ -using SizeType32 = tensorrt_llm::runtime::SizeType32; -using TensorPtr = runtime::ITensor::SharedPtr; - -PromptTuningBuffers::PromptTuningBuffers(SizeType32 maxBatchSize, runtime::BufferManager const& manager, - runtime::ModelConfig const& modelConfig, runtime::WorldConfig const& worldConfig) -{ - auto maxPromptEmbeddingTableSize = modelConfig.getMaxPromptEmbeddingTableSize(); - auto const hiddenSize = modelConfig.getHiddenSize() * worldConfig.getTensorParallelism(); - - // vocabSize and mMaxPromptVocabSize - mPromptTuningParams.vocabSize = manager.gpu(runtime::ITensor::makeShape({1}), nvinfer1::DataType::kINT32); - mMaxPromptVocabSize = maxPromptEmbeddingTableSize / maxBatchSize; - - auto promptVocabSizeHost - = runtime::BufferManager::pinned(runtime::ITensor::makeShape({1}), nvinfer1::DataType::kINT32); - auto promptVocabSizeHostData = runtime::bufferCast(*promptVocabSizeHost); - promptVocabSizeHostData[0] = mMaxPromptVocabSize; - manager.copy(*promptVocabSizeHost, *mPromptTuningParams.vocabSize); - - // embeddingTable - mPromptTuningParams.embeddingTable = manager.gpu( - runtime::ITensor::makeShape({maxPromptEmbeddingTableSize, hiddenSize}), modelConfig.getDataType()); - - // tasks - mPromptTuningParams.tasks = manager.emptyTensor(runtime::MemoryType::kGPU, nvinfer1::DataType::kINT32); -} - -PromptTuningBuffers::PromptTuningBuffers(SizeType32 maxBatchSize, runtime::BufferManager const& manager, - runtime::ModelConfig const& modelConfig, runtime::WorldConfig const& worldConfig, bool promptTableOffloading) -{ - auto maxPromptEmbeddingTableSize = modelConfig.getMaxPromptEmbeddingTableSize(); - auto const hiddenSize = modelConfig.getHiddenSize() * worldConfig.getTensorParallelism(); - - // vocabSize and mMaxPromptVocabSize - mPromptTuningParams.vocabSize = manager.gpu(runtime::ITensor::makeShape({1}), nvinfer1::DataType::kINT32); - mMaxPromptVocabSize = maxPromptEmbeddingTableSize / maxBatchSize; - mPromptTableOffloading = promptTableOffloading; - - auto promptVocabSizeHost - = runtime::BufferManager::pinned(runtime::ITensor::makeShape({1}), nvinfer1::DataType::kINT32); - auto promptVocabSizeHostData = runtime::bufferCast(*promptVocabSizeHost); - promptVocabSizeHostData[0] = mMaxPromptVocabSize; - manager.copy(*promptVocabSizeHost, *mPromptTuningParams.vocabSize); - - // embeddingTable - mPromptTuningParams.embeddingTable = manager.gpu( - runtime::ITensor::makeShape({maxPromptEmbeddingTableSize, hiddenSize}), modelConfig.getDataType()); - - // tasks - mPromptTuningParams.tasks = manager.emptyTensor(runtime::MemoryType::kGPU, nvinfer1::DataType::kINT32); -} - -void PromptTuningBuffers::validate( - std::optional const& optReqPromptEmbeddingTable, std::optional const& optReqPromptVocabSize) -{ - // Need to copy request embeddingTable to promptEmbeddingTable - if (optReqPromptEmbeddingTable.has_value()) - { - - auto reqPromptEmbeddingTable = optReqPromptEmbeddingTable.value(); - auto reqPromptVocabSize = optReqPromptVocabSize.value(); - - if (reqPromptVocabSize > mMaxPromptVocabSize) - { - std::string errStr = "Prompt vocab size" + std::to_string(reqPromptVocabSize) - + " is larger than max prompt vocab size of " + std::to_string(mMaxPromptVocabSize) - + ". Max prompt vocab size is computed from max_prompt_embedding_table_size / max_batch_size. "; - TLLM_LOG_ERROR(errStr); - throw std::runtime_error(errStr); - } - else - { - // Check that type matches model weights - if (reqPromptEmbeddingTable->getDataType() != mPromptTuningParams.embeddingTable->getDataType()) - { - std::string errStr = "Request embedding table data type doesn't match model weight data type."; - TLLM_LOG_ERROR(errStr); - throw std::runtime_error(errStr); - } - - if (reqPromptEmbeddingTable->getShape().d[1] != reqPromptVocabSize) - { - std::string errStr - = "First dimension of request embedding table is expected to be equal to prompt vocab size"; - TLLM_LOG_ERROR(errStr); - throw std::runtime_error(errStr); - } - } - } -} - -void PromptTuningBuffers::fill(RequestVector const& contextRequests, RequestVector const& genRequests, - runtime::BufferManager const& manager, bool packed) -{ - NVTX3_SCOPED_RANGE_WITH_NAME(range, "PromptTuningBuffers::fill"); - - auto const numContextRequests = static_cast(contextRequests.size()); - - std::vector reqBeamWidths; - std::vector reqPromptLengths; - mPromptTuningParams.promptTuningEnabled.clear(); - - SizeType32 batchIdx{0}; - for (auto const& requests : {contextRequests, genRequests}) - { - for (auto const& llmReq : requests) - { - reqBeamWidths.push_back(llmReq->mSamplingConfig.beamWidth); - if (batchIdx < numContextRequests) - { - SizeType32 numContextTokens = 0; - auto const draftLength = llmReq->isLastContextChunk() ? llmReq->getNumDraftTokens() : 0; - auto const contextChunkSize = llmReq->getContextChunkSize(); - numContextTokens += contextChunkSize + draftLength; - reqPromptLengths.push_back(numContextTokens); - } - - std::optional optReqPromptEmbeddingTable = std::nullopt; - std::optional optReqPromptVocabSize = std::nullopt; - - if (mPromptTableOffloading) - { - optReqPromptEmbeddingTable = getChunkPtableBuffer(getChunkPtableCurrentIndex()); - optReqPromptVocabSize = getChunkPtableBufferSliceSize(getChunkPtableCurrentIndex(), batchIdx); - } - else - { - optReqPromptEmbeddingTable = llmReq->getPromptEmbeddingTable(); - optReqPromptVocabSize = llmReq->getPromptVocabSize(); - } - - mPromptTuningParams.promptTuningEnabled.push_back(optReqPromptEmbeddingTable.has_value()); - - // If context request & has embedding table, validate it - if (optReqPromptEmbeddingTable.has_value()) - { - // If a context request, validate prompt tensors and move to GPU - if (batchIdx < numContextRequests) - { - if (mPromptTableOffloading) - { - // Need to slice the ptable since we don't need the entire buffer - // The size depends on optReqPromptVocabSize which stores how many fake prompts are in the chunk - auto slicedPtable = runtime::ITensor::slice( - optReqPromptEmbeddingTable.value(), 0, optReqPromptVocabSize.value()); - slicedPtable->unsqueeze(0); - optReqPromptEmbeddingTable = std::move(slicedPtable); - } - else - { - // Move to GPU - llmReq->movePromptEmbeddingTableToGpu(manager); - optReqPromptEmbeddingTable = llmReq->getPromptEmbeddingTable(); - } - - // Validate the table, prompt_vocab_size - validate(optReqPromptEmbeddingTable, optReqPromptVocabSize); - } - - auto const reqPromptEmbeddingTable = optReqPromptEmbeddingTable.value(); - auto const reqPromptVocabSize = optReqPromptVocabSize.value(); - - // TODO: Use invokeCopyBatch to avoid multiple bs1 copies - // Copy into large prompt embedding table - TensorPtr reqPromptEmbeddingTableView = runtime::ITensor::view(reqPromptEmbeddingTable); - reqPromptEmbeddingTableView->squeeze(0); - auto const promptEmbeddingTableSlice = runtime::ITensor::slice( - mPromptTuningParams.embeddingTable, batchIdx * mMaxPromptVocabSize, reqPromptVocabSize); - manager.copy(*reqPromptEmbeddingTable, *promptEmbeddingTableSlice); - // TODO: src: 2007040 (llmReq->getPromptEmbeddingTable()) != dst: 1003520 (reqPromptVocabSize) - // (original shape passed from - // python == 196 * 5120, fp16) - // VILA mode 1 , 2 images in one request - } - ++batchIdx; - } - } - - auto const batchSize = batchIdx; - std::vector tasksHostVec(batchSize); - std::iota(tasksHostVec.begin(), tasksHostVec.end(), 0); - - // Create a tensor that wraps the vector and convert unique_ptr to shared_ptr - auto tasksHost = std::shared_ptr( - runtime::ITensor::wrap(tasksHostVec, runtime::ITensor::makeShape({batchSize})).release()); - - mPromptTuningParams.fillTasksTensor( - tasksHost, batchSize, numContextRequests, reqBeamWidths, reqPromptLengths, manager, packed); -} - -void PromptTuningBuffers::initializeChunkPtableBuffers(runtime::BufferManager const& manager, - runtime::ModelConfig const& modelConfig, SizeType32 contextChunkSize, std::shared_ptr const& llmReq) -{ - if (mChunkPtableInitialized) - { - return; - } - - std::array buffers; - std::vector> startPositions(2); - for (int i = 0; i < 2; i++) - { - startPositions[i].emplace_back(0); - auto memType = llmReq->getPromptEmbeddingTable().value()->getDataType(); - buffers[i] = manager.gpu(runtime::ITensor::makeShape({contextChunkSize, modelConfig.getHiddenSize()}), memType); - } - - mChunkPtableBuffers = std::move(buffers); - mChunkPtableBufferStartPositions = std::move(startPositions); - - mChunkPtableCurrentIndex = 0; - mChunkPtableInitialized = true; -} - -void PromptTuningBuffers::switchChunkPtableBuffer() -{ - mChunkPtableCurrentIndex = 1 - mChunkPtableCurrentIndex; - clearBufferStartPositions(mChunkPtableCurrentIndex); -} - -size_t PromptTuningBuffers::getChunkPtableCurrentIndex() -{ - return mChunkPtableCurrentIndex; -} - -TensorPtr& PromptTuningBuffers::getChunkPtableBuffer(size_t index) -{ - if (!mChunkPtableBuffers.has_value()) - { - TLLM_THROW("Chunk ptable buffers not initialized"); - } - if (!mChunkPtableBuffers.value()[index]) - { - TLLM_THROW("Chunk ptable buffer at index %zu is null", index); - } - return mChunkPtableBuffers.value()[index]; -} - -SizeType32 PromptTuningBuffers::getChunkPtableBufferSliceSize(size_t index, size_t batchIdx) -{ - if (!mChunkPtableBufferStartPositions.has_value()) - { - return 0; - } - - if (batchIdx + 1 >= mChunkPtableBufferStartPositions.value()[index].size()) - { - TLLM_THROW("Batch index %zu + 1 out of bounds for buffer %zu (size: %zu)", batchIdx, index, - mChunkPtableBufferStartPositions.value()[index].size()); - } - - return mChunkPtableBufferStartPositions.value()[index][batchIdx + 1] - - mChunkPtableBufferStartPositions.value()[index][batchIdx]; -} - -SizeType32 PromptTuningBuffers::getChunkPtableBufferStartPosition(size_t index, size_t batchIdx) -{ - if (!mChunkPtableBufferStartPositions.has_value()) - { - return 0; - } - - if (batchIdx >= mChunkPtableBufferStartPositions.value()[index].size()) - { - TLLM_THROW("Batch index %zu out of bounds for buffer %zu (size: %zu)", batchIdx, index, - mChunkPtableBufferStartPositions.value()[index].size()); - } - - // For first batch, return the value directly - if (batchIdx == 0) - { - return mChunkPtableBufferStartPositions.value()[index][0]; - } - - // For other batches, return difference from previous position - return mChunkPtableBufferStartPositions.value()[index][batchIdx] - - mChunkPtableBufferStartPositions.value()[index][batchIdx - 1]; -} - -void PromptTuningBuffers::updateBufferStartPosition(size_t index, SizeType32 numRows) -{ - if (!mChunkPtableBufferStartPositions.has_value()) - { - return; - } - auto& positions = mChunkPtableBufferStartPositions.value()[index]; - positions.push_back(positions.back() + numRows); -} - -void PromptTuningBuffers::clearBufferStartPositions(size_t index) -{ - if (mChunkPtableBufferStartPositions.has_value()) - { - mChunkPtableBufferStartPositions.value()[index].clear(); - mChunkPtableBufferStartPositions.value()[index].emplace_back(0); - } -} - -} // namespace tensorrt_llm::batch_manager diff --git a/cpp/tensorrt_llm/batch_manager/rnnCacheTransBuffer.cpp b/cpp/tensorrt_llm/batch_manager/rnnCacheTransBuffer.cpp index 37af8e31baf9..f9c04200e8d2 100644 --- a/cpp/tensorrt_llm/batch_manager/rnnCacheTransBuffer.cpp +++ b/cpp/tensorrt_llm/batch_manager/rnnCacheTransBuffer.cpp @@ -21,6 +21,7 @@ #include "tensorrt_llm/common/dataType.h" #include "tensorrt_llm/common/envUtils.h" #include "tensorrt_llm/common/logger.h" +#include "tensorrt_llm/common/tllmDataType.h" #include @@ -95,7 +96,7 @@ size_t RnnCacheTransBufferManager::computeTransferBufferSizeFromPool( RnnCacheTransBufferManager::RnnCacheTransBufferManager(kv_cache_manager::BaseKVCacheManager* kvCacheManager, executor::kv_cache::CacheState const& cacheState, std::optional maxNumTokens) : BaseTransBufferManager(computeTransferBufferSizeFromPool(kvCacheManager, cacheState, maxNumTokens), - nvinfer1::DataType::kUINT8, maxNumTokens) + tensorrt_llm::DataType::kUINT8, maxNumTokens) { TLLM_CHECK(kvCacheManager != nullptr); TLLM_LOG_INFO("RnnCacheTransBufferManager created for unified pool RNN cache"); diff --git a/cpp/tensorrt_llm/batch_manager/rnnStateBuffers.cpp b/cpp/tensorrt_llm/batch_manager/rnnStateBuffers.cpp deleted file mode 100644 index 6fc7977ef8f1..000000000000 --- a/cpp/tensorrt_llm/batch_manager/rnnStateBuffers.cpp +++ /dev/null @@ -1,78 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2025 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. - */ - -#include "rnnStateBuffers.h" - -#include "tensorrt_llm/batch_manager/llmRequest.h" -#include "tensorrt_llm/batch_manager/rnnStateManager.h" -#include "tensorrt_llm/batch_manager/runtimeBuffers.h" -#include "tensorrt_llm/common/nvtxUtils.h" -#include "tensorrt_llm/runtime/tllmRuntime.h" - -using namespace tensorrt_llm::runtime; - -namespace tensorrt_llm::batch_manager -{ - -RnnStateBuffers::RnnStateBuffers(SizeType32 maxBatchSize, runtime::TllmRuntime const& runtime) -{ - auto const& manager = runtime.getBufferManager(); - - slotMappingHost = BufferManager::cpu(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kINT32); - slotMappingDevice = manager.gpu(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kINT32); -} - -void RnnStateBuffers::reshape(SizeType32 numSequences) -{ - slotMappingHost->reshape(ITensor::makeShape({numSequences})); - slotMappingDevice->reshape(ITensor::makeShape({numSequences})); -} - -void RnnStateBuffers::fillSlotMappings( - RequestVector const& contextRequests, rnn_state_manager::RnnStateManager* rnnStateManager) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(rnnStateBuffersFillSlotMappings); - - SizeType32 batchIdx{0}; - for (auto const& llmReq : contextRequests) - { - auto const seqSlot = llmReq->mSeqSlot.value(); - auto const reqBeamWidth = llmReq->mSamplingConfig.beamWidth; - rnnStateManager->fillSlotMapping(*slotMappingHost, batchIdx, seqSlot, reqBeamWidth); - ++batchIdx; - } - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void RnnStateBuffers::copySlotMappingH2D(runtime::TllmRuntime const& runtime) -{ - auto const& manager = runtime.getBufferManager(); - manager.copy(*slotMappingHost, *slotMappingDevice); -} - -void RnnStateBuffers::getBuffers(TensorMap& inputBuffers) const -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(rnnStateBuffersGetBuffers); - - inputBuffers.insert_or_assign("slot_mapping", slotMappingDevice); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -} // namespace tensorrt_llm::batch_manager diff --git a/cpp/tensorrt_llm/batch_manager/rnnStateBuffers.h b/cpp/tensorrt_llm/batch_manager/rnnStateBuffers.h deleted file mode 100644 index e25df47382a1..000000000000 --- a/cpp/tensorrt_llm/batch_manager/rnnStateBuffers.h +++ /dev/null @@ -1,58 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2025 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. - */ - -#pragma once - -#include "tensorrt_llm/batch_manager/common.h" -#include "tensorrt_llm/runtime/iTensor.h" - -namespace tensorrt_llm::runtime -{ -class TllmRuntime; -} // namespace tensorrt_llm::runtime - -namespace tensorrt_llm::batch_manager -{ - -namespace rnn_state_manager -{ -class RnnStateManager; -} - -class RnnStateBuffers -{ -public: - using SizeType32 = tensorrt_llm::runtime::SizeType32; - using TensorPtr = runtime::ITensor::SharedPtr; - using TensorMap = runtime::StringPtrMap; - - // others should be in rnnStateManager, we only need slotMapping here. - TensorPtr slotMappingHost; // [batch_size] - TensorPtr slotMappingDevice; // [batch_size] - - RnnStateBuffers(SizeType32 maxBatchSize, runtime::TllmRuntime const& runtime); - - void reshape(SizeType32 numSequences); - - void fillSlotMappings(RequestVector const& contextRequests, rnn_state_manager::RnnStateManager* rnnStateManager); - - void copySlotMappingH2D(runtime::TllmRuntime const& runtime); - - void getBuffers(TensorMap& inputBuffers) const; -}; - -} // namespace tensorrt_llm::batch_manager diff --git a/cpp/tensorrt_llm/batch_manager/rnnStateManager.cpp b/cpp/tensorrt_llm/batch_manager/rnnStateManager.cpp index 7608079fb396..7d032a268fdd 100644 --- a/cpp/tensorrt_llm/batch_manager/rnnStateManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/rnnStateManager.cpp @@ -17,6 +17,7 @@ #include "tensorrt_llm/batch_manager/rnnStateManager.h" #include "tensorrt_llm/common/assert.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/runtime/cudaStream.h" #include "tensorrt_llm/runtime/utils/runtimeUtils.h" @@ -80,7 +81,7 @@ RnnStateManager::RnnStateManager(SizeType32 maxNumSequences, tensorrt_llm::runti {localNbLayers, mMaxNumSequences * mBeamSlotsPerSequence, convKernel - 1, rnnConvDimSize}); mDtype = dataType; - mSsmCacheDtype = nvinfer1::DataType::kFLOAT; + mSsmCacheDtype = tensorrt_llm::DataType::kFLOAT; // Store RNN model config for CacheTransceiver mDState = stateSize; @@ -117,7 +118,7 @@ RnnStateManager::RnnStateManager(SizeType32 maxNumSequences, tensorrt_llm::runti RnnStateManager::RnnStateManager(SizeType32 dState, SizeType32 dConv, SizeType32 numHeads, SizeType32 nGroups, SizeType32 headDim, SizeType32 maxBatchSize, WorldConfig const& worldConfig, int64_t stream, - nvinfer1::DataType dtype, nvinfer1::DataType ssmCacheDtype, std::vector const& ppLayers, + tensorrt_llm::DataType dtype, tensorrt_llm::DataType ssmCacheDtype, std::vector const& ppLayers, SizeType32 numLayers) : mMaxNumSequences(maxBatchSize) , mMaxBeamWidth{1} @@ -297,12 +298,12 @@ RnnStateManager::TensorPtr RnnStateManager::getSsmStates() const return pagedRnnStates; } -nvinfer1::DataType RnnStateManager::getConvStateDataType() const noexcept +tensorrt_llm::DataType RnnStateManager::getConvStateDataType() const noexcept { return mDtype; } -nvinfer1::DataType RnnStateManager::getSsmStateDataType() const noexcept +tensorrt_llm::DataType RnnStateManager::getSsmStateDataType() const noexcept { return mSsmCacheDtype; } diff --git a/cpp/tensorrt_llm/batch_manager/runtimeBuffers.cpp b/cpp/tensorrt_llm/batch_manager/runtimeBuffers.cpp deleted file mode 100644 index ea5b9b06a96e..000000000000 --- a/cpp/tensorrt_llm/batch_manager/runtimeBuffers.cpp +++ /dev/null @@ -1,1029 +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. - */ - -#include "tensorrt_llm/batch_manager/runtimeBuffers.h" - -#include "tensorrt_llm/batch_manager/encoderBuffers.h" -#include "tensorrt_llm/batch_manager/kvCacheManager.h" -#include "tensorrt_llm/batch_manager/loraBuffers.h" -#include "tensorrt_llm/batch_manager/medusaBuffers.h" -#include "tensorrt_llm/batch_manager/promptTuningBuffers.h" -#include "tensorrt_llm/batch_manager/rnnStateBuffers.h" -#include "tensorrt_llm/batch_manager/rnnStateManager.h" -#include "tensorrt_llm/batch_manager/transformerBuffers.h" -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/logger.h" -#include "tensorrt_llm/common/nvtxUtils.h" -#include "tensorrt_llm/common/stlUtils.h" -#include "tensorrt_llm/runtime/bufferManager.h" -#include "tensorrt_llm/runtime/common.h" -#include "tensorrt_llm/runtime/decoderState.h" -#include "tensorrt_llm/runtime/iBuffer.h" -#include "tensorrt_llm/runtime/iTensor.h" -#include "tensorrt_llm/runtime/runtimeKernels.h" -#include "tensorrt_llm/runtime/tllmRuntime.h" - -#include -#include -#include -#include -#include - -using namespace tensorrt_llm::runtime; - -namespace tensorrt_llm::batch_manager -{ - -RuntimeBuffers::RuntimeBuffers(SizeType32 maxBatchSize, SizeType32 maxBeamWidth, - std::vector const& maxAttentionWindowVec, SizeType32 maxAttentionWindow, SizeType32 sinkTokenLen, - TllmRuntime const& runtime, ModelConfig const& modelConfig, WorldConfig const& worldConfig, - executor::DecodingConfig const& decodingConfig, bool gatherGenerationLogits, std::optional maxNumTokens, - std::optional> const& additionalModelOutputs, - bool promptTableOffloadingParam) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - promptTableOffloading = promptTableOffloadingParam; - - create(maxBatchSize, maxBeamWidth, maxAttentionWindowVec, maxAttentionWindow, sinkTokenLen, runtime, modelConfig, - worldConfig, decodingConfig, gatherGenerationLogits, additionalModelOutputs); - - // pre-allocate - setMaxBufferSizes(maxBatchSize, maxBeamWidth, modelConfig, maxNumTokens); - reshape(runtime, modelConfig, worldConfig, gatherGenerationLogits); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -RuntimeBuffers::~RuntimeBuffers() = default; - -void RuntimeBuffers::create(SizeType32 maxBatchSize, SizeType32 maxBeamWidth, - std::vector const& maxAttentionWindowVec, SizeType32 maxAttentionWindow, SizeType32 sinkTokenLen, - TllmRuntime const& runtime, ModelConfig const& modelConfig, WorldConfig const& worldConfig, - executor::DecodingConfig const& decodingConfig, bool gatherGenerationLogits, - std::optional> const& additionalModelOutputs) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - auto const& manager = runtime.getBufferManager(); - auto const& engine = runtime.getEngine(); - - if (modelConfig.isTransformerBased()) - { - transformerBuffers = std::make_unique(maxBatchSize, maxBeamWidth, maxAttentionWindowVec, - maxAttentionWindow, sinkTokenLen, runtime, modelConfig, worldConfig); - } - if (modelConfig.isRnnBased()) - { - rnnStateBuffers = std::make_unique(maxBatchSize, runtime); - } - - auto constexpr nvTokenIdType = TRTDataType::value; - inputsIds = manager.emptyTensor(MemoryType::kGPU, nvTokenIdType); - - mropeRotaryCosSin = manager.emptyTensor(MemoryType::kGPU, nvinfer1::DataType::kFLOAT); - mropePositionDeltas = manager.emptyTensor(MemoryType::kGPU, nvinfer1::DataType::kINT32); - - if (worldConfig.isLastPipelineParallelRank()) - { - auto const logitsType = engine.getTensorDataType(batch_manager::RuntimeBuffers::kLogitsTensorName); - logits = manager.emptyTensor(MemoryType::kGPU, logitsType); - } - - // TODO: check which tensors can be allocated as pinned for max size - requestTypes = manager.emptyTensor(MemoryType::kCPU, TRTDataType::value); - - contextLengthsHost = manager.emptyTensor(MemoryType::kCPU, nvinfer1::DataType::kINT32); - contextLengthsDevice = manager.emptyTensor(MemoryType::kGPU, nvinfer1::DataType::kINT32); - sequenceLengthsHost = manager.emptyTensor(MemoryType::kCPU, nvinfer1::DataType::kINT32); - sequenceLengthsDevice = manager.emptyTensor(MemoryType::kGPU, nvinfer1::DataType::kINT32); - - lastTokenIdsHost = manager.emptyTensor(MemoryType::kCPU, nvinfer1::DataType::kINT32); - lastTokenIdsDevice = manager.emptyTensor(MemoryType::kGPU, nvinfer1::DataType::kINT32); - logitsIdsHost = manager.emptyTensor(MemoryType::kCPU, nvinfer1::DataType::kINT32); - - inputsIds = manager.emptyTensor(MemoryType::kGPU, nvinfer1::DataType::kINT32); - - if (worldConfig.isPipelineParallel()) - { - hiddenStates = manager.emptyTensor(MemoryType::kGPU, modelConfig.getDataType()); - } - - auto const maxBatchSizeShape = ITensor::makeShape({maxBatchSize}); - seqSlots = tensorrt_llm::runtime::BufferManager::pinnedPool(maxBatchSizeShape, nvinfer1::DataType::kINT32); - seqSlotsDevice = manager.gpu(maxBatchSizeShape, nvinfer1::DataType::kINT32); - - cacheIndirDecoderIOBatchedCopySrcOffsets - = tensorrt_llm::runtime::BufferManager::pinnedPool(maxBatchSizeShape, nvinfer1::DataType::kINT64); - cacheIndirDecoderIOBatchedCopyDstOffsets - = tensorrt_llm::runtime::BufferManager::pinnedPool(maxBatchSizeShape, nvinfer1::DataType::kINT64); - cacheIndirDecoderIOBatchedCopySizes - = tensorrt_llm::runtime::BufferManager::pinnedPool(maxBatchSizeShape, nvinfer1::DataType::kINT64); - mCacheIndirDecoderIOBatchedCopySrcOffsetsSliceDevice = manager.gpu(maxBatchSizeShape, nvinfer1::DataType::kINT64); - mCacheIndirDecoderIOBatchedCopyDstOffsetsSliceDevice = manager.gpu(maxBatchSizeShape, nvinfer1::DataType::kINT64); - mCacheIndirDecoderIOBatchedCopyCopySizesDevice = manager.gpu(maxBatchSizeShape, nvinfer1::DataType::kINT64); - - // Pre-allocate buffer for saving generation logits for model w/o draft tokens - if (gatherGenerationLogits - && (modelConfig.getSpeculativeDecodingMode().isDraftTokensExternal() - || modelConfig.getSpeculativeDecodingMode().isNone()) - && worldConfig.isLastPipelineParallelRank()) - { - auto const vocabSizePadded = modelConfig.getVocabSizePadded(worldConfig.getSize()); - auto const logitsType = engine.getTensorDataType(batch_manager::RuntimeBuffers::kLogitsTensorName); - - generationLogitsCache.transposedLogits = manager.gpu( - ITensor::makeShape({maxBeamWidth, GenerationLogitsCache::kCACHE_LENGTH, vocabSizePadded}), logitsType); - generationLogitsCache.logits = manager.gpu( - ITensor::makeShape({GenerationLogitsCache::kCACHE_LENGTH, maxBatchSize * maxBeamWidth, vocabSizePadded}), - logitsType); - - generationLogitsCache.fragmentPointerDevice = manager.gpu( - ITensor::makeShape({maxBatchSize, GenerationLogitsCache::kCACHE_LENGTH}), nvinfer1::DataType::kINT64); - generationLogitsCache.fragmentPointerHost = tensorrt_llm::runtime::BufferManager::pinnedPool( - ITensor::makeShape({maxBatchSize, GenerationLogitsCache::kCACHE_LENGTH}), nvinfer1::DataType::kINT64); - } - - if (modelConfig.useCrossAttention()) - { - encoderBuffers = std::make_unique(); - encoderBuffers->create(maxBatchSize, modelConfig, runtime); - } - - if (modelConfig.usePromptTuning()) - { - promptTuningBuffers = std::make_unique( - maxBatchSize, manager, modelConfig, worldConfig, promptTableOffloading); - } - - if (modelConfig.useLoraPlugin()) - { - loraBuffers = std::make_unique(maxBatchSize, maxBeamWidth, runtime, modelConfig, worldConfig); - } - - if (modelConfig.getSpeculativeDecodingMode().isMedusa()) - { - mMedusaBuffers = std::make_unique( - maxBatchSize, maxBeamWidth, manager, modelConfig, worldConfig, decodingConfig, runtime); - } - else if (modelConfig.getSpeculativeDecodingMode().isLookaheadDecoding()) - { - mLookaheadBuffers = std::make_unique( - maxBatchSize, maxBeamWidth, manager, modelConfig, worldConfig, decodingConfig, runtime); - } - else if (modelConfig.getSpeculativeDecodingMode().isExplicitDraftTokens()) - { - mExplicitDraftTokensBuffers = std::make_unique( - maxBatchSize, maxBeamWidth, manager, modelConfig, worldConfig); - } - else if (modelConfig.getSpeculativeDecodingMode().isEagle()) - { - mEagleBuffers = std::make_unique( - maxBatchSize, maxBeamWidth, manager, modelConfig, worldConfig, decodingConfig); - } - - if (modelConfig.useLanguageAdapter()) - { - languageAdapterRoutings = manager.emptyTensor(MemoryType::kGPU, TRTDataType::value); - } - - for (auto const& output : additionalModelOutputs.value_or(std::vector{})) - { - auto const& engine = runtime.getEngine(); - auto const dataType = engine.getTensorDataType(output.name.c_str()); - mAdditionalOutputTensors.emplace(output.name, manager.emptyTensor(runtime::MemoryType::kGPU, dataType)); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void RuntimeBuffers::setMaxBufferSizes(SizeType32 maxBatchSize, SizeType32 maxBeamWidth, - runtime::ModelConfig const& modelConfig, std::optional maxNumRuntimeTokens) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - // `maxNumSequences` is reached when all requests are in generation - numContextRequests = 0; - numGenRequests = maxBatchSize; - numGenSequences = maxBatchSize * maxBeamWidth; - - auto const maxDraftTokens = modelConfig.getMaxDecodingDraftTokens(); - // Draft-Tokens and Beam-Search are mutually exclusive - numLogits = maxBatchSize * std::max(1 + maxDraftTokens, maxBeamWidth); - auto const maxNumModelTokens = modelConfig.getMaxNumTokens(); - auto const maxNumContextTokens = maxBatchSize * modelConfig.getMaxInputLen(); - auto const maxNumGenTokens = numLogits; - // For pre-allocation - numContextTokens = 0; // Set in `setBufferSizes` rather than here for `computeContextLogits` - numGenTokens - = maxNumRuntimeTokens.value_or(maxNumModelTokens.value_or(std::max(maxNumContextTokens, maxNumGenTokens))); - - if (modelConfig.useCrossAttention()) - { - encoderBuffers->setMaxBufferSizes(maxBatchSize, modelConfig); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void RuntimeBuffers::setBufferSizes(RequestVector const& contextRequests, RequestVector const& genRequests) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(runtimeBuffersSetBufferSizes); - - // set context sizes - numContextRequests = static_cast(contextRequests.size()); - auto numContextLogits = numContextRequests; - numContextTokens = 0; - maxContextLength = 0; - for (auto const& llmReq : contextRequests) - { - auto const draftLength = llmReq->isLastContextChunk() ? llmReq->getNumDraftTokens() : 0; - numContextLogits += draftLength; - - auto const contextChunkSize = llmReq->getContextChunkSize(); - numContextTokens += contextChunkSize + draftLength; - if (maxContextLength < llmReq->mPromptLen) - { - maxContextLength = llmReq->mPromptLen; - } - } - - // set generation sizes - numGenRequests = static_cast(genRequests.size()); - numGenSequences = 0; - numGenTokens = 0; - for (auto const& llmReq : genRequests) - { - auto const reqBeamWidth = llmReq->getBeamWidthByIter(); - numGenSequences += reqBeamWidth; - auto const draftLen = llmReq->getNumDraftTokens(); - numGenTokens += draftLen + reqBeamWidth; - } - - numLogits = numContextLogits + numGenTokens; - - if (encoderBuffers) - { - encoderBuffers->setBufferSizes(contextRequests, genRequests); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void RuntimeBuffers::reshape(TllmRuntime const& runtime, ModelConfig const& modelConfig, WorldConfig const& worldConfig, - bool gatherGenerationLogits) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(runtimeBuffersReshape); - - if (worldConfig.isLastPipelineParallelRank()) - { - auto const vocabSizePadded = modelConfig.getVocabSizePadded(worldConfig.getSize()); - - if (modelConfig.computeContextLogits() && (numContextRequests > 0)) - { - // Only when need to return context logits, and there are new requests will execute context phase, - // logits buffer need to be re-allocated with size of [numContextTokens + numGenSequences, vocabSizePadded] - auto const& engine = runtime.getEngine(); - auto const& manager = runtime.getBufferManager(); - auto const logitsType = engine.getTensorDataType(kLogitsTensorName); - logits = manager.gpu(ITensor::makeShape({numContextTokens + numGenSequences, vocabSizePadded}), logitsType); - } - else if (gatherGenerationLogits && modelConfig.getSpeculativeDecodingMode().isNone()) - { - // If need to return generation logits, re-point the logit buffer to avoid overwrite, - // so we could write back GenerationLogitsCache::kCACHE_LENGTH steps' logits together - // logits shape: [1, maxBatchSize * maxBeamWidth, vocabSizePadded] - // which is large enough to cover both numContextRequests and numGenSequences - logits = ITensor::slice(generationLogitsCache.logits, generationLogitsCache.offset, 1); - generationLogitsCache.offset = (generationLogitsCache.offset + 1) % GenerationLogitsCache::kCACHE_LENGTH; - logits->squeeze(0); - } - else - { - logits->reshape(ITensor::makeShape({numLogits, vocabSizePadded})); - } - } - - auto const numSequences = getNumSequences(); - auto const numSequencesShape = ITensor::makeShape({numSequences}); - requestTypes->reshape(numSequencesShape); - contextLengthsHost->reshape(numSequencesShape); - contextLengthsDevice->reshape(numSequencesShape); - sequenceLengthsHost->reshape(numSequencesShape); - sequenceLengthsDevice->reshape(numSequencesShape); - - auto const numLogitsShape = ITensor::makeShape({numLogits}); - lastTokenIdsHost->reshape(numLogitsShape); - lastTokenIdsDevice->reshape(numLogitsShape); - logitsIdsHost->reshape(numLogitsShape); - - if (transformerBuffers) - { - transformerBuffers->reshape(numSequences, numContextTokens + numGenTokens); - } - - if (rnnStateBuffers) - { - rnnStateBuffers->reshape(numSequences); - } - - if (modelConfig.useCrossAttention()) - { - encoderBuffers->reshape(); - } - - if (modelConfig.useLoraPlugin()) - { - loraBuffers->reshape(numSequences); - } - - if (mMedusaBuffers) - { - mMedusaBuffers->reshape( - numContextRequests, numGenRequests, modelConfig.getSpeculativeDecodingModulePtr()->getMaxDecodingTokens()); - } - - if (mLookaheadBuffers && modelConfig.getSpeculativeDecodingMode().isLookaheadDecoding()) - { - mLookaheadBuffers->reshape( - numContextRequests, numGenRequests, modelConfig.getSpeculativeDecodingModulePtr()->getMaxDecodingTokens()); - } - - if (mExplicitDraftTokensBuffers) - { - mExplicitDraftTokensBuffers->reshape(numContextRequests, numGenRequests, modelConfig); - } - - if (mEagleBuffers) - { - mEagleBuffers->reshape(numContextRequests, numGenRequests, modelConfig); - } - - auto const numRequests = getNumRequests(); - auto const numRequestsShape = ITensor::makeShape({numRequests}); - seqSlots->reshape(numRequestsShape); - seqSlotsDevice->reshape(numRequestsShape); - - auto const numTokens = getNumTokens(); - inputsIds->reshape(ITensor::makeShape({numTokens})); - - if (modelConfig.useMrope()) - { - auto const mropeRotaryCosSinSize = modelConfig.getMaxPositionEmbeddings() * modelConfig.getRotaryEmbeddingDim(); - mropeRotaryCosSin->reshape(ITensor::makeShape({numSequences, mropeRotaryCosSinSize})); - mropePositionDeltas->reshape(ITensor::makeShape({numSequences, 1})); - } - - if (worldConfig.isPipelineParallel()) - { - auto const hiddenSize = (!modelConfig.getPpReduceScatter() || worldConfig.isFirstPipelineParallelRank()) - ? modelConfig.getHiddenSize() * worldConfig.getTensorParallelism() - : modelConfig.getHiddenSize(); - - auto const hiddenStatesShape = ITensor::makeShape({numTokens, hiddenSize}); - hiddenStates->reshape(hiddenStatesShape); - } - - if (modelConfig.useLanguageAdapter()) - { - languageAdapterRoutings->reshape(ITensor::makeShape({numTokens, 1})); - } - - for (auto const& outputTensor : mAdditionalOutputTensors) - { - auto const& [name, tensor] = outputTensor; - auto const& engine = runtime.getEngine(); - auto shape = engine.getTensorShape(name.c_str()); - TLLM_CHECK_WITH_INFO( - shape.d[0] == -1, "First dimension of additional output tensor '%s' must be dynamic", name.c_str()); - shape.d[0] = numTokens; - tensor->reshape(shape); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void RuntimeBuffers::prepareBuffersForCudaGraph(SizeType32 maxSequenceLength) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(prepareBuffersForCudaGraph); - - TLLM_CHECK(numContextRequests == 0); - - if (transformerBuffers) - { - // Set pastKeyValueLength for graph capturing. This way we will capture graph with - // maxKvCacheLengthRounded rounded to the next kKV_CACHE_LEN_CUDA_GRAPH_ROUND_SIZE. - // MMHA will launch excessive amount of blocks and some of them will exit early during the actual launch. - // We can reuse the same graph for the next kKV_CACHE_LEN_CUDA_GRAPH_ROUND_SIZE iterations. - - // make sure the size does not overflow the max allowed pastKvCacheLength - auto const pastKvCacheLength = std::min(maxSequenceLength - 1, maxKvCacheLengthRounded); - - auto* pastKeyValueLengthsPtr = bufferCast(*transformerBuffers->pastKeyValueLengths); - std::fill_n(pastKeyValueLengthsPtr, getNumSequences(), pastKvCacheLength); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void RuntimeBuffers::setFromInputs(RequestVector const& contextRequests, RequestVector const& genRequests, - SizeType32 maxBeamWidth, SizeType32 maxAttentionWindow, runtime::decoder::DecoderState const& decoderState, - kv_cache_manager::BaseKVCacheManager* kvCacheManagerPtr, - kv_cache_manager::BaseKVCacheManager* crossKvCacheManagerPtr, - rnn_state_manager::RnnStateManager* rnnStateManagerPtr, PeftTable const& peftTable, - runtime::TllmRuntime const& runtime, runtime::ModelConfig const& modelConfig, - runtime::WorldConfig const& worldConfig, bool trtOverlap, OptionalRef newOutputTokens) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(runtimeBuffersSetFromInputs); - - auto const& manager = runtime.getBufferManager(); - auto const& stream = runtime.getStream(); - - // Fill requestTypes - { - auto* hostRequestTypes = bufferCast(*requestTypes); - std::fill_n(hostRequestTypes, numContextRequests, runtime::RequestType::kCONTEXT); - std::fill_n(hostRequestTypes + numContextRequests, numGenSequences, runtime::RequestType::kGENERATION); - } - - SizeType32 totalInputSize = 0; - std::vector inputHost; - std::vector positionIdsHost; - std::vector positionIdsHostRow2; - std::vector mropePositionDeltasHost; - std::vector languageAdapterRoutingsHost; - - auto* contextLengthsHostPtr = bufferCast(*contextLengthsHost); - auto* sequenceLengthsHostPtr = bufferCast(*sequenceLengthsHost); - auto* pastKeyValueLengthsPtr - = transformerBuffers ? bufferCast(*transformerBuffers->pastKeyValueLengths) : nullptr; - SizeType32 totalNumLogits{0}; - auto* logitsIdsHostPtr = bufferCast(*logitsIdsHost); - bool const isChatGlm = modelConfig.getModelVariant() == ModelConfig::ModelVariant::kChatGlm; - bool const isGlm = modelConfig.getModelVariant() == ModelConfig::ModelVariant::kGlm; - auto const mropeRotaryCosSinSize = modelConfig.getMaxPositionEmbeddings() * modelConfig.getRotaryEmbeddingDim(); - - { - NVTX3_SCOPED_RANGE(seqSlotsLoop); - auto* seqSlotIndices = bufferCast(*seqSlots); - - SizeType32 batchIdx{0}; - for (auto const& requests : {contextRequests, genRequests}) - { - for (auto const& llmReq : requests) - { - // Get position of the current sequence in the decoder - auto const seqSlot = llmReq->mSeqSlot.value(); - seqSlotIndices[batchIdx] = seqSlot; - ++batchIdx; - } - } - - TLLM_CHECK(seqSlots->getSize() == static_cast(batchIdx)); - manager.copy(*seqSlots, *seqSlotsDevice); - } - - // context preparation loop - if (!contextRequests.empty()) - { - NVTX3_SCOPED_RANGE(contextPrepareLoop); - numContextLogits.resize(contextRequests.size()); - - SizeType32 batchIdx{0}; - for (auto const& llmReq : contextRequests) - { - TLLM_CHECK_WITH_INFO(llmReq->isContextInitState() || llmReq->isDisaggGenerationTransmissionComplete(), - "The request should be in context phase or disaggregated generation tranmissionComplete phase."); - TLLM_CHECK_WITH_INFO( - llmReq->getMaxNumGeneratedTokens() == 0, "Context request should not have generated tokens."); - - auto const& reqTokens = llmReq->getTokens(0); - auto const& draftTokens = llmReq->getDraftTokens(); - auto const draftLength = llmReq->getNumDraftTokens(); - auto const& positionIds = llmReq->getPositionIds(); - - auto const contextChunkSize = llmReq->getContextChunkSize(); - auto const beginCompute = llmReq->getContextCurrentPosition(); - auto const endCompute = beginCompute + contextChunkSize; - inputHost.insert(inputHost.end(), reqTokens.begin() + beginCompute, reqTokens.begin() + endCompute); - - logitsIdsHostPtr[totalNumLogits++] = contextChunkSize; - numContextLogits.at(batchIdx) = modelConfig.computeContextLogits() ? contextChunkSize : 1; - - if (llmReq->isLastContextChunk()) - { - inputHost.insert(inputHost.end(), draftTokens->begin(), draftTokens->end()); - std::fill_n(logitsIdsHostPtr + totalNumLogits, draftLength, 1); - totalNumLogits += draftLength; - } - auto const inputLength = contextChunkSize + (llmReq->isLastContextChunk() ? draftLength : 0); - contextLengthsHostPtr[batchIdx] = inputLength; - auto const sequenceLen = inputLength + llmReq->getContextCurrentPosition(); - sequenceLengthsHostPtr[batchIdx] = sequenceLen; - - if (static_cast(pastKeyValueLengthsPtr)) - { - pastKeyValueLengthsPtr[batchIdx] = beginCompute + inputLength; - } - - if (positionIds.has_value()) - { - TLLM_CHECK_WITH_INFO(!(isChatGlm || isGlm), "ChatGLM-6B and Glm only use the default initialization"); - positionIdsHost.insert(positionIdsHost.end(), positionIds.value()->begin() + beginCompute, - positionIds.value()->begin() + endCompute); - } - else - { - if (isChatGlm) - { - // Specialize for ChatGLM-6B with 2D-Position-Embedding - positionIdsHost.resize(totalInputSize + inputLength); - std::iota(std::begin(positionIdsHost) + totalInputSize, std::end(positionIdsHost), 0); - positionIdsHost.back() = positionIdsHost.back() - 1; - - positionIdsHostRow2.resize(totalInputSize + inputLength); - positionIdsHostRow2.back() = 1; - } - else if (isGlm) - { - // Specialize for GLM-10B with 2D-Position-Embedding and special value of the mask id position - auto start = inputHost.begin() + totalInputSize; - auto end = start + inputLength; - auto it = std::find_if( - start, end, [](SizeType32 id) { return id == 50260 || id == 50263 || id == 50264; }); - llmReq->mMaskPosition = (it != end) ? std::distance(start, it) : maxContextLength; - - positionIdsHost.resize(totalInputSize + inputLength); - std::iota(std::begin(positionIdsHost) + totalInputSize, std::end(positionIdsHost), 0); - positionIdsHost.back() = llmReq->mMaskPosition; - - positionIdsHostRow2.resize(totalInputSize + inputLength); - positionIdsHostRow2.back() = 1; - } - else - { - // Other models - positionIdsHost.resize(totalInputSize + inputLength); - std::iota(std::begin(positionIdsHost) + totalInputSize, - std::begin(positionIdsHost) + totalInputSize + inputLength, beginCompute); - } - } - if (modelConfig.useMrope()) - { - auto optMropeRotaryCosSin = llmReq->getMropeRotaryCosSin().value(); - TLLM_CHECK_WITH_INFO(optMropeRotaryCosSin->getShape().d[0] == mropeRotaryCosSinSize, - "Provided MropeRotarySinCos is %ld and expected is %d.\n", optMropeRotaryCosSin->getShape().d[0], - int(mropeRotaryCosSinSize)); - - auto const mropeRotaryCosSinCtx = ITensor::slice(mropeRotaryCosSin, batchIdx, 1); - manager.copy(*optMropeRotaryCosSin, *mropeRotaryCosSinCtx); - } - - if (modelConfig.useLanguageAdapter()) - { - auto const languageAdapterRouting = llmReq->getLanguageAdapterRouting( - modelConfig.getNumLanguages().value(), endCompute - beginCompute); - languageAdapterRoutingsHost.insert(languageAdapterRoutingsHost.end(), - std::begin(languageAdapterRouting), std::end(languageAdapterRouting)); - } - totalInputSize += inputLength; - ++batchIdx; - } - - if (rnnStateBuffers) - { - rnnStateBuffers->fillSlotMappings(contextRequests, rnnStateManagerPtr); - } - } - - // generation preparation loop - if (!genRequests.empty()) - { - NVTX3_SCOPED_RANGE(genPrepareLoop); - - auto const numContextRequests = static_cast(contextRequests.size()); - auto numSequences = numContextRequests; - for (auto const& llmReq : genRequests) - { - auto const reqBeamWidth = llmReq->getBeamWidthByIter(); - auto const draftLength = llmReq->getNumDraftTokens(); - auto const& draftTokens = llmReq->getDraftTokens(); - auto const numLogits = draftLength + reqBeamWidth; - TLLM_CHECK(draftLength == 0 || reqBeamWidth == 1); - - auto const promptLen = llmReq->mPromptLen; - auto const sequenceLen - = promptLen + llmReq->getMaxNumGeneratedTokens() + static_cast(trtOverlap); - auto const& positionIds = llmReq->getPositionIds(); - for (int beam = 0; beam < reqBeamWidth; ++beam) - { - auto const numTokens = llmReq->getNumTokens(beam) + static_cast(trtOverlap); - // TODO: can this be removed completely? - if (!trtOverlap) - { - auto const lastToken = llmReq->getLastTokens(beam); - inputHost.push_back(lastToken); - if (draftLength > 0) - { - inputHost.insert(inputHost.end(), draftTokens->begin(), draftTokens->end()); - } - } - - // If model updates generation position ids do not append them here. - if (!modelConfig.getSpeculativeDecodingMode().updatesPositionIds()) - { - if (positionIds.has_value()) - { - TLLM_CHECK_WITH_INFO( - !(isChatGlm || isGlm), "ChatGLM-6B and Glm only use the default initialization"); - auto last_context_position_id = positionIds.value()->back(); - positionIdsHost.push_back( - static_cast(last_context_position_id + sequenceLen - promptLen)); - } - else - { - if (isChatGlm) // ChatGLM-6B - { - positionIdsHost.push_back(static_cast(promptLen - 2)); - positionIdsHostRow2.push_back(static_cast(sequenceLen - promptLen + 1)); - } - else if (isGlm) - { - positionIdsHost.push_back(llmReq->mMaskPosition); - positionIdsHostRow2.push_back(static_cast(sequenceLen - promptLen + 1)); - } - else // GPT / ChatGLM2-6B / ChatGLM3-6B / BART - { - // positionIds is just the size of tokens -1 - positionIdsHost.push_back(numTokens - 1); - } - } - } - - if (modelConfig.useMrope()) - { - auto optMropePositionDeltas = llmReq->getMropePositionDeltas().value(); - mropePositionDeltasHost.push_back(optMropePositionDeltas); - } - - if (modelConfig.useLanguageAdapter()) - { - // Generation requests only have one token per sequence - auto const languageAdapterRouting - = llmReq->getLanguageAdapterRouting(modelConfig.getNumLanguages().value(), 1); - languageAdapterRoutingsHost.insert(languageAdapterRoutingsHost.end(), - std::begin(languageAdapterRouting), std::end(languageAdapterRouting)); - } - } - - if (static_cast(pastKeyValueLengthsPtr)) - { - SizeType32 pastKeyValueLength = sequenceLen - 1; - std::fill_n(pastKeyValueLengthsPtr + numSequences, reqBeamWidth, pastKeyValueLength); - } - totalInputSize += numLogits; - - std::fill_n(logitsIdsHostPtr + totalNumLogits, numLogits, 1); - - totalNumLogits += numLogits; - - if (rnnStateBuffers) - { - auto const seqSlot = llmReq->mSeqSlot.value(); - auto& rnnStateManager = *rnnStateManagerPtr; - rnnStateManager.fillSlotMapping(*rnnStateBuffers->slotMappingHost, numSequences, seqSlot, reqBeamWidth); - } - numSequences += reqBeamWidth; - } - - if (transformerBuffers && maxBeamWidth > 1) - { - transformerBuffers->copyCacheIndirection(genRequests, decoderState.getCacheIndirectionOutput(), stream); - } - - numSequences = numContextRequests; - for (auto const& llmReq : genRequests) - { - auto const reqBeamWidth = llmReq->getBeamWidthByIter(); - auto const draftLength = llmReq->getNumDraftTokens(); - - auto const contextQLength = llmReq->mPromptLen + draftLength; - auto const sequenceLen - = contextQLength + llmReq->getMaxNumGeneratedTokens() + static_cast(trtOverlap); - - std::fill_n(contextLengthsHostPtr + numSequences, reqBeamWidth, contextQLength); - std::fill_n(sequenceLengthsHostPtr + numSequences, reqBeamWidth, sequenceLen); - numSequences += reqBeamWidth; - } - if (modelConfig.getSpeculativeDecodingMode().isLookaheadDecoding()) - { - // copy from lookahead decoding buffer - mLookaheadBuffers->setFromInputs(numContextRequests, numGenRequests, *requestTypes, *seqSlots, - decoderState.getLookaheadBuffers(), runtime, modelConfig, worldConfig); - } - } - - // check skipCrossAttnBlocks - if (transformerBuffers && modelConfig.skipCrossAttnBlocks()) - { - bool isSkipCrossAttn = true; - for (auto const& requests : {contextRequests, genRequests}) - { - for (auto const& llmReq : requests) - { - bool tmpValue = false; - if (llmReq->getSkipCrossAttnBlocks() != nullptr) - { - manager.copy(*llmReq->getSkipCrossAttnBlocks(), &tmpValue); - } - isSkipCrossAttn &= tmpValue; - } - } - transformerBuffers->copySkipCrossAttnBlocks(isSkipCrossAttn, runtime); - } - - if (isChatGlm || isGlm) - { - positionIdsHost.reserve(totalInputSize * 2); - positionIdsHost.insert(positionIdsHost.end(), positionIdsHostRow2.begin(), positionIdsHostRow2.end()); - } - - if (modelConfig.useCrossAttention()) - { - encoderBuffers->fill(contextRequests, genRequests, manager); - } - if (modelConfig.usePromptTuning()) - { - promptTuningBuffers->fill(contextRequests, genRequests, manager, modelConfig.usePackedInput()); - } - if (modelConfig.useLoraPlugin()) - { - loraBuffers->fill(contextRequests, genRequests, peftTable, manager, modelConfig, worldConfig); - } - if (modelConfig.useMrope()) - { - if (!mropePositionDeltasHost.empty()) - { - auto mropePositionDeltasGen = ITensor::slice(mropePositionDeltas, 0, numGenSequences); - manager.copy(mropePositionDeltasHost.data(), *mropePositionDeltasGen); - } - } - - { - NVTX3_SCOPED_RANGE(bufferCopies); - if (trtOverlap) - { - auto contextInputsIds = ITensor::slice(inputsIds, 0, numContextTokens); - manager.copy(inputHost.data(), *contextInputsIds); - - if (!genRequests.empty()) - { - auto generationInputsIds = ITensor::slice(inputsIds, numContextTokens); - auto seqSlotsDeviceSlice = ITensor::slice(seqSlotsDevice, numContextRequests); - runtime::kernels::invokeGatherBatch( - *generationInputsIds, *newOutputTokens, *seqSlotsDeviceSlice, maxBeamWidth, stream); - } - } - else - { - manager.copy(inputHost.data(), *inputsIds); - } - // In generation phase, device ptr of context lengths need to be tiled. - manager.copy(*contextLengthsHost, *contextLengthsDevice); - manager.copy(*sequenceLengthsHost, *sequenceLengthsDevice); - auto const logitsIdsHostRange = BufferRange(*logitsIdsHost); - auto lastTokenIdsHostRange = BufferRange(*lastTokenIdsHost); - common::stl_utils::inclusiveScan( - logitsIdsHostRange.begin(), logitsIdsHostRange.end(), lastTokenIdsHostRange.begin()); - manager.copy(*lastTokenIdsHost, *lastTokenIdsDevice); - if (transformerBuffers) - { - TensorPtr decoderPositionIds = modelConfig.getSpeculativeDecodingMode().isLookaheadDecoding() - ? mLookaheadBuffers->positionIdsDevice - : nullptr; - transformerBuffers->copyPositionIds(runtime, positionIdsHost, isChatGlm || isGlm, decoderPositionIds); - } - if (rnnStateBuffers) - { - rnnStateBuffers->copySlotMappingH2D(runtime); - } - if (modelConfig.useLanguageAdapter()) - { - manager.copy(languageAdapterRoutingsHost.data(), *languageAdapterRoutings); - } - } - - if (transformerBuffers && static_cast(kvCacheManagerPtr)) - { - transformerBuffers->copyKvBlockOffsets( - contextRequests, genRequests, kvCacheManagerPtr, crossKvCacheManagerPtr, manager); - } - - if (modelConfig.useCrossAttention()) - { - transformerBuffers->copyCrossAttentionMasks(contextRequests, genRequests, contextLengthsDevice, - encoderBuffers->inputLengths, maxContextLength, encoderBuffers->getMaxInputLengthInBatch(), runtime); - } - - maxKvCacheLengthRounded = 0; - if (static_cast(pastKeyValueLengthsPtr)) - { - auto const maxKvCacheLength - = *std::max_element(pastKeyValueLengthsPtr, pastKeyValueLengthsPtr + getNumSequences()); - // Round up kv cache length - maxKvCacheLengthRounded = common::ceilDiv(maxKvCacheLength, kKV_CACHE_LEN_CUDA_GRAPH_ROUND_SIZE) - * kKV_CACHE_LEN_CUDA_GRAPH_ROUND_SIZE; - } - - if (modelConfig.getSpeculativeDecodingMode().needsDecoderPrologue()) - { - if (modelConfig.getSpeculativeDecodingMode().isExplicitDraftTokens()) - { - prepareExplicitDraftTokenBuffers( - decoderState.getExplicitDraftTokensBuffers(), runtime, modelConfig, worldConfig); - } - if (modelConfig.getSpeculativeDecodingMode().isEagle()) - { - prepareEagleBuffers( - contextRequests, genRequests, decoderState.getEagleBuffers(), runtime, modelConfig, worldConfig); - } - } - - sync_check_cuda_error(stream.get()); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void RuntimeBuffers::prepareExplicitDraftTokenBuffers( - runtime::ExplicitDraftTokensBuffers::Inputs const& explicitDraftTokensBuffers, TllmRuntime const& runtime, - ModelConfig const& modelConfig, WorldConfig const& worldConfig) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - TLLM_CHECK(mExplicitDraftTokensBuffers); - - mExplicitDraftTokensBuffers->setFromInputs(numContextRequests, numGenRequests, *requestTypes, *seqSlots, - explicitDraftTokensBuffers, *transformerBuffers->positionIds, modelConfig, worldConfig, - runtime.getBufferManager(), runtime.getStream()); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void RuntimeBuffers::prepareEagleBuffers(RequestVector const& contextRequests, RequestVector const& genRequests, - runtime::EagleBuffers::Inputs const& eagleBuffers, TllmRuntime const& runtime, ModelConfig const& modelConfig, - WorldConfig const& worldConfig) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - TLLM_CHECK(mEagleBuffers); - - mEagleBuffers->setFromInputs(contextRequests, genRequests, *requestTypes, *seqSlots, eagleBuffers, - runtime.getBufferManager(), modelConfig, worldConfig); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -std::tuple RuntimeBuffers::prepareStep( - RequestVector const& contextRequests, RequestVector const& genRequests, SizeType32 maxBeamWidth, - SizeType32 maxAttentionWindow, runtime::decoder::DecoderState const& decoderState, - kv_cache_manager::BaseKVCacheManager* kvCacheManager, kv_cache_manager::BaseKVCacheManager* crossKvCacheManager, - rnn_state_manager::RnnStateManager* rnnStateManager, PeftTable const& peftTable, TllmRuntime const& runtime, - ModelConfig const& modelConfig, WorldConfig const& worldConfig, bool gatherGenerationLogits, bool trtOverlap, - OptionalRef newOutputTokens) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(runtimeBuffersPrepareStep); - - setBufferSizes(contextRequests, genRequests); - reshape(runtime, modelConfig, worldConfig, gatherGenerationLogits); - - setFromInputs(contextRequests, genRequests, maxBeamWidth, maxAttentionWindow, decoderState, kvCacheManager, - crossKvCacheManager, rnnStateManager, peftTable, runtime, modelConfig, worldConfig, trtOverlap, - newOutputTokens); - - fillIOMaps(modelConfig, worldConfig); - - auto const numTokens = getNumTokens(); - auto const optProfileId = runtime.getOptProfileId(numTokens, ModelConfig::getOptProfilesSplitPoints()); - setContextIndex(optProfileId); - TLLM_LOG_DEBUG("numTokens: %d, optProfileId: %d", numTokens, optProfileId); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); - return {optProfileId, inputMap, outputMap}; -} - -void RuntimeBuffers::fillIOMaps(ModelConfig const& modelConfig, WorldConfig const& worldConfig) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(runtimeBuffersFillIOMaps); - - inputMap.clear(); - outputMap.clear(); - - if (transformerBuffers) - { - transformerBuffers->getBuffers(inputMap, outputMap, modelConfig); - } - if (rnnStateBuffers) - { - rnnStateBuffers->getBuffers(inputMap); - } - - if (worldConfig.isLastPipelineParallelRank()) - { - // feed a view to TensorRT runtime so reshaping does not change logits buffer - outputMap.insert_or_assign(kLogitsTensorName, ITensor::view(logits)); - } - else - { - outputMap.insert_or_assign(kHiddenStatesOutputTensorName, hiddenStates); - } - - if (worldConfig.isFirstPipelineParallelRank()) - { - inputMap.insert_or_assign(kInputIdsTensorName, inputsIds); - } - else - { - inputMap.insert_or_assign(kHiddenStatesInputTensorName, hiddenStates); - } - - inputMap.insert_or_assign(kLastTokenIdsTensorName, lastTokenIdsDevice); - - inputMap.insert_or_assign(kHostRequestTypesTensorName, requestTypes); - // In the generation phase, we still pass context lengths. - inputMap.insert_or_assign(kContextLengthsTensorName, contextLengthsDevice); - inputMap.insert_or_assign(kHostContextLengthsTensorName, contextLengthsHost); - inputMap.insert_or_assign(kSequenceLengthsTensorName, sequenceLengthsDevice); - - if (modelConfig.useCrossAttention()) - { - encoderBuffers->insertInputTensors(inputMap); - } - if (modelConfig.usePromptTuning()) - { - auto const& promptTuningParams = promptTuningBuffers->mPromptTuningParams; - inputMap.insert_or_assign(kPromptEmbeddingTableTensorName, promptTuningParams.embeddingTable); - inputMap.insert_or_assign(kTasksTensorName, promptTuningParams.tasks); - inputMap.insert_or_assign(kPromptVocabSizeTensorName, promptTuningParams.vocabSize); - } - if (modelConfig.useMrope()) - { - - inputMap.insert_or_assign(kMRopeRotaryCosSinTensorName, mropeRotaryCosSin); - inputMap.insert_or_assign(kMRopePositionDeltasTensorName, mropePositionDeltas); - } - if (modelConfig.useLoraPlugin()) - { - loraBuffers->insertInputTensors(inputMap, loraBuffers->mLoraWeightsPointersHost, - loraBuffers->mLoraAdapterSizesHost, modelConfig, worldConfig); - } - if (modelConfig.useLanguageAdapter()) - { - inputMap.insert_or_assign("language_adapter_routings", languageAdapterRoutings); - } - - if (mMedusaBuffers) - { - mMedusaBuffers->insertInputTensors(inputMap, outputMap, worldConfig); - } - if (mLookaheadBuffers) - { - mLookaheadBuffers->insertInputTensors(inputMap, outputMap, worldConfig); - } - if (mExplicitDraftTokensBuffers) - { - mExplicitDraftTokensBuffers->insertInputTensors(inputMap, outputMap, worldConfig); - } - if (mEagleBuffers) - { - mEagleBuffers->insertInputTensors(inputMap, outputMap, worldConfig); - } - - for (auto const& outputTensor : mAdditionalOutputTensors) - { - outputMap.insert_or_assign(outputTensor.first, outputTensor.second); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -} // namespace tensorrt_llm::batch_manager diff --git a/cpp/tensorrt_llm/batch_manager/transformerBuffers.cpp b/cpp/tensorrt_llm/batch_manager/transformerBuffers.cpp deleted file mode 100644 index 4f81c8926682..000000000000 --- a/cpp/tensorrt_llm/batch_manager/transformerBuffers.cpp +++ /dev/null @@ -1,679 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2025 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. - */ - -#include "tensorrt_llm/batch_manager/transformerBuffers.h" - -#include "tensorrt_llm/batch_manager/kvCacheManager.h" -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/logger.h" -#include "tensorrt_llm/common/nvtxUtils.h" -#include "tensorrt_llm/kernels/attentionMask.h" -#include "tensorrt_llm/kernels/contextFusedMultiHeadAttention/fmhaPackedMask.h" -#include "tensorrt_llm/runtime/bufferManager.h" -#include "tensorrt_llm/runtime/common.h" -#include "tensorrt_llm/runtime/iTensor.h" -#include "tensorrt_llm/runtime/modelConfig.h" -#include "tensorrt_llm/runtime/runtimeKernels.h" -#include "tensorrt_llm/runtime/tllmBuffers.h" -#include "tensorrt_llm/runtime/tllmRuntime.h" -#include - -using namespace tensorrt_llm::runtime; -namespace tk = tensorrt_llm::kernels; - -namespace tensorrt_llm::batch_manager -{ - -TransformerBuffers::TransformerBuffers(SizeType32 maxBatchSize, SizeType32 maxBeamWidth, - std::vector const& maxAttentionWindowVec, SizeType32 maxAttentionWindow, SizeType32 sinkTokenLen, - runtime::TllmRuntime const& runtime, runtime::ModelConfig const& modelConfig, - runtime::WorldConfig const& worldConfig) - : maxInputLen(modelConfig.getMaxInputLen()) - , maxEncoderOutputLen(modelConfig.getMaxEncoderLen()) -{ - auto const& manager = runtime.getBufferManager(); - auto const& engine = runtime.getEngine(); - - positionIds = manager.emptyTensor(MemoryType::kGPU, nvinfer1::DataType::kINT32); - - auto const localNbAttnLayers - = modelConfig.getNbAttentionLayers(worldConfig.getPipelineParallelism(), worldConfig.getPipelineParallelRank()); - // find the index of the first attention layer in the current rank - auto const firstLayerId = modelConfig.countLowerRankLayers(runtime::ModelConfig::LayerType::kATTENTION, - worldConfig.getPipelineParallelism(), worldConfig.getPipelineParallelRank()); - - cacheIndirection - = manager.gpu(ITensor::makeShape({maxBatchSize, maxBeamWidth, maxAttentionWindow}), nvinfer1::DataType::kINT32); - - if (!modelConfig.getMaxNumTokens().has_value()) - { - TLLM_THROW("Model must configure a max number of tokens."); - } - maxNumTokens = modelConfig.getMaxNumTokens().value(); - - if (modelConfig.isKVCacheEnabled()) - { - auto const kvCacheBlockOffsetsType = engine.getTensorDataType("kv_cache_block_offsets"); - kvCacheBlockOffsetsHost = manager.emptyTensor(MemoryType::kPINNEDPOOL, kvCacheBlockOffsetsType); - kvCacheBlockOffsetsDevice = manager.emptyTensor(MemoryType::kGPU, kvCacheBlockOffsetsType); - - if (modelConfig.useCrossAttention()) - { - crossKvCacheBlockOffsetsHost = manager.emptyTensor(MemoryType::kPINNEDPOOL, kvCacheBlockOffsetsType); - crossKvCacheBlockOffsetsDevice = manager.emptyTensor(MemoryType::kGPU, kvCacheBlockOffsetsType); - crossAttentionMaskDevice = manager.emptyTensor(MemoryType::kGPU, nvinfer1::DataType::kBOOL); - crossAttentionMaskPinnedHost = tensorrt_llm::runtime::BufferManager::pinnedPool( - ITensor::makeShape({maxNumTokens, maxEncoderOutputLen}), nvinfer1::DataType::kBOOL); - crossAttentionPackedMaskDevice = manager.emptyTensor(MemoryType::kGPU, nvinfer1::DataType::kINT32); - crossAttentionCuQSeqLensDevice = manager.emptyTensor(MemoryType::kGPU, nvinfer1::DataType::kINT32); - crossAttentionPackedMaskCuMaskRowsDevice - = manager.emptyTensor(MemoryType::kGPU, nvinfer1::DataType::kINT32); - - // Pinned memory for batch copy of attention masks. - // There will be paddings in the dim1, so copy it by tokens. - crossAttentionMaskCopySrcOffsets = tensorrt_llm::runtime::BufferManager::pinnedPool( - ITensor::makeShape({maxNumTokens}), nvinfer1::DataType::kINT64); - crossAttentionMaskCopyDstOffsets = tensorrt_llm::runtime::BufferManager::pinnedPool( - ITensor::makeShape({maxNumTokens}), nvinfer1::DataType::kINT64); - crossAttentionMaskCopySizes = tensorrt_llm::runtime::BufferManager::pinnedPool( - ITensor::makeShape({maxNumTokens}), nvinfer1::DataType::kINT64); - } - } - - fillValuesAlt = tensorrt_llm::runtime::BufferManager::pinnedPool( - ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kINT32); - fillValuesAltDevice = manager.gpu(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kINT32); - seqSlotsAlt = tensorrt_llm::runtime::BufferManager::pinnedPool( - ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kINT32); - seqSlotsAltDevice = manager.gpu(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kINT32); - - cacheIndirBatchedCopySrcOffsets = tensorrt_llm::runtime::BufferManager::pinnedPool( - ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kINT64); - cacheIndirBatchedCopyDstOffsets = tensorrt_llm::runtime::BufferManager::pinnedPool( - ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kINT64); - cacheIndirBatchedCopySizes = tensorrt_llm::runtime::BufferManager::pinnedPool( - ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kINT64); - skipCrossAttnBlocks - = tensorrt_llm::runtime::BufferManager::pinnedPool(ITensor::makeShape({1}), nvinfer1::DataType::kBOOL); - - pastKeyValueLengths = manager.emptyTensor(MemoryType::kCPU, nvinfer1::DataType::kINT32); - - maxAttentionWindows = BufferManager::cpu(ITensor::makeShape({localNbAttnLayers}), nvinfer1::DataType::kINT32); - auto* maxAttentionWindowsPtr = bufferCast(*maxAttentionWindows); - auto const attentionWindowLength = maxAttentionWindowVec.size(); - for (SizeType32 i = 0; i < localNbAttnLayers; ++i) - { - maxAttentionWindowsPtr[i] = maxAttentionWindowVec[(firstLayerId + i) % attentionWindowLength]; - } - - sinkTokenLengths = BufferManager::cpu(ITensor::makeShape({1}), nvinfer1::DataType::kINT32); - bufferCast(*sinkTokenLengths)[0] = sinkTokenLen; - - contextProgressHost = BufferManager::cpu(ITensor::makeShape({1}), nvinfer1::DataType::kINT64); - bufferCast(*contextProgressHost)[0] = 0; - - if (modelConfig.useGemmAllReducePlugin() && worldConfig.isTensorParallel()) - { - nvinfer1::DataType ARType = modelConfig.getGemmAllReduceDtype(); - - auto hiddenSize = modelConfig.getHiddenSize() * worldConfig.getTensorParallelism(); - - auto tpGroup = worldConfig.getTensorParallelGroup(); - std::set tpGroupSet(tpGroup.begin(), tpGroup.end()); - - auto outputDims = ITensor::makeShape({modelConfig.getMaxNumTokens().value() * hiddenSize}); - - gemmAllReduceOutput = std::make_shared(outputDims, ARType, tpGroupSet); - } -} - -void TransformerBuffers::reshape(SizeType32 numSequences, SizeType32 numInputTokens) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - pastKeyValueLengths->reshape(ITensor::makeShape({numSequences})); - - if (kvCacheBlockOffsetsHost) - { - auto cacheBlockOffsetsShape = kvCacheBlockOffsetsHost->getShape(); - if (cacheBlockOffsetsShape.nbDims > 0) - { - cacheBlockOffsetsShape.d[1] = numSequences; - kvCacheBlockOffsetsHost->reshape(cacheBlockOffsetsShape); - kvCacheBlockOffsetsDevice->reshape(cacheBlockOffsetsShape); - } - else - { - TLLM_LOG_DEBUG("kvCacheBlockOffsets not allocated yet"); - } - } - - if (crossKvCacheBlockOffsetsHost) - { - TLLM_CHECK_WITH_INFO( - crossKvCacheBlockOffsetsDevice, "crossKvCacheBlockOffsetsDevice is empty for model with cross attention!"); - auto crossCacheBlockOffsetsShape = crossKvCacheBlockOffsetsHost->getShape(); - if (crossCacheBlockOffsetsShape.nbDims > 0) - { - crossCacheBlockOffsetsShape.d[1] = numSequences; - crossKvCacheBlockOffsetsHost->reshape(crossCacheBlockOffsetsShape); - crossKvCacheBlockOffsetsDevice->reshape(crossCacheBlockOffsetsShape); - } - else - { - TLLM_LOG_DEBUG("crossKvCacheBlockOffsets not allocated yet"); - } - } - - if (crossAttentionMaskDevice) - { - auto crossAttentionMaskShape = crossAttentionMaskDevice->getShape(); - if (crossAttentionMaskShape.nbDims > 0) - { - crossAttentionMaskShape.d[0] = numInputTokens; - crossAttentionMaskDevice->reshape(crossAttentionMaskShape); - crossAttentionMaskPinnedHost->reshape(crossAttentionMaskShape); - crossAttentionMaskCopySrcOffsets->reshape(ITensor::makeShape({numInputTokens})); - crossAttentionMaskCopyDstOffsets->reshape(ITensor::makeShape({numInputTokens})); - crossAttentionMaskCopySizes->reshape(ITensor::makeShape({numInputTokens})); - } - else - { - TLLM_LOG_DEBUG("crossAttentionMaskDevice not allocated yet"); - } - } - - if (crossAttentionPackedMaskDevice) - { - auto crossAttentionMaskPackedShape = crossAttentionPackedMaskDevice->getShape(); - if (crossAttentionMaskPackedShape.nbDims > 0) - { - crossAttentionMaskPackedShape.d[0] = numInputTokens; - crossAttentionPackedMaskDevice->reshape(crossAttentionMaskPackedShape); - } - else - { - TLLM_LOG_DEBUG("crossAttentionPackedMaskDevice not allocated yet"); - } - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TransformerBuffers::reshapeKvTensors(SizeType32 maxBatchSize, SizeType32 maxBeamWidth, SizeType32 maxBlocksPerSeq, - kv_cache_manager::CacheType kvCacheType, SizeType32 numPools, BufferManager const& manager) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - // allocate with max shape during init - if (kvCacheType == kv_cache_manager::CacheType::kSELF) - { - auto const cacheBlockOffsetsShape - = ITensor::makeShape({numPools, maxBatchSize * maxBeamWidth, 2, maxBlocksPerSeq}); - - kvCacheBlockOffsetsHost->reshape(cacheBlockOffsetsShape); - manager.setZero(*kvCacheBlockOffsetsHost); - - kvCacheBlockOffsetsDevice->reshape(cacheBlockOffsetsShape); - manager.setZero(*kvCacheBlockOffsetsDevice); - } - else if (kvCacheType == kv_cache_manager::CacheType::kCROSS) - { - auto const crossCacheBlockOffsetsShape - = ITensor::makeShape({numPools, maxBatchSize * maxBeamWidth, 2, maxBlocksPerSeq}); - - crossKvCacheBlockOffsetsHost->reshape(crossCacheBlockOffsetsShape); - manager.setZero(*crossKvCacheBlockOffsetsHost); - - crossKvCacheBlockOffsetsDevice->reshape(crossCacheBlockOffsetsShape); - manager.setZero(*crossKvCacheBlockOffsetsDevice); - - crossAttentionMaskDevice->reshape(ITensor::makeShape({maxNumTokens, maxEncoderOutputLen})); - manager.setZero(*crossAttentionMaskDevice); - manager.setZero(*crossAttentionMaskPinnedHost); - - // Only context attention needs this, so allocate it by shape [maxBatchSize, maxInputLen, maxEncoderOutputLen]. - auto [packedMaskM, packedMaskN] = tk::roundUpPackedMaskMNDims(maxInputLen, maxEncoderOutputLen); - crossAttentionPackedMaskDevice->reshape(ITensor::makeShape({maxBatchSize * packedMaskM, packedMaskN})); - manager.setZero(*crossAttentionPackedMaskDevice); - - crossAttentionCuQSeqLensDevice->reshape(ITensor::makeShape({maxBatchSize + 1})); - manager.setZero(*crossAttentionCuQSeqLensDevice); - - crossAttentionPackedMaskCuMaskRowsDevice->reshape(ITensor::makeShape({maxBatchSize + 1})); - manager.setZero(*crossAttentionPackedMaskCuMaskRowsDevice); - } - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TransformerBuffers::getBuffers( - TensorMap& inputBuffers, TensorMap& outputBuffers, runtime::ModelConfig const& modelConfig) const -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(transformerBuffersGetBuffers); - - inputBuffers.insert_or_assign(kPositionIdsTensorName, positionIds); - inputBuffers.insert_or_assign(kHostPastKeyValueLengthsTensorName, pastKeyValueLengths); - inputBuffers.insert_or_assign(kCacheIndirectionsTensorName, cacheIndirection); - inputBuffers.insert_or_assign(kHostSinkTokenLengthTensorName, sinkTokenLengths); - - inputBuffers.insert_or_assign(kHostMaxAttentionWindowSizesTensorName, maxAttentionWindows); - inputBuffers.insert_or_assign(kKvCacheBlockOffsetsTensorName, kvCacheBlockOffsetsDevice); - inputBuffers.insert_or_assign(kHostKvCacheBlockOffsetsTensorName, kvCacheBlockOffsetsHost); - inputBuffers.insert_or_assign(kHostContextProgressTensorName, contextProgressHost); - - if (crossKvCacheBlockOffsetsHost) - { - inputBuffers.insert_or_assign(kCrossKvCacheBlockOffsetsTensorName, crossKvCacheBlockOffsetsDevice); - inputBuffers.insert_or_assign(kHostCrossKvCacheBlockOffsetsTensorName, crossKvCacheBlockOffsetsHost); - inputBuffers.insert_or_assign(kHostCrossKvCachePoolPointersTensorName, crossKvCacheBlockPoolPointers); - inputBuffers.insert_or_assign(kHostCrossKvCachePoolMappingTensorName, crossKvCacheBlockPoolMapping); - inputBuffers.insert_or_assign(kCrossAttentionMaskTensorName, crossAttentionMaskDevice); - inputBuffers.insert_or_assign(kCrossAttentionPackedMaskTensorName, crossAttentionPackedMaskDevice); - } - - if (skipCrossAttnBlocks) - { - inputBuffers.insert_or_assign(kSkipCrossAttentionBlocksTensorName, skipCrossAttnBlocks); - } - - if (modelConfig.useGemmAllReducePlugin()) - { - for (int idx = 0; idx < modelConfig.getNbAttentionLayers() * 2; ++idx) - { - // XXX (xsimmons): this is a bit hacky as it assumes - // 2x RowLinear layers per attention block. - // This will be fixed soon when I remove coupling between model - // and runtime. - auto gemmARViewUC = gemmAllReduceOutput->getTensorView(MulticastTensorView::ViewType::kUNICAST); - auto gemmARViewMC = gemmAllReduceOutput->getTensorView(MulticastTensorView::ViewType::kMULTICAST); - auto gemmARViewIpc = gemmAllReduceOutput->getTensorView(MulticastTensorView::ViewType::kIPC_LIST); - - outputBuffers.insert_or_assign("gemm_allreduce_uc_out_" + std::to_string(idx), gemmARViewUC); - outputBuffers.insert_or_assign("gemm_allreduce_mc_out_" + std::to_string(idx), gemmARViewMC); - outputBuffers.insert_or_assign("gemm_allreduce_ipc_out_" + std::to_string(idx), gemmARViewIpc); - } - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TransformerBuffers::copyPositionIds(runtime::TllmRuntime const& runtime, - std::vector const& positionIdsHost, bool isChatGlm, TensorPtr const& decoderPositionIds) -{ - auto const& manager = runtime.getBufferManager(); - if (isChatGlm) - { - positionIds->reshape(ITensor::makeShape({2, static_cast(positionIdsHost.size()) / 2})); - manager.copy(positionIdsHost.data(), *positionIds); - } - else if (decoderPositionIds == nullptr) - { - positionIds->reshape(ITensor::makeShape({static_cast(positionIdsHost.size())})); - manager.copy(positionIdsHost.data(), *positionIds); - } - else - { - // concat context phase and generation phase positionIds. - auto const contextPositionIdsLen = static_cast(positionIdsHost.size()); - auto const generationPositionIdsLen = ITensor::volume(decoderPositionIds->getShape()); - positionIds->reshape(ITensor::makeShape({contextPositionIdsLen + generationPositionIdsLen})); - manager.copy(positionIdsHost.data(), *ITensor::slice(positionIds, 0, contextPositionIdsLen)); - manager.copy(*decoderPositionIds, *ITensor::slice(positionIds, contextPositionIdsLen)); - } -} - -void TransformerBuffers::copyKvBlockOffsets(RequestVector const& contextRequests, RequestVector const& genRequests, - kv_cache_manager::BaseKVCacheManager const* kvCacheManager, - kv_cache_manager::BaseKVCacheManager const* crossKvCacheManager, BufferManager const& manager) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(copyKvBlockOffsets); - - auto const& cudaStream = manager.getStream(); - - SizeType32 constexpr contextBeamWidth{1}; - SizeType32 numSequences{0}; - SizeType32 maxBlockCount{0}; - SizeType32 maxCrossBlockCount{0}; - for (auto const& requests : {contextRequests, genRequests}) - { - for (auto const& llmReq : requests) - { - auto const requestId = llmReq->mRequestId; - auto const isContextRequest = llmReq->isContextInitState(); - auto const beamWidth = isContextRequest ? contextBeamWidth : llmReq->getBeamWidthByIter(); - auto const maxBeamBlockCount - = kvCacheManager->copyBlockOffsets(*kvCacheBlockOffsetsHost, numSequences, requestId); - maxBlockCount = std::max(maxBlockCount, maxBeamBlockCount); - if (crossKvCacheBlockOffsetsHost) - { - auto const maxCrossBeamBlockCount - = crossKvCacheManager->copyBlockOffsets(*crossKvCacheBlockOffsetsHost, numSequences, requestId); - maxCrossBlockCount = std::max(maxCrossBlockCount, maxCrossBeamBlockCount); - } - numSequences += beamWidth; - } - } - - // requests' block offsets collected as [totalNumSequences, 2, maxBlocksPerSeq], copy to device - auto copyOffsetsToDevice = [&cudaStream](TensorPtr& offsetsHost, TensorPtr& offsetsDevice, SizeType32 maxBlockCount) - { - // shape should be [totalNumSequences, 2, maxBlocksPerSeq] - auto const& offsetsShape = offsetsHost->getShape(); - auto const maxBlocksPerSeq = offsetsShape.d[3]; - auto const offsetsTypeSize = tensorrt_llm::common::getDTypeSize(offsetsHost->getDataType()); - auto const copyPitch = maxBlocksPerSeq * offsetsTypeSize; - auto const copyHeight = offsetsShape.d[0] * offsetsShape.d[1] * offsetsShape.d[2]; - auto const copyWidth = maxBlockCount * offsetsTypeSize; - auto* srcPtr = bufferCast(*offsetsHost); - auto* dstPtr = bufferCast(*offsetsDevice); - - TLLM_CUDA_CHECK(cudaMemcpy2DAsync( - dstPtr, copyPitch, srcPtr, copyPitch, copyWidth, copyHeight, cudaMemcpyHostToDevice, cudaStream.get())); - }; - - copyOffsetsToDevice(kvCacheBlockOffsetsHost, kvCacheBlockOffsetsDevice, maxBlockCount); - if (crossKvCacheBlockOffsetsHost) - { - copyOffsetsToDevice(crossKvCacheBlockOffsetsHost, crossKvCacheBlockOffsetsDevice, maxCrossBlockCount); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TransformerBuffers::copyCacheIndirection( - RequestVector const& genRequests, TensorPtr const& decoderCacheIndirectionOutput, CudaStream const& stream) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(copyCacheIndirection); - - auto const numGenerationRequests = genRequests.size(); - - auto batchedCopySrcOffsets = BufferRange(*cacheIndirBatchedCopySrcOffsets); - auto batchedCopyDstOffsets = BufferRange(*cacheIndirBatchedCopyDstOffsets); - auto batchedCopySizes = BufferRange(*cacheIndirBatchedCopySizes); - - auto cacheIndirShape = decoderCacheIndirectionOutput->getShape(); - - // At present, all requests of a batch must have the same beam width in one generation step (or they will not - // be batched together). So, the beam width of the first request is taken here to reshape the buffer. - // Corresponding changes must be done if Diverse-Beam-Width-Search (DBWS, requests with diverse beam width in - // a batch in one generation step) is supported in the future. - auto reqBeamWidth = genRequests[0]->getBeamWidthByIter(); - - // Get size of copying from shape of `CacheIndirectionOutput` - cacheIndirShape.d[0] = 1; - cacheIndirShape.d[1] = reqBeamWidth; // Use beam width of current step rather than max beam width as dst offset - auto const copySize = static_cast(ITensor::volume(cacheIndirShape)); - - std::transform(genRequests.begin(), genRequests.end(), batchedCopySrcOffsets.begin(), - [copySize](auto const& llmReq) { return llmReq->mSeqSlot.value() * copySize; }); - std::generate_n( - batchedCopyDstOffsets.begin(), numGenerationRequests, [copySize, i = 0]() mutable { return (i++) * copySize; }); - std::fill_n(batchedCopySizes.begin(), numGenerationRequests, copySize); - - auto const batchedCopySrcOffsetsSlice = ITensor::slice(cacheIndirBatchedCopySrcOffsets, 0, numGenerationRequests); - auto const batchedCopyDstOffsetsSlice = ITensor::slice(cacheIndirBatchedCopyDstOffsets, 0, numGenerationRequests); - auto const batchedCopySizesSlice = ITensor::slice(cacheIndirBatchedCopySizes, 0, numGenerationRequests); - runtime::kernels::invokeCopyBatch(*decoderCacheIndirectionOutput, *cacheIndirection, *batchedCopySrcOffsetsSlice, - *batchedCopyDstOffsetsSlice, *batchedCopySizesSlice, copySize, stream); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TransformerBuffers::copyCrossAttentionMasks(RequestVector const& contextRequests, RequestVector const& genRequests, - TensorPtr const& decoderContextLengthsDevice, TensorPtr const& encoderInputLengths, - SizeType32 maxDecoderContextLength, SizeType32 maxEncoderInputLengthInBatch, TllmRuntime const& runtime) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - auto const& manager = runtime.getBufferManager(); - - // Reshape the tensor to make sure the dim1 matches maxEncoderInputLengthInBatch. - auto crossAttentionMaskShape = crossAttentionMaskDevice->getShape(); - crossAttentionMaskShape.d[1] = maxEncoderInputLengthInBatch; - crossAttentionMaskDevice->reshape(crossAttentionMaskShape); - // Set crossAttentionMask to true by default if it is not provided. - manager.setMem(*crossAttentionMaskDevice, 1); - - // Check if all context requests have cross attention mask. - bool allContextCrossAttentionMaskProvided = true; - for (auto const& llmReq : contextRequests) - { - auto const& crossAttentionMaskRequest = llmReq->getCrossAttentionMask(); - if (bufferCastOrNull(crossAttentionMaskRequest) == nullptr) - { - allContextCrossAttentionMaskProvided = false; - break; - } - } - // If not all requests have cross attention mask, let us create the default ones. - auto const& stream = runtime.getStream(); - if (!allContextCrossAttentionMaskProvided) - { - TLLM_LOG_WARNING("Default padding attention mask will be used as not all requests have cross attention mask."); - tk::AttentionMaskParams attentionMaskParams; - memset((void*) &attentionMaskParams, 0, sizeof(attentionMaskParams)); - // Set parameters. - attentionMaskParams.mask = bufferCastOrNull(crossAttentionMaskDevice); - attentionMaskParams.cuQSeqLens = bufferCastOrNull(crossAttentionCuQSeqLensDevice); - attentionMaskParams.actualQSeqLens = bufferCastOrNull(decoderContextLengthsDevice); - attentionMaskParams.actualKvSeqLens = bufferCastOrNull(encoderInputLengths); - attentionMaskParams.attentionMaskType = tk::AttentionMaskType::PADDING; - attentionMaskParams.batchSize = static_cast(contextRequests.size()); - attentionMaskParams.maxQSeqLen = maxDecoderContextLength; - attentionMaskParams.maxKvSeqLen = maxEncoderInputLengthInBatch; - // Launch the kernel. - tk::invokeBuildAttentionMask(attentionMaskParams, stream.get()); - sync_check_cuda_error(stream.get()); - } - // Use the first request's cross attention mask tensor's pointer address as the primary source pointer. - auto const& attentionMaskSrc = !contextRequests.empty() ? contextRequests[0]->getCrossAttentionMask() - : genRequests[0]->getCrossAttentionMask(); - bool const* primarySrcPtr = bufferCastOrNull(attentionMaskSrc); - - // Pinned-memory buffer preparation for batch copy. - auto batchedCopySrcOffsets = BufferRange(*crossAttentionMaskCopySrcOffsets); - auto batchedCopyDstOffsets = BufferRange(*crossAttentionMaskCopyDstOffsets); - auto batchedCopySizes = BufferRange(*crossAttentionMaskCopySizes); - // Requests with cross-attention-mask don't need to copy. - manager.setZero(*crossAttentionMaskCopySizes); - sync_check_cuda_error(stream.get()); - - SizeType32 numTokens = 0; - SizeType32 numCopiedTokens = 0; - bool* pinnedMemPtr = bufferCastOrNull(crossAttentionMaskPinnedHost); - for (auto const& llmReq : contextRequests) - { - auto const& crossAttentionMaskRequest = llmReq->getCrossAttentionMask(); - auto const position = llmReq->getContextCurrentPosition(); - auto const size = llmReq->getContextChunkSize(); - if (bufferCastOrNull(crossAttentionMaskRequest) != nullptr) - { - auto memType = crossAttentionMaskRequest->getMemoryType(); - auto const crossAttentionMaskRequestDim0 - = static_cast(crossAttentionMaskRequest->getShape().d[0]); - auto const crossAttentionMaskRequestDim1 - = static_cast(crossAttentionMaskRequest->getShape().d[1]); - TLLM_LOG_DEBUG("copyCrossAttentionMasks (shape [%d, %d]) from contextRequests position %d chunkSize %d", - crossAttentionMaskRequestDim0, crossAttentionMaskRequestDim1, position, size); - if ((position + size - 1) >= crossAttentionMaskRequestDim0) - { - TLLM_LOG_WARNING( - "The provided crossAttentionMask input is not complete for context phases, the last row " - "will be " - "used by default."); - } - // copy it to pinned memory if it is a cpu tensor. - if (memType == MemoryType::kCPU) - { - TLLM_LOG_DEBUG("CrossAttentionMask tensor is on CPU."); - auto const copiedPosition - = std::min(crossAttentionMaskRequestDim0 - 1, static_cast(position)); - auto const copiedSize - = std::min(crossAttentionMaskRequestDim0 - copiedPosition, static_cast(size)); - SizeType64 inputMaskOffset = (copiedPosition * crossAttentionMaskRequestDim1); - SizeType64 inputMaskSize = (copiedSize * crossAttentionMaskRequestDim1); - std::memcpy( - pinnedMemPtr, bufferCastOrNull(crossAttentionMaskRequest) + inputMaskOffset, inputMaskSize); - pinnedMemPtr += inputMaskSize; - for (SizeType32 tokenId = position; tokenId < position + size; tokenId++) - { - SizeType64 tokenIdInPinnedMem - = std::min(copiedSize - 1, static_cast(tokenId - position)); - batchedCopySrcOffsets.begin()[numCopiedTokens] - = (pinnedMemPtr - primarySrcPtr) + tokenIdInPinnedMem * crossAttentionMaskRequestDim1; - batchedCopyDstOffsets.begin()[numCopiedTokens] - = numTokens * static_cast(maxEncoderInputLengthInBatch); - batchedCopySizes.begin()[numCopiedTokens] = crossAttentionMaskRequestDim1; - numCopiedTokens++; - numTokens++; - } - } - else - { - TLLM_LOG_DEBUG("CrossAttentionMask tensor is on GPU."); - for (SizeType32 tokenId = position; tokenId < position + size; tokenId++) - { - batchedCopySrcOffsets.begin()[numCopiedTokens] - = static_cast(bufferCastOrNull(crossAttentionMaskRequest) - primarySrcPtr) - + std::min(crossAttentionMaskRequestDim0 - 1, static_cast(tokenId)) - * crossAttentionMaskRequestDim1; - batchedCopyDstOffsets.begin()[numCopiedTokens] - = numTokens * static_cast(maxEncoderInputLengthInBatch); - batchedCopySizes.begin()[numCopiedTokens] = crossAttentionMaskRequestDim1; - numCopiedTokens++; - numTokens++; - } - } - } - else - { - numTokens += size; - TLLM_LOG_WARNING( - "CrossAttentionMask is not provided for the request. Default padding attention mask will be " - "created."); - } - } - sync_check_cuda_error(stream.get()); - - for (auto const& llmReq : genRequests) - { - auto const promptLen = llmReq->mPromptLen; - auto const decodingIter = llmReq->getDecodingIter(); - auto const& crossAttentionMaskRequest = llmReq->getCrossAttentionMask(); - if (bufferCastOrNull(crossAttentionMaskRequest) != nullptr) - { - auto const memType = crossAttentionMaskRequest->getMemoryType(); - auto const crossAttentionMaskRequestDim0 - = static_cast(crossAttentionMaskRequest->getShape().d[0]); - auto const crossAttentionMaskRequestDim1 - = static_cast(crossAttentionMaskRequest->getShape().d[1]); - TLLM_LOG_DEBUG("copyCrossAttentionMasks (shape [%d, %d]) from genRequests decodingIter %d", - crossAttentionMaskRequestDim0, crossAttentionMaskRequestDim1, decodingIter); - if (promptLen + decodingIter - 1 >= crossAttentionMaskRequestDim0) - { - TLLM_LOG_WARNING( - "The provided crossAttentionMask input is not complete for generation phases, the last row " - "will be " - "used by default."); - } - // copy it to pinned memory if it is a cpu tensor. - if (memType == MemoryType::kCPU) - { - TLLM_LOG_DEBUG("CrossAttentionMask tensor is on CPU."); - SizeType64 copiedPosition = std::min( - crossAttentionMaskRequestDim0 - 1, static_cast(promptLen + decodingIter - 1)); - SizeType64 inputMaskOffset = (copiedPosition * crossAttentionMaskRequestDim1); - SizeType64 inputMaskSize = crossAttentionMaskRequestDim1; - std::memcpy( - pinnedMemPtr, bufferCastOrNull(crossAttentionMaskRequest) + inputMaskOffset, inputMaskSize); - pinnedMemPtr += inputMaskSize; - batchedCopySrcOffsets.begin()[numCopiedTokens] = static_cast(pinnedMemPtr - primarySrcPtr); - batchedCopyDstOffsets.begin()[numCopiedTokens] - = numTokens * static_cast(maxEncoderInputLengthInBatch); - batchedCopySizes.begin()[numCopiedTokens] = crossAttentionMaskRequestDim1; - } - else - { - TLLM_LOG_DEBUG("CrossAttentionMask tensor is on GPU."); - batchedCopySrcOffsets.begin()[numCopiedTokens] - = static_cast(bufferCastOrNull(crossAttentionMaskRequest) - primarySrcPtr) - + std::min(crossAttentionMaskRequestDim0 - 1, static_cast(promptLen + decodingIter - 1)) - * crossAttentionMaskRequestDim1; - batchedCopyDstOffsets.begin()[numCopiedTokens] - = numTokens * static_cast(maxEncoderInputLengthInBatch); - batchedCopySizes.begin()[numCopiedTokens] = crossAttentionMaskRequestDim1; - } - numCopiedTokens++; - numTokens++; - } - else - { - numTokens++; - TLLM_LOG_WARNING( - "CrossAttentionMask is not provided for the generation request. Full valid attentionMask will " - "be used " - "by default."); - } - } - sync_check_cuda_error(stream.get()); - - // Copy all requests' attention mask in one kernel. - if (attentionMaskSrc != nullptr) - { - crossAttentionMaskCopySrcOffsets->reshape(ITensor::makeShape({numCopiedTokens})); - crossAttentionMaskCopyDstOffsets->reshape(ITensor::makeShape({numCopiedTokens})); - crossAttentionMaskCopySizes->reshape(ITensor::makeShape({numCopiedTokens})); - runtime::kernels::invokeCopyBatch(*attentionMaskSrc, *crossAttentionMaskDevice, - *crossAttentionMaskCopySrcOffsets, *crossAttentionMaskCopyDstOffsets, *crossAttentionMaskCopySizes, - maxEncoderInputLengthInBatch, stream); - } - sync_check_cuda_error(stream.get()); - - // The packed mask is only needed by context requests now. - if (!contextRequests.empty()) - { - // Set the parameters for creating packed mask for context FMHA. - tk::PackedMaskParams maskParams{}; - maskParams.maskInput = bufferCastOrNull(crossAttentionMaskDevice); - maskParams.cuQSeqLens = bufferCastOrNull(crossAttentionCuQSeqLensDevice); - maskParams.packedMask = bufferCastOrNull(crossAttentionPackedMaskDevice); - maskParams.cuMaskRows = bufferCastOrNull(crossAttentionPackedMaskCuMaskRowsDevice); - maskParams.actualQSeqLens = bufferCastOrNull(decoderContextLengthsDevice); - maskParams.actualKvSeqLens = bufferCastOrNull(encoderInputLengths); - maskParams.batchSize = contextRequests.size(); - maskParams.maxQSeqLen = maxDecoderContextLength; - maskParams.maxKvSeqLen = maxEncoderInputLengthInBatch; - maskParams.attentionMaskType = tk::ContextAttentionMaskType::CUSTOM_MASK; - maskParams.validPosVal = true; - - // Launch the pack mask kernel. - tk::invokeBuildPackedMask(maskParams, stream.get()); - sync_check_cuda_error(stream.get()); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TransformerBuffers::copySkipCrossAttnBlocks(bool const& _skipCrossAttnBlocks, runtime::TllmRuntime const& runtime) -{ - auto const& manager = runtime.getBufferManager(); - manager.copy(&_skipCrossAttnBlocks, *skipCrossAttnBlocks); -} - -} // namespace tensorrt_llm::batch_manager diff --git a/cpp/tensorrt_llm/batch_manager/trtEncoderModel.cpp b/cpp/tensorrt_llm/batch_manager/trtEncoderModel.cpp deleted file mode 100644 index 0d7dbfde42e6..000000000000 --- a/cpp/tensorrt_llm/batch_manager/trtEncoderModel.cpp +++ /dev/null @@ -1,618 +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. - */ - -#include "trtEncoderModel.h" -#include "encoderBuffers.h" -#include "tensorrt_llm/batch_manager/capacityScheduler.h" -#include "tensorrt_llm/batch_manager/microBatchScheduler.h" -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/logger.h" -#include "tensorrt_llm/common/nvtxUtils.h" -#include "tensorrt_llm/executor/executor.h" -#include "tensorrt_llm/runtime/iTensor.h" -#include "tensorrt_llm/runtime/tllmLogger.h" -#include "tensorrt_llm/runtime/tllmRuntime.h" -#include "tensorrt_llm/runtime/utils/runtimeUtils.h" - -#include -#include -#include - -using namespace tensorrt_llm::runtime; -using namespace tensorrt_llm::mpi; - -namespace tensorrt_llm::batch_manager -{ - -TrtEncoderModel::TrtEncoderModel(runtime::ModelConfig const& modelConfig, WorldConfig const& worldConfig, - runtime::RawEngine const& rawEngine, std::shared_ptr logger, - executor::ExecutorConfig const& executorConfig) - : TrtGptModel(modelConfig, worldConfig, executorConfig) - , mModelConfig{modelConfig} - , mWorldConfig{worldConfig} - , mDevice{runtime::utils::initDevice(worldConfig)} - , mLogger{logger ? std::move(logger) : std::make_shared()} - , mRuntime{std::make_shared( - rawEngine, mLogger.get(), executorConfig.getUseGpuDirectStorage(), executorConfig.getGpuWeightsPercent())} - , mNumMicroBatches{1} - , mNumBuffers{mNumMicroBatches} - , mCopyBufferManager{std::make_shared()} -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - TLLM_CHECK_WITH_INFO( - !mWorldConfig.isPipelineParallel(), "Pipeline parallelism is currently not supported for encoder models."); - - createRuntimeContexts(); - - createBuffers(); - - if (mWorldConfig.isPipelineParallel()) - { - auto const& commSession = COMM_SESSION; - mMpiCommPipelinePara = std::make_shared( - commSession.split(mWorldConfig.getTensorParallelRank(), mWorldConfig.getPipelineParallelRank())); - } - - mMicroBatchScheduledRequests.resize(mNumMicroBatches); - // mEncoderWaitEvents.resize(mNumMicroBatches); - - // set noScheduleUntilState to LlmRequestState::kENCODER_INIT for encoder model - // when null kv cache manager is given, request scheduler will use MaxRequests as capacity scheduler, i.e. no - // handling of maximizing utilization or pause/evict - // TODO: finer control on encoder requests scheduling - mCapacityScheduler = std::make_unique( - getMaxBatchSize() * mNumMicroBatches, executorConfig.getSchedulerConfig().getCapacitySchedulerPolicy(), - /*hasKvCacheManager=*/false, /*twoStepsLookAhead=*/false, - /*noScheduleUntilState=*/LlmRequestState::kENCODER_INIT, - /*noScheduleAfterState=*/LlmRequestState::kCONTEXT_INIT, - /*enablePrefixAwareScheduling=*/executorConfig.getSchedulerConfig().getEnablePrefixAwareScheduling()); - - mMicroBatchScheduler = std::make_unique( - std::nullopt, mModelConfig.getMaxInputLen(), LlmRequestState::kENCODER_INIT, LlmRequestState::kCONTEXT_INIT); - - mHiddenSize = modelConfig.getHiddenSize(); - - mMaxInputLen = mModelConfig.getMaxInputLen(); - TLLM_LOG_INFO("TRTEncoderModel mMaxInputLen: reset to %d from build config.", mMaxInputLen); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -BufferManager const& TrtEncoderModel::getBufferManager() const -{ - return mRuntime->getBufferManager(); -} - -BufferManager::CudaStreamPtr TrtEncoderModel::getRuntimeStreamPtr() const -{ - return mRuntime->getStreamPtr(); -} - -nvinfer1::DataType TrtEncoderModel::getTensorDataType(std::string const& name) const -{ - auto const& engine = mRuntime->getEngine(); - return engine.getTensorDataType(name.c_str()); -} - -nvinfer1::Dims TrtEncoderModel::getTensorShape(std::string const& name) const -{ - auto const& engine = mRuntime->getEngine(); - return engine.getTensorShape(name.c_str()); -} - -void TrtEncoderModel::getCurrentIterationStats(executor::IterationStats& stats) const -{ - stats.iter = mIterCounter; -} - -void TrtEncoderModel::getCurrentRequestStats(executor::RequestStatsPerIteration& stats) const -{ - stats.iter = mIterCounter; -} - -executor::DebugTensorsPerIteration TrtEncoderModel::getCurrentDebugTensors() const -{ - executor::DebugTensorsPerIteration debugTensors; - debugTensors.iter = mIterCounter; - - TLLM_LOG_WARNING("TrtEncoderModel doesn't support getting debug tensors."); - - return debugTensors; -} - -void TrtEncoderModel::setLayerProfiler() -{ - TLLM_CHECK(mRuntime); - mRuntime->setLayerProfiler(); -} - -std::string TrtEncoderModel::getLayerProfileInfo() const -{ - TLLM_CHECK(mRuntime); - return mRuntime->getLayerProfileInfo(); -} - -void TrtEncoderModel::createRuntimeContexts() -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - mRuntime->clearContexts(); - auto const numProfiles = mRuntime->getNbProfiles(); - TLLM_CHECK_WITH_INFO(numProfiles == 1, "Encoder only expects one optimization profile"); - for (auto i = 0; i < numProfiles; ++i) - { - mRuntime->addContext(i); - } - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TrtEncoderModel::executeContext(SizeType32 runtimeContextId) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(executeContext); - auto enqueueSuccessful = mRuntime->executeContext(runtimeContextId); - if (!enqueueSuccessful) - { - throw std::runtime_error("Executing TRT engine failed!"); - } - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TrtEncoderModel::createBuffers() -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - for (SizeType32 i = 0; i < mNumBuffers; ++i) - { - mBuffers.emplace_back( - std::make_shared(getMaxBatchSize(), mModelConfig, mWorldConfig, *mRuntime)); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TrtEncoderModel::executeBatch(ScheduledRequests const& scheduledRequests) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(executeBatch); - - // encoder model only have one optimization profile for now, so no optimization profile switch - SizeType32 optProfileIndex = 0; - auto const bufferId = getBufferId(); - if (!scheduledRequests.contextRequests.empty()) - { - // engine I/O - auto [inputMap, outputMap] - = mBuffers[bufferId]->prepareIO(scheduledRequests.contextRequests, mModelConfig, mWorldConfig, *mRuntime); - mRuntime->setInputTensors(optProfileIndex, inputMap); - mRuntime->setOutputTensors(optProfileIndex, outputMap); - - // engine run - executeContext(optProfileIndex); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TrtEncoderModel::rearrangeOutputs(ScheduledRequests const& scheduledRequests) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(rearrangeOutputs); - - auto const bufferId = getBufferId(); - if (!scheduledRequests.contextRequests.empty()) - { - mBuffers[bufferId]->rearrangeOutputs(scheduledRequests.contextRequests, mModelConfig, mWorldConfig, *mRuntime); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TrtEncoderModel::forwardSync() -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE_WITH_NAME(range, "TrtEncoderModel::forwardSync"); - - auto const device = mWorldConfig.getDevice(); - TLLM_CUDA_CHECK(cudaSetDevice(device)); - - auto& currRequests = mMicroBatchScheduledRequests.at(mMicroBatchId); - // auto& encoderWaitEvent = mEncoderWaitEvents.at(mMicroBatchId); - - if (!currRequests.empty()) - { - if (!mWorldConfig.isPipelineParallel() || !mWorldConfig.isLastPipelineParallelRank()) - { - // TLLM_CHECK_WITH_INFO(mEncStepAsyncSndHdl.get() == nullptr, "encoderSync handle must be nullptr."); - // // Wait for encoding for requests in flight for the current micro batch - // mEncStepAsyncSndHdl = encoderSync(currRequests, encoderWaitEvent); - } - else - { - } - - NVTX3_SCOPED_RANGE(pauseFlaggedCurrRequests); - for (auto const& requests : {currRequests.contextRequests}) - { - for (auto const& llmReq : requests) - { - auto const reqId = llmReq->mRequestId; - mInflightReqIds.erase(reqId); - TLLM_LOG_DEBUG("request ID %u removed from ENCODER inflight set", reqId); - - // If a request in encoder phase had been flagged to be paused, pause it right away - if (mReqIdsToPause.find(reqId) != mReqIdsToPause.end()) - { - terminateRequest(llmReq, true); - mReqIdsToPause.erase(reqId); - } - } - } - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TrtEncoderModel::forwardAsync(RequestList const& activeRequests) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE_WITH_NAME(range, "TrtEncoderModel::ForwardAsync"); - auto const device = mWorldConfig.getDevice(); - TLLM_CUDA_CHECK(cudaSetDevice(device)); - - try - { - auto& currRequests = mMicroBatchScheduledRequests.at(mMicroBatchId); - // auto& encoderWaitEvent = mEncoderWaitEvents.at(mMicroBatchId); - - // Get a new set of requests for encoder - // The scheduler will not include any requests that are already in flight for encoder models - // TODO: add pause handling logic - TLLM_LOG_DEBUG("Running ENCODER request scheduler"); - - auto [fittingRequests, fittingDisaggeGenInitReuqests, requestsToPause] = (*mCapacityScheduler)(activeRequests); - - TLLM_CHECK_WITH_INFO( - fittingDisaggeGenInitReuqests.empty(), "Disaggregated servering is not support by encoder model."); - - std::tie(currRequests.contextRequests, std::ignore) = (*mMicroBatchScheduler)( - fittingRequests, mInflightReqIds, getMaxBatchSize(), mModelConfig.getMaxNumTokens()); - - { - NVTX3_SCOPED_RANGE(pauseRequestsFlaggedByScheduler); - // Loop over requests flagged to be paused, and if not in flight pause it right away - for (auto const& llmReq : requestsToPause) - { - auto const reqId = llmReq->mRequestId; - if (mInflightReqIds.find(reqId) == mInflightReqIds.end()) - { - // Not in flight, can terminate right away - terminateRequest(llmReq, true); - } - else - { - // In flight, add to set for pausing later - mReqIdsToPause.insert(reqId); - } - } - } - - TLLM_CHECK(currRequests.size() <= static_cast(getMaxBatchSize())); - - if (!currRequests.empty()) - { - TLLM_LOG_DEBUG("Running ENCODER model with batch size: %u", currRequests.size()); - { - NVTX3_SCOPED_RANGE(updateInflightReqIds); - // Add to set of requests in flight - for (auto const& requests : {currRequests.contextRequests}) - { - for (auto const& llmReq : requests) - { - TLLM_LOG_DEBUG("request ID %u added to ENCODER inflight set", llmReq->mRequestId); - mInflightReqIds.insert(llmReq->mRequestId); - } - } - } - - executeBatch(currRequests); - - sync_check_cuda_error(mRuntime->getStream().get()); - - rearrangeOutputs(currRequests); - - sync_check_cuda_error(mRuntime->getStream().get()); - - // encoderWaitEvent = encoderStepAsync(currRequests); - - for (auto const& requests : {currRequests.contextRequests}) - { - for (auto const& llmReq : requests) - { - if (llmReq->isEncoderInitState()) - { - llmReq->setState(LlmRequestState::kCONTEXT_INIT); - TLLM_LOG_DEBUG("request ID: %u finishes encoder phase", llmReq->mRequestId); - } - } - } - } - - // TODO: PP handling - if (!currRequests.empty()) - { - if (mWorldConfig.isPipelineParallel() && mWorldConfig.isLastPipelineParallelRank()) - { - // TLLM_CHECK_WITH_INFO(mEncStepAsyncSndHdl.get() == nullptr, "decoderSync handle must be nullptr."); - // Wait for encoding for requests in flight for the current micro batch - // mEncStepAsyncSndHdl = encoderSync(currRequests, encoderWaitEvent); - } - } - - // Update the micro batch ID - mMicroBatchId = (mMicroBatchId + 1) % mNumMicroBatches; - } - // In case of error, we need to free the batch slot associated with those requests - catch (std::exception const& e) - { - for (auto const& llmReq : activeRequests) - { - terminateRequest(llmReq); - } - throw; - } - - ++mIterCounter; - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TrtEncoderModel::terminateRequest(std::shared_ptr const& llmReq, bool pause) -{ - // For encoder-only models, just change req state here. might need to do more when using an asynced forward - // For enc-dec models, only remove cross kv cache after decoder - // genenration has finished - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - if (llmReq->isEncoderInitState()) - { - llmReq->setState(LlmRequestState::kCONTEXT_INIT); - } - else - { - TLLM_LOG_DEBUG("Non-encoder request terminated in encoder model: id %lu", llmReq->mRequestId); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TrtEncoderModel::terminateRequestSync( - std::shared_ptr const& llmReq, executor::FinishReason finishReason) -{ - terminateRequest(llmReq, false); - llmReq->finishByReason(finishReason); - llmReq->clearGeneratedTokens(); -} - -void TrtEncoderModel::fillEncoderOutputSync(RequestVector const& requestList, TensorMap outputTensors) -{ - auto const totalTokensNb = outputTensors["encoder_output"]->getShape().d[0]; - auto const encoderOutputDtype = mRuntime->getEngine().getTensorDataType("encoder_output"); - SizeType32 const bytesPerValue = (encoderOutputDtype == nvinfer1::DataType::kFLOAT) ? 4 : 2; - std::vector encoderOutputHost( - totalTokensNb * mHiddenSize * bytesPerValue * mWorldConfig.getTensorParallelism()); - TLLM_CHECK_WITH_INFO(encoderOutputHost.size() > 0, "Encoder output size is 0!"); - getBufferManager().copy(*(outputTensors["encoder_output"]), reinterpret_cast(encoderOutputHost.data())); - getBufferManager().getStream().synchronize(); // TODO: change engine call to async to improve perf. Also - // need to store output buffers, cuda events, etc. - - auto encoderOutputHostPtr = encoderOutputHost.data(); - for (auto const& llmReq : requestList) - { - SizeType32 const seqLen = llmReq->getEncoderOutputLen(); - TensorPtr currentEncoderOutput - = mCopyBufferManager.copyFrom(reinterpret_cast(encoderOutputHostPtr), - ITensor::makeShape({seqLen, mHiddenSize * mWorldConfig.getTensorParallelism()}), MemoryType::kCPU); - llmReq->setEncoderOutputHost(currentEncoderOutput); - encoderOutputHostPtr += seqLen * mHiddenSize * bytesPerValue * mWorldConfig.getTensorParallelism(); - - if (llmReq->isEncoderInitState()) - { - llmReq->setState(LlmRequestState::kCONTEXT_INIT); - } - else - { - TLLM_LOG_DEBUG("Non-encoder request terminated in encoder model: id %lu", llmReq->mRequestId); - } - } -} - -void TrtEncoderModel::executeBatch(RequestVector const& requestList) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(executeBatch); - - auto const modelName = mModelConfig.getModelName(); - TLLM_CHECK_WITH_INFO(modelName == "EncoderModel" || modelName == "WhisperEncoder", "Model not supported."); - TensorMap inputTensors; - TensorMap outputTensors; - TensorPtr rankOutput; - - std::vector inputIdsHost; - std::vector positionIdsHost; - SizeType32 totalOutputLength = 0; - SizeType32 totalInputLength = 0; - std::vector inputLengthsHost; - std::vector inputFeaturesHost; - - inputLengthsHost.reserve(requestList.size()); - SizeType32 maxInputLengthHost = 0; - - for (auto const& llmReq : requestList) - { - SizeType32 length = 0; - if (mModelConfig.getModelName() == "EncoderModel") - { - auto const& reqTokens = *(llmReq->getEncoderTokens().value()); - length = reqTokens.size(); - - inputIdsHost.insert(inputIdsHost.end(), reqTokens.begin(), reqTokens.end()); - maxInputLengthHost = std::max(maxInputLengthHost, static_cast(length)); - } - else if (mModelConfig.getModelName() == "WhisperEncoder") - { - auto const& reqFeatures = llmReq->getEncoderInputFeatures(); // [length, featureDim] - length = reqFeatures->getShape().d[0]; - - auto const curFeatureBytes = reqFeatures->getSizeInBytes(); - auto const srcPtr = reinterpret_cast(reqFeatures->data()); - inputFeaturesHost.insert(inputFeaturesHost.end(), srcPtr, srcPtr + curFeatureBytes); - } - positionIdsHost.reserve(positionIdsHost.size() + length); - auto const newReqPosBegin = positionIdsHost.end(); - positionIdsHost.resize(positionIdsHost.size() + length); - std::iota(newReqPosBegin, positionIdsHost.end(), 0); - - totalOutputLength += llmReq->getEncoderOutputLen(); - totalInputLength += length; - inputLengthsHost.push_back(length); - } - - TensorPtr hiddenStatesInput; - TensorPtr inputLengths = getBufferManager().copyFrom( - inputLengthsHost, ITensor::makeShape({static_cast(inputLengthsHost.size())}), MemoryType::kGPU); - inputTensors.emplace("input_lengths", inputLengths); - - if (mModelConfig.getModelName() == "EncoderModel") - { - // use shape of maxInputLength to indicates max length, content is not important - TensorPtr maxInputLength - = getBufferManager().gpu(ITensor::makeShape({maxInputLengthHost}), nvinfer1::DataType::kINT32); - inputTensors.emplace("max_input_length", maxInputLength); - } - - // engine outputs - rankOutput = getBufferManager().gpu( - ITensor::makeShape({totalOutputLength, mHiddenSize * mWorldConfig.getTensorParallelism()}), - mModelConfig.getDataType()); - - if (mWorldConfig.isFirstPipelineParallelRank()) - { - if (mModelConfig.getModelName() == "EncoderModel") - { - // Engine inputs - TensorPtr inputIds - = getBufferManager().copyFrom(inputIdsHost, ITensor::makeShape({totalInputLength}), MemoryType::kGPU); - TensorPtr positionIds = getBufferManager().copyFrom( - positionIdsHost, ITensor::makeShape({totalInputLength}), MemoryType::kGPU); - inputTensors.emplace("input_ids", inputIds); - inputTensors.emplace("position_ids", positionIds); - } - else if (mModelConfig.getModelName() == "WhisperEncoder") - { - auto inputFeaturesHostPtr = inputFeaturesHost.data(); - auto const featureDim = requestList.front()->getEncoderInputFeatures()->getShape().d[1]; - auto const dtype = requestList.front()->getEncoderInputFeatures()->getDataType(); - TensorPtr inputFeatures = getBufferManager().gpu(ITensor::makeShape({totalInputLength, featureDim}), dtype); - getBufferManager().copy( - reinterpret_cast(inputFeaturesHostPtr), *inputFeatures, runtime::MemoryType::kCPU); - TensorPtr positionIds = getBufferManager().copyFrom( - positionIdsHost, ITensor::makeShape({totalOutputLength}), MemoryType::kGPU); - inputTensors.emplace("input_features", inputFeatures); - inputTensors.emplace("position_ids", positionIds); - } - } - else - { - SizeType32 length = mModelConfig.getModelName() == "WhisperEncoder" ? totalOutputLength : totalInputLength; - hiddenStatesInput - = getBufferManager().gpu(ITensor::makeShape({length, mHiddenSize * mWorldConfig.getTensorParallelism()}), - mModelConfig.getDataType()); - - inputTensors.emplace("hidden_states_input", hiddenStatesInput); - } - - auto const outputName = mWorldConfig.isLastPipelineParallelRank() ? "encoder_output" : "hidden_states_output"; - outputTensors.emplace(outputName, rankOutput); - - // Set input / output tensors to context, encoder model only have one context - mRuntime->setInputTensors(0, inputTensors); - mRuntime->setOutputTensors(0, outputTensors); - - executeContext(0); - - // copy encoder output to llmRequest, if last PP rank - // dispatch result to each llmReq, only needed by the last PP rank - // TODO: more dtypes support - if (mWorldConfig.isLastPipelineParallelRank()) - { - fillEncoderOutputSync(requestList, outputTensors); - } - else - { - getBufferManager().getStream().synchronize(); - } - - // Update the micro batch ID for next microbatches - mMicroBatchId = (mMicroBatchId + 1) % mWorldConfig.getPipelineParallelism(); - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TrtEncoderModel::forward(RequestVector& activeRequests) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - auto const device = mWorldConfig.getDevice(); - TLLM_CUDA_CHECK(cudaSetDevice(device)); - - try - { - if (activeRequests.empty()) - { - return; - } - - executeBatch(activeRequests); - } - catch (std::exception const& e) - { - for (auto& req : activeRequests) - { - terminateRequest(req); - } - throw; - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TrtEncoderModel::setLogitsPostProcessorBatched( - std::optional logitsPostProcessorBatched) -{ - TLLM_CHECK_WITH_INFO(!logitsPostProcessorBatched.has_value(), "TrtEncoderModel does not use logits processor."); -} - -void TrtEncoderModel::setReplicateLogitsPostProcessor(bool replicateLogitsPostProcessor) -{ - TLLM_THROW("TrtEncoderModel does not use logits processor."); -} - -bool TrtEncoderModel::getReplicateLogitsPostProcessor() const -{ - TLLM_THROW("TrtEncoderModel does not use logits processor."); -} - -TrtEncoderModel::~TrtEncoderModel() = default; - -} // namespace tensorrt_llm::batch_manager diff --git a/cpp/tensorrt_llm/batch_manager/trtEncoderModel.h b/cpp/tensorrt_llm/batch_manager/trtEncoderModel.h deleted file mode 100644 index 31f7d3d0c89b..000000000000 --- a/cpp/tensorrt_llm/batch_manager/trtEncoderModel.h +++ /dev/null @@ -1,205 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2025 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. - */ - -#pragma once - -#include "tensorrt_llm/runtime/rawEngine.h" -#include "tensorrt_llm/runtime/utils/mpiUtils.h" -#include "trtGptModel.h" - -#include - -namespace tensorrt_llm::runtime -{ -class TllmRuntime; -class NcclCommunicator; -} // namespace tensorrt_llm::runtime - -namespace tensorrt_llm::batch_manager -{ -class CapacityScheduler; -class MicroBatchScheduler; -class EncoderBuffers; - -class TrtEncoderModel : public TrtGptModel -{ -public: - using SizeType32 = tensorrt_llm::runtime::SizeType32; - using TokenIdType = tensorrt_llm::runtime::TokenIdType; - using BufferManager = tensorrt_llm::runtime::BufferManager; - using TensorMap = runtime::StringPtrMap; - using TensorPtr = runtime::ITensor::SharedPtr; - - TrtEncoderModel(runtime::ModelConfig const& modelConfig, runtime::WorldConfig const& worldConfig, - runtime::RawEngine const& rawEngine, std::shared_ptr logger, - executor::ExecutorConfig const& executorConfig); - - ~TrtEncoderModel() override; - - void terminateRequest(std::shared_ptr const& llmRequest, bool pause = false) override; - void terminateRequestSync( - std::shared_ptr const& llmRequest, executor::FinishReason finishReason) override; - - void forward(RequestVector& activeRequests); - - void forwardSync() override; - - void forwardAsync(RequestList const& activeRequests) override; - - [[nodiscard]] runtime::BufferManager const& getBufferManager() const override; - [[nodiscard]] runtime::BufferManager::CudaStreamPtr getRuntimeStreamPtr() const override; - - runtime::ModelConfig const& getModelConfig() const override - { - return mModelConfig; - } - - [[nodiscard]] bool getGatherGenerationLogits() const override - { - return getModelConfig().computeGenerationLogits(); - } - - runtime::WorldConfig const& getWorldConfig() const override - { - return mWorldConfig; - } - - [[nodiscard]] SizeType32 getHiddenSize() const override - { - return mHiddenSize; - } - - [[nodiscard]] SizeType32 getMaxInputLen() const override - { - return mMaxInputLen; - } - - [[nodiscard]] SizeType32 getNumMicroBatches() const override - { - return mNumMicroBatches; - } - - [[nodiscard]] nvinfer1::DataType getLogitDataType() const override - { - return getModelConfig().getDataType(); - } - - nvinfer1::DataType getTensorDataType(std::string const& name) const override; - nvinfer1::Dims getTensorShape(std::string const& name) const override; - - [[nodiscard]] TrtGptModelType getModelType() const override - { - throw std::runtime_error("TrtEncoderModel does not have model type."); // FIXME: - } - - [[nodiscard]] executor::IterationType getIterCounter() const noexcept override - { - return mIterCounter; - } - - void updatePeftCache(std::shared_ptr const& /*llmRequest*/) override - { - throw std::runtime_error("TrtEncoderModel does not have Peft Cache."); - } - - void getCurrentIterationStats(executor::IterationStats& stats) const override; - void getCurrentRequestStats(executor::RequestStatsPerIteration& stats) const override; - [[nodiscard]] executor::DebugTensorsPerIteration getCurrentDebugTensors() const override; - - void setLayerProfiler() override; - std::string getLayerProfileInfo() const override; - - void setLogitsPostProcessorBatched(std::optional logitsPostProcessorBatched) override; - void setReplicateLogitsPostProcessor(bool replicateLogitsPostProcessor) override; - [[nodiscard]] bool getReplicateLogitsPostProcessor() const override; - - void resetIterationStats() override {} - - [[nodiscard]] SizeType32 getMaxCapacityBatchSize(SizeType32 inputLength, SizeType32 outputLength) const override - { - return 0; - }; - -protected: - std::shared_ptr getKVCacheManager() override - { - throw std::runtime_error("TrtEncoderModel does not have KVCache."); - } - - [[nodiscard]] std::shared_ptr getKVCacheManager() const override - { - throw std::runtime_error("TrtEncoderModel does not have KVCache."); - } - - [[nodiscard]] std::shared_ptr getPeftCacheManager() override - { - throw std::runtime_error("TrtEncoderModel does not use PEFT."); - } - - [[nodiscard]] std::shared_ptr getPeftCacheManager() const override - { - throw std::runtime_error("TrtEncoderModel does not use PEFT."); - } - -private: - [[nodiscard]] SizeType32 getBufferId() const - { - return mMicroBatchId; - } - - void createRuntimeContexts(); - void executeContext(SizeType32 runtimeContextId); - void createBuffers(); - void executeBatch(RequestVector const& requestList); - void executeBatch(ScheduledRequests const& scheduledRequests); - void rearrangeOutputs(ScheduledRequests const& scheduledRequests); - void createCustomAllReduceWorkspace(); - void fillEncoderOutputSync(RequestVector const& requestList, TensorMap outputTensors); - - runtime::ModelConfig const mModelConfig; - runtime::WorldConfig const mWorldConfig; - int mDevice{-1}; - std::shared_ptr mMpiCommPipelinePara; - - std::shared_ptr mLogger; - std::shared_ptr mRuntime; - - SizeType32 mMicroBatchId{0}; - - // TODO: Add runtime buffers for async PP - std::vector> mBuffers; - - SizeType32 mNumMicroBatches; - SizeType32 mNumBuffers; - - std::vector mMicroBatchScheduledRequests; - ReqIdsSet mInflightReqIds; - ReqIdsSet mReqIdsToPause; - - std::unique_ptr mCapacityScheduler; - std::unique_ptr mMicroBatchScheduler; - - SizeType32 mHiddenSize; // already divided by Tensor Parallelism - SizeType32 mMaxInputLen; // WAR for max_input_len == max_seq_len at all circumstances - - runtime::BufferManager mCopyBufferManager; - - // Iteration counter used to distinguish debug output - executor::IterationType mIterCounter{0}; -}; - -} // namespace tensorrt_llm::batch_manager diff --git a/cpp/tensorrt_llm/batch_manager/trtGptModel.h b/cpp/tensorrt_llm/batch_manager/trtGptModel.h deleted file mode 100644 index 54ad36b13895..000000000000 --- a/cpp/tensorrt_llm/batch_manager/trtGptModel.h +++ /dev/null @@ -1,339 +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. - */ - -#pragma once - -#include "tensorrt_llm/batch_manager/peftCacheManager.h" -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/stlUtils.h" -#include "tensorrt_llm/executor/executor.h" -#include "tensorrt_llm/executor/model.h" -#include "tensorrt_llm/runtime/common.h" -#include "tensorrt_llm/runtime/modelConfig.h" -#include "tensorrt_llm/runtime/worldConfig.h" - -#include - -namespace tc = tensorrt_llm::common; - -namespace tensorrt_llm::batch_manager -{ -enum class TrtGptModelType -{ - InflightBatching, - InflightFusedBatching -}; - -class LlmRequest; - -namespace kv_cache_manager -{ -class BaseKVCacheManager; -} // namespace kv_cache_manager - -class TrtGptModel : public executor::Model -{ -public: - using SizeType32 = tensorrt_llm::runtime::SizeType32; - - TrtGptModel(runtime::ModelConfig const& modelConfig, runtime::WorldConfig const& worldConfig, - executor::ExecutorConfig const& executorConfig) - : mMaxBatchSize{executorConfig.getMaxBatchSize().value_or(modelConfig.getMaxBatchSize())} - , mMaxBeamWidth{executorConfig.getMaxBeamWidth()} - , mMaxSequenceLen{modelConfig.getMaxSequenceLen()} - , mMaxDraftLen{modelConfig.getMaxDecodingDraftTokens()} - , mVocabSizePadded{modelConfig.getVocabSizePadded(worldConfig.getSize())} - , mNormalizeLogProbs{executorConfig.getNormalizeLogProbs()} - , mEnableTrtOverlap{executorConfig.getEnableTrtOverlap()} - , mCudaGraphMode{executorConfig.getExtendedRuntimePerfKnobConfig().getCudaGraphMode()} - { - TLLM_CHECK_WITH_INFO(mMaxBeamWidth <= modelConfig.getMaxBeamWidth(), - "Runtime configured max beam width (%d) must not exceed engine max beam width (%d)", mMaxBeamWidth, - modelConfig.getMaxBeamWidth()); - TLLM_CHECK_WITH_INFO(mMaxBatchSize <= modelConfig.getMaxBatchSize(), - "Runtime configured max batch size (%d) must not exceed engine max batch size (%d)", mMaxBatchSize, - modelConfig.getMaxBatchSize()); - if (executorConfig.getEnableTrtOverlap()) - { - if (mMaxBeamWidth > 1) - { - mEnableTrtOverlap = false; - TLLM_LOG_WARNING( - "TRT overlap is not supported with beam search (maxBeamWidth is set to %d) and will be disabled.", - mMaxBeamWidth); - } - if (!modelConfig.getSpeculativeDecodingMode().isNone()) - { - mEnableTrtOverlap = false; - TLLM_LOG_WARNING("TRT overlap is not supported with speculative decoding and will be disabled."); - } - } - - mMaxAttentionWindow = 0; - if (executorConfig.getKvCacheConfig().getMaxAttentionWindowVec().has_value()) - { - bool warning = false; - auto const& maxAttentionWindowVec = executorConfig.getKvCacheConfig().getMaxAttentionWindowVec(); - for (int maxAttenWin : maxAttentionWindowVec.value()) - { - mMaxAttentionWindowVec.push_back(std::min(maxAttenWin, mMaxSequenceLen)); - mMaxAttentionWindow = std::max(mMaxAttentionWindow, mMaxAttentionWindowVec.back()); - if (maxAttenWin > mMaxSequenceLen) - { - warning = true; - } - TLLM_CHECK_WITH_INFO(mMaxAttentionWindowVec.back() > 0, - "Attention window sizes (elements in maxAttentionWindowVec) must be > 0"); - } - if (warning) - { - TLLM_LOG_WARNING( - "The value of maxAttentionWindow cannot exceed mMaxSequenceLen. " - "Therefore, it has been adjusted to match the value of mMaxSequenceLen."); - } - } - else - { - mMaxAttentionWindowVec.push_back(mMaxSequenceLen); - mMaxAttentionWindow = mMaxSequenceLen; - } - - mSinkTokenLen = executorConfig.getKvCacheConfig().getSinkTokenLength().has_value() - ? executorConfig.getKvCacheConfig().getSinkTokenLength().value() - : 0; - - mMaxNumSequences = mMaxBatchSize * worldConfig.getPipelineParallelism(); - - auto const numTotalAttenLayers = modelConfig.getNbAttentionLayers(); - auto const numRepeatsAttenWindow = numTotalAttenLayers / mMaxAttentionWindowVec.size(); - auto const numRemainsAttenWindow = numTotalAttenLayers % mMaxAttentionWindowVec.size(); - std::string attenWindowRemainInfo = numRemainsAttenWindow > 0 - ? " + " + tc::arr2str(mMaxAttentionWindowVec.data(), numRemainsAttenWindow) - : ""; - - TLLM_LOG_INFO("TRTGptModel maxNumSequences: %d", mMaxNumSequences); - TLLM_LOG_INFO("TRTGptModel maxBatchSize: %d", mMaxBatchSize); - TLLM_LOG_INFO("TRTGptModel maxBeamWidth: %d", mMaxBeamWidth); - TLLM_LOG_INFO("TRTGptModel maxSequenceLen: %d", mMaxSequenceLen); - TLLM_LOG_INFO("TRTGptModel maxDraftLen: %d", mMaxDraftLen); - TLLM_LOG_INFO("TRTGptModel mMaxAttentionWindowSize: %s * %d%s", tc::vec2str(mMaxAttentionWindowVec).c_str(), - numRepeatsAttenWindow, attenWindowRemainInfo.c_str()); - TLLM_LOG_INFO("TRTGptModel enableTrtOverlap: %d", mEnableTrtOverlap); - TLLM_LOG_INFO("TRTGptModel normalizeLogProbs: %d", mNormalizeLogProbs); - - mMaxNumTokens = modelConfig.getMaxNumTokens(); - if (executorConfig.getMaxNumTokens().has_value() && mMaxNumTokens) - { - if (executorConfig.getMaxNumTokens().value() > mMaxNumTokens.value()) - { - TLLM_LOG_WARNING( - "Runtime configured max num tokens (%d) is larger than model max num tokens (%d) and will be " - "ignored.", - executorConfig.getMaxNumTokens().value(), mMaxNumTokens.value()); - } - else - { - mMaxNumTokens = executorConfig.getMaxNumTokens(); - } - } - if (mMaxNumTokens) - { - TLLM_LOG_INFO("TRTGptModel maxNumTokens: %d", mMaxNumTokens.value()); - } - - if (executorConfig.getEnableChunkedContext()) - { - mMaxInputLen = mMaxSequenceLen - 1; - TLLM_LOG_INFO( - "TRTGptModel maxInputLen: %d = maxSequenceLen - 1 since chunked context is enabled", mMaxInputLen); - TLLM_LOG_INFO( - "TRTGptModel If model type is encoder, maxInputLen would be reset in trtEncoderModel to maxInputLen: " - "%d = maxSequenceLen.", - mMaxSequenceLen); - } - else if (modelConfig.getContextFMHA() && modelConfig.usePackedInput()) - { - TLLM_CHECK_WITH_INFO( - mMaxNumTokens, "Max number of tokens has to be set for context FMHA and usePackedInput case."); - mMaxInputLen = std::min(mMaxSequenceLen - 1, mMaxNumTokens.value()); - TLLM_LOG_INFO( - "TRTGptModel maxInputLen: %d = min(maxSequenceLen - 1, maxNumTokens) since context FMHA " - "and usePackedInput are enabled", - mMaxInputLen); - TLLM_LOG_INFO( - "TRTGptModel If model type is encoder, maxInputLen would be reset in trtEncoderModel to maxInputLen: " - "min(maxSequenceLen, maxNumTokens)."); - } - else - { - mMaxInputLen = modelConfig.getMaxInputLen(); - TLLM_LOG_INFO("TRTGptModel maxInputLen: %d = max_input_len (in trtllm-build args)", mMaxInputLen); - } - - using tensorrt_llm::common::stl_utils::toString; - - TLLM_LOG_INFO("Capacity Scheduler Policy: %s", - toString(executorConfig.getSchedulerConfig().getCapacitySchedulerPolicy()).c_str()); - TLLM_LOG_INFO("Context Chunking Scheduler Policy: %s", - toString(executorConfig.getSchedulerConfig().getContextChunkingPolicy()).c_str()); - } - - [[nodiscard]] std::optional getMaxNumTokens() const - { - return mMaxNumTokens; - } - - [[nodiscard]] SizeType32 getMaxNumSequences() const override - { - return mMaxNumSequences; - } - - [[nodiscard]] SizeType32 getMaxBatchSize() const - { - return mMaxBatchSize; - } - - [[nodiscard]] SizeType32 getMaxInputLen() const override - { - return mMaxInputLen; - } - - [[nodiscard]] SizeType32 getHiddenSize() const override - { - return getModelConfig().getHiddenSize(); - }; - - [[nodiscard]] SizeType32 getMaxSequenceLen() const override - { - return mMaxSequenceLen; - } - - [[nodiscard]] virtual TrtGptModelType getModelType() const = 0; - - [[nodiscard]] SizeType32 getVocabSizePadded() const override - { - return mVocabSizePadded; - } - - [[nodiscard]] SizeType32 getMaxDraftLen() const override - { - return mMaxDraftLen; - } - - [[nodiscard]] SizeType32 getOperatingBeamWidth() const override - { - return mMaxBeamWidth; - } - - [[nodiscard]] bool hasSpeculativeDecodingFastLogits() const noexcept override - { - return false; - } - - [[nodiscard]] bool hasGuidedDecoder() const noexcept override - { - return false; - } - - virtual void setLayerProfiler() = 0; - [[nodiscard]] virtual std::string getLayerProfileInfo() const = 0; - - [[nodiscard]] bool hasKVCacheManager() const - { - return getKVCacheManager() != nullptr; - } - -protected: - [[nodiscard]] SizeType32 getMaxBeamWidth() const - { - return mMaxBeamWidth; - } - - [[nodiscard]] std::vector getMaxAttentionWindowVec() const - { - return mMaxAttentionWindowVec; - } - - [[nodiscard]] SizeType32 getMaxAttentionWindow() const - { - return mMaxAttentionWindow; - } - - [[nodiscard]] SizeType32 getSinkTokenLen() const - { - return mSinkTokenLen; - } - - [[nodiscard]] bool isNormalizeLogProbs() const - { - return mNormalizeLogProbs; - } - - [[nodiscard]] bool isTrtOverlap() const - { - return mEnableTrtOverlap; - } - - [[nodiscard]] bool isCudaGraphMode() const - { - return mCudaGraphMode; - } - - void setMaxAttentionWindowVec(std::vector const& maxAttentionWindowVec) - { - TLLM_CHECK_WITH_INFO(maxAttentionWindowVec.size() == mMaxAttentionWindowVec.size(), - "The size of maxAttentionWindowVec must match the size of mMaxAttentionWindowVec"); - mMaxAttentionWindowVec = maxAttentionWindowVec; - mMaxAttentionWindow = *std::max_element(std::begin(mMaxAttentionWindowVec), std::end(mMaxAttentionWindowVec)); - } - - void setMaxSequenceLen(SizeType32 maxSequenceLen) - { - mMaxSequenceLen = maxSequenceLen; - } - - void setMaxInputLen(SizeType32 maxInputLen) - { - mMaxInputLen = maxInputLen; - } - - [[nodiscard]] std::shared_ptr getKVCacheManager() override = 0; - [[nodiscard]] std::shared_ptr getKVCacheManager() const override = 0; - - [[nodiscard]] virtual std::shared_ptr getPeftCacheManager() = 0; - [[nodiscard]] virtual std::shared_ptr getPeftCacheManager() const = 0; - -private: - std::optional mMaxNumTokens; - SizeType32 mMaxNumSequences; - SizeType32 mMaxBatchSize; - SizeType32 mMaxBeamWidth; - SizeType32 mMaxInputLen; - SizeType32 mMaxSequenceLen; - SizeType32 mMaxDraftLen; - - SizeType32 mVocabSizePadded; - std::vector mMaxAttentionWindowVec; - SizeType32 mMaxAttentionWindow; - SizeType32 mSinkTokenLen; - - bool mNormalizeLogProbs; - bool mEnableTrtOverlap; - bool mCudaGraphMode; -}; - -} // namespace tensorrt_llm::batch_manager diff --git a/cpp/tensorrt_llm/batch_manager/trtGptModelFactory.h b/cpp/tensorrt_llm/batch_manager/trtGptModelFactory.h deleted file mode 100644 index bd4d7c767378..000000000000 --- a/cpp/tensorrt_llm/batch_manager/trtGptModelFactory.h +++ /dev/null @@ -1,98 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2025 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. - */ - -#pragma once - -#include "tensorrt_llm/batch_manager/trtGptModel.h" -#include "tensorrt_llm/batch_manager/trtGptModelInflightBatching.h" -#include "tensorrt_llm/executor/executor.h" -#include "tensorrt_llm/runtime/gptJsonConfig.h" -#include "tensorrt_llm/runtime/modelConfig.h" -#include "tensorrt_llm/runtime/rawEngine.h" -#include "tensorrt_llm/runtime/tllmLogger.h" -#include "tensorrt_llm/runtime/worldConfig.h" - -#include - -#include -#include - -namespace tensorrt_llm::batch_manager -{ - -class TrtGptModelFactory -{ -public: - using SizeType32 = tensorrt_llm::runtime::SizeType32; - - static std::shared_ptr create(std::filesystem::path const& trtEnginePath, TrtGptModelType modelType, - executor::ExecutorConfig const& executorConfig, bool isLeaderInOrchMode) - { - auto const jsonConfig = runtime::GptJsonConfig::parse(trtEnginePath / "config.json"); - auto const& deviceIds = executorConfig.getParallelConfig().value_or(executor::ParallelConfig()).getDeviceIds(); - auto const worldConfig = getWorldConfig(jsonConfig, deviceIds); - auto const enginePath = trtEnginePath / jsonConfig.engineFilename(worldConfig); - - auto const& modelConfig = jsonConfig.getModelConfig(); - return create( - runtime::RawEngine(enginePath), modelConfig, worldConfig, modelType, executorConfig, isLeaderInOrchMode); - } - - static std::shared_ptr create(std::filesystem::path const& trtEnginePath, TrtGptModelType modelType, - runtime::GptJsonConfig const& jsonConfig, runtime::WorldConfig const& worldConfig, - executor::ExecutorConfig const& executorConfig, bool isLeaderInOrchMode) - { - auto const enginePath = trtEnginePath / jsonConfig.engineFilename(worldConfig); - auto const& modelConfig = jsonConfig.getModelConfig(); - return create( - runtime::RawEngine(enginePath), modelConfig, worldConfig, modelType, executorConfig, isLeaderInOrchMode); - } - - static std::shared_ptr create(runtime::RawEngine const& rawEngine, - runtime::ModelConfig const& modelConfig, runtime::WorldConfig const& worldConfig, TrtGptModelType modelType, - executor::ExecutorConfig const& executorConfig, bool isLeaderInOrchMode) - { - auto logger = std::make_shared(); - auto const device = worldConfig.getDevice(); - auto const rank = worldConfig.getRank(); - TLLM_LOG_INFO("Rank %d is using GPU %d", rank, device); - TLLM_CUDA_CHECK(cudaSetDevice(device)); - - if ((modelType == TrtGptModelType::InflightBatching) || (modelType == TrtGptModelType::InflightFusedBatching)) - { - executor::ExecutorConfig const& fixedExecutorConfig - = TrtGptModelInflightBatching::executorConfigIsValid(modelConfig, executorConfig) - ? executorConfig - : TrtGptModelInflightBatching::fixExecutorConfig(modelConfig, executorConfig); - bool const ctxGenFusion = modelType == TrtGptModelType::InflightFusedBatching; - return std::make_shared( - logger, modelConfig, worldConfig, rawEngine, ctxGenFusion, fixedExecutorConfig, isLeaderInOrchMode); - } - - throw std::runtime_error("Invalid modelType in trtGptModelFactory"); - } - -private: - static runtime::WorldConfig getWorldConfig( - runtime::GptJsonConfig const& json, std::optional> const& deviceIds) - { - return runtime::WorldConfig::mpi(json.getGpusPerNode(), json.getTensorParallelism(), - json.getPipelineParallelism(), json.getContextParallelism(), deviceIds); - } -}; - -} // namespace tensorrt_llm::batch_manager diff --git a/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp b/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp deleted file mode 100644 index 7a0d78beb8a0..000000000000 --- a/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp +++ /dev/null @@ -1,3136 +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. - */ - -#include "trtGptModelInflightBatching.h" - -#include "tensorrt_llm/batch_manager/allocateKvCache.h" -#include "tensorrt_llm/batch_manager/assignReqSeqSlots.h" -#include "tensorrt_llm/batch_manager/cacheTransceiver.h" -#include "tensorrt_llm/batch_manager/capacityScheduler.h" -#include "tensorrt_llm/batch_manager/common.h" -#include "tensorrt_llm/batch_manager/contextProgress.h" -#include "tensorrt_llm/batch_manager/createNewDecoderRequests.h" -#include "tensorrt_llm/batch_manager/decoderBuffers.h" -#include "tensorrt_llm/batch_manager/disaggTransferAdmissionController.h" -#include "tensorrt_llm/batch_manager/guidedDecoder.h" -#include "tensorrt_llm/batch_manager/handleContextLogits.h" -#include "tensorrt_llm/batch_manager/handleGenerationLogits.h" -#include "tensorrt_llm/batch_manager/kvCacheEventManager.h" -#include "tensorrt_llm/batch_manager/kvCacheManager.h" -#include "tensorrt_llm/batch_manager/llmRequest.h" -#include "tensorrt_llm/batch_manager/logitsPostProcessor.h" -#include "tensorrt_llm/batch_manager/makeDecodingBatchInputOutput.h" -#include "tensorrt_llm/batch_manager/microBatchScheduler.h" -#include "tensorrt_llm/batch_manager/pauseRequests.h" -#include "tensorrt_llm/batch_manager/peftCacheManager.h" -#include "tensorrt_llm/batch_manager/promptTuningBuffers.h" -#include "tensorrt_llm/batch_manager/rnnStateManager.h" -#include "tensorrt_llm/batch_manager/runtimeBuffers.h" -#include "tensorrt_llm/batch_manager/sequenceSlotManager.h" -#include "tensorrt_llm/batch_manager/transformerBuffers.h" -#include "tensorrt_llm/batch_manager/updateDecoderBuffers.h" -#include "tensorrt_llm/batch_manager/utils/debugUtils.h" -#include "tensorrt_llm/batch_manager/utils/inflightBatchingUtils.h" -#include "tensorrt_llm/batch_manager/utils/logitsThread.h" -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/cudaUtils.h" -#include "tensorrt_llm/common/envUtils.h" -#include "tensorrt_llm/common/logger.h" -#include "tensorrt_llm/common/memoryUtils.h" -#include "tensorrt_llm/common/nvtxUtils.h" -#include "tensorrt_llm/common/timestampUtils.h" -#include "tensorrt_llm/kernels/decodingCommon.h" -#include "tensorrt_llm/layers/defaultDecodingParams.h" -#include "tensorrt_llm/runtime/common.h" -#include "tensorrt_llm/runtime/gptDecoderBatched.h" -#include "tensorrt_llm/runtime/iBuffer.h" -#include "tensorrt_llm/runtime/iTensor.h" -#include "tensorrt_llm/runtime/ipcUtils.h" -#include "tensorrt_llm/runtime/lookaheadModule.h" -#include "tensorrt_llm/runtime/memoryCounters.h" -#include "tensorrt_llm/runtime/runtimeKernels.h" -#include "tensorrt_llm/runtime/tllmLogger.h" -#include "tensorrt_llm/runtime/tllmRuntime.h" -#include "tensorrt_llm/runtime/utils/mpiUtils.h" -#include "tensorrt_llm/runtime/utils/runtimeUtils.h" - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -using namespace tensorrt_llm::runtime; -namespace tc = tensorrt_llm::common; -namespace tk = tensorrt_llm::kernels; - -using tensorrt_llm::batch_manager::CacheTransceiverFactory; - -namespace tensorrt_llm::batch_manager -{ - -std::map TrtGptModelInflightBatching::calculateCacheSizePerTokenForDisagg( - ModelConfig const& modelConfig, WorldConfig const& worldConfig, - std::vector const& maxAttentionWindowVec, bool isCrossAttention, SizeType32 kvFactor) -{ - // These are the number of attention layers on this PP rank. - auto const numLocalAttnLayers - = modelConfig.getNbAttentionLayers(worldConfig.getPipelineParallelism(), worldConfig.getPipelineParallelRank()); - // These are the number of attention layers on all previous PP ranks. - auto const numLowerRankAttnLayers = modelConfig.countLowerRankLayers(ModelConfig::LayerType::kATTENTION, - worldConfig.getPipelineParallelism(), worldConfig.getPipelineParallelRank()); - // Use global ranks of attention layers to lookup from maxAttentionWindowVec. - auto const startAttnLayerId = numLowerRankAttnLayers; - auto const endAttnLayerId = numLowerRankAttnLayers + numLocalAttnLayers; - auto const numNonUniqueWindowSizes = static_cast(maxAttentionWindowVec.size()); - std::map> uniqueWindowSizeToLayers; - for (SizeType32 layerIdx = startAttnLayerId; layerIdx < endAttnLayerId; layerIdx++) - { - // maxAttentionWindowVec may or may not be stretched to the length of numLayers yet. - // If not stretched yet, we cycle through the window sizes. - auto const windowSize = maxAttentionWindowVec.at(layerIdx % numNonUniqueWindowSizes); - uniqueWindowSizeToLayers[windowSize].push_back(layerIdx); - } - std::map cacheSizeBytesPerTokenPerWindow; - for (auto const& [windowSize, globalLayerIds] : uniqueWindowSizeToLayers) - { - auto const nkvh = modelConfig.getNumKvHeadsForGivenLayers(globalLayerIds, isCrossAttention); - auto const sumLocalHeads = std::reduce(nkvh.cbegin(), nkvh.cend()); - auto const cacheSizePerToken = sumLocalHeads * kvFactor * modelConfig.getSizePerHead(); - auto const cacheSizeBytesPerToken = cacheSizePerToken * BufferDataType(modelConfig.getKvDataType()).getSize(); - cacheSizeBytesPerTokenPerWindow[windowSize] = cacheSizeBytesPerToken; - } - - return cacheSizeBytesPerTokenPerWindow; -}; - -bool TrtGptModelInflightBatching::executorConfigIsValid( - ModelConfig const& modelConfig, executor::ExecutorConfig const& executorConfig) -{ - // Make sure logic in this function matches fixExecutorConfig - if (executorConfig.getKvCacheConfig().getEnableBlockReuse()) - { - if (!modelConfig.getPagedContextFMHA()) - { - return false; - } - // Context logits cannot be returned for reused tokens, so disable reuse - if (modelConfig.computeContextLogits()) - { - return false; - } - } - return true; -} - -executor::ExecutorConfig TrtGptModelInflightBatching::fixExecutorConfig( - ModelConfig const& modelConfig, executor::ExecutorConfig const& executorConfig) -{ - // Make sure logic in this function matches executorConfigIsValid - if (executorConfig.getKvCacheConfig().getEnableBlockReuse()) - { - auto kvCacheConfig = executorConfig.getKvCacheConfig(); - - if (!modelConfig.getPagedContextFMHA()) - { - TLLM_LOG_WARNING( - "Fixing executorConfig: KV cache reuse disabled because model was not built with paged context FMHA " - "support"); - kvCacheConfig.setEnableBlockReuse(false); - } - if (modelConfig.computeContextLogits()) - { - TLLM_LOG_WARNING( - "Fixing executorConfig: KV cache reuse disabled because model was built to return context logits"); - kvCacheConfig.setEnableBlockReuse(false); - } - - auto fixedExecutorConfig = executor::ExecutorConfig(executorConfig); - fixedExecutorConfig.setKvCacheConfig(kvCacheConfig); - return fixedExecutorConfig; - } - return executorConfig; -} - -TrtGptModelInflightBatching::TrtGptModelInflightBatching(std::shared_ptr logger, - ModelConfig const& modelConfig, WorldConfig const& worldConfig, RawEngine const& rawEngine, bool ctxGenFusion, - executor::ExecutorConfig const& executorConfig, bool isLeaderInOrchMode) - : TrtGptModel(modelConfig, worldConfig, executorConfig) - , mModelConfig(modelConfig) - , mWorldConfig(worldConfig) - , mDevice{runtime::utils::initDevice(worldConfig)} - , mDecodingConfig{executorConfig.getDecodingConfig().value_or(executor::DecodingConfig{})} - , mExtendedRuntimePerfKnobConfig{executorConfig.getExtendedRuntimePerfKnobConfig()} - , mDebugConfig{executorConfig.getDebugConfig()} - , mAdditionalModelOutputs{worldConfig.isLastPipelineParallelRank() ? executorConfig.getAdditionalModelOutputs() - : std::nullopt} - , mLogger{logger ? std::move(logger) : std::make_shared()} - , mRuntime{std::make_unique(rawEngine, mLogger.get(), executorConfig.getUseGpuDirectStorage(), - executorConfig.getGpuWeightsPercent(), modelConfig.useShapeInference())} - , mCopyBufferManager{std::make_shared()} - , mCtxGenFusion(ctxGenFusion) - , mOperatingBeamWidth{getMaxBeamWidth()} - , mGatherGenerationLogits{executorConfig.getGatherGenerationLogits()} - , mPromptTableOffloading{executorConfig.getPromptTableOffloading()} - , mIsLeaderInOrchMode{isLeaderInOrchMode} -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - TLLM_LOG_INFO("gatherContextLogits: %d", mModelConfig.computeContextLogits()); - TLLM_LOG_INFO("gatherGenerationLogits: %d", getGatherGenerationLogits()); - - if (!(mModelConfig.supportsInflightBatching())) - { - throw std::runtime_error( - "TrtGptModelInflightBatching requires GPT attention/Mamba Conv 1d plugin with " - "packed input and paged KV cache."); - } - if (mWorldConfig.isTensorParallel()) - { - mRuntime->initializeUserBuffer(mWorldConfig, mModelConfig.getMaxBatchSize(), mModelConfig.getMaxBeamWidth(), - mModelConfig.getMaxSequenceLen(), mModelConfig.getHiddenSize(), getMaxNumTokens()); - } - if (mWorldConfig.isPipelineParallel()) - { - mNumMicroBatches = mWorldConfig.getPipelineParallelism(); - } - else - { - mNumMicroBatches = isTrtOverlap() ? 2 : 1; - } - - mNumBuffers = (mCtxGenFusion ? 1 : 2) * mNumMicroBatches; - - auto const& kvCacheConfig = executorConfig.getKvCacheConfig(); - - if (mModelConfig.getSpeculativeDecodingMode().isDraftTokensExternal()) - { - TLLM_CHECK_WITH_INFO(kvCacheConfig.getEnableBlockReuse(), - "KV cache block reuse must be enabled for speculative decoding target model"); - } - - if (mCtxGenFusion) - { - TLLM_CHECK_WITH_INFO(!mModelConfig.isRnnBased(), "RNN based model doesn't support context generation fusion."); - TLLM_CHECK_WITH_INFO( - mModelConfig.isTransformerBased(), "Only transformer based model support context generation fusion now."); - } - - if (mModelConfig.getSpeculativeDecodingMode().isLookaheadDecoding()) - { - mSeamlessLADMaxDraftLen = modelConfig.getMaxDecodingDraftTokens(); - // TODO: enable it when speculativeDecodingMode is None and run with '--lookahead_config' - mUseSeamlessLookahead = false; - } - - setupSpeculativeDecodingModule(mDecodingConfig); - - if (mWorldConfig.isLastPipelineParallelRank() && executorConfig.getGuidedDecodingConfig()) - { - mGuidedDecoder = std::make_unique(executorConfig.getGuidedDecodingConfig().value(), - getMaxNumSequences(), mModelConfig.getVocabSizePadded(mWorldConfig.getSize()), - mModelConfig.getLogitsDtype(), mRuntime->getBufferManager()); - } - - createRuntimeContexts(); - - if (mWorldConfig.isTensorParallel()) - { - createCustomAllReduceWorkspace(); - } - - if (mModelConfig.isTransformerBased()) - { - createRuntimePerfKnobsTensor(mExtendedRuntimePerfKnobConfig); - } - - auto& memCounter = MemoryCounters::getInstance(); - auto const gpuUsage1 = memCounter.getGpu(); - createBuffers(mDecodingConfig, mAdditionalModelOutputs); - auto const gpuUsage2 = memCounter.getGpu(); - TLLM_LOG_INFO("[MemUsageChange] Allocated %s GPU memory for runtime buffers.", - memCounter.bytesToString(gpuUsage2 - gpuUsage1).c_str()); - - createDecoder(mDecodingConfig.getDecodingMode()); - auto const gpuUsage3 = memCounter.getGpu(); - TLLM_LOG_INFO("[MemUsageChange] Allocated %s GPU memory for decoder.", - memCounter.bytesToString(gpuUsage3 - gpuUsage2).c_str()); - - if (modelConfig.getManageWeightsType() != ModelConfig::ManageWeightsType::kDisabled) - { - mRuntime->loadManagedWeights(rawEngine, worldConfig.getLocalRank()); - } - - if (mModelConfig.useLoraPlugin()) - { - auto const peftCacheManagerConfig - = PeftCacheManagerConfig(executorConfig.getPeftCacheConfig().value_or(executor::PeftCacheConfig())); - mPeftCacheManager = std::make_shared( - peftCacheManagerConfig, mModelConfig, mWorldConfig, mRuntime->getBufferManager()); - } - else - { - mPeftCacheManager = std::make_shared(); - } - - if (mModelConfig.isRnnBased()) - { - createRnnStateManager(); - } - if (mModelConfig.isTransformerBased() && modelConfig.isKVCacheEnabled()) - { - auto cacheTransceiverConfig - = executorConfig.getCacheTransceiverConfig().value_or(executor::CacheTransceiverConfig()); - - auto const cacheSizeBytesPerTokenPerWindow = calculateCacheSizePerTokenForDisagg( - mModelConfig, mWorldConfig, getMaxAttentionWindowVec(), mModelConfig.useCrossAttention(), 2); - auto cacheTransPreAllocaSize = kv_cache_manager::CacheTransBufferManager::preAllocBufferSize( - cacheSizeBytesPerTokenPerWindow, mModelConfig.getTokensPerBlock(), cacheTransceiverConfig); - - auto const [freePrimaryMemBytes, freeSecondaryMemBytes] - = BaseKVCacheManager::calculateFreeMemBytes(mRuntime->getBufferManager(), kvCacheConfig); - if (mModelConfig.useCrossAttention()) - { - TLLM_CHECK_WITH_INFO(kvCacheConfig.getCrossKvCacheFraction().has_value(), - "Must set crossKvCacheFraction for encoder-decoder model"); - auto const crossKvCacheFraction = kvCacheConfig.getCrossKvCacheFraction().value(); - mKvCacheManager = createKvCacheManager(kvCacheConfig, KvCacheType::kSELF, - freePrimaryMemBytes * (1.0f - crossKvCacheFraction), - freeSecondaryMemBytes * (1.0f - crossKvCacheFraction), cacheTransPreAllocaSize, - executorConfig.getFailFastOnAttentionWindowTooLarge()); - mCrossKvCacheManager = createKvCacheManager(kvCacheConfig, KvCacheType::kCROSS, - freePrimaryMemBytes * crossKvCacheFraction, freeSecondaryMemBytes * crossKvCacheFraction, - cacheTransPreAllocaSize, executorConfig.getFailFastOnAttentionWindowTooLarge()); - TLLM_LOG_INFO("This is an Encoder-Decoder model, set %0.1f cross KV cache fraction based on the config.", - crossKvCacheFraction); - } - else - { - TLLM_CHECK_WITH_INFO(!kvCacheConfig.getCrossKvCacheFraction().has_value(), - "Do not set crossKvCacheFraction for decoder-only model"); - mKvCacheManager = createKvCacheManager(kvCacheConfig, KvCacheType::kSELF, freePrimaryMemBytes, - freeSecondaryMemBytes, cacheTransPreAllocaSize, executorConfig.getFailFastOnAttentionWindowTooLarge()); - } - - mCacheTransceiver - = CacheTransceiverFactory::createCacheTransceiver(mKvCacheManager.get(), mModelConfig, mWorldConfig, - executor::kv_cache::CacheState::AttentionType::kDEFAULT, executorConfig.getCacheTransceiverConfig()); - mDisaggTransferAdmissionController = std::make_unique( - cacheTransceiverConfig.getMaxTokensInBuffer(), mModelConfig.getTokensPerBlock()); - } - - if (mModelConfig.getSpeculativeDecodingMode().needsKVCacheRewind()) - { - TLLM_CHECK_WITH_INFO( - mModelConfig.isKVCacheEnabled(), "When needsKVCacheRewind() returns true, KV cache needs to be enabled."); - auto const& blockManager = mKvCacheManager->getBlockManager(); - - TLLM_CHECK_WITH_INFO(blockManager.getNumPools() == 1, - "Rewinding KV cache blocks for models with multiple pools is not supported"); - - // Two "redundant" checks given the pool size check above, but those below don't rely on an implementation - // detail I guess. - TLLM_CHECK_WITH_INFO( - !blockManager.isVariableWindow(), "Rewinding KV cache blocks for variable SWA models isn't supported"); - auto const maxBlocksPerSeq = blockManager.getMaxBlockPerSeqWhenSingleWindowSize(); - - // TODO(oargov): VGQA is not supported, assume all layers have the same num_kv_heads - TLLM_CHECK_WITH_INFO( - !blockManager.isVariableGQA(), "Rewinding KV cache blocks for variable GQA models isn't supported"); - auto const numKvHeads = mModelConfig.getNbKvHeads(0); - - mRewindInputs = RewindInputs{maxBlocksPerSeq, /*isUseOneMoreBlock*/ false, numKvHeads}; - } - - if (mWorldConfig.isPipelineParallel()) - { - mAsyncSendWaitThread = std::make_unique( - "asyncSendWaitThread", - [this]() - { - mDecStepAsyncSndHdls.clear(); - mDecSlotAsyncSndHdls.clear(); - }, - [this]() { TLLM_CUDA_CHECK(cudaSetDevice(mWorldConfig.getDevice())); }); - - auto const& commSession = COMM_SESSION; - mMpiCommPipelinePara = std::make_unique( - commSession.split(mWorldConfig.getTensorParallelRank(), mWorldConfig.getPipelineParallelRank())); - mDecSlotAsyncSndHdls.reserve(getMaxBatchSize()); - } - if (mWorldConfig.isTensorParallel()) - { - auto const& commSession = COMM_SESSION; - mMpiCommTensorPara = std::make_unique( - commSession.split(mWorldConfig.getPipelineParallelRank(), mWorldConfig.getTensorParallelRank())); - } - - mSeqSlotManager - = std::make_shared(getMaxNumSequences(), executorConfig.getMaxSeqIdleMicroseconds()); - - mMicroBatchScheduledRequests.resize(mNumMicroBatches); - mDecoderFinishedEvents.resize(mNumMicroBatches); - mPeftTables.resize(mNumMicroBatches); - - if (modelConfig.isRnnBased()) - { - TLLM_CHECK_WITH_INFO(modelConfig.getMaxBeamWidth() == 1, "RNN based model doesn't support beam search now."); - TLLM_CHECK_WITH_INFO( - !executorConfig.getEnableChunkedContext(), "RNN based model doesn't support Chunked Context now."); - TLLM_CHECK_WITH_INFO( - modelConfig.getSpeculativeDecodingMode().isNone(), "RNN based model doesn't support speculative decoding."); - } - - std::optional ctxChunkConfig; - if (executorConfig.getEnableChunkedContext()) - { - TLLM_CHECK_WITH_INFO(modelConfig.isKVCacheEnabled() && mModelConfig.getPagedContextFMHA(), - "Chunked context requires context FMHA, paged kv_cache and paged context FMHA all enabled at the same " - "time."); - SizeType32 chunkUnitSize = mKvCacheManager->getTokensPerBlock(); - // If sliding window attention is used, then make sure the unit size aligns with the paged context fmha's kv - // step size. - if (getMaxInputLen() > getMaxAttentionWindow()) // TODO(nhaber): minAttentionWindow - { - chunkUnitSize = std::max(/* maxKvStepSizeInFmha */ 256, chunkUnitSize); - TLLM_LOG_INFO("ChunkUnitSize is set to %d as sliding window attention is used.", chunkUnitSize); - } - ctxChunkConfig = batch_scheduler::ContextChunkingConfig{ - executorConfig.getSchedulerConfig().getContextChunkingPolicy().value_or( - executor::ContextChunkingPolicy::kFIRST_COME_FIRST_SERVED), - chunkUnitSize}; - } - - auto maxNumTokens = getMaxNumTokens(); - TLLM_CHECK_WITH_INFO(maxNumTokens, "Max number of tokens is not set in model config."); - - // Max context size is limited by `max_num_tokens` for chunked-context or context-FMHA, - // or by `max_input_len` of the model. - auto const maxContextLength = (executorConfig.getEnableChunkedContext() || mModelConfig.getContextFMHA()) - ? maxNumTokens - : std::make_optional(mModelConfig.getMaxInputLen()); - - mMaxBatchSizeTunerRecommended = 0; - mMaxBatchSizeRuntime = getMaxBatchSize(); - mMaxNumTokensStatic = maxNumTokens; - mMaxNumTokensTunerRecommended = 0; - mMaxNumTokensRuntime = maxNumTokens; - - if (mKvCacheManager && ctxChunkConfig) - { - TLLM_CHECK_WITH_INFO(ctxChunkConfig.value().chunkUnitSize % mKvCacheManager->getTokensPerBlock() == 0, - "To prevent cache fragmentation, the context chunk unit size (%d) should be divisible by the number of " - "tokens per kv-cache block (%d).", - ctxChunkConfig.value().chunkUnitSize, mKvCacheManager->getTokensPerBlock()); - } - - mCapacityScheduler = std::make_unique(getMaxNumSequences(), - executorConfig.getSchedulerConfig().getCapacitySchedulerPolicy(), mKvCacheManager != nullptr, - /*twoStepsLookAhead=*/mWorldConfig.isPipelineParallel(), - /*noScheduleUntilState=*/LlmRequestState::kCONTEXT_INIT, - /*noScheduleAfterState=*/LlmRequestState::kGENERATION_COMPLETE, - /*enablePrefixAwareScheduling=*/executorConfig.getSchedulerConfig().getEnablePrefixAwareScheduling()); - - mMicroBatchScheduler = std::make_unique(ctxChunkConfig, maxContextLength); - - if (ctxChunkConfig) - { - if (maxContextLength) - { - ctxChunkConfig.value().chunkUnitSize - = std::min(ctxChunkConfig.value().chunkUnitSize, maxContextLength.value()); - } - TLLM_CHECK_WITH_INFO(ctxChunkConfig.value().chunkUnitSize > 0, - "Context chunk size (%d) must be a positive integer.", maxContextLength.value()); - } - else - { - if (maxContextLength && maxNumTokens) - { - TLLM_CHECK_WITH_INFO(maxContextLength.value() <= maxNumTokens.value(), - "Without enabling chunked context, the max context length (%d) needs to be less than or equal to the " - "max number of tokens (%d).", - maxContextLength.value(), maxNumTokens.value()); - } - } - - mPauseRequests = std::make_unique(getMaxInputLen()); - mAssignReqSeqSlots = std::make_unique(); - mAllocateKvCache = std::make_unique(); - - if (isCudaGraphMode()) - { - // Limit cuda graph cache size. Depending on the model one graph is 4-10MB of GPU memory. - SizeType32 cudaGraphCacheSize - = std::min(getMaxBatchSize(), std::max(mExtendedRuntimePerfKnobConfig.getCudaGraphCacheSize(), 1)); - // We can't have common cache for all microbatches as cuda graph is tied to the memory pointers of the runtime - // buffers. - mCudaGraphExecutorCaches.resize(mNumBuffers, utils::CudaGraphExecutorCache(cudaGraphCacheSize)); - } - - mSpeculativeDecodingFastLogits - = executorConfig.getSpecDecConfig().has_value() && executorConfig.getSpecDecConfig()->fastLogits; - if (mSpeculativeDecodingFastLogits && modelConfig.getSpeculativeDecodingMode().isNone() && mIsLeaderInOrchMode) - { - mDraftModelSendLogitsThread - = std::make_unique(&utils::draftModelSendLogitsThread, mDevice, &mDraftModelThreadShouldExit, - &mDraftRequestsWaitingToSendLogits, &mDraftRequestsDoneSendingLogits, &mDraftRequestsMtx); - } - - mCreateNewDecoderRequests = std::make_unique( - mSpeculativeDecodingFastLogits, mIsLeaderInOrchMode, isNormalizeLogProbs()); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -TrtGptModelInflightBatching::~TrtGptModelInflightBatching() -{ - if (mCacheTransceiver) - { - mCacheTransceiver->checkContextTransferStatus(1, true); - TLLM_CHECK_WITH_INFO(mCacheTransceiver->checkGenTransferComplete(), "Generation transfer not complete"); - } - if (mAsyncSendWaitThread) - { - mAsyncSendWaitThread.reset(nullptr); - } - if (mDraftModelSendLogitsThread) - { - mDraftModelThreadShouldExit = true; - mDraftModelSendLogitsThread->join(); - mDraftModelSendLogitsThread.reset(nullptr); - } -} - -void TrtGptModelInflightBatching::setupSpeculativeDecodingModule(executor::DecodingConfig const& decodingConfig) -{ - if (mModelConfig.getSpeculativeDecodingMode().isExplicitDraftTokens() - || mModelConfig.getSpeculativeDecodingMode().isEagle()) - { - TLLM_CHECK_WITH_INFO(mCtxGenFusion, "Current speculative decoding mode requires context-gen fusion IFB"); - } - - if (mModelConfig.getSpeculativeDecodingMode().isLookaheadDecoding() && decodingConfig.getLookaheadDecodingConfig()) - { - // FIXME choose defaults - auto maxLookaheadConfig = decodingConfig.getLookaheadDecodingConfig().value(); - - SizeType32 maxDraftTokens{0}; - SizeType32 maxDraftPathLen{0}; - std::tie(std::ignore, std::ignore, maxDraftTokens, maxDraftPathLen) - = maxLookaheadConfig.calculateSpeculativeResource(); - TLLM_CHECK(maxDraftTokens <= mModelConfig.getMaxDecodingDraftTokens()); - mModelConfig.getSpeculativeDecodingModulePtr()->setMaxDraftTokens(maxDraftTokens); - mModelConfig.getSpeculativeDecodingModulePtr()->setMaxDraftPathLen(maxDraftPathLen); - - auto lookaheadModulePtr - = std::dynamic_pointer_cast(mModelConfig.getSpeculativeDecodingModulePtr()); - lookaheadModulePtr->setExecutionConfig(maxLookaheadConfig); - } -} - -void TrtGptModelInflightBatching::reshapeKvTensors(OffsetTableDimensions const& dims) -{ - TLLM_CHECK(mBuffers.size() == static_cast(mNumBuffers)); - auto const& manager = mRuntime->getBufferManager(); - for (auto& buffers : mBuffers) - { - TLLM_CHECK(buffers->transformerBuffers); - // any method that operates on transformerBuffers must distinguish between self and cross cache, because - // transformerBuffers is not managed by KVCacheManager same rule applies to kv pool pointers below - buffers->transformerBuffers->reshapeKvTensors( - getMaxBatchSize(), mOperatingBeamWidth, dims.maxBlocksPerSeq, dims.cacheType, dims.numPools, manager); - } -} - -using BlocksPerWindow = std::map>; - -std::pair> -TrtGptModelInflightBatching::clampWindowSizesToFitAtLeastOneSequence( - BlocksPerWindow const& blocksPerWindow, bool const failFastOnAttentionWindowTooLarge) -{ - // At this point, we can only validate that the cheapest sequence in terms of kv-cache resources still fits. More - // validation is needed on a per-request basis, once the prompt / output lengths and the actual beam width are - // known. - auto const promptLength = getMaxInputLen(); - auto const outputLength - = getMaxSequenceLen() - promptLength; // This makes it the best case scenario, as context tokens are 'cheaper' - // in terms of kv-cache resources on average. - auto const sinkTokenLength = getSinkTokenLen(); - auto const maxBeamWidth = getMaxBeamWidth(); - auto const tokensPerBlock = mModelConfig.getTokensPerBlock(); - auto const& oldMaxAttentionWindowVec = getMaxAttentionWindowVec(); - std::vector newMaxAttentionWindowVec; - BlocksPerWindow newBlocksPerWindow; - - newMaxAttentionWindowVec.reserve(oldMaxAttentionWindowVec.size()); - for (auto const windowSize : oldMaxAttentionWindowVec) - { - auto const bestCaseBlockRequirements = kv_cache_manager::KVCacheManager::calculateMaxBlockRequirements( - promptLength, outputLength, sinkTokenLength, windowSize, maxBeamWidth, tokensPerBlock); - auto const [numPrimaryBlocks, numSecondaryBlocks] = blocksPerWindow.at(windowSize); - if (bestCaseBlockRequirements > numPrimaryBlocks) - { - auto const newMaxAttentionWindow = KVCacheManager::calculateMaxAttentionWindow( - promptLength, outputLength, sinkTokenLength, numPrimaryBlocks, maxBeamWidth, tokensPerBlock); - newMaxAttentionWindowVec.push_back(newMaxAttentionWindow); - newBlocksPerWindow[newMaxAttentionWindow] = std::make_tuple(numPrimaryBlocks, numSecondaryBlocks); - } - else - { - newMaxAttentionWindowVec.push_back(windowSize); - newBlocksPerWindow[windowSize] = std::make_tuple(numPrimaryBlocks, numSecondaryBlocks); - } - } - if (newMaxAttentionWindowVec == getMaxAttentionWindowVec()) - { - return {blocksPerWindow, newMaxAttentionWindowVec}; - } - TLLM_LOG_WARNING("maxAttentionWindowVec too large to fit at least one sequence in kvCache. Old: %s, New: %s", - common::vec2str(getMaxAttentionWindowVec()).c_str(), common::vec2str(newMaxAttentionWindowVec).c_str()); - - if (failFastOnAttentionWindowTooLarge) - { - throw std::runtime_error( - "Attention window too large to fit even a single sequence in the KV cache. Failing fast rather than " - "attempting an adjustment of the window sizes. " - "Old: " - + common::vec2str(getMaxAttentionWindowVec()) + ", New: " + common::vec2str(newMaxAttentionWindowVec)); - } - - setMaxAttentionWindowVec(newMaxAttentionWindowVec); - if (getMaxSequenceLen() > getMaxAttentionWindow()) - { - TLLM_LOG_WARNING("maxSequenceLen is reduced to maxAttentionWindow: %d", getMaxAttentionWindow()); - setMaxSequenceLen(getMaxAttentionWindow()); - if (getMaxInputLen() > getMaxSequenceLen() - 1) - { - setMaxInputLen(getMaxSequenceLen() - 1); - TLLM_LOG_WARNING("maxInputLen is reduced to %d", getMaxInputLen()); - } - } - // createBuffers depends on: - // maxAttentionWindow; maxAttentionWindowVec; maxSequenceLen; - // TODO: This is problematic, as createBuffers edits the state of trtGptModelInflightBatching, but - // what if there are different window values for cross+self etc. in encoder+decoder scenario... - createBuffers(mDecodingConfig, mAdditionalModelOutputs); - createDecoder(mDecodingConfig.getDecodingMode()); - return {newBlocksPerWindow, newMaxAttentionWindowVec}; -} - -std::unique_ptr TrtGptModelInflightBatching::createKvCacheManager( - KvCacheConfig const& kvCacheConfig, KvCacheType kvCacheType, uint64_t freePrimaryMemBytes, - uint64_t freeSecondaryMemBytes, size_t extraCostMemory, bool const failFastOnAttentionWindowTooLarge) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - bool isCrossAttention = kvCacheType == KvCacheType::kCROSS; - TLLM_CHECK_WITH_INFO( - mModelConfig.isTransformerBased(), "KvCacheManager is only needed by transformer based model."); - - auto const tokensPerBlock = mModelConfig.getTokensPerBlock(); - auto const kvDtype = mModelConfig.getKvDataType(); - - // init KV cache block manager - auto [numKvHeadsPerLayerBegin, numKvHeadsPerLayerEnd] = mModelConfig.getNumKvHeadsPerLayerLocalRange( - mWorldConfig.getPipelineParallelism(), mWorldConfig.getPipelineParallelRank(), isCrossAttention); - auto numKvHeadsPerLayer = std::vector(numKvHeadsPerLayerBegin, numKvHeadsPerLayerEnd); - - auto maxAttentionWindowVec = getMaxAttentionWindowVec(); - if (kvCacheType != KvCacheType::kSELF) // TODO(nhaber): more foolproof way of initing cross-kvcache-manager - { - maxAttentionWindowVec = std::vector{mModelConfig.getMaxEncoderLen()}; - } - - auto const numLayers = static_cast(numKvHeadsPerLayer.size()); - auto const windowSizeToLayers = KVCacheManager::groupLayersByWindowSize(maxAttentionWindowVec, numLayers); - auto const sizePerHead = mModelConfig.getSizePerHead(); - auto blocksPerWindow = KVCacheManager::calculateMaxNumBlocks(kvCacheConfig, kvDtype, numKvHeadsPerLayer, - sizePerHead, tokensPerBlock, mWorldConfig, windowSizeToLayers, freePrimaryMemBytes, freeSecondaryMemBytes, - extraCostMemory, 2, getMaxBatchSize()); - - // now we check if any of the window sizes is too large for at least one sequence to fit in kvCache - // this can happen if e.g. maxSeqLen is deduced from the model and is too large - // and user also didn't provide maxAttentionWindow, which leads it to be equal to maxSeqLen - if (kvCacheType == KvCacheType::kSELF) - { - std::tie(blocksPerWindow, maxAttentionWindowVec) - = clampWindowSizesToFitAtLeastOneSequence(blocksPerWindow, failFastOnAttentionWindowTooLarge); - } - - if (kvCacheType == KvCacheType::kCROSS && kvCacheConfig.getEnableBlockReuse()) - { - TLLM_LOG_INFO( - "Cross KV cache does not support reuse because cross attention depends on encoder and decoder input ids. " - "Thus, KV cache reuse is disabled for cross KV cache."); - } - auto const enableBlockReuse = kvCacheType == KvCacheType::kSELF ? kvCacheConfig.getEnableBlockReuse() : false; - - auto kvCacheManager = std::make_unique(numKvHeadsPerLayer, sizePerHead, tokensPerBlock, - blocksPerWindow, getMaxNumSequences(), getMaxBeamWidth(), maxAttentionWindowVec, kvDtype, getSinkTokenLen(), - mRuntime->getStreamPtr(), - kvCacheType == KvCacheType::kCROSS ? mModelConfig.getMaxEncoderLen() : getMaxSequenceLen(), - getMaxNumTokens().value(), enableBlockReuse, kvCacheType, kvCacheConfig.getSecondaryOffloadMinPriority(), - kvCacheConfig.getEventBufferMaxSize() > 0 - ? std::make_unique(kvCacheConfig.getEventBufferMaxSize()) - : nullptr, - kvCacheConfig.getEnablePartialReuse(), kvCacheConfig.getCopyOnPartialReuse()); - - reshapeKvTensors(kvCacheManager->getOffsetTableDimensions()); - - kvCacheManager->allocatePools(kvCacheConfig.getUseUvm()); - - TensorMap inputBuffers; - TensorPtr poolPointers = kvCacheManager->getBlockPoolPointers(); - TensorPtr poolMapping = kvCacheManager->getLayerToPoolMapping(); - - if (kvCacheType == KvCacheType::kSELF) - { - inputBuffers.insert_or_assign("host_kv_cache_pool_pointers", std::move(poolPointers)); - inputBuffers.insert_or_assign("host_kv_cache_pool_mapping", std::move(poolMapping)); - } - else - { - inputBuffers.insert_or_assign("host_cross_kv_cache_pool_pointers", std::move(poolPointers)); - inputBuffers.insert_or_assign("host_cross_kv_cache_pool_mapping", std::move(poolMapping)); - } - mRuntime->setStaticInputTensors(inputBuffers); - - // Emit the `created` event - kvCacheManager->flushIterationEvents(); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); - return kvCacheManager; -} - -void TrtGptModelInflightBatching::createRnnStateManager() -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - TLLM_CHECK_WITH_INFO(mModelConfig.isRnnBased(), "RnnStateManager is only needed by RNN based model."); - - mRnnStateManager = std::make_unique( - getMaxNumSequences(), mModelConfig, mWorldConfig, mRuntime->getBufferManager()); - - TensorMap inputBuffers; - mRnnStateManager->getPtrBuffers(inputBuffers, mModelConfig, mWorldConfig); - mRuntime->setStaticInputTensors(inputBuffers); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TrtGptModelInflightBatching::createCustomAllReduceWorkspace() -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - TLLM_CHECK(mWorldConfig.isTensorParallel()); - - auto const& manager = mRuntime->getBufferManager(); - auto const hiddenSize = mModelConfig.getHiddenSize(); - - mAllReduceBuffers = std::make_unique(getMaxBatchSize(), getMaxBeamWidth(), getMaxSequenceLen(), - hiddenSize, manager, mWorldConfig, mRuntime->isUserBufferEnabled()); - - TensorMap inputBuffers; - inputBuffers.insert_or_assign("all_reduce_workspace", mAllReduceBuffers->mAllReduceCommPtrs); - mRuntime->setStaticInputTensors(inputBuffers); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TrtGptModelInflightBatching::createRuntimePerfKnobsTensor( - executor::ExtendedRuntimePerfKnobConfig const& extendedRuntimePerfKnobConfig) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - SizeType32 constexpr perfKnobSize{16}; - mExtendedRuntimePerfKnobsHost = BufferManager::cpu(ITensor::makeShape({perfKnobSize}), nvinfer1::DataType::kINT64); - auto* runtimePerfKnobsHostPtr = bufferCast(*mExtendedRuntimePerfKnobsHost); - std::fill_n(runtimePerfKnobsHostPtr, perfKnobSize, -1); - SizeType32 multiBlockModeVal = extendedRuntimePerfKnobConfig.getMultiBlockMode() ? 1 : 0; - SizeType32 enableContextFMHAFP32AccVal = extendedRuntimePerfKnobConfig.getEnableContextFMHAFP32Acc() ? 1 : 0; - runtimePerfKnobsHostPtr[0] = multiBlockModeVal; - runtimePerfKnobsHostPtr[1] = enableContextFMHAFP32AccVal; - - TensorMap inputBuffers; - inputBuffers.insert_or_assign("host_runtime_perf_knobs", mExtendedRuntimePerfKnobsHost); - mRuntime->setStaticInputTensors(inputBuffers); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TrtGptModelInflightBatching::terminateRequest(LlmRequestPtr const& llmReq, bool pause) -{ - utils::terminateRequest( - *mSeqSlotManager, *llmReq, getMaxInputLen(), mKvCacheManager, mCrossKvCacheManager, mPeftCacheManager, pause); -} - -void TrtGptModelInflightBatching::terminateRequestSync( - LlmRequestPtr const& llmRequest, executor::FinishReason finishReason) -{ - TLLM_LOG_DEBUG("Registering termination for request %lu with finish reason %d", llmRequest->mRequestId, - static_cast(finishReason)); - mReqIdsToTerminate.try_emplace(llmRequest->mRequestId, finishReason); -} - -TrtGptModelInflightBatching::IterationStatsIFB TrtGptModelInflightBatching::fillIterationStats( - ScheduledRequests const& scheduledRequests, RequestVector const& requestsToPause) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(fillIterationStats); - - IterationStatsIFB iterationStatsIfb{mMicroBatchId}; - iterationStatsIfb.numCtxRequests = scheduledRequests.contextRequests.size(); - iterationStatsIfb.numGenRequests = scheduledRequests.generationRequests.size(); - iterationStatsIfb.avgNumDecodedTokensPerIter = 0; - - auto const contextBufferId = mCtxGenFusion ? getFusedBufferId() : getContextBufferId(); - auto const& buffers = mBuffers.at(contextBufferId); - iterationStatsIfb.numCtxTokens = buffers->getNumContextTokens(); - - for (auto const& llmReq : scheduledRequests.contextRequests) - { - iterationStatsIfb.scheduledRequests.insert(llmReq->mRequestId); - } - for (auto const& llmReq : scheduledRequests.generationRequests) - { - iterationStatsIfb.scheduledRequests.insert(llmReq->mRequestId); - iterationStatsIfb.avgNumDecodedTokensPerIter += llmReq->getAvgDecodedTokensPerIter(); - } - if (iterationStatsIfb.numGenRequests > 0) - { - iterationStatsIfb.avgNumDecodedTokensPerIter /= iterationStatsIfb.numGenRequests; - TLLM_LOG_DEBUG( - "iterationStatsIfb.avgNumDecodedTokensPerIter = %.2f", iterationStatsIfb.avgNumDecodedTokensPerIter); - } - for (auto const& llmReq : requestsToPause) - { - iterationStatsIfb.pausedRequests.insert(llmReq->mRequestId); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); - return iterationStatsIfb; -} - -void TrtGptModelInflightBatching::forwardSync() -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE_WITH_NAME(range, "TrtGptModelInflightBatching::forwardSync"); - - TLLM_CUDA_CHECK(cudaSetDevice(mWorldConfig.getDevice())); - - if (!mWorldConfig.isLastPipelineParallelRank()) - { - mAsyncSendWaitThread->waitStop(); - } - - auto& currRequests = mMicroBatchScheduledRequests.at(mMicroBatchId); - - if (!currRequests.empty()) - { - if (!mWorldConfig.isPipelineParallel() || !mWorldConfig.isLastPipelineParallelRank()) - { - for (auto& hdl : mDecStepAsyncSndHdls) - { - TLLM_CHECK_WITH_INFO(hdl.get() == nullptr, "decoderSync handle must be nullptr."); - } - // Wait for decoding for requests in flight for the current micro batch - auto& decoderWaitEvent = mDecoderFinishedEvents.at(mMicroBatchId); - mDecStepAsyncSndHdls = decoderSync(currRequests, decoderWaitEvent); - decoderWaitEvent.reset(); - - if (!mWorldConfig.isLastPipelineParallelRank()) - { - mAsyncSendWaitThread->notifyStart(); - } - } - else - { - for (auto const& requests : {currRequests.contextRequests, currRequests.generationRequests}) - { - for (auto const& llmReq : requests) - { - for (SizeType32 beam = 0; beam < llmReq->mSamplingConfig.beamWidth; ++beam) - { - llmReq->setNumPreDecodedTokens(0, beam); - } - if (llmReq->isGenerationToCompleteState()) - { - llmReq->setState(LlmRequestState::kGENERATION_COMPLETE); - terminateRequest(llmReq); - } - } - } - } - - (*mPauseRequests)(currRequests.generationRequests, mInflightReqIds, mReqIdsToPause, true, *mSeqSlotManager, - mKvCacheManager, mCrossKvCacheManager, mPeftCacheManager); - - if (!mReqIdsToTerminate.empty()) - { - for (auto const& requests : {currRequests.contextRequests, currRequests.generationRequests}) - { - for (auto const& llmReq : requests) - { - if (mReqIdsToTerminate.count(llmReq->mRequestId) != 0U) - { - if (!llmReq->isGenerationCompleteState()) - { - TLLM_LOG_DEBUG("Terminating request %lu with finish reason %d", llmReq->mRequestId, - static_cast(mReqIdsToTerminate[llmReq->mRequestId])); - terminateRequest(llmReq); - llmReq->finishByReason(mReqIdsToTerminate[llmReq->mRequestId]); - llmReq->clearGeneratedTokens(); - } - mReqIdsToTerminate.erase(llmReq->mRequestId); - } - } - } - } - - // Terminate draft requests whose logits have been sent by the background thread. - { - RequestVector doneSending; - { - std::lock_guard lk(mDraftRequestsMtx); - doneSending.swap(mDraftRequestsDoneSendingLogits); - } - for (auto const& llmReq : doneSending) - { - terminateRequest(llmReq); - } - } - - // Finished context requests have been moved to generationRequests by moveFinishedContextRequestsToGeneration - for (auto const& llmReq : currRequests.generationRequests) - { - // If a context-only request is finished, send its KV cache and mark it. - if (llmReq->isContextOnlyRequest() && llmReq->isContextFinished()) - { - // TODO: skip if sending layer-wise - { - TLLM_CHECK_WITH_INFO(mCacheTransceiver, - "Disaggregated serving is not enabled, please check the configuration of " - "cacheTransceiverConfig."); - mCacheTransceiver->respondAndSendAsync(llmReq); - } - mSeqSlotManager->freeSequenceSlot(llmReq->mRequestId); - } - } - } - // report profile data - auto const bufferId = getFusedBufferId(); - auto const contextId = mBuffers[bufferId]->getContextIndex(); - if (mRuntime->hasLayerProfiler(contextId)) - { - mRuntime->reportToProfiler(contextId); - } - if (mCacheTransceiver) - { - mCacheTransceiver->checkContextTransferStatus(0, true); - } - ++mIterCounter; - - if (mKvCacheManager) - { - mKvCacheManager->flushIterationEvents(); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TrtGptModelInflightBatching::storeContextBlocks(std::shared_ptr const& llmReq) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - // TMJ - Note - // Make context blocks reusable immediately after context phase finishes. - // For chunked contexts, this occurs in step that processes last context chunk. - // isLastContextChunk() is always true for non-chunked contexts. - // This check is made in code that calls storeContextBlocks, so omitted here. - if (mKvCacheManager) - { - mKvCacheManager->storeContextBlocks(*llmReq); - } - if (mCrossKvCacheManager) - { - mCrossKvCacheManager->storeContextBlocks(*llmReq); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TrtGptModelInflightBatching::storeNewBlock(std::shared_ptr const& llmReq) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - // TMJ - Note - // Make context blocks reusable immediately after each generation step. - - if (mKvCacheManager) - { - mKvCacheManager->storeNewBlock(*llmReq); - } - if (mCrossKvCacheManager) - { - mCrossKvCacheManager->storeNewBlock(*llmReq); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TrtGptModelInflightBatching::resetIterationStats() -{ - mLastIterationStatsIFB = IterationStatsIFB{mMicroBatchId}; -} - -void TrtGptModelInflightBatching::forwardAsync(RequestList const& activeRequests) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE_WITH_NAME(range, "TrtGptModelInflightBatching::forwardAsync"); - - TLLM_CUDA_CHECK(cudaSetDevice(mWorldConfig.getDevice())); - - try - { - verifyRequests(activeRequests); - if (mModelConfig.isTransformerBased() && getKVCacheManager() && mCacheTransceiver) - { - checkDisaggGenTransferStatus(activeRequests); - } - auto& currRequests = mMicroBatchScheduledRequests.at(mMicroBatchId); - - // Get a new set of requests for that context - // The scheduler will not include any requests that are (i) still in encoder state if encoder-decoder models OR - // (ii) already in flight for decoder models - TLLM_LOG_DEBUG("Running DECODER request scheduler"); - auto [fittingRequests, fittingDisaggGenInitRequests, requestsToPause] - = (*mCapacityScheduler)(activeRequests, mKvCacheManager, mPeftCacheManager, mCrossKvCacheManager); - // Remove from fitting requests the requests that cannot be scheduled due to disagg KV cache transfer - bool waitForDisaggGenTransferProgress = false; - if (mModelConfig.isTransformerBased() && getKVCacheManager() && mCacheTransceiver) - { - if (mDisaggTransferAdmissionController && mDisaggTransferAdmissionController->enabled() - && !fittingDisaggGenInitRequests.empty()) - { - auto admissionResult - = mDisaggTransferAdmissionController->select(activeRequests, fittingDisaggGenInitRequests); - waitForDisaggGenTransferProgress = admissionResult.isBlockedByActiveTransfers(); - if (admissionResult.deferredRequestCount > 0) - { - TLLM_LOG_DEBUG( - "Disagg transfer admission deferred %zu requests; active transfer blocks=%zu, admitted " - "transfer blocks=%zu, budget=%zu", - admissionResult.deferredRequestCount, admissionResult.activeTransferBlocks, - admissionResult.admittedTransferBlocks, - mDisaggTransferAdmissionController->getMaxTransferBlocks().value_or(0)); - } - fittingDisaggGenInitRequests = std::move(admissionResult.admittedRequests); - } - prepareDisaggGenInitRequests(activeRequests, fittingDisaggGenInitRequests); - } - if (fittingRequests.empty() && fittingDisaggGenInitRequests.empty()) - { - TLLM_LOG_WARNING( - "CapacityScheduler didn't schedule any requests in iteration %lu, " - "probably because of insufficient resources such as KV cache, " - "will try wait for KV cache transfer to complete", - mIterCounter); - if (mCacheTransceiver) - { - if (waitForDisaggGenTransferProgress) - { - TLLM_LOG_DEBUG("Waiting for generation KV cache transfer progress to free disagg admission budget"); - mCacheTransceiver->checkGenTransferStatus(1); - } - else - { - mCacheTransceiver->checkContextTransferStatus(1, true); - // will free kvCache in next iteration. - } - } - } - std::tie(currRequests.contextRequests, currRequests.generationRequests) - = (*mMicroBatchScheduler)(fittingRequests, mInflightReqIds, mMaxBatchSizeRuntime, mMaxNumTokensRuntime); - TLLM_CHECK(currRequests.size() <= static_cast(getMaxBatchSize())); - - (*mPauseRequests)(requestsToPause, mInflightReqIds, mReqIdsToPause, false, *mSeqSlotManager, mKvCacheManager, - mCrossKvCacheManager, mPeftCacheManager); - - if (mUseSeamlessLookahead) - { - changeSpecDecMode(currRequests); - } - - if (!currRequests.empty()) - { - TLLM_LOG_DEBUG("Running DECODER model with batch size: %lu", currRequests.size()); - // For overlap don't store inflight requests, so they are not skipped in scheduler - if (!isTrtOverlap()) - { - NVTX3_SCOPED_RANGE(updateInflightReqIds); - // Add requests to in-flight set, so they can be skipped in other micro batches - for (auto const& llmReq : currRequests.contextRequests) - { - // Context requests that are chunking are not added to inflight set, so they are scheduled in the - // next micro batch. - if (llmReq->isLastContextChunk()) - { - TLLM_LOG_DEBUG( - "Context request with ID %lu added to DECODER model inflight set", llmReq->mRequestId); - mInflightReqIds.insert(llmReq->mRequestId); - } - } - for (auto const& llmReq : currRequests.generationRequests) - { - TLLM_LOG_DEBUG( - "Generation request with ID %lu added to DECODER model inflight set", llmReq->mRequestId); - mInflightReqIds.insert(llmReq->mRequestId); - } - } - - (*mAssignReqSeqSlots)(*mSeqSlotManager, currRequests.contextRequests, currRequests.generationRequests); - - if (mKvCacheManager) - { - (*mAllocateKvCache)(*mKvCacheManager, currRequests.contextRequests, currRequests.generationRequests, - mModelConfig, mCrossKvCacheManager); - } - - mPeftTables.at(mMicroBatchId) - = mPeftCacheManager->ensureBatch(currRequests.contextRequests, currRequests.generationRequests, true); - - // Do decoder setup before context phase if model needs to setup buffers for the context phase. - if (mModelConfig.getSpeculativeDecodingMode().needsDecoderPrologue()) - { - auto const contextBufferId = mCtxGenFusion ? getFusedBufferId() : getContextBufferId(); - setupDecoderStep(currRequests.contextRequests, *mBuffers.at(contextBufferId), - mDecoderInputBuffers.at(getFusedBufferId())); - // WAR: Sync to ensure that the decoder setup is complete before the context phase starts. - // Without this, there may be a race condition between the decoder setup and the context phase - // which also leads to spurious test failure in trtGptModelRealDecoderTest. - mRuntime->getStream().synchronize(); - } - else - { - prepareDistGenBufferAndDecoder(currRequests.generationRequests); - } - sync_check_cuda_error(mRuntime->getStream().get()); - - executeBatch(currRequests); - if (mWorldConfig.isLastPipelineParallelRank() && mGuidedDecoder) - { - // XGrammar: build maskcache for context requests and perform maskgen for all requests - // These need to be overlapped with the kernel execution of forward step - mGuidedDecoder->build(currRequests); - } - - sync_check_cuda_error(mRuntime->getStream().get()); - - // Postpone decoder setup if model does not need to setup buffers for the context phase. - if (!mModelConfig.getSpeculativeDecodingMode().needsDecoderPrologue()) - { - auto const contextBufferId = mCtxGenFusion ? getFusedBufferId() : getContextBufferId(); - setupDecoderStep(currRequests.contextRequests, *mBuffers.at(contextBufferId), - mDecoderInputBuffers.at(getFusedBufferId())); - } - - sync_check_cuda_error(mRuntime->getStream().get()); - - if (isTrtOverlap()) - { - // WAR: Because the decoder is not stateless (yet) a sync is needed between - // decoder execution and next decoder step preparation. - auto const prevMicroBatchId = getPrevMicroBatchId(mMicroBatchId); - auto& prevDecoderFinishedEvent = mDecoderFinishedEvents.at(prevMicroBatchId); - if (prevDecoderFinishedEvent) - { - prevDecoderFinishedEvent->synchronize(); - } - } - - auto& decoderFinishedEvent = mDecoderFinishedEvents.at(mMicroBatchId); - TLLM_CHECK_WITH_INFO(!decoderFinishedEvent.has_value(), "decoderFinishedEvent must be nullopt."); - decoderFinishedEvent = mWorldConfig.isLastPipelineParallelRank() - ? std::make_optional(decoderStepAsync(currRequests)) - : std::nullopt; - - sync_check_cuda_error(mRuntime->getStream().get()); - - mLastIterationStatsIFB = fillIterationStats(currRequests, requestsToPause); - for (auto const& requests : {currRequests.contextRequests, currRequests.generationRequests}) - { - for (auto const& llmReq : requests) - { - if (llmReq->isContextInitState()) - { - llmReq->moveToNextContextChunk(); - if (llmReq->getContextRemainingLength() == 0) - { - TLLM_LOG_DEBUG("[RANK %d] request with ID %lu finishes decoder ctx phase", - COMM_SESSION.getRank(), llmReq->mRequestId); - - llmReq->setState(LlmRequestState::kGENERATION_IN_PROGRESS); - - // for encoder-decoder models, free encoder output buffers after decoder context phase is - // completed - if (llmReq->getEncoderTokens().has_value()) - { - llmReq->freeEncoderOutputBuffers(); - } - storeContextBlocks(llmReq); - - if (isTrtOverlap() && llmReq->willCompleteNextIteration()) - { - // This prohibits the request from being scheduled for another iteration if only one - // iteration is expected. - llmReq->setState(LlmRequestState::kGENERATION_TO_COMPLETE); - } - } - } - else if (llmReq->isGenerationInProgressState()) - { - storeNewBlock(llmReq); - TLLM_LOG_DEBUG("request with ID %lu forwards a step in decoder gen phase", llmReq->mRequestId); - } - } - } - - utils::moveFinishedContextRequestsToGeneration(currRequests); - } - else - { - mLastIterationStatsIFB = IterationStatsIFB{mMicroBatchId}; - } - - if (mWorldConfig.isPipelineParallel() && mWorldConfig.isLastPipelineParallelRank()) - { - mAsyncSendWaitThread->waitStop(); - if (!currRequests.empty()) - { - for (auto& hdl : mDecStepAsyncSndHdls) - { - TLLM_CHECK_WITH_INFO(hdl.get() == nullptr, "decoderSync handle must be nullptr."); - } - // Wait for decoding for requests in flight for the current micro batch - auto& decoderFinishedEvent = mDecoderFinishedEvents.at(mMicroBatchId); - mDecStepAsyncSndHdls = decoderSync(currRequests, decoderFinishedEvent); - decoderFinishedEvent.reset(); - - mAsyncSendWaitThread->notifyStart(); - } - } - - // Update the micro batch ID - mMicroBatchId = getNextMicroBatchId(mMicroBatchId); - } - // In case of error, we need to free the batch slot associated with those requests - catch (std::exception const&) - { - try - { - for (auto const& llmReq : activeRequests) - { - // Remove from mInflightReqIds so changeBeamWidth can proceed on the next iteration. - // terminateRequest frees seqSlot/KV cache but does not clean up mInflightReqIds. - mInflightReqIds.erase(llmReq->mRequestId); - terminateRequest(llmReq); - } - // Force buffer/decoder reset to clean up any partial state from the aborted batch - // (e.g. partially-filled cross-KV block offsets from mid-context-chunk processing). - // Guard on mInflightReqIds.empty(): in pipeline-parallel multi-micro-batch mode, - // other micro-batches may still have requests tracked here; changeBeamWidth asserts - // emptiness so we skip the reset and let the next successful forwardAsync iteration - // perform it when the set is clear. - if (mWorldConfig.isLastPipelineParallelRank() && mInflightReqIds.empty()) - { - changeBeamWidth(mOperatingBeamWidth); - } - } - catch (std::exception const& e) - { - TLLM_LOG_ERROR("forwardAsync catch-all catch block that runs `terminateRequest` has failed with:"); - TLLM_LOG_EXCEPTION(e); - TLLM_LOG_ERROR("Rethrowing *outer* exception:"); - } - throw; - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TrtGptModelInflightBatching::setRuntimeBatchSize(SizeType32 runtimeMaxBatchSize) -{ - mMaxBatchSizeTunerRecommended = runtimeMaxBatchSize; - mMaxBatchSizeRuntime = std::min(getMaxBatchSize(), runtimeMaxBatchSize); -} - -SizeType32 TrtGptModelInflightBatching::getRuntimeBatchSize() const -{ - return mMaxBatchSizeRuntime; -} - -void TrtGptModelInflightBatching::setRuntimeMaxNumTokens(SizeType32 runtimeMaxNumTokens) -{ - mMaxNumTokensTunerRecommended = runtimeMaxNumTokens; - mMaxNumTokensRuntime - = (mMaxNumTokensStatic) ? std::min(mMaxNumTokensStatic.value(), runtimeMaxNumTokens) : runtimeMaxNumTokens; -} - -void TrtGptModelInflightBatching::updatePeftCache(std::shared_ptr const& llmRequest) -{ - mPeftCacheManager->addRequestPeft(llmRequest, true); -} - -runtime::BufferManager const& TrtGptModelInflightBatching::getBufferManager() const -{ - return mRuntime->getBufferManager(); -} - -BufferManager::CudaStreamPtr TrtGptModelInflightBatching::getRuntimeStreamPtr() const -{ - return mRuntime->getStreamPtr(); -} - -void TrtGptModelInflightBatching::executeContext(SizeType32 runtimeContextId, SizeType32 bufferId) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(executeContext); - - auto const& currBatchState = mBuffers[bufferId]->getBatchState(); - - bool hasCudaGraph = false; - // If batch state is context only, do not capture/launch graph and execute the engine as is. - if (isCudaGraphMode() && !currBatchState.isAnyContext()) - { - auto cudaGraphOpt = mCudaGraphExecutorCaches[bufferId].get(currBatchState); - // If graph exists for current batch state, launch it. - if (cudaGraphOpt.has_value()) - { - hasCudaGraph = true; - } - } - - // If there is no graph for current state, execute the engine. - if (!hasCudaGraph) - { - auto enqueueSuccessful = mRuntime->executeContext(runtimeContextId); - if (!enqueueSuccessful) - { - throw std::runtime_error("Executing TRT engine failed!"); - } - } - else - { - // Launch graph. - auto cudaGraphOpt = mCudaGraphExecutorCaches[bufferId].get(currBatchState); - cudaGraphOpt.value()->launch(mRuntime->getStream()); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TrtGptModelInflightBatching::setLayerProfiler() -{ - mRuntime->setLayerProfiler(); -} - -std::string TrtGptModelInflightBatching::getLayerProfileInfo() const -{ - return mRuntime->getLayerProfileInfo(); -} - -void TrtGptModelInflightBatching::verifyRequests(RequestList const& activeRequests) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(verifyRequests); - - if (activeRequests.empty()) - { - return; - } - - auto const& firstRequest = activeRequests.front(); - auto const firstRequestId = firstRequest->mRequestId; - auto const firstBeamWidth = firstRequest->mSamplingConfig.beamWidth; - - for (auto const& llmReq : activeRequests) - { - auto const beamWidth = llmReq->mSamplingConfig.beamWidth; - auto const draftLength = llmReq->getNumDraftTokens(); - auto const maxDraftLength = mModelConfig.getMaxDecodingDraftTokens(); - - TLLM_CHECK_WITH_INFO(beamWidth == 1 || draftLength == 0, "Can't use speculative decoding with beam search."); - TLLM_CHECK_WITH_INFO(draftLength <= maxDraftLength, - "Number of draft tokens (%d) is larger than maximum number of draft tokens (%d)", draftLength, - maxDraftLength); - - // FIXME: Remove this check when varying beam width is supported - { - TLLM_CHECK_WITH_INFO(beamWidth == firstBeamWidth, - "All active requests must have same beam width, " - "but request %lu with beam width %d differs from first request %lu with beam width %d", - llmReq->mRequestId, beamWidth, firstRequestId, firstBeamWidth); - } - } - - if (firstBeamWidth != mOperatingBeamWidth) - { - changeBeamWidth(firstBeamWidth); - } - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TrtGptModelInflightBatching::executeBatch(ScheduledRequests const& scheduledRequests) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(executeBatch); - - if (!mCtxGenFusion) - { - if (!scheduledRequests.contextRequests.empty()) - { - auto const bufferId = getContextBufferId(); - executeStep(scheduledRequests.contextRequests, {}, bufferId); - } - if (!scheduledRequests.generationRequests.empty()) - { - auto const bufferId = getGenerationBufferId(); - executeStep({}, scheduledRequests.generationRequests, bufferId); - } - } - else - { - auto const bufferId = getFusedBufferId(); - executeStep(scheduledRequests.contextRequests, scheduledRequests.generationRequests, bufferId); - } - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TrtGptModelInflightBatching::createRuntimeContexts() -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - mRuntime->clearContexts(); - auto const numProfiles = mRuntime->getNbProfiles(); - for (auto i = 0; i < numProfiles; ++i) - { - mRuntime->addContext(i); - } - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -namespace -{ -// TODO: move this somewhere else? -/** - * This function logic is also implemented in tensorrt_llm/_torch/pyexecutor/_util.py get_decoding_mode(). - */ -executor::DecodingMode getDecodingMode(SpeculativeDecodingMode specDecodingMode, - std::optional const& decodingModeOpt, runtime::SizeType32 const beamWidth) -{ - auto getDefaultDecodingMode = [beamWidth](std::optional const& decodingModeOpt) - { - if (decodingModeOpt.has_value() && !decodingModeOpt->isAuto()) - { - return decodingModeOpt.value(); - } - return (beamWidth == 1) ? executor::DecodingMode::TopKTopP() : executor::DecodingMode::BeamSearch(); - }; - - auto decodingMode = getDefaultDecodingMode(decodingModeOpt); - // Variable-Beam-Width-Search (special mode of Beam-Search) is enabled. - if (decodingMode.isBeamSearch() && decodingMode.isUseVariableBeamWidthSearch()) - { - TLLM_LOG_INFO("Variable-Beam-Width-Search is enabled"); - } - // Overwrite decoding mode when beam width is one. - if (beamWidth == 1 && decodingMode.isBeamSearch()) - { - TLLM_LOG_WARNING( - "Beam width is set to 1, but decoding mode is BeamSearch. Overwriting decoding mode to TopKTopP."); - decodingMode = executor::DecodingMode::TopKTopP(); - } - // Overwrite decoding mode when Medusa is used. - if (specDecodingMode.isMedusa() && !decodingMode.isMedusa()) - { - TLLM_LOG_WARNING("Model is Medusa, but decoding mode is not Medusa. Overwriting decoding mode to Medusa."); - decodingMode = executor::DecodingMode::Medusa(); - } - // Overwrite decoding mode when Medusa is not used. - if (!specDecodingMode.isMedusa() && decodingMode.isMedusa()) - { - TLLM_LOG_WARNING("Model is not Medusa, but decoding mode is Medusa. Overwriting decoding mode."); - decodingMode = getDefaultDecodingMode(decodingModeOpt); - } - // Overwrite decoding mode when lookahead decoding is used. - if (specDecodingMode.isLookaheadDecoding() && !decodingMode.isLookahead()) - { - TLLM_LOG_WARNING( - "Model is Lookahead, but decoding mode is not Lookahead. Overwriting decoding mode to Lookahead."); - decodingMode = executor::DecodingMode::Lookahead(); - } - // Overwrite decoding mode when lookahead decoding is not used. - if (!specDecodingMode.isLookaheadDecoding() && decodingMode.isLookahead()) - { - TLLM_LOG_WARNING( - "Model is not built with Lookahead decoding, but decoding mode is Lookahead. Overwriting decoding " - "mode."); - decodingMode = getDefaultDecodingMode(decodingModeOpt); - } - // Overwrite decoding mode when 'explicit draft tokens' is used. - if (specDecodingMode.isExplicitDraftTokens() && !decodingMode.isExplicitDraftTokens()) - { - TLLM_LOG_WARNING( - "Model is built with 'explicit draft tokens' decoding, but decoding mode is something else. Overwriting " - "decoding mode."); - decodingMode = executor::DecodingMode::ExplicitDraftTokens(); - } - // Overwrite decoding mode when 'explicit draft tokens' is not used. - if (!specDecodingMode.isExplicitDraftTokens() && decodingMode.isExplicitDraftTokens()) - { - TLLM_LOG_WARNING( - "Model is not built with 'explicit draft tokens' decoding, but decoding mode is set to it. Overwriting " - "decoding " - "mode to default."); - decodingMode = getDefaultDecodingMode(decodingModeOpt); - } - // Overwrite decoding mode when EAGLE is used. - if (specDecodingMode.isEagle() && !decodingMode.isEagle()) - { - TLLM_LOG_WARNING("Model is Eagle, but decoding mode is not Eagle. Overwriting decoding mode to Eagle."); - decodingMode = executor::DecodingMode::Eagle(); - } - // Overwrite decoding mode when Eagle is not used. - if (!specDecodingMode.isEagle() && decodingMode.isEagle()) - { - TLLM_LOG_WARNING("Model is not Eagle, but decoding mode is Eagle. Overwriting decoding mode."); - decodingMode = getDefaultDecodingMode(decodingModeOpt); - } - if (specDecodingMode.isDraftTokensExternal()) - { - TLLM_LOG_WARNING("Overwriting decoding mode to external draft token"); - decodingMode = executor::DecodingMode::ExternalDraftTokens(); - } - TLLM_LOG_DEBUG("DecodingMode: %s", decodingMode.getName()); - return decodingMode; -} -} // namespace - -void TrtGptModelInflightBatching::createDecoder(std::optional const& decodingModeOpt) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - mDecoderState = std::make_unique(); - - if (mWorldConfig.isLastPipelineParallelRank()) - { - auto decoderType = mRuntime->getEngine().getTensorDataType("logits"); - - auto const decodingMode - = getDecodingMode(mModelConfig.getSpeculativeDecodingMode(), decodingModeOpt, mOperatingBeamWidth); - - if (decodingMode.isExplicitDraftTokens()) - { - // There are no logits in Explicit draft tokens model. - decoderType = mModelConfig.getDataType(); - // Decoder is not instantiated for bf16. We use half to get the same data size - // and explicitly pass dtype to redrafter that has bf16 kernels. - if (decoderType == nvinfer1::DataType::kBF16) - { - decoderType = nvinfer1::DataType::kHALF; - } - } - - mDecoder = std::make_unique(mRuntime->getStreamPtr()); - mDecoder->setup( - decodingMode, getMaxNumSequences(), mOperatingBeamWidth, decoderType, mModelConfig, mWorldConfig); - - mDecoderState->setup(getMaxNumSequences(), mOperatingBeamWidth, getMaxAttentionWindow(), getSinkTokenLen(), - getMaxSequenceLen(), decoderType, mModelConfig, mWorldConfig, mRuntime->getBufferManager()); - - if (!mModelConfig.getSpeculativeDecodingMode().isNone()) - { - mDecoderState->setupSpeculativeDecoding(mModelConfig.getSpeculativeDecodingMode(), - mModelConfig.getMaxDecodingTokens(), decoderType, mModelConfig, mWorldConfig, - mRuntime->getBufferManager()); - } - } - else - { - mDecoderState->setupCacheIndirection( - getMaxNumSequences(), mOperatingBeamWidth, getMaxAttentionWindow(), mRuntime->getBufferManager()); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TrtGptModelInflightBatching::createBuffers(executor::DecodingConfig const& decodingConfig, - std::optional> const& additionalModelOutputs) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - mBuffers.clear(); - for (SizeType32 i = 0; i < mNumBuffers; ++i) - { - mBuffers.emplace_back( - std::make_unique(getMaxBatchSize(), mOperatingBeamWidth, getMaxAttentionWindowVec(), - getMaxAttentionWindow(), getSinkTokenLen(), *mRuntime, mModelConfig, mWorldConfig, decodingConfig, - getGatherGenerationLogits(), getMaxNumTokens(), additionalModelOutputs, mPromptTableOffloading)); - } - - mDecoderInputBuffers.clear(); - mDecoderOutputBuffers.clear(); - for (SizeType32 i = 0; i < mNumMicroBatches; ++i) - { - mDecoderInputBuffers.emplace_back( - getMaxBatchSize(), mModelConfig.getMaxDecodingTokens(), mRuntime->getBufferManager()); - mDecoderInputBuffers.back().setupMedusaLogits(getMaxNumSequences(), mModelConfig); - mDecoderOutputBuffers.emplace_back(getMaxNumSequences(), mOperatingBeamWidth, getMaxSequenceLen(), - mModelConfig.getMaxDecodingTokens(), mRuntime->getBufferManager()); - mDecoderOutputBuffers.back().setupSpeculativeDecoding( - getMaxNumSequences(), mModelConfig.getMaxDecodingTokens(), mModelConfig); - } - - mSlotDecoderBuffers.clear(); - for (SizeType32 i = 0; i < getMaxNumSequences(); ++i) - { - mSlotDecoderBuffers.emplace_back(std::make_unique( - mOperatingBeamWidth, getMaxSequenceLen(), mRuntime->getBufferManager())); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TrtGptModelInflightBatching::prepareDisaggGenInitRequests( - RequestList const& activeRequests, RequestVector& newGenReqs) -{ - NVTX3_SCOPED_RANGE(prepareDisaggGenInitRequests); - - // Allocate KV cache by treating them as context requests - (*mAllocateKvCache)(*mKvCacheManager, newGenReqs, {}, mModelConfig, mCrossKvCacheManager); - - // Initiate KV cache transfer - auto timeStart = std::chrono::steady_clock::now(); - - if (tc::getEnvDisaggBenchmarkGenOnly()) - { - TLLM_LOG_DEBUG("Disaggregated generation only benchmark mode is enabled"); - for (auto& req : newGenReqs) - { - req->setState(LlmRequestState::kDISAGG_GENERATION_TRANS_COMPLETE); - } - return; - } - - auto const genInitReqNum = std::count_if(activeRequests.begin(), activeRequests.end(), - [](auto const& req) { return req->isDisaggGenerationInitState(); }); - - // Loop over the new disagg gen requests and trigger receive of KV cache - for (auto& newGenReq : newGenReqs) - { - TLLM_CHECK_WITH_INFO( - mCacheTransceiver, "Disaggregated serving is not enabled, please check the configuration."); - if (common::getEnvDisableKVCacheTransferOverlap()) - { - mCacheTransceiver->requestAndReceiveSync(newGenReq); - } - else - { - mCacheTransceiver->requestAndReceiveAsync(newGenReq); - } - } - if (!common::getEnvDisableKVCacheTransferOverlap()) - { - auto const blockTransfer = std::all_of(activeRequests.begin(), activeRequests.end(), - [](auto const& req) { return req->isDisaggGenerationTransmissionInProgress(); }); - TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), - "newGenReqs.size():%ld requests, activeRequests.size():%ld allTransferInProgress:%d original " - "gen_only_requests_num:%ld", - newGenReqs.size(), activeRequests.size(), blockTransfer, genInitReqNum); - mCacheTransceiver->checkGenTransferStatus(0); - auto timeEnd = std::chrono::steady_clock::now(); - auto duration = std::chrono::duration(timeEnd - timeStart).count(); - TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), - "receiveDisaggGenCache time:%f ms, " - "blockTransfer:%d,genInitReqNum:%ld,newGenReqs.size():%ld,activeRequests.size():%ld", - duration, blockTransfer, genInitReqNum, newGenReqs.size(), activeRequests.size()); - } - - return; -} - -void TrtGptModelInflightBatching::checkDisaggGenTransferStatus(RequestList const& activeRequests) -{ - NVTX3_SCOPED_RANGE(checkDisaggGenTransferStatus); - - if (common::getEnvDisableKVCacheTransferOverlap()) - { - return; - } - - auto timeStart = std::chrono::steady_clock::now(); - - // TODO: - auto const needCheck = std::any_of(activeRequests.begin(), activeRequests.end(), - [](auto const& req) { return req->isDisaggGenerationTransmissionInProgress(); }); - - if (needCheck) - { - mCacheTransceiver->checkGenTransferStatus(0); - - auto timeEnd = std::chrono::steady_clock::now(); - auto duration = std::chrono::duration(timeEnd - timeStart).count(); - TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), - "no Prepare checkDisaggGenTransferStatus time:%f ms, " - "needCheck:%d,activeRequests.size():%ld", - duration, needCheck, activeRequests.size()); - } -} - -void TrtGptModelInflightBatching::prepareDistGenBufferAndDecoder(RequestVector const& generationRequests) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - // set decoderStep for disagg_generation - RequestVector cacheTransCompleteRequests; - for (auto const& request : generationRequests) - { - if (request->isDisaggGenerationTransmissionComplete()) - { - cacheTransCompleteRequests.push_back((request)); - } - } - if (!cacheTransCompleteRequests.empty()) - { - auto timeStart = std::chrono::steady_clock::now(); - auto const bufferId = getFusedBufferId(); - auto& runtimeBuffers = *mBuffers[bufferId]; - runtimeBuffers.prepareStep(cacheTransCompleteRequests, {}, getMaxBeamWidth(), getMaxAttentionWindow(), - *mDecoderState, mKvCacheManager.get(), mCrossKvCacheManager.get(), mRnnStateManager.get(), - mPeftTables[mMicroBatchId], *mRuntime, mModelConfig, mWorldConfig, getGatherGenerationLogits(), - isTrtOverlap()); - auto const contextBufferId = mCtxGenFusion ? getFusedBufferId() : getContextBufferId(); - setupDecoderStep( - cacheTransCompleteRequests, *mBuffers.at(contextBufferId), mDecoderInputBuffers.at(getFusedBufferId())); - sync_check_cuda_error(mRuntime->getStream().get()); - auto timeEnd = std::chrono::steady_clock::now(); - auto duration = std::chrono::duration(timeEnd - timeStart).count(); - TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), - "prepareDistGenBufferAndDecoder time:%f ms , cacheTransCompleteRequests.size():%ld", duration, - cacheTransCompleteRequests.size()); - } - for (auto& request : cacheTransCompleteRequests) - { - request->setState(LlmRequestState::kGENERATION_IN_PROGRESS); - request->setContextCurrentPosition(request->mPromptLen); - request->setDecodingIter(1); - auto const reqBeamWidth = request->mSamplingConfig.beamWidth; - auto firstGenTokens = request->getContextPhaseParams().value().getFirstGenTokens(); - for (SizeType32 beam = 0; beam < reqBeamWidth; ++beam) - { - request->addNewToken(firstGenTokens.at(beam), beam); - } - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TrtGptModelInflightBatching::debugIOTensors(RequestVector const& contextRequests, - RequestVector const& generationRequests, TensorMap const& inputMap, TensorMap const& outputMap) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - TLLM_CHECK(mDebugConfig); - - auto const& manager = mRuntime->getBufferManager(); - auto requestIds = utils::collectRequestIds(contextRequests, generationRequests); - - if (mDebugConfig->getDebugTensorsMaxIterations() > 0) - { - mLastIterationDebugTensors.clear(); - mLastIterationDebugTensors = utils::storeIOTensors(*mDebugConfig, requestIds, inputMap, outputMap, manager); - } - else - { - utils::dumpIOTensors(*mDebugConfig, mIterCounter, requestIds, inputMap, outputMap, mWorldConfig, manager); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -std::tuple const&, runtime::StringPtrMap&> -TrtGptModelInflightBatching::prepareBuffers( - RequestVector const& contextRequests, RequestVector const& generationRequests, SizeType32 bufferId) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(prepareBuffers); - - auto& runtimeBuffers = *mBuffers.at(bufferId); - - auto allNewTokens = mWorldConfig.isLastPipelineParallelRank() - ? RuntimeBuffers::OptionalRef(mDecoderState->getAllNewTokens()) - : std::nullopt; - - auto [optProfileId, inputMap, outputMap] = runtimeBuffers.prepareStep(contextRequests, generationRequests, - mOperatingBeamWidth, getMaxAttentionWindow(), *mDecoderState, mKvCacheManager.get(), mCrossKvCacheManager.get(), - mRnnStateManager.get(), mPeftTables[bufferId], *mRuntime, mModelConfig, mWorldConfig, - getGatherGenerationLogits(), isTrtOverlap(), allNewTokens); - - // For Variable-Beam-Width-Search - mRuntime->setCurrentBeamWidths( - tensorrt_llm::batch_manager::utils::getRequestBeamWidths(contextRequests, generationRequests)); - - mRuntime->setInputTensors(optProfileId, inputMap); - mRuntime->setOutputTensors(optProfileId, outputMap); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); - return {optProfileId, inputMap, outputMap}; -} - -void TrtGptModelInflightBatching::prepareGraph(SizeType32 bufferId, SizeType32 optProfileId) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(prepareGraph); - - auto const nextBatchState = mBuffers[bufferId]->getBatchState(); - auto cudaGraphOpt = mCudaGraphExecutorCaches[bufferId].get(nextBatchState); - // If graph is not found in the cache, capture it. - if (!cudaGraphOpt.has_value()) - { - // We need to prepare some tensors once again to properly set values for graph capture. - // Graph capture requires setting some tensors (e.g. past_kv_len) - // to the round_up(max_kv_cache_len, kKV_CACHE_LEN_CUDA_GRAPH_ROUND_SIZE) - // in order to capture the kernels with the large enough grid. - mBuffers[bufferId]->prepareBuffersForCudaGraph(getMaxSequenceLen()); - - auto cudaGraph = std::make_shared(); - cudaGraph->prepareNextGraph(mRuntime, optProfileId); - mCudaGraphExecutorCaches[bufferId].put(nextBatchState, cudaGraph); - } - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TrtGptModelInflightBatching::executeStep( - RequestVector const& contextRequests, RequestVector const& generationRequests, SizeType32 bufferId) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE_WITH_NAME(range, - "executeStep: " + std::to_string(contextRequests.size()) + " ctx reqs, " - + std::to_string(generationRequests.size()) + " gen reqs"); - - if (mPromptTableOffloading) - { - prefetchNextPromptTableChunk(contextRequests, /* isFirstChunk */ true, bufferId); - } - - auto [optProfileId, inputMap, outputMap] = prepareBuffers(contextRequests, generationRequests, bufferId); - - if (mBuffers[bufferId]->transformerBuffers) - { - // Creation of context progress, or remains nullptr if not needed - std::shared_ptr progress = nullptr; - RequestVector layerWiseRequests; - if (common::getEnvDisaggLayerwise()) - { - for (auto const& request : contextRequests) - { - bool const enableLayerWise = request->isContextOnlyRequest() && request->isLastContextChunk(); - if (enableLayerWise) - { - layerWiseRequests.push_back(request); - } - } - } - // TODO: support layer-wise cross kv cache in encoder-decoder models - if (!layerWiseRequests.empty() && !mModelConfig.useCrossAttention()) - { - int const numLayers = mModelConfig.getNbAttentionLayers( - mWorldConfig.getPipelineParallelism(), mWorldConfig.getPipelineParallelRank()); - progress = std::make_shared(numLayers); - } - bufferCast(*mBuffers[bufferId]->transformerBuffers->contextProgressHost)[0] = progress.get(); - if (progress) - { - TLLM_CHECK_WITH_INFO(mCacheTransceiver, - "Disaggregated serving is not enabled, please check the configuration of cacheTransceiverConfig."); - mCacheTransceiver->respondAndSendLayerWise(layerWiseRequests, progress); - } - } - - if (mPromptTableOffloading) - { - prefetchNextPromptTableChunk(contextRequests, /* isFirstChunk */ false, bufferId); - } - - executeContext(optProfileId, bufferId); - - // If batch state has any context request, do not capture this graph. - if (isCudaGraphMode() && contextRequests.empty()) - { - // Capture graph of current batch state during engine execution. - // This is based on the assumptions that - // a) We can hide CPU graph capture behind the GPU engine execution. - // b) Batch size in the next iterations won't change and we can reuse the graph multiple times. - prepareGraph(bufferId, optProfileId); - } - - if (mDebugConfig) - { - debugIOTensors(contextRequests, generationRequests, inputMap, outputMap); - } - - if (mAdditionalModelOutputs.has_value() && !mAdditionalModelOutputs.value().empty()) - { - utils::copyAdditionalOutputs( - mAdditionalModelOutputs.value(), contextRequests, generationRequests, outputMap, getBufferManager()); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TrtGptModelInflightBatching::setupDecoderStep( - RequestVector const& contextRequests, RuntimeBuffers const& buffers, DecoderInputBuffers& inputBuffers) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(setupDecoderStep); - - if (mWorldConfig.isLastPipelineParallelRank() && !contextRequests.empty()) - { - auto const logitsType = mRuntime->getEngine().getTensorDataType("logits"); - - auto [batchSlots, samplingConfigs, lookaheadPrompt, lookaheadAlgoConfigs] - = (*mCreateNewDecoderRequests)(mModelConfig, mWorldConfig, mDecodingConfig, contextRequests, logitsType, - inputBuffers, *mDecoderState, mRuntime->getStream(), *mDecoder->getDecoderStream(), getMaxSequenceLen(), - mOperatingBeamWidth, buffers.mMedusaBuffers); - - auto const localBatchSize = batchSlots->getSize(); - if (localBatchSize > 0) - { - auto samplingConfig = SamplingConfig(samplingConfigs); - mDecoder->getUnderlyingDecoder().setup(samplingConfig, localBatchSize, batchSlots, - {mDecoderState->getJointDecodingOutput()}, mModelConfig.getDataType(), lookaheadPrompt, - lookaheadAlgoConfigs); - - auto const& stream = mDecoder->getDecoderStream(); - CudaEvent event{}; - stream->record(event); - mRuntime->getStreamPtr()->wait(event); - } - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TrtGptModelInflightBatching::postProcessRequest( - LlmRequest& llmReq, std::vector const& numDroppedTokens) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - auto const seqSlot = llmReq.mSeqSlot.value(); - auto const reqBeamWidth = llmReq.getBeamWidthByIter(true); - auto const& bufferManager = getBufferManager(); - - if (llmReq.getReturnGenerationLogits() && !llmReq.getGenerationLogitsFragments().empty()) - { - TLLM_CHECK(!llmReq.isStreaming()); - auto const genBufferId = mCtxGenFusion ? getFusedBufferId() : getGenerationBufferId(); - auto& genRuntimeBuffers = *mBuffers.at(genBufferId); - - auto constexpr beforeDecoder = false; - utils::copyGenerationLogits( - genRuntimeBuffers.generationLogitsCache, bufferManager, llmReq, beforeDecoder, numDroppedTokens); - - bufferManager.getStream().synchronize(); - } - - if (reqBeamWidth == 1) - { - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); - return; - } - - // Update mDecoderBuffers->slotOutputIdsHost and synchronize - getDecoderSlotHostOutputs(seqSlot, llmReq.returnLogProbs(), llmReq.mSamplingConfig, llmReq.isStreaming()); - - auto const* outputIdsHostData = bufferCast(*mSlotDecoderBuffers[seqSlot]->outputIdsHost); - auto const* sequenceLengthsHostData = bufferCast(*mSlotDecoderBuffers[seqSlot]->sequenceLengthsHost); - auto const* cumLogProbsHostData = bufferCast(*mSlotDecoderBuffers[seqSlot]->cumLogProbsHost); - auto logProbsHost = mSlotDecoderBuffers[seqSlot]->logProbsHost; - auto const* logProbsHostData = bufferCast(*logProbsHost); - - auto const& outputIdsShape = mSlotDecoderBuffers[seqSlot]->outputIdsHost->getShape(); - auto const maxSeqLength = outputIdsShape.d[1]; - - std::vector> generatedTokens(reqBeamWidth); - for (SizeType32 beam = 0; beam < reqBeamWidth; ++beam) - { - auto const* const begin = outputIdsHostData + tc::flat_index2(beam, llmReq.mPromptLen, maxSeqLength); - auto const generatedLength = sequenceLengthsHostData[beam] - llmReq.mPromptLen; - auto const* const end = begin + generatedLength; - generatedTokens[beam].assign(begin, end); - - if (llmReq.returnLogProbs()) - { - llmReq.setCumLogProb(cumLogProbsHostData[beam], beam); - - auto const beginLogProbsOffset = reqBeamWidth == 1 ? llmReq.mPromptLen : 0; - auto const* const begin = logProbsHostData + beam * logProbsHost->getShape().d[1] + beginLogProbsOffset; - auto const* const end = begin + generatedLength; - LlmRequest::VecLogProbs logProbs(begin, end); - llmReq.setLogProbs(logProbs, beam); - } - } - - // store the generated tokens into the mTokensGathered buffer - llmReq.setGeneratedTokens(generatedTokens); - - if (llmReq.getReturnGenerationLogits() && llmReq.getGenerationLogitsHost() - && mWorldConfig.isLastPipelineParallelRank()) - { - reorderGenerationLogitsForBeamSearch( - llmReq, seqSlot, reqBeamWidth, maxSeqLength, outputIdsHostData, sequenceLengthsHostData); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TrtGptModelInflightBatching::reorderGenerationLogitsForBeamSearch(LlmRequest& llmReq, SizeType32 seqSlot, - SizeType32 reqBeamWidth, SizeType32 maxSeqLength, TokenIdType const* outputIdsHostData, - SizeType32 const* sequenceLengthsHostData) -{ - // Reorder generation logits to match the gathered (finalized) beam ordering. - // During generation, logits are stored indexed by beam SLOT position. After beam search - // finalization (gatherTree), output_ids are reordered by tracing parentIds to reconstruct - // the correct beam paths. However, generation_logits are NOT reordered by gatherTree. - // We fix this here by tracing parentIds on the host to build the beam-slot mapping, - // then reindexing the logits accordingly. - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - auto const promptLen = llmReq.mPromptLen; - - // Copy parentIds and ids (ungathered step IDs) from GPU to temporary host buffers. - // parentIds[slot][t] = the parent slot of beam slot `slot` at position t. - // ids[slot][t] = the token in beam slot `slot` at position t (before gather). - auto parentIdsDevice = ITensor::at(mDecoderState->getParentIds(), {seqSlot}); - auto idsDevice = mDecoderState->getIds(seqSlot); - - auto parentIdsHost = runtime::BufferManager::pinnedPool(parentIdsDevice->getShape(), nvinfer1::DataType::kINT32); - auto idsHost = runtime::BufferManager::pinnedPool(idsDevice->getShape(), nvinfer1::DataType::kINT32); - - mCopyBufferManager.copy(*parentIdsDevice, *parentIdsHost); - mCopyBufferManager.copy(*idsDevice, *idsHost); - mCopyBufferManager.getStream().synchronize(); - - auto const* parentIdsData = bufferCast(*parentIdsHost); - auto const* idsData = bufferCast(*idsHost); - - // For each final beam b, find the beam slot at the last generated step, then - // trace back through parentIds to build the slot trace for every generation step. - // slotTrace[beam][genStep] = the beam slot that produced the logits at that step. - auto const generationLogitsHost = llmReq.getGenerationLogitsHost(); - auto const& logitsShape = generationLogitsHost->getShape(); - // Non-streaming shape: [beamWidth, maxNewTokens, vocabSizePadded] - TLLM_CHECK_WITH_INFO(logitsShape.d[0] == reqBeamWidth, - "Generation logits beam dimension (%ld) does not match beam width (%d).", logitsShape.d[0], reqBeamWidth); - auto const maxNewTokens = logitsShape.d[1]; - auto const vocabSizePadded = logitsShape.d[2]; - - std::vector> slotTrace(reqBeamWidth, std::vector(maxNewTokens, 0)); - bool anyReorderNeeded = false; - - for (SizeType32 beam = 0; beam < reqBeamWidth; ++beam) - { - auto const seqLen = sequenceLengthsHostData[beam]; - auto const genLen = seqLen - promptLen; - if (genLen <= 0) - { - continue; - } - - // Find the starting beam slot at the last generated step by matching the - // backtracked token sequence against the gathered (finalized) output. - SizeType32 startSlot = -1; - for (SizeType32 s = 0; s < reqBeamWidth; ++s) - { - SizeType32 slot = s; - bool matches = true; - for (SizeType32 t = seqLen - 1; t >= promptLen; --t) - { - if (idsData[slot * maxSeqLength + t] != outputIdsHostData[beam * maxSeqLength + t]) - { - matches = false; - break; - } - if (t > promptLen) - { - slot = parentIdsData[slot * maxSeqLength + t]; - } - } - if (matches) - { - startSlot = s; - break; - } - } - - TLLM_CHECK_WITH_INFO(startSlot >= 0, - "Could not determine beam slot mapping for beam %d during generation logits reordering.", beam); - - // Build the slot trace: slotTrace[beam][g] = the pre-reassignment slot whose - // logits correspond to generation step g of this beam. - // - // The model runs BEFORE beam search reassigns beams to slots, so - // generationLogits[slot][g] was produced by the pre-reassignment slot — - // i.e. the slot the beam occupied in the *previous* step. - // parentIds[postSlot][promptLen+g] gives exactly that pre-reassignment slot, - // so taking the parentIds lookup before storing (rather than after) yields - // the correct source slot in a single pass. - SizeType32 slot = startSlot; - for (SizeType32 t = seqLen - 1; t >= promptLen; --t) - { - slot = parentIdsData[slot * maxSeqLength + t]; - slotTrace[beam][t - promptLen] = slot; - } - - // Check if any reordering is actually needed for this beam - auto& slotTraceIds = slotTrace[beam]; - anyReorderNeeded |= std::any_of( - slotTraceIds.begin(), slotTraceIds.begin() + genLen, [beam](SizeType32 s) { return s != beam; }); - } - - // Reorder the generation logits in-place using a per-step temporary buffer. - if (anyReorderNeeded) - { - auto const logitsDataType = generationLogitsHost->getDataType(); - auto const elemSize = runtime::BufferDataType(logitsDataType).getSize(); - auto const stepSize = static_cast(vocabSizePadded) * elemSize; - - // Temp buffer for one generation step across all beams: [beamWidth, vocabSizePadded] - auto tempLogits - = runtime::BufferManager::pinnedPool(ITensor::makeShape({reqBeamWidth, vocabSizePadded}), logitsDataType); - - auto* logitsPtr = static_cast(generationLogitsHost->data()); - auto* tempPtr = static_cast(tempLogits->data()); - - std::vector genLens(reqBeamWidth); - SizeType32 maxGenLen = 0; - for (SizeType32 b = 0; b < reqBeamWidth; ++b) - { - genLens[b] = std::max(SizeType32{0}, sequenceLengthsHostData[b] - promptLen); - maxGenLen = std::max(maxGenLen, genLens[b]); - } - - for (SizeType32 g = 0; g < maxGenLen; ++g) - { - // Check if any beam that generated this step needs reordering - bool stepNeedsReorder = false; - for (SizeType32 b = 0; b < reqBeamWidth; ++b) - { - if (g < genLens[b] && slotTrace[b][g] != b) - { - stepNeedsReorder = true; - break; - } - } - if (!stepNeedsReorder) - { - continue; - } - - // Copy all beams' logits at this step to the temp buffer - for (SizeType32 b = 0; b < reqBeamWidth; ++b) - { - // logits layout: [beamWidth, maxNewTokens, vocabSizePadded] - auto const offset = (static_cast(b) * maxNewTokens + g) * stepSize; - std::memcpy(tempPtr + static_cast(b) * stepSize, logitsPtr + offset, stepSize); - } - - // Reorder: logits[b][g] = temp[slotTrace[b][g]] - for (SizeType32 b = 0; b < reqBeamWidth; ++b) - { - if (g >= genLens[b]) - { - continue; - } - auto const dstOffset = (static_cast(b) * maxNewTokens + g) * stepSize; - auto const srcSlot = slotTrace[b][g]; - std::memcpy(logitsPtr + dstOffset, tempPtr + static_cast(srcSlot) * stepSize, stepSize); - } - } - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TrtGptModelInflightBatching::getDecoderSlotHostOutputs( - SizeType32 seqSlot, bool returnLogProbs, SamplingConfig const& samplingConfig, bool streaming) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - if (mWorldConfig.isLastPipelineParallelRank()) - { - auto event = mDecoder->finalize(*mDecoderState, seqSlot, samplingConfig, streaming); - // Make sure that postprocessing is done before copying outputIds - mCopyBufferManager.getStream().wait(event.get()); - - auto sequenceLengths = mDecoderState->getSequenceLengths(seqSlot); - auto outputIds = mDecoderState->getGatheredIds(seqSlot); - auto cumLogProbs = mDecoderState->getCumLogProbs(seqSlot); - auto logProbs = mDecoderState->getLogProbs(seqSlot); - - mCopyBufferManager.copy(*sequenceLengths, *mSlotDecoderBuffers[seqSlot]->sequenceLengths); - mCopyBufferManager.copy(*outputIds, *mSlotDecoderBuffers[seqSlot]->outputIds); - if (returnLogProbs) - { - mCopyBufferManager.copy(*cumLogProbs, *mSlotDecoderBuffers[seqSlot]->cumLogProbs); - mCopyBufferManager.copy(*logProbs, *mSlotDecoderBuffers[seqSlot]->logProbs); - } - - if (mWorldConfig.isPipelineParallel()) - { - // Make sure that postprocessing is done before sending outputIds - event.synchronize(); - - auto const peerSend = 0; - mDecSlotAsyncSndHdls.emplace_back(std::make_unique( - outputIds, sequenceLengths, cumLogProbs, logProbs, returnLogProbs, *mMpiCommPipelinePara, peerSend)); - } - } - else - { - auto const peerRecv = mWorldConfig.getPipelineParallelRank() == 0 ? mWorldConfig.getPipelineParallelism() - 1 - : mWorldConfig.getPipelineParallelRank() - 1; - DecoderSlotAsyncSend::recv(*mSlotDecoderBuffers[seqSlot], returnLogProbs, *mMpiCommPipelinePara, peerRecv); - - auto const peerSend = mWorldConfig.getPipelineParallelRank() + 1; - if (peerSend != mWorldConfig.getPipelineParallelism() - 1) - { - mDecSlotAsyncSndHdls.emplace_back(std::make_unique( - *mSlotDecoderBuffers[seqSlot], returnLogProbs, *mMpiCommPipelinePara, peerSend)); - } - } - sync_check_cuda_error(mRuntime->getStream().get()); - - // Here copy stream is synchronized after receiving decoderSlotOutputIdsView either by copy or by receive - // before copying to host on copy stream - runtime::CudaEvent beforeEvent{}; - mRuntime->getStreamPtr()->record(beforeEvent); - mCopyBufferManager.getStream().wait(beforeEvent); - mCopyBufferManager.copy(*mSlotDecoderBuffers[seqSlot]->outputIds, *mSlotDecoderBuffers[seqSlot]->outputIdsHost); - mCopyBufferManager.copy( - *mSlotDecoderBuffers[seqSlot]->sequenceLengths, *mSlotDecoderBuffers[seqSlot]->sequenceLengthsHost); - - if (returnLogProbs) - { - mCopyBufferManager.copy( - *mSlotDecoderBuffers[seqSlot]->cumLogProbs, *mSlotDecoderBuffers[seqSlot]->cumLogProbsHost); - mCopyBufferManager.copy(*mSlotDecoderBuffers[seqSlot]->logProbs, *mSlotDecoderBuffers[seqSlot]->logProbsHost); - } - - // Make sure copy is done before continuing on host - mCopyBufferManager.getStream().synchronize(); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -namespace -{ -// Check if one of the request needs log probs, need to get from decoder and communicate -bool batchReturnLogProbs(ScheduledRequests const& scheduledRequests) -{ - auto pred = [](auto const& llmReq) { return llmReq->returnLogProbs(); }; - return std::any_of(scheduledRequests.contextRequests.begin(), scheduledRequests.contextRequests.end(), pred) - || std::any_of(scheduledRequests.generationRequests.begin(), scheduledRequests.generationRequests.end(), pred); -} -} // namespace - -runtime::CudaEvent TrtGptModelInflightBatching::decoderStepAsync(ScheduledRequests const& scheduledRequests) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(decoderStepAsync); - - auto& decoderInputBuffers = mDecoderInputBuffers.at(getFusedBufferId()); - - auto const contextBufferId = mCtxGenFusion ? getFusedBufferId() : getContextBufferId(); - auto& contextRuntimeBuffers = mBuffers.at(contextBufferId); - auto const logitsIndex = (*mHandleContextLogits)(decoderInputBuffers, scheduledRequests.contextRequests, - contextRuntimeBuffers->logits, contextRuntimeBuffers->numContextLogits, mModelConfig, - mRuntime->getBufferManager(), contextRuntimeBuffers->mMedusaBuffers); - - auto const genLogitsIndex = mCtxGenFusion ? logitsIndex : 0; - auto const genBufferId = mCtxGenFusion ? getFusedBufferId() : getGenerationBufferId(); - auto& genRuntimeBuffers = mBuffers.at(genBufferId); - (*mHandleGenerationLogits)(decoderInputBuffers, scheduledRequests.generationRequests, genRuntimeBuffers->logits, - genLogitsIndex, mModelConfig, mRuntime->getBufferManager(), *genRuntimeBuffers, - genRuntimeBuffers->mMedusaBuffers); - - if (mOperatingBeamWidth > 1) - { - copyCacheIndirectionFromOutputsToInputs(scheduledRequests, genBufferId); - } - - mLogitsPostProcessorIsApplied = (*mLogitsPostProcessor)(decoderInputBuffers, mReplicateLogitsPostProcessor, - mWorldConfig, mRuntime->getStreamPtr(), mLogitsPostProcessorBatched); - - if (mGuidedDecoder) - { - mGuidedDecoder->execute(decoderInputBuffers, mRuntime->getBufferManager()); - } - - auto const fusedBufferId = getFusedBufferId(); - auto& fusedRuntimeBuffers = mBuffers.at(fusedBufferId); - - (*mMakeDecodingBatchInputOutput)(decoderInputBuffers, *mDecoderState, mModelConfig, *fusedRuntimeBuffers); - - auto decoderFinishEvent = mDecoder->forwardAsync(*mDecoderState, decoderInputBuffers); - - auto const returnLogProbs = batchReturnLogProbs(scheduledRequests); - auto updateDecoderBuffersEvent = (*mUpdateDecoderBuffers)(mModelConfig, mDecoderOutputBuffers.at(fusedBufferId), - mRuntime->getBufferManager(), *mDecoderState, returnLogProbs, decoderFinishEvent); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); - return updateDecoderBuffersEvent; -} - -void TrtGptModelInflightBatching::copyCacheIndirectionFromOutputsToInputs( - ScheduledRequests const& scheduledRequests, SizeType32 genBufferId) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(copyCacheIndirectionFromOutputsToInputs); - - auto& genRuntimeBuffers = *mBuffers.at(genBufferId); - auto* srcOffsetsPtr = bufferCast(*genRuntimeBuffers.cacheIndirDecoderIOBatchedCopySrcOffsets); - auto* dstOffsetsPtr = bufferCast(*genRuntimeBuffers.cacheIndirDecoderIOBatchedCopyDstOffsets); - auto* copySizesPtr = bufferCast(*genRuntimeBuffers.cacheIndirDecoderIOBatchedCopySizes); - - // Only `cacheIndirShape.d[2]` is used - auto const& cacheIndirShape = mDecoderState->getCacheIndirectionOutput()->getShape(); - auto const maxBeamWidth = cacheIndirShape.d[1]; - auto const maxAttentionWindow = cacheIndirShape.d[2]; - auto const slotOffset = maxBeamWidth * maxAttentionWindow; - - SizeType32 batchIdx{0}; - SizeType64 maxCopySize{0}; - auto& manager = mRuntime->getBufferManager(); - for (auto const& requests : {scheduledRequests.contextRequests, scheduledRequests.generationRequests}) - { - for (auto const& llmReq : requests) - { - auto const reqBeamWidth = llmReq->getBeamWidthByIter(); - auto const seqSlot = llmReq->mSeqSlot.value(); - auto const copySize = reqBeamWidth * maxAttentionWindow; - srcOffsetsPtr[batchIdx] = seqSlot * slotOffset; - dstOffsetsPtr[batchIdx] = seqSlot * slotOffset; - copySizesPtr[batchIdx] = copySize; - maxCopySize = std::max(maxCopySize, copySize); - batchIdx++; - } - } - if (batchIdx != 0) - { - auto const srcOffsetsSlice - = ITensor::slice(genRuntimeBuffers.cacheIndirDecoderIOBatchedCopySrcOffsets, 0, batchIdx); - auto const srcOffsetsSliceDeviceSlice - = ITensor::slice(genRuntimeBuffers.mCacheIndirDecoderIOBatchedCopySrcOffsetsSliceDevice, 0, batchIdx); - manager.copy(srcOffsetsSlice->data(), *srcOffsetsSliceDeviceSlice, - runtime::MemoryType::kGPU); // Explicitly move to device for faster access. - auto const dstOffsetsSlice - = ITensor::slice(genRuntimeBuffers.cacheIndirDecoderIOBatchedCopyDstOffsets, 0, batchIdx); - auto const dstOffsetsSliceDeviceSlice - = ITensor::slice(genRuntimeBuffers.mCacheIndirDecoderIOBatchedCopyDstOffsetsSliceDevice, 0, batchIdx); - manager.copy(dstOffsetsSlice->data(), *dstOffsetsSliceDeviceSlice, - runtime::MemoryType::kGPU); // Explicitly move to device for faster access. - auto const sizesSlice = ITensor::slice(genRuntimeBuffers.cacheIndirDecoderIOBatchedCopySizes, 0, batchIdx); - auto const copySizesDeviceSlice - = ITensor::slice(genRuntimeBuffers.mCacheIndirDecoderIOBatchedCopyCopySizesDevice, 0, batchIdx); - manager.copy(sizesSlice->data(), *copySizesDeviceSlice); // Explicitly move to device for faster access. - runtime::kernels::invokeCopyBatch(*mDecoderState->getCacheIndirectionOutput(), - *mDecoderState->getCacheIndirectionInput(), *srcOffsetsSliceDeviceSlice, *dstOffsetsSliceDeviceSlice, - *copySizesDeviceSlice, maxCopySize, manager.getStream()); - } - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -std::vector> TrtGptModelInflightBatching::communicateDecoderBuffers( - bool returnLogProbs) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(communicateDecoderBuffers); - - auto& decoderOutputBuffers = mDecoderOutputBuffers.at(getFusedBufferId()); - - std::vector> asyncHandles; - if (mWorldConfig.isLastPipelineParallelRank()) - { - if (broadcastPostDecoder()) - { - DecoderStepAsyncSend::bcast(decoderOutputBuffers, *mDecoderState, returnLogProbs, mOperatingBeamWidth, - mModelConfig.getSpeculativeDecodingMode().needsKVCacheRewind(), *mMpiCommTensorPara, 0); - } - - if (mWorldConfig.isPipelineParallel()) - { - auto const peerSend = 0; - asyncHandles.emplace_back(std::make_unique(decoderOutputBuffers, *mDecoderState, - returnLogProbs, mOperatingBeamWidth, mModelConfig.getSpeculativeDecodingMode().needsKVCacheRewind(), - *mMpiCommPipelinePara, peerSend)); - } - } - else - { - auto const peerRecv = mWorldConfig.isFirstPipelineParallelRank() ? mWorldConfig.getPipelineParallelism() - 1 - : mWorldConfig.getPipelineParallelRank() - 1; - DecoderStepAsyncSend::recv(decoderOutputBuffers, *mDecoderState, returnLogProbs, mOperatingBeamWidth, - mModelConfig.getSpeculativeDecodingMode().needsKVCacheRewind(), *mMpiCommPipelinePara, peerRecv); - auto const peerSend = mWorldConfig.getPipelineParallelRank() + 1; - if (peerSend != mWorldConfig.getPipelineParallelism() - 1) - { - asyncHandles.emplace_back(std::make_unique(decoderOutputBuffers, *mDecoderState, - returnLogProbs, mOperatingBeamWidth, mModelConfig.getSpeculativeDecodingMode().needsKVCacheRewind(), - *mMpiCommPipelinePara, peerSend)); - } - } - TLLM_CHECK_WITH_INFO(asyncHandles.size() <= 2, "Up to two decoder step async handles expected"); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); - return asyncHandles; -} - -void TrtGptModelInflightBatching::updateRequests(ScheduledRequests const& scheduledRequests) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(updateRequests); - - auto const& decoderOutputBuffers = mDecoderOutputBuffers.at(getFusedBufferId()); - - auto const hostNewOutputTokensShape = decoderOutputBuffers.newOutputTokensHost->getShape(); - auto const* const hostNewOutputTokensData - = bufferCast(*decoderOutputBuffers.newOutputTokensHost); - auto const* const sequenceLengthsHostData = bufferCast(*decoderOutputBuffers.sequenceLengthsHost); - auto const* const decoderFinishedSumPtr = bufferCast(*decoderOutputBuffers.finishedSumHost); - auto const* const cumLogProbsPtr = bufferCast(*decoderOutputBuffers.cumLogProbsHost); - auto const* const logProbsPtr = bufferCast(*decoderOutputBuffers.logProbsHost); - auto const* const finishReasonsHostData - = bufferCast(*decoderOutputBuffers.finishReasonsHost); - - // Update only requests that ran through the decoder - for (auto const& llmReq : scheduledRequests.generationRequests) - { - if (llmReq->isGenerationCompleteState()) - { - continue; - } - auto const reqBeamWidth = llmReq->getBeamWidthByIter(true); - auto const seqSlot = llmReq->mSeqSlot.value(); - auto const currentNumOfTokens = llmReq->getMaxBeamNumTokens(); - - // Save the accepted token logits from target model - if (mModelConfig.getSpeculativeDecodingMode().isDraftTokensExternal() && llmReq->getReturnGenerationLogits() - && llmReq->hasDraftTokens()) - { - TLLM_CHECK_WITH_INFO(reqBeamWidth == 1, "Speculative decoding only works for beam width == 1"); - - SizeType32 numAcceptedTokens - = sequenceLengthsHostData[seqSlot * mOperatingBeamWidth + 0] - llmReq->getMaxBeamNumTokens(); - - auto const& generationLogitsHost = llmReq->getGenerationLogitsHost(); - auto shape = generationLogitsHost->getShape(); - shape.d[1] = numAcceptedTokens; - generationLogitsHost->reshape(shape); - } - - std::vector numNewTokens(reqBeamWidth); - std::vector numDroppedTokens(reqBeamWidth); - - // numGeneratedTokens is the number of tokens generated by the decoder. - // Some tokens might be dropped due to end token or rejected draft tokens. - auto const numGeneratedTokens = llmReq->getNumDraftTokens() + 1; - - for (SizeType32 beam = 0; beam < reqBeamWidth; ++beam) - { - // Sequence length is only advanced for accepted tokens. - auto const seqLen = sequenceLengthsHostData[seqSlot * mOperatingBeamWidth + beam]; - // Actual number of tokens that should be added to the request. - auto const numNewOutputTokens = seqLen - llmReq->getNumTokens(beam); - if (reqBeamWidth == 1) - { - TLLM_CHECK_WITH_INFO(numGeneratedTokens >= numNewOutputTokens, - "numNewOutputTokens must not be greater than numGeneratedTokens: " - "numGeneratedTokens %d < numNewOutputTokens %d", - numGeneratedTokens, numNewOutputTokens); - } - numNewTokens[beam] = std::min(numGeneratedTokens, numNewOutputTokens); - numDroppedTokens[beam] = numGeneratedTokens - numNewTokens[beam]; - for (SizeType32 step = 0; step < numNewTokens[beam]; ++step) - { - auto const newTokenIdx = tc::flat_index(hostNewOutputTokensShape.d, step, seqSlot, beam); - auto const newToken = hostNewOutputTokensData[newTokenIdx]; - llmReq->addNewToken(newToken, beam); - TLLM_LOG_DEBUG("request ID %ld beam %d newToken %d", llmReq->mRequestId, beam, newToken); - - if (llmReq->returnLogProbs()) - { - auto const cumLogProb = cumLogProbsPtr[seqSlot * mOperatingBeamWidth + beam]; - llmReq->setCumLogProb(cumLogProb, beam); - - auto const beginLogProbsOffset = reqBeamWidth == 1 ? llmReq->mPromptLen : 0; - SizeType32 offset - = (seqSlot * mOperatingBeamWidth + beam) * getMaxSequenceLen() + beginLogProbsOffset; - auto const generatedLength = seqLen - llmReq->mPromptLen; - std::vector logProbs(logProbsPtr + offset, logProbsPtr + offset + generatedLength); - llmReq->setLogProbs(logProbs, beam); - } - } - - auto const finishReason = finishReasonsHostData[seqSlot * mOperatingBeamWidth + beam]; - llmReq->setFinishedReason(finishReason.toFinishReason(), beam); - - TLLM_LOG_DEBUG("[RANK %d] decoderSync: request ID %lu beam %d tokens %s finished %d", - COMM_SESSION.getRank(), llmReq->mRequestId, beam, common::vec2str(llmReq->getTokens(beam)).c_str(), - static_cast(finishReason.toFinishReason())); - } - - // Set number of tokens predicted per runtime iteration. Will be > 1 for speculative decoding. - llmReq->updateNumTokensPerIteration(llmReq->getMaxBeamNumTokens() - currentNumOfTokens, mModelConfig); - - // Fill new draft tokens for the next step - if (decoderFinishedSumPtr[seqSlot] != reqBeamWidth - && (mModelConfig.getSpeculativeDecodingMode().predictsDraftTokens() - || mModelConfig.getSpeculativeDecodingMode().needsKVCacheRewind())) - { - auto const maxDraftTokensLen = mModelConfig.getMaxDecodingDraftTokens(); - auto prevDraftTokensLen = llmReq->getNumDraftTokens(); - - // We overallocate KV cache for EAGLE to the maxDecodingTokens + maxPathLen in order to fit both - // Base model verification (needs up to maxDecodingTokens) and - // Drafter (needs up to maxPathLen of accepted tokens and maxDecodingDraftTokens for new draft tokens). - if (mModelConfig.getSpeculativeDecodingMode().isEagle()) - { - prevDraftTokensLen = mModelConfig.getSpeculativeDecodingModule().getMaxDecodingTokens() - + mModelConfig.getSpeculativeDecodingModule().getMaxPathLen() - 1; - } - - auto nextDraftTokensLen = mModelConfig.getSpeculativeDecodingModule().getMaxDecodingDraftTokens(); - if (mModelConfig.getSpeculativeDecodingMode().variableDraftLength()) - { - auto const* const nextDraftTokensLengthsHostData - = bufferCast(*decoderOutputBuffers.nextDraftTokensLengthsHost); - nextDraftTokensLen = nextDraftTokensLengthsHostData[seqSlot]; - } - TLLM_CHECK(nextDraftTokensLen <= maxDraftTokensLen); - - auto const* const nextDraftTokensHostData - = bufferCast(*decoderOutputBuffers.nextDraftTokensHost); - auto draftTokensShared - = std::make_shared>(nextDraftTokensHostData + seqSlot * maxDraftTokensLen, - nextDraftTokensHostData + seqSlot * maxDraftTokensLen + nextDraftTokensLen); - - llmReq->setDraftTokens(draftTokensShared); - - // For all phases except context that does not have draft tokens - if (!llmReq->isGenerationCompleteState() && prevDraftTokensLen != 0 - && mModelConfig.getSpeculativeDecodingMode().needsKVCacheRewind()) - { - // -1 here is for current 'main' token - auto const acceptedTokensLen = llmReq->getMaxBeamNumTokens() - currentNumOfTokens - 1; - auto const rewindLength = prevDraftTokensLen - acceptedTokensLen; - - TLLM_LOG_DEBUG("request ID %lu (seqSlot %d): accepted %d of %d draft tokens, rewind %d tokens", - llmReq->mRequestId, seqSlot, acceptedTokensLen, prevDraftTokensLen, rewindLength); - TLLM_CHECK(0 <= acceptedTokensLen && acceptedTokensLen <= prevDraftTokensLen); - - // At this point, KV cache rows are already gathered and moved to the right location. - // We can safely rewind (draft - accepted) tokens - mKvCacheManager->rewindKVCache(llmReq->mRequestId, rewindLength); - } - } - - // Terminate if request has finished or if it is speculative decoding target model - if (decoderFinishedSumPtr[seqSlot] == reqBeamWidth - || (mModelConfig.getSpeculativeDecodingMode().isDraftTokensExternal() && llmReq->hasDraftTokens())) - { - postProcessRequest(*llmReq, numDroppedTokens); - - if (!mWorldConfig.isPipelineParallel() || !mWorldConfig.isLastPipelineParallelRank()) - { - if (llmReq->getReturnGenerationLogits() && mSpeculativeDecodingFastLogits && mIsLeaderInOrchMode) - { - std::lock_guard lk(mDraftRequestsMtx); - mDraftRequestsWaitingToSendLogits.push_back(llmReq); - } - else - { - terminateRequest(llmReq); - } - llmReq->setState(LlmRequestState::kGENERATION_COMPLETE); - } - else - { - llmReq->setState(LlmRequestState::kGENERATION_TO_COMPLETE); - } - } - else - { - // gather tokens in the case of streaming and beam search - if (llmReq->isStreaming() && llmReq->mSamplingConfig.beamWidth > 1) - { - postProcessRequest(*llmReq, numDroppedTokens); - } - if (llmReq->isContextInitState()) - { - llmReq->setState(LlmRequestState::kGENERATION_IN_PROGRESS); - } - - if (isTrtOverlap() && llmReq->willCompleteNextIteration()) - { - // This state prohibits the request from being scheduled for another iteration. It assumes that the next - // iteration has already been scheduled and the request can finish in the next call to updateRequests(). - llmReq->setState(LlmRequestState::kGENERATION_TO_COMPLETE); - } - } - - if (llmReq->getReturnPerfMetrics()) - { - llmReq->updatePerfMetrics(mIterCounter); - } - - llmReq->advanceDecodingIter(); - - if (mWorldConfig.isPipelineParallel() && mWorldConfig.isLastPipelineParallelRank()) - { - for (SizeType32 beam = 0; beam < reqBeamWidth; ++beam) - { - llmReq->setNumPreDecodedTokens(numNewTokens[beam], beam); - } - } - } - - if (mModelConfig.getSpeculativeDecodingMode().needsKVCacheRewind()) - { - SizeType32 numSequences{0}; - for (auto const& requests : {scheduledRequests.contextRequests, scheduledRequests.generationRequests}) - { - for (auto const& llmReq : requests) - { - auto const reqBeamWidth = llmReq->mSamplingConfig.beamWidth; - numSequences += reqBeamWidth; - } - } - - TLLM_CHECK_WITH_INFO(mCtxGenFusion, "Current speculative decoding mode requires context-gen fusion IFB"); - rewindKVCacheBlocks(numSequences); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -std::vector> TrtGptModelInflightBatching::decoderSync( - ScheduledRequests const& scheduledRequests, std::optional const& decoderFinishEvent) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(decoderSync); - - if (mWorldConfig.isLastPipelineParallelRank()) - { - decoderFinishEvent->synchronize(); - } - - auto const returnLogProbs = batchReturnLogProbs(scheduledRequests); - auto asyncHandles = communicateDecoderBuffers(returnLogProbs); - - updateRequests(scheduledRequests); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); - return asyncHandles; -} - -void TrtGptModelInflightBatching::rewindKVCacheBlocks(SizeType32 numSequences) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - auto const bufferId = getFusedBufferId(); - auto& runtimeBuffers = *mBuffers.at(bufferId); - auto& decoderOutputBuffers = mDecoderOutputBuffers.at(bufferId); - - auto localNbLayers = mModelConfig.getNbAttentionLayers( - mWorldConfig.getPipelineParallelism(), mWorldConfig.getPipelineParallelRank()); - if (mWorldConfig.isLastPipelineParallelRank() && mModelConfig.getSpeculativeDecodingMode().isEagle()) - { - // Do not correct the last kv caches, which are for EagleNet drafter. Those KV caches are managed separately. - auto eagleModulePtr - = std::dynamic_pointer_cast(mModelConfig.getSpeculativeDecodingModulePtr()); - localNbLayers -= eagleModulePtr->getNumTransformerLayers(); - } - - auto const tokensPerBlock = mModelConfig.getTokensPerBlock(); - auto const elemSize = BufferDataType(mModelConfig.getKvDataType()).getSize(); - auto const sizeInBytesPerKVHead = mModelConfig.getSizePerHead() * elemSize; - - auto const poolPointers = mKvCacheManager->getBlockPoolPointers(); - auto* const* pointerArrayPtr = bufferCast(*poolPointers); - auto const* offsetArrayPtr - = bufferCast(*runtimeBuffers.transformerBuffers->kvCacheBlockOffsetsDevice); - - auto commonRewindLen = mModelConfig.getSpeculativeDecodingModule().getMaxDecodingDraftTokens(); - SizeType32 const* rewindLens = nullptr; - if (mModelConfig.getSpeculativeDecodingMode().variableDraftLength()) - { - commonRewindLen = 0; - rewindLens = bufferCast(*decoderOutputBuffers.prevDraftTokensLengthsHost); - } - - tensorrt_llm::runtime::kernels::invokeUpdateKVBlockArrayDraftTokenLocation( - *mDecoderState->getAcceptedLengthsCumSum(), *mDecoderState->getAcceptedPackedPaths(), - *runtimeBuffers.sequenceLengthsDevice, pointerArrayPtr, offsetArrayPtr, localNbLayers, numSequences, - mRewindInputs.numKvHeads, sizeInBytesPerKVHead, commonRewindLen, rewindLens, *runtimeBuffers.seqSlots, - getMaxAttentionWindow(), mRewindInputs.maxBlocksPerSeq, tokensPerBlock, mRewindInputs.isUseOneMoreBlock, - mRuntime->getStreamPtr()->get()); - - sync_check_cuda_error(mRuntime->getStream().get()); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -nvinfer1::DataType TrtGptModelInflightBatching::getLogitDataType() const -{ - return mModelConfig.getLogitsDtype(); -} - -TrtGptModelInflightBatching::SizeType32 TrtGptModelInflightBatching::numCachedCudaGraphs() const -{ - return std::accumulate(mCudaGraphExecutorCaches.begin(), mCudaGraphExecutorCaches.end(), SizeType32{0}, - [](SizeType32 sum, auto const& cache) { return sum + cache.size(); }); -} - -void TrtGptModelInflightBatching::changeBeamWidth(SizeType32 beamWidth) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - TLLM_CHECK(mInflightReqIds.empty()); - - TLLM_CHECK_WITH_INFO(beamWidth <= getMaxBeamWidth(), - "Requested beam width %d is larger than configured max beam width %d", beamWidth, getMaxBeamWidth()); - TLLM_LOG_DEBUG("Changing operating beam width from %d to %d", mOperatingBeamWidth, beamWidth); - mOperatingBeamWidth = beamWidth; - - if (isCudaGraphMode()) - { - for (auto& cache : mCudaGraphExecutorCaches) - { - cache.clear(); - } - } - createBuffers(mDecodingConfig, mAdditionalModelOutputs); - createDecoder(mDecodingConfig.getDecodingMode()); - - if (static_cast(mKvCacheManager)) - { - auto const dims = mKvCacheManager->getOffsetTableDimensions(); - reshapeKvTensors(dims); - } - if (static_cast(mCrossKvCacheManager)) - { - auto const dims = mCrossKvCacheManager->getOffsetTableDimensions(); - reshapeKvTensors(dims); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TrtGptModelInflightBatching::changeSpecDecMode(ScheduledRequests const& scheduledRequests) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - if ((!mModelConfig.getSpeculativeDecodingMode().isLookaheadDecoding() - && !mModelConfig.getSpeculativeDecodingMode().isNone()) - || scheduledRequests.empty() || mSeamlessLADMaxDraftLen == 0 || getGatherGenerationLogits() - || mModelConfig.isRnnBased()) - { - return; - } - - bool canUseLookahead = false; - auto maxNumRequestForLad = mDecodingConfig.getLookaheadDecodingMaxNumRequest(); - SizeType32 numRequests = scheduledRequests.contextRequests.size() + scheduledRequests.generationRequests.size(); - if (numRequests > maxNumRequestForLad) - { - if (mModelConfig.getSpeculativeDecodingMode().isLookaheadDecoding()) - { - canUseLookahead = false; - } - else - { - return; - } - } - { - bool useTopKTopP = false; - bool useBanWords = false; - bool useTempAccVocabPenalties = false; // use temperature and penalties that need to accumulate #vocab. - SizeType32 beamWidth = 1; - for (auto const& requests : {scheduledRequests.contextRequests, scheduledRequests.generationRequests}) - { - for (auto const& llmReq : requests) - { - useTopKTopP |= !(llmReq->mSamplingConfig.useDefaultValues( - llmReq->mSamplingConfig.topK, layers::DefaultDecodingParams::getTopK()) - || llmReq->mSamplingConfig.useDefaultValues(llmReq->mSamplingConfig.topK, 1)); - useTopKTopP |= !llmReq->mSamplingConfig.useDefaultValues( - llmReq->mSamplingConfig.topP, layers::DefaultDecodingParams::getTopP()); - useBanWords |= llmReq->getBadWordsList().has_value(); - useBanWords |= !llmReq->mSamplingConfig.useDefaultValues( - llmReq->mSamplingConfig.noRepeatNgramSize, layers::DefaultDecodingParams::getNoRepeatNgramSize()); - useTempAccVocabPenalties |= !llmReq->mSamplingConfig.useDefaultValues( - llmReq->mSamplingConfig.temperature, layers::DefaultDecodingParams::getTemperature()); - useTempAccVocabPenalties |= !llmReq->mSamplingConfig.useDefaultValues( - llmReq->mSamplingConfig.repetitionPenalty, layers::DefaultDecodingParams::getRepetitionPenalty()); - useTempAccVocabPenalties |= !llmReq->mSamplingConfig.useDefaultValues( - llmReq->mSamplingConfig.presencePenalty, layers::DefaultDecodingParams::getPresencePenalty()); - useTempAccVocabPenalties |= !llmReq->mSamplingConfig.useDefaultValues( - llmReq->mSamplingConfig.frequencyPenalty, layers::DefaultDecodingParams::getFrequencyPenalty()); - beamWidth = llmReq->mSamplingConfig.beamWidth; - if (useTopKTopP || useBanWords || useTempAccVocabPenalties || beamWidth > 1) - { - break; - } - } - canUseLookahead = !(useTopKTopP || useBanWords || useTempAccVocabPenalties || beamWidth > 1); - } - } - - // Change speculative decoding mode - auto const bufferId = mCtxGenFusion - ? getFusedBufferId() - : (!scheduledRequests.contextRequests.empty() ? getContextBufferId() : getGenerationBufferId()); - // TODO: enable lookahead for generation requests. - bool canChangeToLookahead = scheduledRequests.generationRequests.empty(); - if (mModelConfig.getSpeculativeDecodingMode().isNone() && canUseLookahead && canChangeToLookahead) - { - // None -> Lookahead - mModelConfig.enableSeamlessLookaheadDecoding(mSeamlessLADMaxDraftLen); - mDecodingConfig.enableSeamlessLookaheadDecoding(); - setupSpeculativeDecodingModule(mDecodingConfig); - mBuffers.at(bufferId)->mLookaheadBuffers->enableLookaheadDecoding( - getMaxBatchSize(), mModelConfig.getMaxDecodingTokens()); - mDecoderOutputBuffers.at(getFusedBufferId()) - .enableLookaheadDecoding(getMaxNumSequences(), mModelConfig.getMaxDecodingTokens()); - createDecoder(mDecodingConfig.getDecodingMode()); - } - else if (mModelConfig.getSpeculativeDecodingMode().isLookaheadDecoding() - && (!canUseLookahead || numRequests > maxNumRequestForLad)) - { - // Lookahead -> None - mModelConfig.disableSeamlessLookaheadDecoding(); - mDecodingConfig.setDecodingMode(executor::DecodingMode::Auto()); - mBuffers.at(bufferId)->mLookaheadBuffers->disableLookaheadDecoding(); - mDecoderOutputBuffers.at(getFusedBufferId()).disableLookaheadDecoding(getMaxNumSequences()); - mDecoder->disableLookahead( - scheduledRequests.generationRequests, mDecoderInputBuffers.at(getFusedBufferId()).setupBatchSlots); - mDecoderState->disableLookahead(scheduledRequests.generationRequests); - for (auto const& llmReq : scheduledRequests.generationRequests) - { - if (llmReq->getNumDraftTokens() > 0) - { - llmReq->discardDraftTokens(llmReq->getNumDraftTokens()); - } - } - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TrtGptModelInflightBatching::getCurrentIterationStats(executor::IterationStats& stats) const -{ - stats.iter = mIterCounter; - - // Max batch size and max num tokens can be tuned at runtime - stats.maxBatchSizeStatic = getMaxBatchSize(); - stats.maxBatchSizeTunerRecommended = mMaxBatchSizeTunerRecommended; - stats.maxBatchSizeRuntime = mMaxBatchSizeRuntime; - stats.maxNumTokensStatic = mMaxNumTokensStatic.value_or(0); - stats.maxNumTokensTunerRecommended = mMaxNumTokensTunerRecommended; - stats.maxNumTokensRuntime = mMaxNumTokensRuntime.value_or(0); - - // KVCacheManager statistics - auto const& kvCacheManager = getKVCacheManager(); - if (kvCacheManager) - { - executor::KvCacheStats kvStats{}; - auto kvCacheStats = kvCacheManager->getKvCacheStats(); - kvStats.maxNumBlocks = kvCacheStats.maxNumBlocks; - kvStats.freeNumBlocks = kvCacheStats.freeNumBlocks; - kvStats.usedNumBlocks = kvCacheStats.usedNumBlocks; - kvStats.tokensPerBlock = kvCacheStats.toksPerBlock; - kvStats.allocTotalBlocks = kvCacheStats.allocTotalBlocks; - kvStats.allocNewBlocks = kvCacheStats.allocNewBlocks; - kvStats.reusedBlocks = kvCacheStats.reusedBlocks; - kvStats.missedBlocks = kvCacheStats.missedBlocks; - kvStats.cacheHitRate = kvCacheStats.cacheHitRate; - stats.kvCacheStats = kvStats; - } - auto const& crossKvCacheManager = getCrossKVCacheManager(); - if (crossKvCacheManager) - { - executor::KvCacheStats kvStats{}; - auto kvCacheStats = crossKvCacheManager->getKvCacheStats(); - kvStats.maxNumBlocks = kvCacheStats.maxNumBlocks; - kvStats.freeNumBlocks = kvCacheStats.freeNumBlocks; - kvStats.usedNumBlocks = kvCacheStats.usedNumBlocks; - kvStats.tokensPerBlock = kvCacheStats.toksPerBlock; - kvStats.allocTotalBlocks = kvCacheStats.allocTotalBlocks; - kvStats.allocNewBlocks = kvCacheStats.allocNewBlocks; - kvStats.reusedBlocks = kvCacheStats.reusedBlocks; - kvStats.missedBlocks = kvCacheStats.missedBlocks; - kvStats.cacheHitRate = kvCacheStats.cacheHitRate; - stats.crossKvCacheStats = kvStats; - } - executor::InflightBatchingStats modelStats{}; - modelStats.numScheduledRequests = mLastIterationStatsIFB.scheduledRequests.size(); - modelStats.numContextRequests = mLastIterationStatsIFB.numCtxRequests; - modelStats.numGenRequests = mLastIterationStatsIFB.numGenRequests; - modelStats.numPausedRequests = mLastIterationStatsIFB.pausedRequests.size(); - modelStats.avgNumDecodedTokensPerIter = mLastIterationStatsIFB.avgNumDecodedTokensPerIter; - modelStats.numCtxTokens = mLastIterationStatsIFB.numCtxTokens; - modelStats.microBatchId = mLastIterationStatsIFB.microBatchId; - stats.inflightBatchingStats = modelStats; -} - -void TrtGptModelInflightBatching::getCurrentRequestStats(executor::RequestStatsPerIteration& stats) const -{ - stats.iter = mIterCounter; - for (auto& requestStat : stats.requestStats) - { - requestStat.scheduled - = mLastIterationStatsIFB.scheduledRequests.count(static_cast(requestStat.id)); - requestStat.paused = mLastIterationStatsIFB.pausedRequests.count(static_cast(requestStat.id)); - } -} - -executor::DebugTensorsPerIteration TrtGptModelInflightBatching::getCurrentDebugTensors() const -{ - executor::DebugTensorsPerIteration debugTensors; - debugTensors.iter = mIterCounter; - - for (auto const& [name, tensor] : mLastIterationDebugTensors) - { - debugTensors.debugTensors.emplace(name, executor::detail::ofITensor(tensor)); - } - - return debugTensors; -} - -nvinfer1::DataType TrtGptModelInflightBatching::getTensorDataType(std::string const& name) const -{ - auto const& engine = mRuntime->getEngine(); - return engine.getTensorDataType(name.c_str()); -} - -nvinfer1::Dims TrtGptModelInflightBatching::getTensorShape(std::string const& name) const -{ - auto const& engine = mRuntime->getEngine(); - return engine.getTensorShape(name.c_str()); -} - -SizeType32 TrtGptModelInflightBatching::getMaxCapacityBatchSize(SizeType32 inputLength, SizeType32 outputLength) const -{ - return mKvCacheManager->getMaxCapacityBatchSize(inputLength, outputLength); -} - -/* - * Manages prefetching of prompt table chunks using a double-buffer strategy - * - * Function Flow: - * 1. First Chunk Processing (isFirstChunk == true): - * - Uses blocking prefetch on main runtime stream - * - Ensures initial data is ready before computation starts - * - * 2. Subsequent Chunks (isFirstChunk == false): - * - Uses non-blocking prefetch on separate copy stream - * - Overlaps data transfer with computation - * - * Synchronization: - * - First prefetch: No wait needed (fresh start) - * - Later prefetches: Wait for previous copy to complete - * - Uses mPtableCopyDoneEvent to track completion - * - * Key Functions: - * 1. prefetchNextPromptTableChunk: - * - Calls the correct function based on position in code (before or after prepareBuffers()) - * - Waits for previous copy to complete if not the first chunk - * - * 2. remapInputTokensForPromptTable: - * - Identifies tokens that need prompt table embeddings (tokens that are greater than vocabSize) - * - Remaps IDs to match chunked prompt table layout - * - * 3. copyPromptTableToGpuInChunk: - * - Handles actual transfer from CPU pinned memory to GPU - * - Uses appropriate buffer manager based on isFirstChunk - */ -void TrtGptModelInflightBatching::prefetchNextPromptTableChunk( - RequestVector const& contextRequests, bool isFirstChunk, SizeType32 bufferId) -{ - auto& promptTuningBuffers = mBuffers[bufferId]->promptTuningBuffers; - - if (!isFirstChunk) - { - // Only switch buffer after prepareBuffer() - promptTuningBuffers->switchChunkPtableBuffer(); - } - - SizeType32 contextId = 0; - for (auto const& llmReq : contextRequests) - { - if (llmReq->isFirstContextChunk() && isFirstChunk) - { - // For first chunk: Blocking prefetch on runtime stream to ensure data is ready - remapInputTokensForPromptTable(llmReq, true, bufferId, contextId); - } - else if (!isFirstChunk) // prefetching for subsequent chunks - { - // For the first prefetch chunk, don't need to wait for previous prefetch to complete - // For subsequent chunks: Need to wait for previous prefetch to complete - if (!llmReq->isFirstContextChunk()) - { - mRuntime->getBufferManager().getStream().wait(mPtableCopyDoneEvent); - } - - // Non-blocking prefetch on copy stream to prepare next chunk in pong buffer - if (llmReq->getContextRemainingLength() > 0) - { - remapInputTokensForPromptTable(llmReq, false, bufferId, contextId); - } - } - - ++contextId; - } -} - -void TrtGptModelInflightBatching::remapInputTokensForPromptTable( - std::shared_ptr const& llmReq, bool isFirstChunk, SizeType32 bufferId, SizeType32 contextId) -{ - NVTX3_SCOPED_RANGE_WITH_NAME(range, "remapInputTokensForPromptTable"); - auto& promptTuningBuffers = mBuffers[bufferId]->promptTuningBuffers; - auto const chunkSize = llmReq->getContextChunkSize(); - auto& inputTokensMutable = llmReq->getTokensMutable(0); - auto vocabSize = mModelConfig.getVocabSize(); - - if (isFirstChunk) - { - promptTuningBuffers->initializeChunkPtableBuffers( - mRuntime->getBufferManager(), mModelConfig, chunkSize, llmReq); - } - - size_t processChunkSize; - size_t beginPos; - - if (!isFirstChunk) - { - processChunkSize = std::min(chunkSize, llmReq->getContextRemainingLength() - chunkSize); - } - else - { - processChunkSize = std::min(chunkSize, llmReq->getContextRemainingLength()); - } - - if (!isFirstChunk) - { - // For prefetching next chunk - if (llmReq->getContextRemainingLength() - chunkSize <= 0) - { - promptTuningBuffers->updateBufferStartPosition(promptTuningBuffers->getChunkPtableCurrentIndex(), 0); - return; // No more chunks to prefetch - } - beginPos = llmReq->getContextCurrentPosition() + chunkSize; - } - else - { - // For current chunk - beginPos = llmReq->getContextCurrentPosition(); - } - - TLLM_CHECK_WITH_INFO(beginPos + processChunkSize <= inputTokensMutable.size(), - "Invalid chunk access: beginPos(%zu) + processChunkSize(%zu) > totalSize(%zu)", beginPos, processChunkSize, - inputTokensMutable.size()); - - auto inputTokensChunk = inputTokensMutable.begin() + beginPos; - std::vector outOfVocabTokens; - SizeType32 ptableTokenId = vocabSize; - for (size_t i = 0; i < processChunkSize; i++) - { - if (inputTokensChunk[i] >= vocabSize) - { - outOfVocabTokens.push_back(inputTokensChunk[i]); - inputTokensChunk[i] = ptableTokenId++; - } - } - - copyPromptTableToGpuInChunk(llmReq, outOfVocabTokens, isFirstChunk, bufferId, contextId); -} - -void TrtGptModelInflightBatching::copyPromptTableToGpuInChunk(std::shared_ptr const& llmReq, - std::vector const& outOfVocabTokens, bool isFirstChunk, SizeType32 bufferId, SizeType32 contextId) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE_WITH_NAME(range, "copyPromptTableToGpuInChunk"); - auto& promptTuningBuffers = mBuffers[bufferId]->promptTuningBuffers; - - if (outOfVocabTokens.empty()) - { - return; - } - - auto const& promptTable = llmReq->getPromptEmbeddingTable(); - TLLM_CHECK_WITH_INFO(promptTable.has_value(), "promptTable is empty but there's fake_prompt"); - TLLM_CHECK_WITH_INFO(promptTable.value() != nullptr, "promptTable value is null but there's fake_prompt"); - - auto currentBufferManager = isFirstChunk ? mRuntime->getBufferManager() : mCopyBufferManager; - auto const hiddenSize = mModelConfig.getHiddenSize(); - auto numRows = outOfVocabTokens.size(); - std::size_t sliceSize = static_cast(numRows * hiddenSize); - auto currentIndex = promptTuningBuffers->getChunkPtableCurrentIndex(); - - // Calculate the offset based on current position - size_t srcOffset = llmReq->mPtableCurrentPosition * hiddenSize; - size_t dstOffset = promptTuningBuffers->getChunkPtableBufferStartPosition(currentIndex, contextId); - - auto gpuBuffer = promptTuningBuffers->getChunkPtableBuffer(currentIndex); - - // First view as 1D tensor of elements - auto totalElements = promptTable.value()->getSize(); - auto table1D = runtime::ITensor::view( - promptTable.value(), runtime::ITensor::makeShape({static_cast(totalElements)})); - - TLLM_CHECK_WITH_INFO(srcOffset + sliceSize <= totalElements, - "Buffer bounds violation: Trying to access up to %zu elements but buffer only has %zu elements (offset: %zu, " - "slice size: %zu)", - srcOffset + sliceSize, totalElements, srcOffset, sliceSize); - - auto table1DShared = runtime::ITensor::SharedPtr(table1D.release()); - auto pTableView = runtime::ITensor::slice(table1DShared, srcOffset, sliceSize); - - auto gpuBufferSlice = runtime::ITensor::slice(gpuBuffer, dstOffset, numRows); - - currentBufferManager.copy(*pTableView, *gpuBufferSlice); - - promptTuningBuffers->updateBufferStartPosition(currentIndex, outOfVocabTokens.size()); - - llmReq->mPtableCurrentPosition += outOfVocabTokens.size(); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -} // namespace tensorrt_llm::batch_manager diff --git a/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.h b/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.h deleted file mode 100644 index d6550281a758..000000000000 --- a/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.h +++ /dev/null @@ -1,639 +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. - */ - -#pragma once - -#include "tensorrt_llm/batch_manager/common.h" -#include "tensorrt_llm/batch_manager/kvCacheType.h" -#include "tensorrt_llm/executor/executor.h" -#include "tensorrt_llm/executor/types.h" -#include "tensorrt_llm/runtime/modelConfig.h" -#include "tensorrt_llm/runtime/rawEngine.h" -#include "tensorrt_llm/runtime/utils/mpiUtils.h" -#include "tensorrt_llm/runtime/worldConfig.h" -#include "trtGptModel.h" - -#include - -namespace tensorrt_llm::runtime -{ -class TllmRuntime; -class GptDecoderBatched; -class AllReduceBuffers; -class NcclCommunicator; -class SpeculativeDecodingMode; - -namespace decoder -{ -class DecoderState; -} // namespace decoder - -namespace decoder_batch -{ -class Input; -class Output; -} // namespace decoder_batch - -} // namespace tensorrt_llm::runtime - -namespace tensorrt_llm::mpi -{ -class MpiWaitThread; -} // namespace tensorrt_llm::mpi - -namespace tensorrt_llm::batch_manager -{ -class BaseCacheTransceiver; -} - -namespace tensorrt_llm::batch_manager -{ - -namespace kv_cache_manager -{ -class KVCacheManager; -struct OffsetTableDimensions; -} // namespace kv_cache_manager - -namespace rnn_state_manager -{ -class RnnStateManager; -} // namespace rnn_state_manager - -class SequenceSlotManager; -class DecoderStepAsyncSend; -class DecoderSlotAsyncSend; -class DecoderInputBuffers; -class DecoderOutputBuffers; -class SlotDecoderBuffers; -class LlmRequest; -class RuntimeBuffers; -class BasePeftCacheManager; -class GuidedDecoder; -class TrtGptModelTest; - -// Algorithms -class CapacityScheduler; -class DisaggTransferAdmissionController; -class MicroBatchScheduler; -class PauseRequests; -class AssignReqSeqSlots; -class AllocateKvCache; -class HandleContextLogits; -class HandleGenerationLogits; -class GenerateRequestOptions; -class LogitsPostProcessor; -class MakeDecodingBatchInputOutput; -class CreateNewDecoderRequests; -class UpdateDecoderBuffers; - -namespace utils -{ -class CudaGraphExecutorCache; -} // namespace utils - -struct RewindInputs -{ - SizeType32 maxBlocksPerSeq; - bool isUseOneMoreBlock; - SizeType32 numKvHeads; -}; - -class TrtGptModelInflightBatching : public TrtGptModel -{ - using BaseKVCacheManager = kv_cache_manager::BaseKVCacheManager; - using OffsetTableDimensions = kv_cache_manager::OffsetTableDimensions; - using KVCacheManager = kv_cache_manager::KVCacheManager; - using KvCacheType = kv_cache_manager::CacheType; - using KvCacheConfig = executor::KvCacheConfig; - using RnnStateManager = rnn_state_manager::RnnStateManager; - using LlmRequestPtr = std::shared_ptr; - -public: - class IterationStatsIFB - { - public: - explicit IterationStatsIFB(SizeType32 microBatchId) - : microBatchId{microBatchId} - { - } - - SizeType32 microBatchId; - SizeType32 numCtxRequests{}; - SizeType32 numGenRequests{}; - SizeType32 numCtxTokens{}; - float avgNumDecodedTokensPerIter{}; - ReqIdsSet scheduledRequests; - ReqIdsSet pausedRequests; - }; - - using SizeType32 = tensorrt_llm::runtime::SizeType32; - using TokenIdType = tensorrt_llm::runtime::TokenIdType; - using BufferManager = tensorrt_llm::runtime::BufferManager; - using PeftTable = PeftCacheManager::PeftTable; - using TensorMap = runtime::StringPtrMap; - using TensorPtr = runtime::ITensor::SharedPtr; - - TrtGptModelInflightBatching(std::shared_ptr logger, runtime::ModelConfig const& modelConfig, - runtime::WorldConfig const& worldConfig, runtime::RawEngine const& rawEngine, bool ctxGenFusion, - executor::ExecutorConfig const& executorConfig, bool isLeaderInOrchMode); - - ~TrtGptModelInflightBatching() override; - - /// @brief Calculate the cache size per token for the disaggregated serving. - /// @param modelConfig Model configuration. - /// @param worldConfig World configuration. - /// @param maxAttentionWindowVec Maximum attention window vector. (may have fewer elements than numLayers, in which - /// case it cycles) - /// @param isCrossAttention Whether the attention is cross attention. - /// @param kvFactor KV factor. - /// @return Cache size per token for the disaggregated layers. Note that window size is not included in the result - /// here. - [[nodiscard]] static std::map calculateCacheSizePerTokenForDisagg( - runtime::ModelConfig const& modelConfig, runtime::WorldConfig const& worldConfig, - std::vector const& maxAttentionWindowVec, bool isCrossAttention, SizeType32 kvFactor); - - void terminateRequest(LlmRequestPtr const& llmRequest, bool pause = false) override; - - /// @brief Terminate request in the next forwardSync call that includes the request. - /// @details This function does not terminate requests immediately. It will add the requests to the - /// mReqIdsToTerminate set. The requests will be terminated in the next forwardSync call that - /// includes the request in the batch. - void terminateRequestSync(LlmRequestPtr const& llmRequest, executor::FinishReason finishReason) override; - - /// @brief Function that waits for the decoding of requests in flight. - /// When the requests have finished or using speculative decoding, the state of requests - /// will become LlmRequestState::kGENERATION_COMPLETE. Else, it will be set to - /// LlmRequestState::kGENERATION_IN_PROGRESS. - void forwardSync() override; - - /// @brief Function that tries to advance the active requests. - /// Depending on resources available, it's possible that not all requests will get advanced. - /// Requests that may be in state LlmRequestState::kCONTEXT_INIT become - /// LlmRequestState::kGENERATION_IN_PROGRESS or LlmRequestState::kGENERATION_TO_COMPLETE. - /// @param activeRequests The list of request to try to advance. - void forwardAsync(RequestList const& activeRequests) override; - - /// @brief Override the runtime batch size for the model - void setRuntimeBatchSize(SizeType32 runtimeMaxBatchSize) override; - - /// @brief Get the runtime batch size for the model - [[nodiscard]] SizeType32 getRuntimeBatchSize() const override; - - /// @brief Override the runtime max num tokens for the model - void setRuntimeMaxNumTokens(SizeType32 runtimeMaxNumTokens) override; - - void updatePeftCache(std::shared_ptr const& llmRequest) override; - - [[nodiscard]] IterationStatsIFB getLastIterationStats() const - { - return mLastIterationStatsIFB; - } - - [[nodiscard]] TrtGptModelType getModelType() const override - { - return mCtxGenFusion ? TrtGptModelType::InflightFusedBatching : TrtGptModelType::InflightBatching; - }; - - [[nodiscard]] runtime::BufferManager const& getBufferManager() const override; - [[nodiscard]] runtime::BufferManager::CudaStreamPtr getRuntimeStreamPtr() const override; - - void getCurrentIterationStats(executor::IterationStats& stats) const override; - void getCurrentRequestStats(executor::RequestStatsPerIteration& stats) const override; - [[nodiscard]] executor::DebugTensorsPerIteration getCurrentDebugTensors() const override; - - [[nodiscard]] executor::IterationType getIterCounter() const noexcept override - { - return mIterCounter; - } - - [[nodiscard]] static bool executorConfigIsValid( - runtime::ModelConfig const& modelConfig, executor::ExecutorConfig const& executorConfig); - [[nodiscard]] static executor::ExecutorConfig fixExecutorConfig( - runtime::ModelConfig const& modelConfig, executor::ExecutorConfig const& executorConfig); - - void prepareDisaggGenInitRequests(RequestList const& activeRequests, RequestVector& newGenReques); - void checkDisaggGenTransferStatus(RequestList const& activeRequests); - void prepareDistGenBufferAndDecoder(RequestVector const& generationRequests); - - void resetIterationStats() override; - - runtime::SpeculativeDecodingMode getSpeculativeDecodingMode() const noexcept - { - return mModelConfig.getSpeculativeDecodingMode(); - } - - [[nodiscard]] SizeType32 numCachedCudaGraphs() const; - -private: - friend class TrtGptModelTest; - - [[nodiscard]] SizeType32 getContextBufferId() const - { - return mMicroBatchId; - } - - [[nodiscard]] SizeType32 getGenerationBufferId() const - { - return mNumMicroBatches + mMicroBatchId; - } - - [[nodiscard]] SizeType32 getFusedBufferId() const - { - return mMicroBatchId; - } - - [[nodiscard]] SizeType32 getNextMicroBatchId(SizeType32 bufferId) const - { - return (bufferId + 1) % mNumMicroBatches; - } - - [[nodiscard]] SizeType32 getPrevMicroBatchId(SizeType32 bufferId) const - { - return (bufferId + mNumMicroBatches - 1) % mNumMicroBatches; - } - - //! @brief Store full kv cache blocks contributed by req. - //! These blocks become reusable from next step. - void storeContextBlocks(std::shared_ptr const& req); - - //! @brief Store newest kv cache block for reuse. - //! The block become reusable from next step. - void storeNewBlock(std::shared_ptr const& req); - - //! @brief Set LayerProfiler to collect performance per layer. - void setLayerProfiler() override; - - //! @brief Print profile information per layer. - std::string getLayerProfileInfo() const override; - - std::tuple prepareBuffers( - RequestVector const& contextRequests, RequestVector const& generationRequests, SizeType32 bufferId); - - //! @brief Capture graph of current batch state during engine execution. - //! This is based on the assumptions that - //! a) We can hide CPU graph capture behind the GPU engine execution. - //! b) Batch size in the next iterations won't change and we can reuse the graph multiple times. - void prepareGraph(SizeType32 bufferId, SizeType32 optProfileId); - - void executeContext(SizeType32 runtimeContextId, SizeType32 bufferId); - void executeBatch(ScheduledRequests const& scheduledRequests); - void executeStep( - RequestVector const& contextRequests, RequestVector const& generationRequests, SizeType32 bufferId); - - void debugIOTensors(RequestVector const& contextRequests, RequestVector const& generationRequests, - TensorMap const& inputMap, TensorMap const& outputMap); - - void createRuntimeContexts(); - void createDecoder(std::optional const& decodingModeOpt); - void createBuffers(executor::DecodingConfig const& decodingConfig, - std::optional> const& additionalModelOutputs); - std::unique_ptr createKvCacheManager(KvCacheConfig const& kvCacheConfig, KvCacheType kvCacheType, - uint64_t freePrimaryMemBytes, uint64_t freeSecondaryMemBytes, size_t extraCostMemory, - bool const failFastOnAttentionWindowTooLarge = false); - void createRnnStateManager(); - void createCustomAllReduceWorkspace(); - void createRuntimePerfKnobsTensor(executor::ExtendedRuntimePerfKnobConfig const& extendedRuntimePerfKnobConfig); - - /// @brief Verify draft token length and beam width of all active requests. - /// May change operating beam width if all requests agree on same beam width. - void verifyRequests(RequestList const& activeRequests); - - /// @brief Change the operating beam width. - /// Only possible if no requests are currently in-flight. - /// @param beamWidth New operating beam width. Must be smaller than initial maxBeamWidth. - void changeBeamWidth(SizeType32 beamWidth); - - SizeType32 getOperatingBeamWidth() const override - { - return mOperatingBeamWidth; - } - - /// @details Should be called after setting up the current batch in executeBatch to get the correct number of - /// context tokens. - IterationStatsIFB fillIterationStats( - ScheduledRequests const& scheduledRequests, RequestVector const& requestsToPause); - - /// @brief Function that sets up the TensorRT execution context that is going to be used for execution. If multiple - /// TensorRT optimization profiles are built in the engine, it selects the corresponding context that is going to be - /// used, and prepares the input and output tensors so that both buffers and the context is ready for the execution. - /// @return The TensorRT execution context index that has been setup. - void setupContext( - RequestVector const& contextRequests, RequestVector const& generationRequests, SizeType32 bufferId); - - void setupDecoderStep( - RequestVector const& contextRequests, RuntimeBuffers const& buffers, DecoderInputBuffers& inputBuffers); - runtime::CudaEvent decoderStepAsync(ScheduledRequests const& scheduledRequests); - std::vector> decoderSync( - ScheduledRequests const& scheduledRequests, std::optional const& decoderFinishEvent); - - std::vector> communicateDecoderBuffers(bool returnLogProbs); - void updateRequests(ScheduledRequests const& scheduledRequests); - - /// @brief It gathers the logits if they need to be returned, calls getDecoderSlotHostOutputs, - /// and overwrites the llmRequest tokens buffer. - /// Called either on request finishing, or at every step when doing beam search and streaming. - void postProcessRequest(LlmRequest& llmReq, std::vector const& numDroppedTokens); - /// @brief Reorders generation logits to match finalized beam paths after gatherTree. - /// During beam search, logits are stored by beam slot. After finalization, output_ids are - /// reordered by parentIds, but logits are not. This method traces parentIds on the host - /// to build the slot mapping and reindexes the logits accordingly. - void reorderGenerationLogitsForBeamSearch(LlmRequest& llmReq, SizeType32 seqSlot, SizeType32 reqBeamWidth, - SizeType32 maxSeqLength, TokenIdType const* outputIdsHostData, SizeType32 const* sequenceLengthsHostData); - /// @brief Calls gatherTree (via finalize) and transmits the received data across ranks if PP>1 - void getDecoderSlotHostOutputs( - SizeType32 seqSlot, bool returnLogProbs, runtime::SamplingConfig const& samplingConfig, bool streaming); - void rewindKVCacheBlocks(SizeType32 numSequences); - void setupSpeculativeDecodingModule(executor::DecodingConfig const& decodingConfig); - - /// @brief Copies the content of the cache indirection outputs to the cache indirection inputs. - /// @param[in] scheduledRequests The requests to copy the cache indirections for. - /// @param[in] genBufferId The id of the generation buffers for those requests. - void copyCacheIndirectionFromOutputsToInputs(ScheduledRequests const& scheduledRequests, SizeType32 genBufferId); - - [[nodiscard]] bool getGatherGenerationLogits() const override - { - return getModelConfig().computeGenerationLogits() || mGatherGenerationLogits; - } - - [[nodiscard]] runtime::ModelConfig const& getModelConfig() const override - { - return mModelConfig; - } - - [[nodiscard]] runtime::WorldConfig const& getWorldConfig() const override - { - return mWorldConfig; - } - - [[nodiscard]] SizeType32 getNumMicroBatches() const override - { - return mNumMicroBatches; - } - - [[nodiscard]] nvinfer1::DataType getLogitDataType() const override; - - [[nodiscard]] nvinfer1::DataType getTensorDataType(std::string const& name) const override; - - [[nodiscard]] nvinfer1::Dims getTensorShape(std::string const& name) const override; - - void reshapeKvTensors(OffsetTableDimensions const& dims); - - [[nodiscard]] bool hasSpeculativeDecodingFastLogits() const noexcept override - { - return mSpeculativeDecodingFastLogits; - } - - [[nodiscard]] bool hasGuidedDecoder() const noexcept override - { - return static_cast(mGuidedDecoder); - } - - using BlocksPerWindow = std::map>; - /// @brief Based on the KV-cache manager's capacity and configuration, we adjust the maximum supported attention - /// window. - /// - /// @param blocksPerWindow map of window size to number of blocks. - /// @param failFastOnAttentionWindowTooLarge if true, the function will report a runtime error if the attention - /// window is too large to fit even a single sequence in the KV cache. - /// @return pair of new blocks per window and new maxAttentionWindowVec - [[nodiscard]] std::pair> clampWindowSizesToFitAtLeastOneSequence( - BlocksPerWindow const& blocksPerWindow, bool const failFastOnAttentionWindowTooLarge = false); - - /// @brief Change the speculative decoding mode. - void changeSpecDecMode(ScheduledRequests const& scheduledRequests); - - void prefetchNextPromptTableChunk(RequestVector const& contextRequests, bool isFirstChunk, SizeType32 bufferId); - - void remapInputTokensForPromptTable( - std::shared_ptr const& llmReq, bool isCurrentChunk, SizeType32 bufferId, SizeType32 contextId); - - void copyPromptTableToGpuInChunk(std::shared_ptr const& llmReq, - std::vector const& outOfVocabTokens, bool useCurrentBuffer, SizeType32 bufferId, SizeType32 contextId); - -protected: - std::shared_ptr getKVCacheManager() override - { - return mKvCacheManager; - } - - [[nodiscard]] std::shared_ptr getKVCacheManager() const override - { - return mKvCacheManager; - } - - std::shared_ptr getCrossKVCacheManager() - { - return mCrossKvCacheManager; - } - - [[nodiscard]] std::shared_ptr getCrossKVCacheManager() const - { - return mCrossKvCacheManager; - } - - [[nodiscard]] std::shared_ptr getPeftCacheManager() override - { - return mPeftCacheManager; - } - - [[nodiscard]] std::shared_ptr getPeftCacheManager() const override - { - return mPeftCacheManager; - } - - void setLogitsPostProcessorBatched(std::optional logitsPostProcessorBatched) override - { - mLogitsPostProcessorBatched = logitsPostProcessorBatched; - } - - void setReplicateLogitsPostProcessor(bool replicateLogitsPostProcessor) override - { - mReplicateLogitsPostProcessor = replicateLogitsPostProcessor; - } - - [[nodiscard]] bool getReplicateLogitsPostProcessor() const override - { - return mReplicateLogitsPostProcessor; - } - - SizeType32 getMaxCapacityBatchSize(SizeType32 inputLength, SizeType32 outputLength) const override; - -private: - /******************** Configs ********************/ - // Parameters of the model (TRT engine) - runtime::ModelConfig mModelConfig; - // Parameters of the execution environment - runtime::WorldConfig mWorldConfig; - // Device ID of this instance - int mDevice{-1}; - // Config for (speculative) decoding - executor::DecodingConfig mDecodingConfig; - // Performance knobs for the engine. - executor::ExtendedRuntimePerfKnobConfig mExtendedRuntimePerfKnobConfig; - TensorPtr mExtendedRuntimePerfKnobsHost; - // Config for debugging output - std::optional mDebugConfig; - // List of additional outputs for each request - std::optional> mAdditionalModelOutputs; - - /******************** Components ********************/ - std::shared_ptr mLogger; - // Runner for the TRT engine. The engine produces logits. - std::unique_ptr mRuntime; - // Decoder that generates new tokens from the logits. - std::unique_ptr mDecoder; - // Decoder state for all requests - std::unique_ptr mDecoderState; - // Synchronization handles for decoder - std::vector> mDecoderFinishedEvents; - - // Manager that maps requests to slots - std::shared_ptr mSeqSlotManager; - // KV cache manager for attention layers (optional) - std::shared_ptr mKvCacheManager; - // KV cache manager for cross attention in enc-dec models (optional) - std::shared_ptr mCrossKvCacheManager = nullptr; - // RNN state manager for recurrent layers (optional) - std::unique_ptr mRnnStateManager; - // PEFT cache manager for LoRA tasks (optional) - std::shared_ptr mPeftCacheManager; - // BufferManager using a separate stream for async copy operations. - runtime::BufferManager mCopyBufferManager; - // Event for async data transfers - runtime::CudaEvent mPtableCopyDoneEvent; - - /******************** Logits Post-Processor ********************/ - std::optional mLogitsPostProcessorBatched; - bool mReplicateLogitsPostProcessor{true}; - // Set if any request invoked a logits processor in current step - bool mLogitsPostProcessorIsApplied{false}; - - constexpr bool broadcastPostDecoder() - { - return mWorldConfig.isTensorParallel() && !mReplicateLogitsPostProcessor && mLogitsPostProcessorIsApplied; - } - - std::unique_ptr mGuidedDecoder; - - /******************** Pipeline parallelism ********************/ - std::unique_ptr mMpiCommPipelinePara; - std::vector> mDecStepAsyncSndHdls; - std::vector> mDecSlotAsyncSndHdls; - std::unique_ptr mAsyncSendWaitThread; - - /******************** Tensor parallelism ********************/ - std::unique_ptr mMpiCommTensorPara; - std::unique_ptr mAllReduceBuffers; - - /******************** Runtime parameters ********************/ - // Flag to select fused or unfused context+generation execution - bool mCtxGenFusion; - // ID of current micro batch, changes after each iteration - SizeType32 mMicroBatchId{0}; - // Number of micro batches. Multiple batches are used for overlapping setup and execution, - // and in pipeline parallelism. - SizeType32 mNumMicroBatches; - // Number of buffers to be added to mBuffers. - SizeType32 mNumBuffers; - // Current operating beam width. Can be changed with changeBeamWidth function. - SizeType32 mOperatingBeamWidth; - // Runtime batch size optimized during execution for microBatchScheduler: - /// The max batch size recommended by the dynamic tuner - SizeType32 mMaxBatchSizeTunerRecommended; - /// The min of mMaxBatchSize and mMaxBatchSizeTunerRecommended - SizeType32 mMaxBatchSizeRuntime; - // Runtime max num tokens optimized during execution for microBatchScheduler: - /// Build time max num tokens - std::optional mMaxNumTokensStatic; - /// The max num tokens recommended by the dynamic tuner - SizeType32 mMaxNumTokensTunerRecommended; - /// The min of mMaxNumTokens and mMaxNumTokensTunerRecommended - std::optional mMaxNumTokensRuntime; - // Controls if generation logits should be gathered, so that returnGenerationLogits can be requested. - bool mGatherGenerationLogits{false}; - // offloading and prefetching the prompt tuning table (only effective in chunked prefill mode) - bool mPromptTableOffloading; - - /******************** Buffers ********************/ - // Buffers for each micro batch. Unfused path (mCtxGenFusion==false) uses two times the buffers. - std::vector> mBuffers; - // Decoder input buffers for each micro batch. - std::vector mDecoderInputBuffers; - // Decoder output buffers for each micro batch. - std::vector mDecoderOutputBuffers; - // Buffers for each slot in the decoder - std::vector> mSlotDecoderBuffers; - // PEFT table for each micro batch - std::vector mPeftTables; - - /******************** Book keeping ********************/ - // List of requests in each micro batch - std::vector mMicroBatchScheduledRequests; - // Set of in-flight requests of *all* micro batches - ReqIdsSet mInflightReqIds; - // Requests that should be terminated (requested from outside the model) - std::unordered_map mReqIdsToTerminate; - // Requests that the scheduler selected to be paused - ReqIdsSet mReqIdsToPause; - // Stats collected in last iteration - IterationStatsIFB mLastIterationStatsIFB{-1}; - // Iteration counter used to distinguish debug output - executor::IterationType mIterCounter{0}; - // Debug tensors of last itreation - TensorMap mLastIterationDebugTensors; - // Cuda graph instances for each microbatch. - std::vector mCudaGraphExecutorCaches; - - /******************** Cache transceiver ********************/ - std::unique_ptr mCacheTransceiver; - std::unique_ptr mDisaggTransferAdmissionController; - - /******************** Spec dec ***********************/ - std::unique_ptr mDraftModelSendLogitsThread; - bool mSpeculativeDecodingFastLogits; - std::atomic mDraftModelThreadShouldExit{false}; - bool mIsLeaderInOrchMode{false}; - // List of completed draft requests which logits will need to be sent to the target model. - // Guarded by mDraftRequestsMtx (shared with the background logits sender thread). - RequestVector mDraftRequestsWaitingToSendLogits; - // Draft requests whose logits have been sent — pending termination by main thread. - // Guarded by mDraftRequestsMtx. - RequestVector mDraftRequestsDoneSendingLogits; - std::mutex mDraftRequestsMtx; - SizeType32 mSeamlessLADMaxDraftLen{0}; - bool mUseSeamlessLookahead{false}; - RewindInputs mRewindInputs; - - /******************** Algorithms ********************/ - // Algorithms are reentrant, they are assigned a state at - // construction time and it is not modified through execution, hence they are const. - // Schedulers that select which requests to run in each iteration - std::unique_ptr mCapacityScheduler; - std::unique_ptr mMicroBatchScheduler; - std::unique_ptr mPauseRequests; - std::unique_ptr mAssignReqSeqSlots; - std::unique_ptr mAllocateKvCache; - std::unique_ptr mHandleContextLogits; - std::unique_ptr mHandleGenerationLogits; - std::unique_ptr mLogitsPostProcessor; - std::unique_ptr mMakeDecodingBatchInputOutput; - std::unique_ptr mCreateNewDecoderRequests; - std::unique_ptr mUpdateDecoderBuffers; -}; - -} // namespace tensorrt_llm::batch_manager diff --git a/cpp/tensorrt_llm/batch_manager/updateDecoderBuffers.cpp b/cpp/tensorrt_llm/batch_manager/updateDecoderBuffers.cpp deleted file mode 100644 index ead120135f3a..000000000000 --- a/cpp/tensorrt_llm/batch_manager/updateDecoderBuffers.cpp +++ /dev/null @@ -1,78 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2025 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. - */ - -#include "tensorrt_llm/batch_manager/updateDecoderBuffers.h" -#include "tensorrt_llm/batch_manager/decoderBuffers.h" -#include "tensorrt_llm/common/nvtxUtils.h" -#include "tensorrt_llm/runtime/decoderState.h" -#include "tensorrt_llm/runtime/iTensor.h" - -namespace tensorrt_llm::batch_manager -{ - -using BufferManager = tensorrt_llm::runtime::BufferManager; -using TensorPtr = runtime::ITensor::SharedPtr; -using ITensor = runtime::ITensor; -using SizeType32 = tensorrt_llm::runtime::SizeType32; - -runtime::CudaEvent UpdateDecoderBuffers::operator()(runtime::ModelConfig const& modelConfig, - DecoderOutputBuffers& decoderOutputBuffers, runtime::BufferManager const& copyBufferManager, - runtime::decoder::DecoderState const& decoderState, bool returnLogProbs, - runtime::CudaEvent const& decoderFinishEvent) const -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(updateDecoderBuffers); - - // Chain copy after decoder event, using a different stream - copyBufferManager.getStream().wait(decoderFinishEvent); - - copyBufferManager.copy(*decoderState.getAllNewTokens(), *decoderOutputBuffers.newOutputTokensHost); - copyBufferManager.copy(*decoderState.getSequenceLengths(), *decoderOutputBuffers.sequenceLengthsHost); - - auto const finishedSumDevice = decoderState.getFinishedSum(); - copyBufferManager.copy(*finishedSumDevice, *decoderOutputBuffers.finishedSumHost); - auto const finishReasonsDevice = decoderState.getFinishReasons(); - copyBufferManager.copy(*finishReasonsDevice, *decoderOutputBuffers.finishReasonsHost); - - if (returnLogProbs) - { - copyBufferManager.copy(*decoderState.getCumLogProbs(), *decoderOutputBuffers.cumLogProbsHost); - copyBufferManager.copy(*decoderState.getLogProbs(), *decoderOutputBuffers.logProbsHost); - } - - if (modelConfig.getSpeculativeDecodingMode().predictsDraftTokens()) - { - // TODO: keep data on device for next iteration - copyBufferManager.copy(*decoderState.getNextDraftTokens(), *decoderOutputBuffers.nextDraftTokensHost); - - if (modelConfig.getSpeculativeDecodingMode().variableDraftLength()) - { - copyBufferManager.copy( - *decoderState.getNextDraftTokensLengths(), *decoderOutputBuffers.nextDraftTokensLengthsHost); - copyBufferManager.copy( - *decoderState.getPrevDraftTokensLengths(), *decoderOutputBuffers.prevDraftTokensLengthsHost); - } - } - - runtime::CudaEvent copyEvent{}; - copyBufferManager.getStream().record(copyEvent); - // Store the event for later sync. Sync stream before calling next decoder. Sync host before updating requests. - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); - return copyEvent; -} - -} // namespace tensorrt_llm::batch_manager diff --git a/cpp/tensorrt_llm/batch_manager/utils/debugUtils.h b/cpp/tensorrt_llm/batch_manager/utils/debugUtils.h index e4732a75f649..c041c7a71de8 100644 --- a/cpp/tensorrt_llm/batch_manager/utils/debugUtils.h +++ b/cpp/tensorrt_llm/batch_manager/utils/debugUtils.h @@ -23,11 +23,6 @@ #include "tensorrt_llm/runtime/iTensor.h" #include "tensorrt_llm/runtime/worldConfig.h" -namespace tensorrt_llm::runtime -{ -class TllmRuntime; -} // namespace tensorrt_llm::runtime - namespace tensorrt_llm::batch_manager::utils { diff --git a/cpp/tensorrt_llm/batch_manager/utils/inflightBatchingUtils.cpp b/cpp/tensorrt_llm/batch_manager/utils/inflightBatchingUtils.cpp index 416235f347b8..a3e54a6b0f9b 100644 --- a/cpp/tensorrt_llm/batch_manager/utils/inflightBatchingUtils.cpp +++ b/cpp/tensorrt_llm/batch_manager/utils/inflightBatchingUtils.cpp @@ -16,7 +16,6 @@ */ #include "inflightBatchingUtils.h" -#include "tensorrt_llm/runtime/runtimeKernels.h" namespace tensorrt_llm::batch_manager::utils { @@ -88,170 +87,6 @@ void moveFinishedContextRequestsToGeneration(ScheduledRequests& scheduledRequest TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); } -void copyGenerationLogits(RuntimeBuffers::GenerationLogitsCache& generationLogitsCache, - runtime::BufferManager const& bufferManager, LlmRequest& llmReq, bool beforeDecoder, - std::vector const& numDroppedTokens) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - TLLM_CHECK_WITH_INFO( - !beforeDecoder || numDroppedTokens.empty(), "numDroppedTokens are only possible after decoder."); - - auto const reqBeamWidth = llmReq.getBeamWidthByIter(); - TLLM_CHECK_WITH_INFO(numDroppedTokens.empty() || numDroppedTokens.size() == static_cast(reqBeamWidth), - "Dropped tokens have to be defined for all beams."); - - auto const fragmentSize = llmReq.getGenerationLogitsFragmentsSize(); - - // Merge logits fragments on device. getFragmentPointerSlot() returns the matching host and - // device rows for the current workIdx and advances the index atomically, so concurrent flushes - // for different requests in the same batch never clobber each other's pointer arrays. - auto const& transposeBufferPtr = generationLogitsCache.transposedLogits; - auto [cachePointerHost, cachePointerDevice] = generationLogitsCache.getFragmentPointerSlot(); - tensorrt_llm::runtime::kernels::mergeLogitsFragments(bufferManager, *transposeBufferPtr, - llmReq.getGenerationLogitsFragments(), *cachePointerDevice, *cachePointerHost, 0, 1, reqBeamWidth, - bufferManager.getStream(), 0); - llmReq.clearGenerationLogitsFragments(); - - // Copy logits to host - for (SizeType32 beam = 0; beam < reqBeamWidth; beam++) - { - auto const droppedSize = !numDroppedTokens.empty() ? numDroppedTokens.at(beam) : 0; - // Ignore logits of dropped tokens - auto const beamFragmentSize = fragmentSize - droppedSize; - // If this function is called before the decoder, the request does not contain the generated token of the - // current iteration, so we add 1 to the number of tokens. - auto const numGenerationToken - = static_cast(beforeDecoder) + llmReq.getNumTokens(beam) - llmReq.mPromptLen; - auto const hostOffset = numGenerationToken - beamFragmentSize; - - // [beamWidth, GENERATION_LOGITS_BUFFER_LENGTH, vocabSizePadded] -> [beamFragmentSize, vocabSizePadded] - auto beamDeviceTensorPtr = ITensor::slice(transposeBufferPtr, {beam, 0}, beamFragmentSize); - // [beamWidth, mMaxNewTokens, vocabSizePadded] -> [beamFragmentSize, vocabSizePadded] - auto beamHostTensorPtr = ITensor::slice(llmReq.getGenerationLogitsHost(), {beam, hostOffset}, beamFragmentSize); - bufferManager.copy(*beamDeviceTensorPtr, *beamHostTensorPtr); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -namespace -{ - -std::pair findOutputTensor(std::string const& outputTensorName, - std::vector const& additionalModelOutputs, - RuntimeBuffers::TensorMap const& outputMap, bool isContext) -{ - auto const aoIter = std::find_if(additionalModelOutputs.cbegin(), additionalModelOutputs.cend(), - [&outputTensorName](auto const& ao) { return ao.name == outputTensorName; }); - TLLM_CHECK_WITH_INFO(aoIter != additionalModelOutputs.cend(), "Additional %s output tensor not found: %s", - isContext ? "context" : "generation", outputTensorName.c_str()); - - auto const gatherContext = aoIter->gatherContext; - if (isContext) - { - TLLM_CHECK_WITH_INFO( - gatherContext, "Additional context output tensor not gathered: %s", outputTensorName.c_str()); - } - - auto const tensorIt = outputMap.find(outputTensorName); - TLLM_CHECK_WITH_INFO(tensorIt != outputMap.end(), "Additional %s output tensor not found: %s", - isContext ? "context" : "generation", outputTensorName.c_str()); - - return {tensorIt->second, gatherContext}; -} - -} // namespace - -void copyAdditionalOutputs(std::vector const& additionalModelOutputs, - RequestVector const& contextRequests, RequestVector const& generationRequests, - RuntimeBuffers::TensorMap const& outputMap, runtime::BufferManager const& manager) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - // One index shared across all output tensors that have gatherContext - SizeType32 srcTensorIndexWithContext{0}; - // One index shared across all output tensors that do not have gatherContext - SizeType32 srcTensorIndexWithoutContext{0}; - - for (auto const& llmReq : contextRequests) - { - auto numContextTokens = llmReq->getContextChunkSize(); - for (auto const& outputTensor : llmReq->getAdditionalContextOutputs()) - { - auto const& [tensor, gatherContext] - = findOutputTensor(outputTensor.first, additionalModelOutputs, outputMap, true); - - auto const srcTensorIndex = srcTensorIndexWithContext; - auto srcView = ITensor::slice(tensor, srcTensorIndex, numContextTokens); - auto dstView = ITensor::slice(outputTensor.second, llmReq->getContextCurrentPosition(), numContextTokens); - manager.copy(*srcView, *dstView); - } - srcTensorIndexWithContext += numContextTokens; - srcTensorIndexWithoutContext += 1; - - // Copy output of last token to generation outputs - if (llmReq->isLastContextChunk()) - { - for (auto const& outputTensor : llmReq->getAdditionalGenerationOutputs()) - { - auto const& [tensor, gatherContext] - = findOutputTensor(outputTensor.first, additionalModelOutputs, outputMap, false); - - auto const srcTensorIndex = gatherContext ? srcTensorIndexWithContext : srcTensorIndexWithoutContext; - auto srcView = ITensor::slice(tensor, srcTensorIndex - 1, 1); - for (SizeType32 beam = 0; beam < llmReq->getBeamWidthByIter(); beam++) - { - auto dstView = ITensor::slice(outputTensor.second, {beam, 0}, 1); - manager.copy(*srcView, *dstView); - } - } - } - } - - for (auto const& llmReq : generationRequests) - { - auto const reqBeamWidth = llmReq->getBeamWidthByIter(); - for (auto const& outputTensor : llmReq->getAdditionalGenerationOutputs()) - { - auto const& [tensor, gatherContext] - = findOutputTensor(outputTensor.first, additionalModelOutputs, outputMap, false); - - auto const srcTensorIndex = gatherContext ? srcTensorIndexWithContext : srcTensorIndexWithoutContext; - for (SizeType32 beam = 0; beam < reqBeamWidth; beam++) - { - auto const generatedLength = llmReq->getNumTokens(beam) - llmReq->getPromptLen(); - TLLM_CHECK(generatedLength >= 1); - auto srcView = ITensor::slice(tensor, srcTensorIndex + beam, 1); - auto dstView = ITensor::slice(outputTensor.second, {beam, generatedLength}, 1); - manager.copy(*srcView, *dstView); - } - } - srcTensorIndexWithContext += reqBeamWidth; - srcTensorIndexWithoutContext += reqBeamWidth; - } - - // Check final indices - for (auto const& outputTensor : additionalModelOutputs) - { - auto const& outputTensorName = outputTensor.name; - auto const gatherContext = outputTensor.gatherContext; - - auto const tensorIt = outputMap.find(outputTensorName); - TLLM_CHECK_WITH_INFO( - tensorIt != outputMap.end(), "Additional output tensor not found: %s", outputTensorName.c_str()); - - auto const& outputShape = tensorIt->second->getShape(); - auto const outputSize = outputShape.d[0]; - auto const finalIndex = gatherContext ? srcTensorIndexWithContext : srcTensorIndexWithoutContext; - - TLLM_CHECK_WITH_INFO(finalIndex == outputSize, "Additional %s output tensor final index mismatch %d != %ld: %s", - gatherContext ? "context" : "generation", finalIndex, outputSize, outputTensorName.c_str()); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - void terminateRequest(SequenceSlotManager& seqSlotManager, LlmRequest& llmReq, SizeType32 maxInputLen, OptionalRef kvCacheManager, OptionalRef crossKvCacheManager, @@ -300,109 +135,4 @@ std::vector getRequestBeamWidths( return beamWidths; } -void CudaGraphExecutor::create(cudaGraph_t const& graph) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - assert(mInstance == nullptr); - TLLM_CUDA_CHECK(cudaGraphInstantiate(&mInstance, graph, nullptr, nullptr, 0)); - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void CudaGraphExecutor::uploadToStream(runtime::CudaStream const& stream) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - assert(hasInstance()); - TLLM_CUDA_CHECK(cudaGraphUpload(mInstance, stream.get())); - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void CudaGraphExecutor::launch(runtime::CudaStream const& stream) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - TLLM_CUDA_CHECK(cudaGraphLaunch(mInstance, stream.get())); - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -bool CudaGraphExecutor::update(cudaGraph_t const& graph) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - return cudaGraphExecUpdate(mInstance, graph, nullptr) != cudaSuccess; -} - -void CudaGraphExecutor::clear() -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - if (mInstance != nullptr) - { - TLLM_CUDA_CHECK(cudaGraphExecDestroy(mInstance)); - mInstance = nullptr; - } - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void CudaGraphExecutor::prepareNextGraph(std::unique_ptr& runtime, SizeType32 nextContextId) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - auto& stream = runtime->getStream(); - - cudaGraph_t nextGraph; - TLLM_CUDA_CHECK(cudaStreamBeginCapture(stream.get(), cudaStreamCaptureModeThreadLocal)); - runtime->executeContext(nextContextId); - TLLM_CUDA_CHECK(cudaStreamEndCapture(stream.get(), &nextGraph)); - - if (hasInstance()) - { - if (update(nextGraph)) - { - clear(); - create(nextGraph); - } - } - else - { - create(nextGraph); - } - - TLLM_CUDA_CHECK(cudaGraphDestroy(nextGraph)); - uploadToStream(stream); - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -std::optional> CudaGraphExecutorCache::get(BatchState const& state) -{ - auto it = mMap.find(state); - if (it == mMap.end()) - { - return std::nullopt; - } - mCache.splice(mCache.begin(), mCache, it->second); - return it->second->second; -} - -void CudaGraphExecutorCache::put(BatchState const& state, std::shared_ptr const& value) -{ - auto it = mMap.find(state); - if (it != mMap.end()) - { - mCache.erase(it->second); - } - mCache.emplace_front(state, value); - mMap[state] = mCache.begin(); - - if (static_cast(mMap.size()) > mCapacity) - { - auto lastState = mCache.back().first; - mCache.pop_back(); - mMap.erase(lastState); - } -} - -void CudaGraphExecutorCache::clear() -{ - // Releasing the shared_ptrs runs ~CudaGraphExecutor, which calls - // cudaGraphExecDestroy on each cached instance. - mMap.clear(); - mCache.clear(); -} - } // namespace tensorrt_llm::batch_manager::utils diff --git a/cpp/tensorrt_llm/batch_manager/utils/inflightBatchingUtils.h b/cpp/tensorrt_llm/batch_manager/utils/inflightBatchingUtils.h index fe0c4e505218..374ae398f781 100644 --- a/cpp/tensorrt_llm/batch_manager/utils/inflightBatchingUtils.h +++ b/cpp/tensorrt_llm/batch_manager/utils/inflightBatchingUtils.h @@ -20,7 +20,6 @@ #include "tensorrt_llm/batch_manager/common.h" #include "tensorrt_llm/batch_manager/kvCacheManager.h" #include "tensorrt_llm/batch_manager/peftCacheManager.h" -#include "tensorrt_llm/batch_manager/runtimeBuffers.h" #include "tensorrt_llm/batch_manager/sequenceSlotManager.h" #include "tensorrt_llm/common/optionalRef.h" #include "tensorrt_llm/runtime/iTensor.h" @@ -50,17 +49,6 @@ void sortRequests(RequestVector& contextRequests, RequestVector& generationReque //! @param scheduledRequests The scheduled context and generation requests. void moveFinishedContextRequestsToGeneration(ScheduledRequests& scheduledRequests); -//! @param beforeDecoder Whether the function is called before the decoder. If it is true, correct the output offset. -//! @param numDroppedTokens The number of dropped tokens for each beam (e.g. when the requests finished early). -//! Generation logits for dropped tokens are ignored. -void copyGenerationLogits(RuntimeBuffers::GenerationLogitsCache& generationLogitsCache, - runtime::BufferManager const& bufferManager, LlmRequest& llmReq, bool beforeDecoder, - std::vector const& numDroppedTokens = {}); - -void copyAdditionalOutputs(std::vector const& additionalModelOutputs, - RequestVector const& contextRequests, RequestVector const& generationRequests, - RuntimeBuffers::TensorMap const& outputMap, runtime::BufferManager const& manager); - void terminateRequest(SequenceSlotManager& seqSlotManager, LlmRequest& llmRequest, SizeType32 maxInputLen, OptionalRef kvCacheManager = std::nullopt, OptionalRef crossKvCacheManager = std::nullopt, @@ -68,66 +56,4 @@ void terminateRequest(SequenceSlotManager& seqSlotManager, LlmRequest& llmReques std::vector getRequestBeamWidths( RequestVector const& contextRequests, RequestVector const& generationRequests); - -class CudaGraphExecutor -{ -public: - CudaGraphExecutor() = default; - - ~CudaGraphExecutor() - { - try - { - clear(); - } - catch (std::exception& e) - { - TLLM_LOG_EXCEPTION(e); - } - } - - bool hasInstance() const - { - return mInstance != nullptr; - } - - void clear(); - void prepareNextGraph(std::unique_ptr& runtime, SizeType32 nextContextId); - void launch(runtime::CudaStream const& stream); - -private: - void create(cudaGraph_t const& graph); - bool update(cudaGraph_t const& graph); - void uploadToStream(runtime::CudaStream const& stream); - - cudaGraphExec_t mInstance = nullptr; -}; - -class CudaGraphExecutorCache -{ - /// @brief LRU cache to store cuda graph instances. -public: - explicit CudaGraphExecutorCache(runtime::SizeType32 capacity) - : mCapacity(capacity) - { - } - - std::optional> get(BatchState const& state); - - void put(BatchState const& state, std::shared_ptr const& value); - - void clear(); - - [[nodiscard]] runtime::SizeType32 size() const noexcept - { - return static_cast(mCache.size()); - } - -private: - using BatchStateGraphExecutorPair = std::pair>; - using GraphExecutorLruCache = std::list; - SizeType32 mCapacity; - GraphExecutorLruCache mCache; - std::unordered_map mMap; -}; } // namespace tensorrt_llm::batch_manager::utils diff --git a/cpp/tensorrt_llm/batch_manager/utils/logitsThread.cpp b/cpp/tensorrt_llm/batch_manager/utils/logitsThread.cpp index 941c1b655073..4f978a302187 100644 --- a/cpp/tensorrt_llm/batch_manager/utils/logitsThread.cpp +++ b/cpp/tensorrt_llm/batch_manager/utils/logitsThread.cpp @@ -19,6 +19,7 @@ #include "tensorrt_llm/batch_manager/llmRequest.h" #include "tensorrt_llm/common/logger.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/executor/executor.h" #include "tensorrt_llm/runtime/utils/mpiTags.h" #include "tensorrt_llm/runtime/utils/mpiUtils.h" @@ -124,7 +125,7 @@ void draftModelSendLogitsThread(int device, std::atomic* draftModelThreadS } void targetModelReceiveLogits(runtime::ITensor::SharedPtr& draftLogitsHost, - executor::SpeculativeDecodingFastLogitsInfo const& fastLogitsInfo, nvinfer1::DataType logitsDtype) + executor::SpeculativeDecodingFastLogitsInfo const& fastLogitsInfo, tensorrt_llm::DataType logitsDtype) { #if ENABLE_MULTI_DEVICE auto const& worldComm = tensorrt_llm::mpi::MpiComm::world(); diff --git a/cpp/tensorrt_llm/batch_manager/utils/logitsThread.h b/cpp/tensorrt_llm/batch_manager/utils/logitsThread.h index 637f8a850610..7af3b1762d20 100644 --- a/cpp/tensorrt_llm/batch_manager/utils/logitsThread.h +++ b/cpp/tensorrt_llm/batch_manager/utils/logitsThread.h @@ -18,6 +18,7 @@ #pragma once #include "tensorrt_llm/batch_manager/common.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/executor/executor.h" #include "tensorrt_llm/runtime/common.h" #include "tensorrt_llm/runtime/iTensor.h" @@ -46,6 +47,6 @@ void draftModelSendLogitsThread(int device, std::atomic* draftModelThreadS std::mutex* draftRequestsMtx); void targetModelReceiveLogits(runtime::ITensor::SharedPtr& draftLogitsHost, - executor::SpeculativeDecodingFastLogitsInfo const& fastLogitsInfo, nvinfer1::DataType logitsDtype); + executor::SpeculativeDecodingFastLogitsInfo const& fastLogitsInfo, tensorrt_llm::DataType logitsDtype); } // namespace tensorrt_llm::batch_manager::utils diff --git a/cpp/tensorrt_llm/common/attentionOp.cpp b/cpp/tensorrt_llm/common/attentionOp.cpp index b91c0ef98df6..856f3e5a6155 100644 --- a/cpp/tensorrt_llm/common/attentionOp.cpp +++ b/cpp/tensorrt_llm/common/attentionOp.cpp @@ -22,6 +22,7 @@ #include "tensorrt_llm/common/logger.h" #include "tensorrt_llm/common/memoryUtils.h" #include "tensorrt_llm/common/sageQuant.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/kernels/decoderMaskedMultiheadAttention.h" #include "tensorrt_llm/kernels/decoderMaskedMultiheadAttention/cascadeAttentionKernel.h" #include "tensorrt_llm/kernels/flashMLA/flash_mla.h" @@ -758,8 +759,8 @@ size_t AttentionOp::getFmhaMultiCtasKvScratchSize() const noexcept return partialStatsSize + partialOSize; } -size_t AttentionOp::getWorkspaceSizeForContext(nvinfer1::DataType type, int32_t max_num_seq, int32_t input_seq_length, - int32_t cross_kv_length, int32_t max_num_tokens, int32_t total_kv_len) const noexcept +size_t AttentionOp::getWorkspaceSizeForContext(tensorrt_llm::DataType type, int32_t max_num_seq, + int32_t input_seq_length, int32_t cross_kv_length, int32_t max_num_tokens, int32_t total_kv_len) const noexcept { if (max_num_tokens == 0) { @@ -911,7 +912,7 @@ size_t AttentionOp::getWorkspaceSizeForContext(nvinfer1::DataType type, int32_t return context_workspace_size; } -size_t AttentionOp::getWorkspaceSizeForGeneration(nvinfer1::DataType type, int32_t max_num_seq, +size_t AttentionOp::getWorkspaceSizeForGeneration(tensorrt_llm::DataType type, int32_t max_num_seq, int32_t max_attention_window_size, int32_t max_num_tokens, int32_t max_blocks_per_sequence) const noexcept { if (max_num_tokens == 0) @@ -2818,7 +2819,7 @@ int AttentionOp::initialize() noexcept if (mEnableContextFMHA) { mEnableContextFMHA = false; - if (!(mType == nvinfer1::DataType::kHALF || mType == nvinfer1::DataType::kBF16)) + if (!(mType == tensorrt_llm::DataType::kHALF || mType == tensorrt_llm::DataType::kBF16)) { TLLM_LOG_WARNING("Fall back to unfused MHA because of unsupported data type."); } @@ -2863,7 +2864,7 @@ int AttentionOp::initialize() noexcept "mFP8ContextFMHA must enable if FP4 KV cache is enabled"); TLLM_CHECK(isRoPE() == (mRotaryEmbeddingDim != 0)); - TLLM_CHECK_WITH_INFO((mSM >= 80) || (mType != nvinfer1::DataType::kBF16), + TLLM_CHECK_WITH_INFO((mSM >= 80) || (mType != tensorrt_llm::DataType::kBF16), "Unsupported data type, pre SM 80 GPUs do not support bfloat16"); // Pre-check whether the head size is supported by MMHA. @@ -2915,11 +2916,11 @@ int AttentionOp::initialize() noexcept // Pre-checked during constructing. Data_type data_type, data_type_kv; - if (mType == nvinfer1::DataType::kHALF) + if (mType == tensorrt_llm::DataType::kHALF) { data_type = DATA_TYPE_FP16; } - else if (mType == nvinfer1::DataType::kBF16) + else if (mType == tensorrt_llm::DataType::kBF16) { data_type = DATA_TYPE_BF16; } @@ -3079,13 +3080,13 @@ int AttentionOp::initialize() noexcept Data_type kvDataType = DATA_TYPE_FP32; Data_type outputDataType = DATA_TYPE_FP32; - if (mType == nvinfer1::DataType::kHALF) + if (mType == tensorrt_llm::DataType::kHALF) { qDataType = DATA_TYPE_FP16; kvDataType = DATA_TYPE_FP16; outputDataType = DATA_TYPE_FP16; } - else if (mType == nvinfer1::DataType::kBF16) + else if (mType == tensorrt_llm::DataType::kBF16) { qDataType = DATA_TYPE_BF16; kvDataType = DATA_TYPE_BF16; @@ -3175,7 +3176,7 @@ int AttentionOp::initialize() noexcept } mEnableXQA = (mEnableXQA || mIsSpecDecodingEnabled) - && (mType == nvinfer1::DataType::kHALF || mType == nvinfer1::DataType::kBF16) && mUseKVCache; + && (mType == tensorrt_llm::DataType::kHALF || mType == tensorrt_llm::DataType::kBF16) && mUseKVCache; if (mEnableXQA) { @@ -3185,12 +3186,12 @@ int AttentionOp::initialize() noexcept fixedParams.isMLA = mIsGenerationMLA; // TODO: support more combinations. // Update Q and O dtype. - if (mType == nvinfer1::DataType::kHALF) + if (mType == tensorrt_llm::DataType::kHALF) { fixedParams.inputDataType = DATA_TYPE_FP16; fixedParams.outputDataType = DATA_TYPE_FP16; } - else if (mType == nvinfer1::DataType::kBF16) + else if (mType == tensorrt_llm::DataType::kBF16) { fixedParams.inputDataType = DATA_TYPE_BF16; fixedParams.outputDataType = DATA_TYPE_BF16; @@ -3258,10 +3259,6 @@ int AttentionOp::initialize() noexcept reserveSemaphoreArray(mNbMultiBlockSemaphores); } - if (isBuilding()) - { - return 0; - } #if ENABLE_MULTI_DEVICE if (mCpSize > 1 && COMM_SESSION.getSize() > 1) { diff --git a/cpp/tensorrt_llm/common/attentionOp.h b/cpp/tensorrt_llm/common/attentionOp.h index f7337c9c9cb2..eba32cb52de8 100644 --- a/cpp/tensorrt_llm/common/attentionOp.h +++ b/cpp/tensorrt_llm/common/attentionOp.h @@ -20,6 +20,7 @@ #include "tensorrt_llm/common/cublasMMWrapper.h" #include "tensorrt_llm/common/opUtils.h" #include "tensorrt_llm/common/quantization.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/kernels/contextFusedMultiHeadAttention/fused_multihead_attention_common.h" #include "tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_gemm.h" #include "tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderXQARunner.h" @@ -56,10 +57,11 @@ class AttentionOp [[nodiscard]] size_t getFmhaMultiCtasKvScratchSize() const noexcept; [[nodiscard]] int getHeadSize(bool checkInit = true) const; [[nodiscard]] int getMaxNumSeqLenTile(int batch_beam_size = 1) const; - [[nodiscard]] size_t getWorkspaceSizeForContext(nvinfer1::DataType type, int32_t nbReq, int32_t max_input_length, - int32_t cross_kv_length = 0, int32_t max_num_tokens = 0, int32_t total_kv_len = 0) const noexcept; + [[nodiscard]] size_t getWorkspaceSizeForContext(tensorrt_llm::DataType type, int32_t nbReq, + int32_t max_input_length, int32_t cross_kv_length = 0, int32_t max_num_tokens = 0, + int32_t total_kv_len = 0) const noexcept; // total_num_seq is the sum of beam_width for multiple requests - [[nodiscard]] size_t getWorkspaceSizeForGeneration(nvinfer1::DataType type, int32_t total_num_seq, + [[nodiscard]] size_t getWorkspaceSizeForGeneration(tensorrt_llm::DataType type, int32_t total_num_seq, int32_t max_attention_window_size, int32_t max_num_tokens, int32_t max_blocks_per_sequence) const noexcept; template @@ -181,14 +183,14 @@ class AttentionOp if (this->context_lengths && batch_size > 0) { ss << "context_lengths: " - << *(runtime::ITensor::wrap((void*) this->context_lengths, nvinfer1::DataType::kINT32, + << *(runtime::ITensor::wrap((void*) this->context_lengths, tensorrt_llm::DataType::kINT32, runtime::ITensor::makeShape({batch_size}))) << std::endl; } if (this->sequence_lengths && batch_size > 0) { ss << "sequence_lengths: " - << *(runtime::ITensor::wrap((void*) this->sequence_lengths, nvinfer1::DataType::kINT32, + << *(runtime::ITensor::wrap((void*) this->sequence_lengths, tensorrt_llm::DataType::kINT32, runtime::ITensor::makeShape({batch_size}))) << std::endl; } @@ -478,7 +480,7 @@ class AttentionOp int mTpSize = 1; int mTpRank = 0; bool mUnfuseQkvGemm = false; - nvinfer1::DataType mType; + tensorrt_llm::DataType mType; int32_t mMaxContextLength = 0; int32_t mMaxSeqLen = 0; int32_t mMaxNumRequests = 0; diff --git a/cpp/tensorrt_llm/common/opUtils.cpp b/cpp/tensorrt_llm/common/opUtils.cpp index ff9b57cdd099..560750c2ba93 100644 --- a/cpp/tensorrt_llm/common/opUtils.cpp +++ b/cpp/tensorrt_llm/common/opUtils.cpp @@ -17,6 +17,7 @@ #include "tensorrt_llm/common/opUtils.h" #include "tensorrt_llm/common/ncclUtils.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/runtime/ipcNvlsMemory.h" #include "tensorrt_llm/runtime/utils/mpiTags.h" #include "tensorrt_llm/runtime/utils/mpiUtils.h" @@ -33,18 +34,18 @@ TRTLLM_NAMESPACE_BEGIN #if ENABLE_MULTI_DEVICE -std::unordered_map* getDtypeMap() +std::unordered_map* getDtypeMap() { - static std::unordered_map dtypeMap = { - {nvinfer1::DataType::kFLOAT, ncclFloat32}, - {nvinfer1::DataType::kHALF, ncclFloat16}, - {nvinfer1::DataType::kBF16, ncclBfloat16}, - {nvinfer1::DataType::kFP8, ncclInt8}, - {nvinfer1::DataType::kBOOL, ncclInt8}, - {nvinfer1::DataType::kINT32, ncclInt32}, - {nvinfer1::DataType::kINT64, ncclInt64}, - {nvinfer1::DataType::kUINT8, ncclUint8}, - {nvinfer1::DataType::kINT8, ncclInt8}, + static std::unordered_map dtypeMap = { + {tensorrt_llm::DataType::kFLOAT, ncclFloat32}, + {tensorrt_llm::DataType::kHALF, ncclFloat16}, + {tensorrt_llm::DataType::kBF16, ncclBfloat16}, + {tensorrt_llm::DataType::kFP8, ncclInt8}, + {tensorrt_llm::DataType::kBOOL, ncclInt8}, + {tensorrt_llm::DataType::kINT32, ncclInt32}, + {tensorrt_llm::DataType::kINT64, ncclInt64}, + {tensorrt_llm::DataType::kUINT8, ncclUint8}, + {tensorrt_llm::DataType::kINT8, ncclInt8}, }; return &dtypeMap; } diff --git a/cpp/tensorrt_llm/common/opUtils.h b/cpp/tensorrt_llm/common/opUtils.h index 72e5a5ea3e09..22169843a9f5 100644 --- a/cpp/tensorrt_llm/common/opUtils.h +++ b/cpp/tensorrt_llm/common/opUtils.h @@ -21,7 +21,7 @@ #include "tensorrt_llm/common/cublasMMWrapper.h" #include "tensorrt_llm/common/workspace.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include #include #include @@ -62,14 +62,14 @@ void read(char const*& buffer, T& val) buffer += sizeof(T); } -inline cudaDataType_t trtToCublasDtype(nvinfer1::DataType type) +inline cudaDataType_t trtToCublasDtype(tensorrt_llm::DataType type) { switch (type) { - case nvinfer1::DataType::kFLOAT: return CUDA_R_32F; - case nvinfer1::DataType::kHALF: return CUDA_R_16F; + case tensorrt_llm::DataType::kFLOAT: return CUDA_R_32F; + case tensorrt_llm::DataType::kHALF: return CUDA_R_16F; #if defined(NV_TENSORRT_MAJOR) && NV_TENSORRT_MAJOR >= 9 - case nvinfer1::DataType::kBF16: return CUDA_R_16BF; + case tensorrt_llm::DataType::kBF16: return CUDA_R_16BF; #endif default: TLLM_THROW("Not supported data type for cuBLAS"); } @@ -185,13 +185,6 @@ struct hash void const* getCommSessionHandle(); } // namespace common::op -inline bool isBuilding() -{ - auto constexpr key = "IS_BUILDING"; - auto const val = getenv(key); - return val != nullptr && std::string(val) == "1"; -} - #if ENABLE_MULTI_DEVICE #define NCCLCHECK(cmd) \ do \ @@ -214,7 +207,7 @@ inline bool isBuilding() } \ } while (0) -std::unordered_map* getDtypeMap(); +std::unordered_map* getDtypeMap(); std::shared_ptr getComm(std::set const& group); diff --git a/cpp/tensorrt_llm/common/safetensors.cpp b/cpp/tensorrt_llm/common/safetensors.cpp index 9171f79e44e5..8bd91ccfbd51 100644 --- a/cpp/tensorrt_llm/common/safetensors.cpp +++ b/cpp/tensorrt_llm/common/safetensors.cpp @@ -18,7 +18,7 @@ #include "nlohmann/json.hpp" #include "tensorrt_llm/common/assert.h" #include "tensorrt_llm/common/config.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include #include #include @@ -30,7 +30,7 @@ TRTLLM_NAMESPACE_BEGIN namespace common::safetensors { -using nvinfer1::DataType; +using tensorrt_llm::DataType; static DataType convertDataTypeStrToEnum(std::string const& str) { diff --git a/cpp/tensorrt_llm/common/safetensors.h b/cpp/tensorrt_llm/common/safetensors.h index e31225f1be24..bdecf95e5909 100644 --- a/cpp/tensorrt_llm/common/safetensors.h +++ b/cpp/tensorrt_llm/common/safetensors.h @@ -18,7 +18,7 @@ #include "tensorrt_llm/common/assert.h" #include "tensorrt_llm/common/config.h" #include "tensorrt_llm/common/logger.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include #include #include @@ -34,13 +34,13 @@ class INdArray [[nodiscard]] virtual void const* data() const = 0; [[nodiscard]] virtual int ndim() const = 0; [[nodiscard]] virtual std::vector const& dims() const = 0; - [[nodiscard]] virtual nvinfer1::DataType dtype() const = 0; + [[nodiscard]] virtual tensorrt_llm::DataType dtype() const = 0; - [[nodiscard]] nvinfer1::Dims trtDims() const + [[nodiscard]] tensorrt_llm::Dims trtDims() const { - nvinfer1::Dims dims; + tensorrt_llm::Dims dims; dims.nbDims = ndim(); - TLLM_CHECK(dims.nbDims <= nvinfer1::Dims::MAX_DIMS); + TLLM_CHECK(dims.nbDims <= tensorrt_llm::Dims::MAX_DIMS); memset(dims.d, 0, sizeof(dims.d)); for (int i = 0; i < dims.nbDims; ++i) { diff --git a/cpp/tensorrt_llm/executor/CMakeLists.txt b/cpp/tensorrt_llm/executor/CMakeLists.txt index ca1ab298d7f4..358494e9989c 100644 --- a/cpp/tensorrt_llm/executor/CMakeLists.txt +++ b/cpp/tensorrt_llm/executor/CMakeLists.txt @@ -27,9 +27,7 @@ set(SRCS contextPhaseParams.cpp debugConfig.cpp decodingConfig.cpp - executor.cpp executorConfig.cpp - executorImpl.cpp executorKVCacheEventManager.cpp extendedRuntimePerfKnobConfig.cpp guidedDecodingConfig.cpp @@ -52,7 +50,6 @@ set(SRCS response.cpp samplingConfig.cpp dynamicBatchConfig.cpp - dynamicBatchTuner.cpp schedulerConfig.cpp serialization.cpp speculativeDecodingConfig.cpp @@ -60,7 +57,6 @@ set(SRCS types.cpp requestUtils.cpp contextPhaseParams.cpp - disaggServerUtil.cpp cacheTransceiverConfig.cpp) if(NOT WIN32) diff --git a/cpp/tensorrt_llm/executor/cache_transmission/cacheSplitConcat.cu b/cpp/tensorrt_llm/executor/cache_transmission/cacheSplitConcat.cu index 6e8c68d7efa7..d57ed7530e28 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/cacheSplitConcat.cu +++ b/cpp/tensorrt_llm/executor/cache_transmission/cacheSplitConcat.cu @@ -20,6 +20,7 @@ #include "tensorrt_llm/common/cudaFp8Utils.h" #include "tensorrt_llm/common/cudaUtils.h" #include "tensorrt_llm/common/reduceKernelUtils.cuh" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/executor/dataTransceiverState.h" #include "tensorrt_llm/executor/tensor.h" #include "tensorrt_llm/executor/types.h" @@ -27,7 +28,6 @@ #include "tensorrt_llm/runtime/iBuffer.h" #include "tensorrt_llm/runtime/iTensor.h" #include "tensorrt_llm/runtime/utils/mpiUtils.h" -#include #include #include #include @@ -467,7 +467,7 @@ void concatKVCache(runtime::ITensor::SharedPtr* inputBlocks, int inputBlockNum, blockInfos[outputBlockNum * inputAllRankNum + oi] = fillBlockInfo(oCacheState, outputBlocks[oi], oRank); } runtime::BufferManager::IBufferPtr blockInfosDeviceBuffer - = bufferManager.gpu(sizeof(BlockInfo) * (blockInfos.size()), nvinfer1::DataType::kUINT8); + = bufferManager.gpu(sizeof(BlockInfo) * (blockInfos.size()), tensorrt_llm::DataType::kUINT8); bufferManager.copy((blockInfos.data()), *blockInfosDeviceBuffer, runtime::MemoryType::kCPU); BlockInfo* iBlockInfoDevice = static_cast*>(blockInfosDeviceBuffer->data()); @@ -594,7 +594,7 @@ void concatKVCacheDispatch(runtime::ITensor::SharedPtr* inputBlocks, int inputBl } } -nvinfer1::Dims makeShapeFromCacheState(kv_cache::CacheState const& cacheState) +tensorrt_llm::Dims makeShapeFromCacheState(kv_cache::CacheState const& cacheState) { int64_t blockSize = static_cast(cacheState.getModelConfig().mNbKvHeadsPerLayer[0] @@ -1130,8 +1130,8 @@ void splitKVCache(std::map> std::vector layersInWindow; size_t cacheBlockSizeSum = 0; size_t inputBlockLayerNumSum = 0; - auto cacheDataType - = isIndexerKCache ? nvinfer1::DataType::kUINT8 : kVCacheBlocksPerWindow.begin()->second.front()->getDataType(); + auto cacheDataType = isIndexerKCache ? tensorrt_llm::DataType::kUINT8 + : kVCacheBlocksPerWindow.begin()->second.front()->getDataType(); for (auto const& [window, blocks] : kVCacheBlocksPerWindow) { @@ -1170,7 +1170,7 @@ void splitKVCache(std::map> bool const isWindow = windowSizes.size() > 1; runtime::BufferManager::IBufferPtr PtrsDeviceBuffer - = bufferManager.gpu(cachePtrs.size(), nvinfer1::DataType::kINT64); + = bufferManager.gpu(cachePtrs.size(), tensorrt_llm::DataType::kINT64); TLLM_CHECK(PtrsDeviceBuffer->getSizeInBytes() == cachePtrs.size() * sizeof(T*)); bufferManager.copy(cachePtrs.data(), *PtrsDeviceBuffer, runtime::MemoryType::kCPU); @@ -1182,7 +1182,7 @@ void splitKVCache(std::map> windowInfoHostBuffer.insert(windowInfoHostBuffer.end(), blockNumInwindow.begin(), blockNumInwindow.end()); windowInfoHostBuffer.insert(windowInfoHostBuffer.end(), layersInWindow.begin(), layersInWindow.end()); - windowInfoDeviceBuffer = bufferManager.gpu(windowInfoHostBuffer.size(), nvinfer1::DataType::kINT32); + windowInfoDeviceBuffer = bufferManager.gpu(windowInfoHostBuffer.size(), tensorrt_llm::DataType::kINT32); bufferManager.copy(windowInfoHostBuffer.data(), *windowInfoDeviceBuffer, runtime::MemoryType::kCPU); for (auto layerNum : layersInWindow) @@ -1404,8 +1404,8 @@ void splitKVCacheDispatch(std::mapsecond.front()->getDataType(); + auto dataType = isIndexerKCache ? tensorrt_llm::DataType::kUINT8 + : kVCacheBlocksPerWindow.begin()->second.front()->getDataType(); auto dataSize = tensorrt_llm::common::getDTypeSize(dataType); switch (dataSize) @@ -1513,7 +1513,7 @@ void concatKVCache(std::vector const& inputSplitBlo } cachePtrs.insert(cachePtrs.end(), prefixLayerNum.begin(), prefixLayerNum.end()); runtime::BufferManager::IBufferPtr PtrsDeviceBuffer - = bufferManager.gpu(cachePtrs.size(), nvinfer1::DataType::kINT64); + = bufferManager.gpu(cachePtrs.size(), tensorrt_llm::DataType::kINT64); TLLM_CHECK(PtrsDeviceBuffer->getSizeInBytes() == cachePtrs.size() * sizeof(uint64_t)); bufferManager.copy(cachePtrs.data(), *PtrsDeviceBuffer, runtime::MemoryType::kCPU); bool const isWindow = windowSizes.size() > 1; @@ -1525,7 +1525,7 @@ void concatKVCache(std::vector const& inputSplitBlo windowInfoHostBuffer.insert(windowInfoHostBuffer.end(), blockNumInwindow.begin(), blockNumInwindow.end()); windowInfoHostBuffer.insert(windowInfoHostBuffer.end(), layersInWindow.begin(), layersInWindow.end()); - windowInfoDeviceBuffer = bufferManager.gpu(windowInfoHostBuffer.size(), nvinfer1::DataType::kINT32); + windowInfoDeviceBuffer = bufferManager.gpu(windowInfoHostBuffer.size(), tensorrt_llm::DataType::kINT32); bufferManager.copy(windowInfoHostBuffer.data(), *windowInfoDeviceBuffer, runtime::MemoryType::kCPU); } constexpr int subWarpSize = 8; diff --git a/cpp/tensorrt_llm/executor/cache_transmission/cacheSplitConcat.h b/cpp/tensorrt_llm/executor/cache_transmission/cacheSplitConcat.h index c7036b219612..80816bc5c3a0 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/cacheSplitConcat.h +++ b/cpp/tensorrt_llm/executor/cache_transmission/cacheSplitConcat.h @@ -27,7 +27,7 @@ #include "tensorrt_llm/runtime/bufferManager.h" #include "tensorrt_llm/runtime/iTensor.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" namespace tensorrt_llm::executor::kv_cache { @@ -78,7 +78,7 @@ void concatKVCacheDispatch(runtime::ITensor::SharedPtr* inputBlocks, int inputBl runtime::ITensor::SharedPtr* outputBlocks, int outputBlockNum, int selfRank, kv_cache::CacheState const& selfCacheState, runtime::BufferManager const& bufferManager); -nvinfer1::Dims makeShapeFromCacheState(kv_cache::CacheState const& cacheState); +tensorrt_llm::Dims makeShapeFromCacheState(kv_cache::CacheState const& cacheState); void splitKVCacheDispatch(std::map> const& kVCacheBlocksPerWindow, std::vector& ouputSplitBlocks, kv_cache::CacheState const& peerCacheState, @@ -147,7 +147,7 @@ void concatRnnSsmStateDispatch(std::vector const& i void splitUnifiedPoolSsmDispatch(runtime::ITensor::SharedPtr const& pool, std::vector const& realBlockIndices, std::vector& outputSplitBlocks, kv_cache::CacheState const& destCacheState, kv_cache::CacheState const& selfCacheState, int selfIdx, - size_t ssmBytes, size_t blockSizeBytes, nvinfer1::DataType ssmDataType, + size_t ssmBytes, size_t blockSizeBytes, tensorrt_llm::DataType ssmDataType, runtime::BufferManager const& bufferManager); /** @@ -156,7 +156,7 @@ void splitUnifiedPoolSsmDispatch(runtime::ITensor::SharedPtr const& pool, void splitUnifiedPoolConvDispatch(runtime::ITensor::SharedPtr const& pool, std::vector const& realBlockIndices, std::vector& outputSplitBlocks, kv_cache::CacheState const& destCacheState, kv_cache::CacheState const& selfCacheState, int selfIdx, - size_t ssmBytes, size_t blockSizeBytes, nvinfer1::DataType convDataType, + size_t ssmBytes, size_t blockSizeBytes, tensorrt_llm::DataType convDataType, runtime::BufferManager const& bufferManager); /** @@ -165,7 +165,7 @@ void splitUnifiedPoolConvDispatch(runtime::ITensor::SharedPtr const& pool, void concatUnifiedPoolSsmDispatch(runtime::ITensor::SharedPtr const& pool, std::vector const& realBlockIndices, std::vector const& inputSplitBlocks, kv_cache::CacheState const& srcCacheState, kv_cache::CacheState const& selfCacheState, int selfIdx, size_t ssmBytes, - size_t blockSizeBytes, nvinfer1::DataType ssmDataType, runtime::BufferManager const& bufferManager); + size_t blockSizeBytes, tensorrt_llm::DataType ssmDataType, runtime::BufferManager const& bufferManager); /** * @brief Concat conv state from per-source buffers into unified pool blocks (section-aware). @@ -173,6 +173,6 @@ void concatUnifiedPoolSsmDispatch(runtime::ITensor::SharedPtr const& pool, void concatUnifiedPoolConvDispatch(runtime::ITensor::SharedPtr const& pool, std::vector const& realBlockIndices, std::vector const& inputSplitBlocks, kv_cache::CacheState const& srcCacheState, kv_cache::CacheState const& selfCacheState, int selfIdx, size_t ssmBytes, - size_t blockSizeBytes, nvinfer1::DataType convDataType, runtime::BufferManager const& bufferManager); + size_t blockSizeBytes, tensorrt_llm::DataType convDataType, runtime::BufferManager const& bufferManager); } // namespace tensorrt_llm::executor::rnn_cache diff --git a/cpp/tensorrt_llm/executor/cache_transmission/rnnCacheSplitConcat.cu b/cpp/tensorrt_llm/executor/cache_transmission/rnnCacheSplitConcat.cu index fb6837e8d188..f464c41aadc9 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/rnnCacheSplitConcat.cu +++ b/cpp/tensorrt_llm/executor/cache_transmission/rnnCacheSplitConcat.cu @@ -24,6 +24,7 @@ #include "tensorrt_llm/common/logger.h" #include "tensorrt_llm/common/memoryUtils.h" #include "tensorrt_llm/common/reduceKernelUtils.cuh" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/executor/dataTransceiverState.h" #include "tensorrt_llm/executor/tensor.h" #include "tensorrt_llm/executor/types.h" @@ -31,7 +32,6 @@ #include "tensorrt_llm/runtime/iBuffer.h" #include "tensorrt_llm/runtime/iTensor.h" #include "tensorrt_llm/runtime/utils/mpiUtils.h" -#include #include #include #include @@ -366,7 +366,7 @@ void splitRnnConvState(std::vector const& inputConv // Allocate and copy pointer array to device runtime::BufferManager::IBufferPtr PtrsDeviceBuffer - = bufferManager.gpu(cachePtrs.size(), nvinfer1::DataType::kINT64); + = bufferManager.gpu(cachePtrs.size(), tensorrt_llm::DataType::kINT64); TLLM_CHECK(PtrsDeviceBuffer->getSizeInBytes() == cachePtrs.size() * sizeof(uint64_t)); bufferManager.copy(cachePtrs.data(), *PtrsDeviceBuffer, runtime::MemoryType::kCPU); @@ -497,7 +497,7 @@ void splitRnnSsmState(std::vector const& inputSsmBl cachePtrs.insert(cachePtrs.end(), prefixLayerNum.begin(), prefixLayerNum.end()); runtime::BufferManager::IBufferPtr PtrsDeviceBuffer - = bufferManager.gpu(cachePtrs.size(), nvinfer1::DataType::kINT64); + = bufferManager.gpu(cachePtrs.size(), tensorrt_llm::DataType::kINT64); TLLM_CHECK(PtrsDeviceBuffer->getSizeInBytes() == cachePtrs.size() * sizeof(uint64_t)); bufferManager.copy(cachePtrs.data(), *PtrsDeviceBuffer, runtime::MemoryType::kCPU); @@ -625,7 +625,7 @@ void concatRnnConvState(std::vector const& inputSpl cachePtrs.insert(cachePtrs.end(), prefixLayerNum.begin(), prefixLayerNum.end()); runtime::BufferManager::IBufferPtr PtrsDeviceBuffer - = bufferManager.gpu(cachePtrs.size(), nvinfer1::DataType::kINT64); + = bufferManager.gpu(cachePtrs.size(), tensorrt_llm::DataType::kINT64); TLLM_CHECK(PtrsDeviceBuffer->getSizeInBytes() == cachePtrs.size() * sizeof(uint64_t)); bufferManager.copy(cachePtrs.data(), *PtrsDeviceBuffer, runtime::MemoryType::kCPU); @@ -747,7 +747,7 @@ void concatRnnSsmState(std::vector const& inputSpli cachePtrs.insert(cachePtrs.end(), prefixLayerNum.begin(), prefixLayerNum.end()); runtime::BufferManager::IBufferPtr PtrsDeviceBuffer - = bufferManager.gpu(cachePtrs.size(), nvinfer1::DataType::kINT64); + = bufferManager.gpu(cachePtrs.size(), tensorrt_llm::DataType::kINT64); TLLM_CHECK(PtrsDeviceBuffer->getSizeInBytes() == cachePtrs.size() * sizeof(uint64_t)); bufferManager.copy(cachePtrs.data(), *PtrsDeviceBuffer, runtime::MemoryType::kCPU); @@ -1348,7 +1348,7 @@ void splitUnifiedPoolSsm(runtime::ITensor::SharedPtr const& pool, std::vector(PtrsDeviceBuffer->data()); @@ -1481,10 +1481,10 @@ void splitUnifiedPoolConv(runtime::ITensor::SharedPtr const& pool, std::vector(PtrsDeviceBuffer->data()); @@ -1606,7 +1606,7 @@ void concatUnifiedPoolSsm(runtime::ITensor::SharedPtr const& pool, std::vector(PtrsDeviceBuffer->data()); @@ -1731,10 +1731,10 @@ void concatUnifiedPoolConv(runtime::ITensor::SharedPtr const& pool, std::vector< sectionInfo.insert(sectionInfo.end(), sectionDimsDomainTP.begin(), sectionDimsDomainTP.end()); sectionInfo.insert(sectionInfo.end(), sectionOffsetsLocal.begin(), sectionOffsetsLocal.end()); - auto PtrsDeviceBuffer = bufferManager.gpu(allPtrs.size(), nvinfer1::DataType::kINT64); + auto PtrsDeviceBuffer = bufferManager.gpu(allPtrs.size(), tensorrt_llm::DataType::kINT64); bufferManager.copy(allPtrs.data(), *PtrsDeviceBuffer, runtime::MemoryType::kCPU); - auto sectionInfoBuffer = bufferManager.gpu(sectionInfo.size(), nvinfer1::DataType::kINT32); + auto sectionInfoBuffer = bufferManager.gpu(sectionInfo.size(), tensorrt_llm::DataType::kINT32); bufferManager.copy(sectionInfo.data(), *sectionInfoBuffer, runtime::MemoryType::kCPU); T** outputPtrsDev = static_cast(PtrsDeviceBuffer->data()); @@ -1810,7 +1810,8 @@ void concatUnifiedPoolConv(runtime::ITensor::SharedPtr const& pool, std::vector< void splitUnifiedPoolSsmDispatch(runtime::ITensor::SharedPtr const& pool, std::vector const& realBlockIndices, std::vector& outputSplitBlocks, kv_cache::CacheState const& destCacheState, kv_cache::CacheState const& selfCacheState, int selfIdx, - size_t ssmBytes, size_t blockSizeBytes, nvinfer1::DataType ssmDataType, runtime::BufferManager const& bufferManager) + size_t ssmBytes, size_t blockSizeBytes, tensorrt_llm::DataType ssmDataType, + runtime::BufferManager const& bufferManager) { auto dataSize = tensorrt_llm::common::getDTypeSize(ssmDataType); switch (dataSize) @@ -1834,7 +1835,7 @@ void splitUnifiedPoolSsmDispatch(runtime::ITensor::SharedPtr const& pool, void splitUnifiedPoolConvDispatch(runtime::ITensor::SharedPtr const& pool, std::vector const& realBlockIndices, std::vector& outputSplitBlocks, kv_cache::CacheState const& destCacheState, kv_cache::CacheState const& selfCacheState, int selfIdx, - size_t ssmBytes, size_t blockSizeBytes, nvinfer1::DataType convDataType, + size_t ssmBytes, size_t blockSizeBytes, tensorrt_llm::DataType convDataType, runtime::BufferManager const& bufferManager) { auto dataSize = tensorrt_llm::common::getDTypeSize(convDataType); @@ -1859,7 +1860,7 @@ void splitUnifiedPoolConvDispatch(runtime::ITensor::SharedPtr const& pool, void concatUnifiedPoolSsmDispatch(runtime::ITensor::SharedPtr const& pool, std::vector const& realBlockIndices, std::vector const& inputSplitBlocks, kv_cache::CacheState const& srcCacheState, kv_cache::CacheState const& selfCacheState, int selfIdx, size_t ssmBytes, - size_t blockSizeBytes, nvinfer1::DataType ssmDataType, runtime::BufferManager const& bufferManager) + size_t blockSizeBytes, tensorrt_llm::DataType ssmDataType, runtime::BufferManager const& bufferManager) { auto dataSize = tensorrt_llm::common::getDTypeSize(ssmDataType); switch (dataSize) @@ -1883,7 +1884,7 @@ void concatUnifiedPoolSsmDispatch(runtime::ITensor::SharedPtr const& pool, void concatUnifiedPoolConvDispatch(runtime::ITensor::SharedPtr const& pool, std::vector const& realBlockIndices, std::vector const& inputSplitBlocks, kv_cache::CacheState const& srcCacheState, kv_cache::CacheState const& selfCacheState, int selfIdx, size_t ssmBytes, - size_t blockSizeBytes, nvinfer1::DataType convDataType, runtime::BufferManager const& bufferManager) + size_t blockSizeBytes, tensorrt_llm::DataType convDataType, runtime::BufferManager const& bufferManager) { auto dataSize = tensorrt_llm::common::getDTypeSize(convDataType); switch (dataSize) diff --git a/cpp/tensorrt_llm/executor/disaggServerUtil.cpp b/cpp/tensorrt_llm/executor/disaggServerUtil.cpp deleted file mode 100644 index 6be2e4fb8ae8..000000000000 --- a/cpp/tensorrt_llm/executor/disaggServerUtil.cpp +++ /dev/null @@ -1,555 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2025 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. - */ - -#include "tensorrt_llm/executor/disaggServerUtil.h" -#include "tensorrt_llm/common/utils.h" -#include "tensorrt_llm/executor/types.h" -#include "tensorrt_llm/runtime/utils/mpiUtils.h" - -#include - -namespace tensorrt_llm::executor::disagg_executor -{ - -class DisaggExecutorOrchestrator::Impl -{ -public: - Impl(std::vector const& ctxEnginePaths, - std::vector const& genEnginePaths, - std::vector const& ctxExecutorConfigs, - std::vector const& genExecutorConfigs, bool hasContextAwaitThreads, - bool hasGenAwaitThreads) - : mhasContextAwaitThreads(hasContextAwaitThreads) - , mhasGenAwaitThreads(hasGenAwaitThreads) - { - TLLM_CHECK(ctxEnginePaths.size() == ctxExecutorConfigs.size()); - TLLM_CHECK(genEnginePaths.size() == genExecutorConfigs.size()); - TLLM_CHECK(!(ctxEnginePaths.empty() || genEnginePaths.empty())); - int worldRank = tensorrt_llm::mpi::MpiComm::world().getRank(); - mIsOrchestrator = (worldRank == 0); - auto contextNum = ctxEnginePaths.size(); - mContextReqIdToGlobalId = std::vector>(contextNum); - mContextMapMutexs = std::vector(contextNum); - auto genNum = genEnginePaths.size(); - mGenerationReqIdToGlobalId = std::vector>(genNum); - mGenerationMapMutexs = std::vector(genNum); - - for (size_t cN = 0; cN < contextNum; cN++) - { - mContextExecutors.push_back(std::make_unique( - ctxEnginePaths[cN], texec::ModelType::kDECODER_ONLY, ctxExecutorConfigs[cN])); - } - - for (size_t gN = 0; gN < genNum; gN++) - { - mGenerationExecutors.push_back(std::make_unique( - genEnginePaths[gN], texec::ModelType::kDECODER_ONLY, genExecutorConfigs[gN])); - } - - if (mIsOrchestrator) - { - if (mhasContextAwaitThreads) - { - for (size_t contextIdx = 0; contextIdx < contextNum; contextIdx++) - { - mContextThreads.emplace_back( - [this, contextIdx]() { this->waitResponseAndAppendThreadFun(true, contextIdx); }); - } - } - if (mhasGenAwaitThreads) - { - - for (size_t genIdx = 0; genIdx < genNum; genIdx++) - { - mGenerationThreads.emplace_back( - [this, genIdx]() { this->waitResponseAndAppendThreadFun(false, genIdx); }); - } - } - } - tensorrt_llm::mpi::MpiComm::world().barrier(); - } - - std::vector enqueueContext(std::vector const& requests, - std::optional selectContextId = std::nullopt, bool batch = false) - { - - std::vector globalReqIds; - for (auto const& request : requests) - { - globalReqIds.push_back(generatedGlobalId()); - TLLM_CHECK(request.getRequestType() == tensorrt_llm::executor::RequestType::REQUEST_TYPE_CONTEXT_ONLY); - } - - if (batch) - { - size_t contextId = selectContextId.has_value() ? selectContextId.value() : selectContextExecutor(); - auto contextReqIds = mContextExecutors[contextId]->enqueueRequests(requests); - { - std::scoped_lock lock{mContextMapMutexs[contextId]}; - for (size_t i = 0; i < requests.size(); ++i) - { - mContextReqIdToGlobalId[contextId][contextReqIds[i]] = globalReqIds[i]; - } - } - } - else - { - for (size_t i = 0; i < requests.size(); ++i) - { - size_t contextId = selectContextId.has_value() ? selectContextId.value() : selectContextExecutor(); - - auto contextReqId = mContextExecutors[contextId]->enqueueRequest(requests[i]); - { - std::scoped_lock lock{mContextMapMutexs[contextId]}; - mContextReqIdToGlobalId[contextId][contextReqId] = globalReqIds[i]; - } - } - } - return globalReqIds; - } - - void enqueueGeneration(std::vector const& requests, std::vector const& globalRequestIds, - std::optional selectGenIdx = std::nullopt, bool batch = false) - { - - TLLM_CHECK(globalRequestIds.size() == requests.size()); - - for (auto const& request : requests) - { - - TLLM_CHECK(request.getRequestType() == tensorrt_llm::executor::RequestType::REQUEST_TYPE_GENERATION_ONLY); - } - if (batch) - { - size_t genIdx = selectGenIdx.has_value() ? selectGenIdx.value() : selectGenerationExecutor(); - auto genReqIds = mGenerationExecutors[genIdx]->enqueueRequests(requests); - { - std::scoped_lock lock{mGenerationMapMutexs[genIdx]}; - for (size_t i = 0; i < requests.size(); ++i) - { - mGenerationReqIdToGlobalId[genIdx][genReqIds[i]] = globalRequestIds[i]; - } - } - } - else - { - for (size_t i = 0; i < requests.size(); ++i) - { - size_t genIdx = selectGenIdx.has_value() ? selectGenIdx.value() : selectGenerationExecutor(); - - auto genReqId = mGenerationExecutors[genIdx]->enqueueRequest(requests[i]); - { - std::scoped_lock lock{mGenerationMapMutexs[genIdx]}; - mGenerationReqIdToGlobalId[genIdx][genReqId] = globalRequestIds[i]; - } - } - } - } - - std::vector awaitContextResponses( - std::optional contextIdx, std::optional const& timeout) - { - - std::vector responses; - - if (mhasContextAwaitThreads) - { - - std::unique_lock lock(mResponsesContextMtx); - auto pred = [&mShutdown = mShutdown, &resp = this->mContextResponses]() -> bool - { return !resp.empty() || mShutdown; }; - auto storeResponses = [&resp = this->mContextResponses, &responses]() - { - responses = std::move(resp); - resp.clear(); - }; - if (timeout) - { - if (mContextResponsesCV.wait_for(lock, timeout.value(), pred)) - { - storeResponses(); - } - } - else - { - mContextResponsesCV.wait(lock, pred); - storeResponses(); - } - TLLM_CHECK_WITH_INFO( - !contextIdx.has_value(), "contextIdx should not be provided when mhasContextAwaitThreads is true"); - - return responses; - } - - if (contextIdx.has_value()) - { - TLLM_CHECK(!mhasContextAwaitThreads); - auto responseFromExecutor = mContextExecutors[contextIdx.value()]->awaitResponses(timeout); - for (auto&& resp : responseFromExecutor) - { - - auto reqId = resp.getRequestId(); - IdType globalId{0}; - - { - std::scoped_lock lock{mContextMapMutexs.at(contextIdx.value())}; - globalId = mContextReqIdToGlobalId.at(contextIdx.value()).at(reqId); - } - TLLM_CHECK(globalId != 0); - responses.emplace_back(std::move(resp), globalId); - } - return responses; - } - TLLM_CHECK(timeout.has_value()); - auto timeouP = timeout.value() / mContextExecutors.size(); - for (size_t ci = 0; ci < mContextExecutors.size(); ci++) - { - auto responseFromExecutor = mContextExecutors.at(ci)->awaitResponses(timeouP); - for (auto&& resp : responseFromExecutor) - { - auto reqId = resp.getRequestId(); - IdType globalId{0}; - - { - std::scoped_lock lock{mContextMapMutexs.at(ci)}; - globalId = mContextReqIdToGlobalId.at(ci).at(reqId); - } - TLLM_CHECK(globalId != 0); - responses.emplace_back(std::move(resp), globalId); - } - } - - return responses; - }; - - std::vector awaitGenerationResponses( - std::optional genIdx, std::optional const& timeout) - { - - std::vector responses; - - if (mhasGenAwaitThreads) - { - - std::unique_lock lock(mResponseGenerationMtx); - auto pred = [&mShutdown = mShutdown, &resp = this->mGenerationResponses]() -> bool - { return !resp.empty() || mShutdown; }; - auto storeResponses = [&resp = this->mGenerationResponses, &responses]() - { - responses = std::move(resp); - resp.clear(); - }; - if (timeout) - { - if (mGenerationResponsesCv.wait_for(lock, timeout.value(), pred)) - { - storeResponses(); - } - } - else - { - mGenerationResponsesCv.wait(lock, pred); - storeResponses(); - } - TLLM_CHECK_WITH_INFO(!genIdx.has_value(), "genIdx should not be provided when mhasGenAwaitThreads is true"); - return responses; - } - - if (genIdx.has_value()) - { - TLLM_CHECK(!mhasGenAwaitThreads); - auto responseFromExecutor = mGenerationExecutors[genIdx.value()]->awaitResponses(timeout); - for (auto&& resp : responseFromExecutor) - { - - auto reqId = resp.getRequestId(); - IdType globalId{0}; - - { - std::scoped_lock lock{mGenerationMapMutexs.at(genIdx.value())}; - globalId = mGenerationReqIdToGlobalId.at(genIdx.value()).at(reqId); - } - TLLM_CHECK(globalId != 0); - responses.emplace_back(std::move(resp), globalId); - } - return responses; - } - TLLM_CHECK(timeout.has_value()); - auto timeouP = timeout.value() / mGenerationExecutors.size(); - - for (size_t gi = 0; gi < mGenerationExecutors.size(); gi++) - { - auto responseFromExecutor = mGenerationExecutors.at(gi)->awaitResponses(timeouP); - for (auto&& resp : responseFromExecutor) - { - - auto reqId = resp.getRequestId(); - IdType globalId{0}; - - { - std::scoped_lock lock{mGenerationMapMutexs.at(gi)}; - globalId = mGenerationReqIdToGlobalId.at(gi).at(reqId); - } - TLLM_CHECK(globalId != 0); - responses.emplace_back(std::move(resp), globalId); - } - } - - return responses; - }; - - [[nodiscard]] bool canEnqueue() const - { - return mIsOrchestrator; - } - - [[nodiscard]] std::vector> const& getContextExecutors() const - { - return mContextExecutors; - } - - [[nodiscard]] std::vector> const& getGenExecutors() const - { - return mGenerationExecutors; - } - - ~Impl() - { - - mShutdown = true; - - mContextResponsesCV.notify_all(); - mGenerationResponsesCv.notify_all(); - for (auto&& executor : mContextExecutors) - { - executor->shutdown(); - } - for (auto&& executor : mGenerationExecutors) - { - executor->shutdown(); - } - - if (mIsOrchestrator) - { - if (mhasContextAwaitThreads) - { - for (auto&& contextThread : mContextThreads) - { - if (contextThread.joinable()) - { - contextThread.join(); - } - } - } - if (mhasGenAwaitThreads) - { - for (auto&& genThread : mGenerationThreads) - { - if (genThread.joinable()) - { - genThread.join(); - } - } - } - } - } - -private: - IdType generatedGlobalId() - { - return (++mLastId % UINT64_MAX); - }; - - size_t selectContextExecutor() - { - static size_t selectContextId = 0; - auto contextId = (selectContextId++) % mContextExecutors.size(); - if (selectContextId >= mContextExecutors.size()) - { - selectContextId = 0; - } - return contextId; - } - - size_t selectGenerationExecutor() - { - static size_t selectGenerationId = 0; - auto generationIdx = (selectGenerationId++) % mGenerationExecutors.size(); - if (selectGenerationId >= mGenerationExecutors.size()) - { - selectGenerationId = 0; - } - return generationIdx; - } - - void appendNewContextResponse(std::vector&& newResponses) - { - { - std::scoped_lock lock(mResponsesContextMtx); - for (auto&& response : newResponses) - { - mContextResponses.emplace_back(std::move(response)); - } - } - mContextResponsesCV.notify_all(); - } - - void appendNewGenerationResponse(std::vector&& newResponses) - { - { - std::scoped_lock lock(mResponseGenerationMtx); - for (auto&& response : newResponses) - { - mGenerationResponses.emplace_back(std::move(response)); - } - } - mGenerationResponsesCv.notify_all(); - } - - void waitResponseAndAppendThreadFun(bool isContext, int executorIdx) - { - - tensorrt_llm::common::setThreadName("waitResponseAndAppendThreadFun"); - - auto& executor = isContext ? mContextExecutors[executorIdx] : mGenerationExecutors[executorIdx]; - - while (!mShutdown) - { - auto responses = executor->awaitResponses(); - - if (responses.empty()) - { - continue; - } - std::vector responseWithIds; - if (isContext) - { - for (auto&& response : responses) - { - auto reqId = response.getRequestId(); - IdType globalId{0}; - - { - std::scoped_lock lock{mContextMapMutexs.at(executorIdx)}; - globalId = mContextReqIdToGlobalId.at(executorIdx).at(reqId); - } - TLLM_CHECK(globalId != 0); - responseWithIds.emplace_back(std::move(response), globalId); - } - if (responseWithIds.size() > 0) - { - appendNewContextResponse(std::move(responseWithIds)); - } - } - else - { - - for (auto&& response : responses) - { - auto reqId = response.getRequestId(); - IdType globalId{0}; - - { - std::scoped_lock lock{mGenerationMapMutexs.at(executorIdx)}; - globalId = mGenerationReqIdToGlobalId.at(executorIdx).at(reqId); - } - TLLM_CHECK(globalId != 0); - responseWithIds.emplace_back(std::move(response), globalId); - } - if (responseWithIds.size() > 0) - { - appendNewGenerationResponse(std::move(responseWithIds)); - } - } - } - }; - - std::vector> mContextExecutors; - std::vector> mGenerationExecutors; - std::vector mContextThreads; - std::vector mGenerationThreads; - - std::atomic mLastId{0}; - std::vector> mContextReqIdToGlobalId; - std::vector> mGenerationReqIdToGlobalId; - std::vector mContextMapMutexs; - std::vector mGenerationMapMutexs; - std::vector mContextResponses; - std::condition_variable mContextResponsesCV; - std::mutex mResponsesContextMtx; - - std::vector mGenerationResponses; - std::condition_variable mGenerationResponsesCv; - std::mutex mResponseGenerationMtx; - std::atomic mShutdown{false}; - std::atomic mhasContextAwaitThreads{false}; - std::atomic mhasGenAwaitThreads{false}; - bool mIsOrchestrator{false}; -}; - -DisaggExecutorOrchestrator::DisaggExecutorOrchestrator(std::vector const& ctxEnginePaths, - std::vector const& genEnginePaths, - std::vector const& ctxExecutorConfigs, - std::vector const& genExecutorConfigs, bool hasContextAwaitThreads, - bool hasGenAwaitThreads) - : mImpl(std::make_unique(ctxEnginePaths, genEnginePaths, ctxExecutorConfigs, - genExecutorConfigs, hasContextAwaitThreads, hasGenAwaitThreads)) -{ -} - -std::vector DisaggExecutorOrchestrator::enqueueContext( - std::vector const& requests, std::optional selectContextId, bool batch) -{ - return mImpl->enqueueContext(requests, selectContextId, batch); -} - -void DisaggExecutorOrchestrator::enqueueGeneration(std::vector const& requests, - std::vector const& globalRequestIds, std::optional selectGenIdx, bool batch) -{ - mImpl->enqueueGeneration(requests, globalRequestIds, selectGenIdx, batch); -} - -std::vector DisaggExecutorOrchestrator::awaitContextResponses( - std::optional const& timeout, std::optional contextIdx) -{ - return mImpl->awaitContextResponses(contextIdx, timeout); -} - -std::vector DisaggExecutorOrchestrator::awaitGenerationResponses( - std::optional const& timeout, std::optional genIdx) -{ - return mImpl->awaitGenerationResponses(genIdx, timeout); -} - -bool DisaggExecutorOrchestrator::canEnqueue() const -{ - return mImpl->canEnqueue(); -}; - -std::vector> const& DisaggExecutorOrchestrator::getContextExecutors() const -{ - return mImpl->getContextExecutors(); -} - -std::vector> const& DisaggExecutorOrchestrator::getGenExecutors() const -{ - return mImpl->getGenExecutors(); -} - -DisaggExecutorOrchestrator::~DisaggExecutorOrchestrator() = default; - -} // namespace tensorrt_llm::executor::disagg_executor diff --git a/cpp/tensorrt_llm/executor/dynamicBatchTuner.cpp b/cpp/tensorrt_llm/executor/dynamicBatchTuner.cpp deleted file mode 100644 index b7cd49f5430b..000000000000 --- a/cpp/tensorrt_llm/executor/dynamicBatchTuner.cpp +++ /dev/null @@ -1,113 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2025 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. - */ - -#include "tensorrt_llm/executor/dynamicBatchTuner.h" -#include "tensorrt_llm/common/logger.h" -#include - -namespace -{ -using namespace tensorrt_llm::executor; - -void updateStats(SizeType32 value, std::deque& stats, int64_t& sum, SizeType32 windowSize) -{ - while (static_cast(stats.size()) >= windowSize) - { - sum -= stats.front(); - stats.pop_front(); - } - stats.push_back(value); - sum += value; -} -} // namespace - -namespace tensorrt_llm::executor -{ - -DynamicBatchTuner::DynamicBatchTuner(DynamicBatchConfig const& config) - : mEnableBatchSizeTuning(config.getEnableBatchSizeTuning()) - , mEnableMaxNumTokensTuning(config.getEnableMaxNumTokensTuning()) - , mDynamicBatchMovingAverageWindow(config.getDynamicBatchMovingAverageWindow()) - , mBatchSizeTable(config.getBatchSizeTable()) -{ - TLLM_CHECK_WITH_INFO(!mBatchSizeTable.empty(), "Batch size table is empty."); - for (size_t i = 1; i < mBatchSizeTable.size(); ++i) - { - TLLM_CHECK_WITH_INFO(mBatchSizeTable[i - 1].first < mBatchSizeTable[i].first, - "Batch size table is not sorted in ascending order."); - } -} - -void DynamicBatchTuner::updateStats(SizeType32 inputLength, SizeType32 outputLength) -{ - ::updateStats(inputLength, mInputLengthStats, mInputLengthSum, mDynamicBatchMovingAverageWindow); - ::updateStats(outputLength, mOutputLengthStats, mOutputLengthSum, mDynamicBatchMovingAverageWindow); -} - -double DynamicBatchTuner::getAverageInputLength() const -{ - return mInputLengthStats.empty() ? 0 : static_cast(mInputLengthSum) / mInputLengthStats.size(); -} - -double DynamicBatchTuner::getAverageOutputLength() const -{ - return mOutputLengthStats.empty() ? 0 : static_cast(mOutputLengthSum) / mOutputLengthStats.size(); -} - -SizeType32 DynamicBatchTuner::getRuntimeBatchSize(SizeType32 maxCapacityBatchSize) const -{ - for (auto const& [batchSizeLimit, batchSize] : mBatchSizeTable) - { - if (maxCapacityBatchSize < batchSizeLimit) - { - return batchSize; - } - } - SizeType32 threshold = maxCapacityBatchSize / kBatchSizeFallbackGranularity * kBatchSizeFallbackGranularity; - if (maxCapacityBatchSize < (threshold + kBatchSizeFallbackThreshold)) - { - return threshold; - } - return maxCapacityBatchSize; -} - -SizeType32 DynamicBatchTuner::getRuntimeMaxNumTokens(SizeType32 maxRuntimeBatchSize) const -{ - // calculate max num token in fully overlapped case - SizeType32 adjustedNumTokens - = 1.0 * (maxRuntimeBatchSize * getAverageInputLength() / getAverageOutputLength() + maxRuntimeBatchSize); - SizeType32 tokenThreshold; - // context heavy (avg ISL/OSL > kMaxNumTokensRatioContextHeavy) - if (getAverageInputLength() / getAverageOutputLength() > kMaxNumTokensRatioContextHeavy) - { - tokenThreshold = kMaxNumTokensThresholdContextHeavy; - } - // balanced case (kMaxNumTokensRatioBalanced < avg ISL/OSL < kMaxNumTokensRatioContextHeavy) - else if (getAverageInputLength() / getAverageOutputLength() > kMaxNumTokensRatioBalanced) - { - tokenThreshold = kMaxNumTokensThresholdBalanced; - } - // gen heavy (avg ISL/OSL < kMaxNumTokensRatioBalanced) - else - { - tokenThreshold = kMaxNumTokensThresholdGenHeavy; - } - // pad it to pow of 2 and max of this value and threshold. - return (std::max(1 << int(ceil(log2(adjustedNumTokens))), tokenThreshold)); -} - -} // namespace tensorrt_llm::executor diff --git a/cpp/tensorrt_llm/executor/dynamicBatchTuner.h b/cpp/tensorrt_llm/executor/dynamicBatchTuner.h deleted file mode 100644 index df38cd157fde..000000000000 --- a/cpp/tensorrt_llm/executor/dynamicBatchTuner.h +++ /dev/null @@ -1,87 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2025 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. - */ - -#pragma once - -#include "tensorrt_llm/executor/executor.h" -#include "tensorrt_llm/executor/types.h" - -#include - -namespace tensorrt_llm::executor -{ - -/// @brief A class that maintains runtime input and output length statistics and computes runtime dynamic batch size. -class DynamicBatchTuner -{ -public: - explicit DynamicBatchTuner(DynamicBatchConfig const& config); - - /// @brief Check if dynamic batch size tuning is enabled. - [[nodiscard]] bool isBatchSizeTuningEnabled() const - { - return mEnableBatchSizeTuning; - } - - /// @brief Check if max num tokens tuning is enabled. - [[nodiscard]] bool isMaxNumTokensTuningEnabled() const - { - return mEnableMaxNumTokensTuning; - } - - /// @brief Update current stats given the input and output length from a single request. - void updateStats(SizeType32 inputLen, SizeType32 outputLen); - - /// @brief Get average input length. - [[nodiscard]] double getAverageInputLength() const; - - /// @brief Get average output length. - [[nodiscard]] double getAverageOutputLength() const; - - /// @brief Get the dynamic batch size based on the current statistics. - [[nodiscard]] SizeType32 getRuntimeBatchSize(SizeType32 maxCapacityBatchSize) const; - - /// @brief Get the dynamic max num tokens based on the current statistics. - [[nodiscard]] SizeType32 getRuntimeMaxNumTokens(SizeType32 runtimeBatchSize) const; - -private: - bool mEnableBatchSizeTuning = false; - - bool mEnableMaxNumTokensTuning = false; - - SizeType32 mDynamicBatchMovingAverageWindow = 0; - - std::vector> mBatchSizeTable; - - int64_t mInputLengthSum = 0; - std::deque mInputLengthStats; - - int64_t mOutputLengthSum = 0; - std::deque mOutputLengthStats; - - static SizeType32 const kBatchSizeFallbackGranularity = 512; - static SizeType32 const kBatchSizeFallbackThreshold = 128; - - static double constexpr kMaxNumTokensRatioContextHeavy = 2.0; - static double constexpr kMaxNumTokensRatioBalanced = 0.5; - - static SizeType32 const kMaxNumTokensThresholdContextHeavy = 8192; - static SizeType32 const kMaxNumTokensThresholdBalanced = 4096; - static SizeType32 const kMaxNumTokensThresholdGenHeavy = 2048; -}; - -} // namespace tensorrt_llm::executor diff --git a/cpp/tensorrt_llm/executor/executor.cpp b/cpp/tensorrt_llm/executor/executor.cpp deleted file mode 100644 index 091bb5128230..000000000000 --- a/cpp/tensorrt_llm/executor/executor.cpp +++ /dev/null @@ -1,144 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2025 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. - */ - -#include -#include - -namespace tensorrt_llm::executor -{ - -Executor::Executor(std::filesystem::path const& modelPath, ModelType modelType, ExecutorConfig const& executorConfig) - : mImpl(std::make_unique(modelPath, std::nullopt, modelType, executorConfig)) -{ -} - -Executor::Executor(std::filesystem::path const& encoderModelPath, std::filesystem::path const& decoderModelPath, - ModelType modelType, ExecutorConfig const& executorConfig) - : mImpl(std::make_unique(decoderModelPath, encoderModelPath, modelType, executorConfig)) -{ -} - -Executor::Executor(BufferView const& engineBuffer, std::string const& jsonConfigStr, ModelType modelType, - ExecutorConfig const& executorConfig, std::optional> const& managedWeights) - : mImpl(std::make_unique( - engineBuffer, jsonConfigStr, std::nullopt, std::nullopt, modelType, executorConfig, managedWeights)) -{ -} - -Executor::Executor(BufferView const& encoderEngineBuffer, std::string const& encoderJsonConfigStr, - BufferView const& decoderEngineBuffer, std::string const& decoderJsonConfigStr, ModelType modelType, - ExecutorConfig const& executorConfig) - : mImpl(std::make_unique(decoderEngineBuffer, decoderJsonConfigStr, encoderEngineBuffer, - encoderJsonConfigStr, modelType, executorConfig, std::nullopt)) -{ -} - -Executor::Executor(std::shared_ptr model, ExecutorConfig const& executorConfig) - : mImpl(std::make_unique(std::move(model), std::nullopt, executorConfig)) -{ -} - -Executor::Executor( - std::shared_ptr encoderModel, std::shared_ptr decoderModel, ExecutorConfig const& executorConfig) - : mImpl(std::make_unique(std::move(decoderModel), std::move(encoderModel), executorConfig)) -{ -} - -Executor::~Executor() = default; - -IdType Executor::enqueueRequest(Request const& llmRequest) -{ - return mImpl->enqueueRequest(llmRequest); -} - -std::vector Executor::enqueueRequests(std::vector const& llmRequests) -{ - return mImpl->enqueueRequests(llmRequests); -} - -std::vector Executor::awaitResponses(std::optional const& timeout) -{ - return mImpl->awaitResponses(timeout); -} - -std::vector Executor::awaitResponses( - IdType const& requestId, std::optional const& timeout) -{ - return mImpl->awaitResponses(requestId, timeout); -} - -std::vector> Executor::awaitResponses( - std::vector const& requestIds, std::optional const& timeout) -{ - return mImpl->awaitResponses(requestIds, timeout); -} - -SizeType32 Executor::getNumResponsesReady(std::optional const& requestId) const -{ - return mImpl->getNumResponsesReady(requestId); -} - -void Executor::cancelRequest(IdType requestId) -{ - return mImpl->cancelRequest(requestId); -} - -void Executor::shutdown() -{ - return mImpl->shutdown(); -} - -std::deque Executor::getLatestIterationStats() -{ - return mImpl->getLatestIterationStats(); -} - -std::deque Executor::getLatestRequestStats() -{ - return mImpl->getLatestRequestStats(); -} - -std::deque Executor::getLatestDebugTensors() -{ - return mImpl->getLatestDebugTensors(); -} - -bool Executor::canEnqueueRequests() const -{ - return mImpl->canEnqueueRequests(); -} - -bool Executor::isParticipant() const -{ - return mImpl->isParticipant(); -} - -std::optional> Executor::getKVCacheEventManager() const -{ - return mImpl->getKVCacheEventManager(); -} - -KVCacheEvent::KVCacheEvent( - size_t eventId, KVCacheEventData data, SizeType32 windowSize, std::optional attentionDpRank) - : eventId{eventId} - , data{std::move(data)} - , windowSize{windowSize} - , attentionDpRank{attentionDpRank} -{ -} - -} // namespace tensorrt_llm::executor diff --git a/cpp/tensorrt_llm/executor/executorImpl.cpp b/cpp/tensorrt_llm/executor/executorImpl.cpp deleted file mode 100644 index 9f7fb654a2d5..000000000000 --- a/cpp/tensorrt_llm/executor/executorImpl.cpp +++ /dev/null @@ -1,2791 +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. - */ - -#include "tensorrt_llm/executor/executorImpl.h" -#include "tensorrt_llm/batch_manager/trtEncoderModel.h" -#include "tensorrt_llm/batch_manager/trtGptModelFactory.h" -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/cudaProfilerUtils.h" -#include "tensorrt_llm/common/logger.h" -#include "tensorrt_llm/common/nvtxUtils.h" -#include "tensorrt_llm/common/timestampUtils.h" -#include "tensorrt_llm/common/utils.h" -#include "tensorrt_llm/executor/dataTransceiverState.h" -#include "tensorrt_llm/executor/executor.h" -#include "tensorrt_llm/executor/orchestratorUtils.h" -#include "tensorrt_llm/executor/requestUtils.h" -#include "tensorrt_llm/executor/serialization.h" -#include "tensorrt_llm/executor/serializeUtils.h" -#include "tensorrt_llm/executor/types.h" -#include "tensorrt_llm/executor/version.h" -#include "tensorrt_llm/runtime/loraCache.h" -#include "tensorrt_llm/runtime/memoryCounters.h" -#include "tensorrt_llm/runtime/utils/mpiTags.h" -#include "tensorrt_llm/runtime/utils/mpiUtils.h" - -#include -#include -#include -#include -#include -#include -#include -#include - -namespace tensorrt_llm::executor -{ - -namespace -{ - -[[nodiscard]] bool executorConfigIsValid( - ::tensorrt_llm::executor::ExecutorConfig const& executorConfig, runtime::ModelConfig const& modelConfig) -{ - // Make sure logic in this function matches fixExecutorConfig - if (executorConfig.getEnableChunkedContext()) - { - if (modelConfig.isRnnBased() || !modelConfig.isKVCacheEnabled() || !modelConfig.getPagedContextFMHA()) - { - return false; - } - } - return true; -} - -[[nodiscard]] ::tensorrt_llm::executor::ExecutorConfig fixExecutorConfig( - ::tensorrt_llm::executor::ExecutorConfig const& executorConfig, runtime::ModelConfig const& modelConfig) -{ - // Make sure logic in this function matches executorConfigIsValid - auto fixedExecutorConfig = executorConfig; - // Disable chunked context when not supported - if (executorConfig.getEnableChunkedContext()) - { - if (modelConfig.isRnnBased() || !modelConfig.isKVCacheEnabled() || !modelConfig.getPagedContextFMHA()) - { - fixedExecutorConfig.setEnableChunkedContext(false); - TLLM_LOG_WARNING( - "Chunked context is not supported for this configuration and will be disabled. " - "Related configs: RNNBased: %d, KVCacheEnabled: %d, PagedContextFMHA: %d", - modelConfig.isRnnBased(), modelConfig.isKVCacheEnabled(), modelConfig.getPagedContextFMHA()); - } - } - return fixedExecutorConfig; -} - -[[nodiscard]] bool statsBufferIsEnabled(SizeType32 maxIterations) -{ - return maxIterations != 0; -} - -[[nodiscard]] bool statsBufferIsBounded(SizeType32 maxIterations) -{ - return maxIterations > 0; -} - -SizeType32 getNumChildRequests(Request const& request) -{ - auto samplingConfig = request.getSamplingConfig(); - return samplingConfig.getBeamWidth() > 1 ? 0 : samplingConfig.getNumReturnSequences().value_or(1) - 1; -} - -} // namespace - -/// @brief Version of TRT-LLM as defined in tensorrt_llm/version.py -char const* version() noexcept -{ - return kTensorRtLlmVersion; -} - -class CancelledRequestsAsyncSend -{ -public: - CancelledRequestsAsyncSend(std::shared_ptr const& commSession, - std::unordered_set const& cancelledReqIds, int peer) - { - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - mNumReq = static_cast(cancelledReqIds.size()); - TLLM_LOG_DEBUG("start send %ld cancelled requests to rank %d", mNumReq, peer); - mRequest1 - = commSession->sendAsync(&mNumReq, 1, mpi::MpiType::kINT64, peer, mpi::MpiTag::kCancelledRequestsNumReq); - if (mNumReq > 0) - { - mIds.assign(cancelledReqIds.begin(), cancelledReqIds.end()); - mRequest2 = commSession->sendAsync( - mIds.data(), mIds.size(), mpi::MpiType::kUINT64, peer, mpi::MpiTag::kCancelledRequestsIds); - } - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); - } - - ~CancelledRequestsAsyncSend() - { - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - mRequest1->wait(); - if (mRequest2) - { - mRequest2->wait(); - } - TLLM_LOG_DEBUG("end send cancelled requests"); - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); - } - - CancelledRequestsAsyncSend(CancelledRequestsAsyncSend const& executor) = delete; - CancelledRequestsAsyncSend& operator=(CancelledRequestsAsyncSend const& executor) = delete; - CancelledRequestsAsyncSend(CancelledRequestsAsyncSend&&) = delete; - CancelledRequestsAsyncSend& operator=(CancelledRequestsAsyncSend&&) = delete; - - static std::unordered_set cancelledRequestsRecv( - std::shared_ptr const& commSession, int peer) - { - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - TLLM_LOG_DEBUG("start recv cancelled requests from rank %d", peer); - std::unordered_set cancelledReqIds; - int64_t numReq{0}; - commSession->recv(&numReq, 1, mpi::MpiType::kINT64, peer, mpi::MpiTag::kCancelledRequestsNumReq); - TLLM_LOG_DEBUG("recv %ld cancelled requests", numReq); - if (numReq > 0) - { - std::vector buffer(numReq); - commSession->recv( - buffer.data(), buffer.size(), mpi::MpiType::kUINT64, peer, mpi::MpiTag::kCancelledRequestsIds); - cancelledReqIds = std::unordered_set(buffer.begin(), buffer.end()); - } - TLLM_LOG_DEBUG("end recv cancelled requests from rank %d", peer); - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); - return cancelledReqIds; - } - -private: - int64_t mNumReq; - std::vector mIds; - std::shared_ptr mRequest1; - std::shared_ptr mRequest2; -}; - -class RequestWithIdAsyncSend -{ -public: - RequestWithIdAsyncSend(std::shared_ptr const& commSession, - std::vector const& reqWithIds, int peer) - { - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - TLLM_LOG_DEBUG("start send requests to rank %d", peer); - mNumReq = static_cast(reqWithIds.size()); - mRequest1 = commSession->sendAsync(&mNumReq, 1, mpi::MpiType::kINT64, peer, mpi::MpiTag::kRequestWithIdNumReq); - if (mNumReq > 0) - { - mPacked = RequestWithId::serializeReqWithIds(reqWithIds); - mVecSize = static_cast(mPacked.size()); - mRequest2 - = commSession->sendAsync(&mVecSize, 1, mpi::MpiType::kINT64, peer, mpi::MpiTag::kRequestWithIdVecSize); - mRequest3 = commSession->sendAsync( - mPacked.data(), mPacked.size(), mpi::MpiType::kCHAR, peer, mpi::MpiTag::kRequestWithIdPacked); - } - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); - } - - ~RequestWithIdAsyncSend() - { - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - mRequest1->wait(); - if (mRequest2) - { - mRequest2->wait(); - } - if (mRequest3) - { - mRequest3->wait(); - } - TLLM_LOG_DEBUG("end send requests"); - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); - } - - RequestWithIdAsyncSend(RequestWithIdAsyncSend const& executor) = delete; - RequestWithIdAsyncSend& operator=(RequestWithIdAsyncSend const& executor) = delete; - RequestWithIdAsyncSend(RequestWithIdAsyncSend&&) = delete; - RequestWithIdAsyncSend& operator=(RequestWithIdAsyncSend&&) = delete; - - static std::vector requestWithIdRecv( - std::shared_ptr const& commSession, int peer) - { - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - TLLM_LOG_DEBUG("start recv requests from rank %d", peer); - std::vector reqWithIds; - int64_t numReq{0}; - commSession->recv(&numReq, 1, mpi::MpiType::kINT64, peer, mpi::MpiTag::kRequestWithIdNumReq); - if (numReq > 0) - { - std::vector buffer; - int64_t vecSize = 0; - commSession->recv(&vecSize, 1, mpi::MpiType::kINT64, peer, mpi::MpiTag::kRequestWithIdVecSize); - buffer.resize(vecSize); - commSession->recv( - buffer.data(), buffer.size(), mpi::MpiType::kCHAR, peer, mpi::MpiTag::kRequestWithIdPacked); - reqWithIds = RequestWithId::deserializeReqWithIds(buffer); - } - TLLM_LOG_DEBUG("end recv requests from rank %d", peer); - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); - return reqWithIds; - } - -private: - int64_t mNumReq; - int64_t mVecSize; - std::vector mPacked; - std::shared_ptr mRequest1; - std::shared_ptr mRequest2; - std::shared_ptr mRequest3; -}; - -void Executor::Impl::loadModel(std::optional const& modelPathOpt, - std::optional const& engineBufferOpt, runtime::GptJsonConfig const& jsonConfig, - ::tensorrt_llm::executor::ExecutorConfig const& executorConfig, bool isEncoder, - std::optional> const& managedWeightsOpt) -{ - auto const gpusPerNode = jsonConfig.getGpusPerNode(); - auto const tp = jsonConfig.getTensorParallelism(); - auto const pp = jsonConfig.getPipelineParallelism(); - auto const cp = jsonConfig.getContextParallelism(); - auto parallelConfig = executorConfig.getParallelConfig().value_or(ParallelConfig()); - auto worldConfig = runtime::WorldConfig::mpi(gpusPerNode, tp, pp, cp, parallelConfig.getDeviceIds()); - - TLLM_CHECK_WITH_INFO(modelPathOpt.has_value() || engineBufferOpt.has_value(), - "Either engine path or deserialized engine buffer should be given to load the model properly."); - auto rawEngine = engineBufferOpt.has_value() - ? runtime::RawEngine(engineBufferOpt.value().data(), engineBufferOpt.value().size()) - : runtime::RawEngine(modelPathOpt.value() / jsonConfig.engineFilename(worldConfig)); - - if (rawEngine.getType() != tensorrt_llm::runtime::RawEngine::FilePath) - { - if (modelPathOpt.has_value()) - { - rawEngine.setPath(modelPathOpt.value() / jsonConfig.engineFilename(worldConfig)); - if (managedWeightsOpt.has_value()) - { - TLLM_LOG_WARNING( - "Executor::Impl::loadModel: managedWeightsOpt argument is ignored when loading engine from file."); - } - } - else if (managedWeightsOpt.has_value()) - { - rawEngine.setManagedWeightsMap(managedWeightsOpt.value()); - } - } - - auto const& modelConfig = jsonConfig.getModelConfig(); - - if (isEncoder) - { - mEncoderModel = createEncoderModel(rawEngine, modelConfig, worldConfig, executorConfig); - } - else - { - mModel = createModel(rawEngine, modelConfig, worldConfig, executorConfig); - } -}; - -Executor::Impl::Impl(std::filesystem::path const& modelPath, - std::optional const& encoderModelPath, ModelType const modelType, - ::tensorrt_llm::executor::ExecutorConfig const& executorConfig) -{ - auto decoderJsonConfig = runtime::GptJsonConfig::parse(modelPath / "config.json"); - - // for now, assume encoder & decoder models share the same MPI config - auto const tp = decoderJsonConfig.getTensorParallelism(); - auto const pp = decoderJsonConfig.getPipelineParallelism(); - auto const cp = decoderJsonConfig.getContextParallelism(); - initializeCommAndWorkers(tp, pp, cp, executorConfig, modelType, modelPath, std::nullopt, decoderJsonConfig); - - if (mIsWorker) - { - if (modelType == ModelType::kENCODER_DECODER) - { - if (encoderModelPath.has_value()) - { - auto const encoderJsonConfig = runtime::GptJsonConfig::parse(encoderModelPath.value() / "config.json"); - - auto const encoderMaxInputLen = encoderJsonConfig.getModelConfig().getMaxInputLen(); - auto const encoderHiddenSize = encoderJsonConfig.getModelConfig().getHiddenSize() - * encoderJsonConfig.getTensorParallelism(); // recover full hidden size - // add encoder info to decoder for encoder-decoder models - // note: GptJsonConfig can no longer have modelConfig as const member since it must be mutable here - decoderJsonConfig.getModelConfigMutable().setMaxEncoderLen(encoderMaxInputLen); - decoderJsonConfig.getModelConfigMutable().setEncoderHiddenSize(encoderHiddenSize); - - loadModel( - encoderModelPath.value(), std::nullopt, encoderJsonConfig, executorConfig, true, std::nullopt); - } - else - { - TLLM_LOG_WARNING("Encoder model path not provided. Skipping Encoder Run."); - } - } - loadModel(modelPath, std::nullopt, decoderJsonConfig, executorConfig, false, std::nullopt); - } - initialize(executorConfig); -} - -Executor::Impl::Impl(BufferView const& engineBufferView, std::string const& jsonConfigStr, - std::optional const& encoderEngineBufferView, std::optional const& encoderJsonConfigStr, - ModelType const modelType, ::tensorrt_llm::executor::ExecutorConfig const& executorConfig, - std::optional> const& managedWeightsOpt) -{ - auto decoderJsonConfig = runtime::GptJsonConfig::parse(jsonConfigStr); - - // for now, assume encoder & decoder models share the same MPI config - auto const tp = decoderJsonConfig.getTensorParallelism(); - auto const pp = decoderJsonConfig.getPipelineParallelism(); - auto const cp = decoderJsonConfig.getContextParallelism(); - initializeCommAndWorkers(tp, pp, cp, executorConfig, modelType, std::nullopt, std::nullopt, decoderJsonConfig); - - if (mIsWorker) - { - if (modelType == ModelType::kENCODER_DECODER) - { - TLLM_CHECK(encoderEngineBufferView.has_value() && encoderJsonConfigStr.has_value()); - TLLM_CHECK_WITH_INFO( - !managedWeightsOpt.has_value(), "Managed weights are not supported for enc-dec models"); - - auto const encoderJsonConfig = runtime::GptJsonConfig::parse(encoderJsonConfigStr.value()); - - auto const encoderMaxInputLen = encoderJsonConfig.getModelConfig().getMaxInputLen(); - auto const encoderHiddenSize = encoderJsonConfig.getModelConfig().getHiddenSize() - * encoderJsonConfig.getTensorParallelism(); // recover full hidden size - // add encoder info to decoder for encoder-decoder models - // note: GptJsonConfig can no longer have modelConfig as const member since it must be mutable here - decoderJsonConfig.getModelConfigMutable().setMaxEncoderLen(encoderMaxInputLen); - decoderJsonConfig.getModelConfigMutable().setEncoderHiddenSize(encoderHiddenSize); - - loadModel( - std::nullopt, encoderEngineBufferView.value(), encoderJsonConfig, executorConfig, true, std::nullopt); - } - loadModel(std::nullopt, engineBufferView, decoderJsonConfig, executorConfig, false, managedWeightsOpt); - } - initialize(executorConfig); -} - -Executor::Impl::Impl(std::shared_ptr model, std::optional> encoderModel, - ::tensorrt_llm::executor::ExecutorConfig const& executorConfig) -{ - auto const& worldConfig = model->getWorldConfig(); - auto const tp = worldConfig.getTensorParallelism(); - auto const pp = worldConfig.getPipelineParallelism(); - auto const cp = worldConfig.getContextParallelism(); - auto const modelType = encoderModel.has_value() ? ModelType::kENCODER_DECODER : ModelType::kDECODER_ONLY; - initializeCommAndWorkers(tp, pp, cp, executorConfig, modelType, std::nullopt, worldConfig); - if (modelType == ModelType::kENCODER_DECODER) - { - mEncoderModel = encoderModel.value(); - } - mModel = std::move(model); - initialize(executorConfig); -} - -Executor::Impl::~Impl() -{ - shutdown(); -} - -void Executor::Impl::initialize(::tensorrt_llm::executor::ExecutorConfig const& executorConfig) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - mShutdown = false; - mShutdownCalled = false; - mIterStatsMaxIterations = executorConfig.getIterStatsMaxIterations(); - mRequestStatsMaxIterations = executorConfig.getRequestStatsMaxIterations(); - mDebugTensorsMaxIterations - = executorConfig.getDebugConfig() ? executorConfig.getDebugConfig()->getDebugTensorsMaxIterations() : 0; - TLLM_CHECK_WITH_INFO(mDebugTensorsMaxIterations == 0 || mCommMode == CommunicationMode::kLEADER, - "debugTensorsMaxIterations > 0 is only allowed in leader mode."); - mBatchingType = executorConfig.getBatchingType(); - mIsSchedulerMaxUtilization = (executorConfig.getSchedulerConfig().getCapacitySchedulerPolicy() - == CapacitySchedulerPolicy::kMAX_UTILIZATION); - mIsSchedulerGuaranteedNoEvict = (executorConfig.getSchedulerConfig().getCapacitySchedulerPolicy() - == CapacitySchedulerPolicy::kGUARANTEED_NO_EVICT); - mIsChunkedContext = executorConfig.getEnableChunkedContext(); - mPromptTableOffloading = executorConfig.getPromptTableOffloading(); - mMaxQueueSize = executorConfig.getMaxQueueSize(); - - mLastReqId = 1; - - auto const& logitsProcConfig = executorConfig.getLogitsPostProcessorConfig(); - if (logitsProcConfig.has_value()) - { - mLogitsPostProcessorMap = logitsProcConfig.value().getProcessorMap().value_or(LogitsPostProcessorMap{}); - initializeLogitsPostProcessorBatched(logitsProcConfig.value()); - if (!logitsProcConfig.value().getReplicate()) - { - mModel->setReplicateLogitsPostProcessor(false); - } - } - - auto const& commComm = COMM_SESSION; - int32_t const commSize = commComm.getSize(); - if (mIsWorker) - { - if (commSize > 1) - { - auto const& worldConfig = mModel->getWorldConfig(); - auto const& commSession = COMM_SESSION; - auto const& rank = commSession.getRank(); - auto const& tp = worldConfig.getTensorParallelism(); - auto const& cp = worldConfig.getContextParallelism(); - - mCommTensorParallel = std::make_shared( - commSession.split(rank / tp, worldConfig.getTensorParallelRank())); - mCommContextParallel = std::make_shared( - commSession.split(rank / (tp * cp) * tp + rank % tp, worldConfig.getContextParallelRank())); - mCommPipelineParallel = std::make_shared( - commSession.split(rank % (tp * cp), worldConfig.getPipelineParallelRank())); - - if (worldConfig.isPipelineParallel()) - { - mRequestWithIdWaitThread = std::make_unique( - "requestWithIdWaitThread", [this]() { mRequestWithIdAsyncSndHdl.reset(nullptr); }); - mCancelledRequestsWaitThread = std::make_unique( - "cancelledRequestsWaitThread", [this]() { mCancelledRequestsAsyncSndHdl.reset(nullptr); }); - if (mIsLeader) - { - mRequestWithIdLeaderThread - = std::make_unique(&Executor::Impl::requestWithIdLeaderThread, this); - mCancelledRequestsLeaderThread - = std::make_unique(&Executor::Impl::cancelledRequestsLeaderThread, this); - } - } - } - // Launch the execution thread - mMaxNumActiveRequests = mModel->getMaxNumSequences(); - mExecutionThread = std::thread(&Impl::executionLoop, this); - } - - mEnableBlockReuse = executorConfig.getKvCacheConfig().getEnableBlockReuse(); - - auto const& dynamicBatchConfig = executorConfig.getSchedulerConfig().getDynamicBatchConfig(); - if (dynamicBatchConfig) - { - if (mIsWorker) - { - if (mModel->getModelConfig().isTransformerBased() && mModel->getModelConfig().isKVCacheEnabled()) - { - mDynamicBatchTuner = std::make_shared(dynamicBatchConfig.value()); - } - else - { - TLLM_LOG_WARNING("Dynamic batch tuner can only support transformer models that use KV cache."); - } - } - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -std::shared_ptr Executor::Impl::createModel(runtime::RawEngine const& rawEngine, - runtime::ModelConfig const& modelConfig, runtime::WorldConfig const& worldConfig, - ::tensorrt_llm::executor::ExecutorConfig const& executorConfig) -{ - auto const gptModelType = [&executorConfig, &modelConfig]() - { - switch (executorConfig.getBatchingType()) - { - case BatchingType::kSTATIC: - TLLM_THROW( - "Static batching type is deprecated. Please use in-flight batching with " - "CapacitySchedulerPolicy::kSTATIC_BATCH instead."); - case BatchingType::kINFLIGHT: - return modelConfig.isRnnBased() ? batch_manager::TrtGptModelType::InflightBatching - : batch_manager::TrtGptModelType::InflightFusedBatching; - default: TLLM_THROW("Invalid batching strategy"); - } - }(); - - bool const isLeaderInOrchMode = (mCommMode == CommunicationMode::kORCHESTRATOR) && mIsLeader; - auto const& fixedExecutorConfig = executorConfigIsValid(executorConfig, modelConfig) - ? executorConfig - : fixExecutorConfig(executorConfig, modelConfig); - - return batch_manager::TrtGptModelFactory::create( - rawEngine, modelConfig, worldConfig, gptModelType, fixedExecutorConfig, isLeaderInOrchMode); -} - -std::shared_ptr Executor::Impl::createEncoderModel(runtime::RawEngine const& rawEngine, - runtime::ModelConfig const& modelConfig, runtime::WorldConfig const& worldConfig, - ::tensorrt_llm::executor::ExecutorConfig const& executorConfig) -{ - auto fixedExecutorConfig = ExecutorConfig{}; - fixedExecutorConfig.setSchedulerConfig(executorConfig.getSchedulerConfig()); - return std::make_shared( - modelConfig, worldConfig, rawEngine, std::make_shared(), fixedExecutorConfig); -} - -void Executor::Impl::setOrchLeaderComm( - SizeType32 tp, SizeType32 pp, SizeType32 cp, ParallelConfig const& parallelConfig) -{ -#if ENABLE_MULTI_DEVICE - auto optOrchestratorConfig = parallelConfig.getOrchestratorConfig(); - if (optOrchestratorConfig.value().getIsOrchestrator()) - { - TLLM_CHECK_WITH_INFO(mWorldRank == 0, "Rank 0 must be orchestrator"); - } - - TLLM_CHECK_WITH_INFO(parallelConfig.getParticipantIds(), - "When not spawning processes in orchestrator mode, participant IDs must be provided"); - auto participantIds = parallelConfig.getParticipantIds().value(); - - TLLM_CHECK_WITH_INFO(static_cast(participantIds.size()) == tp * pp * cp, - "When specifying participantIds, participantIds size must be equal to tp*pp*cp"); - - bool isLeader = (mWorldRank == participantIds.front()); - bool isOrchestrator = (mWorldRank == 0); - - // OrchLeaderComm rank 0 is orchestrator, rank 1 is leader - mOrchRank = 0; - mLeaderRank = 1; - - // Create a leaderOrch comm - std::vector leaderOrchRanks{0, participantIds.front()}; - - MPI_Group worldGroup = nullptr; - MPICHECK(MPI_Comm_group(MPI_COMM_WORLD, &worldGroup)); // NOLINT - int worldGroupRank = 0; - MPI_Group_rank(worldGroup, &worldGroupRank); - - int worldSize = 0; - MPICHECK(MPI_Group_size(worldGroup, &worldSize)); // NOLINT - TLLM_CHECK_WITH_INFO(participantIds.front() < worldSize, "Not enough ranks in world"); - - MPI_Group leaderOrchCommGroup = nullptr; - MPICHECK( - MPI_Group_incl(worldGroup, leaderOrchRanks.size(), leaderOrchRanks.data(), &leaderOrchCommGroup)); // NOLINT - int leaderOrchGroupRank = 0; - int leaderOrchGroupSize = 0; - MPI_Group_rank(leaderOrchCommGroup, &leaderOrchGroupRank); - MPI_Group_size(leaderOrchCommGroup, &leaderOrchGroupSize); - - if (isOrchestrator || isLeader) - { - MPI_Comm leaderOrchComm = nullptr; - MPICHECK(MPI_Comm_create_group( - MPI_COMM_WORLD, leaderOrchCommGroup, participantIds.front(), &leaderOrchComm)); // NOLINT - mOrchLeaderComm = std::make_shared(leaderOrchComm, false); - } - else - { - mOrchLeaderComm = nullptr; - } -#endif // ENABLE_MULTI_DEVICE -} - -void Executor::Impl::initializeCommAndWorkers(SizeType32 tp, SizeType32 pp, SizeType32 cp, - ::tensorrt_llm::executor::ExecutorConfig const& executorConfig, std::optional modelType, - std::optional const& modelPath, std::optional const& worldConfig, - std::optional const& decoderGptJsonConfig) -{ - if (modelType.has_value() && modelType.value() == ModelType::kENCODER_DECODER) - { - TLLM_CHECK_WITH_INFO(pp == 1, - "Encoder-Decoder C++ runtime doesn't support Pipeline Parallelism currently. Please switch to Python " - "runtime for PP mode, if necessary."); - } - - tensorrt_llm::mpi::initialize(tensorrt_llm::mpi::MpiThreadSupport::THREAD_MULTIPLE); - mWorldRank = tensorrt_llm::mpi::MpiComm::world().getRank(); - mUsePipelineParallel = pp > 1; - - auto parallelConfig = executorConfig.getParallelConfig().value_or(ParallelConfig()); - validateParallelConfig(parallelConfig, modelType, modelPath); - - mCommMode = parallelConfig.getCommunicationMode(); - auto optOrchestratorConfig = parallelConfig.getOrchestratorConfig(); - - mRecvPollPeriodMs = executorConfig.getRecvPollPeriodMs(); - - // Need to create communicator between orchestrator and leader if not spawning processes in orchestrator mode - if (mCommMode == CommunicationMode::kORCHESTRATOR && !optOrchestratorConfig.value().getSpawnProcesses()) - { - setOrchLeaderComm(tp, pp, cp, parallelConfig); - } - - if (mCommMode == CommunicationMode::kORCHESTRATOR && optOrchestratorConfig.value().getIsOrchestrator()) - { - initializeOrchestrator(tp, pp, cp, executorConfig, parallelConfig, modelType.value(), modelPath.value()); - } - else - { - initializeWorkers(tp, pp, cp, parallelConfig, worldConfig, decoderGptJsonConfig); - } -} - -void Executor::Impl::validateParallelConfig(ParallelConfig const& parallelConfig, std::optional modelType, - std::optional const& modelPath) -{ - TLLM_CHECK_WITH_INFO(parallelConfig.getCommunicationType() == CommunicationType::kMPI, - "Only CommunicationType kMPI is supported for now."); - - auto optOrchestratorConfig = parallelConfig.getOrchestratorConfig(); - - if (parallelConfig.getCommunicationMode() == CommunicationMode::kORCHESTRATOR) - { - TLLM_CHECK_WITH_INFO( - optOrchestratorConfig, "OrchestratorConfig must be set when using ORCHESTRATOR communication mode."); - - TLLM_CHECK_WITH_INFO(modelPath, "OrchestratorMode only supports reading model weight from disk currently."); - - TLLM_CHECK_WITH_INFO(modelType, "OrchestratorMode requires modelType to be specified."); - } -} - -void Executor::Impl::initializeOrchestrator(SizeType32 tp, SizeType32 pp, SizeType32 cp, - ::tensorrt_llm::executor::ExecutorConfig const& executorConfig, ParallelConfig parallelConfig, ModelType modelType, - std::filesystem::path const& modelPath) -{ -#if ENABLE_MULTI_DEVICE - namespace su = tensorrt_llm::executor::serialize_utils; - - auto const& worldComm = tensorrt_llm::mpi::MpiComm::world(); - int32_t const worldSize = worldComm.getSize(); - - auto orchestratorConfig = parallelConfig.getOrchestratorConfig().value(); - - mIsWorker = false; - mIsLeader = false; - mIsPipelineLeader = false; - mIsOrchestrator = true; - - // Verify that worldSize is 1 - if (orchestratorConfig.getSpawnProcesses()) - { - TLLM_CHECK_WITH_INFO(worldSize == 1, - "When using the orchestrator mode and isOrchestrator is true, expect MPI worldSize to be 1."); - - // Spawn the worker threads - auto workerExecPath = orchestratorConfig.getWorkerExecutablePath(); - MPI_Comm intercomm = nullptr; - MPI_Info mpiInfo = nullptr; - MPICHECK(MPI_Info_create(&mpiInfo)); - MPICHECK(MPI_Info_set(mpiInfo, "env", "FORCE_NCCL_ALL_REDUCE_STRATEGY")); - - // Binding policy is not inherited for dynamically spawned jobs, resulting in the worker being bound - // to a single core. Override the setting to avoid perf issue - see https://nvbugs/4574329 - MPICHECK(MPI_Info_set(mpiInfo, "bind_to", "none")); - - MPICHECK(MPI_Comm_spawn(workerExecPath.c_str(), MPI_ARGV_NULL, tp * pp * cp, mpiInfo, 0, MPI_COMM_SELF, - &intercomm, MPI_ERRCODES_IGNORE)); - - mOrchLeaderComm = std::make_shared(intercomm, true); - // With intercomm, leader is rank 0 in the local group - mLeaderRank = 0; - mOrchRank = 0; - - // Copy the executor config, but set the orchestrator flag to false - auto newOrchConfig = OrchestratorConfig(false, orchestratorConfig.getWorkerExecutablePath()); - parallelConfig.setOrchestratorConfig(newOrchConfig); - auto execConfig = executorConfig; - execConfig.setParallelConfig(parallelConfig); - - // Serialize and send the executorConfig, the modelType and the modelPath - std::ostringstream oStream; - su::serialize(modelPath.string(), oStream); - su::serialize(modelType, oStream); - su::serialize(execConfig, oStream); - - auto str = oStream.str(); - std::vector buffer(str.begin(), str.end()); - auto bufferSize = static_cast(buffer.size()); - mOrchLeaderComm->bcast(&bufferSize, 1, mpi::MpiType::kINT64, MPI_ROOT); - mOrchLeaderComm->bcast(buffer.data(), buffer.size(), mpi::MpiType::kCHAR, MPI_ROOT); - - // Wait for workers to have created their executor instance - MPICHECK(MPI_Barrier(intercomm)); - } - - // Spawn the thread responsible for sending new requests to the leader of the model - mOrchSendReqThread = std::thread(&Impl::orchSendReqThread, this); - - // Spawn the thread responsible for receiving new responses from the leader of the model - mOrchRecvThread - = std::thread([&]() { this->orchRecvThread(mpi::MpiTag::kOrchestratorId, mpi::MpiTag::kOrchestratorData); }); - -#endif // ENABLE_MULTI_DEVICE -} - -void Executor::Impl::initializeWorkers(SizeType32 tp, SizeType32 pp, SizeType32 cp, ParallelConfig& parallelConfig, - std::optional const& worldConfig, - std::optional const& decoderGptJsonConfig) -{ - auto const& worldComm = tensorrt_llm::mpi::MpiComm::world(); - int32_t const worldSize = worldComm.getSize(); - - auto const& orchestratorConfig = parallelConfig.getOrchestratorConfig(); - mIsOrchestrator = mCommMode == CommunicationMode::kORCHESTRATOR && orchestratorConfig.value().getIsOrchestrator(); - - TLLM_CHECK_WITH_INFO(mCommMode != CommunicationMode::kORCHESTRATOR || orchestratorConfig.has_value(), - "When using ORCHESTRATOR mode, orchestrator config must be set"); - - if (mCommMode == CommunicationMode::kORCHESTRATOR && !orchestratorConfig.value().getSpawnProcesses()) - { - TLLM_CHECK_WITH_INFO(parallelConfig.getParticipantIds(), - "When not spawning processes in orchestrator mode, participant IDs must be provided"); - - // Check that rank 0 is reserved for the orchestrator - auto const participantIds = parallelConfig.getParticipantIds().value(); - for (auto const& participantId : participantIds) - { - TLLM_CHECK_WITH_INFO(participantId != 0, "Rank 0 is reserved for the orchestrator"); - } - } - - // Participant ids - std::vector participantIds; - if (!parallelConfig.getParticipantIds()) - { - TLLM_CHECK_WITH_INFO(worldSize == tp * pp * cp, - "With communicationMode kLEADER, MPI worldSize is expected to be equal to tp*pp*cp when " - "participantIds are not specified"); - - participantIds.resize(tp * pp * cp); - std::iota(participantIds.begin(), participantIds.end(), 0); - } - else - { - if (mCommMode == CommunicationMode::kORCHESTRATOR && orchestratorConfig.value().getSpawnProcesses()) - { - TLLM_THROW( - "Participant ids should not be set when using CommunicationMode::kORCHESTRATOR with " - "spawnProcesses=true"); - } - participantIds = parallelConfig.getParticipantIds().value(); - TLLM_CHECK_WITH_INFO(static_cast(participantIds.size()) == tp * pp * cp, - tensorrt_llm::common::fmtstr("When specifying participantIds, participantIds size (%lu) must be equal to " - "tp*pp*cp (tp is %u, pp is %u, cp is %u)", - participantIds.size(), tp, pp, cp)); - } - - // If deviceIds are specified, check that they match tp*pp*cp - if (parallelConfig.getDeviceIds()) - { - auto deviceIds = parallelConfig.getDeviceIds().value(); - auto const hasNumNodes = parallelConfig.getNumNodes().has_value(); - if (hasNumNodes || static_cast(deviceIds.size()) != tp * pp * cp) - { - auto const numNodes = hasNumNodes ? parallelConfig.getNumNodes().value() : tensorrt_llm::mpi::getNumNodes(); - TLLM_CHECK_WITH_INFO(static_cast(deviceIds.size() * numNodes) == tp * pp * cp, - tensorrt_llm::common::fmtstr("When specifying deviceIds, deviceIds (%lu) * numNodes (%u) must be equal " - "to tp*pp*cp (tp is %u, pp is %u, cp is %u)", - deviceIds.size(), numNodes, tp, pp, cp)); - } - } - - // Bool that indicates if current process is worker for this model or not - auto participantIt = std::find(participantIds.begin(), participantIds.end(), mWorldRank); - mIsWorker = participantIt != participantIds.end(); - // Bool that indicates if current ranks is leader for this model - mIsLeader = (mWorldRank == participantIds.front()); - mIsPipelineLeader = (mWorldRank == participantIds[tp * (pp - 1)]); - -#if ENABLE_MULTI_DEVICE - if (mIsWorker) - { - // Create a session, but only assign to COMM_SESSION for ranks participating in this model - MPI_Group worldGroup = MPI_GROUP_NULL; - MPICHECK(MPI_Comm_group(MPI_COMM_WORLD, &worldGroup)); // NOLINT - MPI_Group sessionGroup = MPI_GROUP_NULL; - if (pp > 1) - { - // reverse participantIds to move leader to last pp rank. retain order in each tp group - std::reverse(participantIds.begin(), participantIds.end()); - if (tp > 1) - { - for (SizeType32 ppRank = 0; ppRank < pp; ppRank++) - { - std::reverse(participantIds.begin() + ppRank * tp, participantIds.begin() + (ppRank + 1) * tp); - } - } - } - MPICHECK(MPI_Group_incl(worldGroup, participantIds.size(), participantIds.data(), &sessionGroup)); // NOLINT - MPI_Comm sessionComm = MPI_COMM_NULL; - MPICHECK( - MPI_Comm_create_group(MPI_COMM_WORLD, sessionGroup, 1000 + participantIds.front(), &sessionComm)); // NOLINT - - tensorrt_llm::mpi::MpiComm::setSession(tensorrt_llm::mpi::MpiComm(sessionComm, false)); - } - - if (mIsLeader && mCommMode == CommunicationMode::kORCHESTRATOR) - { - auto optOrchestratorConfig = parallelConfig.getOrchestratorConfig(); - if (orchestratorConfig.has_value() && orchestratorConfig.value().getSpawnProcesses()) - { - mOrchLeaderComm = optOrchestratorConfig.value().getOrchLeaderComm(); - } - else - { - // mOrchLeaderComm has already been created - } - TLLM_CHECK(mOrchLeaderComm.get() != nullptr); - - TLLM_CHECK(worldConfig.has_value() || decoderGptJsonConfig.has_value()); - if (worldConfig.has_value()) - { - mDeviceId = worldConfig->getDevice(); - } - else - { - auto gpusPerNode = decoderGptJsonConfig->getGpusPerNode(); - auto worldConfig = runtime::WorldConfig::mpi(gpusPerNode, tp, pp, cp, parallelConfig.getDeviceIds()); - mDeviceId = worldConfig.getDevice(); - } - // Spawn the thread responsible for receiving new requests from the orchestrator - mLeaderRecvReqThread = std::thread(&Impl::leaderRecvReqThread, this); - - // Spawn the thread responsible for sending new responses to the orchestrator - mLeaderSendThread = std::thread([&]() - { this->leaderSendThread(mSendQueue, mpi::MpiTag::kOrchestratorId, mpi::MpiTag::kOrchestratorData); }); - } -#endif // ENABLE_MULTI_DEVICE -} - -void Executor::Impl::initializeLogitsPostProcessorBatched(LogitsPostProcessorConfig const& logitsProcConfig) -{ - if (logitsProcConfig.getProcessorBatched().has_value()) - { - mLogitsPostProcessorBatched - = [cb = logitsProcConfig.getProcessorBatched().value()]( - std::vector const& reqIdsVec, - std::vector& logitsVec, - std::vector> const& beamTokensVec, - CudaStreamPtr const& cudaStreamPtr, - std::vector> const& clientIdsVec) - { - std::vector cbLogitsVec; - cbLogitsVec.reserve(logitsVec.size()); - for (auto& logits : logitsVec) - { - cbLogitsVec.emplace_back(executor::detail::ofITensor(logits)); - } - - cb(reqIdsVec, cbLogitsVec, beamTokensVec, cudaStreamPtr, clientIdsVec); - }; - - mModel->setLogitsPostProcessorBatched(mLogitsPostProcessorBatched); - } -} - -IdType Executor::Impl::enqueueRequest(Request const& request) -{ - return enqueueRequests({&request, 1}).at(0); -} - -std::vector Executor::Impl::enqueueRequests(std::vector const& requests) -{ - return enqueueRequests({requests.data(), requests.size()}); -} - -std::vector Executor::Impl::enqueueRequests(common::ArrayView const& requests) -{ - TLLM_CHECK_WITH_INFO(!mShutdownCalled, "Shutdown called, cannot enqueue requests"); - checkParallelApiUsage(__func__); - - TLLM_LOG_DEBUG("Enqueuing %lu requests", requests.size()); - std::vector requestWithIds; - requestWithIds.reserve(requests.size()); - - // First check valid of request in enqueue thread, so Exceptions can be thrown to user. - for (auto const& req : requests) - { - auto logitsPostProcessorName = req.getLogitsPostProcessorName(); - if (logitsPostProcessorName && logitsPostProcessorName.value() != Request::kBatchedPostProcessorName) - { - getLogitsPostProcessor(*logitsPostProcessorName); - } - } - - std::vector ids; - { - auto now = std::chrono::steady_clock::now(); - for (auto const& req : requests) - { - ids.emplace_back(generateReqId(req)); - TLLM_LOG_DEBUG("Enqueue new request with id %d", ids.back()); - - std::vector childReqIds; - auto numChildRequests = getNumChildRequests(req); - if (numChildRequests > 0) - { - childReqIds.reserve(numChildRequests); - for (int childId = 0; childId < numChildRequests; childId++) - { - childReqIds.emplace_back(generateLocalReqId()); - TLLM_LOG_DEBUG("Add new child request with id %d", childReqIds.back()); - } - } - requestWithIds.emplace_back(RequestWithId{req, ids.back(), std::move(childReqIds), now}); - } - } - - if (mCommMode == CommunicationMode::kLEADER) - { - { - std::scoped_lock const lck(mQueuedReqMtx); - if (mMaxQueueSize) - { - auto const maxQueueSize = mMaxQueueSize.value(); - - auto totalRequestSize = 0; - for (auto&& reqWithId : requestWithIds) - { - totalRequestSize += (getNumChildRequests(reqWithId.req) + 1); - } - - if (maxQueueSize > 0 && mQueuedRequests.size() + totalRequestSize > static_cast(maxQueueSize)) - { - TLLM_THROW("Maximum queue size of %d has been reached, please try again later", maxQueueSize); - } - } - - for (auto&& req : requestWithIds) - { - insertRequestInOrder(mQueuedRequests, std::move(req)); - } - } - mQueuedReqCv.notify_one(); - } - else if (mCommMode == CommunicationMode::kORCHESTRATOR) - { - MpiMessage message(MpiId::PENDING_REQUEST); - message.data = PendingRequestData{std::move(requestWithIds)}; - mSendQueue.push(std::move(message)); - } - return ids; -} - -std::vector Executor::Impl::awaitResponses(std::optional const& timeout) -{ - TLLM_CHECK_WITH_INFO(!mShutdownCalled, "Shutdown called"); - checkParallelApiUsage(__func__); - std::unique_lock lck(mResponsesMtx); - auto pred = [this]() -> bool { return !mResponses.empty() || mShutdown; }; - auto storeResponses = [this]() - { - std::vector responses; - for (auto it = mResponses.begin(); it != mResponses.end();) - { - responses.insert(responses.end(), it->second.begin(), it->second.end()); - addTerminatedReqId(it->second, it->first); - it = mResponses.erase(it); - } - return responses; - }; - - std::vector responses; - if (timeout) - { - if (mResponsesCv.wait_for(lck, timeout.value(), pred)) - { - responses = storeResponses(); - } - } - else - { - mResponsesCv.wait(lck, pred); - responses = storeResponses(); - } - return responses; -} - -std::vector Executor::Impl::awaitResponses( - IdType const& reqId, std::optional const& timeout) -{ - TLLM_CHECK_WITH_INFO(!mShutdownCalled, "Shutdown called"); - checkParallelApiUsage(__func__); - std::unique_lock lck(mResponsesMtx); - auto pred = [this, reqId]() -> bool - { return (mResponses.find(reqId) != mResponses.end() && !mResponses.at(reqId).empty()) || mShutdown; }; - auto storeIdResponse = [this, reqId]() - { - std::vector responses; - responses.swap(mResponses.at(reqId)); - mResponses.erase(reqId); - addTerminatedReqId(responses, reqId); - return responses; - }; - - // We don't process a terminated request again. Terminated request is defined as a response - // with isFinal = true for a given requestId. - if (mTerminatedReqIds.contains(reqId)) - { - if (mResponses.find(reqId) != mResponses.end()) - { - TLLM_THROW("ReqId should already be removed from responses!"); - } - std::string const err = "ReqId " + std::to_string(reqId) + " has already been processed and was terminated."; - TLLM_LOG_ERROR("%s", err.c_str()); - - return {Response(reqId, err)}; - } - - std::vector responses; - if (timeout) - { - if (mResponsesCv.wait_for(lck, timeout.value(), pred)) - { - responses = storeIdResponse(); - } - } - else - { - mResponsesCv.wait(lck, pred); - responses = storeIdResponse(); - } - return responses; -} - -std::vector> Executor::Impl::awaitResponses( - std::vector const& requestIds, std::optional const& timeout) -{ - TLLM_CHECK_WITH_INFO(!mShutdownCalled, "Shutdown called"); - checkParallelApiUsage(__func__); - std::vector> responses; - responses.reserve(requestIds.size()); - if (timeout) - { - auto const start_time = std::chrono::high_resolution_clock::now(); - for (auto const requestId : requestIds) - { - auto const elapsed_ms = std::chrono::duration_cast( - std::chrono::high_resolution_clock::now() - start_time); - responses.emplace_back(awaitResponses( - requestId, timeout.value() > elapsed_ms ? timeout.value() - elapsed_ms : std::chrono::milliseconds{0})); - } - } - else - { - for (auto const requestId : requestIds) - { - responses.emplace_back(awaitResponses(requestId)); - } - } - return responses; -} - -SizeType32 Executor::Impl::getNumResponsesReady(std::optional const& optId) const -{ - TLLM_CHECK_WITH_INFO(!mShutdownCalled, "Shutdown called"); - checkParallelApiUsage(__func__); - std::scoped_lock lck(mResponsesMtx); - SizeType32 numResponsesReady = 0; - if (optId) - { - auto const reqId = optId.value(); - auto const respIt = mResponses.find(reqId); - if (respIt != mResponses.end()) - { - numResponsesReady = static_cast(respIt->second.size()); - } - } - else - { - for (auto const& [id, responses] : mResponses) - { - numResponsesReady += static_cast(responses.size()); - } - } - return numResponsesReady; -} - -void Executor::Impl::shutdown() -{ - // Cannot call shutdown multiple times - if (mShutdownCalled) - { - return; - } - mShutdownCalled = true; - - if (!mShutdown) - { - if (mCommMode == CommunicationMode::kLEADER && mIsLeader) - { - // Enqueue a request to indicate to other ranks to terminate - enqueueTerminateRequest(); - } - else if (mCommMode == CommunicationMode::kORCHESTRATOR) - { - if (mIsOrchestrator) - { - // Send to the leader the termination signal - mShutdown = true; - mResponsesCv.notify_all(); - - mSendQueue.push(MpiMessage(MpiId::TERMINATION)); - - // Wait for sender thread to exit - if (mOrchSendReqThread.joinable()) - { - mOrchSendReqThread.join(); - } - // Wait for recv response thread to exit - if (mOrchRecvThread.joinable()) - { - mOrchRecvThread.join(); - } - } - else if (mIsLeader) - { - // Wait for sender thread to exit - if (mLeaderRecvReqThread.joinable()) - { - mLeaderRecvReqThread.join(); - } - // Wait for send response thread to exit - if (mLeaderSendThread.joinable()) - { - mLeaderSendThread.join(); - } - } - } - } - - // Wait for execution thread to terminate - if (mExecutionThread.joinable()) - { - mExecutionThread.join(); - } - - // If we overwrote COMM_SESSION with split, free it now. Otherwise, since - // COMM_SESSION is a global static object, it will be destroyed in an - // undefined order and can cause crashes on program exit. - if (mIsWorker) - { - tensorrt_llm::mpi::MpiComm::setSession(tensorrt_llm::mpi::MpiComm(MPI_COMM_WORLD, false)); - } -} - -void Executor::Impl::cancelRequest(IdType requestId) -{ - TLLM_CHECK_WITH_INFO(!mShutdownCalled, "Shutdown called"); - checkParallelApiUsage(__func__); - - // Check if the request is terminated already. If so, return - { - std::scoped_lock lckResp(mResponsesMtx); - if (mTerminatedReqIds.contains(requestId)) - { - TLLM_LOG_INFO("Ignoring already terminated request %lu", requestId); - return; - } - } - - if (mCommMode == CommunicationMode::kLEADER) - { - std::scoped_lock lck(mCancelReqMtx); - auto& selCancelledReqIds = mUsePipelineParallel ? mPipelineCancelledReqIds : mCancelledReqIds; - selCancelledReqIds.insert(requestId); - } - else if (mCommMode == CommunicationMode::kORCHESTRATOR) - { - MpiMessage message(MpiId::CANCEL_REQUEST); - std::vector cancelledReqIds{requestId}; - message.data = RequestIdsData{std::move(cancelledReqIds)}; - mSendQueue.push(std::move(message)); - } -} - -std::deque Executor::Impl::getLatestIterationStats() -{ - TLLM_CHECK_WITH_INFO(!mShutdownCalled, "Shutdown called"); - checkParallelApiUsage(__func__); - std::scoped_lock lck(mIterStatsMtx); - return std::exchange(mIterationStats, {}); -} - -std::deque Executor::Impl::getLatestRequestStats() -{ - TLLM_CHECK_WITH_INFO(!mShutdownCalled, "Shutdown called"); - checkParallelApiUsage(__func__); - - std::scoped_lock lck(mRequestStatsMtx); - return std::exchange(mRequestStats, {}); -} - -std::deque Executor::Impl::getLatestDebugTensors() -{ - TLLM_CHECK_WITH_INFO(!mShutdownCalled, "Shutdown called"); - if (mCommMode == CommunicationMode::kORCHESTRATOR) - { - TLLM_LOG_WARNING("getLatestDebugTensors is not supported in ORCHESTRATOR mode yet"); - return {}; - } - if (mEncoderModel) - { - TLLM_LOG_WARNING("getLatestDebugTensors is not supported for encoder model yet"); - } - std::scoped_lock lck(mDebugTensorsMtx); - return std::exchange(mDebugTensors, {}); -} - -bool Executor::Impl::canEnqueueRequests() const -{ - return !mShutdownCalled - && ((mCommMode == CommunicationMode::kLEADER && mIsLeader) - || (mCommMode == CommunicationMode::kORCHESTRATOR && mIsOrchestrator)); -} - -bool Executor::Impl::isParticipant() const -{ - return mIsWorker; -} - -std::optional> Executor::Impl::getKVCacheEventManager() const -{ - if (!mModel) - { - return std::nullopt; - } - auto cacheEventManager = mModel->getKVCacheManager(); - return cacheEventManager ? std::optional(std::make_shared(cacheEventManager)) : std::nullopt; -} - -void Executor::Impl::requestWithIdLeaderThread() -{ - TLLM_CUDA_CHECK(cudaSetDevice(mModel->getWorldConfig().getDevice())); - auto constexpr peer = 0; - while (true) - { - int64_t numActiveRequests; - mCommPipelineParallel->recv( - &numActiveRequests, 1, mpi::MpiType::kINT64, peer, mpi::MpiTag::kExecutorNumActiveRequests); - if (numActiveRequests < 0) - { - break; - } - - bool lowestPriorityActiveHasValue; - std::optional lowestPriorityActive; - mCommPipelineParallel->recv(&lowestPriorityActiveHasValue, 1, mpi::MpiType::kBOOL, peer, - mpi::MpiTag::kExecutorLowestPriorityActiveHasValue); - if (lowestPriorityActiveHasValue) - { - PriorityType lowestPriorityActiveValue; - mCommPipelineParallel->recv( - &lowestPriorityActiveValue, 1, mpi::MpiType::kFLOAT, peer, mpi::MpiTag::kExecutorLowestPriorityActive); - lowestPriorityActive = lowestPriorityActiveValue; - } - - auto reqWithIds = getLeaderNewReqWithIds(numActiveRequests, lowestPriorityActive); - setupDynamicLogitsPostProcessors(reqWithIds); - auto requestWithIdAsyncSndHdl - = std::make_unique(mCommPipelineParallel, reqWithIds, peer); - requestWithIdAsyncSndHdl.reset(nullptr); - } -} - -void Executor::Impl::cancelledRequestsLeaderThread() -{ - TLLM_CUDA_CHECK(cudaSetDevice(mModel->getWorldConfig().getDevice())); - auto constexpr peer = 0; - while (true) - { - bool shouldExit; - mCommPipelineParallel->recv(&shouldExit, 1, mpi::MpiType::kBOOL, peer, mpi::MpiTag::kExecutorShouldExit); - if (shouldExit) - { - break; - } - - std::unique_ptr cancelledRequestsAsyncSndHdl; - { - std::scoped_lock lck(mCancelReqMtx); - cancelledRequestsAsyncSndHdl - = std::make_unique(mCommPipelineParallel, mPipelineCancelledReqIds, peer); - mPipelineCancelledReqIds.clear(); - } - cancelledRequestsAsyncSndHdl.reset(nullptr); - } -} - -std::vector Executor::Impl::getLeaderNewReqWithIds( - SizeType32 numActiveRequests, std::optional lowestPriorityActive) -{ - std::unique_lock lck(mQueuedReqMtx); - mQueuedReqCv.wait(lck, [&]() { return (!mQueuedRequests.empty() || numActiveRequests > 0 || mShutdown); }); - - std::vector reqWithIds; - - if (mQueuedRequests.empty() || mShutdown) - { - return reqWithIds; - } - - if (mQueuedRequests.front().id == kTerminateReqId) - { - reqWithIds.emplace_back(std::move(mQueuedRequests.front())); - mQueuedRequests.pop_front(); - return reqWithIds; - } - - auto const& firstRequest = mQueuedRequests.front(); - auto const firstBeamWidth = firstRequest.req.getSamplingConfig().getBeamWidth(); - auto const operatingBeamWidth = numActiveRequests > 0 ? mModel->getOperatingBeamWidth() : firstBeamWidth; - - auto const tryInsertQueuedRequestIntoReqWithIds = [this, &reqWithIds, operatingBeamWidth]() -> bool - { - auto& nextRequest = mQueuedRequests.front(); - auto const beamWidth = nextRequest.req.getSamplingConfig().getBeamWidth(); - if (beamWidth != operatingBeamWidth) - { - TLLM_LOG_INFO( - "Can't dequeue request with ID %ld because beam width %d differs from operating beam width %d.", - nextRequest.id, beamWidth, operatingBeamWidth); - return false; - } - - TLLM_LOG_DEBUG("Dequeue request with ID %ld", nextRequest.id); - reqWithIds.emplace_back(std::move(nextRequest)); - mQueuedRequests.pop_front(); - return true; - }; - - auto const maxNewRequests = static_cast(std::max(mMaxNumActiveRequests - numActiveRequests, 0)); - for (size_t req = 0; !mQueuedRequests.empty() && req < maxNewRequests;) - { - req += (getNumChildRequests(mQueuedRequests.front().req) + 1); - if (req > maxNewRequests) - { - break; - } - if (!tryInsertQueuedRequestIntoReqWithIds()) - { - break; - } - } - - if (lowestPriorityActive) - { - while (!mQueuedRequests.empty() && mQueuedRequests.front().req.getPriority() > (*lowestPriorityActive)) - { - if (!tryInsertQueuedRequestIntoReqWithIds()) - { - break; - } - } - } - return reqWithIds; -} - -std::vector Executor::Impl::getNewReqWithIds( - SizeType32 numActiveRequests, std::optional lowestPriorityActive) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - auto const& worldConfig = mModel->getWorldConfig(); - - if (worldConfig.isPipelineParallel()) - { - mRequestWithIdWaitThread->waitStop(); - } - - TLLM_CUDA_CHECK(cudaSetDevice(mModel->getWorldConfig().getDevice())); - std::vector reqWithIds; - if (mIsPipelineLeader) - { - if (!worldConfig.isPipelineParallel()) - { - reqWithIds = getLeaderNewReqWithIds(numActiveRequests, lowestPriorityActive); - setupDynamicLogitsPostProcessors(reqWithIds); - } - else - { - auto const peer = worldConfig.getPipelineParallelism() - 1; - auto numActiveRequestsValue = static_cast(numActiveRequests); - auto request1 = mCommPipelineParallel->sendAsync( - &numActiveRequestsValue, 1, mpi::MpiType::kINT64, peer, mpi::MpiTag::kExecutorNumActiveRequests); - bool lowestPriorityActiveHasValue = lowestPriorityActive.has_value(); - auto request2 = mCommPipelineParallel->sendAsync(&lowestPriorityActiveHasValue, 1, mpi::MpiType::kBOOL, - peer, mpi::MpiTag::kExecutorLowestPriorityActiveHasValue); - auto request3 = lowestPriorityActiveHasValue - ? mCommPipelineParallel->sendAsync(&lowestPriorityActive.value(), 1, mpi::MpiType::kFLOAT, peer, - mpi::MpiTag::kExecutorLowestPriorityActive) - : nullptr; - request1->wait(); - request2->wait(); - if (request3) - { - request3->wait(); - } - reqWithIds = RequestWithIdAsyncSend::requestWithIdRecv(mCommPipelineParallel, peer); - } - if (worldConfig.isTensorParallel() || worldConfig.isContextParallel()) - { - auto packed = RequestWithId::serializeReqWithIds(reqWithIds); - if (worldConfig.isTensorParallel()) - { - mCommTensorParallel->bcast(packed, 0); - } - if (worldConfig.isContextParallel()) - { - mCommContextParallel->bcast(packed, 0); - } - } - } - else - { - if (worldConfig.isFirstPipelineParallelRank()) - { - std::vector buffer; - mCommTensorParallel->bcast(buffer, 0); - mCommContextParallel->bcast(buffer, 0); - reqWithIds = RequestWithId::deserializeReqWithIds(buffer); - } - else - { - auto const peer = worldConfig.getPipelineParallelRank() - 1; - reqWithIds = RequestWithIdAsyncSend::requestWithIdRecv(mCommPipelineParallel, peer); - } - } - if (!worldConfig.isLastPipelineParallelRank()) - { - auto const peer = worldConfig.getPipelineParallelRank() + 1; - mRequestWithIdAsyncSndHdl = std::make_unique(mCommPipelineParallel, reqWithIds, peer); - mRequestWithIdWaitThread->notifyStart(); - } - TLLM_CUDA_CHECK(cudaSetDevice(mModel->getWorldConfig().getDevice())); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); - return reqWithIds; -} - -std::tuple Executor::Impl::fetchNewRequests( - SizeType32 numActiveRequests, std::optional lowestPriorityActive) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(fetchNewRequests); - - // If grab requests from queue, do exchange between ranks - auto reqWithIds = getNewReqWithIds(numActiveRequests, lowestPriorityActive); - RequestList newRequests; - double newActiveRequestsQueueLatencyMS{0.}; - for (auto& reqWithId : reqWithIds) - { - if (reqWithId.id == kTerminateReqId) - { - mShutdown = true; - mResponsesCv.notify_all(); - return {}; - } - - try - { - std::optional llmRequestLogitsPostProcessor; - bool applyLogitsPostProcessorBatched{false}; - if (mModel->getWorldConfig().isLastPipelineParallelRank()) - { - auto logitsPostProcessorName = reqWithId.req.getLogitsPostProcessorName(); - if (logitsPostProcessorName) - { - if (logitsPostProcessorName.value() == Request::kBatchedPostProcessorName) - { - TLLM_CHECK_WITH_INFO( - mLogitsPostProcessorBatched, "Batched logits post processor is not defined."); - applyLogitsPostProcessorBatched = true; - } - else - { - if (logitsPostProcessorName->compare(0, - std::char_traits::length(Request::kDynamicPostProcessorNamePrefix), - Request::kDynamicPostProcessorNamePrefix) - == 0) - { - TLLM_CHECK_WITH_INFO(!mModel->getReplicateLogitsPostProcessor() - || mModel->getWorldConfig().getTensorParallelism() == 1, - "Dynamic logits postprocessor must be used with replicate=false or no tensor " - "parallelism."); - } - if (mModel->getWorldConfig().isFirstTensorParallelRank() - || mModel->getReplicateLogitsPostProcessor()) - { - llmRequestLogitsPostProcessor = getLogitsPostProcessor(logitsPostProcessorName.value()); - } - else - { - llmRequestLogitsPostProcessor - = [](IdType reqId, RtTensorPtr& logits, BeamTokens const& beamTokens, - CudaStreamPtr const& cudaStreamPtr, std::optional clientId) {}; - } - } - } - } - auto newLlmReq = std::make_shared( - reqWithId.id, reqWithId.req, llmRequestLogitsPostProcessor, applyLogitsPostProcessorBatched); - - auto numReturnSequences = newLlmReq->getNumSubRequests(); - if (numReturnSequences > 1) - { - TLLM_CHECK(reqWithId.childReqIds.size() == static_cast(numReturnSequences - 1)); - mChildReqIdsMap[reqWithId.id] = reqWithId.childReqIds; - } - - for (auto seqIdx = 0; seqIdx < numReturnSequences; seqIdx++) - { - auto newReq - = seqIdx == 0 ? newLlmReq : newLlmReq->createChildRequest(reqWithId.childReqIds.at(seqIdx - 1)); - - // If static batching and streaming, disable streaming and exclude input - if (mBatchingType == BatchingType::kSTATIC && newReq->isStreaming()) - { - newReq->setStreaming(false); - newReq->setExcludeInputFromOutput(true); - } - - // Validate the request parameters - newReq->validate(mModel->getMaxInputLen(), mModel->getMaxSequenceLen(), mModel->getMaxDraftLen(), - mModel->getVocabSizePadded(), - mEncoderModel ? std::optional(mEncoderModel->getMaxInputLen()) : std::nullopt, - mEnableBlockReuse); - - TLLM_CHECK_WITH_INFO(!mEncoderModel || !mIsSchedulerMaxUtilization, - "Encoder or Encoder-Decoder model don't support max utilization scheduler yet. Only max requests " - "or guaranteed no evict."); - - // When streaming is enabled and scheduling policy permits evict/restart, need to guard against the case - // where the sequence is truncated on eviction (to respect maxInputLen limits), resulting in loss of - // some tokens that have been streamed out. In this case, resuming generation may result in different - // completion for locations whose tokens have already been returned. There is no way to protect against - // this, so disallowing. - if (newReq->isStreaming() && !mIsSchedulerGuaranteedNoEvict && !mIsChunkedContext) - { - auto const maxReqSeqLen = newReq->mPromptLen + newReq->mMaxNewTokens; - auto const maxRestartLen = maxReqSeqLen - 1; - TLLM_CHECK_WITH_INFO(maxRestartLen <= mModel->getMaxInputLen(), - "Request sequence length is potentially greater than max input length. This cannot be run " - "unless streaming is disabled, context chunking is enabled or the GUARANTEED_NO_EVICT " - "scheduling policy is used"); - } - - // Create the encoder output tensor - if (mEncoderModel) - { - TLLM_CHECK_WITH_INFO(mModel || (!mModel && newReq->getReturnEncoderOutput()), - "Encoder-Decoder models allow optionally returning encoder output. But if it is Encoder-only " - "models, please make sure returnEncoderOutput is always true."); - - // gpu buffers for passing to the next phase - newReq->allocEncoderOutput(mEncoderModel->getBufferManager(), mEncoderModel->getLogitDataType()); - newReq->allocEncoderHiddenStates( - mEncoderModel->getBufferManager(), mEncoderModel->getLogitDataType()); - // pinned buffers for returning results to host - if (newReq->getReturnEncoderOutput()) - { - newReq->allocEncoderOutputHost( - mEncoderModel->getHiddenSize() * mEncoderModel->getWorldConfig().getTensorParallelism(), - mEncoderModel->getLogitDataType()); - } - } - - if (!mEncoderModel && newReq->getEncoderInputFeatures()) - { - TLLM_LOG_INFO("Allocating buffers for encoder output"); - // gpu buffers for passing to the next phase - newReq->allocEncoderOutput(mModel->getBufferManager(), mModel->getLogitDataType()); - newReq->allocEncoderHiddenStates(mModel->getBufferManager(), mModel->getLogitDataType()); - } - - // Create the context logits tensor - if (newReq->getReturnContextLogits()) - { - TLLM_CHECK_WITH_INFO(mModel->getModelConfig().computeContextLogits(), - "Return context logit need to build engine with gather_context_logits"); - newReq->allocContextLogitsHost(mModel->getVocabSizePadded(), mModel->getLogitDataType()); - } - - // Create the generation logits tensor - if (newReq->getReturnGenerationLogits()) - { - TLLM_CHECK_WITH_INFO(mModel->getGatherGenerationLogits(), - "To return generation logits, gather_generation_logits must be enabled in ExecutorConfig"); - - if (mModel->getModelConfig().getSpeculativeDecodingMode().isDraftTokensExternal() - && newReq->hasDraftTokens()) - { - newReq->allocTargetModelAcceptedTokenLogitsHost( - mModel->getVocabSizePadded(), mModel->getLogitDataType()); - } - else - { - newReq->allocGenerationLogitsHost(mModel->getVocabSizePadded(), mModel->getLogitDataType()); - } - } - - if (mModel->getWorldConfig().isLastPipelineParallelRank() && newReq->getGuidedDecodingParams()) - { - TLLM_CHECK_WITH_INFO(mModel->hasGuidedDecoder(), - "Request is specified with GuidedDecodingParams, but GuidedDecoder is not setup. Please " - "provide a valid GuidedDecodingConfig to setup GuidedDecoder."); - } - - if (mModel->getWorldConfig().isLastPipelineParallelRank() && newReq->hasAdditionalOutputs()) - { - newReq->allocAdditionalOutputs([this](std::string const& name) - { return mModel->getTensorDataType(name); }, - [this](std::string const& name) { return mModel->getTensorShape(name); }); - } - - mModel->updatePeftCache(newReq); - - newRequests.emplace_back(std::move(newReq)); - } - - auto queuedEnd = std::chrono::steady_clock::now(); - auto reqQueueLatencyMS - = std::chrono::duration(queuedEnd - reqWithId.queuedStart).count(); - newActiveRequestsQueueLatencyMS += reqQueueLatencyMS; - } - catch (runtime::LoraExpectedException const& e) - { - if (mIsLeader) - { - // In case of an expected LoRA exception (e.g. cache full, cache miss), log a warning and enqueue - // response - TLLM_LOG_WARNING("%s", e.what()); - enqueueNewResponses({{reqWithId.id, e.what(), reqWithId.req.getClientId()}}); - } - } - catch (std::exception const& e) - { - if (mIsLeader) - { - // In case of error, create a response with error for this request - auto err = std::string("Encountered an error when fetching new request: ") + e.what(); - TLLM_LOG_ERROR("%s", err.c_str()); - enqueueNewResponses({{reqWithId.id, err, reqWithId.req.getClientId()}}); - } - } - } - TLLM_LOG_DEBUG("[RANK %d] num new requests fetched from queue: %d", COMM_SESSION.getRank(), newRequests.size()); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); - return {newRequests, newActiveRequestsQueueLatencyMS}; -} - -void Executor::Impl::terminateActiveRequests(RequestList& activeRequests, std::string const& err) -{ - TLLM_LOG_ERROR("%s", err.c_str()); - - // Create a response for all requests and add to queue - for (auto it = activeRequests.cbegin(); it != activeRequests.cend();) - { - auto llmReq = (*it); - - llmReq->setState(batch_manager::LlmRequestState::kGENERATION_COMPLETE); - mModel->terminateRequest(llmReq); - - if (mIsLeader) - { - enqueueNewResponses({{llmReq->mRequestId, err, llmReq->mClientId}}); - } - - // Remove from the requestList - it = activeRequests.erase(it); - } -} - -void Executor::Impl::forwardSync(RequestList& activeRequests) -{ - TLLM_LOG_TRACE("[RANK %d] %s start", COMM_SESSION.getRank(), __PRETTY_FUNCTION__); - try - { - if (mEncoderModel) - { - mEncoderModel->forwardSync(); - } - mModel->forwardSync(); - } - catch (std::exception const& e) - { - std::string const err = std::string("Encountered an error in forwardSync function: ") + e.what(); - terminateActiveRequests(activeRequests, err); - } - TLLM_LOG_TRACE("[RANK %d] %s stop", COMM_SESSION.getRank(), __PRETTY_FUNCTION__); -} - -// The function is used to change the state of a request to context_init from encoder_init for enc-dec model whose -// encoder is skipped. The encoder output is populated accordingly with input features given through model executor of -// decoder. -void Executor::Impl::prepRequestsForEncoderSkip(RequestList& activeRequests) -{ - - for (auto& req : activeRequests) - { - - if (req->isEncoderInitState() && req->getEncoderInputFeatures()) - { - TLLM_LOG_INFO("Changing state of request and setting encoder output to skip encoder run"); - req->setState(batch_manager::LlmRequestState::kCONTEXT_INIT); - req->setEncoderOutput(req->getEncoderInputFeatures()); - } - } -} - -void Executor::Impl::finishTimedOutRequests(RequestList const& activeRequests) -{ - if (mIsLeader) - { - for (auto const& request : activeRequests) - { - if (request->isTimedOut() && !request->isFinished()) - { - // workaround to cancelRequest since it throws an error if - // mCommMode == CommunicationMode::kORCHESTRATOR && !mIsOrchestrator - { - std::scoped_lock lck(mCancelReqMtx); - auto& selCancelledReqIds = mUsePipelineParallel ? mPipelineCancelledReqIds : mCancelledReqIds; - selCancelledReqIds.insert(request->mRequestId); - } - } - } - } -} - -void Executor::Impl::forwardAsync(RequestList& activeRequests) -{ - try - { - TLLM_LOG_DEBUG("num active requests in scope: %d", activeRequests.size()); - - if (mDynamicBatchTuner) - { - auto const averageInputLength = static_cast(mDynamicBatchTuner->getAverageInputLength()); - auto const averageOutputLength = static_cast(mDynamicBatchTuner->getAverageOutputLength()); - auto const maxCapacityBatchSize = mModel->getMaxCapacityBatchSize(averageInputLength, averageOutputLength); - - if (mDynamicBatchTuner->isBatchSizeTuningEnabled()) - { - auto runtimeBatchSize = mDynamicBatchTuner->getRuntimeBatchSize(maxCapacityBatchSize); - mModel->setRuntimeBatchSize(runtimeBatchSize); - } - - if (mDynamicBatchTuner->isMaxNumTokensTuningEnabled()) - { - auto runtimeBatchSize = mModel->getRuntimeBatchSize(); - auto runtimeMaxNumTokens = mDynamicBatchTuner->getRuntimeMaxNumTokens(runtimeBatchSize); - mModel->setRuntimeMaxNumTokens(runtimeMaxNumTokens); - } - } - - if (mEncoderModel) - { - mEncoderModel->forwardAsync(activeRequests); - auto const& encoderStream = *(mEncoderModel->getRuntimeStreamPtr()); - auto const& decoderStream = *(mModel->getRuntimeStreamPtr()); - runtime::CudaEvent encoderFinished; - encoderStream.record(encoderFinished); - decoderStream.wait(encoderFinished); - } - - if (!mEncoderModel) - { - prepRequestsForEncoderSkip(activeRequests); - } - - mModel->forwardAsync(activeRequests); - } - catch (std::exception const& e) - { - std::string err = std::string("Encountered an error in forwardAsync function: ") + e.what(); - terminateActiveRequests(activeRequests, err); - } -} - -IterationStats Executor::Impl::getCurrentIterationStats(RequestList const& activeRequests, double iterLatencyMS, - SizeType32 numNewActiveRequests, double newActiveRequestsQueueLatencyMS, SizeType32 numCompletedRequests) -{ - IterationStats stats; - // Timestamp - stats.timestamp = tensorrt_llm::common::getCurrentTimestamp(); - stats.numNewActiveRequests = numNewActiveRequests; - stats.iterLatencyMS = iterLatencyMS; - stats.newActiveRequestsQueueLatencyMS = newActiveRequestsQueueLatencyMS; - // Active request count - stats.numActiveRequests = static_cast(activeRequests.size()); - // Queued request count - { - std::scoped_lock lck(mQueuedReqMtx); - stats.numQueuedRequests = static_cast(mQueuedRequests.size()); - } - stats.numCompletedRequests = numCompletedRequests; - // Max number of requests - stats.maxNumActiveRequests = mMaxNumActiveRequests; - // Runtime memory allocation statistics - auto const& memoryCounters = runtime::MemoryCounters::getInstance(); - stats.gpuMemUsage = memoryCounters.getGpu(); - stats.cpuMemUsage = memoryCounters.getCpu(); - stats.pinnedMemUsage = memoryCounters.getPinned(); - - // Model specific stats - mModel->getCurrentIterationStats(stats); - return stats; -} - -RequestStatsPerIteration Executor::Impl::getCurrentRequestStats( - RequestList const& activeRequests, RequestList const& finishedRequests) -{ - std::vector requestStatsVec; - - auto includeDisServingStats = [](LlmRequestPtr const& request, tensorrt_llm::executor::RequestStats& requestStats) - { - auto requestType = request->getLlmRequestType(); - if (requestType == batch_manager::LlmRequestType::LLMREQUEST_TYPE_CONTEXT_ONLY - || requestType == batch_manager::LlmRequestType::LLMREQUEST_TYPE_GENERATION_ONLY) - { - requestStats.disServingStats - = executor::DisServingRequestStats{request->getKvCacheTransferTimeMS(), request->getKvCacheSize()}; - } - }; - - for (auto const& request : activeRequests) - { - RequestStats requestStats; - requestStats.id = request->mRequestId; - requestStats.stage = request->getRequestStage(); - requestStats.contextPrefillPosition = request->getContextCurrentPosition(); - requestStats.numGeneratedTokens = request->getMaxBeamNumTokens() - request->getOrigPromptLen(); - requestStats.avgNumDecodedTokensPerIter = request->getAvgDecodedTokensPerIter(); - includeDisServingStats(request, requestStats); - requestStats.allocTotalBlocksPerRequest = request->getAllocTotalBlocksPerRequest(); - requestStats.allocNewBlocksPerRequest = request->getAllocNewBlocksPerRequest(); - requestStats.reusedBlocksPerRequest = request->getReusedBlocksPerRequest(); - requestStats.missedBlocksPerRequest = request->getMissedBlocksPerRequest(); - requestStats.kvCacheHitRatePerRequest = request->getKVCacheHitRatePerRequest(); - requestStatsVec.emplace_back(requestStats); - } - - { - std::unique_lock lck(mQueuedReqMtx); - for (auto const& request : mQueuedRequests) - { - // Still waiting for the first scheduling - RequestStats requestStats; - requestStats.id = static_cast(request.id); - requestStats.stage = executor::RequestStage::kQUEUED; - requestStats.contextPrefillPosition = 0; - requestStats.numGeneratedTokens = 0; - requestStats.avgNumDecodedTokensPerIter = 0; - requestStats.allocTotalBlocksPerRequest = 0; - requestStats.allocNewBlocksPerRequest = 0; - requestStats.reusedBlocksPerRequest = 0; - requestStats.missedBlocksPerRequest = 0; - requestStats.kvCacheHitRatePerRequest = 0; - requestStatsVec.emplace_back(requestStats); - } - } - - for (auto const& request : finishedRequests) - { - // Still waiting for the first scheduling - RequestStats requestStats; - requestStats.id = static_cast(request->mRequestId); - requestStats.stage = executor::RequestStage::kGENERATION_COMPLETE; - requestStats.contextPrefillPosition = request->getContextCurrentPosition(); - requestStats.numGeneratedTokens = request->getMaxBeamNumTokens() - request->getOrigPromptLen(); - requestStats.avgNumDecodedTokensPerIter = request->getAvgDecodedTokensPerIter(); - includeDisServingStats(request, requestStats); - requestStats.allocTotalBlocksPerRequest = request->getAllocTotalBlocksPerRequest(); - requestStats.allocNewBlocksPerRequest = request->getAllocNewBlocksPerRequest(); - requestStats.reusedBlocksPerRequest = request->getReusedBlocksPerRequest(); - requestStats.missedBlocksPerRequest = request->getMissedBlocksPerRequest(); - requestStats.kvCacheHitRatePerRequest = request->getKVCacheHitRatePerRequest(); - requestStatsVec.emplace_back(requestStats); - } - - RequestStatsPerIteration stats{0, std::move(requestStatsVec)}; - - // Model specific stats - mModel->getCurrentRequestStats(stats); - return stats; -} - -void Executor::Impl::appendCurrentIterStats(IterationStats&& currentIterStats) -{ - std::scoped_lock lck(mIterStatsMtx); - if (statsBufferIsBounded(mIterStatsMaxIterations)) - { - auto const maxIterStats = static_cast(mIterStatsMaxIterations); - if (mIterationStats.size() >= maxIterStats) - { - mIterationStats.pop_front(); - } - } - mIterationStats.emplace_back(std::move(currentIterStats)); -} - -void Executor::Impl::appendMultipleIterStats(std::vector&& currentIterStatsVec) -{ - std::scoped_lock lck(mIterStatsMtx); - mIterationStats.insert(mIterationStats.end(), std::make_move_iterator(currentIterStatsVec.begin()), - std::make_move_iterator(currentIterStatsVec.end())); - if (statsBufferIsBounded(mIterStatsMaxIterations)) - { - auto const maxIterStats = static_cast(mIterStatsMaxIterations); - while (mIterationStats.size() > maxIterStats) - { - mIterationStats.pop_front(); - } - } -} - -void Executor::Impl::updateIterationStats(RequestList const& activeRequests, double iterLatencyMS, - SizeType32 numNewActiveRequests, double newActiveRequestsQueueLatencyMS, SizeType32 numCompletedRequests, - bool flushToOrchestrator) -{ - NVTX3_SCOPED_RANGE(updateIterationStats); - if (statsBufferIsEnabled(mIterStatsMaxIterations) && mIsLeader) - { - auto currentIterStats = getCurrentIterationStats( - activeRequests, iterLatencyMS, numNewActiveRequests, newActiveRequestsQueueLatencyMS, numCompletedRequests); - // Send the stats to the orchestrator - if (mCommMode == CommunicationMode::kORCHESTRATOR) - { - bool hasSchedThisIter = (currentIterStats.inflightBatchingStats - && currentIterStats.inflightBatchingStats->numScheduledRequests > 0) - || (currentIterStats.staticBatchingStats - && currentIterStats.staticBatchingStats->numScheduledRequests > 0); - appendCurrentIterStats(std::move(currentIterStats)); - if (hasSchedThisIter || flushToOrchestrator) - { - std::deque iterStatsQueue; - { - std::scoped_lock lck(mIterStatsMtx); - iterStatsQueue = std::exchange(mIterationStats, {}); - } - MpiMessage message(MpiId::ITER_STATS); - std::vector iterStates( - std::make_move_iterator(iterStatsQueue.begin()), std::make_move_iterator(iterStatsQueue.end())); - message.data = IterStatsData{std::move(iterStates)}; - mSendQueue.push(std::move(message)); - } - } - else - { - // Add current iteration stats - appendCurrentIterStats(std::move(currentIterStats)); - } - } -} - -void Executor::Impl::appendCurrentRequestStats(RequestStatsPerIteration&& currentRequestStats) -{ - std::scoped_lock lck(mRequestStatsMtx); - if (statsBufferIsBounded(mRequestStatsMaxIterations)) - { - auto const maxRequestStats = static_cast(mRequestStatsMaxIterations); - if (mRequestStats.size() >= maxRequestStats) - { - mRequestStats.pop_front(); - } - } - mRequestStats.emplace_back(std::move(currentRequestStats)); -} - -void Executor::Impl::appendMultipleRequestStats(std::vector&& currentRequestStatsVec) -{ - std::scoped_lock lck(mRequestStatsMtx); - mRequestStats.insert(mRequestStats.end(), std::make_move_iterator(currentRequestStatsVec.begin()), - std::make_move_iterator(currentRequestStatsVec.end())); - if (statsBufferIsBounded(mRequestStatsMaxIterations)) - { - auto const maxRequestStats = static_cast(mRequestStatsMaxIterations); - while (mRequestStats.size() > maxRequestStats) - { - mRequestStats.pop_front(); - } - } -} - -void Executor::Impl::updateRequestStats( - RequestList const& activeRequests, RequestList const& finishedRequests, bool flushToOrchestrator) -{ - NVTX3_SCOPED_RANGE(updateRequestStats); - if (statsBufferIsEnabled(mRequestStatsMaxIterations) && mIsLeader) - { - // Add current iteration request stats - auto currentRequestStats = getCurrentRequestStats(activeRequests, finishedRequests); - // Send the stats to the orchestrator - if (mCommMode == CommunicationMode::kORCHESTRATOR) - { - bool hasScheduledReqs = false; - if (!flushToOrchestrator) - { - size_t activeSize = activeRequests.size(); - TLLM_CHECK_WITH_INFO(currentRequestStats.requestStats.size() >= activeSize, - "currentRequestStats num is %ld should >= activeRequest num:%zu", - currentRequestStats.requestStats.size(), activeSize); - hasScheduledReqs = std::any_of(currentRequestStats.requestStats.begin(), - currentRequestStats.requestStats.begin() + static_cast(activeSize), - [](RequestStats const& requestStat) { return requestStat.scheduled; }); - } - appendCurrentRequestStats(std::move(currentRequestStats)); - if (hasScheduledReqs || flushToOrchestrator) - { - std::deque requestStatsQueue; - { - std::scoped_lock lck(mRequestStatsMtx); - requestStatsQueue = std::exchange(mRequestStats, {}); - } - std::vector requestIterStates( - std::make_move_iterator(requestStatsQueue.begin()), - std::make_move_iterator(requestStatsQueue.end())); - MpiMessage message(MpiId::REQUEST_ITER_STATS); - message.data = RequestStatsPerIterationData{std::move(requestIterStates)}; - mSendQueue.push(std::move(message)); - } - } - else - { - // Add current iteration stats - appendCurrentRequestStats(std::move(currentRequestStats)); - } - } -} - -void Executor::Impl::appendCurrentDebugTensors() -{ - if (mDebugTensorsMaxIterations > 0) - { - std::scoped_lock lck(mDebugTensorsMtx); - if (mDebugTensors.size() >= mDebugTensorsMaxIterations) - { - mDebugTensors.pop_front(); - } - mDebugTensors.emplace_back(mModel->getCurrentDebugTensors()); - } -} - -void Executor::Impl::terminateCancelledRequests(RequestList& activeRequests) -{ - NVTX3_SCOPED_RANGE(terminateCancelledRequests); - auto const& worldConfig = mModel->getWorldConfig(); - auto const broadcastCancelledRequests = [this, &activeRequests, &worldConfig] - { - auto const& commSession = COMM_SESSION; - - if (worldConfig.isPipelineParallel()) - { - mCancelledRequestsWaitThread->waitStop(); - } - - if (commSession.getSize() > 1 && !activeRequests.empty()) - { - if (mIsPipelineLeader) - { - if (worldConfig.isPipelineParallel()) - { - auto const peer = worldConfig.getPipelineParallelism() - 1; - bool shouldExit = false; - mCommPipelineParallel->send( - &shouldExit, 1, mpi::MpiType::kBOOL, peer, mpi::MpiTag::kExecutorShouldExit); - auto pipelineCancelledReqIds - = CancelledRequestsAsyncSend::cancelledRequestsRecv(mCommPipelineParallel, peer); - mCancelledReqIds.insert(pipelineCancelledReqIds.begin(), pipelineCancelledReqIds.end()); - } - - auto numCancelledRequests = static_cast(mCancelledReqIds.size()); - if (worldConfig.isTensorParallel()) - { - mCommTensorParallel->bcastValue(numCancelledRequests, 0); - if (numCancelledRequests > 0) - { - std::vector cancelledReqIdsVec(mCancelledReqIds.begin(), mCancelledReqIds.end()); - mCommTensorParallel->bcast( - cancelledReqIdsVec.data(), cancelledReqIdsVec.size(), mpi::MpiType::kUINT64, 0); - } - } - if (worldConfig.isContextParallel()) - { - mCommContextParallel->bcastValue(numCancelledRequests, 0); - if (numCancelledRequests > 0) - { - std::vector cancelledReqIdsVec(mCancelledReqIds.begin(), mCancelledReqIds.end()); - mCommContextParallel->bcast( - cancelledReqIdsVec.data(), cancelledReqIdsVec.size(), mpi::MpiType::kUINT64, 0); - } - } - } - // If not leader - else - { - if (worldConfig.isFirstPipelineParallelRank()) - { - int64_t numCancelledRequests = 0; - mCommTensorParallel->bcastValue(numCancelledRequests, 0); - mCommContextParallel->bcastValue(numCancelledRequests, 0); - if (numCancelledRequests > 0) - { - std::vector cancelledReqIdsVec(numCancelledRequests); - mCommTensorParallel->bcast( - cancelledReqIdsVec.data(), cancelledReqIdsVec.size(), mpi::MpiType::kUINT64, 0); - mCommContextParallel->bcast( - cancelledReqIdsVec.data(), cancelledReqIdsVec.size(), mpi::MpiType::kUINT64, 0); - mCancelledReqIds - = std::unordered_set(cancelledReqIdsVec.begin(), cancelledReqIdsVec.end()); - } - } - else - { - auto const peer = worldConfig.getPipelineParallelRank() - 1; - mCancelledReqIds = CancelledRequestsAsyncSend::cancelledRequestsRecv(mCommPipelineParallel, peer); - } - } - if (!worldConfig.isLastPipelineParallelRank()) - { - auto const peer = worldConfig.getPipelineParallelRank() + 1; - mCancelledRequestsAsyncSndHdl - = std::make_unique(mCommPipelineParallel, mCancelledReqIds, peer); - mCancelledRequestsWaitThread->notifyStart(); - } - } - }; - - std::unique_lock lck{mCancelReqMtx, std::defer_lock}; - if (!worldConfig.isPipelineParallel()) - { - lck.lock(); - } - - broadcastCancelledRequests(); - - if (!mCancelledReqIds.empty()) - { - // Loop over active requests and terminate those that have been cancelled - std::unordered_set terminatedReqIds; - for (auto& req : activeRequests) - { - auto reqId = req->isChild() ? req->getParentRequestId() : req->mRequestId; - if (mCancelledReqIds.find(reqId) != mCancelledReqIds.end()) - { - auto finishReason = req->isTimedOut() ? FinishReason::kTIMED_OUT : FinishReason::kCANCELLED; - mModel->terminateRequestSync(req, finishReason); - // Parent and child requests share the same request id. - // Mark it terminated first and remove from the set later. - terminatedReqIds.insert(reqId); - } - } - - for (auto const& reqId : terminatedReqIds) - { - mCancelledReqIds.erase(reqId); - } - } -} - -void Executor::Impl::terminateContextFinishedRequests(InTransList& inTransmissionRequests) -{ - NVTX3_SCOPED_RANGE(terminateContextFinishedRequests); - for (auto it = inTransmissionRequests.begin(); it != inTransmissionRequests.end();) - { - auto& item = *it; - auto req = item.request; - if (req->isDisaggContextCompleteState()) - { - // If pinnedBlockIds were tracked, unpin them. Otherwise, just terminate. - auto kvMgr = mModel->getKVCacheManager(); - if (kvMgr && !item.pinnedBlockIds.empty()) - { - kvMgr->unpinBlocksById(item.pinnedBlockIds); - } - else - { - mModel->terminateRequest(req); - } - it = inTransmissionRequests.erase(it); - } - else - { - ++it; - } - } -} - -void Executor::Impl::appendNewResponses(std::vector&& newResponses) -{ - { - std::scoped_lock lck(mResponsesMtx); - for (auto& response : newResponses) - { - mResponses[response.getRequestId()].emplace_back(std::move(response)); - } - } - mResponsesCv.notify_all(); -} - -Executor::Impl::RequestList Executor::Impl::populateNewResponses( - RequestList& activeRequests, InTransList& inTransmissionRequests, std::vector& newResponses) -{ - NVTX3_SCOPED_RANGE(populateNewResponses); - RequestList finishedRequests; - for (auto it = activeRequests.begin(); it != activeRequests.end();) - { - auto const& llmReq = (*it); - bool const requestDone = llmReq->isFinished(); - // Only leader should store responses - if (mIsLeader) - { - auto response = llmReq->createResponse(mModel->hasSpeculativeDecodingFastLogits(), mWorldRank); - if (response) - { - newResponses.emplace_back(std::move(response.value())); - } - } - // Remove from active requests if last response has been generated - if (requestDone) - { - // move the in transmission requests to another tracker - if (llmReq->isDisaggContextTransmissionState()) - { - std::vector pinnedBlockIds{}; - auto kvMgr = mModel->getKVCacheManager(); - if (kvMgr && kvMgr->isEnableBlockReuse() && !kvMgr->getBlockManager().isVariableWindow()) - { - pinnedBlockIds = kvMgr->storeBlocksForReuse(llmReq->mRequestId, llmReq, /*pinBlocks=*/true); - mModel->terminateRequest(llmReq); - } - inTransmissionRequests.push_back(InTransmissionItem{*it, pinnedBlockIds}); - } - finishedRequests.push_back(*it); - it = activeRequests.erase(it); - } - else - { - ++it; - } - } - return finishedRequests; -} - -void Executor::Impl::executionLoop() -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - tensorrt_llm::common::setThreadName("executionLoop"); - - auto const& worldConfig = mModel->getWorldConfig(); - TLLM_CUDA_CHECK(cudaSetDevice(worldConfig.getDevice())); - - auto const [profileIterIdxs, stopIterIdxs] = tensorrt_llm::common::populateIterationIndexes( - kPROFILE_START_STOP_ENV_VAR_NAME, kLEGACY_PROFILE_START_STOP_ENV_VAR_NAME); - - SizeType32 numNewActiveRequests{0}; - std::chrono::time_point iterStart; - std::chrono::time_point iterEnd; - bool firstIteration{true}; - RequestList activeRequests; - InTransList inTransmissionRequests; - std::vector newResponses; - while (!mShutdown || !activeRequests.empty()) - { - double iterLatencyMS{0.0}; - double newActiveRequestsQueueLatencyMS{0.0}; - bool reportFinishedRequests = true; - RequestList finishedRequests; - if (!activeRequests.empty()) - { - finishTimedOutRequests(activeRequests); - terminateCancelledRequests(activeRequests); - forwardSync(activeRequests); - finishedRequests = populateNewResponses(activeRequests, inTransmissionRequests, newResponses); - cleanupDynamicLogitsPostProcessors(finishedRequests); - auto const iterCounter = mModel->getIterCounter(); - auto const stopIter = !stopIterIdxs.empty() && (stopIterIdxs.count(iterCounter - 1) > 0); - if (stopIter) - { - cudaProfilerStop(); - } - - // When there are no active or inflight requests, we need to update the stats before calling - // fetchNewRequests to make sure that the stats are reported accurately. - if (activeRequests.empty() && (!firstIteration)) - { - mModel->resetIterationStats(); - updateIterationStats(activeRequests, iterLatencyMS, numNewActiveRequests, - newActiveRequestsQueueLatencyMS, static_cast(finishedRequests.size()), true); - updateRequestStats(activeRequests, finishedRequests, true); - reportFinishedRequests = false; - } - if (!newResponses.empty()) - { - enqueueNewResponses(std::move(newResponses)); - newResponses.clear(); - } - iterEnd = std::chrono::steady_clock::now(); - iterLatencyMS = std::chrono::duration(iterEnd - iterStart).count(); - } - - if (!inTransmissionRequests.empty()) - { - terminateContextFinishedRequests(inTransmissionRequests); - } - - if (!mShutdown) - { - auto const iterCounter = mModel->getIterCounter(); - auto const profileIter = !profileIterIdxs.empty() && (profileIterIdxs.count(iterCounter) > 0); - if (profileIter) - { - cudaProfilerStart(); - } - iterStart = std::chrono::steady_clock::now(); - std::optional lowestPriority = std::nullopt; - if (!activeRequests.empty()) - { - lowestPriority = activeRequests.back()->priority(); - } - - auto [newRequests, newActiveRequestsQueueLatency] - = fetchNewRequests(static_cast(activeRequests.size()), lowestPriority); - newActiveRequestsQueueLatencyMS = newActiveRequestsQueueLatency; - numNewActiveRequests = newRequests.size(); - - if (firstIteration) - { - firstIteration = false; - } - - for (auto const& newRequest : newRequests) - { - insertRequestInOrder(activeRequests, newRequest); - } - - // Update dynamic tuning stats - if (mDynamicBatchTuner) - { - for (auto const& req : activeRequests) - { - auto const inputLength = req->mPromptLen; - auto const outputLength = req->mMaxNewTokens; - mDynamicBatchTuner->updateStats(inputLength, outputLength); - } - } - } - if (!activeRequests.empty()) - { - forwardAsync(activeRequests); - updateIterationStats(activeRequests, iterLatencyMS, numNewActiveRequests, newActiveRequestsQueueLatencyMS, - static_cast(finishedRequests.size()), false); - // Finished requests were reported once. Avoid reporting it twice. - if (reportFinishedRequests) - { - updateRequestStats(activeRequests, finishedRequests, false); - } - else - { - updateRequestStats(activeRequests, {}, false); - } - appendCurrentDebugTensors(); - } - } - - if (mCancelledRequestsWaitThread) - { - mCancelledRequestsWaitThread.reset(nullptr); - } - if (mRequestWithIdWaitThread) - { - mRequestWithIdWaitThread.reset(nullptr); - } - if (worldConfig.isPipelineParallel() && mIsPipelineLeader) - { - auto const peer = worldConfig.getPipelineParallelism() - 1; - int64_t numActiveRequests = -1; - mCommPipelineParallel->send( - &numActiveRequests, 1, mpi::MpiType::kINT64, peer, mpi::MpiTag::kExecutorNumActiveRequests); - bool shouldExit = true; - mCommPipelineParallel->send(&shouldExit, 1, mpi::MpiType::kBOOL, peer, mpi::MpiTag::kExecutorShouldExit); - } - if (mRequestWithIdLeaderThread) - { - mRequestWithIdLeaderThread->join(); - mRequestWithIdLeaderThread.reset(nullptr); - } - if (mCancelledRequestsLeaderThread) - { - mCancelledRequestsLeaderThread->join(); - mCancelledRequestsLeaderThread.reset(nullptr); - } - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void Executor::Impl::enqueueTerminateRequest() -{ - { - std::scoped_lock lck(mQueuedReqMtx); - Request dummyReq({1}, 1); - RequestWithId reqWithId{std::move(dummyReq), kTerminateReqId}; - mQueuedRequests.emplace_back(reqWithId); - } - mQueuedReqCv.notify_one(); -} - -void Executor::Impl::enqueueNewResponses(std::vector&& newResponses) -{ - TLLM_CHECK_WITH_INFO(mIsLeader, "Only leader should store responses"); - - if (mCommMode == CommunicationMode::kLEADER) - { - appendNewResponses(std::move(newResponses)); - } - else if (mCommMode == CommunicationMode::kORCHESTRATOR) - { - MpiMessage message(MpiId::RESPONSE); - message.data = ResponseData{std::move(newResponses)}; - mSendQueue.push(std::move(message)); - } -} - -// Orchestrator thread sending new requests to leader of the model -void Executor::Impl::orchSendReqThread() -{ - tensorrt_llm::common::setThreadName("orchSendReq"); - - while (true) - { - auto message = mSendQueue.pop(); - - if (message.id == MpiId::TERMINATION) - { - mOrchLeaderComm->send(&message.id, 1, mpi::MpiType::kUINT64, mLeaderRank, mpi::MpiTag::kOrchestratorId); - TLLM_LOG_INFO("Orchestrator sendReq thread exiting"); - break; - } - if (message.id == MpiId::PENDING_REQUEST) - { - auto& reqWithIds = std::get(message.data); - auto packed = RequestWithId::serializeReqWithIds(reqWithIds.requests); - - TLLM_LOG_DEBUG("Orchestrator sendReq thread sending %d pending requests", reqWithIds.requests.size()); - // Temporary WAR to indicate to client that we cannot send the serialized request - // because it exceeds int32_t size limit. - // TODO: Should fix as part of https://jirasw.nvidia.com/browse/TRTLLM-708 - if (packed.size() > std::numeric_limits::max()) - { - for (auto const& reqWithId : reqWithIds.requests) - { - { - std::scoped_lock lck(mResponsesMtx); - mResponses[reqWithId.id].emplace_back(reqWithId.id, - "Request is too large, or you are enqueuing too many requests at once " - "to be sent via MPI_Send, please try to enqueue the request(s) again. " - "This issue will be resolved in a future version of TRT-LLM."); - } - mResponsesCv.notify_all(); - } - } - else - { - mOrchLeaderComm->send(&message.id, 1, mpi::MpiType::kUINT64, mLeaderRank, mpi::MpiTag::kOrchestratorId); - mOrchLeaderComm->send( - packed.data(), packed.size(), mpi::MpiType::kCHAR, mLeaderRank, mpi::MpiTag::kOrchestratorData); - } - } - else if (message.id == MpiId::CANCEL_REQUEST) - { - auto& data = std::get(message.data); - - mOrchLeaderComm->send(&message.id, 1, mpi::MpiType::kUINT64, mLeaderRank, mpi::MpiTag::kOrchestratorId); - mOrchLeaderComm->send( - data.ids.data(), data.ids.size(), mpi::MpiType::kUINT64, mLeaderRank, mpi::MpiTag::kOrchestratorData); - } - else - { - TLLM_THROW("Invalid message id"); - } - } -} - -// Leader thread receiving new requests from orchestrator -void Executor::Impl::leaderRecvReqThread() -{ - tensorrt_llm::common::setThreadName("leaderRecvReq"); - TLLM_CUDA_CHECK(cudaSetDevice(mDeviceId)); -#if ENABLE_MULTI_DEVICE - auto& selCancelledReqIds = mUsePipelineParallel ? mPipelineCancelledReqIds : mCancelledReqIds; - while (true) - { - if (mRecvPollPeriodMs > 0) - { - mOrchLeaderComm->recvPoll(mOrchRank, mpi::MpiTag::kOrchestratorId, mRecvPollPeriodMs); - } - - // Blocking is okay: terminate message is expected to arrive here - MPI_Message msg = nullptr; - MPI_Status status; - mOrchLeaderComm->mprobe(mOrchRank, mpi::MpiTag::kOrchestratorId, &msg, &status); - - int32_t count = 0; - MPICHECK(MPI_Get_count(&status, MPI_UINT64_T, &count)); // NOLINT - TLLM_CHECK(count == 1); - - MpiId mpiId{}; - MPICHECK(MPI_Mrecv(&mpiId, count, MPI_UINT64_T, &msg, &status)); // NOLINT - - // EXIT condition from receiving TERMINATE msg - if (mpiId == MpiId::TERMINATION) - { - // Enqueue a request to indicate to other ranks to terminate - enqueueTerminateRequest(); - - // Send message to orchestrator to indicate to terminate orch recv thread - mSendQueue.push(MpiMessage(mpiId)); - TLLM_LOG_INFO("Leader recvReq thread exiting"); - break; - } - if (mpiId == MpiId::PENDING_REQUEST) - { - mOrchLeaderComm->mprobe(mOrchRank, mpi::MpiTag::kOrchestratorData, &msg, &status); - MPICHECK(MPI_Get_count(&status, MPI_CHAR, &count)); // NOLINT - std::vector buffer(count); - MPICHECK(MPI_Mrecv(buffer.data(), count, MPI_CHAR, &msg, &status)); // NOLINT - - auto requestWithIds = RequestWithId::deserializeReqWithIds(buffer); - TLLM_LOG_DEBUG("Leader recvReq thread receiving %d pending requests", requestWithIds.size()); - { - std::scoped_lock lck(mQueuedReqMtx); - if (mMaxQueueSize) - { - auto const maxQueueSize = mMaxQueueSize.value(); - if (maxQueueSize > 0 && mQueuedRequests.size() >= static_cast(maxQueueSize)) - { - auto err = tensorrt_llm::common::fmtstr( - "Maximum queue size of %d has been reached, please try again later", maxQueueSize); - TLLM_LOG_ERROR("%s", err.c_str()); - std::vector responses; - responses.reserve(requestWithIds.size()); - for (auto const& reqWithId : requestWithIds) - { - responses.emplace_back(reqWithId.id, err); - } - enqueueNewResponses(std::move(responses)); - continue; - } - } - for (auto&& req : requestWithIds) - { - req.queuedStart = std::chrono::steady_clock::now(); - insertRequestInOrder(mQueuedRequests, std::move(req)); - } - } - mQueuedReqCv.notify_one(); - } - else if (mpiId == MpiId::CANCEL_REQUEST) - { - // Prepare receiving data - mOrchLeaderComm->mprobe(mOrchRank, mpi::MpiTag::kOrchestratorData, &msg, &status); - MPICHECK(MPI_Get_count(&status, MPI_UINT64_T, &count)); // NOLINT - std::vector cancelledReqIds(count); - MPICHECK(MPI_Mrecv(cancelledReqIds.data(), count, MPI_UINT64_T, &msg, &status)); // NOLINT - - std::scoped_lock lck(mCancelReqMtx); - selCancelledReqIds.insert(cancelledReqIds.begin(), cancelledReqIds.end()); - } - else - { - TLLM_THROW("Invalid message id"); - } - } -#endif // ENABLE_MULTI_DEVICE -} - -// Leader thread sending responses to orchestrator -void Executor::Impl::leaderSendThread(MpiMessageQueue& sendQueue, mpi::MpiTag idTag, mpi::MpiTag dataTag) -{ - tensorrt_llm::common::setThreadName("leaderSend"); - TLLM_CUDA_CHECK(cudaSetDevice(mDeviceId)); - -#if ENABLE_MULTI_DEVICE - while (true) - { - auto message = sendQueue.pop(); - - if (message.id == MpiId::TERMINATION) - { - mOrchLeaderComm->send(&message.id, 1, mpi::MpiType::kUINT64, mOrchRank, idTag); - TLLM_LOG_INFO("Leader sendThread exiting"); - break; - } - if (message.id == MpiId::RESPONSE || message.id == MpiId::ITER_STATS - || message.id == MpiId ::REQUEST_ITER_STATS) - { - std::vector buffer; - if (message.id == MpiId::RESPONSE) - { - auto& responseData = std::get(message.data); - TLLM_LOG_DEBUG("Leader sendResp thread sending %d responses", responseData.responses.size()); - buffer = Serialization::serialize(responseData.responses); - } - else if (message.id == MpiId::ITER_STATS) - { - auto& iterStatsData = std::get(message.data); - TLLM_LOG_DEBUG("Leader sendResp thread sending iter stats"); - buffer = Serialization::serialize(iterStatsData.iterStatsVec); - } - else if (message.id == MpiId::REQUEST_ITER_STATS) - { - auto& requestIterStatsData = std::get(message.data); - TLLM_LOG_DEBUG("Leader sendResp thread sending iter request stats"); - buffer = Serialization::serialize(requestIterStatsData.requestStatsPerIterationVec); - } - mOrchLeaderComm->send(&message.id, 1, mpi::MpiType::kUINT64, mOrchRank, idTag); - mOrchLeaderComm->send(buffer.data(), buffer.size(), mpi::MpiType::kCHAR, mOrchRank, dataTag); - } - else - { - TLLM_THROW("Invalid message id"); - } - } -#endif // ENABLE_MULTI_DEVICE -} - -void Executor::Impl::orchRecvThread(mpi::MpiTag idTag, mpi::MpiTag dataTag) -{ - tensorrt_llm::common::setThreadName("orchRecv"); - -#if ENABLE_MULTI_DEVICE - while (true) - { - if (mRecvPollPeriodMs > 0) - { - mOrchLeaderComm->recvPoll(mOrchRank, mpi::MpiTag::kOrchestratorId, mRecvPollPeriodMs); - } - - MPI_Message msg = nullptr; - MPI_Status status; - mOrchLeaderComm->mprobe(mLeaderRank, idTag, &msg, &status); - - int32_t count = 0; - MPICHECK(MPI_Get_count(&status, MPI_UINT64_T, &count)); // NOLINT - TLLM_CHECK(count == 1); - - MpiId mpiId{}; - MPICHECK(MPI_Mrecv(&mpiId, count, MPI_UINT64_T, &msg, &status)); // NOLINT - - if (mpiId == MpiId::TERMINATION) - { - TLLM_LOG_INFO("Orchestrator recv thread exiting"); - break; - } - if (mpiId == MpiId::RESPONSE || mpiId == MpiId::ITER_STATS || mpiId == MpiId::REQUEST_ITER_STATS) - { - mOrchLeaderComm->mprobe(mLeaderRank, dataTag, &msg, &status); - MPICHECK(MPI_Get_count(&status, MPI_CHAR, &count)); // NOLINT - - std::vector buffer(count); - MPICHECK(MPI_Mrecv(buffer.data(), count, MPI_CHAR, &msg, &status)); // NOLINT - - if (mpiId == MpiId::RESPONSE) - { - auto newResponses = Serialization::deserializeResponses(buffer); - TLLM_LOG_DEBUG("Orchestrator recv thread receiving %d responses", newResponses.size()); - appendNewResponses(std::move(newResponses)); - } - else if (mpiId == MpiId::ITER_STATS) - { - appendMultipleIterStats(Serialization::deserializeIterationStatsVec(buffer)); - } - else if (mpiId == MpiId::REQUEST_ITER_STATS) - { - appendMultipleRequestStats(Serialization::deserializeRequestStatsPerIterationVec(buffer)); - } - } - else - { - TLLM_THROW("Invalid message id"); - } - } -#endif // ENABLE_MULTI_DEVICE -} - -Executor::Impl::LlmRequestLogitsPostProcessor Executor::Impl::getLogitsPostProcessor(std::string const& name) -{ - auto const postProcIt = mLogitsPostProcessorMap.find(name); - TLLM_CHECK_WITH_INFO( - postProcIt != mLogitsPostProcessorMap.end(), "LogitsPostProcessor %s not found.", name.c_str()); - auto executorLogitsPostProcessor = postProcIt->second; - return [executorLogitsPostProcessor](IdType reqId, RtTensorPtr& logits, BeamTokens const& beamTokens, - CudaStreamPtr const& cudaStreamPtr, std::optional clientId) - { - auto logitsTensor = executor::detail::ofITensor(logits); - executorLogitsPostProcessor(reqId, logitsTensor, beamTokens, cudaStreamPtr, clientId); - }; -} - -void Executor::Impl::setupDynamicLogitsPostProcessors(std::vector& newReqWithIds) -{ - for (auto& reqWithId : newReqWithIds) - { - auto logitsPostProcessor = reqWithId.req.getLogitsPostProcessor(); - if (logitsPostProcessor) - { - std::string const name = Request::kDynamicPostProcessorNamePrefix + std::to_string(reqWithId.id); - mLogitsPostProcessorMap[name] = logitsPostProcessor.value(); - reqWithId.req.setLogitsPostProcessor(std::nullopt); - reqWithId.req.setLogitsPostProcessorName(name); - } - } -} - -void Executor::Impl::cleanupDynamicLogitsPostProcessors(RequestList const& finishedRequests) -{ - for (auto& req : finishedRequests) - { - std::string const name = Request::kDynamicPostProcessorNamePrefix + std::to_string(req->mRequestId); - auto const postProcIt = mLogitsPostProcessorMap.find(name); - if (postProcIt != mLogitsPostProcessorMap.end()) - { - mLogitsPostProcessorMap.erase(name); - } - } -} - -void Executor::Impl::addTerminatedReqId(std::vector const& responses, IdType const& reqId) -{ - for (auto const& response : responses) - { - if (response.hasError() || (!response.hasError() && response.getResult().isFinal)) - { - mTerminatedReqIds.insert(reqId); - if (mChildReqIdsMap.find(reqId) != mChildReqIdsMap.end()) - { - for (auto childReqId : mChildReqIdsMap.at(reqId)) - { - mTerminatedReqIds.insert(childReqId); - } - mChildReqIdsMap.erase(reqId); - } - } - } -} - -void Executor::Impl::checkParallelApiUsage(std::string const& methodName) const -{ - // If leader mode, and not leader, throw error - if (mCommMode == CommunicationMode::kLEADER && !mIsLeader) - { - // Non-leader are not expected to call cancelRequest - TLLM_THROW("With LEADER communication mode, only leader rank is expected to call %s", methodName.c_str()); - } - if (mCommMode == CommunicationMode::kORCHESTRATOR && !mIsOrchestrator) - { - TLLM_THROW( - "With ORCHESTRATOR communication mode, only orchestrator rank is expected to call %s", methodName.c_str()); - } -} - -} // namespace tensorrt_llm::executor diff --git a/cpp/tensorrt_llm/executor/executorImpl.h b/cpp/tensorrt_llm/executor/executorImpl.h deleted file mode 100644 index f812b55a3fa0..000000000000 --- a/cpp/tensorrt_llm/executor/executorImpl.h +++ /dev/null @@ -1,385 +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. - */ - -#pragma once - -#include "tensorrt_llm/batch_manager/llmRequest.h" -#include "tensorrt_llm/common/arrayView.h" -#include "tensorrt_llm/executor/dynamicBatchTuner.h" -#include "tensorrt_llm/executor/executor.h" -#include "tensorrt_llm/executor/intervalSet.h" -#include "tensorrt_llm/executor/model.h" -#include "tensorrt_llm/executor/orchestratorUtils.h" -#include "tensorrt_llm/executor/requestWithId.h" -#include "tensorrt_llm/executor/types.h" -#include "tensorrt_llm/runtime/gptJsonConfig.h" -#include "tensorrt_llm/runtime/modelConfig.h" -#include "tensorrt_llm/runtime/rawEngine.h" -#include "tensorrt_llm/runtime/utils/mpiUtils.h" -#include "tensorrt_llm/runtime/worldConfig.h" - -#include -#include -#include -#include -#include -#include -#include -#include -#include - -namespace tensorrt_llm::executor -{ - -class RequestWithIdAsyncSend; -class CancelledRequestsAsyncSend; - -class MpiMessageQueue -{ -public: - void push(MpiMessage&& message) - { - std::lock_guard const lock(mMutex); - mQueue.push(std::move(message)); - mCv.notify_one(); - } - - MpiMessage pop() - { - std::unique_lock lock(mMutex); - mCv.wait(lock, [this] { return !mQueue.empty(); }); - MpiMessage message = std::move(mQueue.front()); - mQueue.pop(); - return message; - } - -private: - std::queue mQueue; - std::mutex mMutex; - std::condition_variable mCv; -}; - -class Executor::Impl - -{ - using LlmRequestPtr = std::shared_ptr; - using RequestList = std::list; - - // When block reuse is enabled for context worker for disaggregated serving, - // we need to store the pinned block ids so that we can unpin them when - // the request is finished. - struct InTransmissionItem - { - LlmRequestPtr request; - std::vector pinnedBlockIds; - }; - - using InTransList = std::list; - -public: - Impl(std::filesystem::path const& modelPath, std::optional const& encoderModelPath, - [[maybe_unused]] ModelType modelType, ExecutorConfig const& executorConfig); - - Impl(BufferView const& engineBufferView, std::string const& jsonConfigStr, - std::optional const& encoderEngineBufferView, - std::optional const& encoderJsonConfigStr, [[maybe_unused]] ModelType modelType, - ExecutorConfig const& executorConfig, std::optional> const& managedWeightsOpt); - - Impl(std::shared_ptr model, std::optional> encoderModel, - ExecutorConfig const& executorConfig); - - ~Impl(); - - Impl(Impl const& executor) = delete; - Impl& operator=(Impl const& executor) = delete; - Impl(Impl&&) = delete; - Impl& operator=(Impl&&) = delete; - - IdType enqueueRequest(Request const& request); - - std::vector enqueueRequests(std::vector const& requests); - - std::vector enqueueRequests(common::ArrayView const& requests); - - std::vector awaitResponses(std::optional const& timeout = std::nullopt); - - std::vector awaitResponses( - IdType const& reqId, std::optional const& optTimeout = std::nullopt); - - std::vector> awaitResponses( - std::vector const& requestIds, std::optional const& timeout); - - SizeType32 getNumResponsesReady(std::optional const& optId = std::nullopt) const; - - void cancelRequest(IdType requestId); - - void shutdown(); - - std::deque getLatestIterationStats(); - std::deque getLatestRequestStats(); - std::deque getLatestDebugTensors(); - - bool canEnqueueRequests() const; - - bool isParticipant() const; - - std::optional> getKVCacheEventManager() const; - -private: - using RtTensorPtr = runtime::ITensor::SharedPtr; - using CudaStreamPtr = runtime::BufferManager::CudaStreamPtr; - using LlmRequestLogitsPostProcessor - = std::function)>; - - void initialize(ExecutorConfig const& executorConfig); - - void loadModel(std::optional const& modelPath, std::optional const& engineBuffer, - runtime::GptJsonConfig const& jsonConfig, ExecutorConfig const& executorConfig, bool isEncoder, - std::optional> const& managedWeightsOpt); - - std::shared_ptr createModel(runtime::RawEngine const& rawEngine, runtime::ModelConfig const& modelConfig, - runtime::WorldConfig const& worldConfig, ExecutorConfig const& executorConfig); - - std::shared_ptr createEncoderModel(runtime::RawEngine const& rawEngine, - runtime::ModelConfig const& modelConfig, runtime::WorldConfig const& worldConfig, - ExecutorConfig const& executorConfig); - - void setOrchLeaderComm(SizeType32 tp, SizeType32 pp, SizeType32 cp, ParallelConfig const& parallelConfig); - - void initializeCommAndWorkers(SizeType32 tp, SizeType32 pp, SizeType32 cp, ExecutorConfig const& executorConfig, - std::optional modelType = std::nullopt, - std::optional const& modelPath = std::nullopt, - std::optional const& worldConfig = std::nullopt, - std::optional const& decoderGptJsonConfig = std::nullopt); - - static void validateParallelConfig(ParallelConfig const& parallelConfig, std::optional modelType, - std::optional const& modelPath); - - void initializeOrchestrator(SizeType32 tp, SizeType32 pp, SizeType32 cp, ExecutorConfig const& executorConfig, - ParallelConfig parallelConfig, ModelType modelType, std::filesystem::path const& modelPath); - - void initializeWorkers(SizeType32 tp, SizeType32 pp, SizeType32 cp, ParallelConfig& parallelConfig, - std::optional const& worldConfig = std::nullopt, - std::optional const& decoderGptJsonConfig = std::nullopt); - - void initializeLogitsPostProcessorBatched(LogitsPostProcessorConfig const& logitsProcConfig); - - IdType generateReqId(Request const& request) - { - // If the request has a disaggregated request id, prefer it. - if (request.getDisaggRequestId().has_value() && request.getDisaggRequestId().value() > kMaxLocalReqId) - { - return request.getDisaggRequestId().value(); - } - // Otherwise, generate a local request id in range [1, kMaxLocalReqId). - return generateLocalReqId(); - } - - IdType generateLocalReqId() - { - return (mLastReqId++ % kMaxLocalReqId); - } - - std::vector getLeaderNewReqWithIds( - SizeType32 numActiveRequests, std::optional lowestPriorityActive); - std::vector getNewReqWithIds( - SizeType32 numActiveRequests, std::optional lowestPriorityActive); - - std::tuple fetchNewRequests( - SizeType32 numActiveRequests, std::optional lowestPriorityActive); - - void forwardSync(RequestList& activeRequests); - - void forwardAsync(RequestList& activeRequests); - - void prepRequestsForEncoderSkip(RequestList& activeRequests); - - void terminateActiveRequests(RequestList& activeRequests, std::string const& err); - - IterationStats getCurrentIterationStats(RequestList const& activeRequests, double iterLatencyMS, - SizeType32 numNewActiveRequests, double newActiveRequestsQueueLatencyMS, SizeType32 numCompletedRequests); - - void appendCurrentIterStats(IterationStats&& currentIterStats); - void appendMultipleIterStats(std::vector&& currentIterStatsVec); - void updateIterationStats(RequestList const& activeRequests, double iterLatencyMS, SizeType32 numNewActiveRequests, - double newActiveRequestsQueueLatencyMS, SizeType32 numCompletedRequests, bool flushToOrchestrator); - void appendCurrentRequestStats(RequestStatsPerIteration&& currentRequestStats); - void appendMultipleRequestStats(std::vector&& currentRequestStatsVec); - RequestStatsPerIteration getCurrentRequestStats( - RequestList const& activeRequests, RequestList const& finishedRequests); - void updateRequestStats( - RequestList const& activeRequests, RequestList const& finishedRequests, bool flushToOrchestrator); - - void appendCurrentDebugTensors(); - - void terminateCancelledRequests(RequestList& activeRequests); - - void terminateContextFinishedRequests(InTransList& inTransmissionRequests); - - void appendNewResponses(std::vector&& newResponses); - - /// @brief Populates new responses from active requests. - /// Active requests that have completed are erased from activeRequests - /// and returned for bookkeeping. - /// @return A list of requests that have completed. - RequestList populateNewResponses( - RequestList& activeRequests, InTransList& inTransmissionRequests, std::vector& newResponses); - - void executionLoop(); - - void enqueueTerminateRequest(); - void enqueueNewResponses(std::vector&& newResponses); - - LlmRequestLogitsPostProcessor getLogitsPostProcessor(std::string const& name); - void setupDynamicLogitsPostProcessors(std::vector& newReqWithIds); - void cleanupDynamicLogitsPostProcessors(RequestList const& finishedRequests); - - void orchSendReqThread(); - void orchRecvThread(mpi::MpiTag idTag, mpi::MpiTag dataTag); - void leaderRecvReqThread(); - void leaderSendThread(MpiMessageQueue& sendQueue, mpi::MpiTag idTag, mpi::MpiTag dataTag); - - void addTerminatedReqId(std::vector const& responses, IdType const& reqId); - - // Check that the current process is the leader or orchestrator - void checkParallelApiUsage(std::string const& methodName) const; - - // These functions wait for MPI async sends on separate threads - void requestWithIdWaitThread(); - void cancelledRequestsWaitThread(); - // These functions send data from leader to pipeline leader on separate threads - void requestWithIdLeaderThread(); - void cancelledRequestsLeaderThread(); - - /// @brief mark requests that have timed out before ever being executed as finished. - /// uses cancellation based on communication mode. - /// - /// @param activeRequests [in] List of active requests to check for timeouts - void finishTimedOutRequests(RequestList const& activeRequests); - - // The model to execute - std::shared_ptr mModel = nullptr; - std::shared_ptr mEncoderModel = nullptr; - - // The maximum number of activeRequests - SizeType32 mMaxNumActiveRequests; - - // Thread the executes the main loop - std::thread mExecutionThread; - - // Atomic that indicates threads should shutdown - std::atomic mShutdown; - - // Atomic that indicates if shutdown method has been called - std::atomic mShutdownCalled = false; - - // Queued requests - std::mutex mQueuedReqMtx; - std::condition_variable mQueuedReqCv; - std::deque mQueuedRequests; - std::optional mMaxQueueSize; - - // Cancelled requests - std::mutex mCancelReqMtx; - std::unordered_set mCancelledReqIds; - std::unordered_set mPipelineCancelledReqIds; - - // Ready responses - std::unordered_map> mResponses; - mutable std::mutex mResponsesMtx; - std::condition_variable mResponsesCv; - - // Since the request IDs are generated sequentially, IntervalSet is preferred over unordered_set for its efficient - // memory usage to stores request ID intervals rather than individual request ID numbers. - IntervalSet mTerminatedReqIds; - - std::unordered_map> mChildReqIdsMap; - - // Iteration stats - SizeType32 mIterStatsMaxIterations; - std::mutex mIterStatsMtx; - std::deque mIterationStats; - - // Request stats - SizeType32 mRequestStatsMaxIterations; - std::mutex mRequestStatsMtx; - std::deque mRequestStats; - - // Debug - IterationType mDebugTensorsMaxIterations; - std::mutex mDebugTensorsMtx; - std::deque mDebugTensors; - - IdType mLastReqId = 1; - - static constexpr IdType kTerminateReqId = 0; - // Request id > kMaxLocalReqId is reserved for disaggregated requests. - // This max ID is also in Python side. - static constexpr IdType kMaxLocalReqId = 1ULL << 42U; - - BatchingType mBatchingType; - bool mIsSchedulerMaxUtilization; - bool mIsSchedulerGuaranteedNoEvict; - bool mIsChunkedContext; - bool mPromptTableOffloading; - - CommunicationMode mCommMode; - bool mIsWorker = false; - bool mIsLeader = false; - bool mIsPipelineLeader = false; - bool mUsePipelineParallel = false; - - std::unordered_map mLogitsPostProcessorMap; - std::optional mLogitsPostProcessorBatched; - - bool mIsOrchestrator = false; - std::shared_ptr mOrchLeaderComm; - - std::thread mOrchSendReqThread; - std::thread mOrchRecvThread; - std::thread mLeaderRecvReqThread; - std::thread mLeaderSendThread; - - int32_t mRecvPollPeriodMs = 0; - - int32_t mLeaderRank = -1; - int32_t mOrchRank = 0; - int32_t mWorldRank = -1; - int32_t mDeviceId = 0; - - MpiMessageQueue mSendQueue; - - std::shared_ptr mCommTensorParallel; - std::shared_ptr mCommPipelineParallel; - std::shared_ptr mCommContextParallel; - std::unique_ptr mRequestWithIdAsyncSndHdl; - std::unique_ptr mCancelledRequestsAsyncSndHdl; - std::unique_ptr mRequestWithIdLeaderThread; - std::unique_ptr mCancelledRequestsLeaderThread; - std::unique_ptr mRequestWithIdWaitThread; - std::unique_ptr mCancelledRequestsWaitThread; - - // for validating requests - bool mEnableBlockReuse; - - inline static std::string const kPROFILE_START_STOP_ENV_VAR_NAME = "TLLM_PROFILE_START_STOP"; - inline static std::string const kLEGACY_PROFILE_START_STOP_ENV_VAR_NAME = "TLLM_GPTM_PROFILE_START_STOP"; - - std::shared_ptr mDynamicBatchTuner; -}; - -} // namespace tensorrt_llm::executor diff --git a/cpp/tensorrt_llm/executor/intervalSet.h b/cpp/tensorrt_llm/executor/intervalSet.h deleted file mode 100644 index 01f40a685f46..000000000000 --- a/cpp/tensorrt_llm/executor/intervalSet.h +++ /dev/null @@ -1,141 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2025 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. - */ - -#pragma once - -#include "tensorrt_llm/executor/types.h" - -namespace tensorrt_llm::executor -{ - -/// @brief An interval inclusive on both ends. -/// A single number interval is represented as [num, num]. -template -struct Interval -{ - T lowerEnd; - T upperEnd; -}; - -template -bool operator<(Interval const& a, Interval const& b) -{ - return a.lowerEnd < b.lowerEnd; -} - -/// @brief A container to store unique numbers, represented as a vector of ordered and disjoint intervals. -template -class IntervalSet -{ -public: - /// @brief Check if the given number is in set. - bool contains(NumType num) const - { - // Binary search - SizeType32 left = 0; - SizeType32 right = static_cast(mIntervals.size()) - 1; - while (left <= right) - { - SizeType32 mid = left + (right - left) / 2; - if (mIntervals[mid].lowerEnd <= num && num <= mIntervals[mid].upperEnd) - { - return true; - } - else if (num < mIntervals[mid].lowerEnd) - { - right = mid - 1; - } - else - { - left = mid + 1; - } - } - return false; - } - - /// @brief Insert a number into set. Do nothing if the number is already in the set. - void insert(NumType num) - { - auto intervalToAdd = Interval{num, num}; - - if (mIntervals.size() == 0) - { - mIntervals.insert(mIntervals.begin(), intervalToAdd); - mNumElements++; - return; - } - - // Iter is the first place in mIntervals such that num <= it.lowerEnd - auto iter = std::lower_bound(mIntervals.begin(), mIntervals.end(), intervalToAdd); - - bool iterAtBegin = iter == mIntervals.begin(); - bool iterAtEnd = iter == mIntervals.end(); - - if ((!iterAtEnd && iter->lowerEnd == num) || (!iterAtBegin && num <= (iter - 1)->upperEnd)) - { - // Number falls within the current interval or previous interval. No need to add again. - return; - } - - if (!iterAtBegin && !iterAtEnd && (iter - 1)->upperEnd + 1 == num && iter->lowerEnd - 1 == num) - { - // Merge two adjacent intervals - (iter - 1)->upperEnd = iter->upperEnd; - mIntervals.erase(iter); - } - else if (!iterAtBegin && (iter - 1)->upperEnd + 1 == num) - { - // Number is adjacent to the upper end of the previous interval. Merge left. - (iter - 1)->upperEnd = num; - } - else if (!iterAtEnd && iter->lowerEnd - 1 == num) - { - // Number is adjacent to the lower end of the current interval. Merge right. - iter->lowerEnd = num; - } - else - { - mIntervals.insert(iter, intervalToAdd); - } - mNumElements++; - } - - /// @brief Clear interval set and reset numElements to 0. - void clear() - { - mIntervals.clear(); - mNumElements = 0; - } - - /// @brief Return the size of the set. - SizeType32 getNumElements() const - { - return mNumElements; - } - - /// @brief Return the underlying mIntervals. - std::vector> const& getIntervals() const - { - return mIntervals; - } - -private: - std::vector> mIntervals; - SizeType32 mNumElements{0}; -}; - -} // namespace tensorrt_llm::executor diff --git a/cpp/tensorrt_llm/executor/model.h b/cpp/tensorrt_llm/executor/model.h deleted file mode 100644 index 52fedf1d1113..000000000000 --- a/cpp/tensorrt_llm/executor/model.h +++ /dev/null @@ -1,131 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2025 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. - */ - -#pragma once - -#include "tensorrt_llm/batch_manager/common.h" -#include "tensorrt_llm/batch_manager/logitsPostProcessor.h" -#include "tensorrt_llm/runtime/bufferManager.h" -#include "tensorrt_llm/runtime/modelConfig.h" -#include "tensorrt_llm/runtime/worldConfig.h" - -#include - -namespace tensorrt_llm::executor -{ - -class Model -{ - using LlmRequestPtr = std::shared_ptr; - -public: - Model() = default; - - virtual ~Model() = default; - - /// @brief Function that marks a request Id as complete and cleans up associated state - virtual void terminateRequest(LlmRequestPtr const& llmRequest, bool pause) = 0; - - void terminateRequest(LlmRequestPtr const& llmRequest) - { - terminateRequest(llmRequest, false); - } - - /// @brief Terminate request in the next forwardSync call that includes the request. - virtual void terminateRequestSync(LlmRequestPtr const& llmRequest, FinishReason finishReason) = 0; - - /// @brief Function that synchronizes the decoder - virtual void forwardSync() = 0; - - /// @brief Function that tries to advance the active requests - /// Depending on resources available, it's possible that not all requests will get advanced - /// @param activeRequests The list of request to try to advance - virtual void forwardAsync(batch_manager::RequestList const& activeRequests) = 0; - - /// @brief Override the runtime batch size for the model - virtual void setRuntimeBatchSize(SizeType32 runtimeBatchSize) - { - // By default, we ignore the runtimeBatchSize unless the model actively supports it - } - - /// @brief Get the runtime batch size for the model - [[nodiscard]] virtual SizeType32 getRuntimeBatchSize() const - { - TLLM_CHECK_WITH_INFO(false, "getRuntimeBatchSize is not implemented"); - } - - /// @brieft Override the runtime max num tokens for the model - virtual void setRuntimeMaxNumTokens(SizeType32 runtimeMaxNumTokens) - { - // By default, we ignore the runtimeMaxNumTokens unless the model actively supports it - } - - virtual void updatePeftCache(LlmRequestPtr const& llmRequest) = 0; - - /// @brief Reset the iteration stats when there are no inflight requests - virtual void resetIterationStats() = 0; - - [[nodiscard]] virtual SizeType32 getMaxNumSequences() const = 0; - [[nodiscard]] virtual SizeType32 getMaxInputLen() const = 0; - [[nodiscard]] virtual SizeType32 getHiddenSize() const = 0; - [[nodiscard]] virtual SizeType32 getMaxSequenceLen() const = 0; - [[nodiscard]] virtual SizeType32 getVocabSizePadded() const = 0; - [[nodiscard]] virtual SizeType32 getMaxDraftLen() const = 0; - [[nodiscard]] virtual SizeType32 getNumMicroBatches() const = 0; - [[nodiscard]] virtual SizeType32 getOperatingBeamWidth() const = 0; - [[nodiscard]] virtual nvinfer1::DataType getLogitDataType() const = 0; - [[nodiscard]] virtual runtime::WorldConfig const& getWorldConfig() const = 0; - [[nodiscard]] virtual runtime::ModelConfig const& getModelConfig() const = 0; - [[nodiscard]] virtual runtime::BufferManager const& getBufferManager() const = 0; - [[nodiscard]] virtual runtime::BufferManager::CudaStreamPtr getRuntimeStreamPtr() const = 0; - [[nodiscard]] virtual IterationType getIterCounter() const noexcept = 0; - [[nodiscard]] virtual bool hasSpeculativeDecodingFastLogits() const noexcept = 0; - [[nodiscard]] virtual bool getGatherGenerationLogits() const = 0; - [[nodiscard]] virtual nvinfer1::DataType getTensorDataType(std::string const& name) const = 0; - [[nodiscard]] virtual nvinfer1::Dims getTensorShape(std::string const& name) const = 0; - - /// @brief Function that provides per iteration stats specific to a certain model - /// @param stats The json object to write stats to - virtual void getCurrentIterationStats(IterationStats& stats) const = 0; - - /// @brief Function that provides per request stats specific to a certain model - /// @param stats The request stats to be updated - virtual void getCurrentRequestStats(RequestStatsPerIteration& stats) const = 0; - - [[nodiscard]] virtual DebugTensorsPerIteration getCurrentDebugTensors() const = 0; - - using LogitsPostProcessorBatched = tensorrt_llm::batch_manager::LogitsPostProcessor::LogitsPostProcessorBatched; - - virtual void setLogitsPostProcessorBatched(std::optional logitsPostProcessorBatched) - = 0; - virtual void setReplicateLogitsPostProcessor(bool replicateLogitsPostProcessor) = 0; - [[nodiscard]] virtual bool getReplicateLogitsPostProcessor() const = 0; - - [[nodiscard]] virtual bool hasGuidedDecoder() const noexcept = 0; - - [[nodiscard]] virtual std::shared_ptr - getKVCacheManager() = 0; - [[nodiscard]] virtual std::shared_ptr - getKVCacheManager() const = 0; - - //! \brief Get the batch size that can fill the kv cache to the maximum capacity give the sequence length - //! \param seqLen The sequence length - //! \return The batch size that can fill the kv cache to the maximum capacity. If unsuporrted, return 0. - [[nodiscard]] virtual SizeType32 getMaxCapacityBatchSize(SizeType32 inputLength, SizeType32 outputLength) const = 0; -}; - -} // namespace tensorrt_llm::executor diff --git a/cpp/tensorrt_llm/executor/serialization.cpp b/cpp/tensorrt_llm/executor/serialization.cpp index 020306e03e56..4c2998c25d26 100644 --- a/cpp/tensorrt_llm/executor/serialization.cpp +++ b/cpp/tensorrt_llm/executor/serialization.cpp @@ -17,6 +17,7 @@ #include "tensorrt_llm/executor/serialization.h" #include "tensorrt_llm/batch_manager/kvCacheManager.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/executor/dataTransceiverState.h" #include "tensorrt_llm/executor/executor.h" #include "tensorrt_llm/executor/requestImpl.h" @@ -626,8 +627,8 @@ kv_cache::CacheState Serialization::deserializeCacheState(std::istream& is) auto hasRnnConfig = su::deserialize(is); std::optional rnnModelConfig; std::vector rnnLayerNumPerPP; - nvinfer1::DataType convStateDataType{nvinfer1::DataType::kFLOAT}; - nvinfer1::DataType ssmStateDataType{nvinfer1::DataType::kFLOAT}; + tensorrt_llm::DataType convStateDataType{tensorrt_llm::DataType::kFLOAT}; + tensorrt_llm::DataType ssmStateDataType{tensorrt_llm::DataType::kFLOAT}; if (hasRnnConfig) { CacheState::RnnModelConfig rnnCfg; @@ -641,8 +642,8 @@ kv_cache::CacheState Serialization::deserializeCacheState(std::istream& is) rnnCfg.mNumHeads = su::deserialize(is); rnnCfg.mConvSectionLayout = static_cast(su::deserialize(is)); - convStateDataType = su::deserialize(is); - ssmStateDataType = su::deserialize(is); + convStateDataType = su::deserialize(is); + ssmStateDataType = su::deserialize(is); rnnLayerNumPerPP = su::deserialize>(is); rnnModelConfig = std::move(rnnCfg); } diff --git a/cpp/tensorrt_llm/executor/tensor.cpp b/cpp/tensorrt_llm/executor/tensor.cpp index c38feb0e34b8..9c508c0ec5c3 100644 --- a/cpp/tensorrt_llm/executor/tensor.cpp +++ b/cpp/tensorrt_llm/executor/tensor.cpp @@ -18,6 +18,7 @@ #include "tensorrt_llm/executor/tensor.h" #include "tensorrt_llm/common/assert.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/runtime/bufferManager.h" #include "tensorrt_llm/runtime/iTensor.h" @@ -53,17 +54,17 @@ DataType Tensor::getDataType() const } switch (mTensor->getDataType()) { - case nvinfer1::DataType::kBOOL: return DataType::kBOOL; - case nvinfer1::DataType::kINT8: return DataType::kINT8; - case nvinfer1::DataType::kINT32: return DataType::kINT32; - case nvinfer1::DataType::kUINT8: return DataType::kUINT8; - case nvinfer1::DataType::kFP8: return DataType::kFP8; - case nvinfer1::DataType::kHALF: return DataType::kFP16; - case nvinfer1::DataType::kFLOAT: return DataType::kFP32; - case nvinfer1::DataType::kBF16: return DataType::kBF16; - case nvinfer1::DataType::kINT64: return DataType::kINT64; - case nvinfer1::DataType::kINT4: [[fallthrough]] /* do nothing */; - case nvinfer1::DataType::kFP4: [[fallthrough]] /* do nothing */; + case tensorrt_llm::DataType::kBOOL: return DataType::kBOOL; + case tensorrt_llm::DataType::kINT8: return DataType::kINT8; + case tensorrt_llm::DataType::kINT32: return DataType::kINT32; + case tensorrt_llm::DataType::kUINT8: return DataType::kUINT8; + case tensorrt_llm::DataType::kFP8: return DataType::kFP8; + case tensorrt_llm::DataType::kHALF: return DataType::kFP16; + case tensorrt_llm::DataType::kFLOAT: return DataType::kFP32; + case tensorrt_llm::DataType::kBF16: return DataType::kBF16; + case tensorrt_llm::DataType::kINT64: return DataType::kINT64; + case tensorrt_llm::DataType::kINT4: [[fallthrough]] /* do nothing */; + case tensorrt_llm::DataType::kFP4: [[fallthrough]] /* do nothing */; default: TLLM_THROW("Unsupported data type"); } } @@ -135,19 +136,19 @@ tr::ITensor::Shape toDims(Shape const& shape) return dims; } -nvinfer1::DataType toDataType(DataType dataType) +tensorrt_llm::DataType toDataType(DataType dataType) { switch (dataType) { - case DataType::kBOOL: return nvinfer1::DataType::kBOOL; - case DataType::kUINT8: return nvinfer1::DataType::kUINT8; - case DataType::kINT8: return nvinfer1::DataType::kINT8; - case DataType::kINT32: return nvinfer1::DataType::kINT32; - case DataType::kINT64: return nvinfer1::DataType::kINT64; - case DataType::kBF16: return nvinfer1::DataType::kBF16; - case DataType::kFP8: return nvinfer1::DataType::kFP8; - case DataType::kFP16: return nvinfer1::DataType::kHALF; - case DataType::kFP32: return nvinfer1::DataType::kFLOAT; + case DataType::kBOOL: return tensorrt_llm::DataType::kBOOL; + case DataType::kUINT8: return tensorrt_llm::DataType::kUINT8; + case DataType::kINT8: return tensorrt_llm::DataType::kINT8; + case DataType::kINT32: return tensorrt_llm::DataType::kINT32; + case DataType::kINT64: return tensorrt_llm::DataType::kINT64; + case DataType::kBF16: return tensorrt_llm::DataType::kBF16; + case DataType::kFP8: return tensorrt_llm::DataType::kFP8; + case DataType::kFP16: return tensorrt_llm::DataType::kHALF; + case DataType::kFP32: return tensorrt_llm::DataType::kFLOAT; case DataType::kUNKNOWN: TLLM_THROW("Unsupported data type"); } diff --git a/cpp/tensorrt_llm/executor_worker/CMakeLists.txt b/cpp/tensorrt_llm/executor_worker/CMakeLists.txt deleted file mode 100644 index 2feb6dfe5790..000000000000 --- a/cpp/tensorrt_llm/executor_worker/CMakeLists.txt +++ /dev/null @@ -1,26 +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. -set(SRCS executorWorker.cpp) - -include_directories(${PROJECT_SOURCE_DIR}/include) - -set(EXECUTOR_WORKER_TARGET executorWorker) - -add_executable(${EXECUTOR_WORKER_TARGET} ${SRCS}) - -target_link_libraries(${EXECUTOR_WORKER_TARGET} - PUBLIC ${SHARED_TARGET} nvinfer_plugin_tensorrt_llm) - -target_compile_features(${EXECUTOR_WORKER_TARGET} PRIVATE cxx_std_17) diff --git a/cpp/tensorrt_llm/executor_worker/executorWorker.cpp b/cpp/tensorrt_llm/executor_worker/executorWorker.cpp deleted file mode 100644 index aa1b06c2cb74..000000000000 --- a/cpp/tensorrt_llm/executor_worker/executorWorker.cpp +++ /dev/null @@ -1,92 +0,0 @@ -/* - * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. - * - * 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. - */ - -#include "tensorrt_llm/common/cudaUtils.h" -#include "tensorrt_llm/common/logger.h" -#include "tensorrt_llm/executor/executor.h" -#include "tensorrt_llm/executor/serialization.h" -#include "tensorrt_llm/plugins/api/tllmPlugin.h" -#include "tensorrt_llm/runtime/utils/mpiUtils.h" -#include - -namespace tle = tensorrt_llm::executor; - -int main(int argc, char* argv[]) -{ -#if ENABLE_MULTI_DEVICE - - if (std::getenv("FORCE_NCCL_ALL_REDUCE_STRATEGY") != nullptr) - { - TLLM_LOG_INFO("FORCE_NCCL_ALL_REDUCE_STRATEGY env variable detected in worker"); - } - - // Register the TRT-LLM plugins - initTrtLlmPlugins(); - - tensorrt_llm::mpi::initialize(tensorrt_llm::mpi::MpiThreadSupport::THREAD_MULTIPLE, true); - - MPI_Comm parentComm; - MPI_Comm_get_parent(&parentComm); - if (parentComm == MPI_COMM_NULL) - { - TLLM_LOG_ERROR("TRT-LLM worker has no parent!"); - return -1; - } - - int size; - MPI_Comm_remote_size(parentComm, &size); - if (size != 1) - { - TLLM_LOG_ERROR("Parent size is %d, must be 1", size); - return -1; - } - - // Since parentComm is an intercommunicator, input root - // is the rank of the parent process in his group - // (always 0 as the parent size is checked before) - - // Receive from the parent the executor configuration - int64_t bufferSize; - MPICHECK(MPI_Bcast(&bufferSize, 1, MPI_INT64_T, 0, parentComm)); - std::vector buffer(bufferSize); - MPICHECK(MPI_Bcast(buffer.data(), bufferSize, MPI_CHAR, 0, parentComm)); - std::istringstream is(std::string(buffer.begin(), buffer.end())); - auto modelPath = tle::Serialization::deserializeString(is); - auto modelType = tle::Serialization::deserializeModelType(is); - auto executorConfig = tle::Serialization::deserializeExecutorConfig(is); - - // Create the orchestrator config for workers - auto orchLeaderComm = std::make_shared(parentComm, true); - auto parallelConfig = executorConfig.getParallelConfig(); - TLLM_CHECK_WITH_INFO(parallelConfig.has_value(), "Parallel config should have a value."); - TLLM_CHECK_WITH_INFO( - parallelConfig.value().getOrchestratorConfig().has_value(), "Orchestrator config should have a value."); - auto orchConfig = parallelConfig.value().getOrchestratorConfig().value(); - TLLM_CHECK_WITH_INFO(parallelConfig.has_value(), "Parallel config should have a value."); - auto newOrchConfig = tle::OrchestratorConfig(false, orchConfig.getWorkerExecutablePath(), orchLeaderComm); - parallelConfig.value().setOrchestratorConfig(newOrchConfig); - executorConfig.setParallelConfig(parallelConfig.value()); - // In orchestrator mode, the spawned threads will wait for termination signal from orchestrator - auto executor = tle::Executor(modelPath, modelType, executorConfig); - - // Wait for all workers to have created their instances - MPI_Barrier(parentComm); - TLLM_LOG_INFO("Executor instance created by worker"); - -#endif // ENABLE_MULTI_DEVICE - - return 0; -} diff --git a/cpp/tensorrt_llm/kernels/communicationKernels/MiniMaxReduceRMSKernel.cu b/cpp/tensorrt_llm/kernels/communicationKernels/MiniMaxReduceRMSKernel.cu index 5be8b1c2ff78..cf0ad7040f59 100644 --- a/cpp/tensorrt_llm/kernels/communicationKernels/MiniMaxReduceRMSKernel.cu +++ b/cpp/tensorrt_llm/kernels/communicationKernels/MiniMaxReduceRMSKernel.cu @@ -16,6 +16,7 @@ #include "tensorrt_llm/common/config.h" #include "tensorrt_llm/common/envUtils.h" #include "tensorrt_llm/common/reduceKernelUtils.cuh" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/kernels/communicationKernels/MiniMaxReduceRMSKernel.h" #include "tensorrt_llm/kernels/quantization.cuh" #include @@ -816,7 +817,7 @@ void dispatch_dtype(MiniMaxReduceRMSParams const& params) bool use_float4 = (params.allreduce_in_k != nullptr) && (params.hidden_dim * params.nranks == 6144) && (params.hidden_dim_k * params.nranks == 1024); - if (params.dtype == nvinfer1::DataType::kHALF) + if (params.dtype == tensorrt_llm::DataType::kHALF) { if (use_float4) { @@ -827,7 +828,7 @@ void dispatch_dtype(MiniMaxReduceRMSParams const& params) minimax_reduce_rms_kernel_launcher(params); } } - else if (params.dtype == nvinfer1::DataType::kBF16) + else if (params.dtype == tensorrt_llm::DataType::kBF16) { if (use_float4) { @@ -838,7 +839,7 @@ void dispatch_dtype(MiniMaxReduceRMSParams const& params) minimax_reduce_rms_kernel_launcher<__nv_bfloat16, NRanks>(params); } } - else if (params.dtype == nvinfer1::DataType::kFLOAT) + else if (params.dtype == tensorrt_llm::DataType::kFLOAT) { if (use_float4) { diff --git a/cpp/tensorrt_llm/kernels/communicationKernels/MiniMaxReduceRMSKernel.h b/cpp/tensorrt_llm/kernels/communicationKernels/MiniMaxReduceRMSKernel.h index b0cfd0ca074c..bf5775f96de9 100644 --- a/cpp/tensorrt_llm/kernels/communicationKernels/MiniMaxReduceRMSKernel.h +++ b/cpp/tensorrt_llm/kernels/communicationKernels/MiniMaxReduceRMSKernel.h @@ -15,7 +15,7 @@ */ #pragma once #include "tensorrt_llm/common/assert.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include #include @@ -59,7 +59,7 @@ struct MiniMaxReduceRMSParams { int nranks{}; int rank{}; - nvinfer1::DataType dtype; + tensorrt_llm::DataType dtype; int size_q{}; // numel of Q (num_token * head_dim_q) int hidden_dim{}; // head_dim_q int size_k{}; // numel of K (num_token * head_dim_k) diff --git a/cpp/tensorrt_llm/kernels/communicationKernels/allReduceFusionKernels.cu b/cpp/tensorrt_llm/kernels/communicationKernels/allReduceFusionKernels.cu index 5a3edda04a70..d9fbc9da0a0c 100644 --- a/cpp/tensorrt_llm/kernels/communicationKernels/allReduceFusionKernels.cu +++ b/cpp/tensorrt_llm/kernels/communicationKernels/allReduceFusionKernels.cu @@ -16,6 +16,7 @@ #include "tensorrt_llm/common/config.h" #include "tensorrt_llm/common/envUtils.h" #include "tensorrt_llm/common/reduceKernelUtils.cuh" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/kernels/communicationKernels/allReduceFusionKernels.h" #include "tensorrt_llm/kernels/quantization.cuh" #include @@ -795,15 +796,15 @@ void allreduce_fusion_op(AllReduceFusionParams const& params) } #define DISPATCH_DTYPE(NRanks) \ - if (params.dtype == nvinfer1::DataType::kHALF) \ + if (params.dtype == tensorrt_llm::DataType::kHALF) \ { \ DISPATCH_PATTERN(half, NRanks); \ } \ - else if (params.dtype == nvinfer1::DataType::kBF16) \ + else if (params.dtype == tensorrt_llm::DataType::kBF16) \ { \ DISPATCH_PATTERN(__nv_bfloat16, NRanks); \ } \ - else if (params.dtype == nvinfer1::DataType::kFLOAT) \ + else if (params.dtype == tensorrt_llm::DataType::kFLOAT) \ { \ DISPATCH_PATTERN(float, NRanks); \ } \ diff --git a/cpp/tensorrt_llm/kernels/communicationKernels/allReduceFusionKernels.h b/cpp/tensorrt_llm/kernels/communicationKernels/allReduceFusionKernels.h index 6d2074a6589e..769776273ef6 100644 --- a/cpp/tensorrt_llm/kernels/communicationKernels/allReduceFusionKernels.h +++ b/cpp/tensorrt_llm/kernels/communicationKernels/allReduceFusionKernels.h @@ -16,7 +16,7 @@ #pragma once #include "tensorrt_llm/common/assert.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include #include @@ -124,7 +124,7 @@ struct AllReduceFusionParams { int nranks; int rank; - nvinfer1::DataType dtype; + tensorrt_llm::DataType dtype; int size; int hidden_dim; void** workspace; diff --git a/cpp/tensorrt_llm/kernels/communicationKernels/customLowPrecisionAllReduceKernels.cu b/cpp/tensorrt_llm/kernels/communicationKernels/customLowPrecisionAllReduceKernels.cu index f1d5c08bda6b..09a742494a9a 100644 --- a/cpp/tensorrt_llm/kernels/communicationKernels/customLowPrecisionAllReduceKernels.cu +++ b/cpp/tensorrt_llm/kernels/communicationKernels/customLowPrecisionAllReduceKernels.cu @@ -21,6 +21,7 @@ #include "tensorrt_llm/common/customAllReduceUtils.h" #include "tensorrt_llm/common/dataType.h" #include "tensorrt_llm/common/envUtils.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/kernels/communicationKernels/customLowPrecisionAllReduceKernels.h" #include #include @@ -1357,7 +1358,7 @@ std::vector splitNumber(size_t number) } LowPrecisionAllReduceParams LowPrecisionAllReduceParams::deserialize( - size_t tpSize, size_t tpRank, nvinfer1::DataType dataType, int token_num, int hidden_size) + size_t tpSize, size_t tpRank, tensorrt_llm::DataType dataType, int token_num, int hidden_size) { // Get appropriate static buffer @@ -1401,7 +1402,7 @@ LowPrecisionAllReduceParams LowPrecisionAllReduceParams::deserialize( } LowPrecisionAllReduceParams LowPrecisionAllReduceParams::deserialize_hier( - size_t tpSize, size_t tpRank, nvinfer1::DataType dataType, int token_num, int hidden_size) + size_t tpSize, size_t tpRank, tensorrt_llm::DataType dataType, int token_num, int hidden_size) { // Get appropriate static buffer @@ -1616,7 +1617,7 @@ int32_t max_workspace_size_lowprecision(int32_t tp_size) } void customLowPrecisionAllReduce( - kernels::LowPrecisionAllReduceParams& params, nvinfer1::DataType dataType, cudaStream_t stream) + kernels::LowPrecisionAllReduceParams& params, tensorrt_llm::DataType dataType, cudaStream_t stream) { TLLM_CHECK_WITH_INFO(lowPrecisionConfigurationSupported(params.ranks_per_node, params.elts_total), "Low Precision Custom all-reduce configuration unsupported"); @@ -1625,10 +1626,10 @@ void customLowPrecisionAllReduce( switch (dataType) { - case nvinfer1::DataType::kFLOAT: lowPrecisionAllReduceDispatchType(params, stream); break; - case nvinfer1::DataType::kHALF: lowPrecisionAllReduceDispatchType(params, stream); break; + case tensorrt_llm::DataType::kFLOAT: lowPrecisionAllReduceDispatchType(params, stream); break; + case tensorrt_llm::DataType::kHALF: lowPrecisionAllReduceDispatchType(params, stream); break; #ifdef ENABLE_BF16 - case nvinfer1::DataType::kBF16: lowPrecisionAllReduceDispatchType<__nv_bfloat16>(params, stream); break; + case tensorrt_llm::DataType::kBF16: lowPrecisionAllReduceDispatchType<__nv_bfloat16>(params, stream); break; #endif default: TLLM_THROW("Unsupported dataType for customAllReduce"); } diff --git a/cpp/tensorrt_llm/kernels/communicationKernels/customLowPrecisionAllReduceKernels.h b/cpp/tensorrt_llm/kernels/communicationKernels/customLowPrecisionAllReduceKernels.h index 5fc87ef1a523..62d19039cc4e 100644 --- a/cpp/tensorrt_llm/kernels/communicationKernels/customLowPrecisionAllReduceKernels.h +++ b/cpp/tensorrt_llm/kernels/communicationKernels/customLowPrecisionAllReduceKernels.h @@ -19,8 +19,8 @@ #include "tensorrt_llm/common/assert.h" #include "tensorrt_llm/common/config.h" #include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/kernels/customAllReduceKernels.h" -#include #include #include #include @@ -111,15 +111,15 @@ struct LowPrecisionAllReduceParams uint64_t* ag_notify_peer_inside_numa_flags[LP_ALLREDUCE_MAX_BLOCKS * 4]; // 3*flags , 3 is other rank inside numa static LowPrecisionAllReduceParams deserialize( - size_t tpSize, size_t tpRank, nvinfer1::DataType dataType, int token_num, int hidden_size); + size_t tpSize, size_t tpRank, tensorrt_llm::DataType dataType, int token_num, int hidden_size); static LowPrecisionAllReduceParams deserialize_hier( - size_t tpSize, size_t tpRank, nvinfer1::DataType dataType, int token_num, int hidden_size); + size_t tpSize, size_t tpRank, tensorrt_llm::DataType dataType, int token_num, int hidden_size); }; bool lowPrecisionConfigurationSupported(size_t msg_size, size_t n_ranks); void customLowPrecisionAllReduce( - kernels::LowPrecisionAllReduceParams& params, nvinfer1::DataType dataType, cudaStream_t stream); + kernels::LowPrecisionAllReduceParams& params, tensorrt_llm::DataType dataType, cudaStream_t stream); int32_t max_workspace_size_lowprecision(int32_t tp_size); } // namespace kernels diff --git a/cpp/tensorrt_llm/kernels/communicationKernels/mnnvlAllreduceKernels.cu b/cpp/tensorrt_llm/kernels/communicationKernels/mnnvlAllreduceKernels.cu index eb44f1638a19..a8ff3dd431a8 100644 --- a/cpp/tensorrt_llm/kernels/communicationKernels/mnnvlAllreduceKernels.cu +++ b/cpp/tensorrt_llm/kernels/communicationKernels/mnnvlAllreduceKernels.cu @@ -33,6 +33,7 @@ #include "tensorrt_llm/common/lamportUtils.cuh" #include "tensorrt_llm/common/logger.h" #include "tensorrt_llm/common/reduceKernelUtils.cuh" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/kernels/quantization.cuh" TRTLLM_NAMESPACE_BEGIN @@ -856,9 +857,9 @@ void oneshotAllreduceFusionOp(AllReduceFusionParams const& params) }; #undef LAUNCH_ALLREDUCE_KERNEL #undef DISPATCH_ALLREDUCE_PATTERN - bool launched = (params.dType == nvinfer1::DataType::kBF16 && dispatchImpl((__nv_bfloat16*) nullptr)) - || (params.dType == nvinfer1::DataType::kFLOAT && dispatchImpl((float*) nullptr)) - || (params.dType == nvinfer1::DataType::kHALF && dispatchImpl((__nv_half*) nullptr)); + bool launched = (params.dType == tensorrt_llm::DataType::kBF16 && dispatchImpl((__nv_bfloat16*) nullptr)) + || (params.dType == tensorrt_llm::DataType::kFLOAT && dispatchImpl((float*) nullptr)) + || (params.dType == tensorrt_llm::DataType::kHALF && dispatchImpl((__nv_half*) nullptr)); if (!launched) { TLLM_CHECK_WITH_INFO(false, "Failed to dispatch MNNVL AllReduceOneShot kernel."); @@ -1246,9 +1247,9 @@ void twoshotAllreduceFusionOp(AllReduceFusionParams const& params) #undef LAUNCH_ALLREDUCE_KERNEL - bool launched = (params.dType == nvinfer1::DataType::kFLOAT && dispatchAR((float*) nullptr)) - || (params.dType == nvinfer1::DataType::kBF16 && dispatchAR((__nv_bfloat16*) nullptr)) - || (params.dType == nvinfer1::DataType::kHALF && dispatchAR((__nv_half*) nullptr)); + bool launched = (params.dType == tensorrt_llm::DataType::kFLOAT && dispatchAR((float*) nullptr)) + || (params.dType == tensorrt_llm::DataType::kBF16 && dispatchAR((__nv_bfloat16*) nullptr)) + || (params.dType == tensorrt_llm::DataType::kHALF && dispatchAR((__nv_half*) nullptr)); if (!launched) { TLLM_CHECK_WITH_INFO(false, "[MNNVL AllReduceTwoShot] Failed to dispatch twoshotAllreduce kernel."); @@ -1388,9 +1389,9 @@ void twoshotAllreduceFusionOp(AllReduceFusionParams const& params) return true; }; - launched = (params.dType == nvinfer1::DataType::kFLOAT && dispatchRN((float*) nullptr)) - || (params.dType == nvinfer1::DataType::kBF16 && dispatchRN((__nv_bfloat16*) nullptr)) - || (params.dType == nvinfer1::DataType::kHALF && dispatchRN((__nv_half*) nullptr)); + launched = (params.dType == tensorrt_llm::DataType::kFLOAT && dispatchRN((float*) nullptr)) + || (params.dType == tensorrt_llm::DataType::kBF16 && dispatchRN((__nv_bfloat16*) nullptr)) + || (params.dType == tensorrt_llm::DataType::kHALF && dispatchRN((__nv_half*) nullptr)); if (!launched) { TLLM_CHECK_WITH_INFO(false, "[MNNVL AllReduceTwoShot] Failed to dispatch rmsnorm lamport kernel."); diff --git a/cpp/tensorrt_llm/kernels/communicationKernels/mnnvlAllreduceKernels.h b/cpp/tensorrt_llm/kernels/communicationKernels/mnnvlAllreduceKernels.h index 2a228e815b8d..f2006f52240c 100644 --- a/cpp/tensorrt_llm/kernels/communicationKernels/mnnvlAllreduceKernels.h +++ b/cpp/tensorrt_llm/kernels/communicationKernels/mnnvlAllreduceKernels.h @@ -17,8 +17,8 @@ #define TRTLLM_MNNVL_ALLREDUCE_KERNELS_H #include "tensorrt_llm/common/config.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/kernels/communicationKernels/allReduceFusionKernels.h" -#include #include TRTLLM_NAMESPACE_BEGIN @@ -39,16 +39,16 @@ struct AllReduceFusionParams //! \name Environmental and Auxiliary Data //! @{ - int nRanks; //!< Total number of participating ranks in the AllReduce operation - int rank; //!< Current rank ID - nvinfer1::DataType dType; //!< Data type of the tensors (e.g., FP16, BF16, FP32) - int numTokens; //!< Number of tokens in the input tensor - int tokenDim; //!< Hidden Dimension - void** bufferPtrsDev; //!< Unicast Device pointers to communication buffers for each rank - void* bufferPtrLocal; //!< Local buffer pointer for temporary storage (i.e., bufferPtrsDev[rank]) - void* multicastPtr; //!< Multicast buffer pointer. - uint32_t* bufferFlags; //!< Synchronization flags for coordinating communication phases - bool rmsNormFusion; //!< Whether to fuse RMS normalization with the AllReduce operation + int nRanks; //!< Total number of participating ranks in the AllReduce operation + int rank; //!< Current rank ID + tensorrt_llm::DataType dType; //!< Data type of the tensors (e.g., FP16, BF16, FP32) + int numTokens; //!< Number of tokens in the input tensor + int tokenDim; //!< Hidden Dimension + void** bufferPtrsDev; //!< Unicast Device pointers to communication buffers for each rank + void* bufferPtrLocal; //!< Local buffer pointer for temporary storage (i.e., bufferPtrsDev[rank]) + void* multicastPtr; //!< Multicast buffer pointer. + uint32_t* bufferFlags; //!< Synchronization flags for coordinating communication phases + bool rmsNormFusion; //!< Whether to fuse RMS normalization with the AllReduce operation ar_fusion::AllReduceFusionPattern pattern = ar_fusion::AllReduceFusionPattern::kAllReduce; //!< Fused epilogue pattern diff --git a/cpp/tensorrt_llm/kernels/communicationKernels/moeAllReduceFusionKernels.cu b/cpp/tensorrt_llm/kernels/communicationKernels/moeAllReduceFusionKernels.cu index 306d42677e2f..d1f50a2fb9c8 100644 --- a/cpp/tensorrt_llm/kernels/communicationKernels/moeAllReduceFusionKernels.cu +++ b/cpp/tensorrt_llm/kernels/communicationKernels/moeAllReduceFusionKernels.cu @@ -16,6 +16,7 @@ #include "tensorrt_llm/common/config.h" #include "tensorrt_llm/common/envUtils.h" #include "tensorrt_llm/common/reduceKernelUtils.cuh" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/kernels/communicationKernels/moeAllReduceFusionKernels.h" #include "tensorrt_llm/kernels/quantization.cuh" #include @@ -442,11 +443,11 @@ void moereduction_allreduce_fusion_op(MoeReductionAllReduceFusionParams const& p #define MOE_DISPATCH1(DTYPE, NRANKS, RESIDUAL_OUT, NORM_OUT, QUANT_OUT) \ return moereduction_allreduce_fusion_kernel_launcher(params); #define MOE_DISPATCH0(NRANKS, RESIDUAL_OUT, NORM_OUT, QUANT_OUT) \ - if (params.nranks == NRANKS && params.dtype == nvinfer1::DataType::kHALF) \ + if (params.nranks == NRANKS && params.dtype == tensorrt_llm::DataType::kHALF) \ { \ MOE_DISPATCH1(half, NRANKS, RESIDUAL_OUT, NORM_OUT, QUANT_OUT); \ } \ - else if (params.nranks == NRANKS && params.dtype == nvinfer1::DataType::kBF16) \ + else if (params.nranks == NRANKS && params.dtype == tensorrt_llm::DataType::kBF16) \ { \ MOE_DISPATCH1(__nv_bfloat16, NRANKS, RESIDUAL_OUT, NORM_OUT, QUANT_OUT); \ } @@ -727,13 +728,13 @@ void moefinalize_allreduce_fusion_op(MoeFinalizeAllReduceFusionParams const& par #define MOE_FINALIZE_DISPATCH1(DTYPE, NRANKS, RESIDUAL_OUT, NORM_OUT, QUANT_OUT) \ return moefinalize_allreduce_fusion_kernel_launcher(params); #define MOE_FINALIZE_DISPATCH0(NRANKS, RESIDUAL_OUT, NORM_OUT, QUANT_OUT) \ - if (params.nranks == NRANKS && params.dtype == nvinfer1::DataType::kHALF \ - && params.scale_dtype == nvinfer1::DataType::kHALF) \ + if (params.nranks == NRANKS && params.dtype == tensorrt_llm::DataType::kHALF \ + && params.scale_dtype == tensorrt_llm::DataType::kHALF) \ { \ MOE_FINALIZE_DISPATCH1(half, NRANKS, RESIDUAL_OUT, NORM_OUT, QUANT_OUT); \ } \ - else if (params.nranks == NRANKS && params.dtype == nvinfer1::DataType::kBF16 \ - && params.scale_dtype == nvinfer1::DataType::kBF16) \ + else if (params.nranks == NRANKS && params.dtype == tensorrt_llm::DataType::kBF16 \ + && params.scale_dtype == tensorrt_llm::DataType::kBF16) \ { \ MOE_FINALIZE_DISPATCH1(__nv_bfloat16, NRANKS, RESIDUAL_OUT, NORM_OUT, QUANT_OUT); \ } diff --git a/cpp/tensorrt_llm/kernels/communicationKernels/moeAllReduceFusionKernels.h b/cpp/tensorrt_llm/kernels/communicationKernels/moeAllReduceFusionKernels.h index 556dd4e5cd24..e526a70268b3 100644 --- a/cpp/tensorrt_llm/kernels/communicationKernels/moeAllReduceFusionKernels.h +++ b/cpp/tensorrt_llm/kernels/communicationKernels/moeAllReduceFusionKernels.h @@ -16,7 +16,7 @@ #pragma once #include "tensorrt_llm/common/assert.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include #include @@ -44,7 +44,7 @@ struct AllReduceFusionParams { int nranks; int rank; - nvinfer1::DataType dtype; + tensorrt_llm::DataType dtype; // size = token_num * hidden_dim int size; int hidden_dim; @@ -94,7 +94,7 @@ struct MoeFinalizeAllReduceFusionParams : public AllReduceFusionParams // Refer to kernel implementation on layout of those params // number of active experts on current device int top_k; - nvinfer1::DataType scale_dtype; + tensorrt_llm::DataType scale_dtype; // [num_tokens, top_k] void* expert_scale_factor = nullptr; void* shared_expert_output = nullptr; diff --git a/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu b/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu index 472a5877a80d..74e40dbb2b81 100644 --- a/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu +++ b/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu @@ -17,6 +17,7 @@ #include "tensorrt_llm/common/config.h" #include "tensorrt_llm/common/cudaUtils.h" #include "tensorrt_llm/common/envUtils.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/common/vec_dtypes.cuh" #include "tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h" #include "tensorrt_llm/kernels/quantization.cuh" @@ -129,25 +130,25 @@ using tensorrt_llm::common::launchWithPdlWhenEnabled; #define SWITCH_DTYPE(dtype, TYPE, ...) \ switch (dtype) \ { \ - case nvinfer1::DataType::kHALF: \ + case tensorrt_llm::DataType::kHALF: \ { \ using TYPE = half; \ __VA_ARGS__; \ break; \ } \ - case nvinfer1::DataType::kBF16: \ + case tensorrt_llm::DataType::kBF16: \ { \ using TYPE = __nv_bfloat16; \ __VA_ARGS__; \ break; \ } \ - case nvinfer1::DataType::kFLOAT: \ + case tensorrt_llm::DataType::kFLOAT: \ { \ using TYPE = float; \ __VA_ARGS__; \ break; \ } \ - case nvinfer1::DataType::kFP8: \ + case tensorrt_llm::DataType::kFP8: \ { \ using TYPE = __nv_fp8_e4m3; \ __VA_ARGS__; \ @@ -1403,7 +1404,7 @@ void moe_a2a_combine_launch(MoeA2ACombineParams const& params) // When use_low_precision is set the recv buffers contain FP8 data regardless of params.dtype, // so dispatch the FP8 accumulation kernel in that case. - auto const effective_dtype = params.use_low_precision ? nvinfer1::DataType::kFP8 : params.dtype; + auto const effective_dtype = params.use_low_precision ? tensorrt_llm::DataType::kFP8 : params.dtype; // Launch appropriate kernel with compact macros SWITCH_BOOL(params.enable_rank_mask, ENABLE_RANK_MASK, { diff --git a/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h b/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h index 177293684874..5184878ffc51 100644 --- a/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h +++ b/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h @@ -16,7 +16,7 @@ #pragma once #include "tensorrt_llm/common/config.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include #include @@ -174,8 +174,8 @@ struct MoeA2ACombineParams // Output tensor void* output_data; // Output buffer [local_num_tokens, elements_per_token] // Payload information - int elements_per_token; // Number of elements per token - nvinfer1::DataType dtype; // Data type of the payload (used for combine kernel dispatch) + int elements_per_token; // Number of elements per token + tensorrt_llm::DataType dtype; // Data type of the payload (used for combine kernel dispatch) bool use_low_precision; // If true, prepare kernel quantizes payload→FP8; combine kernel accumulates FP8→output dtype diff --git a/cpp/tensorrt_llm/kernels/cuda_graph_grouped_gemm.cu b/cpp/tensorrt_llm/kernels/cuda_graph_grouped_gemm.cu index 81e947977797..8af43b2b4914 100644 --- a/cpp/tensorrt_llm/kernels/cuda_graph_grouped_gemm.cu +++ b/cpp/tensorrt_llm/kernels/cuda_graph_grouped_gemm.cu @@ -23,6 +23,7 @@ #include "cutlass/cutlass.h" #include "cutlass/gemm/device/gemm_grouped.h" #include "cutlass/gemm/kernel/default_gemm_grouped.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/cutlass_extensions/include/cutlass_extensions/gemm/device/splitk_gemm_grouped.h" #include "tensorrt_llm/cutlass_extensions/include/cutlass_extensions/gemm/kernel/default_splitk_gemm_grouped.h" @@ -117,16 +118,16 @@ void cudaGraphGroupedGemmTemplate(cutlass::gemm::GemmCoord* problemSizesPtr, int template void cudaGraphGroupedGemmType(cutlass::gemm::GemmCoord* problemSizesPtr, int problemCount, void** ptrAGpu, void** ptrBGpu, void** ptrCGpu, void** ptrDGpu, int64_t* ldaGpu, int64_t* ldbGpu, int64_t* ldcGpu, int64_t* lddGpu, - nvinfer1::DataType dataType, cutlass::gemm::GemmCoord* hostMaxProblemSizesPtr, cudaStream_t stream) + tensorrt_llm::DataType dataType, cutlass::gemm::GemmCoord* hostMaxProblemSizesPtr, cudaStream_t stream) { - if (dataType == nvinfer1::DataType::kHALF) + if (dataType == tensorrt_llm::DataType::kHALF) { cudaGraphGroupedGemmTemplate( problemSizesPtr, problemCount, ptrAGpu, ptrBGpu, ptrCGpu, ptrDGpu, ldaGpu, ldbGpu, ldcGpu, lddGpu, hostMaxProblemSizesPtr, stream); } #ifdef ENABLE_BF16 - else if (dataType == nvinfer1::DataType::kBF16) + else if (dataType == tensorrt_llm::DataType::kBF16) { cudaGraphGroupedGemmTemplate( problemSizesPtr, problemCount, ptrAGpu, ptrBGpu, ptrCGpu, ptrDGpu, ldaGpu, ldbGpu, ldcGpu, lddGpu, @@ -141,7 +142,7 @@ void cudaGraphGroupedGemmType(cutlass::gemm::GemmCoord* problemSizesPtr, int pro void cudaGraphGroupedGemm(cutlass::gemm::GemmCoord* problemSizesPtr, int problemCount, void** ptrAGpu, void** ptrBGpu, void** ptrCGpu, void** ptrDGpu, int64_t* ldaGpu, int64_t* ldbGpu, int64_t* ldcGpu, int64_t* lddGpu, bool isLoraIn, - nvinfer1::DataType dataType, int minKN, cutlass::gemm::GemmCoord* hostMaxProblemSizesPtr, cudaStream_t stream) + tensorrt_llm::DataType dataType, int minKN, cutlass::gemm::GemmCoord* hostMaxProblemSizesPtr, cudaStream_t stream) { if (isLoraIn) { @@ -283,17 +284,17 @@ void cudaGraphSplitKGroupedGemmTemplate(cutlass::gemm::GemmCoord* problemSizesPt template void cudaGraphSplitKGroupedGemmType(cutlass::gemm::GemmCoord* problemSizesPtr, int problemCount, void** ptrAGpu, void** ptrBGpu, void** ptrCGpu, void** ptrDGpu, int64_t* ldaGpu, int64_t* ldbGpu, int64_t* ldcGpu, int64_t* lddGpu, - nvinfer1::DataType dataType, int splitKSlices, cutlass::gemm::GemmCoord* hostMaxProblemSizesPtr, + tensorrt_llm::DataType dataType, int splitKSlices, cutlass::gemm::GemmCoord* hostMaxProblemSizesPtr, int64_t* splitKOffsetsGpu, cudaStream_t stream) { - if (dataType == nvinfer1::DataType::kHALF) + if (dataType == tensorrt_llm::DataType::kHALF) { cudaGraphSplitKGroupedGemmTemplate( problemSizesPtr, problemCount, ptrAGpu, ptrBGpu, ptrCGpu, ptrDGpu, ldaGpu, ldbGpu, ldcGpu, lddGpu, splitKSlices, hostMaxProblemSizesPtr, splitKOffsetsGpu, stream); } #ifdef ENABLE_BF16 - else if (dataType == nvinfer1::DataType::kBF16) + else if (dataType == tensorrt_llm::DataType::kBF16) { cudaGraphSplitKGroupedGemmTemplate(problemSizesPtr, problemCount, ptrAGpu, ptrBGpu, ptrCGpu, ptrDGpu, ldaGpu, ldbGpu, ldcGpu, lddGpu, @@ -308,7 +309,7 @@ void cudaGraphSplitKGroupedGemmType(cutlass::gemm::GemmCoord* problemSizesPtr, i void cudaGraphSplitKGroupedGemm(cutlass::gemm::GemmCoord* problemSizesPtr, int problemCount, void** ptrAGpu, void** ptrBGpu, void** ptrCGpu, void** ptrDGpu, int64_t* ldaGpu, int64_t* ldbGpu, int64_t* ldcGpu, int64_t* lddGpu, - bool isLoraIn, nvinfer1::DataType dataType, int splitKSlices, int minKN, + bool isLoraIn, tensorrt_llm::DataType dataType, int splitKSlices, int minKN, cutlass::gemm::GemmCoord* hostMaxProblemSizesPtr, int64_t* splitKOffsetsGpu, cudaStream_t stream) { if (isLoraIn) diff --git a/cpp/tensorrt_llm/kernels/cuda_graph_grouped_gemm.h b/cpp/tensorrt_llm/kernels/cuda_graph_grouped_gemm.h index 0eecccb78852..b447bba3a785 100644 --- a/cpp/tensorrt_llm/kernels/cuda_graph_grouped_gemm.h +++ b/cpp/tensorrt_llm/kernels/cuda_graph_grouped_gemm.h @@ -18,7 +18,7 @@ #include "cutlass/gemm_coord.h" #include "tensorrt_llm/common/config.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include TRTLLM_NAMESPACE_BEGIN @@ -45,7 +45,7 @@ namespace kernels */ void cudaGraphGroupedGemm(cutlass::gemm::GemmCoord* problemSizesPtr, int problemCount, void** ptrAGpu, void** ptrBGpu, void** ptrCGpu, void** ptrDGpu, int64_t* ldaGpu, int64_t* ldbGpu, int64_t* ldcGpu, int64_t* lddGpu, bool isLoraIn, - nvinfer1::DataType dataType, int minKN, cutlass::gemm::GemmCoord* hostMaxProblemSizesPtr, cudaStream_t stream); + tensorrt_llm::DataType dataType, int minKN, cutlass::gemm::GemmCoord* hostMaxProblemSizesPtr, cudaStream_t stream); /** * @brief CUDA Graph compatible wrapper for split-K grouped GEMM operations. @@ -55,7 +55,7 @@ void cudaGraphGroupedGemm(cutlass::gemm::GemmCoord* problemSizesPtr, int problem */ void cudaGraphSplitKGroupedGemm(cutlass::gemm::GemmCoord* problemSizesPtr, int problemCount, void** ptrAGpu, void** ptrBGpu, void** ptrCGpu, void** ptrDGpu, int64_t* ldaGpu, int64_t* ldbGpu, int64_t* ldcGpu, int64_t* lddGpu, - bool isLoraIn, nvinfer1::DataType dataType, int splitKSlices, int minKN, + bool isLoraIn, tensorrt_llm::DataType dataType, int splitKSlices, int minKN, cutlass::gemm::GemmCoord* hostMaxProblemSizesPtr, int64_t* splitKOffsetsGpu, cudaStream_t stream); } // namespace kernels diff --git a/cpp/tensorrt_llm/kernels/customAllReduceKernels.cu b/cpp/tensorrt_llm/kernels/customAllReduceKernels.cu index 9cf2b51eb583..ea217d465be1 100644 --- a/cpp/tensorrt_llm/kernels/customAllReduceKernels.cu +++ b/cpp/tensorrt_llm/kernels/customAllReduceKernels.cu @@ -22,6 +22,7 @@ #include "tensorrt_llm/common/customAllReduceUtils.h" #include "tensorrt_llm/common/dataType.h" #include "tensorrt_llm/common/envUtils.h" +#include "tensorrt_llm/common/tllmDataType.h" #include #include #include @@ -1187,14 +1188,14 @@ bool is_lamport_supported(int token_num, int hidden_size) return true; } -bool is_lamport_supported(nvinfer1::DataType dataType, int token_num, int hidden_size) +bool is_lamport_supported(tensorrt_llm::DataType dataType, int token_num, int hidden_size) { switch (dataType) { - case nvinfer1::DataType::kFLOAT: return is_lamport_supported(token_num, hidden_size); - case nvinfer1::DataType::kHALF: return is_lamport_supported(token_num, hidden_size); + case tensorrt_llm::DataType::kFLOAT: return is_lamport_supported(token_num, hidden_size); + case tensorrt_llm::DataType::kHALF: return is_lamport_supported(token_num, hidden_size); #ifdef ENABLE_BF16 - case nvinfer1::DataType::kBF16: return is_lamport_supported<__nv_bfloat16>(token_num, hidden_size); + case tensorrt_llm::DataType::kBF16: return is_lamport_supported<__nv_bfloat16>(token_num, hidden_size); #endif default: return false; } @@ -1658,7 +1659,7 @@ static __global__ void __launch_bounds__(512, 1) twoShotAllReduceKernel(AllReduc update_barrier_flag(params.barrier_flag_ptr, params.barrier_flag_counter_ptr); } -bool configurationSupported(AllReduceStrategyType algo, size_t msg_size, size_t n_ranks, nvinfer1::DataType type) +bool configurationSupported(AllReduceStrategyType algo, size_t msg_size, size_t n_ranks, tensorrt_llm::DataType type) { size_t elts_per_thread = 16 / common::getDTypeSize(type); int const msg_align = (algo == AllReduceStrategyType::TWOSHOT) ? n_ranks * elts_per_thread : elts_per_thread; @@ -1894,8 +1895,8 @@ void AllReduceDispatchType(AllReduceParams& params, AllReduceStrategyType strat, } } -AllReduceParams AllReduceParams::deserialize(int64_t* buffer, size_t tpSize, size_t tpRank, nvinfer1::DataType dataType, - int token_num, int hidden_size, AllReduceFusionOp op) +AllReduceParams AllReduceParams::deserialize(int64_t* buffer, size_t tpSize, size_t tpRank, + tensorrt_llm::DataType dataType, int token_num, int hidden_size, AllReduceFusionOp op) { void* const* buffer_ptrs = reinterpret_cast(buffer); int flag_offset; @@ -1933,7 +1934,7 @@ AllReduceParams AllReduceParams::deserialize(int64_t* buffer, size_t tpSize, siz return params; } -void customAllReduce(kernels::AllReduceParams& params, nvinfer1::DataType dataType, AllReduceStrategyType strat, +void customAllReduce(kernels::AllReduceParams& params, tensorrt_llm::DataType dataType, AllReduceStrategyType strat, AllReduceStrategyConfig config, AllReduceFusionOp fusionOp, cudaStream_t stream) { TLLM_CHECK_WITH_INFO(configurationSupported(strat, params.elts_total, params.ranks_per_node, dataType), @@ -1943,10 +1944,10 @@ void customAllReduce(kernels::AllReduceParams& params, nvinfer1::DataType dataTy switch (dataType) { - case nvinfer1::DataType::kFLOAT: AllReduceDispatchType(params, strat, config, fusionOp, stream); break; - case nvinfer1::DataType::kHALF: AllReduceDispatchType(params, strat, config, fusionOp, stream); break; + case tensorrt_llm::DataType::kFLOAT: AllReduceDispatchType(params, strat, config, fusionOp, stream); break; + case tensorrt_llm::DataType::kHALF: AllReduceDispatchType(params, strat, config, fusionOp, stream); break; #ifdef ENABLE_BF16 - case nvinfer1::DataType::kBF16: + case tensorrt_llm::DataType::kBF16: AllReduceDispatchType<__nv_bfloat16>(params, strat, config, fusionOp, stream); break; #endif @@ -1991,22 +1992,22 @@ void launchResidualRmsNormKernel(kernels::AllReduceParams& params, cudaStream_t } void residualRmsNorm( - kernels::AllReduceParams& params, nvinfer1::DataType dataType, cudaStream_t stream, AllReduceFusionOp fusionOp) + kernels::AllReduceParams& params, tensorrt_llm::DataType dataType, cudaStream_t stream, AllReduceFusionOp fusionOp) { sync_check_cuda_error(stream); switch (dataType) { - case nvinfer1::DataType::kFLOAT: launchResidualRmsNormKernel(params, stream, fusionOp); break; - case nvinfer1::DataType::kHALF: launchResidualRmsNormKernel(params, stream, fusionOp); break; + case tensorrt_llm::DataType::kFLOAT: launchResidualRmsNormKernel(params, stream, fusionOp); break; + case tensorrt_llm::DataType::kHALF: launchResidualRmsNormKernel(params, stream, fusionOp); break; #ifdef ENABLE_BF16 - case nvinfer1::DataType::kBF16: launchResidualRmsNormKernel<__nv_bfloat16>(params, stream, fusionOp); break; + case tensorrt_llm::DataType::kBF16: launchResidualRmsNormKernel<__nv_bfloat16>(params, stream, fusionOp); break; #endif default: TLLM_THROW("Unsupported dataType for customAllReduce"); } sync_check_cuda_error(stream); } -void lamportInitialize(void* buffer, size_t size, nvinfer1::DataType dataType, cudaStream_t stream) +void lamportInitialize(void* buffer, size_t size, tensorrt_llm::DataType dataType, cudaStream_t stream) { sync_check_cuda_error(stream); if (size == 0) @@ -2015,14 +2016,14 @@ void lamportInitialize(void* buffer, size_t size, nvinfer1::DataType dataType, c } switch (dataType) { - case nvinfer1::DataType::kFLOAT: + case tensorrt_llm::DataType::kFLOAT: reduce_fusion::lamport_initialize_kernel_launcher(buffer, size, stream); break; - case nvinfer1::DataType::kHALF: + case tensorrt_llm::DataType::kHALF: reduce_fusion::lamport_initialize_kernel_launcher(buffer, size, stream); break; #ifdef ENABLE_BF16 - case nvinfer1::DataType::kBF16: + case tensorrt_llm::DataType::kBF16: reduce_fusion::lamport_initialize_kernel_launcher<__nv_bfloat16>(buffer, size, stream); break; #endif diff --git a/cpp/tensorrt_llm/kernels/customAllReduceKernels.h b/cpp/tensorrt_llm/kernels/customAllReduceKernels.h index f7151f1cd0ab..93f67ffdd911 100644 --- a/cpp/tensorrt_llm/kernels/customAllReduceKernels.h +++ b/cpp/tensorrt_llm/kernels/customAllReduceKernels.h @@ -17,7 +17,7 @@ #pragma once #include "tensorrt_llm/common/assert.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include #include #include @@ -178,23 +178,23 @@ struct AllReduceParams AllReduceFusionParams fusion_params; - static AllReduceParams deserialize(int64_t* buffer, size_t tpSize, size_t tpRank, nvinfer1::DataType dataType, + static AllReduceParams deserialize(int64_t* buffer, size_t tpSize, size_t tpRank, tensorrt_llm::DataType dataType, int token_num, int hidden_size, AllReduceFusionOp op); }; -bool configurationSupported(AllReduceStrategyType algo, size_t msg_size, size_t n_ranks, nvinfer1::DataType type); +bool configurationSupported(AllReduceStrategyType algo, size_t msg_size, size_t n_ranks, tensorrt_llm::DataType type); -void customAllReduce(kernels::AllReduceParams& params, nvinfer1::DataType dataType, AllReduceStrategyType strat, +void customAllReduce(kernels::AllReduceParams& params, tensorrt_llm::DataType dataType, AllReduceStrategyType strat, AllReduceStrategyConfig config, AllReduceFusionOp fusionOp, cudaStream_t stream); void residualRmsNorm( - kernels::AllReduceParams& params, nvinfer1::DataType dataType, cudaStream_t stream, AllReduceFusionOp fusionOp); + kernels::AllReduceParams& params, tensorrt_llm::DataType dataType, cudaStream_t stream, AllReduceFusionOp fusionOp); -void lamportInitialize(void* buffer, size_t size, nvinfer1::DataType dataType, cudaStream_t stream); +void lamportInitialize(void* buffer, size_t size, tensorrt_llm::DataType dataType, cudaStream_t stream); namespace reduce_fusion { -bool is_lamport_supported(nvinfer1::DataType dataType, int token_num, int hidden_size); +bool is_lamport_supported(tensorrt_llm::DataType dataType, int token_num, int hidden_size); } } // namespace kernels diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/cutlass_type_conversion.h b/cpp/tensorrt_llm/kernels/cutlass_kernels/cutlass_type_conversion.h index dbbed4e08c97..6632f273cc35 100644 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/cutlass_type_conversion.h +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/cutlass_type_conversion.h @@ -17,7 +17,7 @@ #pragma once #include "tensorrt_llm/common/config.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include "cutlass/half.h" #include @@ -38,34 +38,34 @@ namespace kernels namespace cutlass_kernels { /////////////////////////////////////////////////////////////////////////////////////////////////// -// nvinfer1::DataType to Cutlass +// tensorrt_llm::DataType to Cutlass /////////////////////////////////////////////////////////////////////////////////////////////////// -template +template struct CutlassType { using type = void; }; template <> -struct CutlassType +struct CutlassType { using type = cutlass::half_t; }; template <> -struct CutlassType +struct CutlassType { using type = cutlass::bfloat16_t; }; template <> -struct CutlassType +struct CutlassType { using type = cutlass::float_e4m3_t; }; template <> -struct CutlassType +struct CutlassType { using type = cutlass::float_e2m1_t; }; diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_kernels.h b/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_kernels.h index 24781bec76e7..ab7ed876257d 100644 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_kernels.h +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_kernels.h @@ -27,7 +27,7 @@ #include #endif #include "tensorrt_llm/common/config.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include #include #include @@ -1032,8 +1032,8 @@ struct GemmProfilerBackend using Config = cutlass_extensions::CutlassGemmConfig; using GemmToProfile = MoeGemmId; - void init(CutlassMoeFCRunnerInterface& runner, GemmToProfile gemm_to_profile, nvinfer1::DataType dtype, - nvinfer1::DataType wtype, nvinfer1::DataType otype, int num_experts, int k, int64_t hidden_size, + void init(CutlassMoeFCRunnerInterface& runner, GemmToProfile gemm_to_profile, tensorrt_llm::DataType dtype, + tensorrt_llm::DataType wtype, tensorrt_llm::DataType otype, int num_experts, int k, int64_t hidden_size, int64_t unpadded_hidden_size, int64_t inter_size, int64_t group_size, ActivationType activation_type, bool bias, bool use_lora, bool min_latency_mode, bool need_weights, MOEParallelismConfig parallelism_config, bool const enable_alltoall, bool use_mxfp8_weight_scaling = false) @@ -1061,20 +1061,21 @@ struct GemmProfilerBackend mSM = common::getSMVersion(); mScalingType = TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::NONE; - if (dtype == nvinfer1::DataType::kFP8 - && (wtype == nvinfer1::DataType::kFP4 || wtype == nvinfer1::DataType::kINT64)) + if (dtype == tensorrt_llm::DataType::kFP8 + && (wtype == tensorrt_llm::DataType::kFP4 || wtype == tensorrt_llm::DataType::kINT64)) { mScalingType = TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::MXFPX; } - else if (dtype == nvinfer1::DataType::kFP8 && wtype == nvinfer1::DataType::kFP8 && use_mxfp8_weight_scaling) + else if (dtype == tensorrt_llm::DataType::kFP8 && wtype == tensorrt_llm::DataType::kFP8 + && use_mxfp8_weight_scaling) { // MXFP8 W8A8: e4m3 acts × e4m3 weights with UE8M0 1x32 block scales on both sides. // Profiler must produce MXFPX block-scaled inputs (otherwise the per-expert SF // pointer arrays stay uninitialized and the kernel reads garbage SF addresses). mScalingType = TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::MXFPX; } - else if ((dtype == nvinfer1::DataType::kFP4 || dtype == nvinfer1::DataType::kINT64) - && (wtype == nvinfer1::DataType::kFP4 || wtype == nvinfer1::DataType::kINT64)) + else if ((dtype == tensorrt_llm::DataType::kFP4 || dtype == tensorrt_llm::DataType::kINT64) + && (wtype == tensorrt_llm::DataType::kFP4 || wtype == tensorrt_llm::DataType::kINT64)) { mScalingType = TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::NVFP4; } @@ -1106,9 +1107,9 @@ struct GemmProfilerBackend int mSampleIndex = 0; - nvinfer1::DataType mDType{}; - nvinfer1::DataType mWType{}; - nvinfer1::DataType mOType{}; + tensorrt_llm::DataType mDType{}; + tensorrt_llm::DataType mWType{}; + tensorrt_llm::DataType mOType{}; // This will be a unique value for every iteration of warmup and actual bench constexpr static int64_t NUM_ROUTING_SAMPLES = 16; diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_lora_grouped_gemm.h b/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_lora_grouped_gemm.h index 55ab4e40a3ae..15af49fc1839 100644 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_lora_grouped_gemm.h +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_lora_grouped_gemm.h @@ -18,7 +18,7 @@ #include "tensorrt_llm/common/config.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include #include @@ -56,7 +56,7 @@ struct MoeLoraGroupedGemmModule; // stream: CUDA stream to launch onto. using MoeLoraGroupedGemmRunFn = void (*)(MoeLoraGroupedGemmModule const& mod, int64_t num_permuted_tokens, int64_t in_hidden_size, int64_t max_lora_rank, int64_t dtype_bytes, int64_t splitk_slices, void const* input_base, - void* output_base, nvinfer1::DataType data_type, cudaStream_t stream); + void* output_base, tensorrt_llm::DataType data_type, cudaStream_t stream); // Per-module device-resident scratch for the MoE LoRA capture-safe path. // Pointers refer to device memory unless noted. diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_util_kernels.h b/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_util_kernels.h index e902e2c9d6d3..f3e8940b0c28 100644 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_util_kernels.h +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_util_kernels.h @@ -25,7 +25,6 @@ #ifdef ENABLE_FP4 #include #endif -#include #include #include #include diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu b/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu index b7a32be2e285..8bed9c16b58e 100644 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu @@ -58,6 +58,7 @@ #include "tensorrt_llm/kernels/preQuantScaleKernel.h" #include "tensorrt_llm/kernels/quantization.cuh" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/kernels/cutlass_kernels/include/moe_lora_pointer_expand.h" #include "tensorrt_llm/kernels/cutlass_kernels/include/moe_util_kernels.h" // NOTE: the grouped-GEMM dispatch (cudaGraph(SplitK)GroupedGemm, @@ -3686,7 +3687,7 @@ void CutlassMoeFCRunner -constexpr nvinfer1::DataType moeLoraNvInferType() +constexpr tensorrt_llm::DataType moeLoraDataType() { if constexpr (std::is_same_v) { - return nvinfer1::DataType::kHALF; + return tensorrt_llm::DataType::kHALF; } #if defined(ENABLE_BF16) else if constexpr (std::is_same_v) { - return nvinfer1::DataType::kBF16; + return tensorrt_llm::DataType::kBF16; } #endif else if constexpr (std::is_same_v) { - return nvinfer1::DataType::kFLOAT; + return tensorrt_llm::DataType::kFLOAT; } else { @@ -3925,7 +3926,7 @@ auto CutlassMoeFCRunner(); + tensorrt_llm::DataType const data_type = moeLoraDataType(); // The grouped-GEMM GEMM skips rank-0 rows, but the bias/reorder paths // read lora_fc1_result_ for every valid row. Zero the buffer first so @@ -4020,7 +4021,7 @@ void CutlassMoeFCRunner(); + tensorrt_llm::DataType const data_type = moeLoraDataType(); // As in loraFC1, zero the output so rank-0 rows the GEMM skips do not // feed stale data into the downstream add. @@ -4750,18 +4751,18 @@ std::map> GemmProfilerBackend::getProfile size_t k = mK; size_t num_expanded_tokens = mMinLatencyMode ? maxM * mNumExpertsPerNode : maxM * k; - TLLM_CHECK(mDType != nvinfer1::DataType::kINT4); + TLLM_CHECK(mDType != tensorrt_llm::DataType::kINT4); // nvllm still uses int64 because torch doesn't have fp4 yet. - bool is_4bit_act = mDType == nvinfer1::DataType::kFP4 || mDType == nvinfer1::DataType::kINT64; - bool is_4bit_weight = mWType == nvinfer1::DataType::kINT4 || mWType == nvinfer1::DataType::kFP4 - || mWType == nvinfer1::DataType::kINT64; + bool is_4bit_act = mDType == tensorrt_llm::DataType::kFP4 || mDType == tensorrt_llm::DataType::kINT64; + bool is_4bit_weight = mWType == tensorrt_llm::DataType::kINT4 || mWType == tensorrt_llm::DataType::kFP4 + || mWType == tensorrt_llm::DataType::kINT64; TLLM_CHECK_WITH_INFO(!is_4bit_act || is_4bit_weight, "Cannot have 4-bit activation with non-4-bit weight"); float dtype_bytes = is_4bit_act ? 0.5f - : static_cast(mWType == nvinfer1::DataType::kINT4 ? getDTypeSize(mOType) : getDTypeSize(mDType)); + : static_cast(mWType == tensorrt_llm::DataType::kINT4 ? getDTypeSize(mOType) : getDTypeSize(mDType)); float weight_bytes = is_4bit_weight ? 0.5f : static_cast(getDTypeSize(mWType)); size_t output_bytes = getDTypeSize(mOType); - size_t gemm_output_bytes = (mOType == nvinfer1::DataType::kFP8) + size_t gemm_output_bytes = (mOType == tensorrt_llm::DataType::kFP8) ? sizeof(TmaWarpSpecializedGroupedGemmInput::OutputTypeAdaptor_t<__nv_fp8_e4m3>) : output_bytes; @@ -4803,18 +4804,18 @@ std::map> GemmProfilerBackend::getProfile // TODO Make quant 2 & 4 bigger for FP8 if we ever change to scaling per expert bool is_int_w_quant - = (mWType == nvinfer1::DataType::kINT8 || mWType == nvinfer1::DataType::kINT4) && mGroupSize <= 0; + = (mWType == tensorrt_llm::DataType::kINT8 || mWType == tensorrt_llm::DataType::kINT4) && mGroupSize <= 0; bool is_int_groupwise_w_quant - = (mWType == nvinfer1::DataType::kINT8 || mWType == nvinfer1::DataType::kINT4) && mGroupSize > 0; - bool is_fp8_act_quant = mDType == nvinfer1::DataType::kFP8; - bool is_fp8_w_quant = mWType == nvinfer1::DataType::kFP8; + = (mWType == tensorrt_llm::DataType::kINT8 || mWType == tensorrt_llm::DataType::kINT4) && mGroupSize > 0; + bool is_fp8_act_quant = mDType == tensorrt_llm::DataType::kFP8; + bool is_fp8_w_quant = mWType == tensorrt_llm::DataType::kFP8; // nvllm still uses int64 because torch doesn't have fp4 yet. - // bool is_fp4_act_quant = mDType == nvinfer1::DataType::kFP4 || mDType == nvinfer1::DataType::kINT64; - bool is_fp4_w_quant = mWType == nvinfer1::DataType::kFP4 || mWType == nvinfer1::DataType::kINT64; + // bool is_fp4_act_quant = mDType == tensorrt_llm::DataType::kFP4 || mDType == tensorrt_llm::DataType::kINT64; + bool is_fp4_w_quant = mWType == tensorrt_llm::DataType::kFP4 || mWType == tensorrt_llm::DataType::kINT64; bool is_w4afp8_quant = is_int_groupwise_w_quant && is_fp8_act_quant; // bool is_wfp4afp8_quant = is_fp4_w_quant && is_fp8_act_quant; - bool is_wfp4a16_quant = (mDType == nvinfer1::DataType::kHALF || mDType == nvinfer1::DataType::kBF16) - && mWType == nvinfer1::DataType::kUINT8; + bool is_wfp4a16_quant = (mDType == tensorrt_llm::DataType::kHALF || mDType == tensorrt_llm::DataType::kBF16) + && mWType == tensorrt_llm::DataType::kUINT8; // Int sizes size_t quant_1_size = is_int_w_quant ? fc1_out_size * num_experts_per_node * dtype_bytes : 0; @@ -5047,19 +5048,19 @@ void GemmProfilerBackend::prepareQuantParams(int num_tokens, char* workspace_ptr GET_WS_PTR(float const*, w4a8_alpha); #undef GET_WS_PTR - if ((mWType == nvinfer1::DataType::kINT8 || mWType == nvinfer1::DataType::kINT4 - || mWType == nvinfer1::DataType::kUINT8) + if ((mWType == tensorrt_llm::DataType::kINT8 || mWType == tensorrt_llm::DataType::kINT4 + || mWType == tensorrt_llm::DataType::kUINT8) && mGroupSize < 0) { TLLM_CHECK(quant_1 && quant_2); mQuantParams = QuantParams::Int(quant_1, quant_2); } - else if (mWType == nvinfer1::DataType::kINT4 || mWType == nvinfer1::DataType::kUINT8) + else if (mWType == tensorrt_llm::DataType::kINT4 || mWType == tensorrt_llm::DataType::kUINT8) { TLLM_CHECK(quant_1 && quant_2); - if (mDType == nvinfer1::DataType::kFP8 - || (mWType == nvinfer1::DataType::kUINT8 - && (mDType == nvinfer1::DataType::kHALF || mDType == nvinfer1::DataType::kBF16))) + if (mDType == tensorrt_llm::DataType::kFP8 + || (mWType == tensorrt_llm::DataType::kUINT8 + && (mDType == tensorrt_llm::DataType::kHALF || mDType == tensorrt_llm::DataType::kBF16))) { TLLM_CHECK(w4a8_alpha); mQuantParams = QuantParams::GroupWise( @@ -5070,7 +5071,7 @@ void GemmProfilerBackend::prepareQuantParams(int num_tokens, char* workspace_ptr mQuantParams = QuantParams::GroupWise(mGroupSize, quant_1, quant_2, nullptr, nullptr, quant_3, quant_4); } } - else if (mWType == nvinfer1::DataType::kFP8) + else if (mWType == tensorrt_llm::DataType::kFP8) { if (mUseMxfp8WeightScaling) { @@ -5089,8 +5090,8 @@ void GemmProfilerBackend::prepareQuantParams(int num_tokens, char* workspace_ptr static_cast(quant_3), static_cast(quant_4)); } } - else if (mDType == nvinfer1::DataType::kFP8 - && (mWType == nvinfer1::DataType::kFP4 || mWType == nvinfer1::DataType::kINT64)) + else if (mDType == tensorrt_llm::DataType::kFP8 + && (mWType == tensorrt_llm::DataType::kFP4 || mWType == tensorrt_llm::DataType::kINT64)) { TLLM_CHECK(quant_1 && quant_2 && quant_3 && quant_4 && quant_5 && quant_6); mQuantParams = QuantParams::FP8MXFP4(static_cast(quant_1), @@ -5099,8 +5100,8 @@ void GemmProfilerBackend::prepareQuantParams(int num_tokens, char* workspace_ptr static_cast(quant_5), static_cast(quant_6)); } - else if ((mDType == nvinfer1::DataType::kFP4 || mDType == nvinfer1::DataType::kINT64) - && (mWType == nvinfer1::DataType::kFP4 || mWType == nvinfer1::DataType::kINT64)) + else if ((mDType == tensorrt_llm::DataType::kFP4 || mDType == tensorrt_llm::DataType::kINT64) + && (mWType == tensorrt_llm::DataType::kFP4 || mWType == tensorrt_llm::DataType::kINT64)) { // nvllm still uses int64 because torch doesn't have fp4 yet. TLLM_CHECK(quant_1 && quant_2 && quant_3 && quant_4 && quant_5 && quant_6); @@ -5120,9 +5121,9 @@ void GemmProfilerBackend::prepareTmaWsInputs(int num_tokens, char* workspace_ptr return; } - bool use_w4afp8 = (mDType == nvinfer1::DataType::kFP8 && mWType == nvinfer1::DataType::kINT4); - bool use_wfp4a16 = ((mDType == nvinfer1::DataType::kHALF || mDType == nvinfer1::DataType::kBF16) - && mWType == nvinfer1::DataType::kUINT8); + bool use_w4afp8 = (mDType == tensorrt_llm::DataType::kFP8 && mWType == tensorrt_llm::DataType::kINT4); + bool use_wfp4a16 = ((mDType == tensorrt_llm::DataType::kHALF || mDType == tensorrt_llm::DataType::kBF16) + && mWType == tensorrt_llm::DataType::kUINT8); bool const use_finalize_fusion = fusion == TmaWarpSpecializedGroupedGemmInput::EpilogueFusion::FINALIZE; bool const finalize_fusion_not_supported = !mInterface->use_fused_finalize_ || mMinLatencyMode || use_wfp4a16 || mGemmToProfile != GemmToProfile::GEMM_2; diff --git a/cpp/tensorrt_llm/kernels/gptKernels.h b/cpp/tensorrt_llm/kernels/gptKernels.h index e13e9bca4d6a..d855aade79c2 100644 --- a/cpp/tensorrt_llm/kernels/gptKernels.h +++ b/cpp/tensorrt_llm/kernels/gptKernels.h @@ -16,6 +16,7 @@ #pragma once #include "tensorrt_llm/common/config.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/kernels/contextFusedMultiHeadAttention/fused_multihead_attention_common.h" #include "tensorrt_llm/runtime/iTensor.h" #include @@ -227,11 +228,11 @@ struct BuildDecoderInfoParams std::string toString() const { std::stringstream ss; - auto printTensor = [&ss](char const* name, void* ptr, nvinfer1::Dims shape) + auto printTensor = [&ss](char const* name, void* ptr, tensorrt_llm::Dims shape) { ss << name << ": "; if (ptr) - ss << *(runtime::ITensor::wrap((void*) ptr, nvinfer1::DataType::kINT32, shape)); + ss << *(runtime::ITensor::wrap((void*) ptr, tensorrt_llm::DataType::kINT32, shape)); else ss << "nullptr"; ss << std::endl; diff --git a/cpp/tensorrt_llm/kernels/groupGemm.cu b/cpp/tensorrt_llm/kernels/groupGemm.cu index 5b8c0d929150..b41021ffe6f7 100644 --- a/cpp/tensorrt_llm/kernels/groupGemm.cu +++ b/cpp/tensorrt_llm/kernels/groupGemm.cu @@ -28,6 +28,7 @@ #include "tensorrt_llm/common/cudaUtils.h" #include "tensorrt_llm/common/memoryUtils.h" +#include "tensorrt_llm/common/tllmDataType.h" TRTLLM_NAMESPACE_BEGIN @@ -63,7 +64,7 @@ template problem_sizes, std::vector const& ptrA, std::vector const& ptrB, std::vector const& ptrC, std::vector const& ptrD, void* gemmParamsWorkSpace, int64_t gemmParamsWorkSpaceSize, void* gemmWorkSpace, int64_t gemmWorkspaceSize, - nvinfer1::DataType dataType, cudaStream_t stream) + tensorrt_llm::DataType dataType, cudaStream_t stream) { TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); using ElementA = cutlassType; @@ -178,20 +179,20 @@ template problem_sizes, std::vector const& ptrA, std::vector const& ptrB, std::vector const& ptrC, std::vector const& ptrD, void* gemmParamsWorkSpace, int64_t gemmParamsWorkSpaceSize, void* gemmWorkSpace, int64_t gemmWorkspaceSize, - nvinfer1::DataType dataType, cudaStream_t stream) + tensorrt_llm::DataType dataType, cudaStream_t stream) { - if (dataType == nvinfer1::DataType::kHALF) + if (dataType == tensorrt_llm::DataType::kHALF) { groupedGemm_(problem_sizes, ptrA, ptrB, ptrC, ptrD, gemmParamsWorkSpace, gemmParamsWorkSpaceSize, gemmWorkSpace, gemmWorkspaceSize, dataType, stream); } - else if (dataType == nvinfer1::DataType::kFLOAT) + else if (dataType == tensorrt_llm::DataType::kFLOAT) { TLLM_CHECK_WITH_INFO(false, "not support float input/output"); } #ifdef ENABLE_BF16 - else if (dataType == nvinfer1::DataType::kBF16) + else if (dataType == tensorrt_llm::DataType::kBF16) { groupedGemm_(problem_sizes, ptrA, ptrB, ptrC, ptrD, gemmParamsWorkSpace, gemmParamsWorkSpaceSize, gemmWorkSpace, gemmWorkspaceSize, @@ -203,7 +204,7 @@ void groupedGemmType_(std::vector problem_sizes, std:: void groupedGemm(std::vector problem_sizes, std::vector const& ptrA, std::vector const& ptrB, std::vector const& ptrC, std::vector const& ptrD, void* gemmParamsWorkSpace, int64_t gemmParamsWorkSpaceSize, void* gemmWorkSpace, int64_t gemmWorkspaceSize, - bool isLoraIn, nvinfer1::DataType dataType, int minKN, cudaStream_t stream) + bool isLoraIn, tensorrt_llm::DataType dataType, int minKN, cudaStream_t stream) { TLLM_LOG_TRACE("%s start, isLoraIn: %d, minKN = %d", __PRETTY_FUNCTION__, static_cast(isLoraIn), minKN); if (isLoraIn) diff --git a/cpp/tensorrt_llm/kernels/groupGemm.h b/cpp/tensorrt_llm/kernels/groupGemm.h index dbc1e498b7b2..c526e08e986c 100644 --- a/cpp/tensorrt_llm/kernels/groupGemm.h +++ b/cpp/tensorrt_llm/kernels/groupGemm.h @@ -17,7 +17,7 @@ #include "cutlass/gemm_coord.h" #include "tensorrt_llm/common/config.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" TRTLLM_NAMESPACE_BEGIN @@ -29,7 +29,7 @@ int64_t getGroupedGemmParamsWorkSpaceSize(int64_t problem_count); void groupedGemm(std::vector problem_sizes, std::vector const& ptrA, std::vector const& ptrB, std::vector const& ptrC, std::vector const& ptrD, void* gemmParamsWorkspace, int64_t gemmParamsWorkSpaceSize, void* gemmWorkSpace, int64_t gemmWorkspaceSize, - bool isLoraIn, nvinfer1::DataType dataType, int minKN, cudaStream_t stream); + bool isLoraIn, tensorrt_llm::DataType dataType, int minKN, cudaStream_t stream); } // namespace kernels diff --git a/cpp/tensorrt_llm/kernels/groupRmsNormKernels/groupRmsNormKernels.cu b/cpp/tensorrt_llm/kernels/groupRmsNormKernels/groupRmsNormKernels.cu index 409968bb510d..3332c2918a3f 100644 --- a/cpp/tensorrt_llm/kernels/groupRmsNormKernels/groupRmsNormKernels.cu +++ b/cpp/tensorrt_llm/kernels/groupRmsNormKernels/groupRmsNormKernels.cu @@ -22,6 +22,7 @@ #include "tensorrt_llm/common/envUtils.h" #include "tensorrt_llm/common/memoryUtils.h" #include "tensorrt_llm/common/reduceKernelUtils.cuh" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/kernels/groupRmsNormKernels/groupRmsNormKernels.h" TRTLLM_NAMESPACE_BEGIN @@ -655,9 +656,9 @@ void GroupRMSNormBaseKernelLauncher(GroupRMSParams& params) switch (params.dtype) { - case nvinfer1::DataType::kHALF: GROUP_RMS_NORM_DISPATCH(half); break; - case nvinfer1::DataType::kBF16: GROUP_RMS_NORM_DISPATCH(__nv_bfloat16); break; - case nvinfer1::DataType::kFLOAT: GROUP_RMS_NORM_DISPATCH(float); break; + case tensorrt_llm::DataType::kHALF: GROUP_RMS_NORM_DISPATCH(half); break; + case tensorrt_llm::DataType::kBF16: GROUP_RMS_NORM_DISPATCH(__nv_bfloat16); break; + case tensorrt_llm::DataType::kFLOAT: GROUP_RMS_NORM_DISPATCH(float); break; default: TLLM_CHECK_WITH_INFO(false, "Unsupported data type for GroupRMSNorm"); } @@ -750,9 +751,9 @@ void GroupRMSNormKernelLargeBatchLauncher(GroupRMSParams& params) switch (params.dtype) { - case nvinfer1::DataType::kHALF: GROUP_RMS_NORM_LARGE_BATCH_DISPATCH(half); break; - case nvinfer1::DataType::kBF16: GROUP_RMS_NORM_LARGE_BATCH_DISPATCH(__nv_bfloat16); break; - case nvinfer1::DataType::kFLOAT: GROUP_RMS_NORM_LARGE_BATCH_DISPATCH(float); break; + case tensorrt_llm::DataType::kHALF: GROUP_RMS_NORM_LARGE_BATCH_DISPATCH(half); break; + case tensorrt_llm::DataType::kBF16: GROUP_RMS_NORM_LARGE_BATCH_DISPATCH(__nv_bfloat16); break; + case tensorrt_llm::DataType::kFLOAT: GROUP_RMS_NORM_LARGE_BATCH_DISPATCH(float); break; default: TLLM_CHECK_WITH_INFO(false, "Unsupported data type for GroupRMSNormV2"); } @@ -813,15 +814,15 @@ void GroupRMSNormKernelLauncherWithHeuristic(GroupRMSParams& params) // Choose the appropriate DType switch (params.dtype) { - case nvinfer1::DataType::kHALF: + case tensorrt_llm::DataType::kHALF: base_warps = calculateNumWarpsBase(params); large_batch_warps = calculateNumWarpsLargeBatch(params).num_warps_to_launch; break; - case nvinfer1::DataType::kBF16: + case tensorrt_llm::DataType::kBF16: base_warps = calculateNumWarpsBase<__nv_bfloat16, n>(params); large_batch_warps = calculateNumWarpsLargeBatch<__nv_bfloat16, n>(params).num_warps_to_launch; break; - case nvinfer1::DataType::kFLOAT: + case tensorrt_llm::DataType::kFLOAT: base_warps = calculateNumWarpsBase(params); large_batch_warps = calculateNumWarpsLargeBatch(params).num_warps_to_launch; break; diff --git a/cpp/tensorrt_llm/kernels/groupRmsNormKernels/groupRmsNormKernels.h b/cpp/tensorrt_llm/kernels/groupRmsNormKernels/groupRmsNormKernels.h index 335adf44ed67..70425f924217 100644 --- a/cpp/tensorrt_llm/kernels/groupRmsNormKernels/groupRmsNormKernels.h +++ b/cpp/tensorrt_llm/kernels/groupRmsNormKernels/groupRmsNormKernels.h @@ -15,7 +15,7 @@ */ #pragma once #include "tensorrt_llm/common/assert.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include #include #include @@ -44,7 +44,7 @@ struct GroupRMSParams float eps; float weight_bias; bool enable_weights; - nvinfer1::DataType dtype; + tensorrt_llm::DataType dtype; cudaStream_t stream; }; diff --git a/cpp/tensorrt_llm/kernels/internal_cutlass_kernels/include/moe_kernels.h b/cpp/tensorrt_llm/kernels/internal_cutlass_kernels/include/moe_kernels.h index 132990603db9..0b02686f5e53 100644 --- a/cpp/tensorrt_llm/kernels/internal_cutlass_kernels/include/moe_kernels.h +++ b/cpp/tensorrt_llm/kernels/internal_cutlass_kernels/include/moe_kernels.h @@ -26,7 +26,7 @@ #ifdef ENABLE_FP4 #include #endif -#include +#include "tensorrt_llm/common/tllmDataType.h" #include #include #include @@ -869,8 +869,8 @@ struct GemmProfilerBackend GEMM_2 }; - void init(CutlassMoeFCRunnerInterface& runner, GemmToProfile gemm_to_profile, nvinfer1::DataType dtype, - nvinfer1::DataType wtype, nvinfer1::DataType otype, int num_experts, int k, int64_t hidden_size, + void init(CutlassMoeFCRunnerInterface& runner, GemmToProfile gemm_to_profile, tensorrt_llm::DataType dtype, + tensorrt_llm::DataType wtype, tensorrt_llm::DataType otype, int num_experts, int k, int64_t hidden_size, int64_t inter_size, int64_t group_size, ActivationType activation_type, bool bias, bool use_lora, bool min_latency_mode, bool need_weights, MOEParallelismConfig parallelism_config) { @@ -895,13 +895,13 @@ struct GemmProfilerBackend mSorter.updateNumExperts(mNumExpertsPerNode); mScalingType = TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::NONE; - if (dtype == nvinfer1::DataType::kFP8 - && (wtype == nvinfer1::DataType::kFP4 || wtype == nvinfer1::DataType::kINT64)) + if (dtype == tensorrt_llm::DataType::kFP8 + && (wtype == tensorrt_llm::DataType::kFP4 || wtype == tensorrt_llm::DataType::kINT64)) { mScalingType = TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::MXFPX; } - else if ((dtype == nvinfer1::DataType::kFP4 || dtype == nvinfer1::DataType::kINT64) - && (wtype == nvinfer1::DataType::kFP4 || wtype == nvinfer1::DataType::kINT64)) + else if ((dtype == tensorrt_llm::DataType::kFP4 || dtype == tensorrt_llm::DataType::kINT64) + && (wtype == tensorrt_llm::DataType::kFP4 || wtype == tensorrt_llm::DataType::kINT64)) { mScalingType = TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::NVFP4; } @@ -932,9 +932,9 @@ struct GemmProfilerBackend int mSampleIndex = 0; - nvinfer1::DataType mDType{}; - nvinfer1::DataType mWType{}; - nvinfer1::DataType mOType{}; + tensorrt_llm::DataType mDType{}; + tensorrt_llm::DataType mWType{}; + tensorrt_llm::DataType mOType{}; // This will be a unique value for every iteration of warmup and actual bench constexpr static int64_t NUM_ROUTING_SAMPLES = 16; diff --git a/cpp/tensorrt_llm/kernels/kvCachePartialCopy.cu b/cpp/tensorrt_llm/kernels/kvCachePartialCopy.cu index 3b91cf3f1776..04105721dfca 100644 --- a/cpp/tensorrt_llm/kernels/kvCachePartialCopy.cu +++ b/cpp/tensorrt_llm/kernels/kvCachePartialCopy.cu @@ -15,6 +15,7 @@ */ #include "tensorrt_llm/common/config.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/kernels/kvCachePartialCopy.h" #include #include @@ -89,42 +90,42 @@ void kvCacheBlockPartialCopy(IBuffer& dst, IBuffer const& src, unsigned int numL TLLM_CHECK_WITH_INFO(dataType == dst.getDataType(), "src and dst dataType does not match"); switch (dataType) { - case nvinfer1::DataType::kINT64: + case tensorrt_llm::DataType::kINT64: hostKVCacheBlockPartialCopy( dst, src, numLayers, numHeads, tokensPerBlock, numHidden, numTokensToCopy, kvFactor, stream); break; - case nvinfer1::DataType::kINT32: + case tensorrt_llm::DataType::kINT32: hostKVCacheBlockPartialCopy( dst, src, numLayers, numHeads, tokensPerBlock, numHidden, numTokensToCopy, kvFactor, stream); break; - case nvinfer1::DataType::kFLOAT: + case tensorrt_llm::DataType::kFLOAT: hostKVCacheBlockPartialCopy( dst, src, numLayers, numHeads, tokensPerBlock, numHidden, numTokensToCopy, kvFactor, stream); break; #ifdef ENABLE_BF16 - case nvinfer1::DataType::kBF16: + case tensorrt_llm::DataType::kBF16: hostKVCacheBlockPartialCopy<__nv_bfloat16>( dst, src, numLayers, numHeads, tokensPerBlock, numHidden, numTokensToCopy, kvFactor, stream); break; #endif - case nvinfer1::DataType::kHALF: + case tensorrt_llm::DataType::kHALF: hostKVCacheBlockPartialCopy( dst, src, numLayers, numHeads, tokensPerBlock, numHidden, numTokensToCopy, kvFactor, stream); break; - case nvinfer1::DataType::kBOOL: + case tensorrt_llm::DataType::kBOOL: hostKVCacheBlockPartialCopy( dst, src, numLayers, numHeads, tokensPerBlock, numHidden, numTokensToCopy, kvFactor, stream); break; - case nvinfer1::DataType::kUINT8: + case tensorrt_llm::DataType::kUINT8: hostKVCacheBlockPartialCopy( dst, src, numLayers, numHeads, tokensPerBlock, numHidden, numTokensToCopy, kvFactor, stream); break; - case nvinfer1::DataType::kINT8: + case tensorrt_llm::DataType::kINT8: hostKVCacheBlockPartialCopy( dst, src, numLayers, numHeads, tokensPerBlock, numHidden, numTokensToCopy, kvFactor, stream); break; #ifdef ENABLE_FP8 - case nvinfer1::DataType::kFP8: + case tensorrt_llm::DataType::kFP8: hostKVCacheBlockPartialCopy<__nv_fp8_e4m3>( dst, src, numLayers, numHeads, tokensPerBlock, numHidden, numTokensToCopy, kvFactor, stream); break; diff --git a/cpp/tensorrt_llm/kernels/lora/dora.cpp b/cpp/tensorrt_llm/kernels/lora/dora.cpp index 43dbf4fdccb0..883d02df9291 100644 --- a/cpp/tensorrt_llm/kernels/lora/dora.cpp +++ b/cpp/tensorrt_llm/kernels/lora/dora.cpp @@ -20,13 +20,13 @@ #include "tensorrt_llm/common/memoryUtils.h" #include "tensorrt_llm/kernels/doraScaling.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include #include using tensorrt_llm::kernels::DoraImpl; -DoraImpl::DoraImpl(std::vector const& outHiddenSizes, nvinfer1::DataType type) +DoraImpl::DoraImpl(std::vector const& outHiddenSizes, tensorrt_llm::DataType type) : mType(type) { mCumModuleSizes.resize(outHiddenSizes.size()); @@ -73,14 +73,14 @@ int DoraImpl::run(int64_t numTokens, void const* input, void const* const* loraW auto const* deviceCumModuleSizes = reinterpret_cast(workspace); auto const* deviceScalePtrs = reinterpret_cast((&deviceCumModuleSizes[numModules])); - if (mType == nvinfer1::DataType::kHALF) + if (mType == tensorrt_llm::DataType::kHALF) { tokenPerChannelScale(numel, numModules, numTokens, deviceCumModuleSizes, reinterpret_cast(input), reinterpret_cast(deviceScalePtrs), reinterpret_cast(outputs[0]), stream); } #ifdef ENABLE_BF16 - else if (mType == nvinfer1::DataType::kBF16) + else if (mType == tensorrt_llm::DataType::kBF16) { tokenPerChannelScale(numel, numModules, numTokens, deviceCumModuleSizes, reinterpret_cast(input), reinterpret_cast(deviceScalePtrs), diff --git a/cpp/tensorrt_llm/kernels/lora/dora.h b/cpp/tensorrt_llm/kernels/lora/dora.h index fc21fe669366..02cd68e7c1f0 100644 --- a/cpp/tensorrt_llm/kernels/lora/dora.h +++ b/cpp/tensorrt_llm/kernels/lora/dora.h @@ -17,7 +17,7 @@ #include "tensorrt_llm/common/config.h" #include "tensorrt_llm/common/cudaUtils.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" TRTLLM_NAMESPACE_BEGIN @@ -28,7 +28,7 @@ class DoraImpl public: DoraImpl() = delete; - DoraImpl(std::vector const& outHiddenSizes, nvinfer1::DataType type); + DoraImpl(std::vector const& outHiddenSizes, tensorrt_llm::DataType type); ~DoraImpl() = default; @@ -41,7 +41,7 @@ class DoraImpl private: std::vector mCumModuleSizes; std::vector mHostBuf; - nvinfer1::DataType mType; + tensorrt_llm::DataType mType; }; } // namespace kernels diff --git a/cpp/tensorrt_llm/kernels/lora/lora.cpp b/cpp/tensorrt_llm/kernels/lora/lora.cpp index 61f6af00fedc..7a2b7d330afc 100644 --- a/cpp/tensorrt_llm/kernels/lora/lora.cpp +++ b/cpp/tensorrt_llm/kernels/lora/lora.cpp @@ -19,6 +19,7 @@ #include "tensorrt_llm/common/config.h" #include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/kernels/groupGemm.h" #include "tensorrt_llm/kernels/lora/lora.h" #include "tensorrt_llm/kernels/splitkGroupGemm.h" @@ -48,8 +49,9 @@ void _getProblemParams(cublasOperation_t& transa, cublasOperation_t& transb, int // TODO should reuse the function in gemmPlugin void _runGemm(int const M, int const N, int const K, bool const transA, bool const transB, - nvinfer1::DataType const type, CublasGemmWrapperPtr const& cublasWrapperPtr, void const* act, void const* weight, - void* output, std::optional const& heuristic, void* workspace, cudaStream_t stream) + tensorrt_llm::DataType const type, CublasGemmWrapperPtr const& cublasWrapperPtr, void const* act, + void const* weight, void* output, std::optional const& heuristic, void* workspace, + cudaStream_t stream) { cublasWrapperPtr->setStream(stream); cublasWrapperPtr->setWorkspace(workspace); @@ -65,7 +67,8 @@ void _runGemm(int const M, int const N, int const K, bool const transA, bool con } LoraImpl::LoraImpl(int in_hidden_size, std::vector out_hidden_sizes, bool transA, bool transB, - int num_lora_modules, nvinfer1::DataType type, int max_low_rank, std::shared_ptr cublasWrapper) + int num_lora_modules, tensorrt_llm::DataType type, int max_low_rank, + std::shared_ptr cublasWrapper) : mInHiddenSize(in_hidden_size) , mTransA(transA) , mTransB(transB) @@ -82,16 +85,16 @@ LoraImpl::LoraImpl(int in_hidden_size, std::vector out_hidden_sizes, bool t void LoraImpl::setGemmConfig() { TLLM_LOG_DEBUG("%s", __PRETTY_FUNCTION__); - if (mType == nvinfer1::DataType::kHALF) + if (mType == tensorrt_llm::DataType::kHALF) { mCublasWrapper->setFP16GemmConfig(); } - else if (mType == nvinfer1::DataType::kFLOAT) + else if (mType == tensorrt_llm::DataType::kFLOAT) { mCublasWrapper->setFP32GemmConfig(); } #ifdef ENABLE_BF16 - else if (mType == nvinfer1::DataType::kBF16) + else if (mType == tensorrt_llm::DataType::kBF16) { mCublasWrapper->setBF16GemmConfig(); } @@ -121,7 +124,7 @@ int64_t getGemmWorkSpaceSize(int64_t numTokens, int64_t maxLoraModuleNum, int64_ } size_t LoraImpl::getWorkspaceSize( - int64_t const numTokens, int64_t const numReqs, nvinfer1::DataType const type) const noexcept + int64_t const numTokens, int64_t const numReqs, tensorrt_llm::DataType const type) const noexcept { TLLM_LOG_DEBUG("%s", __PRETTY_FUNCTION__); auto const typeSize = tensorrt_llm::common::getDTypeSize(type); diff --git a/cpp/tensorrt_llm/kernels/lora/lora.h b/cpp/tensorrt_llm/kernels/lora/lora.h index 7215a7af74d4..73a2cbe330be 100644 --- a/cpp/tensorrt_llm/kernels/lora/lora.h +++ b/cpp/tensorrt_llm/kernels/lora/lora.h @@ -19,7 +19,7 @@ #include "tensorrt_llm/common/config.h" #include "tensorrt_llm/common/cublasMMWrapper.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include #include @@ -37,9 +37,10 @@ class LoraImpl { public: LoraImpl(int in_hidden_size, std::vector out_hidden_sizes, bool transA, bool transB, int num_lora_modules, - nvinfer1::DataType type, int max_low_rank, std::shared_ptr cublasWrapper); + tensorrt_llm::DataType type, int max_low_rank, std::shared_ptr cublasWrapper); - [[nodiscard]] size_t getWorkspaceSize(int64_t numTokens, int64_t numReqs, nvinfer1::DataType type) const noexcept; + [[nodiscard]] size_t getWorkspaceSize( + int64_t numTokens, int64_t numReqs, tensorrt_llm::DataType type) const noexcept; void setBestTactic(std::optional config); int run(int64_t numTokens, int64_t numReqs, void const* input, int32_t const* loraRanks, void const* const* loraWeightsPtr, int weightIndex, void* const* outputs, void* workspace, cudaStream_t stream); @@ -54,7 +55,7 @@ class LoraImpl private: bool mTransA; bool mTransB; - nvinfer1::DataType mType; + tensorrt_llm::DataType mType; int mNumLoraModules; // @fixme: seems this is shared across multiple clones. diff --git a/cpp/tensorrt_llm/kernels/lora/loraGroupGEMMParamFillRowReorderFusion.cu b/cpp/tensorrt_llm/kernels/lora/loraGroupGEMMParamFillRowReorderFusion.cu index c3276ea487bc..1165f39b3e5d 100644 --- a/cpp/tensorrt_llm/kernels/lora/loraGroupGEMMParamFillRowReorderFusion.cu +++ b/cpp/tensorrt_llm/kernels/lora/loraGroupGEMMParamFillRowReorderFusion.cu @@ -18,6 +18,7 @@ #include "tensorrt_llm/common/assert.h" #include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/common/tllmDataType.h" #include #include @@ -358,7 +359,7 @@ void launchLoraGroupGEMMParamFillRowReorderFusion(int32_t* in_sizes, int32_t* ou int64_t a_base, int64_t d_base, int64_t d_prime_base, int32_t const* slot_counts, int32_t const* slot_ranks, int64_t const* slot_offsets, int32_t const* module_out_sizes, int64_t const* module_out_prefix, int64_t const* b_ptrs, int64_t const* b_prime_ptrs, void const* input, int64_t const* sorted_ids, - int32_t module_count, nvinfer1::DataType dtype, cudaStream_t stream) + int32_t module_count, tensorrt_llm::DataType dtype, cudaStream_t stream) { // Determine block dimensions (1D) // Requirements: 1) >= max_lora_count * module_count 2) >= 256 3) divisible by 32 diff --git a/cpp/tensorrt_llm/kernels/lora/loraGroupGEMMParamFillRowReorderFusion.h b/cpp/tensorrt_llm/kernels/lora/loraGroupGEMMParamFillRowReorderFusion.h index 3043054ca4b5..835b9f8bed96 100644 --- a/cpp/tensorrt_llm/kernels/lora/loraGroupGEMMParamFillRowReorderFusion.h +++ b/cpp/tensorrt_llm/kernels/lora/loraGroupGEMMParamFillRowReorderFusion.h @@ -17,7 +17,7 @@ #pragma once #include "tensorrt_llm/common/config.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include #include @@ -70,7 +70,7 @@ void launchLoraGroupGEMMParamFillRowReorderFusion(int32_t* in_sizes, int32_t* ou int64_t a_base, int64_t d_base, int64_t d_prime_base, int32_t const* slot_counts, int32_t const* slot_ranks, int64_t const* slot_offsets, int32_t const* module_out_sizes, int64_t const* module_out_prefix, int64_t const* b_ptrs, int64_t const* b_prime_ptrs, void const* input, int64_t const* sorted_ids, - int32_t module_count, nvinfer1::DataType dtype, cudaStream_t stream); + int32_t module_count, tensorrt_llm::DataType dtype, cudaStream_t stream); } // namespace kernels diff --git a/cpp/tensorrt_llm/kernels/rmsNormFp4QuantKernels.cu b/cpp/tensorrt_llm/kernels/rmsNormFp4QuantKernels.cu index b70a9549df86..45b1a38266b1 100644 --- a/cpp/tensorrt_llm/kernels/rmsNormFp4QuantKernels.cu +++ b/cpp/tensorrt_llm/kernels/rmsNormFp4QuantKernels.cu @@ -389,7 +389,7 @@ void launchRmsNormFp4QuantKernel(RmsNormFp4QuantParams const& params, cudaStream } // namespace rms_norm_fp4_quant -void residualRmsNormFp4Quant(RmsNormFp4QuantParams const& params, nvinfer1::DataType dataType, cudaStream_t stream) +void residualRmsNormFp4Quant(RmsNormFp4QuantParams const& params, tensorrt_llm::DataType dataType, cudaStream_t stream) { // The NVFP4 epilogue (cvt_warp_fp16_to_fp4) is compiled only for // __CUDA_ARCH__ >= 1000 and emits zeros otherwise, so this kernel is correct @@ -410,11 +410,11 @@ void residualRmsNormFp4Quant(RmsNormFp4QuantParams const& params, nvinfer1::Data switch (dataType) { #ifdef ENABLE_BF16 - case nvinfer1::DataType::kBF16: + case tensorrt_llm::DataType::kBF16: rms_norm_fp4_quant::launchRmsNormFp4QuantKernel<__nv_bfloat16, /*OutNorm=*/true>(params, stream); break; #endif - case nvinfer1::DataType::kHALF: + case tensorrt_llm::DataType::kHALF: rms_norm_fp4_quant::launchRmsNormFp4QuantKernel(params, stream); break; default: TLLM_THROW("Unsupported dataType for residualRmsNormFp4Quant"); @@ -425,11 +425,11 @@ void residualRmsNormFp4Quant(RmsNormFp4QuantParams const& params, nvinfer1::Data switch (dataType) { #ifdef ENABLE_BF16 - case nvinfer1::DataType::kBF16: + case tensorrt_llm::DataType::kBF16: rms_norm_fp4_quant::launchRmsNormFp4QuantKernel<__nv_bfloat16, /*OutNorm=*/false>(params, stream); break; #endif - case nvinfer1::DataType::kHALF: + case tensorrt_llm::DataType::kHALF: rms_norm_fp4_quant::launchRmsNormFp4QuantKernel(params, stream); break; default: TLLM_THROW("Unsupported dataType for residualRmsNormFp4Quant"); diff --git a/cpp/tensorrt_llm/kernels/rmsNormFp4QuantKernels.h b/cpp/tensorrt_llm/kernels/rmsNormFp4QuantKernels.h index 86e998fea209..0ecba538d567 100644 --- a/cpp/tensorrt_llm/kernels/rmsNormFp4QuantKernels.h +++ b/cpp/tensorrt_llm/kernels/rmsNormFp4QuantKernels.h @@ -17,8 +17,8 @@ #pragma once #include "tensorrt_llm/common/config.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/kernels/quantization.h" -#include #include TRTLLM_NAMESPACE_BEGIN @@ -80,7 +80,7 @@ struct RmsNormFp4QuantParams // fp4_quantize) pair becomes one launch on the attention-DP path. All inputs, // outputs, and layout configuration are carried in params (see the struct // field docs above); dataType selects the fp16/bf16 instantiation. -void residualRmsNormFp4Quant(RmsNormFp4QuantParams const& params, nvinfer1::DataType dataType, cudaStream_t stream); +void residualRmsNormFp4Quant(RmsNormFp4QuantParams const& params, tensorrt_llm::DataType dataType, cudaStream_t stream); } // namespace kernels diff --git a/cpp/tensorrt_llm/kernels/splitkGroupGemm.cu b/cpp/tensorrt_llm/kernels/splitkGroupGemm.cu index 6397396ea6f6..1f63189ec657 100644 --- a/cpp/tensorrt_llm/kernels/splitkGroupGemm.cu +++ b/cpp/tensorrt_llm/kernels/splitkGroupGemm.cu @@ -26,6 +26,7 @@ #include "tensorrt_llm/common/cudaUtils.h" #include "tensorrt_llm/common/memoryUtils.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/cutlass_extensions/include/cutlass_extensions/gemm/device/splitk_gemm_grouped.h" #include "tensorrt_llm/cutlass_extensions/include/cutlass_extensions/gemm/kernel/default_splitk_gemm_grouped.h" #include "tensorrt_llm/cutlass_extensions/include/cutlass_extensions/gemm/kernel/splitk_gemm_grouped.h" @@ -203,20 +204,20 @@ template const& problemSizes, std::vector const& ptrA, std::vector const& ptrB, std::vector const& ptrC, std::vector const& ptrD, void* gemmParamsWorkSpace, int64_t gemmParamsWorkSpaceSize, void* gemmWorkSpace, int64_t gemmWorkSpaceSize, - nvinfer1::DataType dataType, int splitKSlices, cudaStream_t stream) + tensorrt_llm::DataType dataType, int splitKSlices, cudaStream_t stream) { - if (dataType == nvinfer1::DataType::kHALF) + if (dataType == tensorrt_llm::DataType::kHALF) { splitkGroupedGemm_(problemSizes, ptrA, ptrB, ptrC, ptrD, gemmParamsWorkSpace, gemmParamsWorkSpaceSize, gemmWorkSpace, gemmWorkSpaceSize, splitKSlices, stream); } - else if (dataType == nvinfer1::DataType::kFLOAT) + else if (dataType == tensorrt_llm::DataType::kFLOAT) { TLLM_CHECK_WITH_INFO(false, "not support float input/output"); } #ifdef ENABLE_BF16 - else if (dataType == nvinfer1::DataType::kBF16) + else if (dataType == tensorrt_llm::DataType::kBF16) { splitkGroupedGemm_( problemSizes, ptrA, ptrB, ptrC, ptrD, gemmParamsWorkSpace, gemmParamsWorkSpaceSize, gemmWorkSpace, @@ -228,7 +229,7 @@ void splitkGroupedGemmType_(std::vector const& problem void splitkGroupedGemm(std::vector const& problemSizes, std::vector const& ptrA, std::vector const& ptrB, std::vector const& ptrC, std::vector const& ptrD, void* gemmParamsWorkSpace, int64_t gemmParamsWorkSpaceSize, void* gemmWorkSpace, int64_t gemmWorkSpaceSize, - bool isLoraIn, nvinfer1::DataType dataType, int splitKSlices, int minKN, cudaStream_t stream) + bool isLoraIn, tensorrt_llm::DataType dataType, int splitKSlices, int minKN, cudaStream_t stream) { TLLM_LOG_TRACE("%s start, isLoraIn: %d, minKN = %d", __PRETTY_FUNCTION__, static_cast(isLoraIn), minKN); if (isLoraIn) diff --git a/cpp/tensorrt_llm/kernels/splitkGroupGemm.h b/cpp/tensorrt_llm/kernels/splitkGroupGemm.h index 6ada8255292e..bcde457db7d2 100644 --- a/cpp/tensorrt_llm/kernels/splitkGroupGemm.h +++ b/cpp/tensorrt_llm/kernels/splitkGroupGemm.h @@ -17,7 +17,7 @@ #include "cutlass/gemm_coord.h" #include "tensorrt_llm/common/config.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include TRTLLM_NAMESPACE_BEGIN @@ -30,7 +30,7 @@ int64_t getSplitkGroupedGemmParamsWorkSpaceSize(int64_t problem_count); void splitkGroupedGemm(std::vector const& problem_sizes, std::vector const& ptrA, std::vector const& ptrB, std::vector const& ptrC, std::vector const& ptrD, void* gemmParamsWorkspace, int64_t gemmParamsWorkSpaceSize, void* gemmWorkSpace, int64_t gemmWorkspaceSize, - bool isLoraIn, nvinfer1::DataType dataType, int splitKSlices, int minKN, cudaStream_t stream); + bool isLoraIn, tensorrt_llm::DataType dataType, int splitKSlices, int minKN, cudaStream_t stream); } // namespace kernels diff --git a/cpp/tensorrt_llm/kernels/unfusedAttentionKernels.h b/cpp/tensorrt_llm/kernels/unfusedAttentionKernels.h index 302c278f7a2a..04bcdfbbe2ac 100644 --- a/cpp/tensorrt_llm/kernels/unfusedAttentionKernels.h +++ b/cpp/tensorrt_llm/kernels/unfusedAttentionKernels.h @@ -17,6 +17,7 @@ #include "tensorrt_llm/common/config.h" #include "tensorrt_llm/common/quantization.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/kernels/gptKernels.h" #include "tensorrt_llm/kernels/kvCacheUtils.h" #include "tensorrt_llm/kernels/mlaKernels.h" @@ -242,42 +243,42 @@ struct QKVPreprocessingParams { ss << "seq_lens: " << *(runtime::ITensor::wrap( - (void*) seq_lens, nvinfer1::DataType::kINT32, runtime::ITensor::makeShape({batch_size}))); + (void*) seq_lens, tensorrt_llm::DataType::kINT32, runtime::ITensor::makeShape({batch_size}))); } if (cache_seq_lens && batch_size > 0) { ss << "cache_seq_lens: " - << *(runtime::ITensor::wrap( - (void*) cache_seq_lens, nvinfer1::DataType::kINT32, runtime::ITensor::makeShape({batch_size}))); + << *(runtime::ITensor::wrap((void*) cache_seq_lens, tensorrt_llm::DataType::kINT32, + runtime::ITensor::makeShape({batch_size}))); } if (encoder_seq_lens && batch_size > 0) { ss << "encoder_seq_lens: " - << *(runtime::ITensor::wrap( - (void*) encoder_seq_lens, nvinfer1::DataType::kINT32, runtime::ITensor::makeShape({batch_size}))); + << *(runtime::ITensor::wrap((void*) encoder_seq_lens, tensorrt_llm::DataType::kINT32, + runtime::ITensor::makeShape({batch_size}))); } if (cu_seq_lens && batch_size > 0) { ss << "cu_seq_lens: " << *(runtime::ITensor::wrap( - (void*) cu_seq_lens, nvinfer1::DataType::kINT32, runtime::ITensor::makeShape({batch_size}))); + (void*) cu_seq_lens, tensorrt_llm::DataType::kINT32, runtime::ITensor::makeShape({batch_size}))); } if (cu_kv_seq_lens && batch_size > 0) { ss << "cu_kv_seq_lens: " - << *(runtime::ITensor::wrap( - (void*) cu_kv_seq_lens, nvinfer1::DataType::kINT32, runtime::ITensor::makeShape({batch_size}))); + << *(runtime::ITensor::wrap((void*) cu_kv_seq_lens, tensorrt_llm::DataType::kINT32, + runtime::ITensor::makeShape({batch_size}))); } if (sparse_kv_offsets) { ss << "sparse_kv_offsets: " - << *(runtime::ITensor::wrap((void*) sparse_kv_offsets, nvinfer1::DataType::kINT32, + << *(runtime::ITensor::wrap((void*) sparse_kv_offsets, tensorrt_llm::DataType::kINT32, runtime::ITensor::makeShape({batch_size + 1}))); } if (rotary_embedding_inv_freq && batch_size > 0 && rotary_embedding_dim > 0) { ss << "rotary_embedding_inv_freq: " - << *(runtime::ITensor::wrap((void*) rotary_embedding_inv_freq, nvinfer1::DataType::kFLOAT, + << *(runtime::ITensor::wrap((void*) rotary_embedding_inv_freq, tensorrt_llm::DataType::kFLOAT, runtime::ITensor::makeShape({batch_size, rotary_embedding_dim / 2}))); } ss << "rotary_coef_cache_buffer: " << rotary_coef_cache_buffer << std::endl; diff --git a/cpp/tensorrt_llm/kernels/userbuffers/ub_interface.cpp b/cpp/tensorrt_llm/kernels/userbuffers/ub_interface.cpp index 3e19f9ebe72a..d219572d2707 100644 --- a/cpp/tensorrt_llm/kernels/userbuffers/ub_interface.cpp +++ b/cpp/tensorrt_llm/kernels/userbuffers/ub_interface.cpp @@ -16,6 +16,7 @@ #include "ub_interface.h" #include "tensorrt_llm/common/config.h" #include "tensorrt_llm/common/cudaDriverWrapper.h" +#include "tensorrt_llm/common/tllmDataType.h" #include #include @@ -80,13 +81,13 @@ namespace kernels::ub { void allreduce2_userbuff_inplace_launcher(int const handler, size_t const offset, size_t const elements, - nvinfer1::DataType dataType, communicator* comm, cudaStream_t stream) + tensorrt_llm::DataType dataType, communicator* comm, cudaStream_t stream) { allreduce2_userbuff_inplace_impl(handler, offset, elements, dataType, comm, stream); } int allgather2_userbuff_residual_launcher(int const handler, size_t const offset, size_t const elements, - int const hidden_size, void* residual, nvinfer1::DataType dataType, communicator* comm, cudaStream_t stream, + int const hidden_size, void* residual, tensorrt_llm::DataType dataType, communicator* comm, cudaStream_t stream, bool force_enable) { return allgather2_userbuff_residual_impl( @@ -95,7 +96,7 @@ int allgather2_userbuff_residual_launcher(int const handler, size_t const offset int allreduce2_userbuff_rmsnorm_launcher(int const handler, size_t const offset, int const out_handler, size_t const out_offset, size_t const elements, int const hidden_size, void* beta, void* gamma, float eps, - void* residual_in, void* residual_out, nvinfer1::DataType dataType, communicator* comm, cudaStream_t stream) + void* residual_in, void* residual_out, tensorrt_llm::DataType dataType, communicator* comm, cudaStream_t stream) { return allreduce2_userbuff_rmsnorm_impl(handler, offset, out_handler, out_offset, elements, hidden_size, beta, gamma, eps, residual_in, residual_out, dataType, comm, stream); @@ -103,7 +104,7 @@ int allreduce2_userbuff_rmsnorm_launcher(int const handler, size_t const offset, int allreduce2_userbuff_inplace_rmsnorm_quant_launcher(int const handler, size_t const offset, int const out_handler, size_t const out_offset, size_t const elements, int const hidden_size, void* beta, void* gamma, float eps, - float* scalefactor, void* residual_in, void* residual_out, nvinfer1::DataType dataType, communicator* comm, + float* scalefactor, void* residual_in, void* residual_out, tensorrt_llm::DataType dataType, communicator* comm, cudaStream_t stream) { return allreduce2_userbuff_inplace_rmsnorm_quant_impl(handler, offset, out_handler, out_offset, elements, @@ -113,7 +114,7 @@ int allreduce2_userbuff_inplace_rmsnorm_quant_launcher(int const handler, size_t int allreduce2_userbuff_inplace_rmsnorm_quant_fp4_launcher(int const handler, size_t const offset, int const out_handler, size_t const out_offset, int const scale_handler, size_t const scale_offset, size_t const elements, int const hidden_size, void* beta, void* gamma, float eps, float* scalefactor, - void* residual_in, void* residual_out, nvinfer1::DataType dataType, communicator* comm, cudaStream_t stream) + void* residual_in, void* residual_out, tensorrt_llm::DataType dataType, communicator* comm, cudaStream_t stream) { return allreduce2_userbuff_inplace_rmsnorm_quant_fp4_impl(handler, offset, out_handler, out_offset, scale_handler, scale_offset, elements, hidden_size, beta, gamma, eps, scalefactor, residual_in, residual_out, dataType, comm, @@ -165,12 +166,12 @@ TRTLLM_NAMESPACE_BEGIN namespace kernels::ub { void allreduce2_userbuff_inplace_launcher(int const handler, size_t const offset, size_t const elements, - nvinfer1::DataType dataType, communicator* comm, cudaStream_t stream) + tensorrt_llm::DataType dataType, communicator* comm, cudaStream_t stream) { } int allgather2_userbuff_residual_launcher(int const handler, size_t const offset, size_t const elements, - int const hidden_size, void* residual, nvinfer1::DataType dataType, communicator* comm, cudaStream_t stream, + int const hidden_size, void* residual, tensorrt_llm::DataType dataType, communicator* comm, cudaStream_t stream, bool force_enable) { return 0; @@ -178,7 +179,7 @@ int allgather2_userbuff_residual_launcher(int const handler, size_t const offset int allreduce2_userbuff_inplace_rmsnorm_quant_launcher(int const handler, size_t const offset, int const out_handler, size_t const out_offset, size_t const elements, int const hidden_size, void* beta, void* gamma, float eps, - float* scalefactor, void* residual_in, void* residual_out, nvinfer1::DataType dataType, communicator* comm, + float* scalefactor, void* residual_in, void* residual_out, tensorrt_llm::DataType dataType, communicator* comm, cudaStream_t stream) { return 0; @@ -187,7 +188,7 @@ int allreduce2_userbuff_inplace_rmsnorm_quant_launcher(int const handler, size_t int allreduce2_userbuff_inplace_rmsnorm_quant_fp4_launcher(int const handler, size_t const offset, int const out_handler, size_t const out_offset, int const scale_handler, size_t const scale_offset, size_t const elements, int const hidden_size, void* beta, void* gamma, float eps, float* scalefactor, - void* residual_in, void* residual_out, nvinfer1::DataType dataType, communicator* comm, cudaStream_t stream) + void* residual_in, void* residual_out, tensorrt_llm::DataType dataType, communicator* comm, cudaStream_t stream) { return 0; } diff --git a/cpp/tensorrt_llm/kernels/userbuffers/ub_interface.h b/cpp/tensorrt_llm/kernels/userbuffers/ub_interface.h index e8a48e2c680b..dc68154fb462 100644 --- a/cpp/tensorrt_llm/kernels/userbuffers/ub_interface.h +++ b/cpp/tensorrt_llm/kernels/userbuffers/ub_interface.h @@ -18,6 +18,7 @@ #include "tensorrt_llm/common/config.h" #include "tensorrt_llm/common/cudaUtils.h" #include "tensorrt_llm/common/dataType.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "ub_allocator.h" namespace tensorrt_llm::runtime::ub @@ -40,24 +41,24 @@ namespace kernels::ub using ::tensorrt_llm::runtime::ub::communicator; void allreduce2_userbuff_inplace_launcher(int const handler, size_t const offset, size_t const elements, - nvinfer1::DataType dataType, communicator* comm, cudaStream_t stream = 0); + tensorrt_llm::DataType dataType, communicator* comm, cudaStream_t stream = 0); int allgather2_userbuff_residual_launcher(int const handler, size_t const offset, size_t const elements, - int const hidden_size, void* residual, nvinfer1::DataType dataType, communicator* comm, cudaStream_t stream, + int const hidden_size, void* residual, tensorrt_llm::DataType dataType, communicator* comm, cudaStream_t stream, bool force_enable = false); int allreduce2_userbuff_rmsnorm_launcher(int const handler, size_t const offset, int const out_handler, size_t const out_offset, size_t const elements, int const hidden_size, void* beta, void* gamma, float eps, - void* residual_in, void* residual_out, nvinfer1::DataType dataType, communicator* comm, cudaStream_t stream); + void* residual_in, void* residual_out, tensorrt_llm::DataType dataType, communicator* comm, cudaStream_t stream); int allreduce2_userbuff_inplace_rmsnorm_quant_launcher(int const handler, size_t const offset, int const out_handler, size_t const out_offset, size_t const elements, int const hidden_size, void* beta, void* gamma, float eps, - float* scalefactor, void* residual_in, void* residual_out, nvinfer1::DataType dataType, communicator* comm, + float* scalefactor, void* residual_in, void* residual_out, tensorrt_llm::DataType dataType, communicator* comm, cudaStream_t stream); int allreduce2_userbuff_inplace_rmsnorm_quant_fp4_launcher(int const handler, size_t const offset, int const out_handler, size_t const out_offset, int const scale_handler, size_t const scale_offset, size_t const elements, int const hidden_size, void* beta, void* gamma, float eps, float* scalefactor, - void* residual_in, void* residual_out, nvinfer1::DataType dataType, communicator* comm, cudaStream_t stream); + void* residual_in, void* residual_out, tensorrt_llm::DataType dataType, communicator* comm, cudaStream_t stream); } // namespace kernels::ub TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/userbuffers/userbuffers.cu b/cpp/tensorrt_llm/kernels/userbuffers/userbuffers.cu index 8cb5814e0398..a19059e9010c 100644 --- a/cpp/tensorrt_llm/kernels/userbuffers/userbuffers.cu +++ b/cpp/tensorrt_llm/kernels/userbuffers/userbuffers.cu @@ -15,6 +15,7 @@ */ #include "tensorrt_llm/common/config.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/kernels/quantization.cuh" #include "userbuffers.h" #include "utils.h" @@ -1774,11 +1775,11 @@ int allgather2_userbuff_residual(int const handler, size_t const offset, size_t } void allreduce2_userbuff_inplace_impl(int const handler, size_t const offset, size_t const elements, - nvinfer1::DataType dataType, communicator* comm, cudaStream_t stream) + tensorrt_llm::DataType dataType, communicator* comm, cudaStream_t stream) { switch (dataType) { - case nvinfer1::DataType::kHALF: + case tensorrt_llm::DataType::kHALF: { if (kDISABLE_FP32_ACCUMULATION) { @@ -1791,7 +1792,7 @@ void allreduce2_userbuff_inplace_impl(int const handler, size_t const offset, si break; } #ifdef ENABLE_BF16 - case nvinfer1::DataType::kBF16: + case tensorrt_llm::DataType::kBF16: { if (kDISABLE_FP32_ACCUMULATION) { @@ -1809,17 +1810,17 @@ void allreduce2_userbuff_inplace_impl(int const handler, size_t const offset, si } int allgather2_userbuff_residual_impl(int const handler, size_t const offset, size_t const elements, - int const hidden_size, void* residual, nvinfer1::DataType dataType, communicator* comm, cudaStream_t stream, + int const hidden_size, void* residual, tensorrt_llm::DataType dataType, communicator* comm, cudaStream_t stream, bool force_enable) { switch (dataType) { - case nvinfer1::DataType::kHALF: + case tensorrt_llm::DataType::kHALF: return allgather2_userbuff_residual( handler, offset, elements, hidden_size, residual, comm, stream, force_enable); break; #ifdef ENABLE_BF16 - case nvinfer1::DataType::kBF16: + case tensorrt_llm::DataType::kBF16: return allgather2_userbuff_residual<__nv_bfloat16>( handler, offset, elements, hidden_size, residual, comm, stream, force_enable); break; @@ -1830,11 +1831,11 @@ int allgather2_userbuff_residual_impl(int const handler, size_t const offset, si int allreduce2_userbuff_rmsnorm_impl(int const handler, size_t const offset, int const out_handler, size_t const out_offset, size_t const elements, int const hidden_size, void* beta, void* gamma, float eps, - void* residual_in, void* residual_out, nvinfer1::DataType dataType, communicator* comm, cudaStream_t stream) + void* residual_in, void* residual_out, tensorrt_llm::DataType dataType, communicator* comm, cudaStream_t stream) { switch (dataType) { - case nvinfer1::DataType::kHALF: + case tensorrt_llm::DataType::kHALF: { if (kDISABLE_FP32_ACCUMULATION) { @@ -1849,7 +1850,7 @@ int allreduce2_userbuff_rmsnorm_impl(int const handler, size_t const offset, int break; } #ifdef ENABLE_BF16 - case nvinfer1::DataType::kBF16: + case tensorrt_llm::DataType::kBF16: { if (kDISABLE_FP32_ACCUMULATION) { @@ -1870,12 +1871,12 @@ int allreduce2_userbuff_rmsnorm_impl(int const handler, size_t const offset, int int allreduce2_userbuff_inplace_rmsnorm_quant_impl(int const handler, size_t const offset, int const out_handler, size_t const out_offset, size_t const elements, int const hidden_size, void* beta, void* gamma, float eps, - float* scalefactor, void* residual_in, void* residual_out, nvinfer1::DataType dataType, communicator* comm, + float* scalefactor, void* residual_in, void* residual_out, tensorrt_llm::DataType dataType, communicator* comm, cudaStream_t stream) { switch (dataType) { - case nvinfer1::DataType::kHALF: + case tensorrt_llm::DataType::kHALF: { if (kDISABLE_FP32_ACCUMULATION) { @@ -1890,7 +1891,7 @@ int allreduce2_userbuff_inplace_rmsnorm_quant_impl(int const handler, size_t con break; } #ifdef ENABLE_BF16 - case nvinfer1::DataType::kBF16: + case tensorrt_llm::DataType::kBF16: { if (kDISABLE_FP32_ACCUMULATION) { @@ -1914,11 +1915,11 @@ int allreduce2_userbuff_inplace_rmsnorm_quant_impl(int const handler, size_t con int allreduce2_userbuff_inplace_rmsnorm_quant_fp4_impl(int const handler, size_t const offset, int const out_handler, size_t const out_offset, int const scale_handler, size_t const scale_offset, size_t const elements, int const hidden_size, void* beta, void* gamma, float eps, float* scalefactor, void* residual_in, - void* residual_out, nvinfer1::DataType dataType, communicator* comm, cudaStream_t stream) + void* residual_out, tensorrt_llm::DataType dataType, communicator* comm, cudaStream_t stream) { switch (dataType) { - case nvinfer1::DataType::kHALF: + case tensorrt_llm::DataType::kHALF: { if (kDISABLE_FP32_ACCUMULATION) { @@ -1935,7 +1936,7 @@ int allreduce2_userbuff_inplace_rmsnorm_quant_fp4_impl(int const handler, size_t break; } #ifdef ENABLE_BF16 - case nvinfer1::DataType::kBF16: + case tensorrt_llm::DataType::kBF16: { if (kDISABLE_FP32_ACCUMULATION) { diff --git a/cpp/tensorrt_llm/kernels/userbuffers/userbuffers.h b/cpp/tensorrt_llm/kernels/userbuffers/userbuffers.h index 96f21b748282..5d3ffe0cc950 100644 --- a/cpp/tensorrt_llm/kernels/userbuffers/userbuffers.h +++ b/cpp/tensorrt_llm/kernels/userbuffers/userbuffers.h @@ -14,6 +14,7 @@ * limitations under the License. */ #pragma once +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/runtime/utils/mpiUtils.h" #include "tensorrt_llm/runtime/worldConfig.h" @@ -120,25 +121,25 @@ namespace kernels::ub { using namespace ::tensorrt_llm::runtime::ub; void allreduce2_userbuff_inplace_impl(int const handler, size_t const offset, size_t const elements, - nvinfer1::DataType dataType, communicator* comm, cudaStream_t stream = 0); + tensorrt_llm::DataType dataType, communicator* comm, cudaStream_t stream = 0); // for TP-parallelism, only single node is implemented int allgather2_userbuff_residual_impl(int const handler, size_t const offset, size_t const elements, - int const hidden_size, void* residual, nvinfer1::DataType dataType, communicator* comm, cudaStream_t stream, + int const hidden_size, void* residual, tensorrt_llm::DataType dataType, communicator* comm, cudaStream_t stream, bool force_enable); int allreduce2_userbuff_rmsnorm_impl(int const handler, size_t const offset, int const out_handler, size_t const out_offset, size_t const elements, int const hidden_size, void* beta, void* gamma, float eps, - void* residual_in, void* residual_out, nvinfer1::DataType dataType, communicator* comm, cudaStream_t stream); + void* residual_in, void* residual_out, tensorrt_llm::DataType dataType, communicator* comm, cudaStream_t stream); int allreduce2_userbuff_inplace_rmsnorm_quant_impl(int const handler, size_t const offset, int const out_handler, size_t const out_offset, size_t const elements, int const hidden_size, void* beta, void* gamma, float eps, - float* scalefactor, void* residual_in, void* residual_out, nvinfer1::DataType dataType, communicator* comm, + float* scalefactor, void* residual_in, void* residual_out, tensorrt_llm::DataType dataType, communicator* comm, cudaStream_t stream); int allreduce2_userbuff_inplace_rmsnorm_quant_fp4_impl(int const handler, size_t const offset, int const out_handler, size_t const out_offset, int const scale_handler, size_t const scale_offset, size_t const elements, int const hidden_size, void* beta, void* gamma, float eps, float* scalefactor, void* residual_in, - void* residual_out, nvinfer1::DataType dataType, communicator* comm, cudaStream_t stream); + void* residual_out, tensorrt_llm::DataType dataType, communicator* comm, cudaStream_t stream); } // namespace kernels::ub TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv/cudaCoreGemm.h b/cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv/cudaCoreGemm.h index eb939b57c2db..c2bf35175391 100644 --- a/cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv/cudaCoreGemm.h +++ b/cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv/cudaCoreGemm.h @@ -24,8 +24,6 @@ #include "tensorrt_llm/kernels/cutlass_kernels/cutlass_type_conversion.h" #include "tensorrt_llm/runtime/common.h" -#include - #include #include #include diff --git a/cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv/cudaCoreGemmNVFP4.h b/cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv/cudaCoreGemmNVFP4.h index 616f9d25c2bf..6e901846ed25 100644 --- a/cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv/cudaCoreGemmNVFP4.h +++ b/cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv/cudaCoreGemmNVFP4.h @@ -24,8 +24,6 @@ #include "tensorrt_llm/kernels/cutlass_kernels/cutlass_type_conversion.h" #include "tensorrt_llm/runtime/common.h" -#include - #include #include #include diff --git a/cpp/tensorrt_llm/layers/decodingParams.h b/cpp/tensorrt_llm/layers/decodingParams.h index 1e77b8919ca1..76c5cedd637b 100644 --- a/cpp/tensorrt_llm/layers/decodingParams.h +++ b/cpp/tensorrt_llm/layers/decodingParams.h @@ -17,6 +17,7 @@ #pragma once #include "tensorrt_llm/common/assert.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/executor/executor.h" #include "tensorrt_llm/kernels/beamSearchKernels.h" #include "tensorrt_llm/runtime/iTensor.h" @@ -192,9 +193,9 @@ class ExplicitDraftTokensSetupParams : public DecodingSetupParams public: OptVec temperature; // [setupBatchSize] // Hack to init some data for the context phase in the setup. - TensorPtr randomDataSample; // [maxBatchSize], on gpu - TensorPtr temperatures; // [maxBatchSize], on gpu - nvinfer1::DataType dtype; // [1] + TensorPtr randomDataSample; // [maxBatchSize], on gpu + TensorPtr temperatures; // [maxBatchSize], on gpu + tensorrt_llm::DataType dtype; // [1] }; class EagleSetupParams : public DecodingSetupParams @@ -202,9 +203,9 @@ class EagleSetupParams : public DecodingSetupParams public: OptVec temperature; // [setupBatchSize] // Hack to init some data for the context phase in the setup. - TensorPtr randomDataSample; // [maxBatchSize], on gpu - TensorPtr temperatures; // [maxBatchSize], on gpu - nvinfer1::DataType dtype; // [1] + TensorPtr randomDataSample; // [maxBatchSize], on gpu + TensorPtr temperatures; // [maxBatchSize], on gpu + tensorrt_llm::DataType dtype; // [1] }; class DynamicDecodeSetupParams : public BaseSetupParams diff --git a/cpp/tensorrt_llm/layers/explicitDraftTokensLayer.cpp b/cpp/tensorrt_llm/layers/explicitDraftTokensLayer.cpp index e014ee4535e5..aedeb731574f 100644 --- a/cpp/tensorrt_llm/layers/explicitDraftTokensLayer.cpp +++ b/cpp/tensorrt_llm/layers/explicitDraftTokensLayer.cpp @@ -16,6 +16,7 @@ #include "explicitDraftTokensLayer.h" #include "tensorrt_llm/common/nvtxUtils.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/kernels/penaltyTypes.h" #include "tensorrt_llm/kernels/speculativeDecoding/common.h" #include "tensorrt_llm/kernels/speculativeDecoding/explicitDraftTokensKernels.h" @@ -93,15 +94,15 @@ void ExplicitDraftTokensLayer::setup(SizeType32 batchSize, SizeType32 beamWid batchSlots, getLimitsPenalty(DecodingPenaltyType::Temperature), "temperature penalty"); // Dispatch context buffer fill - if (mDecoderDtype == nvinfer1::DataType::kFLOAT) + if (mDecoderDtype == tensorrt_llm::DataType::kFLOAT) { fillContextBuffers(batchSize, batchSlots, *setupParams, workspace); } - else if (mDecoderDtype == nvinfer1::DataType::kHALF) + else if (mDecoderDtype == tensorrt_llm::DataType::kHALF) { fillContextBuffers(batchSize, batchSlots, *setupParams, workspace); } - else if (mDecoderDtype == nvinfer1::DataType::kBF16) + else if (mDecoderDtype == tensorrt_llm::DataType::kBF16) { fillContextBuffers<__nv_bfloat16>(batchSize, batchSlots, *setupParams, workspace); } @@ -126,15 +127,15 @@ void ExplicitDraftTokensLayer::forwardAsync(std::shared_ptr(*outputs, *inputs, workspace); } - else if (mDecoderDtype == nvinfer1::DataType::kHALF) + else if (mDecoderDtype == tensorrt_llm::DataType::kHALF) { splitInputDataToBatchSlots(*outputs, *inputs, workspace); } - else if (mDecoderDtype == nvinfer1::DataType::kBF16) + else if (mDecoderDtype == tensorrt_llm::DataType::kBF16) { splitInputDataToBatchSlots<__nv_bfloat16>(*outputs, *inputs, workspace); } diff --git a/cpp/tensorrt_llm/layers/explicitDraftTokensLayer.h b/cpp/tensorrt_llm/layers/explicitDraftTokensLayer.h index 75883ded6e5a..17fca4513cf1 100644 --- a/cpp/tensorrt_llm/layers/explicitDraftTokensLayer.h +++ b/cpp/tensorrt_llm/layers/explicitDraftTokensLayer.h @@ -16,6 +16,7 @@ #pragma once +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/layers/baseLayer.h" #include "tensorrt_llm/layers/decodingParams.h" #include "tensorrt_llm/runtime/common.h" @@ -83,7 +84,7 @@ class ExplicitDraftTokensLayer : public BaseLayer TensorPtr mTemperature; - std::optional mDecoderDtype{std::nullopt}; + std::optional mDecoderDtype{std::nullopt}; }; } // namespace tensorrt_llm::layers diff --git a/cpp/tensorrt_llm/layers/lookaheadAlgorithm.cpp b/cpp/tensorrt_llm/layers/lookaheadAlgorithm.cpp index 09843fd7ce44..76da89dfec0d 100644 --- a/cpp/tensorrt_llm/layers/lookaheadAlgorithm.cpp +++ b/cpp/tensorrt_llm/layers/lookaheadAlgorithm.cpp @@ -17,6 +17,7 @@ #include "tensorrt_llm/layers/lookaheadAlgorithm.h" #include "tensorrt_llm/common/assert.h" #include "tensorrt_llm/common/logger.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/executor/executor.h" #include "tensorrt_llm/layers/lookaheadDecodingUtils.h" #include "tensorrt_llm/runtime/common.h" @@ -35,14 +36,14 @@ LookaheadAlgorithm::LookaheadAlgorithm( runtime::SizeType32 maxW, runtime::SizeType32 maxN, runtime::SizeType32 maxG, runtime::SizeType32 id) : mPoolManager(maxG) , mPrefillsMax(runtime::BufferManager::cpu( - runtime::ITensor::makeShape({(maxN <= 1 ? 0 : maxN - 2)}), nvinfer1::DataType::kINT32)) + runtime::ITensor::makeShape({(maxN <= 1 ? 0 : maxN - 2)}), tensorrt_llm::DataType::kINT32)) , mPastTokensMax( - runtime::BufferManager::cpu(runtime::ITensor::makeShape({maxW * (maxN - 1)}), nvinfer1::DataType::kINT32)) - , mKeyTokensMax(runtime::BufferManager::cpu(runtime::ITensor::makeShape({maxW}), nvinfer1::DataType::kINT32)) + runtime::BufferManager::cpu(runtime::ITensor::makeShape({maxW * (maxN - 1)}), tensorrt_llm::DataType::kINT32)) + , mKeyTokensMax(runtime::BufferManager::cpu(runtime::ITensor::makeShape({maxW}), tensorrt_llm::DataType::kINT32)) , mGoldenTokensMax( - runtime::BufferManager::cpu(runtime::ITensor::makeShape({maxN * 2 - 1}), nvinfer1::DataType::kINT32)) + runtime::BufferManager::cpu(runtime::ITensor::makeShape({maxN * 2 - 1}), tensorrt_llm::DataType::kINT32)) , mGuessTokensMax( - runtime::BufferManager::cpu(runtime::ITensor::makeShape({maxG * (maxN - 1)}), nvinfer1::DataType::kINT32)) + runtime::BufferManager::cpu(runtime::ITensor::makeShape({maxG * (maxN - 1)}), tensorrt_llm::DataType::kINT32)) , mMaxW(maxW) , mMaxN(maxN) , mMaxG(maxG) @@ -52,12 +53,13 @@ LookaheadAlgorithm::LookaheadAlgorithm( std::tie(maxGeneratedLen, std::ignore, maxDraftLen, std::ignore) = executor::LookaheadDecodingConfig(maxW, maxN, maxG).calculateSpeculativeResource(); mAttentionMask = runtime::BufferManager::cpu( - runtime::ITensor::makeShape({maxDraftLen, maxDraftLen}), nvinfer1::DataType::kBOOL); + runtime::ITensor::makeShape({maxDraftLen, maxDraftLen}), tensorrt_llm::DataType::kBOOL); mDraftTokensMax - = runtime::BufferManager::cpu(runtime::ITensor::makeShape({maxDraftLen}), nvinfer1::DataType::kINT32); + = runtime::BufferManager::cpu(runtime::ITensor::makeShape({maxDraftLen}), tensorrt_llm::DataType::kINT32); mSampledTokensMax - = runtime::BufferManager::cpu(runtime::ITensor::makeShape({maxGeneratedLen}), nvinfer1::DataType::kINT32); - mEncodeMapMax = runtime::BufferManager::cpu(runtime::ITensor::makeShape({maxDraftLen}), nvinfer1::DataType::kINT32); + = runtime::BufferManager::cpu(runtime::ITensor::makeShape({maxGeneratedLen}), tensorrt_llm::DataType::kINT32); + mEncodeMapMax + = runtime::BufferManager::cpu(runtime::ITensor::makeShape({maxDraftLen}), tensorrt_llm::DataType::kINT32); } void LookaheadAlgorithm::setup(TensorConstPtr const& prompt, SizeType32 w, SizeType32 n, SizeType32 g, uint64_t seed) diff --git a/cpp/tensorrt_llm/layers/lookaheadDecodingLayer.cpp b/cpp/tensorrt_llm/layers/lookaheadDecodingLayer.cpp index bf6e15080f3c..986f0e0b978e 100644 --- a/cpp/tensorrt_llm/layers/lookaheadDecodingLayer.cpp +++ b/cpp/tensorrt_llm/layers/lookaheadDecodingLayer.cpp @@ -19,6 +19,7 @@ #include "tensorrt_llm/common/cudaUtils.h" #include "tensorrt_llm/common/logger.h" #include "tensorrt_llm/common/nvtxUtils.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/executor/executor.h" #include "tensorrt_llm/kernels/samplingTopKKernels.h" #include "tensorrt_llm/layers/decodingParams.h" @@ -64,38 +65,42 @@ LookaheadDecodingLayer::CpuAlgorithmResources::CpuAlgorithmResources(DecoderD mPrompts.reserve(maxBatchSize); for (auto bi = 0; bi < maxBatchSize; bi++) { - mPrompts.emplace_back(BufferManager::cpu(ITensor::makeShape({0}), nvinfer1::DataType::kINT32)); + mPrompts.emplace_back(BufferManager::cpu(ITensor::makeShape({0}), tensorrt_llm::DataType::kINT32)); } auto const maxBatchShape1D = ITensor::makeShape({maxBatchSize}); - mBatchSlots = BufferManager::cpu(maxBatchShape1D, nvinfer1::DataType::kINT32); + mBatchSlots = BufferManager::cpu(maxBatchShape1D, tensorrt_llm::DataType::kINT32); mTargetTokens - = BufferManager::cpu(ITensor::makeShape({maxBatchSize, maxTokensPerStep}), nvinfer1::DataType::kINT32); - mTokensPerStep = BufferManager::cpu(maxBatchShape1D, nvinfer1::DataType::kINT32); - mEndIds = BufferManager::cpu(maxBatchShape1D, nvinfer1::DataType::kINT32); + = BufferManager::cpu(ITensor::makeShape({maxBatchSize, maxTokensPerStep}), tensorrt_llm::DataType::kINT32); + mTokensPerStep = BufferManager::cpu(maxBatchShape1D, tensorrt_llm::DataType::kINT32); + mEndIds = BufferManager::cpu(maxBatchShape1D, tensorrt_llm::DataType::kINT32); - mOutputIds = BufferManager::cpu(ITensor::makeShape({maxBatchSize, maxNumNewTokens}), nvinfer1::DataType::kINT32); + mOutputIds + = BufferManager::cpu(ITensor::makeShape({maxBatchSize, maxNumNewTokens}), tensorrt_llm::DataType::kINT32); mNewTokens = BufferManager::cpu( - ITensor::makeShape({maxTokensPerStep, maxBatchSize, beamWidth}), nvinfer1::DataType::kINT32); + ITensor::makeShape({maxTokensPerStep, maxBatchSize, beamWidth}), tensorrt_llm::DataType::kINT32); mPathsOffsets - = BufferManager::cpu(ITensor::makeShape({maxBatchSize, maxAcceptedDraftLen}), nvinfer1::DataType::kINT32); + = BufferManager::cpu(ITensor::makeShape({maxBatchSize, maxAcceptedDraftLen}), tensorrt_llm::DataType::kINT32); mPathsOffsetsBatch - = BufferManager::cpu(ITensor::makeShape({maxBatchSize, maxAcceptedDraftLen}), nvinfer1::DataType::kINT32); - mNumNewTokens = BufferManager::cpu(maxBatchShape1D, nvinfer1::DataType::kINT32); - mNumNewTokensCumSum = BufferManager::cpu(ITensor::makeShape({maxBatchSize + 1}), nvinfer1::DataType::kINT32); - mNextDraftTokens = BufferManager::cpu(ITensor::makeShape({maxBatchSize, maxDraftLen}), nvinfer1::DataType::kINT32); - mNextDraftPosIds = BufferManager::cpu(ITensor::makeShape({maxBatchSize, maxDraftLen}), nvinfer1::DataType::kINT32); - mGenerationLengths = BufferManager::cpu(maxBatchShape1D, nvinfer1::DataType::kINT32); + = BufferManager::cpu(ITensor::makeShape({maxBatchSize, maxAcceptedDraftLen}), tensorrt_llm::DataType::kINT32); + mNumNewTokens = BufferManager::cpu(maxBatchShape1D, tensorrt_llm::DataType::kINT32); + mNumNewTokensCumSum = BufferManager::cpu(ITensor::makeShape({maxBatchSize + 1}), tensorrt_llm::DataType::kINT32); + mNextDraftTokens + = BufferManager::cpu(ITensor::makeShape({maxBatchSize, maxDraftLen}), tensorrt_llm::DataType::kINT32); + mNextDraftPosIds + = BufferManager::cpu(ITensor::makeShape({maxBatchSize, maxDraftLen}), tensorrt_llm::DataType::kINT32); + mGenerationLengths = BufferManager::cpu(maxBatchShape1D, tensorrt_llm::DataType::kINT32); mPositionOffsets - = BufferManager::cpu(ITensor::makeShape({maxBatchSize, maxTokensPerStep}), nvinfer1::DataType::kINT32); - mPositionIds = BufferManager::cpu(ITensor::makeShape({maxBatchSize, maxTokensPerStep}), nvinfer1::DataType::kINT32); + = BufferManager::cpu(ITensor::makeShape({maxBatchSize, maxTokensPerStep}), tensorrt_llm::DataType::kINT32); + mPositionIds + = BufferManager::cpu(ITensor::makeShape({maxBatchSize, maxTokensPerStep}), tensorrt_llm::DataType::kINT32); mAttentionMask - = BufferManager::cpu(ITensor::makeShape({maxTokensPerStep, maxTokensPerStep}), nvinfer1::DataType::kBOOL); + = BufferManager::cpu(ITensor::makeShape({maxTokensPerStep, maxTokensPerStep}), tensorrt_llm::DataType::kBOOL); mPackedMask = BufferManager::cpu(ITensor::makeShape({maxBatchSize, maxTokensPerStep, static_cast(divUp(maxTokensPerStep, 32))}), - nvinfer1::DataType::kINT32); - mNextDraftLengths = BufferManager::cpu(maxBatchShape1D, nvinfer1::DataType::kINT32); - mSequenceLengths = BufferManager::cpu(maxBatchShape1D, nvinfer1::DataType::kINT32); + tensorrt_llm::DataType::kINT32); + mNextDraftLengths = BufferManager::cpu(maxBatchShape1D, tensorrt_llm::DataType::kINT32); + mSequenceLengths = BufferManager::cpu(maxBatchShape1D, tensorrt_llm::DataType::kINT32); } template @@ -117,12 +122,12 @@ LookaheadDecodingLayer::LookaheadDecodingLayer( auto const maxBatchShape2D = ITensor::makeShape({maxBatchSize, maxTokensPerStep}); mWorkspaceSize = getTopKWorkspaceSize(maxBatchSize, maxTokensPerStep, maxTopK, vocabSizePadded); - mTargetTokensDevice = mBufferManager->gpu(maxBatchShape2D, nvinfer1::DataType::kINT32); + mTargetTokensDevice = mBufferManager->gpu(maxBatchShape2D, tensorrt_llm::DataType::kINT32); mCurandStatesDevice - = mBufferManager->gpu(ITensor::makeShape({maxBatchSize, sizeof(curandState_t)}), nvinfer1::DataType::kINT8); + = mBufferManager->gpu(ITensor::makeShape({maxBatchSize, sizeof(curandState_t)}), tensorrt_llm::DataType::kINT8); mSetupWorkspaceSize = DecodingLayerWorkspace::calculateRequiredWorkspaceSize( - std::make_pair(maxBatchShape1D, nvinfer1::DataType::kINT64)); + std::make_pair(maxBatchShape1D, tensorrt_llm::DataType::kINT64)); TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); } diff --git a/cpp/tensorrt_llm/layers/lookaheadDecodingUtils.h b/cpp/tensorrt_llm/layers/lookaheadDecodingUtils.h index 739cf65001ab..8e3e8f6c590d 100644 --- a/cpp/tensorrt_llm/layers/lookaheadDecodingUtils.h +++ b/cpp/tensorrt_llm/layers/lookaheadDecodingUtils.h @@ -16,6 +16,7 @@ #pragma once +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/runtime/bufferManager.h" #include "tensorrt_llm/runtime/iTensor.h" @@ -318,12 +319,12 @@ class DebugTensor { switch (mTensor.getDataType()) { - case nvinfer1::DataType::kBOOL: return values(); - case nvinfer1::DataType::kFLOAT: return values(); - case nvinfer1::DataType::kINT8: return values(); - case nvinfer1::DataType::kINT32: return values(); - case nvinfer1::DataType::kINT64: return values(); - case nvinfer1::DataType::kUINT8: return values(); + case tensorrt_llm::DataType::kBOOL: return values(); + case tensorrt_llm::DataType::kFLOAT: return values(); + case tensorrt_llm::DataType::kINT8: return values(); + case tensorrt_llm::DataType::kINT32: return values(); + case tensorrt_llm::DataType::kINT64: return values(); + case tensorrt_llm::DataType::kUINT8: return values(); default: return std::string(mName + ": Unsupported data type"); } } @@ -376,12 +377,12 @@ class DebugTensor { switch (mTensor.getDataType()) { - case nvinfer1::DataType::kBOOL: return randomize(3); - case nvinfer1::DataType::kFLOAT: return randomize(3); - case nvinfer1::DataType::kINT8: return randomize(3); - case nvinfer1::DataType::kINT32: return randomize(3); - case nvinfer1::DataType::kINT64: return randomize(3); - case nvinfer1::DataType::kUINT8: return randomize(3); + case tensorrt_llm::DataType::kBOOL: return randomize(3); + case tensorrt_llm::DataType::kFLOAT: return randomize(3); + case tensorrt_llm::DataType::kINT8: return randomize(3); + case tensorrt_llm::DataType::kINT32: return randomize(3); + case tensorrt_llm::DataType::kINT64: return randomize(3); + case tensorrt_llm::DataType::kUINT8: return randomize(3); default: return; } } @@ -391,12 +392,12 @@ class DebugTensor { switch (mTensor.getDataType()) { - case nvinfer1::DataType::kBOOL: return randomize(0); - case nvinfer1::DataType::kFLOAT: return randomize(0); - case nvinfer1::DataType::kINT8: return randomize(0); - case nvinfer1::DataType::kINT32: return randomize(0); - case nvinfer1::DataType::kINT64: return randomize(0); - case nvinfer1::DataType::kUINT8: return randomize(0); + case tensorrt_llm::DataType::kBOOL: return randomize(0); + case tensorrt_llm::DataType::kFLOAT: return randomize(0); + case tensorrt_llm::DataType::kINT8: return randomize(0); + case tensorrt_llm::DataType::kINT32: return randomize(0); + case tensorrt_llm::DataType::kINT64: return randomize(0); + case tensorrt_llm::DataType::kUINT8: return randomize(0); default: return; } } @@ -405,12 +406,12 @@ class DebugTensor { switch (mTensor.getDataType()) { - case nvinfer1::DataType::kBOOL: return randomize(1); - case nvinfer1::DataType::kFLOAT: return randomize(1); - case nvinfer1::DataType::kINT8: return randomize(1); - case nvinfer1::DataType::kINT32: return randomize(1); - case nvinfer1::DataType::kINT64: return randomize(1); - case nvinfer1::DataType::kUINT8: return randomize(1); + case tensorrt_llm::DataType::kBOOL: return randomize(1); + case tensorrt_llm::DataType::kFLOAT: return randomize(1); + case tensorrt_llm::DataType::kINT8: return randomize(1); + case tensorrt_llm::DataType::kINT32: return randomize(1); + case tensorrt_llm::DataType::kINT64: return randomize(1); + case tensorrt_llm::DataType::kUINT8: return randomize(1); default: return; } } diff --git a/cpp/tensorrt_llm/layers/lookaheadPoolManager.cpp b/cpp/tensorrt_llm/layers/lookaheadPoolManager.cpp index 5954bc520ad0..397b4262226a 100644 --- a/cpp/tensorrt_llm/layers/lookaheadPoolManager.cpp +++ b/cpp/tensorrt_llm/layers/lookaheadPoolManager.cpp @@ -16,6 +16,7 @@ #include "tensorrt_llm/layers/lookaheadPoolManager.h" #include "tensorrt_llm/common/logger.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/layers/lookaheadDecodingUtils.h" #include @@ -67,7 +68,7 @@ void LookaheadPoolManager::accept(TensorConstPtr const& prompt, SizeType32 level for (SizeType32 ti = 0; ti + level - 1 < length; ti++) { auto key = promptRange[ti]; - TensorPtr ngram = BufferManager::cpu(ITensor::makeShape({level - 1}), nvinfer1::DataType::kINT32); + TensorPtr ngram = BufferManager::cpu(ITensor::makeShape({level - 1}), tensorrt_llm::DataType::kINT32); BufferRange sourceRange(*ITensor::slice(prompt, ti + 1, level - 1)); BufferRange ngramRange(*ngram); std::copy(sourceRange.begin(), sourceRange.end(), ngramRange.begin()); @@ -107,7 +108,7 @@ void LookaheadPoolManager::update(TensorConstPtr const& keyTokens, TensorConstPt for (SizeType32 wi = 0; wi < window; wi++) { TensorConstPtr source = ITensor::at(ngramTokens, {wi}); - TensorPtr ngram = BufferManager::cpu(source->getShape(), nvinfer1::DataType::kINT32); + TensorPtr ngram = BufferManager::cpu(source->getShape(), tensorrt_llm::DataType::kINT32); BufferRange sourceRange(*source); BufferRange ngramRange(*ngram); std::copy(sourceRange.begin(), sourceRange.end(), ngramRange.begin()); diff --git a/cpp/tensorrt_llm/layers/medusaDecodingLayer.cpp b/cpp/tensorrt_llm/layers/medusaDecodingLayer.cpp index 40eff62c17d6..9e4098b34ebf 100644 --- a/cpp/tensorrt_llm/layers/medusaDecodingLayer.cpp +++ b/cpp/tensorrt_llm/layers/medusaDecodingLayer.cpp @@ -16,6 +16,7 @@ #include "medusaDecodingLayer.h" #include "tensorrt_llm/common/nvtxUtils.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/kernels/decodingCommon.h" #include "tensorrt_llm/kernels/samplingTopKKernels.h" #include "tensorrt_llm/kernels/speculativeDecoding/medusaDecodingKernels.h" @@ -88,10 +89,10 @@ void MedusaDecodingLayer::allocateBuffer() mTiledBatchSlotsSetup = BufferManager::pinnedPool( ITensor::makeShape({static_cast(mDecoderDomain.getBatchSize() * maxDraftPathLen)}), - nvinfer1::DataType::kINT32); + tensorrt_llm::DataType::kINT32); mTiledBatchSlotsForward = BufferManager::pinnedPool( ITensor::makeShape({static_cast(mDecoderDomain.getBatchSize() * maxDraftPathLen)}), - nvinfer1::DataType::kINT32); + tensorrt_llm::DataType::kINT32); mMedusaInputLogitsPtrs = BufferManager::pinnedPool( ITensor::makeShape({static_cast(mDecoderDomain.getBatchSize() * maxDraftPathLen)}), TRTDataType::value); diff --git a/cpp/tensorrt_llm/layers/penaltyLayer.cpp b/cpp/tensorrt_llm/layers/penaltyLayer.cpp index c6c57ca5034d..c72b8e463bc6 100644 --- a/cpp/tensorrt_llm/layers/penaltyLayer.cpp +++ b/cpp/tensorrt_llm/layers/penaltyLayer.cpp @@ -18,6 +18,7 @@ #include "penaltyLayer.h" #include "tensorrt_llm/common/cudaUtils.h" #include "tensorrt_llm/common/nvtxUtils.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/kernels/penaltyKernels.h" #include "tensorrt_llm/kernels/penaltyTypes.h" #include "tensorrt_llm/layers/defaultDecodingParams.h" @@ -84,11 +85,11 @@ void PenaltyLayer::allocateWorkspace() auto const workspaceSize = mDecoderDomain.getBatchSize() * mDecoderDomain.getMaxDecodingTokens() * mConfiguredBeamWidth * mDecoderDomain.getVocabSize() * 2; - mPenaltyWorkspaceDevice = mBufferManager->gpu(workspaceSize, nvinfer1::DataType::kINT32); + mPenaltyWorkspaceDevice = mBufferManager->gpu(workspaceSize, tensorrt_llm::DataType::kINT32); if (mDecodingMode.isBeamSearch()) { - mPenaltyWorkspacePrevDevice = mBufferManager->gpu(workspaceSize, nvinfer1::DataType::kINT32); + mPenaltyWorkspacePrevDevice = mBufferManager->gpu(workspaceSize, tensorrt_llm::DataType::kINT32); } } @@ -111,27 +112,27 @@ void PenaltyLayer::allocateBuffer() if (mDecodingMode.isUseTemperature()) { - mTemperatureDevice = mBufferManager->gpu(batchSizeShape, nvinfer1::DataType::kFLOAT); + mTemperatureDevice = mBufferManager->gpu(batchSizeShape, tensorrt_llm::DataType::kFLOAT); } if (mDecodingMode.isUseRepetitionPenalty()) { - mRepetitionPenaltyDevice = mBufferManager->gpu(batchSizeShape, nvinfer1::DataType::kFLOAT); + mRepetitionPenaltyDevice = mBufferManager->gpu(batchSizeShape, tensorrt_llm::DataType::kFLOAT); } if (mDecodingMode.isUsePresencePenalty()) { - mPresencePenaltyDevice = mBufferManager->gpu(batchSizeShape, nvinfer1::DataType::kFLOAT); + mPresencePenaltyDevice = mBufferManager->gpu(batchSizeShape, tensorrt_llm::DataType::kFLOAT); } if (mDecodingMode.isUseFrequencyPenalty()) { - mFrequencyPenaltyDevice = mBufferManager->gpu(batchSizeShape, nvinfer1::DataType::kFLOAT); + mFrequencyPenaltyDevice = mBufferManager->gpu(batchSizeShape, tensorrt_llm::DataType::kFLOAT); } if (mDecodingMode.isUseMinLength()) { - mMinLengthDevice = mBufferManager->gpu(batchSizeShape, nvinfer1::DataType::kINT32); + mMinLengthDevice = mBufferManager->gpu(batchSizeShape, tensorrt_llm::DataType::kINT32); } if (mDecodingMode.isUseOccurrencePenalty()) { - mPromptIgnoreLengthDevice = mBufferManager->gpu(batchSizeShape, nvinfer1::DataType::kINT32); + mPromptIgnoreLengthDevice = mBufferManager->gpu(batchSizeShape, tensorrt_llm::DataType::kINT32); } auto const logitsPtrDeviceDesc = std::make_pair(batchSizeShape, TRTDataType::value); diff --git a/cpp/tensorrt_llm/nanobind/CMakeLists.txt b/cpp/tensorrt_llm/nanobind/CMakeLists.txt index b523ae193871..4d6fbf9c2607 100755 --- a/cpp/tensorrt_llm/nanobind/CMakeLists.txt +++ b/cpp/tensorrt_llm/nanobind/CMakeLists.txt @@ -14,7 +14,6 @@ set(SRCS batch_manager/llmRequest.cpp common/tllmExceptions.cpp executor/bindings.cpp - executor/executor.cpp executor/executorConfig.cpp executor/request.cpp process_group/bindings.cpp @@ -23,7 +22,6 @@ set(SRCS runtime/moeBindings.cpp suffixAutomaton/bindings.cpp testing/kvCacheManagerTestUtilBinding.cpp - testing/modelSpecBinding.cpp userbuffers/bindings.cpp thop/bindings.cpp ../runtime/ipcNvlsMemory.cu diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/algorithms.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/algorithms.cpp index 4070811b2d72..c13466565342 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/algorithms.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/algorithms.cpp @@ -23,11 +23,11 @@ #include "tensorrt_llm/batch_manager/createNewDecoderRequests.h" #include "tensorrt_llm/batch_manager/kvCacheManager.h" #include "tensorrt_llm/batch_manager/llmRequest.h" -#include "tensorrt_llm/batch_manager/logitsPostProcessor.h" #include "tensorrt_llm/batch_manager/medusaBuffers.h" #include "tensorrt_llm/batch_manager/microBatchScheduler.h" #include "tensorrt_llm/batch_manager/pauseRequests.h" #include "tensorrt_llm/batch_manager/peftCacheManager.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/nanobind/common/customCasters.h" #include "tensorrt_llm/runtime/decoderState.h" #include "tensorrt_llm/runtime/torch.h" @@ -129,13 +129,6 @@ void tensorrt_llm::nanobind::batch_manager::algorithms::initBindings(nb::module_ nb::call_guard()) .def("name", [](AllocateKvCache const&) { return AllocateKvCache::name; }); - nb::class_(m, LogitsPostProcessor::name) - .def(nb::init<>()) - .def("__call__", &LogitsPostProcessor::operator(), nb::arg("decoder_input_buffers"), - nb::arg("replicate_logits_post_processor"), nb::arg("world_config"), nb::arg("stream"), - nb::arg("logits_post_processor_batched") = std::nullopt) - .def("name", [](LogitsPostProcessor const&) { return LogitsPostProcessor::name; }); - nb::class_(m, CreateNewDecoderRequests::name) .def(nb::init(), nb::arg("speculative_decoding_fast_logits"), nb::arg("is_leader_in_orch_mode"), nb::arg("is_normalize_log_probs")) @@ -143,7 +136,7 @@ void tensorrt_llm::nanobind::batch_manager::algorithms::initBindings(nb::module_ "__call__", [](CreateNewDecoderRequests& self, tr::ModelConfig const& modelConfig, tr::WorldConfig const& worldConfig, executor::DecodingConfig const& decodingConfig, RequestVector const& contextRequests, - nvinfer1::DataType logitsType, DecoderInputBuffers& inputBuffers, + tensorrt_llm::DataType logitsType, DecoderInputBuffers& inputBuffers, runtime::decoder::DecoderState& decoderState, tensorrt_llm::runtime::CudaStream const& runtimeStream, tensorrt_llm::runtime::CudaStream const& decoderStream, SizeType32 maxSequenceLength, SizeType32 beamWidth) diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp index 0846663dafad..f8750a718df7 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp @@ -24,6 +24,7 @@ #include "tensorrt_llm/batch_manager/peftCacheManager.h" #include "tensorrt_llm/batch_manager/rnnStateManager.h" #include "tensorrt_llm/batch_manager/sequenceSlotManager.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/nanobind/common/bindTypes.h" #include "tensorrt_llm/runtime/gptDecoderBatched.h" #include "tensorrt_llm/runtime/runtimeKernels.h" @@ -491,7 +492,7 @@ void initBindings(nb::module_& m) nb::arg("max_num_sequences"), nb::arg("model_config"), nb::arg("world_config"), nb::arg("buffer_manager"), nb::call_guard()) .def(nb::init const&, tr::SizeType32>(), nb::arg("d_state"), nb::arg("d_conv"), nb::arg("num_heads"), nb::arg("n_groups"), nb::arg("head_dim"), nb::arg("max_batch_size"), nb::arg("world_config"), nb::arg("stream"), nb::arg("dtype"), diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp index 70be9c2f353e..d11838f68af5 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp @@ -20,6 +20,7 @@ #include "tensorrt_llm/batch_manager/kvCacheManager.h" #include "tensorrt_llm/batch_manager/rnnStateManager.h" #include "tensorrt_llm/common/bindingUtils.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/executor/executor.h" #include "tensorrt_llm/nanobind/common/customCasters.h" #include @@ -120,7 +121,7 @@ void tb::CacheTransceiverBindings::initBindings(nb::module_& m) nb::class_(m, "CacheTransceiver") .def(nb::init, SizeType32, SizeType32, - runtime::WorldConfig, std::vector, nvinfer1::DataType, + runtime::WorldConfig, std::vector, tensorrt_llm::DataType, executor::kv_cache::CacheState::AttentionType, std::optional, std::vector>(), nb::arg("cache_manager"), nb::arg("num_kv_heads_per_layer"), nb::arg("size_per_head"), diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp index 12c352e4a427..b1c4391c0e26 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp @@ -18,6 +18,7 @@ #include "kvCacheManager.h" #include "tensorrt_llm/batch_manager/kvCacheManager.h" #include "tensorrt_llm/batch_manager/peftCacheManager.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/nanobind/common/bindTypes.h" #include "tensorrt_llm/nanobind/common/customCasters.h" #include "tensorrt_llm/runtime/torch.h" @@ -348,8 +349,8 @@ void tb::kv_cache_manager::KVCacheManagerBindings::initBindings(nb::module_& m) nb::class_(m, "PoolConfiguration") .def(nb::init<>()) - .def(nb::init(), nb::arg("window_size"), nb::arg("size_per_head"), - nb::arg("dtype")) + .def(nb::init(), nb::arg("window_size"), + nb::arg("size_per_head"), nb::arg("dtype")) .def_rw("window_size", &tbk::PoolConfiguration::windowSize) .def_rw("size_per_head", &tbk::PoolConfiguration::sizePerHead) .def_rw("dtype", &tbk::PoolConfiguration::dtype); @@ -661,8 +662,8 @@ void tb::kv_cache_manager::KVCacheManagerBindings::initBindings(nb::module_& m) nb::class_(m, "KVCacheManager") .def(nb::init const&, SizeType32, SizeType32, std::map> const&, SizeType32, SizeType32, - std::vector const&, nvinfer1::DataType, SizeType32, int64_t, SizeType32, SizeType32, bool, - tbk::CacheType, std::optional, + std::vector const&, tensorrt_llm::DataType, SizeType32, int64_t, SizeType32, SizeType32, + bool, tbk::CacheType, std::optional, std::shared_ptr, bool, bool, std::shared_ptr, bool, SizeType32, SizeType32, bool, std::optional, std::vector const&>(), diff --git a/cpp/tensorrt_llm/nanobind/bindings.cpp b/cpp/tensorrt_llm/nanobind/bindings.cpp index db263fa639f9..79dd147092bc 100644 --- a/cpp/tensorrt_llm/nanobind/bindings.cpp +++ b/cpp/tensorrt_llm/nanobind/bindings.cpp @@ -32,6 +32,7 @@ #include "tensorrt_llm/batch_manager/peftCacheManagerConfig.h" #include "tensorrt_llm/common/quantization.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/nanobind/batch_manager/algorithms.h" #include "tensorrt_llm/nanobind/batch_manager/bindings.h" #include "tensorrt_llm/nanobind/batch_manager/buffers.h" @@ -46,7 +47,6 @@ #include "tensorrt_llm/nanobind/runtime/bindings.h" #include "tensorrt_llm/nanobind/suffixAutomaton/bindings.h" #include "tensorrt_llm/nanobind/testing/kvCacheManagerTestUtilBinding.h" -#include "tensorrt_llm/nanobind/testing/modelSpecBinding.h" #include "tensorrt_llm/nanobind/thop/bindings.h" #include "tensorrt_llm/nanobind/userbuffers/bindings.h" #include "tensorrt_llm/runtime/common.h" @@ -168,17 +168,17 @@ NB_MODULE(TRTLLM_NB_MODULE, m) .def_rw("host_cache_size", &tb::PeftCacheManagerConfig::hostCacheSize) .def_rw("lora_prefetch_dir", &tb::PeftCacheManagerConfig::loraPrefetchDir); - nb::enum_(m, "DataType") - .value("FLOAT", nvinfer1::DataType::kFLOAT) - .value("HALF", nvinfer1::DataType::kHALF) - .value("INT8", nvinfer1::DataType::kINT8) - .value("INT32", nvinfer1::DataType::kINT32) - .value("BOOL", nvinfer1::DataType::kBOOL) - .value("UINT8", nvinfer1::DataType::kUINT8) - .value("FP8", nvinfer1::DataType::kFP8) - .value("BF16", nvinfer1::DataType::kBF16) - .value("INT64", nvinfer1::DataType::kINT64) - .value("NVFP4", nvinfer1::DataType::kFP4) + nb::enum_(m, "DataType") + .value("FLOAT", tensorrt_llm::DataType::kFLOAT) + .value("HALF", tensorrt_llm::DataType::kHALF) + .value("INT8", tensorrt_llm::DataType::kINT8) + .value("INT32", tensorrt_llm::DataType::kINT32) + .value("BOOL", tensorrt_llm::DataType::kBOOL) + .value("UINT8", tensorrt_llm::DataType::kUINT8) + .value("FP8", tensorrt_llm::DataType::kFP8) + .value("BF16", tensorrt_llm::DataType::kBF16) + .value("INT64", tensorrt_llm::DataType::kINT64) + .value("NVFP4", tensorrt_llm::DataType::kFP4) .export_values(); nb::enum_(m, "GptModelVariant") @@ -295,7 +295,7 @@ NB_MODULE(TRTLLM_NB_MODULE, m) .def(nb::self != nb::self); nb::class_(m, "ModelConfig") - .def(nb::init(), + .def(nb::init(), nb::arg("vocab_size"), nb::arg("num_layers"), nb::arg("num_attention_layers"), nb::arg("num_rnn_layers"), nb::arg("num_heads"), nb::arg("hidden_size"), nb::arg("data_type")) .def_prop_ro("vocab_size", &tr::ModelConfig::getVocabSize) @@ -512,7 +512,6 @@ NB_MODULE(TRTLLM_NB_MODULE, m) tensorrt_llm::nanobind::process_group::initBindings(mInternalProcessGroup); tpb::Buffers::initBindings(mInternalBatchManager); tensorrt_llm::nanobind::runtime::initBindings(mInternalRuntime); - tensorrt_llm::nanobind::testing::initBindings(mInternalTesting); tensorrt_llm::nanobind::testing::initKvCacheTestUtilBindings(mInternalTesting); tpb::initBindings(mInternalBatchManager); diff --git a/cpp/tensorrt_llm/nanobind/executor/bindings.cpp b/cpp/tensorrt_llm/nanobind/executor/bindings.cpp index b0ad31b7347e..a8d2301fa43d 100644 --- a/cpp/tensorrt_llm/nanobind/executor/bindings.cpp +++ b/cpp/tensorrt_llm/nanobind/executor/bindings.cpp @@ -16,7 +16,6 @@ */ #include "bindings.h" -#include "executor.h" #include "executorConfig.h" #include "request.h" #include "tensorrt_llm/executor/executor.h" @@ -287,7 +286,6 @@ void initBindings(nb::module_& m) tensorrt_llm::nanobind::executor::initRequestBindings(m); tensorrt_llm::nanobind::executor::initConfigBindings(m); - tensorrt_llm::nanobind::executor::Executor::initBindings(m); } } // namespace tensorrt_llm::nanobind::executor diff --git a/cpp/tensorrt_llm/nanobind/executor/executor.cpp b/cpp/tensorrt_llm/nanobind/executor/executor.cpp deleted file mode 100644 index 34cc8182d1bb..000000000000 --- a/cpp/tensorrt_llm/nanobind/executor/executor.cpp +++ /dev/null @@ -1,225 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2022-2025 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. - */ - -#include "executor.h" -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/logger.h" -#include "tensorrt_llm/executor/tensor.h" -#include "tensorrt_llm/nanobind/common/customCasters.h" - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -namespace nb = nanobind; -namespace tle = tensorrt_llm::executor; - -namespace nanobind::detail -{ - -template <> -struct dtype_traits -{ - static constexpr dlpack::dtype value{ - (uint8_t) dlpack::dtype_code::Float, // type code - 16, // size in bits - 1 // lanes (simd), usually set to 1 - }; - static constexpr auto name = const_name("float16"); -}; -} // namespace nanobind::detail - -namespace -{ -tle::Tensor numpyToTensor(nb::object const& object) -{ - std::string dtype_name = nb::cast(object.attr("dtype").attr("name")); - nb::object metadata = object.attr("dtype").attr("metadata"); - - tle::DataType dtype; - if (dtype_name == "float16") - { - dtype = tle::DataType::kFP16; - } - else if (dtype_name == "float32") - { - dtype = tle::DataType::kFP32; - } - else if (dtype_name == "int8") - { - dtype = tle::DataType::kINT8; - } - else if (dtype_name == "int32") - { - dtype = tle::DataType::kINT32; - } - else if (dtype_name == "int64") - { - dtype = tle::DataType::kINT64; - } - else if (dtype_name == "void8" && !metadata.is_none() && nb::cast(metadata["dtype"]) == "float8") - { - dtype = tle::DataType::kFP8; - } - else if (dtype_name == "void16" && !metadata.is_none() && nb::cast(metadata["dtype"]) == "bfloat16") - { - dtype = tle::DataType::kBF16; - } - else - { - TLLM_THROW("Unsupported numpy dtype."); - } - - nb::object array_interface = object.attr("__array_interface__"); - nb::object shape_obj = array_interface["shape"]; - std::vector dims; - dims.reserve(nb::len(shape_obj)); - - for (size_t i = 0; i < nb::len(shape_obj); ++i) - { - dims.push_back(nb::cast(shape_obj[i])); - } - - nb::object data_obj = array_interface["data"]; - uintptr_t addr = nb::cast(data_obj[0]); - void* data_ptr = reinterpret_cast(addr); - tle::Shape shape(dims.data(), dims.size()); - return tle::Tensor::of(dtype, data_ptr, shape); -} - -} // namespace - -namespace tensorrt_llm::nanobind::executor -{ - -Executor::Executor( - std::filesystem::path const& modelPath, tle::ModelType modelType, tle::ExecutorConfig const& executorConfig) -{ - mExecutor = std::make_unique(modelPath, modelType, executorConfig); -} - -Executor::Executor(std::filesystem::path const& encoderModelPath, std::filesystem::path const& decoderModelPath, - tle::ModelType modelType, tle::ExecutorConfig const& executorConfig) -{ - mExecutor = std::make_unique(encoderModelPath, decoderModelPath, modelType, executorConfig); -} - -Executor::Executor(nb::bytes const& engineBuffer, std::string const& jsonConfigStr, tle::ModelType modelType, - tle::ExecutorConfig const& executorConfig, std::optional managedWeights) -{ - uint8_t const* data = static_cast(engineBuffer.data()); - size_t size = engineBuffer.size(); - std::optional> managedWeightsMap = std::nullopt; - if (managedWeights.has_value() && !managedWeights.value().empty()) - { - managedWeightsMap = std::map(); - for (auto const& [rawName, rawArray] : managedWeights.value()) - { - std::string name = nb::cast(rawName); - nb::object array_obj = nb::cast(rawArray); - managedWeightsMap->emplace(name, numpyToTensor(array_obj)); - } - } - mExecutor = std::make_unique( - tle::BufferView(data, size), jsonConfigStr, modelType, executorConfig, managedWeightsMap); -} - -Executor::Executor(std::string const& encoderEngineBuffer, std::string const& encoderJsonConfigStr, - std::string const& decoderEngineBuffer, std::string const& decoderJsonConfigStr, tle::ModelType modelType, - tle::ExecutorConfig const& executorConfig) -{ - uint8_t const* encoderData = reinterpret_cast(encoderEngineBuffer.data()); - size_t encoderSize = encoderEngineBuffer.size(); - uint8_t const* decoderData = reinterpret_cast(decoderEngineBuffer.data()); - size_t decoderSize = decoderEngineBuffer.size(); - mExecutor = std::make_unique(tle::BufferView(encoderData, encoderSize), encoderJsonConfigStr, - tle::BufferView(decoderData, decoderSize), decoderJsonConfigStr, modelType, executorConfig); -} - -nb::object Executor::enter() -{ - TLLM_CHECK(static_cast(mExecutor)); - return nb::cast(this); -} - -void Executor::exit( - [[maybe_unused]] nb::handle type, [[maybe_unused]] nb::handle value, [[maybe_unused]] nb::handle traceback) -{ - shutdown(); - mExecutor = nullptr; -} - -void Executor::shutdown() -{ - // NOTE: we must release the GIL here. Executor has spawned a thread for the execution loop. That thread must be - // able to do forward progress for the shutdown process to succeed. It takes the GIL during its callbacks, so - // we release it now. Note that we shouldn't do anything related to python objects after that. - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - nb::gil_scoped_release release; - mExecutor->shutdown(); - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void Executor::initBindings(nb::module_& m) -{ - nb::class_(m, "Executor") - .def(nb::init(), - nb::arg("model_path"), nb::arg("model_type"), nb::arg("executor_config")) - .def(nb::init(), - nb::arg("encoder_model_path"), nb::arg("decoder_model_path"), nb::arg("model_type"), - nb::arg("executor_config")) - .def(nb::init(), - nb::arg("engine_buffer"), nb::arg("json_config_str"), nb::arg("model_type"), nb::arg("executor_config"), - nb::arg("managed_weights") = nb::dict()) - .def(nb::init(), - nb::arg("encoder_engine_buffer"), nb::arg("encoder_json_config_str"), nb::arg("decoder_engine_buffer"), - nb::arg("decoder_json_config_str"), nb::arg("model_type"), nb::arg("executor_config")) - .def("shutdown", &Executor::shutdown) - .def("__enter__", &Executor::enter) - .def("__exit__", &Executor::exit, nb::arg("type").none(), nb::arg("value").none(), nb::arg("traceback").none()) - .def("enqueue_request", &Executor::enqueueRequest, nb::arg("request")) - .def("enqueue_requests", &Executor::enqueueRequests, nb::arg("requests")) - .def("await_responses", - nb::overload_cast const&>(&Executor::awaitResponses), - nb::arg("timeout") = nb::none()) - .def("await_responses", - nb::overload_cast const&>( - &Executor::awaitResponses), - nb::arg("id"), nb::arg("timeout") = nb::none()) - .def("await_responses", - nb::overload_cast const&, std::optional const&>( - &Executor::awaitResponses), - nb::arg("ids"), nb::arg("timeout") = nb::none()) - .def("get_num_responses_ready", &Executor::getNumResponsesReady, nb::arg("id") = nb::none()) - .def("cancel_request", &Executor::cancelRequest, nb::arg("id") = nb::none()) - .def("get_latest_iteration_stats", &Executor::getLatestIterationStats) - .def("get_latest_request_stats", &Executor::getLatestRequestStats) - .def("get_latest_debug_tensors", &Executor::getLatestDebugTensors) - .def("can_enqueue_requests", &Executor::canEnqueueRequests) - .def("get_kv_cache_event_manager", &Executor::getKVCacheEventManager); -} - -} // namespace tensorrt_llm::nanobind::executor diff --git a/cpp/tensorrt_llm/nanobind/executor/executor.h b/cpp/tensorrt_llm/nanobind/executor/executor.h deleted file mode 100644 index 22c24abb4bfd..000000000000 --- a/cpp/tensorrt_llm/nanobind/executor/executor.h +++ /dev/null @@ -1,129 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2022-2025 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. - */ - -#pragma once - -#include "tensorrt_llm/executor/executor.h" -#include "tensorrt_llm/executor/types.h" -#include - -namespace nb = nanobind; -namespace tle = tensorrt_llm::executor; - -namespace tensorrt_llm::nanobind::executor -{ - -class Executor -{ -public: - Executor( - std::filesystem::path const& modelPath, tle::ModelType modelType, tle::ExecutorConfig const& executorConfig); - - Executor(std::filesystem::path const& encoderModelPath, std::filesystem::path const& decoderModelPath, - tle::ModelType modelType, tle::ExecutorConfig const& executorConfig); - - Executor(nb::bytes const& engineBuffer, std::string const& jsonConfigStr, tle::ModelType modelType, - tle::ExecutorConfig const& executorConfig, std::optional managedWeights); - - Executor(std::string const& encoderEngineBuffer, std::string const& encoderJsonConfigStr, - std::string const& decoderEngineBuffer, std::string const& decoderJsonConfigStr, tle::ModelType modelType, - tle::ExecutorConfig const& executorConfig); - - nb::object enter(); - void exit( - [[maybe_unused]] nb::handle type, [[maybe_unused]] nb::handle value, [[maybe_unused]] nb::handle traceback); - void shutdown(); - - [[nodiscard]] tle::IdType enqueueRequest(tle::Request const& request) - { - return mExecutor->enqueueRequest(request); - } - - [[nodiscard]] std::vector enqueueRequests(std::vector const& requests) - { - return mExecutor->enqueueRequests(requests); - } - - [[nodiscard]] std::vector awaitResponses( - std::optional const& timeout = std::nullopt) - { - // Await responses blocks until a response is received. Release GIL so that it can be ran in a background - // thread. - nb::gil_scoped_release release; - return mExecutor->awaitResponses(timeout); - } - - [[nodiscard]] std::vector awaitResponses( - tle::IdType const& requestId, std::optional const& timeout = std::nullopt) - { - // Await responses blocks until a response is received. Release GIL so that it can be ran in a background - // thread. - nb::gil_scoped_release release; - return mExecutor->awaitResponses(requestId, timeout); - } - - [[nodiscard]] std::vector> awaitResponses(std::vector const& requestIds, - std::optional const& timeout = std::nullopt) - { - // Await responses blocks until a response is received. Release GIL so that it can be ran in a background - // thread. - nb::gil_scoped_release release; - return mExecutor->awaitResponses(requestIds, timeout); - } - - [[nodiscard]] tle::SizeType32 getNumResponsesReady(std::optional const& requestId = std::nullopt) const - { - return mExecutor->getNumResponsesReady(requestId); - } - - void cancelRequest(tle::IdType requestId) - { - mExecutor->cancelRequest(requestId); - } - - std::deque getLatestIterationStats() - { - return mExecutor->getLatestIterationStats(); - } - - std::deque getLatestRequestStats() - { - return mExecutor->getLatestRequestStats(); - } - - std::deque getLatestDebugTensors() - { - return mExecutor->getLatestDebugTensors(); - } - - [[nodiscard]] bool canEnqueueRequests() const - { - return mExecutor->canEnqueueRequests(); - } - - [[nodiscard]] std::optional> getKVCacheEventManager() const - { - return mExecutor->getKVCacheEventManager(); - } - - static void initBindings(nb::module_& m); - -private: - std::unique_ptr mExecutor; -}; - -} // namespace tensorrt_llm::nanobind::executor diff --git a/cpp/tensorrt_llm/nanobind/executor/executorConfig.cpp b/cpp/tensorrt_llm/nanobind/executor/executorConfig.cpp index 830f30ab9c67..acd33c0df769 100644 --- a/cpp/tensorrt_llm/nanobind/executor/executorConfig.cpp +++ b/cpp/tensorrt_llm/nanobind/executor/executorConfig.cpp @@ -171,6 +171,9 @@ void initConfigBindings(nb::module_& m) .def("__getstate__", kvCacheConfigGetstate) .def("__setstate__", kvCacheConfigSetstate); + // Deprecated: orchestrator mode is non-functional (its executorWorker binary was + // removed with the TensorRT backend); binding kept for compatibility, removal is a + // follow-up pending API-stability review. nb::class_(m, "OrchestratorConfig") .def(nb::init, bool>(), nb::arg("is_orchestrator") = true, nb::arg("worker_executable_path") = "", nb::arg("orch_leader_comm").none() = nullptr, diff --git a/cpp/tensorrt_llm/nanobind/runtime/bindings.cpp b/cpp/tensorrt_llm/nanobind/runtime/bindings.cpp index 6d5d70aafb6b..eec3cd79bac1 100644 --- a/cpp/tensorrt_llm/nanobind/runtime/bindings.cpp +++ b/cpp/tensorrt_llm/nanobind/runtime/bindings.cpp @@ -18,6 +18,7 @@ #include "bindings.h" #include "hostfunc.h" #include "moeBindings.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/kernels/communicationKernels/allReduceWorkspace.h" #include "tensorrt_llm/kernels/communicationKernels/customLowPrecisionAllReduceKernels.h" #include "tensorrt_llm/kernels/customAllReduceKernels.h" @@ -39,7 +40,6 @@ #include "tensorrt_llm/runtime/loraCache.h" #include "tensorrt_llm/runtime/mcastGPUBuffer.h" #include "tensorrt_llm/runtime/speculativeDecodingMode.h" -#include "tensorrt_llm/runtime/tllmRuntime.h" #include "tensorrt_llm/runtime/torchView.h" #include "tensorrt_llm/runtime/virtualMemory.h" @@ -68,7 +68,7 @@ class PyIGptDecoder : public tr::IGptDecoder void setup(tr::SamplingConfig const& samplingConfig, size_t batchSize, tr::DecodingInput::TensorConstPtr const& batchSlots, std::optional const& output = std::nullopt, - std::optional explicitDraftTokensDType = std::nullopt, + std::optional explicitDraftTokensDType = std::nullopt, std::optional> const& lookaheadPrompt = std::nullopt, std::optional> const& lookaheadAlgoConfigs = std::nullopt) override { @@ -125,47 +125,6 @@ void initBindings(nb::module_& m) .def("materialize_with_tag", &tr::CudaVirtualMemoryManager::materializeWithTag, nb::arg("tag"), nb::call_guard()); - nb::class_(m, "TllmRuntime") - .def( - "__init__", - [](tr::TllmRuntime* self, std::filesystem::path engine_path, float gpu_weights_percent = 1.0f, - bool use_shape_inference = true) - { - // Using default logger by passing nullptr - new (self) - tr::TllmRuntime(tr::RawEngine(engine_path), nullptr, gpu_weights_percent, use_shape_inference); - }, - nb::arg("engine_path"), nb::arg("gpu_weights_percent") = 1.0f, nb::arg("use_shape_inference") = true) - .def( - "__init__", - [](tr::TllmRuntime* self, nb::ndarray engine_buffer, float gpu_weights_percent = 1.0f, - bool use_shape_inference = true) - { - if (engine_buffer.ndim() != 1) - throw std::runtime_error("Expected 1-D array for engine buffer"); - new (self) tr::TllmRuntime(tr::RawEngine(engine_buffer.data(), engine_buffer.size()), nullptr, - gpu_weights_percent, use_shape_inference); - }, - nb::arg("engine_buffer"), nb::arg("gpu_weights_percent") = 1.0f, nb::arg("use_shape_inference") = true) - .def_prop_ro("num_contexts", &tr::TllmRuntime::getNbContexts) - .def_prop_ro("num_profiles", &tr::TllmRuntime::getNbProfiles) - .def("get_opt_profile_id", &tr::TllmRuntime::getOptProfileId, nb::arg("num_tokens"), nb::arg("split_points"), - nb::call_guard()) - .def("clear_contexts", &tr::TllmRuntime::clearContexts, nb::call_guard()) - .def("execute_context", &tr::TllmRuntime::executeContext, nb::arg("context_id"), - nb::call_guard()) - .def_prop_ro("stream_ptr", &tr::TllmRuntime::getStreamPtr) - .def_prop_ro("buffer_manager", - static_cast(&tr::TllmRuntime::getBufferManager)) - .def("set_layer_profiler", &tr::TllmRuntime::setLayerProfiler, nb::call_guard()) - .def("has_layer_profiler", &tr::TllmRuntime::hasLayerProfiler, nb::arg("context_id"), - nb::call_guard()) - .def_prop_ro("layer_profiler_info", &tr::TllmRuntime::getLayerProfileInfo) - .def("report_to_profiler", &tr::TllmRuntime::reportToProfiler, nb::arg("context_id"), - nb::call_guard()) - .def_prop_ro("logits_dtype_from_engine", - [](tr::TllmRuntime& self) { return self.getEngine().getTensorDataType("logits"); }); - nb::class_(m, "LookaheadDecodingBuffers") .def(nb::init(), nb::arg("max_num_sequences"), nb::arg("max_tokens_per_step"), nb::arg("buffer_manager"), nb::call_guard()) @@ -204,7 +163,7 @@ void initBindings(nb::module_& m) "setup", [](tr::IGptDecoder& self, tr::SamplingConfig const& samplingConfig, size_t batchSize, at::Tensor const& batchSlots, std::optional const& output = std::nullopt, - std::optional explicitDraftTokensDType = std::nullopt, + std::optional explicitDraftTokensDType = std::nullopt, std::optional> const& lookaheadPrompt = std::nullopt, std::optional> const& lookaheadAlgoConfigs = std::nullopt) { diff --git a/cpp/tensorrt_llm/nanobind/testing/modelSpecBinding.cpp b/cpp/tensorrt_llm/nanobind/testing/modelSpecBinding.cpp deleted file mode 100644 index caef94c5defd..000000000000 --- a/cpp/tensorrt_llm/nanobind/testing/modelSpecBinding.cpp +++ /dev/null @@ -1,87 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2023-2025 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. - */ - -#include "modelSpecBinding.h" -#include "tensorrt_llm/nanobind/common/customCasters.h" -#include "tensorrt_llm/testing/modelSpec.h" - -#include - -namespace nb = nanobind; -using tensorrt_llm::testing::ModelSpec; -using tensorrt_llm::testing::KVCacheType; -using tensorrt_llm::testing::QuantMethod; -using tensorrt_llm::testing::OutputContentType; - -namespace tensorrt_llm::nanobind::testing -{ - -void initBindings(nb::module_& m) -{ - nb::enum_(m, "QuantMethod", nb::is_arithmetic(), "Quantization Method") - .value("NONE", QuantMethod::kNONE, "No Quantization") - .value("SMOOTH_QUANT", QuantMethod::kSMOOTH_QUANT, "Smooth Quantization"); - - nb::enum_(m, "OutputContentType", nb::is_arithmetic(), "Output Content Type") - .value("NONE", OutputContentType::kNONE, "No Output Content") - .value("CONTEXT_LOGITS", OutputContentType::kCONTEXT_LOGITS, "Context Logits") - .value("GENERATION_LOGITS", OutputContentType::kGENERATION_LOGITS, "Generation Logits") - .value("LOG_PROBS", OutputContentType::kLOG_PROBS, "Log Probs") - .value("CUM_LOG_PROBS", OutputContentType::kCUM_LOG_PROBS, "Cumulative Log"); - - nb::class_(m, "ModelSpec") - .def(nb::init()) - .def("use_gpt_plugin", &ModelSpec::useGptAttentionPlugin, nb::rv_policy::reference_internal) - .def("use_packed_input", &ModelSpec::usePackedInput, nb::rv_policy::reference_internal) - .def("set_kv_cache_type", &ModelSpec::setKVCacheType, nb::rv_policy::reference_internal) - .def("use_decoder_per_request", &ModelSpec::useDecoderPerRequest, nb::rv_policy::reference_internal) - .def("use_tensor_parallelism", &ModelSpec::useTensorParallelism, nb::rv_policy::reference_internal) - .def("use_pipeline_parallelism", &ModelSpec::usePipelineParallelism, nb::rv_policy::reference_internal) - .def("use_context_parallelism", &ModelSpec::useContextParallelism, nb::rv_policy::reference_internal) - .def("set_draft_tokens", &ModelSpec::setDraftTokens, nb::rv_policy::reference_internal) - .def("use_accept_by_logits", &ModelSpec::useAcceptByLogits, nb::rv_policy::reference_internal) - .def("use_mamba_plugin", &ModelSpec::useMambaPlugin, nb::rv_policy::reference_internal) - .def("gather_logits", &ModelSpec::gatherLogits, nb::rv_policy::reference_internal) - .def("replace_logits", &ModelSpec::replaceLogits, nb::rv_policy::reference_internal) - .def("return_log_probs", &ModelSpec::returnLogProbs, nb::rv_policy::reference_internal) - .def("smoke_test", &ModelSpec::smokeTest, nb::rv_policy::reference_internal) - .def("use_medusa", &ModelSpec::useMedusa, nb::rv_policy::reference_internal) - .def("use_eagle", &ModelSpec::useEagle, nb::rv_policy::reference_internal) - .def("use_lookahead_decoding", &ModelSpec::useLookaheadDecoding, nb::rv_policy::reference_internal) - .def("use_explicit_draft_tokens_decoding", &ModelSpec::useExplicitDraftTokensDecoding, - nb::rv_policy::reference_internal) - .def("use_draft_tokens_external_decoding", &ModelSpec::useDraftTokensExternalDecoding, - nb::rv_policy::reference_internal) - .def("use_logits", &ModelSpec::useLogits) - .def("use_multiple_profiles", &ModelSpec::useMultipleProfiles, nb::rv_policy::reference_internal) - .def("set_max_input_length", &ModelSpec::setMaxInputLength, nb::rv_policy::reference_internal) - .def("set_max_output_length", &ModelSpec::setMaxOutputLength, nb::rv_policy::reference_internal) - .def("set_quant_method", &ModelSpec::setQuantMethod, nb::rv_policy::reference_internal) - .def("use_lora_plugin", &ModelSpec::useLoraPlugin, nb::rv_policy::reference_internal) - .def("get_input_file", &ModelSpec::getInputFile) - .def("get_model_path", &ModelSpec::getModelPath) - .def("get_results_file", &ModelSpec::getResultsFile) - .def("get_generation_logits_file", &ModelSpec::getGenerationLogitsFile) - .def("get_context_logits_file", &ModelSpec::getContextLogitsFile) - .def("get_cum_log_probs_file", &ModelSpec::getCumLogProbsFile) - .def("get_log_probs_file", &ModelSpec::getLogProbsFile) - .def("enable_context_fmha_fp32_acc", &ModelSpec::enableContextFMHAFp32Acc, nb::rv_policy::reference_internal) - .def("get_enable_context_fmha_fp32_acc", &ModelSpec::getEnableContextFMHAFp32Acc) - .def("__copy__", [](ModelSpec const& self) { return ModelSpec(self); }); -} - -} // namespace tensorrt_llm::nanobind::testing diff --git a/cpp/tensorrt_llm/nanobind/testing/modelSpecBinding.h b/cpp/tensorrt_llm/nanobind/testing/modelSpecBinding.h deleted file mode 100644 index 1aababc6ff89..000000000000 --- a/cpp/tensorrt_llm/nanobind/testing/modelSpecBinding.h +++ /dev/null @@ -1,29 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2022-2025 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. - */ - -#pragma once - -#include - -namespace nb = nanobind; - -namespace tensorrt_llm::nanobind::testing -{ - -void initBindings(nb::module_& m); - -} // namespace tensorrt_llm::nanobind::testing diff --git a/cpp/tensorrt_llm/plugins/CMakeLists.txt b/cpp/tensorrt_llm/plugins/CMakeLists.txt deleted file mode 100755 index 8b89cccdc813..000000000000 --- a/cpp/tensorrt_llm/plugins/CMakeLists.txt +++ /dev/null @@ -1,183 +0,0 @@ -# -# SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. -# - -set(PLUGIN_TARGET_NAME nvinfer_plugin_tensorrt_llm) -set(PLUGIN_SHARED_TARGET ${PLUGIN_TARGET_NAME}) - -set(TARGET_DIR ${CMAKE_CURRENT_SOURCE_DIR}) -set(PLUGIN_EXPORT_MAP ${TARGET_DIR}/exports.map) # Linux -set(PLUGIN_EXPORT_DEF ${TARGET_DIR}/exports.def) # Windows - -if(${CMAKE_BUILD_TYPE} MATCHES "Debug") - set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -g") -endif() - -set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} --Wno-deprecated-declarations") -set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} --diag-suppress 997") - -if(NOT WIN32) - # additional warnings - # - # Ignore overloaded-virtual warning. We intentionally change parameters of - # some methods in derived class. - set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wall -Wno-overloaded-virtual") - if(WARNING_IS_ERROR) - message(STATUS "Treating warnings as errors in GCC compilation") - set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Werror") - endif() -else() # Windows - # warning level 4 - set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} /W4") -endif() - -set(PLUGIN_SOURCES) -set(PLUGIN_CU_SOURCES) - -set(PLUGIN_LISTS - bertAttentionPlugin - cpSplitPlugin - fusedLayernormPlugin - gptAttentionCommon - gptAttentionPlugin - identityPlugin - gemmPlugin - gemmSwigluPlugin - fp8RowwiseGemmPlugin - smoothQuantGemmPlugin - fp4GemmPlugin - quantizePerTokenPlugin - quantizeTensorPlugin - quantizeToFP4Plugin - layernormQuantizationPlugin - rmsnormQuantizationPlugin - weightOnlyGroupwiseQuantMatmulPlugin - weightOnlyQuantMatmulPlugin - lookupPlugin - loraPlugin - doraPlugin - mixtureOfExperts - selectiveScanPlugin - mambaConv1dPlugin - lruPlugin - cumsumLastDimPlugin - topkLastDimPlugin - lowLatencyGemmPlugin - eaglePlugin - lowLatencyGemmSwigluPlugin - qserveGemmPlugin - cudaStreamPlugin - gemmAllReducePlugin) - -foreach(PLUGIN_ITER ${PLUGIN_LISTS}) - include_directories(${PLUGIN_ITER}) - add_subdirectory(${PLUGIN_ITER}) -endforeach(PLUGIN_ITER) - -if(ENABLE_MULTI_DEVICE) - include_directories(ncclPlugin) - add_subdirectory(ncclPlugin) -endif() -include_directories(common) -add_subdirectory(common) - -# Set gencodes -list(APPEND PLUGIN_SOURCES "${PLUGIN_CU_SOURCES}") - -list(APPEND PLUGIN_SOURCES "${CMAKE_CURRENT_SOURCE_DIR}/api/tllmPlugin.cpp") - -# ################################# SHARED LIBRARY -# ############################################################################## - -if(WIN32) - set(CMAKE_WINDOWS_EXPORT_ALL_SYMBOLS 1) -endif() - -add_library(${PLUGIN_SHARED_TARGET} SHARED ${PLUGIN_SOURCES}) -add_cuda_architectures(${PLUGIN_SHARED_TARGET} 89) - -target_include_directories( - ${PLUGIN_SHARED_TARGET} - PUBLIC ${CUDA_INSTALL_DIR}/include - PUBLIC - $ - PRIVATE ${TARGET_DIR}) - -if(USING_OSS_CUTLASS_FP4_GEMM) - target_compile_definitions(${PLUGIN_SHARED_TARGET} - PUBLIC USING_OSS_CUTLASS_FP4_GEMM) -endif() - -if(USING_OSS_CUTLASS_ALLREDUCE_GEMM) - target_compile_definitions(${PLUGIN_SHARED_TARGET} - PUBLIC USING_OSS_CUTLASS_ALLREDUCE_GEMM) -endif() - -if(USING_OSS_CUTLASS_MOE_GEMM) - target_compile_definitions(${PLUGIN_SHARED_TARGET} - PUBLIC USING_OSS_CUTLASS_MOE_GEMM) -endif() - -if(ENABLE_MULTI_DEVICE) - target_include_directories(${PLUGIN_SHARED_TARGET} - PUBLIC ${MPI_C_INCLUDE_DIRS}) -endif() - -if(CUDA_VERSION VERSION_LESS 11.0) - target_include_directories(${PLUGIN_SHARED_TARGET} PUBLIC ${CUB_ROOT_DIR}) -endif() - -set_target_properties( - ${PLUGIN_SHARED_TARGET} - PROPERTIES CXX_STANDARD "17" - CXX_STANDARD_REQUIRED "YES" - CXX_EXTENSIONS "NO" - ARCHIVE_OUTPUT_DIRECTORY "${TRT_OUT_DIR}" - LIBRARY_OUTPUT_DIRECTORY "${TRT_OUT_DIR}" - RUNTIME_OUTPUT_DIRECTORY "${TRT_OUT_DIR}") - -if(WIN32) - set_target_properties( - ${PLUGIN_SHARED_TARGET} - PROPERTIES LINK_FLAGS "/DEF:${PLUGIN_EXPORT_DEF} ${UNDEFINED_FLAG}") -else() - set_target_properties( - ${PLUGIN_SHARED_TARGET} - PROPERTIES - LINK_FLAGS - "-Wl,--exclude-libs,ALL -Wl,--version-script=${PLUGIN_EXPORT_MAP} -Wl,-rpath,'$ORIGIN' ${AS_NEEDED_FLAG} ${UNDEFINED_FLAG}" - ) -endif() - -set_property(TARGET ${PLUGIN_SHARED_TARGET} PROPERTY CUDA_STANDARD 17) - -target_link_libraries( - ${PLUGIN_SHARED_TARGET} - ${CUBLAS_LIB} - ${CUBLASLT_LIB} - ${TRT_LIB} - ${CUDA_DRV_LIB} - ${CUDA_RT_LIB} - ${CMAKE_DL_LIBS} - ${SHARED_TARGET}) - -if(WIN32) - target_link_libraries(${PLUGIN_SHARED_TARGET} context_attention_src) -endif() - -if(ENABLE_MULTI_DEVICE) - target_link_libraries(${PLUGIN_SHARED_TARGET} ${MPI_C_LIBRARIES} ${NCCL_LIB}) -endif() diff --git a/cpp/tensorrt_llm/plugins/api/tllmPlugin.cpp b/cpp/tensorrt_llm/plugins/api/tllmPlugin.cpp deleted file mode 100644 index f0dceb2f4a99..000000000000 --- a/cpp/tensorrt_llm/plugins/api/tllmPlugin.cpp +++ /dev/null @@ -1,313 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. - */ -#include "tensorrt_llm/plugins/api/tllmPlugin.h" - -#include "tensorrt_llm/common/stringUtils.h" -#include "tensorrt_llm/runtime/tllmLogger.h" - -#include "tensorrt_llm/plugins/bertAttentionPlugin/bertAttentionPlugin.h" -#include "tensorrt_llm/plugins/doraPlugin/doraPlugin.h" -#include "tensorrt_llm/plugins/fp8RowwiseGemmPlugin/fp8RowwiseGemmPlugin.h" -#include "tensorrt_llm/plugins/fusedLayernormPlugin/fusedLayernormPlugin.h" -#include "tensorrt_llm/plugins/gemmPlugin/gemmPlugin.h" -#include "tensorrt_llm/plugins/gemmSwigluPlugin/gemmSwigluPlugin.h" -#include "tensorrt_llm/plugins/gptAttentionPlugin/gptAttentionPlugin.h" -#include "tensorrt_llm/plugins/identityPlugin/identityPlugin.h" -#include "tensorrt_llm/plugins/layernormQuantizationPlugin/layernormQuantizationPlugin.h" -#include "tensorrt_llm/plugins/lookupPlugin/lookupPlugin.h" -#include "tensorrt_llm/plugins/loraPlugin/loraPlugin.h" -#include "tensorrt_llm/plugins/lruPlugin/lruPlugin.h" -#include "tensorrt_llm/plugins/mambaConv1dPlugin/mambaConv1dPlugin.h" -#include "tensorrt_llm/plugins/mixtureOfExperts/mixtureOfExpertsPlugin.h" -#include "tensorrt_llm/plugins/quantizeToFP4Plugin/quantizeToFP4Plugin.h" -#if ENABLE_MULTI_DEVICE -#include "tensorrt_llm/plugins/cpSplitPlugin/cpSplitPlugin.h" -#include "tensorrt_llm/plugins/gemmAllReducePlugin/gemmAllReducePlugin.h" -#include "tensorrt_llm/plugins/ncclPlugin/allgatherPlugin.h" -#include "tensorrt_llm/plugins/ncclPlugin/allreducePlugin.h" -#include "tensorrt_llm/plugins/ncclPlugin/recvPlugin.h" -#include "tensorrt_llm/plugins/ncclPlugin/reduceScatterPlugin.h" -#include "tensorrt_llm/plugins/ncclPlugin/sendPlugin.h" -#endif // ENABLE_MULTI_DEVICE -#include "tensorrt_llm/plugins/cudaStreamPlugin/cudaStreamPlugin.h" -#include "tensorrt_llm/plugins/cumsumLastDimPlugin/cumsumLastDimPlugin.h" -#include "tensorrt_llm/plugins/eaglePlugin/eagleDecodeDraftTokensPlugin.h" -#include "tensorrt_llm/plugins/eaglePlugin/eaglePrepareDrafterInputsPlugin.h" -#include "tensorrt_llm/plugins/eaglePlugin/eagleSampleAndAcceptDraftTokensPlugin.h" -#include "tensorrt_llm/plugins/fp4GemmPlugin/fp4GemmPlugin.h" -#include "tensorrt_llm/plugins/lowLatencyGemmPlugin/lowLatencyGemmPlugin.h" -#include "tensorrt_llm/plugins/lowLatencyGemmSwigluPlugin/lowLatencyGemmSwigluPlugin.h" -#include "tensorrt_llm/plugins/qserveGemmPlugin/qserveGemmPlugin.h" -#include "tensorrt_llm/plugins/quantizePerTokenPlugin/quantizePerTokenPlugin.h" -#include "tensorrt_llm/plugins/quantizeTensorPlugin/quantizeTensorPlugin.h" -#include "tensorrt_llm/plugins/rmsnormQuantizationPlugin/rmsnormQuantizationPlugin.h" -#include "tensorrt_llm/plugins/selectiveScanPlugin/selectiveScanPlugin.h" -#include "tensorrt_llm/plugins/smoothQuantGemmPlugin/smoothQuantGemmPlugin.h" -#include "tensorrt_llm/plugins/topkLastDimPlugin/topkLastDimPlugin.h" -#include "tensorrt_llm/plugins/weightOnlyGroupwiseQuantMatmulPlugin/weightOnlyGroupwiseQuantMatmulPlugin.h" -#include "tensorrt_llm/plugins/weightOnlyQuantMatmulPlugin/weightOnlyQuantMatmulPlugin.h" - -#include -#include - -#include - -namespace tc = tensorrt_llm::common; - -namespace -{ - -nvinfer1::IPluginCreator* creatorPtr(nvinfer1::IPluginCreator& creator) -{ - return &creator; -} - -nvinfer1::IPluginCreatorInterface* creatorInterfacePtr(nvinfer1::IPluginCreatorInterface& creator) -{ - return &creator; -} - -auto tllmLogger = tensorrt_llm::runtime::TllmLogger(); - -nvinfer1::ILogger* gLogger{&tllmLogger}; - -class GlobalLoggerFinder : public nvinfer1::ILoggerFinder -{ -public: - nvinfer1::ILogger* findLogger() override - { - return gLogger; - } -}; - -GlobalLoggerFinder gGlobalLoggerFinder{}; - -#if !defined(_MSC_VER) -[[maybe_unused]] __attribute__((constructor)) -#endif -void initOnLoad() -{ - auto constexpr kLoadPlugins = "TRT_LLM_LOAD_PLUGINS"; - auto const loadPlugins = std::getenv(kLoadPlugins); - if (loadPlugins && loadPlugins[0] == '1') - { - initTrtLlmPlugins(gLogger); - } -} - -bool pluginsInitialized = false; - -} // namespace - -namespace tensorrt_llm::plugins::api -{ - -LoggerManager& tensorrt_llm::plugins::api::LoggerManager::getInstance() noexcept -{ - static LoggerManager instance; - return instance; -} - -void LoggerManager::setLoggerFinder(nvinfer1::ILoggerFinder* finder) -{ - std::lock_guard lk(mMutex); - if (mLoggerFinder == nullptr && finder != nullptr) - { - mLoggerFinder = finder; - } -} - -[[maybe_unused]] nvinfer1::ILogger* LoggerManager::logger() -{ - std::lock_guard lk(mMutex); - if (mLoggerFinder != nullptr) - { - return mLoggerFinder->findLogger(); - } - return nullptr; -} - -nvinfer1::ILogger* LoggerManager::defaultLogger() noexcept -{ - return gLogger; -} -} // namespace tensorrt_llm::plugins::api - -// New Plugin APIs - -extern "C" -{ - bool initTrtLlmPlugins(void* logger, char const* libNamespace) - { - if (pluginsInitialized) - { - return true; - } - - if (logger) - { - gLogger = static_cast(logger); - } - setLoggerFinder(&gGlobalLoggerFinder); - - auto registry = getPluginRegistry(); - - { - std::int32_t nbCreators; - auto creators = getPluginCreators(nbCreators); - - for (std::int32_t i = 0; i < nbCreators; ++i) - { - auto const creator = creators[i]; - creator->setPluginNamespace(libNamespace); - registry->registerCreator(*creator, libNamespace); - if (gLogger) - { - auto const msg = tc::fmtstr("Registered plugin creator %s version %s in namespace %s", - creator->getPluginName(), creator->getPluginVersion(), libNamespace); - gLogger->log(nvinfer1::ILogger::Severity::kVERBOSE, msg.c_str()); - } - } - } - - { - std::int32_t nbCreators; - auto creators = getCreators(nbCreators); - - for (std::int32_t i = 0; i < nbCreators; ++i) - { - auto const creator = creators[i]; - registry->registerCreator(*creator, libNamespace); - } - } - - pluginsInitialized = true; - return true; - } - - [[maybe_unused]] void setLoggerFinder([[maybe_unused]] nvinfer1::ILoggerFinder* finder) - { - tensorrt_llm::plugins::api::LoggerManager::getInstance().setLoggerFinder(finder); - } - - [[maybe_unused]] nvinfer1::IPluginCreator* const* getPluginCreators(std::int32_t& nbCreators) - { - static tensorrt_llm::plugins::IdentityPluginCreator identityPluginCreator; - static tensorrt_llm::plugins::BertAttentionPluginCreator bertAttentionPluginCreator; - static tensorrt_llm::plugins::FusedLayernormPluginCreator fusedLayernormPluginCreator; - static tensorrt_llm::plugins::GPTAttentionPluginCreator gptAttentionPluginCreator; - static tensorrt_llm::plugins::GemmPluginCreator gemmPluginCreator; - static tensorrt_llm::plugins::GemmSwigluPluginCreator gemmSwigluPluginCreator; - static tensorrt_llm::plugins::Fp8RowwiseGemmPluginCreator fp8RowwiseGemmPluginCreator; - static tensorrt_llm::plugins::MixtureOfExpertsPluginCreator moePluginCreator; -#if ENABLE_MULTI_DEVICE - static tensorrt_llm::plugins::SendPluginCreator sendPluginCreator; - static tensorrt_llm::plugins::RecvPluginCreator recvPluginCreator; - static tensorrt_llm::plugins::AllreducePluginCreator allreducePluginCreator; - static tensorrt_llm::plugins::AllgatherPluginCreator allgatherPluginCreator; - static tensorrt_llm::plugins::ReduceScatterPluginCreator reduceScatterPluginCreator; - static tensorrt_llm::plugins::GemmAllReducePluginCreator gemmAllReducePluginCreator; -#endif // ENABLE_MULTI_DEVICE - static tensorrt_llm::plugins::SmoothQuantGemmPluginCreator smoothQuantGemmPluginCreator; - static tensorrt_llm::plugins::QServeGemmPluginCreator qserveGemmPluginCreator; - static tensorrt_llm::plugins::LayernormQuantizationPluginCreator layernormQuantizationPluginCreator; - static tensorrt_llm::plugins::QuantizeToFP4PluginCreator quantizeToFP4PluginCreator; - static tensorrt_llm::plugins::QuantizePerTokenPluginCreator quantizePerTokenPluginCreator; - static tensorrt_llm::plugins::QuantizeTensorPluginCreator quantizeTensorPluginCreator; - static tensorrt_llm::plugins::RmsnormQuantizationPluginCreator rmsnormQuantizationPluginCreator; - static tensorrt_llm::plugins::WeightOnlyGroupwiseQuantMatmulPluginCreator - weightOnlyGroupwiseQuantMatmulPluginCreator; - static tensorrt_llm::plugins::WeightOnlyQuantMatmulPluginCreator weightOnlyQuantMatmulPluginCreator; - static tensorrt_llm::plugins::LookupPluginCreator lookupPluginCreator; - static tensorrt_llm::plugins::LoraPluginCreator loraPluginCreator; - static tensorrt_llm::plugins::SelectiveScanPluginCreator selectiveScanPluginCreator; - static tensorrt_llm::plugins::Fp4GemmPluginCreator fp4GemmPluginCreator; - static tensorrt_llm::plugins::MambaConv1dPluginCreator mambaConv1DPluginCreator; - static tensorrt_llm::plugins::lruPluginCreator lruPluginCreator; - static tensorrt_llm::plugins::CumsumLastDimPluginCreator cumsumLastDimPluginCreator; - static tensorrt_llm::plugins::TopkLastDimPluginCreator topkLastDimPluginCreator; - static tensorrt_llm::plugins::LowLatencyGemmPluginCreator lowLatencyGemmPluginCreator; - static tensorrt_llm::plugins::LowLatencyGemmSwigluPluginCreator lowLatencyGemmSwigluPluginCreator; - static tensorrt_llm::plugins::EagleDecodeDraftTokensPluginCreator eagleDecodeDraftTokensPluginCreator; - static tensorrt_llm::plugins::EagleSampleAndAcceptDraftTokensPluginCreator - eagleSampleAndAcceptDraftTokensPluginCreator; - static tensorrt_llm::plugins::CudaStreamPluginCreator cudaStreamPluginCreator; - - static std::array pluginCreators - = { creatorPtr(identityPluginCreator), - creatorPtr(bertAttentionPluginCreator), - creatorPtr(gptAttentionPluginCreator), - creatorPtr(gemmPluginCreator), - creatorPtr(gemmSwigluPluginCreator), - creatorPtr(fp8RowwiseGemmPluginCreator), - creatorPtr(moePluginCreator), -#if ENABLE_MULTI_DEVICE - creatorPtr(sendPluginCreator), - creatorPtr(recvPluginCreator), - creatorPtr(allreducePluginCreator), - creatorPtr(allgatherPluginCreator), - creatorPtr(reduceScatterPluginCreator), - creatorPtr(gemmAllReducePluginCreator), -#endif // ENABLE_MULTI_DEVICE - creatorPtr(fusedLayernormPluginCreator), - creatorPtr(smoothQuantGemmPluginCreator), - creatorPtr(qserveGemmPluginCreator), - creatorPtr(layernormQuantizationPluginCreator), - creatorPtr(quantizeToFP4PluginCreator), - creatorPtr(quantizePerTokenPluginCreator), - creatorPtr(quantizeTensorPluginCreator), - creatorPtr(rmsnormQuantizationPluginCreator), - creatorPtr(weightOnlyGroupwiseQuantMatmulPluginCreator), - creatorPtr(weightOnlyQuantMatmulPluginCreator), - creatorPtr(lookupPluginCreator), - creatorPtr(loraPluginCreator), - creatorPtr(selectiveScanPluginCreator), - creatorPtr(fp4GemmPluginCreator), - creatorPtr(mambaConv1DPluginCreator), - creatorPtr(lruPluginCreator), - creatorPtr(cumsumLastDimPluginCreator), - creatorPtr(topkLastDimPluginCreator), - creatorPtr(lowLatencyGemmPluginCreator), - creatorPtr(eagleDecodeDraftTokensPluginCreator), - creatorPtr(eagleSampleAndAcceptDraftTokensPluginCreator), - creatorPtr(lowLatencyGemmSwigluPluginCreator), - creatorPtr(cudaStreamPluginCreator), - }; - nbCreators = pluginCreators.size(); - return pluginCreators.data(); - } - - [[maybe_unused]] nvinfer1::IPluginCreatorInterface* const* getCreators(std::int32_t& nbCreators) - { - static tensorrt_llm::plugins::EaglePrepareDrafterInputsPluginCreator eaglePrepareDrafterInputsPluginCreator; -#if ENABLE_MULTI_DEVICE - static tensorrt_llm::plugins::CpSplitPluginCreator cpSplitPluginCreator; -#endif // ENABLE_MULTI_DEVICE - - static tensorrt_llm::plugins::DoraPluginCreator doraPluginCreator; - - static std::array creators - = { creatorInterfacePtr(eaglePrepareDrafterInputsPluginCreator), -#if ENABLE_MULTI_DEVICE - creatorInterfacePtr(cpSplitPluginCreator), -#endif // ENABLE_MULTI_DEVICE - creatorInterfacePtr(doraPluginCreator) }; - - nbCreators = creators.size(); - return creators.data(); - } -} // extern "C" diff --git a/cpp/tensorrt_llm/plugins/bertAttentionPlugin/CMakeLists.txt b/cpp/tensorrt_llm/plugins/bertAttentionPlugin/CMakeLists.txt deleted file mode 100644 index 86876224fccd..000000000000 --- a/cpp/tensorrt_llm/plugins/bertAttentionPlugin/CMakeLists.txt +++ /dev/null @@ -1,21 +0,0 @@ -# -# SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. -# -file(GLOB SRCS *.cpp) -set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS}) -set(PLUGIN_SOURCES - ${PLUGIN_SOURCES} - PARENT_SCOPE) diff --git a/cpp/tensorrt_llm/plugins/bertAttentionPlugin/bertAttentionPlugin.cpp b/cpp/tensorrt_llm/plugins/bertAttentionPlugin/bertAttentionPlugin.cpp deleted file mode 100644 index 6acf0b3a9d25..000000000000 --- a/cpp/tensorrt_llm/plugins/bertAttentionPlugin/bertAttentionPlugin.cpp +++ /dev/null @@ -1,1206 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2025 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. - */ -#include "bertAttentionPlugin.h" -#include "tensorrt_llm/kernels/gptKernels.h" -#include "tensorrt_llm/kernels/recoverFromRingAtten.h" -#include "tensorrt_llm/kernels/sageAttentionKernels.h" -#include "tensorrt_llm/kernels/unfusedAttentionKernels.h" -#include "tensorrt_llm/runtime/iBuffer.h" - -using namespace nvinfer1; -using namespace tensorrt_llm::kernels; -namespace tc = tensorrt_llm::common; - -using tensorrt_llm::plugins::BertAttentionPluginCreator; -using tensorrt_llm::plugins::BertAttentionPlugin; - -static char const* BERT_ATTENTION_PLUGIN_VERSION{"1"}; -static char const* BERT_ATTENTION_PLUGIN_NAME{"BertAttention"}; -PluginFieldCollection BertAttentionPluginCreator::mFC{}; -std::vector BertAttentionPluginCreator::mPluginAttributes; - -BertAttentionPlugin::BertAttentionPlugin(int num_heads, int head_size, float q_scaling, - ContextFMHAType context_fmha_type, nvinfer1::DataType type, bool do_relative_attention, int max_distance, - bool remove_padding, bool sage_attn, int sage_attn_q_block_size, int sage_attn_k_block_size, - int sage_attn_v_block_size, int cp_size, int cp_rank, std::set cp_group) - : mNumHeads(num_heads) - , mHeadSize(head_size) - , mQScaling(q_scaling) - , mType(type) - , mRelativeAttention(do_relative_attention) - , mMaxDistance(max_distance) - , mRemovePadding(remove_padding) - , mEnableContextFMHA(context_fmha_type != ContextFMHAType::DISABLED) - , mFMHAForceFP32Acc(context_fmha_type == ContextFMHAType::ENABLED_WITH_FP32_ACC) - , mSageAttn(sage_attn) - , mCpSize(cp_size) - , mCpRank(cp_rank) - , mCpGroup(std::move(cp_group)) -{ - // pre-check whether FMHA is supported in order to save memory allocation - if (mEnableContextFMHA) - { - mEnableContextFMHA = false; - if (!(mType == DataType::kHALF || mType == DataType::kBF16)) - { - TLLM_LOG_WARNING("Fall back to unfused MHA because of unsupported data type."); - } - else if (mRelativeAttention) - { - TLLM_LOG_WARNING("Fall back to unfused MHA because of relative position embedding."); - } - else - { - mEnableContextFMHA = true; - } - } - - if (mSageAttn) - { - mSageAttnQBlockSize = sage_attn_q_block_size; - mSageAttnKBlockSize = sage_attn_k_block_size; - mSageAttnVBlockSize = sage_attn_v_block_size; - std::vector blockSizeCombination - = {sage_attn_q_block_size, sage_attn_k_block_size, sage_attn_v_block_size}; - if (mSageAttnSupportedBlockSizes.find(blockSizeCombination) == mSageAttnSupportedBlockSizes.end() - || (head_size != 128 && head_size != 72 && head_size != 80)) - { - TLLM_LOG_WARNING(" Q, k ,v quant block size not support. disable sage attention"); - mSageAttn = false; - } - else - { - TLLM_LOG_INFO("SageAttnQBlockSize: %d, SageAttnKBlockSize: %d, SageAttnVBlockSize: %d", mSageAttnQBlockSize, - mSageAttnKBlockSize, mSageAttnVBlockSize); - } - } - - if (cp_group.size() > 1 && !mEnableContextFMHA) - { - TLLM_LOG_ERROR("Unfused MHA do not support context parallel now."); - } -} - -// Parameterized constructor -BertAttentionPlugin::BertAttentionPlugin(void const* data, size_t length) -{ - char const *d = reinterpret_cast(data), *a = d; - read(d, mNumHeads); - read(d, mHeadSize); - read(d, mQScaling); - read(d, mQKHalfAccum); - read(d, mEnableContextFMHA); - read(d, mFMHAForceFP32Acc); - read(d, mType); - read(d, mRelativeAttention); - read(d, mMaxDistance); - read(d, mRemovePadding); - read(d, mSageAttn); - read(d, mSageAttnQBlockSize); - read(d, mSageAttnKBlockSize); - read(d, mSageAttnVBlockSize); - read(d, mCpSize); - read(d, mCpRank); - mCpGroup.clear(); - int groupItem = 0; - while (d != a + length) - { - read(d, groupItem); - mCpGroup.insert(groupItem); - } - - TLLM_CHECK_WITH_INFO(d == a + length, - "Expected length (%d) != real length (%d). This is often " - "caused by using different TensorRT LLM version to build " - "engine and run engine.", - (int) length, (int) (d - a)); -} - -// IPluginV2DynamicExt Methods -nvinfer1::IPluginV2DynamicExt* BertAttentionPlugin::clone() const noexcept -{ - auto* plugin = new BertAttentionPlugin(*this); - plugin->setPluginNamespace(mNamespace.c_str()); - plugin->initialize(); - return plugin; -} - -nvinfer1::DimsExprs BertAttentionPlugin::getOutputDimensions( - int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept -{ - TLLM_CHECK(outputIndex == 0); - auto ret = inputs[0]; - ret.d[mRemovePadding ? 1 : 2] = exprBuilder.constant(ret.d[mRemovePadding ? 1 : 2]->getConstantValue() / 3); - return ret; -} - -bool BertAttentionPlugin::supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept -{ - // inputs: [0] qkv, [1] input_lengths, [2] max_input_length (optional), [3] relative_attention_bias (optional) - // outputs: [X] hidden_states - if (nbInputs == 2) - { // BERT - if (pos == 1) - { - return inOut[pos].type == nvinfer1::DataType::kINT32; - } - - return (inOut[pos].type == mType) && (inOut[pos].format == TensorFormat::kLINEAR); - } - if (nbInputs > 2) - { // Encoder in encoder-decoder - if (pos == 1 || pos == 2) - { - return inOut[pos].type == nvinfer1::DataType::kINT32; - } - - return (inOut[pos].type == mType) && (inOut[pos].format == TensorFormat::kLINEAR); - } - - return false; -} - -void BertAttentionPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept -{ -} - -size_t BertAttentionPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - // if remove padding, inputs[0] "qkv_hidden_states" dim is [num_tokens, 3*hidden_dim] which doesn't have shape - // info should get max_batch_size and max_input_length from inputs[1] "input_lengths" and input[2] - // "max_input_length" - int const batch_size = mRemovePadding ? inputs[1].dims.d[0] : inputs[0].dims.d[0]; - int const input_seq_len = mRemovePadding ? inputs[2].dims.d[0] : inputs[0].dims.d[1]; - int const local_hidden_units_ = inputs[0].dims.d[mRemovePadding ? 1 : 2] / 3; - - auto const size = tensorrt_llm::runtime::BufferDataType(inputs[0].type).getSize(); - - size_t const attention_mask_size = mEnableContextFMHA ? 0 : size * batch_size * input_seq_len * input_seq_len; - size_t const cu_seqlens_size = sizeof(int) * (batch_size + 1); - size_t const q_buf_2_size = mEnableContextFMHA ? 0 : size * batch_size * input_seq_len * local_hidden_units_; - size_t const k_buf_2_size = mEnableContextFMHA ? 0 : size * batch_size * input_seq_len * local_hidden_units_; - size_t const v_buf_2_size = mEnableContextFMHA ? 0 : size * batch_size * input_seq_len * local_hidden_units_; - size_t const qk_buf_size = mEnableContextFMHA ? 0 : size * batch_size * mNumHeads * input_seq_len * input_seq_len; - size_t const qkv_buf_2_size = mEnableContextFMHA ? 0 : size * batch_size * input_seq_len * local_hidden_units_; - size_t const qk_buf_float_size - = mEnableContextFMHA ? 0 : sizeof(float) * batch_size * mNumHeads * input_seq_len * input_seq_len; - size_t const padding_offset_size = mEnableContextFMHA ? 0 : sizeof(int) * batch_size * input_seq_len; - size_t const fmha_scheduler_counter = mEnableContextFMHA ? sizeof(uint32_t) : 0; - int const paddedHeadSize = mSageAttn ? ((mHeadSize + 15) / 16) * 16 : mHeadSize; - const size_t quanted_qkv_size - = mSageAttn ? sizeof(__nv_fp8_e4m3) * batch_size * input_seq_len * mNumHeads * paddedHeadSize * 3 : 0; - const size_t q_scale_size = mSageAttn - ? sizeof(float) * batch_size * ((input_seq_len + mSageAttnQBlockSize - 1) / mSageAttnQBlockSize) * mNumHeads - : 0; - const size_t k_scale_size = mSageAttn - ? sizeof(float) * batch_size * ((input_seq_len + mSageAttnKBlockSize - 1) / mSageAttnKBlockSize) * mNumHeads - : 0; - const size_t v_scale_size = mSageAttn - ? sizeof(float) * batch_size * ((input_seq_len + mSageAttnVBlockSize - 1) / mSageAttnVBlockSize) * mNumHeads - : 0; - const size_t scale_bmm1_device_size = mSageAttn ? sizeof(float) * 2 : 0; - const size_t scale_bmm2_device_size = mSageAttn ? sizeof(float) : 0; - size_t sage_quant_space_size = mSageAttn ? sizeof(float) * batch_size * mNumHeads * mHeadSize : 0; - - if (paddedHeadSize != mHeadSize) - sage_quant_space_size - = sage_quant_space_size < (batch_size * input_seq_len * mNumHeads * paddedHeadSize * sizeof(__nv_bfloat16)) - ? (batch_size * input_seq_len * mNumHeads * paddedHeadSize * sizeof(__nv_bfloat16)) - : sage_quant_space_size; - - // workspace for RingAttention ping-pong buffer - bool const enableRingAttn = (mCpGroup.size() > 1); - const size_t ring_q_buf_size = enableRingAttn ? size * batch_size * input_seq_len * local_hidden_units_ : 0; - const size_t ring_kv_buf_size = enableRingAttn - ? 2 * size * batch_size * input_seq_len * local_hidden_units_ + sizeof(int) * (batch_size + 1) - : 0; - const size_t ring_softmax_stats_buf_size - = enableRingAttn ? 2 * sizeof(float) * batch_size * input_seq_len * mNumHeads : 0; - const size_t ring_softmax_stats_accu_buf_size - = enableRingAttn ? 2 * sizeof(float) * batch_size * input_seq_len * mNumHeads : 0; - const size_t ring_block_output_size = enableRingAttn ? size * batch_size * input_seq_len * local_hidden_units_ : 0; - - int const NUM_BUFFERS = 24; - - size_t workspaces[NUM_BUFFERS]; - workspaces[0] = CUBLAS_WORKSPACE_SIZE; - workspaces[1] = attention_mask_size; - workspaces[2] = cu_seqlens_size; - workspaces[3] = q_buf_2_size; - workspaces[4] = k_buf_2_size; - workspaces[5] = v_buf_2_size; - workspaces[6] = qk_buf_size; - workspaces[7] = qkv_buf_2_size; - workspaces[8] = qk_buf_float_size; - workspaces[9] = padding_offset_size; - workspaces[10] = fmha_scheduler_counter; - workspaces[11] = quanted_qkv_size; - workspaces[12] = q_scale_size; - workspaces[13] = v_scale_size; - workspaces[14] = k_scale_size; - workspaces[15] = scale_bmm1_device_size; - workspaces[16] = scale_bmm2_device_size; - workspaces[17] = sage_quant_space_size; - workspaces[18] = ring_q_buf_size; - workspaces[19] = ring_kv_buf_size; // kv1 - workspaces[20] = ring_kv_buf_size; // kv2 - workspaces[21] = ring_softmax_stats_buf_size; - workspaces[22] = ring_softmax_stats_accu_buf_size; - workspaces[23] = ring_block_output_size; - - return tc::calculateTotalWorkspaceSize(workspaces, NUM_BUFFERS); -} - -template -int BertAttentionPlugin::enqueueImpl(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) -{ - - // inputs - // input_tensor [batch_size, seq_len, local_hidden_size*3] or [num_tokens, local_hidden_size*3] - // input_lengths [batch_size] - // max_input_length [max_input_length] -- use shape dim to represent max value. If remove padding, this records - // the max input length among sequences; otherwise same as input_tensor's padded dim[1] relative_attention_bias - // [num_heads, num_buckets] (optional) - // outputs - // output_tensor [batch_size, seq_len, local_hidden_size] or [num_tokens, local_hidden_size] - - // if remove padding, inputs[0] dim is [num_tokens] which doesn't have workspace info - // should get max_batch_size from inputs[1] and max_input_length from plugin attribute - int const batch_size = mRemovePadding ? inputDesc[1].dims.d[0] : inputDesc[0].dims.d[0]; - int const input_seq_len = mRemovePadding ? inputDesc[2].dims.d[0] : inputDesc[0].dims.d[1]; - int const num_tokens = mRemovePadding ? inputDesc[0].dims.d[0] : batch_size * input_seq_len; - int const request_batch_size = batch_size; - int const request_seq_len = input_seq_len; - int const local_hidden_units_ = inputDesc[0].dims.d[mRemovePadding ? 1 : 2] / 3; - float const q_scaling = mQScaling; - - T const* attention_input = reinterpret_cast(inputs[0]); - int const* input_lengths = reinterpret_cast(inputs[1]); - T const* relative_attn_table = mRelativeAttention ? reinterpret_cast(inputs[3]) : nullptr; - T* context_buf_ = (T*) (outputs[0]); - - auto cublasHandle = mCublasWrapper->getCublasHandle(); - TLLM_CUDA_CHECK(cublasSetStream(cublasHandle, stream)); - mCublasWrapper->setStream(stream); - mCublasWrapper->setWorkspace(workspace); - if (inputDesc[0].type == DataType::kHALF) - { - mCublasWrapper->setFP16GemmConfig(); - } - else if (inputDesc[0].type == DataType::kFLOAT) - { - mCublasWrapper->setFP32GemmConfig(); - } -#ifdef ENABLE_BF16 - else if constexpr (std::is_same_v) - { - mCublasWrapper->setBF16GemmConfig(); - } -#endif - - size_t const attention_mask_size = mEnableContextFMHA ? 0 : sizeof(T) * batch_size * input_seq_len * input_seq_len; - size_t const cu_seqlens_size = sizeof(int) * (batch_size + 1); - size_t const q_buf_2_size = mEnableContextFMHA ? 0 : sizeof(T) * batch_size * input_seq_len * local_hidden_units_; - size_t const k_buf_2_size = mEnableContextFMHA ? 0 : sizeof(T) * batch_size * input_seq_len * local_hidden_units_; - size_t const v_buf_2_size = mEnableContextFMHA ? 0 : sizeof(T) * batch_size * input_seq_len * local_hidden_units_; - size_t const qk_buf_size - = mEnableContextFMHA ? 0 : sizeof(T) * batch_size * mNumHeads * input_seq_len * input_seq_len; - size_t const qkv_buf_2_size = mEnableContextFMHA ? 0 : sizeof(T) * batch_size * input_seq_len * local_hidden_units_; - size_t const qk_buf_float_size - = mEnableContextFMHA ? 0 : sizeof(float) * batch_size * mNumHeads * input_seq_len * input_seq_len; - size_t const padding_offset_size = mEnableContextFMHA ? 0 : sizeof(int) * batch_size * input_seq_len; - size_t const fmha_scheduler_counter = mEnableContextFMHA ? sizeof(uint32_t) : 0; - - int const paddedHeadSize = mSageAttn ? ((mHeadSize + 15) / 16) * 16 : mHeadSize; - const size_t quanted_qkv_size - = mSageAttn ? sizeof(__nv_fp8_e4m3) * batch_size * input_seq_len * mNumHeads * paddedHeadSize * 3 : 0; - const size_t q_scale_size = mSageAttn - ? sizeof(float) * batch_size * ((input_seq_len + mSageAttnQBlockSize - 1) / mSageAttnQBlockSize) * mNumHeads - : 0; - const size_t k_scale_size = mSageAttn - ? sizeof(float) * batch_size * ((input_seq_len + mSageAttnKBlockSize - 1) / mSageAttnKBlockSize) * mNumHeads - : 0; - const size_t v_scale_size = mSageAttn - ? sizeof(float) * batch_size * ((input_seq_len + mSageAttnVBlockSize - 1) / mSageAttnVBlockSize) * mNumHeads - : 0; - const size_t scale_bmm1_device_size = mSageAttn ? sizeof(float) * 2 : 0; - const size_t scale_bmm2_device_size = mSageAttn ? sizeof(float) : 0; - size_t sage_quant_space_size = mSageAttn ? sizeof(float) * batch_size * mNumHeads * mHeadSize : 0; - - if (paddedHeadSize != mHeadSize) - sage_quant_space_size - = sage_quant_space_size < (batch_size * input_seq_len * mNumHeads * paddedHeadSize * sizeof(__nv_bfloat16)) - ? (batch_size * input_seq_len * mNumHeads * paddedHeadSize * sizeof(__nv_bfloat16)) - : sage_quant_space_size; - - bool const enableRingAttn = (mCpGroup.size() > 1); - const size_t ring_q_buf_size = enableRingAttn ? sizeof(T) * batch_size * input_seq_len * local_hidden_units_ : 0; - const size_t ring_kv_buf_size - = enableRingAttn ? 2 * sizeof(T) * batch_size * input_seq_len * local_hidden_units_ : 0; - const size_t ring_softmax_stats_buf_size - = enableRingAttn ? 2 * sizeof(float) * batch_size * input_seq_len * mNumHeads : 0; - const size_t ring_block_output_size - = enableRingAttn ? sizeof(T) * batch_size * input_seq_len * local_hidden_units_ : 0; - - // Workspace pointer shift - int8_t* workspace_byte_ptr = reinterpret_cast(workspace); - size_t offset = CUBLAS_WORKSPACE_SIZE; - - T* attention_mask = reinterpret_cast(tc::nextWorkspacePtr(workspace_byte_ptr, offset, attention_mask_size)); - int* cu_seqlens = reinterpret_cast(tc::nextWorkspacePtr(workspace_byte_ptr, offset, cu_seqlens_size)); - T* q_buf_2_ = reinterpret_cast(tc::nextWorkspacePtr(workspace_byte_ptr, offset, q_buf_2_size)); - T* k_buf_2_ = reinterpret_cast(tc::nextWorkspacePtr(workspace_byte_ptr, offset, k_buf_2_size)); - T* v_buf_2_ = reinterpret_cast(tc::nextWorkspacePtr(workspace_byte_ptr, offset, v_buf_2_size)); - T* qk_buf_ = reinterpret_cast(tc::nextWorkspacePtr(workspace_byte_ptr, offset, qk_buf_size)); - T* qkv_buf_2_ = reinterpret_cast(tc::nextWorkspacePtr(workspace_byte_ptr, offset, qkv_buf_2_size)); - float* qk_buf_float_ - = reinterpret_cast(tc::nextWorkspacePtr(workspace_byte_ptr, offset, qk_buf_float_size)); - int* padding_offset = reinterpret_cast(tc::nextWorkspacePtr(workspace_byte_ptr, offset, padding_offset_size)); - uint32_t* fmha_tile_counter_ptr - = reinterpret_cast(tc::nextWorkspacePtr(workspace_byte_ptr, offset, fmha_scheduler_counter)); - - __nv_fp8_e4m3* quanted_qkv_ptr - = reinterpret_cast<__nv_fp8_e4m3*>(tc::nextWorkspacePtr(workspace_byte_ptr, offset, quanted_qkv_size)); - float* q_scale_ptr = reinterpret_cast(tc::nextWorkspacePtr(workspace_byte_ptr, offset, q_scale_size)); - float* k_scale_ptr = reinterpret_cast(tc::nextWorkspacePtr(workspace_byte_ptr, offset, k_scale_size)); - float* v_scale_ptr = reinterpret_cast(tc::nextWorkspacePtr(workspace_byte_ptr, offset, v_scale_size)); - float* scale_bmm1_ptr - = reinterpret_cast(tc::nextWorkspacePtr(workspace_byte_ptr, offset, scale_bmm1_device_size)); - float* scale_bmm2_ptr - = reinterpret_cast(tc::nextWorkspacePtr(workspace_byte_ptr, offset, scale_bmm2_device_size)); - void* sage_quant_space_ptr - = reinterpret_cast(tc::nextWorkspacePtr(workspace_byte_ptr, offset, sage_quant_space_size)); - - T* ring_q_buf_ = reinterpret_cast(tc::nextWorkspacePtr(workspace_byte_ptr, offset, ring_q_buf_size)); - T* ring_kv_buf_1_ = reinterpret_cast( - tc::nextWorkspacePtr(workspace_byte_ptr, offset, ring_kv_buf_size + sizeof(int) * (batch_size + 1))); - T* ring_kv_buf_2_ = reinterpret_cast( - tc::nextWorkspacePtr(workspace_byte_ptr, offset, ring_kv_buf_size + sizeof(int) * (batch_size + 1))); - float* ring_softmax_stats_buf_ - = reinterpret_cast(tc::nextWorkspacePtr(workspace_byte_ptr, offset, ring_softmax_stats_buf_size)); - float* ring_softmax_accu_stats_buf_ - = reinterpret_cast(tc::nextWorkspacePtr(workspace_byte_ptr, offset, ring_softmax_stats_buf_size)); - T* ring_block_output_ - = reinterpret_cast(tc::nextWorkspacePtr(workspace_byte_ptr, offset, ring_block_output_size)); - - // build attention_mask, cu_seqlens, and padding_offset tensors - BuildDecoderInfoParams params{}; - params.seqQOffsets = cu_seqlens; - params.paddingOffsets = padding_offset; - params.attentionMask = attention_mask; - params.seqQLengths = input_lengths; - params.batchSize = batch_size; - params.maxQSeqLength = input_seq_len; - params.numTokens = num_tokens; - params.attentionMaskType = AttentionMaskType::PADDING; - params.fmhaTileCounter = fmha_tile_counter_ptr; - if (mSageAttn) - { - params.fmhaHostBmm1Scale = 1.0f / (sqrtf(mHeadSize * 1.0f) * q_scaling); - params.fmhaBmm1Scale = scale_bmm1_ptr; - params.fmhaBmm2Scale = scale_bmm2_ptr; - } - invokeBuildDecoderInfo(params, stream); - sync_check_cuda_error(stream); - - auto const gemm_data_type = tc::CudaDataType::value; - int const attention_seq_len_1 = request_seq_len; // q length - int const attention_seq_len_2 = request_seq_len; // kv length - - // If the model has relative attentiona bias, q scaling should be applied in QK gemm stage and use 1 in - // softamax stage (because to get softmax[scale(Q*K) + rel pos bias] here, q_scaling can't be applied during - // softmax phase by qk_scale); otherwise, use 1 in gemm stage and apply scaling in softmax stage - float const qk_scale - = 1.0f / (sqrtf(mHeadSize * 1.0f) * q_scaling); // q_scaling in denominator. by default q_scaling =1.0f - float const qk_scale_gemm = mRelativeAttention ? qk_scale : 1.0f; - T const qk_scale_softmax = static_cast(mRelativeAttention ? 1.0f : qk_scale); - - T* linear_bias_slopes = nullptr; - - // FMHA doesn't apply to MHA with relative attention bias, i.e. softmax(QK + bias) * V - // We update mEnableContextFMHA in constructor to check this condition - if (mEnableContextFMHA) - { - if (enableRingAttn) - { - // make sure the padding part of key/value buffer is 0 - cudaMemsetAsync(ring_kv_buf_1_, 0, - reinterpret_cast(ring_kv_buf_2_) - reinterpret_cast(ring_kv_buf_1_), stream); - - cudaMemcpyAsync(ring_q_buf_, attention_input, ring_q_buf_size, cudaMemcpyDeviceToDevice, stream); - cudaMemcpyAsync(ring_kv_buf_1_, - const_cast(reinterpret_cast(attention_input)) + ring_q_buf_size, ring_kv_buf_size, - cudaMemcpyDeviceToDevice, stream); - cudaMemcpyAsync(reinterpret_cast(ring_kv_buf_1_) + ring_kv_buf_size, cu_seqlens, - sizeof(int) * (batch_size + 1), cudaMemcpyDeviceToDevice, stream); - // init softmax_stats - cudaMemsetAsync(ring_softmax_accu_stats_buf_, 0, ring_softmax_stats_buf_size, stream); - -#if ENABLE_MULTI_DEVICE - // relative position of prev/next rank in cp group - int prev_rank = mCpRank > 0 ? mCpRank - 1 : mCpGroup.size() - 1; - int next_rank = (mCpRank == static_cast(mCpGroup.size() - 1)) ? 0 : mCpRank + 1; -#endif // ENABLE_MULTI_DEVICE - - common::check_cuda_error(cudaStreamCreate(&mNcclStream)); - common::check_cuda_error(cudaStreamSynchronize(stream)); - - uint32_t* fmha_scheduler_counter_h = (uint32_t*) malloc(sizeof(uint32_t)); - cudaMemcpyAsync( - fmha_scheduler_counter_h, fmha_tile_counter_ptr, sizeof(uint32_t), cudaMemcpyDeviceToHost, stream); - for (size_t iter = 0; iter < mCpGroup.size(); ++iter) - { - // KV buffer used by fmha - T* ring_fmha_kv_buf_ = (iter % 2 == 0) ? ring_kv_buf_1_ : ring_kv_buf_2_; -#if ENABLE_MULTI_DEVICE - T* ring_send_kv_buf_ = (iter % 2 == 0) ? ring_kv_buf_1_ : ring_kv_buf_2_; - T* ring_recv_kv_buf_ = (iter % 2 == 0) ? ring_kv_buf_2_ : ring_kv_buf_1_; - if (iter < mCpGroup.size() - 1) - { - NCCLCHECK(ncclGroupStart()); - TLLM_CHECK_WITH_INFO(mNcclComm.get() != nullptr, "mNcclComm should be initialized before used"); - NCCLCHECK(ncclSend(ring_send_kv_buf_, - ring_kv_buf_size / sizeof(T) + sizeof(int) / sizeof(T) * (batch_size + 1), - (*getDtypeMap())[inputDesc[0].type], next_rank, *mNcclComm, mNcclStream)); - NCCLCHECK(ncclRecv(ring_recv_kv_buf_, - ring_kv_buf_size / sizeof(T) + sizeof(int) / sizeof(T) * (batch_size + 1), - (*getDtypeMap())[inputDesc[0].type], prev_rank, *mNcclComm, mNcclStream)); - NCCLCHECK(ncclGroupEnd()); - } -#else - TLLM_LOG_ERROR("Please set ENABLE_MULTI_DEVICE to enable RingAttention"); - return 1; -#endif // ENABLE_MULTI_DEVICE - // Construct the fmha params for running kernels. - MHARunnerParams fmhaParams{}; - fmhaParams.b = request_batch_size; - fmhaParams.qSeqLen = request_seq_len; - fmhaParams.kvSeqLen = request_seq_len; - fmhaParams.totalQSeqLen = request_batch_size * request_seq_len; - // Device buffer pointers. - fmhaParams.qPtr = ring_q_buf_; - fmhaParams.kvPtr = ring_fmha_kv_buf_; - if (iter == 0) - { - fmhaParams.outputPtr = context_buf_; - fmhaParams.softmaxStatsPtr = ring_softmax_accu_stats_buf_; - } - else - { - cudaMemsetAsync(ring_softmax_stats_buf_, 0, ring_softmax_stats_buf_size, stream); - fmhaParams.outputPtr = ring_block_output_; - fmhaParams.softmaxStatsPtr = ring_softmax_stats_buf_; - } - fmhaParams.cuQSeqLenPtr = cu_seqlens; - fmhaParams.cuKvSeqLenPtr - = reinterpret_cast(reinterpret_cast(ring_fmha_kv_buf_) + ring_kv_buf_size); - - fmhaParams.tileCounterPtr = fmha_tile_counter_ptr; - fmhaParams.stream = stream; - // Run the fmha kernel. - cudaMemsetAsync(fmhaParams.outputPtr, 0, ring_block_output_size, stream); - cudaMemcpyAsync(fmhaParams.tileCounterPtr, fmha_scheduler_counter_h, sizeof(uint32_t), - cudaMemcpyHostToDevice, stream); - mFmhaDispatcher->run(fmhaParams); - if (iter != 0) - { - invokeRecoverFromRA((T*) context_buf_, (float*) ring_softmax_accu_stats_buf_, - (T*) ring_block_output_, (float*) ring_softmax_stats_buf_, fmhaParams.b, fmhaParams.qSeqLen, - mNumHeads, mHeadSize, cu_seqlens, stream); - } - cudaStreamSynchronize(stream); - cudaStreamSynchronize(mNcclStream); - } - common::check_cuda_error(cudaStreamDestroy(mNcclStream)); - free(fmha_scheduler_counter_h); - } - - else - { - if (mSageAttn && mHeadSize == 72 && mSageAttnQBlockSize == 64 && mSageAttnKBlockSize == 64 - && mSageAttnVBlockSize == 256) - { - sage_quant<72, 80, 64, 64, 256, __nv_bfloat16, __nv_fp8_e4m3, float>( - // host var - batch_size, mNumHeads, input_seq_len, true, true, - // device var - // q k v - attention_input, attention_input + mNumHeads * mHeadSize, - attention_input + 2 * mNumHeads * mHeadSize, - // stride - 3 * mNumHeads * mHeadSize, 3 * mNumHeads * mHeadSize, 3 * mNumHeads * mHeadSize, cu_seqlens, - cu_seqlens, sage_quant_space_ptr, - // quant q k v - quanted_qkv_ptr, quanted_qkv_ptr + mNumHeads * paddedHeadSize, - quanted_qkv_ptr + 2 * mNumHeads * paddedHeadSize, - // quanted_qkv_ptr, quanted_qkv_ptr + mNumHeads * mHeadSize, context, - 3 * mNumHeads * paddedHeadSize, 3 * mNumHeads * paddedHeadSize, 3 * mNumHeads * paddedHeadSize, - // scales - q_scale_ptr, k_scale_ptr, v_scale_ptr, stream); - - sync_check_cuda_error(stream); - } - if (mSageAttn && mHeadSize == 80 && mSageAttnQBlockSize == 64 && mSageAttnKBlockSize == 64 - && mSageAttnVBlockSize == 256) - { - sage_quant<80, 80, 64, 64, 256, __nv_bfloat16, __nv_fp8_e4m3, float>( - // host var - batch_size, mNumHeads, input_seq_len, true, true, - // device var - // q k v - attention_input, attention_input + mNumHeads * mHeadSize, - attention_input + 2 * mNumHeads * mHeadSize, - // stride - 3 * mNumHeads * mHeadSize, 3 * mNumHeads * mHeadSize, 3 * mNumHeads * mHeadSize, cu_seqlens, - cu_seqlens, sage_quant_space_ptr, - // quant q k v - quanted_qkv_ptr, quanted_qkv_ptr + mNumHeads * paddedHeadSize, - quanted_qkv_ptr + 2 * mNumHeads * paddedHeadSize, - // quanted_qkv_ptr, quanted_qkv_ptr + mNumHeads * mHeadSize, context, - 3 * mNumHeads * paddedHeadSize, 3 * mNumHeads * paddedHeadSize, 3 * mNumHeads * paddedHeadSize, - // scales - q_scale_ptr, k_scale_ptr, v_scale_ptr, stream); - - sync_check_cuda_error(stream); - } - if (mSageAttn && mHeadSize == 128 && mSageAttnQBlockSize == 64 && mSageAttnKBlockSize == 64 - && mSageAttnVBlockSize == 256) - { - sage_quant<128, 128, 64, 64, 256, __nv_bfloat16, __nv_fp8_e4m3, float>( - // host var - batch_size, mNumHeads, input_seq_len, true, true, - // device var - // q k v - attention_input, attention_input + mNumHeads * mHeadSize, - attention_input + 2 * mNumHeads * mHeadSize, - // stride - 3 * mNumHeads * mHeadSize, 3 * mNumHeads * mHeadSize, 3 * mNumHeads * mHeadSize, cu_seqlens, - cu_seqlens, sage_quant_space_ptr, - // quant q k v - quanted_qkv_ptr, quanted_qkv_ptr + mNumHeads * paddedHeadSize, - quanted_qkv_ptr + 2 * mNumHeads * paddedHeadSize, - // quanted_qkv_ptr, quanted_qkv_ptr + mNumHeads * mHeadSize, context, - 3 * mNumHeads * paddedHeadSize, 3 * mNumHeads * paddedHeadSize, 3 * mNumHeads * paddedHeadSize, - // scales - q_scale_ptr, k_scale_ptr, v_scale_ptr, stream); - - sync_check_cuda_error(stream); - } - if (mSageAttn && mHeadSize == 128 && mSageAttnQBlockSize == 64 && mSageAttnKBlockSize == 32 - && mSageAttnVBlockSize == 32) - { - sage_quant<128, 128, 64, 32, 32, __nv_bfloat16, __nv_fp8_e4m3, float>( - // host var - batch_size, mNumHeads, input_seq_len, true, true, - // device var - // q k v - attention_input, attention_input + mNumHeads * mHeadSize, - attention_input + 2 * mNumHeads * mHeadSize, - // stride - 3 * mNumHeads * mHeadSize, 3 * mNumHeads * mHeadSize, 3 * mNumHeads * mHeadSize, cu_seqlens, - cu_seqlens, sage_quant_space_ptr, - // quant q k v - quanted_qkv_ptr, quanted_qkv_ptr + mNumHeads * paddedHeadSize, - quanted_qkv_ptr + 2 * mNumHeads * paddedHeadSize, - // quanted_qkv_ptr, quanted_qkv_ptr + mNumHeads * mHeadSize, context, - 3 * mNumHeads * paddedHeadSize, 3 * mNumHeads * paddedHeadSize, 3 * mNumHeads * paddedHeadSize, - // scales - q_scale_ptr, k_scale_ptr, v_scale_ptr, stream); - - sync_check_cuda_error(stream); - } - if (mSageAttn && mHeadSize == 80 && mSageAttnQBlockSize == 64 && mSageAttnKBlockSize == 32 - && mSageAttnVBlockSize == 32) - { - sage_quant<80, 80, 64, 32, 32, __nv_bfloat16, __nv_fp8_e4m3, float>( - // host var - batch_size, mNumHeads, input_seq_len, true, true, - // device var - // q k v - attention_input, attention_input + mNumHeads * mHeadSize, - attention_input + 2 * mNumHeads * mHeadSize, - // stride - 3 * mNumHeads * mHeadSize, 3 * mNumHeads * mHeadSize, 3 * mNumHeads * mHeadSize, cu_seqlens, - cu_seqlens, sage_quant_space_ptr, - // quant q k v - quanted_qkv_ptr, quanted_qkv_ptr + mNumHeads * paddedHeadSize, - quanted_qkv_ptr + 2 * mNumHeads * paddedHeadSize, - // quanted_qkv_ptr, quanted_qkv_ptr + mNumHeads * mHeadSize, context, - 3 * mNumHeads * paddedHeadSize, 3 * mNumHeads * paddedHeadSize, 3 * mNumHeads * paddedHeadSize, - // scales - q_scale_ptr, k_scale_ptr, v_scale_ptr, stream); - - sync_check_cuda_error(stream); - } - if (mSageAttn && mHeadSize == 72 && mSageAttnQBlockSize == 64 && mSageAttnKBlockSize == 32 - && mSageAttnVBlockSize == 32) - { - sage_quant<72, 80, 64, 32, 32, __nv_bfloat16, __nv_fp8_e4m3, float>( - // host var - batch_size, mNumHeads, input_seq_len, true, true, - // device var - // q k v - attention_input, attention_input + mNumHeads * mHeadSize, - attention_input + 2 * mNumHeads * mHeadSize, - // stride - 3 * mNumHeads * mHeadSize, 3 * mNumHeads * mHeadSize, 3 * mNumHeads * mHeadSize, cu_seqlens, - cu_seqlens, sage_quant_space_ptr, - // quant q k v - quanted_qkv_ptr, quanted_qkv_ptr + mNumHeads * paddedHeadSize, - quanted_qkv_ptr + 2 * mNumHeads * paddedHeadSize, - // quanted_qkv_ptr, quanted_qkv_ptr + mNumHeads * mHeadSize, context, - 3 * mNumHeads * paddedHeadSize, 3 * mNumHeads * paddedHeadSize, 3 * mNumHeads * paddedHeadSize, - // scales - q_scale_ptr, k_scale_ptr, v_scale_ptr, stream); - - sync_check_cuda_error(stream); - } - - // Construct the fmha params for running kernels. - MHARunnerParams fmhaParams{}; - fmhaParams.b = request_batch_size; - fmhaParams.qSeqLen = request_seq_len; - fmhaParams.kvSeqLen = request_seq_len; - fmhaParams.totalQSeqLen = request_batch_size * request_seq_len; - // Device buffer pointers. - fmhaParams.qkvPtr = attention_input; - fmhaParams.outputPtr = context_buf_; - fmhaParams.cuQSeqLenPtr = cu_seqlens; - fmhaParams.cuKvSeqLenPtr = cu_seqlens; - fmhaParams.tileCounterPtr = fmha_tile_counter_ptr; - fmhaParams.stream = stream; - if (mSageAttn) - { - if (paddedHeadSize != mHeadSize) - fmhaParams.outputPtr = sage_quant_space_ptr; - fmhaParams.qkvPtr = quanted_qkv_ptr; - fmhaParams.scaleBmm1Ptr = scale_bmm1_ptr; - fmhaParams.scaleBmm2Ptr = scale_bmm2_ptr; - fmhaParams.qScalePtr = q_scale_ptr; - fmhaParams.kScalePtr = k_scale_ptr; - fmhaParams.vScalePtr = v_scale_ptr; - fmhaParams.qMaxNBlock = (input_seq_len + mSageAttnQBlockSize - 1) / mSageAttnQBlockSize; - fmhaParams.kMaxNBlock = (input_seq_len + mSageAttnKBlockSize - 1) / mSageAttnKBlockSize; - fmhaParams.vMaxNBlock = (input_seq_len + mSageAttnVBlockSize - 1) / mSageAttnVBlockSize; - } - - // Run the fmha kernel. - - // TODO: set it correctly for contiguous kv buffer (cross-attention). - fmhaParams.totalKvSeqLen = num_tokens; - - fmhaParams.cuKvSeqLenPtr = cu_seqlens; - fmhaParams.cuMaskRowsPtr = cu_seqlens; - fmhaParams.tileCounterPtr = fmha_tile_counter_ptr; - - fmhaParams.scaleBmm1Ptr = scale_bmm1_ptr; - fmhaParams.scaleBmm2Ptr = scale_bmm2_ptr; - fmhaParams.forceFp32Acc = mFMHAForceFP32Acc; - mFmhaDispatcher->run(fmhaParams); - sync_check_cuda_error(stream); - if (mSageAttn) - { - if (paddedHeadSize != mHeadSize && mHeadSize == 72) - { - unpadding<80, 72, __nv_bfloat16>(batch_size, mNumHeads, input_seq_len, sage_quant_space_ptr, - mNumHeads * 72, mNumHeads * 80, cu_seqlens, context_buf_, stream); - } - } - } - } - else - { - // FIXME: a temporary solution to make sure the padding part of key/value buffer is 0 - // NOTE: pointer subtraction is used below since there could be some extra gap due to alignment. - // Otherwise, we could do cudaMemsetAsync(k_buf_2_, 0, k_buf_2_size + v_buf_2_size, stream); - // cudaMemsetAsync(k_buf_2_, 0, reinterpret_cast(qk_buf_) - reinterpret_cast(k_buf_2_), - // stream); - // FIXME: the final solution is to change the add_fusedQKV_bias_transpose_kernel to map CTAs corresponding to - // the output shape, and set the padding part to 0. Without zero-initialize guarantee, these workspace buffers - // may contain random NaN values when IFB workload is high. - cudaMemsetAsync(k_buf_2_, 0, - reinterpret_cast(v_buf_2_) - reinterpret_cast(k_buf_2_) + v_buf_2_size, stream); - - // only non-FMHA path needs to split Q,K,V from QKV - invokeAddFusedQKVBiasTranspose(q_buf_2_, k_buf_2_, v_buf_2_, const_cast(attention_input), input_lengths, - mRemovePadding ? padding_offset : nullptr, batch_size, input_seq_len, num_tokens, mNumHeads, mNumHeads, - mHeadSize, 0, 0.0f, RotaryScalingType::kNONE, 0.0f, 0, PositionEmbeddingType::kLEARNED_ABSOLUTE, - (float*) nullptr, 0, stream); - - if (!mQKHalfAccum && gemm_data_type != CUDA_R_32F) - { - mCublasWrapper->stridedBatchedGemm(CUBLAS_OP_T, CUBLAS_OP_N, - attention_seq_len_2, // n - attention_seq_len_1, // m - mHeadSize, // k - qk_scale_gemm, k_buf_2_, gemm_data_type, - mHeadSize, // k - attention_seq_len_2 * mHeadSize, // n * k - q_buf_2_, gemm_data_type, - mHeadSize, // k - attention_seq_len_1 * mHeadSize, // m * k - 0.0f, qk_buf_float_, CUDA_R_32F, - attention_seq_len_2, // n - attention_seq_len_2 * attention_seq_len_1, - request_batch_size * mNumHeads, // global batch size - CUDA_R_32F); - - // add relative position bias - if (mRelativeAttention) - { - // add rel pos bias - // QK is (batch_size, local_head_num, q_length, k_length), rel pos bias is (1, local_head_num, - // max_output_len + 1, max_output_len + 1). broadcast along 1st dim. max_seq_len is already - // max_output_len + 1. In implicit mode, relative_attention_bias is rel attn table - // [num_heads, num_buckets], with necessary params (max_distance, num_buckets) passed at the end - invokeAddRelativeAttentionBiasUnaligned(qk_buf_float_, relative_attn_table, request_batch_size, - mNumHeads, attention_seq_len_1, attention_seq_len_2, stream, mMaxDistance > 0, - inputDesc[3].dims.d[1], mMaxDistance, true /* bidirectional */); - } - - MaskedSoftmaxParam param; - param.attention_score = qk_buf_; // (batch_size, head_num, q_length, k_length) - param.qk = qk_buf_float_; // (batch_size, head_num, q_length, k_length) - param.attention_mask = attention_mask; // (batch_size, q_length, k_length) - param.batch_size = request_batch_size; - param.q_length = attention_seq_len_1; - param.k_length = attention_seq_len_2; - param.num_heads = mNumHeads; - param.qk_scale = qk_scale_softmax; - param.linear_bias_slopes = const_cast(linear_bias_slopes); // (head_num,), optional - invokeMaskedSoftmax(param, stream); - } - else - { - mCublasWrapper->stridedBatchedGemm(CUBLAS_OP_T, CUBLAS_OP_N, attention_seq_len_2, attention_seq_len_1, - mHeadSize, k_buf_2_, mHeadSize, attention_seq_len_2 * mHeadSize, q_buf_2_, mHeadSize, - attention_seq_len_1 * mHeadSize, qk_buf_, attention_seq_len_2, - attention_seq_len_2 * attention_seq_len_1, request_batch_size * mNumHeads, qk_scale_gemm, - 0.0f); // alpha, beta - - // add relative position bias - if (mRelativeAttention) - { - // add rel pos bias - // QK is (batch_size, local_head_num, q_length, k_length), rel pos bias is (1, local_head_num, - // max_output_len + 1, max_output_len + 1). broadcast along 1st dim. max_seq_len is already - // max_output_len + 1. In implicit mode, relative_attention_bias is rel attn table - // [num_heads, num_buckets], with necessary params (max_distance, num_buckets) passed at the end - invokeAddRelativeAttentionBiasUnaligned(qk_buf_, relative_attn_table, request_batch_size, mNumHeads, - attention_seq_len_1, attention_seq_len_2, stream, mMaxDistance > 0, inputDesc[3].dims.d[1], - mMaxDistance, true /* bidirectional */); - } - - MaskedSoftmaxParam param; - param.attention_score = qk_buf_; // (batch_size, head_num, q_length, k_length) - param.qk = qk_buf_; // (batch_size, head_num, q_length, k_length) - param.attention_mask = attention_mask; // (batch_size, q_length, k_length) - param.batch_size = request_batch_size; - param.q_length = attention_seq_len_1; - param.k_length = attention_seq_len_2; - param.num_heads = mNumHeads; - param.qk_scale = qk_scale_softmax; - param.linear_bias_slopes = const_cast(linear_bias_slopes); // (head_num,), optional - invokeMaskedSoftmax(param, stream); - } - - mCublasWrapper->stridedBatchedGemm(CUBLAS_OP_N, CUBLAS_OP_N, mHeadSize, attention_seq_len_1, - attention_seq_len_2, v_buf_2_, mHeadSize, attention_seq_len_2 * mHeadSize, qk_buf_, attention_seq_len_2, - attention_seq_len_1 * attention_seq_len_2, qkv_buf_2_, mHeadSize, attention_seq_len_1 * mHeadSize, - request_batch_size * mNumHeads); - - if (!mRemovePadding) - { - invokeTransposeQKV(context_buf_, qkv_buf_2_, request_batch_size, attention_seq_len_1, mNumHeads, mHeadSize, - (float*) nullptr, 0, stream); - } - else - { - invokeTransposeAttentionOutRemovePadding(qkv_buf_2_, context_buf_, num_tokens, request_batch_size, - request_seq_len, mNumHeads, mHeadSize, padding_offset, (float*) nullptr, 0, stream); - } - } - sync_check_cuda_error(stream); - return 0; -} - -template int BertAttentionPlugin::enqueueImpl(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream); - -template int BertAttentionPlugin::enqueueImpl(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream); - -#ifdef ENABLE_BF16 -template int BertAttentionPlugin::enqueueImpl<__nv_bfloat16>(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream); -#endif - -int BertAttentionPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) noexcept -{ - if (mType == DataType::kHALF) - { - return enqueueImpl(inputDesc, outputDesc, inputs, outputs, workspace, stream); - } - else if (mType == DataType::kFLOAT) - { - return enqueueImpl(inputDesc, outputDesc, inputs, outputs, workspace, stream); - } -#ifdef ENABLE_BF16 - else if (mType == DataType::kBF16) - { - return enqueueImpl<__nv_bfloat16>(inputDesc, outputDesc, inputs, outputs, workspace, stream); - } -#endif - return 0; -} - -// IPluginV2Ext Methods -nvinfer1::DataType BertAttentionPlugin::getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept -{ - TLLM_CHECK(index == 0); - return inputTypes[0]; -} - -// IPluginV2 Methods - -char const* BertAttentionPlugin::getPluginType() const noexcept -{ - return BERT_ATTENTION_PLUGIN_NAME; -} - -char const* BertAttentionPlugin::getPluginVersion() const noexcept -{ - return BERT_ATTENTION_PLUGIN_VERSION; -} - -int BertAttentionPlugin::getNbOutputs() const noexcept -{ - return 1; -} - -int BertAttentionPlugin::initialize() noexcept -{ - auto cublasHandle = getCublasHandle(); - auto cublasLtHandle = getCublasLtHandle(); - mCublasWrapper.reset(new tc::CublasMMWrapper(cublasHandle, cublasLtHandle, nullptr, nullptr)); - if (mEnableContextFMHA) - { - // Pre-checked during constructing. - Data_type data_type; - if (mType == DataType::kHALF) - { - data_type = DATA_TYPE_FP16; - } - else if (mType == DataType::kBF16) - { - data_type = DATA_TYPE_BF16; - } - else - { - TLLM_CHECK_WITH_INFO(false, "GPTAttentionPlugin received wrong data type."); - } - - // Construct the fmha runner. - MHARunnerFixedParams fmhaParams{}; - if (mSageAttn) - { - fmhaParams.dataType = DATA_TYPE_E4M3; - } - else - { - fmhaParams.dataType = data_type; - } - fmhaParams.dataTypeOut = data_type; - fmhaParams.forceFp32Acc = mFMHAForceFP32Acc; - fmhaParams.attentionMaskType = ContextAttentionMaskType::PADDING; - fmhaParams.isSPadded = !mRemovePadding; - fmhaParams.numQHeads = mNumHeads; - fmhaParams.numKvHeads = mNumHeads; - fmhaParams.headSize = mHeadSize; - fmhaParams.qScaling = mQScaling; - fmhaParams.sageBlockSizeQ = mSageAttnQBlockSize; - fmhaParams.sageBlockSizeK = mSageAttnKBlockSize; - fmhaParams.sageBlockSizeV = mSageAttnVBlockSize; - if (mSageAttn) - { - int const paddedHeadSize = ((mHeadSize + 15) / 16) * 16; - fmhaParams.headSize = paddedHeadSize; - } - - if (mCpGroup.size() > 1) - { - fmhaParams.attentionInputLayout = AttentionInputLayout::Q_CONTIGUOUS_KV; - fmhaParams.saveSoftmax = true; - } - - // Load kernels from the pre-compiled cubins. - // The KV input data type. The default is same as dataType. - fmhaParams.dataTypeKv = data_type; - fmhaParams.headSizeV = mHeadSize; - - // Load kernels from the pre-compiled cubins. - mFmhaDispatcher.reset(new FmhaDispatcher(fmhaParams)); - // Fall back to unfused MHA kernels if not supported. - mEnableContextFMHA = mFmhaDispatcher->isSupported(); - } - -#if ENABLE_MULTI_DEVICE - if (mCpGroup.size() > 1 && COMM_SESSION.getSize() > 1) - { - TLLM_LOG_TRACE("%s start for rank %d", __PRETTY_FUNCTION__, COMM_SESSION.getRank()); - mNcclComm = getComm(mCpGroup); - TLLM_LOG_TRACE("%s stop for rank %d", __PRETTY_FUNCTION__, COMM_SESSION.getRank()); - } -#endif // ENABLE_MULTI_DEVICE - - return 0; -} - -void BertAttentionPlugin::destroy() noexcept -{ - delete this; -} - -size_t BertAttentionPlugin::getSerializationSize() const noexcept -{ - return sizeof(mNumHeads) + sizeof(mHeadSize) + sizeof(mQScaling) + sizeof(mQKHalfAccum) + sizeof(mEnableContextFMHA) - + sizeof(mFMHAForceFP32Acc) + sizeof(mType) + sizeof(mRelativeAttention) + sizeof(mMaxDistance) - + sizeof(mRemovePadding) + sizeof(mSageAttn) + sizeof(mSageAttnQBlockSize) + sizeof(mSageAttnKBlockSize) - + sizeof(mSageAttnVBlockSize) + sizeof(mCpSize) + sizeof(mCpRank) + sizeof(int32_t) * mCpGroup.size(); -} - -void BertAttentionPlugin::serialize(void* buffer) const noexcept -{ - char *d = static_cast(buffer), *a = d; - write(d, mNumHeads); - write(d, mHeadSize); - write(d, mQScaling); - write(d, mQKHalfAccum); - write(d, mEnableContextFMHA); - write(d, mFMHAForceFP32Acc); - write(d, mType); - write(d, mRelativeAttention); - write(d, mMaxDistance); - write(d, mRemovePadding); - write(d, mSageAttn); - write(d, mSageAttnQBlockSize); - write(d, mSageAttnKBlockSize); - write(d, mSageAttnVBlockSize); - write(d, mCpSize); - write(d, mCpRank); - for (auto it = mCpGroup.begin(); it != mCpGroup.end(); ++it) - { - write(d, *it); - } - TLLM_CHECK(d == a + getSerializationSize()); -} - -void BertAttentionPlugin::terminate() noexcept {} - -/////////////// - -BertAttentionPluginCreator::BertAttentionPluginCreator() -{ - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - - mPluginAttributes.emplace_back(PluginField("num_heads", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("head_size", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("q_scaling", nullptr, PluginFieldType::kFLOAT32)); - mPluginAttributes.emplace_back(PluginField("context_fmha_type", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("type_id", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("do_relative_attention", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("max_distance", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("remove_padding", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("sage_attn", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("sage_attn_q_block_size", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("sage_attn_k_block_size", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("sage_attn_v_block_size", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("cp_size", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("cp_rank", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("cp_group", nullptr, PluginFieldType::kINT32)); - - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -char const* BertAttentionPluginCreator::getPluginName() const noexcept -{ - return BERT_ATTENTION_PLUGIN_NAME; -} - -char const* BertAttentionPluginCreator::getPluginVersion() const noexcept -{ - return BERT_ATTENTION_PLUGIN_VERSION; -} - -PluginFieldCollection const* BertAttentionPluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -IPluginV2* BertAttentionPluginCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept -{ - PluginField const* fields = fc->fields; - int num_heads{}; - int head_size{}; - ContextFMHAType context_fmha_type{}; - float q_scaling{}; - nvinfer1::DataType type{}; - bool do_relative_attention{}; - int max_distance{}; - bool remove_padding{}; - bool sage_attn{}; - int sage_attn_q_block_size{}; - int sage_attn_k_block_size{}; - int sage_attn_v_block_size{}; - int cp_size{}; - int cp_rank{}; - std::set cp_group{}; - - // Read configurations from each fields - for (int i = 0; i < fc->nbFields; ++i) - { - char const* attrName = fields[i].name; - if (!strcmp(attrName, "num_heads")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - num_heads = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "head_size")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - head_size = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "q_scaling")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kFLOAT32); - q_scaling = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "context_fmha_type")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT8); - context_fmha_type = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "type_id")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - type = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "do_relative_attention")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT8); - do_relative_attention = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "max_distance")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - max_distance = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "remove_padding")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT8); - remove_padding = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "sage_attn")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT8); - sage_attn = static_cast(*(static_cast(fields[i].data))); - if (sage_attn) - { - std::cout << "sage attn true!" << std::endl; - } - } - else if (!strcmp(attrName, "sage_attn_q_block_size")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - sage_attn_q_block_size = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "sage_attn_k_block_size")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - sage_attn_k_block_size = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "sage_attn_v_block_size")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - sage_attn_v_block_size = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "cp_size")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - cp_size = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "cp_rank")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - cp_rank = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "cp_group")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - auto const* r = static_cast(fields[i].data); - for (int j = 0; j < fields[i].length; ++j) - { - cp_group.insert(*r); - ++r; - } - } - } - try - { - auto* obj = new BertAttentionPlugin(num_heads, head_size, q_scaling, context_fmha_type, type, - do_relative_attention, max_distance, remove_padding, sage_attn, sage_attn_q_block_size, - sage_attn_k_block_size, sage_attn_v_block_size, cp_size, cp_rank, cp_group); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -IPluginV2* BertAttentionPluginCreator::deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept -{ - // This object will be deleted when the network is destroyed, which will - // call BertAttentionPlugin::destroy() - try - { - auto* obj = new BertAttentionPlugin(serialData, serialLength); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} diff --git a/cpp/tensorrt_llm/plugins/bertAttentionPlugin/bertAttentionPlugin.h b/cpp/tensorrt_llm/plugins/bertAttentionPlugin/bertAttentionPlugin.h deleted file mode 100644 index 2eb39086a005..000000000000 --- a/cpp/tensorrt_llm/plugins/bertAttentionPlugin/bertAttentionPlugin.h +++ /dev/null @@ -1,142 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2025 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. - */ -#pragma once - -#include "tensorrt_llm/common/cublasMMWrapper.h" -#include "tensorrt_llm/common/quantization.h" -#include "tensorrt_llm/kernels/fmhaDispatcher.h" -#include "tensorrt_llm/kernels/gptKernels.h" -#include "tensorrt_llm/plugins/common/plugin.h" -#include "tensorrt_llm/runtime/utils/mpiUtils.h" -#include -#include -#include -#include -#include - -namespace tensorrt_llm::plugins -{ - -class BertAttentionPlugin : public BasePlugin -{ -public: - BertAttentionPlugin() = delete; - - BertAttentionPlugin(int num_heads, int head_size, float q_scaling, - tensorrt_llm::kernels::ContextFMHAType context_fmha_type, nvinfer1::DataType type, - bool do_relative_attention = false, int max_distance = 0, bool remove_padding = false, bool sage_attn = false, - int sage_attn_q_block_size = 0, int sage_attn_k_block_size = 0, int sage_attn_v_block_size = 0, int cp_size = 1, - int cp_rank = 0, std::set cp_group = {}); - - BertAttentionPlugin(void const* data, size_t length); - - ~BertAttentionPlugin() override = default; - - // IPluginV2DynamicExt Methods - nvinfer1::IPluginV2DynamicExt* clone() const noexcept override; - nvinfer1::DimsExprs getOutputDimensions(int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; - bool supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept override; - void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept override; - size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept override; - int enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - - template - int enqueueImpl(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream); - - // IPluginV2Ext Methods - nvinfer1::DataType getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept override; - - // IPluginV2 Methods - char const* getPluginType() const noexcept override; - char const* getPluginVersion() const noexcept override; - int getNbOutputs() const noexcept override; - int initialize() noexcept override; - void terminate() noexcept override; - size_t getSerializationSize() const noexcept override; - void serialize(void* buffer) const noexcept override; - void destroy() noexcept override; - -private: - const std::string mLayerName; - - int mNumHeads; - int mHeadSize; - float mQScaling; - nvinfer1::DataType mType; - bool mRelativeAttention = false; - int mMaxDistance = 0; - bool mRemovePadding = false; - - // unfused mha - bool mQKHalfAccum = false; - - // fmha runner (disable by default) - bool mEnableContextFMHA = false; - bool mFMHAForceFP32Acc = false; - - // sage attention - bool mSageAttn = false; - int mSageAttnQBlockSize = 0; - int mSageAttnKBlockSize = 0; - int mSageAttnVBlockSize = 0; - std::set> mSageAttnSupportedBlockSizes{{64, 64, 256}, {64, 32, 32}}; - - int mSM = tensorrt_llm::common::getSMVersion(); - - // comm group for RingAttention - int mCpSize = 1; - int mCpRank = 0; - std::set mCpGroup = {}; -#if ENABLE_MULTI_DEVICE - std::shared_ptr mNcclComm; -#endif // ENABLE_MULTI_DEVICE - cudaStream_t mNcclStream; - - // The default copy constructor will leave them as nullptr. clone() shall initialize it. - UniqPtrWNullCopy mFmhaDispatcher; - UniqPtrWNullCopy mCublasWrapper; -}; - -class BertAttentionPluginCreator : public BaseCreator -{ -public: - BertAttentionPluginCreator(); - - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; - - nvinfer1::IPluginV2* deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept override; - -private: - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; -}; - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/common/CMakeLists.txt b/cpp/tensorrt_llm/plugins/common/CMakeLists.txt deleted file mode 100644 index 86876224fccd..000000000000 --- a/cpp/tensorrt_llm/plugins/common/CMakeLists.txt +++ /dev/null @@ -1,21 +0,0 @@ -# -# SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. -# -file(GLOB SRCS *.cpp) -set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS}) -set(PLUGIN_SOURCES - ${PLUGIN_SOURCES} - PARENT_SCOPE) diff --git a/cpp/tensorrt_llm/plugins/common/checkMacrosPlugin.cpp b/cpp/tensorrt_llm/plugins/common/checkMacrosPlugin.cpp deleted file mode 100644 index 2aab6b3675d8..000000000000 --- a/cpp/tensorrt_llm/plugins/common/checkMacrosPlugin.cpp +++ /dev/null @@ -1,35 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. - */ - -#include "checkMacrosPlugin.h" - -#include "tensorrt_llm/common/logger.h" - -namespace tensorrt_llm::plugins -{ - -void caughtError(std::exception const& e) -{ - TLLM_LOG_EXCEPTION(e); -} - -void logError(char const* msg, char const* file, char const* fn, int line) -{ - TLLM_LOG_ERROR("Parameter check failed at: %s::%s::%d, condition: %s", file, fn, line, msg); -} - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/common/checkMacrosPlugin.h b/cpp/tensorrt_llm/plugins/common/checkMacrosPlugin.h deleted file mode 100644 index d8d8af1ef220..000000000000 --- a/cpp/tensorrt_llm/plugins/common/checkMacrosPlugin.h +++ /dev/null @@ -1,29 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. - */ -#pragma once - -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/cudaUtils.h" - -namespace tensorrt_llm::plugins -{ - -void logError(char const* msg, char const* file, char const* fn, int line); - -void caughtError(std::exception const& e); - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/common/gemmPluginProfiler.cpp b/cpp/tensorrt_llm/plugins/common/gemmPluginProfiler.cpp deleted file mode 100644 index e5d6650648ab..000000000000 --- a/cpp/tensorrt_llm/plugins/common/gemmPluginProfiler.cpp +++ /dev/null @@ -1,404 +0,0 @@ -/* - * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved. - * - * 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. - */ - -#include "tensorrt_llm/plugins/common/gemmPluginProfiler.h" -#include "tensorrt_llm/common/cublasMMWrapper.h" -#include "tensorrt_llm/kernels/cutlass_kernels/fp8_rowwise_gemm/fp8_rowwise_gemm.h" -#include "tensorrt_llm/kernels/cutlass_kernels/fpA_intB_gemm/fpA_intB_gemm.h" -#include "tensorrt_llm/kernels/cutlass_kernels/fused_gated_gemm/fused_gated_gemm.h" -#include "tensorrt_llm/kernels/cutlass_kernels/int8_gemm/int8_gemm.h" -#include "tensorrt_llm/plugins/gemmAllReducePlugin/gemmAllReducePlugin.h" -#include "tensorrt_llm/plugins/lowLatencyGemmPlugin/lowLatencyGemmPlugin.h" -#include "tensorrt_llm/plugins/lowLatencyGemmSwigluPlugin/lowLatencyGemmSwigluPlugin.h" -#include "tensorrt_llm/plugins/mixtureOfExperts/mixtureOfExpertsPlugin.h" -#if defined(USING_OSS_CUTLASS_FP4_GEMM) -#include "tensorrt_llm/kernels/cutlass_kernels/include/fp4_gemm.h" -#else -#include "fp4_gemm.h" -#endif -#if defined(USING_OSS_CUTLASS_ALLREDUCE_GEMM) -#include "tensorrt_llm/kernels/cutlass_kernels/include/allreduce_gemm_runner.h" -using GemmAllReduceImplInterface = tensorrt_llm::kernels::opened_cutlass_kernels::GemmAllReduceImplInterface; -#else -#include "allreduce_gemm_runner.h" -using GemmAllReduceImplInterface = tensorrt_llm::kernels::cutlass_kernels::GemmAllReduceImplInterface; -#endif - -#include - -namespace tensorrt_llm::plugins -{ - -template -GemmPluginProfiler::GemmPluginProfiler() -{ - mMNKProfileMap = std::make_shared(); - - // set SKIP_GEMM_PLUGIN_PROFILINGS=1 to avoid tactics profilings - auto const skipEnv = std::getenv("SKIP_GEMM_PLUGIN_PROFILINGS"); - mSkip = (skipEnv != NULL && std::stoi(skipEnv)); - if (mSkip) - { - TLLM_LOG_DEBUG( - "SKIP_GEMM_PLUGIN_PROFILINGS is set. Skipping GEMM plugin profilings. It could result in runtime error " - "if default tactic is not defined."); - } -} - -template -void GemmPluginProfiler::serialize( - char*& buffer, GemmIdType const& gemmId) const -{ - auto mProfileMap = mMNKProfileMap->getMProfileMap(gemmId); - - // Save number of profiles for given GEMM ID - write(buffer, static_cast(mProfileMap->size())); - for (auto const& pair : *mProfileMap) - { - // Save pair of M to the best GEMM config - write(buffer, pair); - } -} - -template -void GemmPluginProfiler::deserialize( - char const*& data, GemmDims& dims, GemmIdType const& gemmId) -{ - // NOTE: this mutex is not needed since each thread owns its private map, but will put here for - // consistency - writer_lock lock(mMNKProfileMap->mutex); - - mDims = dims; - - // GemmId gemmId(dims.n, dims.k); - if (!mMNKProfileMap->existsMProfileMap(gemmId)) - { - // Create GEMM with GEMM ID if it does not exist - mMNKProfileMap->createMProfileMap(gemmId); - } - // Populate map with profiles of GEMM ID - auto profileMap = mMNKProfileMap->getMProfileMap(gemmId); - int selectedMapSize; - read(data, selectedMapSize); - for (int ii = 0; ii < selectedMapSize; ++ii) - { - std::pair> config; - read(data, config); - profileMap->insert(config); - } -} - -template -size_t GemmPluginProfiler::getSerializationSize( - GemmIdType const& gemmId) const -{ - reader_lock lock(mMNKProfileMap->mutex); - return sizeof(int) + // size of the tactics map - mMNKProfileMap->getMProfileMap(gemmId)->size() - * sizeof(std::pair>); // size of the tactics map -} - -template -int GemmPluginProfiler::getMaxProfileM() const -{ - return 8192; -} - -template -void GemmPluginProfiler::initTmpData( - int m, int n, int k, char* workspace, size_t size, cudaStream_t stream) -{ - /* Do nothing */ -} - -template -void GemmPluginProfiler::profileTactics(RunnerPtr const& runner, - nvinfer1::DataType const& type, GemmDims const& dims, GemmIdType const& gemmId, bool hasWeightOnlyCudaKernel) -{ - writer_lock lock(mMNKProfileMap->mutex); - - if (!dims.isInitialized()) - { - return; - } - - mRunner = runner; - mType = type; - - int const maxM = std::min(nextPowerOfTwo(dims.maxM), getMaxProfileM()); - computeTmpSize(maxM, dims.n, dims.k); - - if (!mMNKProfileMap->existsMProfileMap(gemmId)) - { - // Create map for GEMM ID - mMNKProfileMap->createMProfileMap(gemmId); - } - - if (mSkip) - { - return; - } - - auto mProfileMap = mMNKProfileMap->getMProfileMap(gemmId); - bool isAllocated{false}; - - auto profileTactics = [&mProfileMap, &isAllocated, this](int m, int n, int k) - { - if (mProfileMap->count(m) == 0) - { - if (!isAllocated) - { - // Allocate tmp data to run GEMMs - allocateTmpData(); - isAllocated = true; - } - initTmpData(m, n, k, mWorkspaceTmp, mTmpWorkspaceSizeInBytes, mStream); - auto tactics = this->getTactics(m, n, k); - - // Profile different tactics for particular m and insert best config to the map - mProfileMap->insert({m, this->profileTacticsForProblem(m, n, k, tactics)}); - } - }; - - common::check_cuda_error(cudaStreamCreate(&mStream)); - - int const startMinMRounded = nextPowerOfTwo(dims.minM); - - if (hasWeightOnlyCudaKernel) - { - // Profile tactics for finer granularity of M, - // if CUDA kernel is enabled for weight-only plugins - int minM = dims.minM; - for (int m = std::max(1, minM); m < std::min(16, maxM); m += 1) - { - profileTactics(m, dims.n, dims.k); - } - - for (int m = 16; m < maxM; m *= 2) - { - profileTactics(m, dims.n, dims.k); - } - } - else - { - // Profile tactics for CUTLASS kernel only - for (int m = std::max(1, startMinMRounded); m < maxM; m *= 2) - { - profileTactics(m, dims.n, dims.k); - } - } - - profileTactics(maxM, dims.n, dims.k); - - if (isAllocated) - { - // Free tmp data - freeTmpData(); - } - common::check_cuda_error(cudaStreamDestroy(mStream)); -} - -template -std::optional GemmPluginProfiler::getBestConfig( - int m, GemmIdType const& gemmId) const -{ - reader_lock lock(mMNKProfileMap->mutex); - - if (mSkip) - { - TLLM_LOG_TRACE("Skip is set, no best config is set for this instance"); - return std::nullopt; - } - - int const mRounded = std::min(std::max(1, nextPowerOfTwo(m)), getMaxProfileM()); - fflush(stdout); - - if (mMNKProfileMap->getMProfileMap(gemmId)->count(m) > 0) - { - return mMNKProfileMap->getMProfileMap(gemmId)->at(m); - } - else if (mMNKProfileMap->getMProfileMap(gemmId)->count(mRounded) > 0) - { - return mMNKProfileMap->getMProfileMap(gemmId)->at(mRounded); - } - else - { - std::ostringstream msg; - msg << "Cannot find best tactic for m=" << m << " and GEMM ID " << gemmId; - TLLM_LOG_WARNING(msg.str()); - return std::nullopt; - } -} - -template -void GemmPluginProfiler::allocateTmpData() -{ - TLLM_CHECK_WITH_INFO(mTmpWorkspaceSizeInBytes > 0, "tmpWorkspaceSizeInBytes must be larger than 0"); - auto const status = cudaMalloc(&mWorkspaceTmp, mTmpWorkspaceSizeInBytes); - TLLM_CHECK_WITH_INFO(status == cudaSuccess, "Can't allocate tmp workspace for GEMM tactics profiling."); -} - -template -void GemmPluginProfiler::freeTmpData() -{ - auto const status = cudaFree(mWorkspaceTmp); - TLLM_CHECK_WITH_INFO(status == cudaSuccess, "Can't free tmp workspace for GEMM tactics profiling."); -} - -template -std::optional GemmPluginProfiler::profileTacticsForProblem( - int m, int n, int k, std::vector const& tactics) -{ - TLLM_LOG_DEBUG(__PRETTY_FUNCTION__); - - float bestTime = std::numeric_limits::max(); - Config bestConfig; - bool foundOne = false; - - // Iterate over all tactics for given M, N and K - for (size_t ii = 0; ii < tactics.size(); ++ii) - { - Config const& candidateConfig = tactics[ii]; - float time = std::numeric_limits::max(); - try - { - if (!checkTactic(m, n, k, candidateConfig)) - { - continue; - } - // Profile particular tactic for given M, N and K - time = profileTacticForProblem(m, n, k, candidateConfig); - foundOne = true; - } - catch (std::exception const& e) - { - std::ostringstream msg; - msg << "Cannot profile configuration " << ii; - if constexpr (std::is_same_v) - { - msg << ": " << candidateConfig.toString(); - } - msg << "\n (for" - << " m=" << m << ", n=" << n << ", k=" << k << ")" - << ", reason: \"" << e.what() << "\". Skipped"; - TLLM_LOG_TRACE(msg.str()); - cudaGetLastError(); // Reset the last cudaError to cudaSuccess. - continue; - } - - // Choose the fastest tactic - if (time < bestTime) - { - bestConfig = candidateConfig; - bestTime = time; - } - } - - if (!foundOne) - { - std::ostringstream msg; - msg << "Have not found any valid GEMM config for shape (" - << "m=" << m << ", n=" << n << ", k=" << k << "). Will try to use default or fail at runtime"; - TLLM_LOG_WARNING(msg.str()); - return std::nullopt; - } - - return {bestConfig}; -} - -template -float GemmPluginProfiler::profileTacticForProblem( - int m, int n, int k, Config const& tactic) -{ - constexpr int warmup = 5; - constexpr int runs = 10; - - cudaStream_t stream = mStream; - - // Warmup the execution - for (int i = 0; i < warmup; ++i) - { - runTactic(m, n, k, tactic, mWorkspaceTmp, stream); - } - - cudaEvent_t start; - cudaEvent_t stop; - common::check_cuda_error(cudaEventCreate(&start)); - common::check_cuda_error(cudaEventCreate(&stop)); - common::check_cuda_error(cudaStreamSynchronize(stream)); - common::check_cuda_error(cudaEventRecord(start, stream)); - - // Profile GEMM - for (int i = 0; i < runs; ++i) - { - runTactic(m, n, k, tactic, mWorkspaceTmp, stream); - } - - common::check_cuda_error(cudaEventRecord(stop, stream)); - - common::check_cuda_error(cudaEventSynchronize(stop)); - - float elapsed; - common::check_cuda_error(cudaEventElapsedTime(&elapsed, start, stop)); - - common::check_cuda_error(cudaEventDestroy(start)); - common::check_cuda_error(cudaEventDestroy(stop)); - - return elapsed / runs; -} - -template class GemmPluginProfiler, GemmIdCore, - GemmIdCoreHash>; - -template class GemmPluginProfiler, GemmIdCore, - GemmIdCoreHash>; - -template class GemmPluginProfiler, GemmIdCublas, GemmIdCublasHash>; - -// TODO I dont like the dependency on the MOE plugin here, but MOE needs the full context to run profiles -template class GemmPluginProfiler; - -template class GemmPluginProfiler, GemmIdCore, - GemmIdCoreHash>; - -template class GemmPluginProfiler, GemmIdCore, - GemmIdCoreHash>; - -#if defined(USING_OSS_CUTLASS_FP4_GEMM) -template class GemmPluginProfiler, GemmIdCore, GemmIdCoreHash>; -#else -template class GemmPluginProfiler, GemmIdCore, - GemmIdCoreHash>; -#endif - -template class GemmPluginProfiler; - -template class GemmPluginProfiler; - -template class GemmPluginProfiler, - GemmIdCore, GemmIdCoreHash>; - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/common/gemmPluginProfiler.h b/cpp/tensorrt_llm/plugins/common/gemmPluginProfiler.h deleted file mode 100644 index fe85b3b7e456..000000000000 --- a/cpp/tensorrt_llm/plugins/common/gemmPluginProfiler.h +++ /dev/null @@ -1,332 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2023 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. - */ -#pragma once - -#include "pluginUtils.h" - -#include - -#include -#include -#include -#include -#include -#include -#include -#include -#include - -namespace tensorrt_llm::plugins -{ - -struct GemmDims -{ - using DimType64 = utils::DimType64; - - DimType64 minM; - DimType64 maxM; - DimType64 n; - DimType64 k; - - GemmDims() - : minM(-1) - , maxM(-1) - , n(-1) - , k(-1) - { - } - - GemmDims(DimType64 minM_, DimType64 maxM_, DimType64 n_, DimType64 k_) - : minM(minM_) - , maxM(maxM_) - , n(n_) - , k(k_) - { - } - - [[nodiscard]] bool isInitialized() const - { - return minM >= 0 && maxM >= 0 && n >= 0 && k >= 0; - } -}; - -// Unique ID of GEMM -// In our case GEMM is uniqly identified by N and K -class GemmIdCore -{ -public: - int n; - int k; - nvinfer1::DataType dtype; - - GemmIdCore(int n_, int k_, nvinfer1::DataType const& dtype_) - : n(n_) - , k(k_) - , dtype(dtype_) - { - } - - GemmIdCore() - : n(-1) - , k(-1) - , dtype(nvinfer1::DataType::kFLOAT) // dtype does not matter here - { - } - - bool operator==(GemmIdCore const& id) const - { - return isEqual(id); - } - - friend std::ostream& operator<<(std::ostream& out, GemmIdCore const& id) - { - out << "(N;K)=(" << id.n << ";" << id.k << "),"; - out << " type=" << static_cast(id.dtype); - return out; - } - -protected: - bool isEqual(GemmIdCore const& id) const - { - return n == id.n && k == id.k && dtype == id.dtype; - } -}; - -// Hash of GemmId -struct GemmIdCoreHash -{ - std::size_t operator()(GemmIdCore const& id) const - { - auto h1 = std::hash{}(id.n); - auto h2 = std::hash{}(id.k); - auto h3 = std::hash{}(static_cast(id.dtype)); - return h1 ^ h2 ^ h3; - } -}; - -class GemmIdCublas : public GemmIdCore -{ -public: - bool transA{}; - bool transB{}; - nvinfer1::DataType outputDtype; - - GemmIdCublas(int n_, int k_, nvinfer1::DataType const& dtype_, bool transA_, bool transB_, - nvinfer1::DataType const& output_dtype_) - : GemmIdCore(n_, k_, dtype_) - , transA(transA_) - , transB(transB_) - , outputDtype(output_dtype_) - { - } - - GemmIdCublas() {} - - bool operator==(GemmIdCublas const& id) const - { - return isEqual(id) && transA == id.transA && transB == id.transB && outputDtype == id.outputDtype; - } - - friend std::ostream& operator<<(std::ostream& out, GemmIdCublas const& id) - { - out << "(N;K)=(" << id.n << ";" << id.k << "),"; - out << " type=" << static_cast(id.dtype); - out << " transA=" << id.transA; - out << " transB=" << id.transB; - out << " outputDtype=" << static_cast(id.outputDtype); - return out; - } -}; - -// Hash of GemmIdCublas -struct GemmIdCublasHash -{ - std::size_t operator()(GemmIdCublas const& id) const - { - auto h1 = std::hash{}(id.n); - auto h2 = std::hash{}(id.k); - auto h3 = std::hash{}(static_cast(id.dtype)); - auto h4 = std::hash{}(id.transA); - auto h5 = std::hash{}(id.transB); - auto h6 = std::hash{}(static_cast(id.outputDtype)); - return h1 ^ h2 ^ h3 ^ h4 ^ h5 ^ h6; - } -}; - -template -class GemmPluginProfiler -{ -public: - // Map for single GEMM for different Ms (GEMM dimension) to the best config for particular M - using MProfileMap = std::unordered_map>; - using MProfileMapPtr = std::shared_ptr; - - // requires exclusive ownership to write to *this - using reader_lock = std::unique_lock; - // requires shared ownership to read from other - using writer_lock = std::shared_lock; - - // Struct of continuing map if GEMMs to the best profiles for different Ms - struct MNKProfileMap - { - // Mutex guarding map - std::shared_timed_mutex mutex; - // Map from GEMM Id to profile for particular GEMM - std::unordered_map profileMap; - - bool existsMProfileMap(GemmIdType const& id) - { - auto const iter = profileMap.find(id); - return iter != profileMap.end(); - } - - void createMProfileMap(GemmIdType const& id) - { - profileMap[id] = std::make_shared(); - } - - MProfileMapPtr getMProfileMap(GemmIdType const& id) - { - auto const iter = profileMap.find(id); - if (iter == profileMap.end()) - { - std::ostringstream msg; - msg << "Cannot find ID (" << id << ") in the profile map. Abort."; - TLLM_THROW(msg.str()); - } - return iter->second; - } - }; - - using MNKProfileMapPtr = std::shared_ptr; - - GemmPluginProfiler(); - - virtual ~GemmPluginProfiler() = default; - - void serialize(char*& buffer, GemmIdType const& gemmId) const; - - void deserialize(char const*& data, GemmDims& dims, GemmIdType const& gemmId); - size_t getSerializationSize(GemmIdType const& gemmId) const; - - void profileTactics(RunnerPtr const& runner, nvinfer1::DataType const& type, GemmDims const& dims, - GemmIdType const& gemmId, bool hasWeightOnlyCudaKernel = false); - - void setSelectionTactics(MNKProfileMapPtr const& map) - { - mMNKProfileMap = map; - } - - void setTmpWorkspaceSizeInBytes(size_t bytes) - { - mTmpWorkspaceSizeInBytes = bytes; - } - - void setSkip(bool skip) - { - mSkip = mSkip || skip; - } - - std::optional getBestConfig(int m, GemmIdType const& gemmId) const; - - virtual int getMaxProfileM() const; - -protected: - virtual void runTactic(int m, int n, int k, Config const& tactic, char* workspace, cudaStream_t const& stream) = 0; - - virtual void computeTmpSize(size_t maxM, size_t n, size_t k) = 0; - - virtual bool checkTactic(int m, int n, int k, Config const& tactic) const - { - return true; - } - - virtual std::vector getTactics(int m, int n, int k) const = 0; - - virtual void initTmpData(int m, int n, int k, char* workspace, size_t size, cudaStream_t stream); - -private: - void allocateTmpData(); - - void freeTmpData(); - - std::optional profileTacticsForProblem(int m, int n, int k, std::vector const& tactics); - - float profileTacticForProblem(int m, int n, int k, Config const& tactic); - - int nextPowerOfTwo(int v) const - { - --v; - v |= v >> 1; - v |= v >> 2; - v |= v >> 4; - v |= v >> 8; - v |= v >> 16; - return ++v; - } - -protected: - RunnerPtr mRunner{nullptr}; - - nvinfer1::DataType mType{}; - -private: - MNKProfileMapPtr mMNKProfileMap{}; - - size_t mTmpWorkspaceSizeInBytes{0}; - - char* mWorkspaceTmp{nullptr}; - - cudaStream_t mStream; - - GemmDims mDims{}; - - bool mSkip{false}; -}; - -template -class GemmPluginProfilerManager -{ -public: - using MNKProfileMap = typename GemmPluginProfilerType::MNKProfileMap; - using MNKProfileMapPtr = typename GemmPluginProfilerType::MNKProfileMapPtr; - using GemmPluginProfilerPtr = std::shared_ptr; - - GemmPluginProfilerManager() - { - mMNKProfileMap = std::make_shared(); - } - - GemmPluginProfilerPtr createGemmPluginProfiler(bool inference, bool skip = false) - { - auto profiler = std::make_shared(); - profiler->setSkip(skip); - // If the profiler is created during the engine build, - // mMNKProfileMap is shared between different profilers to minimize the time spent on the profiling - // and do not repeat profiling for the GEMMs of the same shape. - if (!inference) - { - profiler->setSelectionTactics(mMNKProfileMap); - } - return profiler; - } - -private: - MNKProfileMapPtr mMNKProfileMap{}; -}; - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/common/plugin.cpp b/cpp/tensorrt_llm/plugins/common/plugin.cpp deleted file mode 100644 index 82c8bf93b13c..000000000000 --- a/cpp/tensorrt_llm/plugins/common/plugin.cpp +++ /dev/null @@ -1,124 +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. - */ -#include "tensorrt_llm/plugins/common/plugin.h" -#include "tensorrt_llm/runtime/utils/mpiUtils.h" - -#include "checkMacrosPlugin.h" -#include -#include - -#ifdef _MSC_VER -#define FN_NAME __FUNCTION__ -#else -#define FN_NAME __func__ -#endif - -PluginFieldParser::PluginFieldParser(int32_t nbFields, nvinfer1::PluginField const* fields) - : mFields{fields} -{ - for (int32_t i = 0; i < nbFields; i++) - { - mMap.emplace(fields[i].name, PluginFieldParser::Record{i}); - } -} - -PluginFieldParser::~PluginFieldParser() -{ - for (auto const& [name, record] : mMap) - { - if (!record.retrieved) - { - std::stringstream ss; - ss << "unused plugin field with name: " << name; - tensorrt_llm::plugins::logError(ss.str().c_str(), __FILE__, FN_NAME, __LINE__); - } - } -} - -template -nvinfer1::PluginFieldType toFieldType(); -#define SPECIALIZE_TO_FIELD_TYPE(T, type) \ - template <> \ - nvinfer1::PluginFieldType toFieldType() \ - { \ - return nvinfer1::PluginFieldType::type; \ - } -SPECIALIZE_TO_FIELD_TYPE(half, kFLOAT16) -SPECIALIZE_TO_FIELD_TYPE(float, kFLOAT32) -SPECIALIZE_TO_FIELD_TYPE(double, kFLOAT64) -SPECIALIZE_TO_FIELD_TYPE(int8_t, kINT8) -SPECIALIZE_TO_FIELD_TYPE(int16_t, kINT16) -SPECIALIZE_TO_FIELD_TYPE(int32_t, kINT32) -SPECIALIZE_TO_FIELD_TYPE(char, kCHAR) -SPECIALIZE_TO_FIELD_TYPE(nvinfer1::Dims, kDIMS) -SPECIALIZE_TO_FIELD_TYPE(void, kUNKNOWN) -#undef SPECIALIZE_TO_FIELD_TYPE - -template -std::optional PluginFieldParser::getScalar(std::string_view const& name) -{ - auto const iter = mMap.find(name); - if (iter == mMap.end()) - { - return std::nullopt; - } - auto& record = mMap.at(name); - auto const& f = mFields[record.index]; - TLLM_CHECK(toFieldType() == f.type && f.length == 1); - record.retrieved = true; - return std::optional{*static_cast(f.data)}; -} - -#define INSTANTIATE_PluginFieldParser_getScalar(T) \ - template std::optional PluginFieldParser::getScalar(std::string_view const&) -INSTANTIATE_PluginFieldParser_getScalar(half); -INSTANTIATE_PluginFieldParser_getScalar(float); -INSTANTIATE_PluginFieldParser_getScalar(double); -INSTANTIATE_PluginFieldParser_getScalar(int8_t); -INSTANTIATE_PluginFieldParser_getScalar(int16_t); -INSTANTIATE_PluginFieldParser_getScalar(int32_t); -INSTANTIATE_PluginFieldParser_getScalar(char); -INSTANTIATE_PluginFieldParser_getScalar(nvinfer1::Dims); -#undef INSTANTIATE_PluginFieldParser_getScalar - -template -std::optional> PluginFieldParser::getSet(std::string_view const& name) -{ - auto const iter = mMap.find(name); - if (iter == mMap.end()) - { - return std::nullopt; - } - auto& record = mMap.at(name); - auto const& f = mFields[record.index]; - TLLM_CHECK(toFieldType() == f.type); - std::set group; - auto const* r = static_cast(f.data); - for (int j = 0; j < f.length; ++j) - { - group.insert(*r); - ++r; - } - - record.retrieved = true; - return std::optional{group}; -} - -#define INSTANTIATE_PluginFieldParser_getVector(T) \ - template std::optional> PluginFieldParser::getSet(std::string_view const&) -INSTANTIATE_PluginFieldParser_getVector(int32_t); -#undef INSTANTIATE_PluginFieldParser_getVector diff --git a/cpp/tensorrt_llm/plugins/common/plugin.h b/cpp/tensorrt_llm/plugins/common/plugin.h deleted file mode 100644 index a7febe4cc13d..000000000000 --- a/cpp/tensorrt_llm/plugins/common/plugin.h +++ /dev/null @@ -1,143 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. - */ - -#pragma once - -#include "tensorrt_llm/common/opUtils.h" -#include "tensorrt_llm/plugins/api/tllmPlugin.h" -#include "tensorrt_llm/plugins/common/checkMacrosPlugin.h" - -#include - -#include -#include -#include -#include -#include -#include -#include - -namespace tensorrt_llm::plugins -{ - -using namespace tensorrt_llm::common::op; - -class BasePlugin : public nvinfer1::IPluginV2DynamicExt -{ -public: - void setPluginNamespace(char const* libNamespace) noexcept override - { - mNamespace = libNamespace; - } - - [[nodiscard]] char const* getPluginNamespace() const noexcept override - { - return mNamespace.c_str(); - } - -protected: - std::string mNamespace{api::kDefaultNamespace}; -}; - -class BasePluginV3 : public nvinfer1::IPluginV3, - public nvinfer1::IPluginV3OneCore, - public nvinfer1::IPluginV3OneBuild, - public nvinfer1::IPluginV3OneRuntime -{ -public: - void setPluginNamespace(char const* libNamespace) noexcept - { - mNamespace = libNamespace; - } - - [[nodiscard]] char const* getPluginNamespace() const noexcept override - { - return mNamespace.c_str(); - } - -protected: - std::string mNamespace{api::kDefaultNamespace}; -}; - -class BaseCreator : public nvinfer1::IPluginCreator -{ -public: - void setPluginNamespace(char const* libNamespace) noexcept override - { - mNamespace = libNamespace; - } - - [[nodiscard]] char const* getPluginNamespace() const noexcept override - { - return mNamespace.c_str(); - } - -protected: - std::string mNamespace{api::kDefaultNamespace}; -}; - -class BaseCreatorV3 : public nvinfer1::IPluginCreatorV3One -{ -public: - void setPluginNamespace(char const* libNamespace) noexcept - { - mNamespace = libNamespace; - } - - [[nodiscard]] char const* getPluginNamespace() const noexcept override - { - return mNamespace.c_str(); - } - -protected: - std::string mNamespace{api::kDefaultNamespace}; -}; - -} // namespace tensorrt_llm::plugins - -// Init with O(n) and retrieve with O(1) -class PluginFieldParser -{ -public: - // field array must remain valid when calling getScalar() later. - PluginFieldParser(int32_t nbFields, nvinfer1::PluginField const* fields); - // delete to remind accidental mis-use (copy) which may result in false-alarm warnings about unused fields. - PluginFieldParser(PluginFieldParser const&) = delete; - PluginFieldParser& operator=(PluginFieldParser const&) = delete; - // check if all fields are retrieved and emit warning if some of them are not. - ~PluginFieldParser(); - template - std::optional getScalar(std::string_view const& name); - template - std::optional> getSet(std::string_view const& name); - -private: - nvinfer1::PluginField const* mFields; - - struct Record - { - Record(int32_t idx) - : index{idx} - { - } - - int32_t const index; - bool retrieved{false}; - }; - - std::unordered_map mMap; -}; diff --git a/cpp/tensorrt_llm/plugins/common/pluginUtils.h b/cpp/tensorrt_llm/plugins/common/pluginUtils.h deleted file mode 100644 index ee3e59d57c6d..000000000000 --- a/cpp/tensorrt_llm/plugins/common/pluginUtils.h +++ /dev/null @@ -1,78 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-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. - */ - -#pragma once - -#include - -#include "tensorrt_llm/common/logger.h" - -namespace tensorrt_llm::plugins::utils -{ -using DimType64 = int64_t; - -inline DimType64 computeMDimension(bool transA, nvinfer1::Dims const& dims) -{ - DimType64 M{1}; - if (transA) - { - for (int i = dims.nbDims - 1; i > 0; --i) - { - M *= dims.d[i]; - } - } - else - { - for (int i = 0; i < dims.nbDims - 1; ++i) - { - M *= dims.d[i]; - } - } - return M; -} - -inline DimType64 computeNDimension(bool transB, nvinfer1::Dims const& dims) -{ - DimType64 N{1}; - if (transB) - { - for (int32_t i = 0; i < dims.nbDims - 1; ++i) - { - N *= dims.d[i]; - } - } - else - { - for (int32_t i = dims.nbDims - 1; i > 0; --i) - { - N *= dims.d[i]; - } - } - return N; -} - -inline std::int32_t logErrorReturn0(char const* variable) -{ - TLLM_LOG_ERROR("Value of %s is out of range for int32_t", variable); - return 0; -} - -#define TLLM_INT32_CAST(value) \ - ((value > 0x7FFFFFFFLL || value < -0x80000000LL) ? tensorrt_llm::plugins::utils::logErrorReturn0(#value) \ - : static_cast(value)) - -} // namespace tensorrt_llm::plugins::utils diff --git a/cpp/tensorrt_llm/plugins/cpSplitPlugin/CMakeLists.txt b/cpp/tensorrt_llm/plugins/cpSplitPlugin/CMakeLists.txt deleted file mode 100644 index 86876224fccd..000000000000 --- a/cpp/tensorrt_llm/plugins/cpSplitPlugin/CMakeLists.txt +++ /dev/null @@ -1,21 +0,0 @@ -# -# SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. -# -file(GLOB SRCS *.cpp) -set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS}) -set(PLUGIN_SOURCES - ${PLUGIN_SOURCES} - PARENT_SCOPE) diff --git a/cpp/tensorrt_llm/plugins/cpSplitPlugin/cpSplitPlugin.cpp b/cpp/tensorrt_llm/plugins/cpSplitPlugin/cpSplitPlugin.cpp deleted file mode 100644 index 221d5dac2da1..000000000000 --- a/cpp/tensorrt_llm/plugins/cpSplitPlugin/cpSplitPlugin.cpp +++ /dev/null @@ -1,356 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-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. - */ - -#include - -#include "cpSplitPlugin.h" - -using namespace nvinfer1; -using namespace tensorrt_llm::common; -using tensorrt_llm::plugins::CpSplitPluginCreator; -using tensorrt_llm::plugins::CpSplitPlugin; - -static char const* CPSPLIT_PLUGIN_VERSION{"1"}; -static char const* CPSPLIT_PLUGIN_NAME{"CpSplit"}; -PluginFieldCollection CpSplitPluginCreator::mFC{}; -std::vector CpSplitPluginCreator::mPluginAttributes; - -CpSplitPlugin::CpSplitPlugin() -{ - initFieldsToSerialize(); -} - -CpSplitPlugin::CpSplitPlugin(int cpSize, int cpRank) - : mCpSize(cpSize) - , mCpRank(cpRank) -{ - initFieldsToSerialize(); -} - -void CpSplitPlugin::initFieldsToSerialize() -{ - mDataToSerialize.clear(); - mDataToSerialize.emplace_back(PluginField("cp_size", &mCpSize, PluginFieldType::kINT32, 1)); - mDataToSerialize.emplace_back(PluginField("cp_rank", &mCpRank, PluginFieldType::kINT32, 1)); - mFCToSerialize.nbFields = mDataToSerialize.size(); - mFCToSerialize.fields = mDataToSerialize.data(); -} - -// IPluginV3 methods -nvinfer1::IPluginCapability* CpSplitPlugin::getCapabilityInterface(nvinfer1::PluginCapabilityType type) noexcept -{ - switch (type) - { - case PluginCapabilityType::kBUILD: return static_cast(this); - case PluginCapabilityType::kRUNTIME: return static_cast(this); - case PluginCapabilityType::kCORE: return static_cast(this); - } - return nullptr; -} - -nvinfer1::IPluginV3* CpSplitPlugin::clone() noexcept -{ - std::unique_ptr plugin{std::make_unique(*this)}; - plugin->setPluginNamespace(mNamespace.c_str()); - plugin->initFieldsToSerialize(); - return plugin.release(); -} - -// IPluginV3OneCore methods -char const* CpSplitPlugin::getPluginName() const noexcept -{ - return CPSPLIT_PLUGIN_NAME; -} - -char const* CpSplitPlugin::getPluginVersion() const noexcept -{ - return CPSPLIT_PLUGIN_VERSION; -} - -// IPluginV3OneBuild methods -int32_t CpSplitPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int32_t nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int32_t nbOutputs) noexcept -{ - return 0; -} - -int32_t CpSplitPlugin::getOutputDataTypes( - DataType* outputTypes, int32_t nbOutputs, DataType const* inputTypes, int32_t nbInputs) const noexcept -{ - outputTypes[0] = inputTypes[0]; - outputTypes[1] = DataType::kINT32; - outputTypes[2] = DataType::kINT32; - return 0; -} - -int32_t CpSplitPlugin::getOutputShapes(DimsExprs const* inputs, int32_t nbInputs, DimsExprs const* shapeInputs, - int32_t nbShapeInputs, DimsExprs* outputs, int32_t nbOutputs, IExprBuilder& exprBuilder) noexcept -{ - outputs[0].nbDims = 1; - - auto cpSize = exprBuilder.constant(mCpSize); - auto upper = inputs[0].d[0]; - auto opt = exprBuilder.operation(DimensionOperation::kCEIL_DIV, *upper, *cpSize); - outputs[0].d[0] = exprBuilder.declareSizeTensor(1, *opt, *upper); - - // We must have such an output size tensor (with dim == 0) to notify the shape of output tensor above - outputs[1].nbDims = 0; - outputs[2].nbDims = 1; - outputs[2].d[0] = upper; - return 0; -} - -bool CpSplitPlugin::supportsFormatCombination( - int32_t pos, nvinfer1::DynamicPluginTensorDesc const* inOut, int32_t nbInputs, int32_t nbOutputs) noexcept -{ - if (pos == IdxEntry::INPUT_IDS) - { - return ((inOut[pos].desc.type == DataType::kINT32) && (inOut[pos].desc.format == TensorFormat::kLINEAR)); - } - else if (pos == IdxEntry::REQUEST_TYPES || pos == IdxEntry::HOST_CONTEXT_LENGTH) - { - return inOut[pos].desc.type == DataType::kINT32; - } - else - { - return ((inOut[pos].desc.type == DataType::kINT32) && (inOut[pos].desc.format == TensorFormat::kLINEAR)); - } - return false; -} - -int32_t CpSplitPlugin::getNbOutputs() const noexcept -{ - return 3; -} - -size_t CpSplitPlugin::getWorkspaceSize(nvinfer1::DynamicPluginTensorDesc const* inputs, int32_t nbInputs, - nvinfer1::DynamicPluginTensorDesc const* outputs, int32_t nbOutputs) const noexcept -{ - return 0; -} - -int32_t CpSplitPlugin::getValidTactics(int32_t* tactics, int32_t nbTactics) noexcept -{ - return 0; -} - -int32_t CpSplitPlugin::getNbTactics() noexcept -{ - return 0; -} - -char const* CpSplitPlugin::getTimingCacheID() noexcept -{ - return nullptr; -} - -int32_t CpSplitPlugin::getFormatCombinationLimit() noexcept -{ - return 1; -} - -char const* CpSplitPlugin::getMetadataString() noexcept -{ - return nullptr; -} - -// IPluginV3OneRuntime methods -int32_t CpSplitPlugin::setTactic(int32_t tactic) noexcept -{ - return 0; -} - -int32_t CpSplitPlugin::onShapeChange(nvinfer1::PluginTensorDesc const* in, int32_t nbInputs, - nvinfer1::PluginTensorDesc const* out, int32_t nbOutputs) noexcept -{ - return 0; -} - -int32_t CpSplitPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) noexcept -{ - // inputs - // @param inputIds [tokenNum] - // @param host_request_types [batchSize]: Tensor = None (On CPU) - // The tensor on the host that indicates if a request is in context or - // generation phase. Its shape is [batch_size]. See Inflight Batching - // in docs/gpt_attention.md, - // @param host_context_lengths [batchSize]: Tensor = None (On CPU) - // A host tensor that contains the lengths of the different inputs - // outputs - // @param outputIds [tokenNum spiltted by cp] - // @param outputLength scalar - // @param joinIdx [tokenNum] - - int64_t tokenNum = 1; - for (int i = 0; i < inputDesc[0].dims.nbDims; ++i) - { - tokenNum *= inputDesc[0].dims.d[i]; - } - - RequestType const* reqTypes = static_cast(inputs[IdxEntry::REQUEST_TYPES]); - int32_t const* hContextLengths = static_cast(inputs[IdxEntry::HOST_CONTEXT_LENGTH]); - int const* inputIds = reinterpret_cast(inputs[IdxEntry::INPUT_IDS]); - int* outputIds = reinterpret_cast(outputs[0]); - int32_t* outputLength = reinterpret_cast(outputs[1]); - int32_t* outputJoinIdx = reinterpret_cast(outputs[2]); - - int32_t const nbSeq = inputDesc[IdxEntry::HOST_CONTEXT_LENGTH].dims.d[0]; - - int32_t* hInputs = new int[inputDesc[IdxEntry::INPUT_IDS].dims.d[0]]; - int32_t* hOutputs = new int[inputDesc[IdxEntry::INPUT_IDS].dims.d[0]]; - int32_t* hOutputJoinIdx = new int[inputDesc[IdxEntry::INPUT_IDS].dims.d[0]]; - cudaMemcpyAsync( - hInputs, inputIds, sizeof(int32_t) * inputDesc[IdxEntry::INPUT_IDS].dims.d[0], cudaMemcpyDeviceToHost, stream); - sync_check_cuda_error(stream); - - int32_t inputIdx = 0; - int32_t outputIdx = 0; - for (int32_t seqIdx = 0; seqIdx < nbSeq; ++seqIdx) - { - if (reqTypes[seqIdx] == RequestType::kCONTEXT) - { - auto const& ctxLength = hContextLengths[seqIdx]; - int32_t partialAverageLength = (ctxLength + mCpSize - 1) / mCpSize; - int32_t partialLength - = mCpRank == mCpSize - 1 ? ctxLength - partialAverageLength * (mCpSize - 1) : partialAverageLength; - for (int i = 0; i < partialLength; i++) - { - hOutputs[outputIdx + i] = hInputs[inputIdx + partialAverageLength * mCpRank + i]; - } - inputIdx += ctxLength; - outputIdx += partialAverageLength; - } - else if (reqTypes[seqIdx] == RequestType::kGENERATION) - { - auto const& genLength = nbSeq - seqIdx; - int32_t partialAverageLength = (genLength + mCpSize - 1) / mCpSize; - int32_t partialLength - = mCpRank == mCpSize - 1 ? genLength - partialAverageLength * (mCpSize - 1) : partialAverageLength; - for (int i = 0; i < partialLength; i++) - { - hOutputs[outputIdx + i] = hInputs[inputIdx + partialAverageLength * mCpRank + i]; - } - outputIdx += partialAverageLength; - break; - } - } - int32_t hOutputLength = outputIdx; - inputIdx = 0; - outputIdx = 0; - for (int32_t seqIdx = 0; seqIdx < nbSeq; ++seqIdx) - { - if (reqTypes[seqIdx] == RequestType::kCONTEXT) - { - auto const& ctxLength = hContextLengths[seqIdx]; - int32_t partialAverageLength = (ctxLength + mCpSize - 1) / mCpSize; - for (int32_t idx = 0; idx < ctxLength; ++idx) - { - hOutputJoinIdx[inputIdx + idx] - = idx % partialAverageLength + idx / partialAverageLength * hOutputLength + outputIdx; - } - inputIdx += ctxLength; - outputIdx += partialAverageLength; - } - else if (reqTypes[seqIdx] == RequestType::kGENERATION) - { - auto const& genLength = nbSeq - seqIdx; - int32_t partialAverageLength = (genLength + mCpSize - 1) / mCpSize; - for (int32_t idx = 0; idx < genLength; ++idx) - { - hOutputJoinIdx[inputIdx + idx] - = idx % partialAverageLength + idx / partialAverageLength * hOutputLength + outputIdx; - } - break; - } - } - cudaMemcpyAsync(outputIds, hOutputs, sizeof(int32_t) * hOutputLength, cudaMemcpyHostToDevice, stream); - cudaMemcpyAsync(outputLength, &hOutputLength, sizeof(int32_t), cudaMemcpyHostToDevice, stream); - cudaMemcpyAsync(outputJoinIdx, hOutputJoinIdx, sizeof(int32_t) * tokenNum, cudaMemcpyHostToDevice, stream); - sync_check_cuda_error(stream); - return 0; -} - -nvinfer1::IPluginV3* CpSplitPlugin::attachToContext(nvinfer1::IPluginResourceContext* context) noexcept -{ - return clone(); -} - -nvinfer1::PluginFieldCollection const* CpSplitPlugin::getFieldsToSerialize() noexcept -{ - return &mFCToSerialize; -} - -CpSplitPluginCreator::CpSplitPluginCreator() -{ - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - mPluginAttributes.emplace_back(PluginField("cp_size", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("cp_rank", nullptr, PluginFieldType::kINT32)); - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -char const* CpSplitPluginCreator::getPluginName() const noexcept -{ - return CPSPLIT_PLUGIN_NAME; -} - -char const* CpSplitPluginCreator::getPluginVersion() const noexcept -{ - return CPSPLIT_PLUGIN_VERSION; -} - -nvinfer1::PluginFieldCollection const* CpSplitPluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -nvinfer1::IPluginV3* CpSplitPluginCreator::createPlugin( - char const* name, nvinfer1::PluginFieldCollection const* fc, nvinfer1::TensorRTPhase phase) noexcept -{ - PluginField const* fields = fc->fields; - int cp_size{}; - int cp_rank{}; - // Read configurations from each fields - for (int i = 0; i < fc->nbFields; ++i) - { - char const* attrName = fields[i].name; - if (!strcmp(attrName, "cp_size")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - cp_size = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "cp_rank")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - cp_rank = static_cast(*(static_cast(fields[i].data))); - } - } - try - { - auto* obj = new CpSplitPlugin(cp_size, cp_rank); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} diff --git a/cpp/tensorrt_llm/plugins/cpSplitPlugin/cpSplitPlugin.h b/cpp/tensorrt_llm/plugins/cpSplitPlugin/cpSplitPlugin.h deleted file mode 100644 index 1dc8c15b355a..000000000000 --- a/cpp/tensorrt_llm/plugins/cpSplitPlugin/cpSplitPlugin.h +++ /dev/null @@ -1,112 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-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. - */ -#pragma once - -#include "tensorrt_llm/plugins/common/plugin.h" -#include -#include -#include -#include - -namespace tensorrt_llm::plugins -{ - -class CpSplitPlugin : public BasePluginV3 -{ -public: - CpSplitPlugin(); - CpSplitPlugin(int cpSize, int cpRank); - CpSplitPlugin(CpSplitPlugin const& p) = default; - void initFieldsToSerialize(); - - // IPluginV3 methods - nvinfer1::IPluginCapability* getCapabilityInterface(nvinfer1::PluginCapabilityType type) noexcept override; - nvinfer1::IPluginV3* clone() noexcept override; - - // IPluginV3OneCore methods - char const* getPluginName() const noexcept override; - char const* getPluginVersion() const noexcept override; - - // IPluginV3OneBuild methods - int32_t configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int32_t nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int32_t nbOutputs) noexcept override; // nochange - int32_t getOutputDataTypes(nvinfer1::DataType* outputTypes, int32_t nbOutputs, nvinfer1::DataType const* inputTypes, - int32_t nbInputs) const noexcept override; // fixed - int32_t getOutputShapes(nvinfer1::DimsExprs const* inputs, int32_t nbInputs, nvinfer1::DimsExprs const* shapeInputs, - int32_t nbShapeInputs, nvinfer1::DimsExprs* outputs, int32_t nbOutputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; // fixed - bool supportsFormatCombination(int32_t pos, nvinfer1::DynamicPluginTensorDesc const* inOut, int32_t nbInputs, - int32_t nbOutputs) noexcept override; // fixed - int32_t getNbOutputs() const noexcept override; // fixed - size_t getWorkspaceSize(nvinfer1::DynamicPluginTensorDesc const* inputs, int32_t nbInputs, - nvinfer1::DynamicPluginTensorDesc const* outputs, int32_t nbOutputs) const noexcept override; // fixed - int32_t getValidTactics(int32_t* tactics, int32_t nbTactics) noexcept override; - int32_t getNbTactics() noexcept override; - char const* getTimingCacheID() noexcept override; - int32_t getFormatCombinationLimit() noexcept override; - char const* getMetadataString() noexcept override; - - // IPluginV3OneRuntime methods - int32_t setTactic(int32_t tactic) noexcept override; - int32_t onShapeChange(nvinfer1::PluginTensorDesc const* in, int32_t nbInputs, nvinfer1::PluginTensorDesc const* out, - int32_t nbOutputs) noexcept override; - int32_t enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) noexcept override; // fixed - nvinfer1::IPluginV3* attachToContext(nvinfer1::IPluginResourceContext* context) noexcept override; - nvinfer1::PluginFieldCollection const* getFieldsToSerialize() noexcept override; - -private: - int mCpSize; - int mCpRank; - std::vector mDataToSerialize; - nvinfer1::PluginFieldCollection mFCToSerialize; - - enum IdxEntry - { - INPUT_IDS, - REQUEST_TYPES, - HOST_CONTEXT_LENGTH, - }; - - enum class RequestType : int32_t - { - kCONTEXT = 0, - kGENERATION = 1 - }; -}; - -class CpSplitPluginCreator : public BaseCreatorV3 -{ -public: - CpSplitPluginCreator(); - - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV3* createPlugin( - char const* name, nvinfer1::PluginFieldCollection const* fc, nvinfer1::TensorRTPhase phase) noexcept override; - -private: - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; -}; - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/cudaStreamPlugin/CMakeLists.txt b/cpp/tensorrt_llm/plugins/cudaStreamPlugin/CMakeLists.txt deleted file mode 100644 index 86876224fccd..000000000000 --- a/cpp/tensorrt_llm/plugins/cudaStreamPlugin/CMakeLists.txt +++ /dev/null @@ -1,21 +0,0 @@ -# -# SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. -# -file(GLOB SRCS *.cpp) -set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS}) -set(PLUGIN_SOURCES - ${PLUGIN_SOURCES} - PARENT_SCOPE) diff --git a/cpp/tensorrt_llm/plugins/cudaStreamPlugin/cudaStreamPlugin.cpp b/cpp/tensorrt_llm/plugins/cudaStreamPlugin/cudaStreamPlugin.cpp deleted file mode 100644 index 802e828c9250..000000000000 --- a/cpp/tensorrt_llm/plugins/cudaStreamPlugin/cudaStreamPlugin.cpp +++ /dev/null @@ -1,293 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. - */ - -#include "cudaStreamPlugin.h" -#include "tensorrt_llm/runtime/iBuffer.h" - -#include - -using namespace nvinfer1; -using tensorrt_llm::plugins::CudaStreamPluginCreator; -using tensorrt_llm::plugins::CudaStreamPlugin; - -static char const* CUDA_STREAM_PLUGIN_VERSION{"1"}; -static char const* CUDA_STREAM_PLUGIN_NAME{"CudaStream"}; -PluginFieldCollection CudaStreamPluginCreator::mFC{}; -std::vector CudaStreamPluginCreator::mPluginAttributes; - -CudaStreamPlugin::CudaStreamPlugin(int sideStreamId, int nbInputs, nvinfer1::DataType type) - : mSideStreamId(sideStreamId) - , mNbInputs(nbInputs) - , mType(type) -{ - init(); -} - -CudaStreamPlugin::CudaStreamPlugin(void const* data, size_t length) -{ - char const *d = reinterpret_cast(data), *a = d; - read(d, mSideStreamId); - read(d, mNbInputs); - read(d, mType); - - init(); - - TLLM_CHECK_WITH_INFO(d == a + length, - "Expected length (%d) != real length (%d). This is often " - "caused by using different TensorRT LLM version to build " - "engine and run engine.", - (int) length, (int) (d - a)); -} - -CudaStreamPlugin::CudaStreamPlugin(CudaStreamPlugin const& other) - : mSideStreamId(other.mSideStreamId) - , mNbInputs(other.mNbInputs) - , mType(other.mType) -{ - init(); -} - -void CudaStreamPlugin::init() -{ - mSideStreamPtr = nullptr; -} - -// IPluginV2DynamicExt Methods -nvinfer1::IPluginV2DynamicExt* CudaStreamPlugin::clone() const noexcept -{ - auto* plugin = new CudaStreamPlugin(*this); - plugin->setPluginNamespace(mNamespace.c_str()); - return plugin; -} - -nvinfer1::DimsExprs CudaStreamPlugin::getOutputDimensions( - int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept -{ - assert(outputIndex == 0); - return inputs[outputIndex]; -} - -bool CudaStreamPlugin::supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept -{ - TLLM_CHECK_WITH_INFO(nbInputs == mNbInputs, "CudaStreamPlugin only accepts mNbInputs inputs"); - TLLM_CHECK_WITH_INFO(nbOutputs == 1, "CudaStreamPlugin only accepts 1 output"); - - auto const& desc = inOut[pos]; - if (desc.format != TensorFormat::kLINEAR) - { - return false; - } - - if (pos > 0 && pos < nbInputs) - { - return true; - } - return desc.type == mType; -} - -void CudaStreamPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept -{ -} - -size_t CudaStreamPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - return 0; -} - -int CudaStreamPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept -{ - if (!mSideStreamPtr) - { - auto const resource_name = nvinfer1::pluginInternal::SideStream::getResourceKey(mSideStreamId); - nvinfer1::pluginInternal::SideStream side_stream{}; - mSideStreamPtr = reinterpret_cast( - getPluginRegistry()->acquirePluginResource(resource_name.c_str(), &side_stream)); - } - mSideStreamPtr->waitSideStreamOnMainStream(stream); - size_t count = 1; - for (int i = 0; i < inputDesc[0].dims.nbDims; ++i) - { - count *= inputDesc[0].dims.d[i]; - } - count *= tensorrt_llm::runtime::BufferDataType(inputDesc[0].type).getSize(); - TLLM_CUDA_CHECK(cudaMemcpyAsync(outputs[0], inputs[0], count, cudaMemcpyDeviceToDevice, stream)); - - return 0; -} - -// IPluginV2Ext Methods -nvinfer1::DataType CudaStreamPlugin::getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept -{ - TLLM_CHECK(index == 0); - return mType; -} - -// IPluginV2 Methods - -char const* CudaStreamPlugin::getPluginType() const noexcept -{ - return CUDA_STREAM_PLUGIN_NAME; -} - -char const* CudaStreamPlugin::getPluginVersion() const noexcept -{ - return CUDA_STREAM_PLUGIN_VERSION; -} - -int CudaStreamPlugin::getNbOutputs() const noexcept -{ - return 1; -} - -int CudaStreamPlugin::initialize() noexcept -{ - return 0; -} - -void CudaStreamPlugin::terminate() noexcept -{ - if (mSideStreamPtr) - { - auto const resource_name = nvinfer1::pluginInternal::SideStream::getResourceKey(mSideStreamId); - getPluginRegistry()->releasePluginResource(resource_name.c_str()); - mSideStreamPtr = nullptr; - } -} - -size_t CudaStreamPlugin::getSerializationSize() const noexcept -{ - return sizeof(mSideStreamId) + sizeof(mNbInputs) + sizeof(mType); -} - -void CudaStreamPlugin::serialize(void* buffer) const noexcept -{ - char *d = static_cast(buffer), *a = d; - write(d, mSideStreamId); - write(d, mNbInputs); - write(d, mType); - TLLM_CHECK(d == a + getSerializationSize()); -} - -void CudaStreamPlugin::destroy() noexcept -{ - delete this; -} - -/////////////// - -CudaStreamPluginCreator::CudaStreamPluginCreator() -{ - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - mPluginAttributes.emplace_back(PluginField("side_stream_id", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("num_inputs", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("type_id", nullptr, PluginFieldType::kINT32)); - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -char const* CudaStreamPluginCreator::getPluginName() const noexcept -{ - return CUDA_STREAM_PLUGIN_NAME; -} - -char const* CudaStreamPluginCreator::getPluginVersion() const noexcept -{ - return CUDA_STREAM_PLUGIN_VERSION; -} - -PluginFieldCollection const* CudaStreamPluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -IPluginV2* CudaStreamPluginCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept -{ - PluginField const* fields = fc->fields; - int sideStreamId; - int nbInputs; - int type; - - // Read configurations from each fields - struct MapPair - { - char const* key; - int& field; - bool optional = false; - bool set = false; - }; - - std::array input_map{ - MapPair{"side_stream_id", std::ref(sideStreamId)}, - MapPair{"num_inputs", std::ref(nbInputs)}, - MapPair{"type_id", std::ref(type)}, - }; - for (int i = 0; i < fc->nbFields; ++i) - { - char const* attrName = fields[i].name; - for (auto& item : input_map) - { - if (!strcmp(item.key, attrName)) - { - TLLM_CHECK(fields[i].type == nvinfer1::PluginFieldType::kINT32); - TLLM_CHECK_WITH_INFO(!item.set, "Parameter %s was set twice", item.key); - item.field = static_cast(*(static_cast(fields[i].data))); - item.set = true; - } - } - } - - for (auto& item : input_map) - { - TLLM_CHECK_WITH_INFO(item.set || item.optional, "Parameter %s is required but not set", item.key); - } - - try - { - auto* obj = new CudaStreamPlugin(sideStreamId, nbInputs, static_cast(type)); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -IPluginV2* CudaStreamPluginCreator::deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept -{ - // This object will be deleted when the network is destroyed, which will - // call CudaStreamPlugin::destroy() - try - { - auto* obj = new CudaStreamPlugin(serialData, serialLength); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} diff --git a/cpp/tensorrt_llm/plugins/cudaStreamPlugin/cudaStreamPlugin.h b/cpp/tensorrt_llm/plugins/cudaStreamPlugin/cudaStreamPlugin.h deleted file mode 100644 index 5b78c3b873bb..000000000000 --- a/cpp/tensorrt_llm/plugins/cudaStreamPlugin/cudaStreamPlugin.h +++ /dev/null @@ -1,265 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. - */ -#pragma once - -#include "NvInferPlugin.h" -#include "tensorrt_llm/plugins/common/plugin.h" -#include "tensorrt_llm/runtime/cudaMemPool.h" -#include "tensorrt_llm/runtime/utils/debugUtils.h" -#include -#include -#include - -namespace nvinfer1 -{ -namespace pluginInternal -{ -class SideWorkspace -{ -public: - SideWorkspace(cudaStream_t stream) - : mWorkspaceSize{0} - , mWorkspacePtr{nullptr} - , mStream{stream} - { - } - - ~SideWorkspace() - { - if (mWorkspacePtr) - { - TLLM_CUDA_CHECK(cudaFreeAsync(mWorkspacePtr, mStream)); - } - } - - void* get(size_t workspaceSize) - { - if (mWorkspacePtr && mWorkspaceSize < workspaceSize) - { - TLLM_CUDA_CHECK(cudaFreeAsync(mWorkspacePtr, mStream)); - mWorkspacePtr = nullptr; - } - if (!mWorkspacePtr) - { - mWorkspaceSize = workspaceSize; - auto pool_ptr - = tensorrt_llm::runtime::CudaMemPool::getPrimaryPoolForDevice(tensorrt_llm::common::getDevice()); - TLLM_CUDA_CHECK(cudaMallocFromPoolAsync(&mWorkspacePtr, mWorkspaceSize, pool_ptr->getPool(), mStream)); - } - return mWorkspacePtr; - } - -private: - size_t mWorkspaceSize; - void* mWorkspacePtr; - cudaStream_t mStream; -}; - -class SideStream : public IPluginResource -{ -public: - SideStream(bool init = false) - : mStream{} - , mMainEvent{} - , mSideEvent{} - , mWorkspace{} - , mInit{init} - { - // The object passed to acquirePluginResource should use the default value init=false - if (init) - { - TLLM_CUDA_CHECK(cudaStreamCreate(&mStream)); - TLLM_CUDA_CHECK(cudaEventCreateWithFlags(&mMainEvent, cudaEventDisableTiming)); - TLLM_CUDA_CHECK(cudaEventCreateWithFlags(&mSideEvent, cudaEventDisableTiming)); - mWorkspace = std::make_shared(mStream); - } - } - - void free() - { - if (mInit) - { - mWorkspace = nullptr; - TLLM_CUDA_CHECK(cudaStreamSynchronize(mStream)); - TLLM_CUDA_CHECK(cudaStreamDestroy(mStream)); - TLLM_CUDA_CHECK(cudaEventDestroy(mMainEvent)); - TLLM_CUDA_CHECK(cudaEventDestroy(mSideEvent)); - mInit = false; - } - } - - int32_t release() noexcept override - { - try - { - free(); - } - catch (std::exception const& e) - { - return -1; - } - return 0; - } - - IPluginResource* clone() noexcept override - { - // An object is cloned only when calling acquirePluginResource for the first time for each key - std::unique_ptr cloned{}; - try - { - if (!mInit) - { - cloned = std::make_unique(/* init */ true); - } - else - { - return nullptr; - } - } - catch (std::exception const& e) - { - return nullptr; - } - return cloned.release(); - } - - ~SideStream() override - { - free(); - } - - void* getWorkspacePtr(size_t workspaceSize) - { - return mWorkspace->get(workspaceSize); - } - - cudaStream_t getStream() const - { - return mStream; - } - - void waitMainStreamOnSideStream(cudaStream_t const stream) const - { - TLLM_CUDA_CHECK(cudaEventRecord(mMainEvent, stream)); - TLLM_CUDA_CHECK(cudaStreamWaitEvent(mStream, mMainEvent)); - } - - void waitSideStreamOnMainStream(cudaStream_t const stream) const - { - TLLM_CUDA_CHECK(cudaEventRecord(mSideEvent, mStream)); - TLLM_CUDA_CHECK(cudaStreamWaitEvent(stream, mSideEvent)); - } - - void stallMainStream(char const* name, cudaStream_t const stream, std::optional delay = std::nullopt) const - { - tensorrt_llm::runtime::utils::stallStream(name, stream, delay); - } - - void stallSideStream(char const* name, std::optional delay = std::nullopt) const - { - tensorrt_llm::runtime::utils::stallStream(name, mStream, delay); - } - - static std::string getResourceKey(int const stream_id) - { - return "side_stream_" + std::to_string(stream_id); - } - -private: - cudaStream_t mStream; - cudaEvent_t mMainEvent; - cudaEvent_t mSideEvent; - std::shared_ptr mWorkspace; - bool mInit; -}; - -} // namespace pluginInternal -} // namespace nvinfer1 - -namespace tensorrt_llm::plugins -{ - -class CudaStreamPlugin : public BasePlugin -{ -public: - CudaStreamPlugin(int sideStreamId, int nbInputs, nvinfer1::DataType type); - - CudaStreamPlugin(void const* data, size_t length); - - CudaStreamPlugin(CudaStreamPlugin const&); - - void init(); - - ~CudaStreamPlugin() override = default; - - // IPluginV2DynamicExt Methods - nvinfer1::IPluginV2DynamicExt* clone() const noexcept override; - nvinfer1::DimsExprs getOutputDimensions(int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; - bool supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept override; - void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept override; - size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept override; - int enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - - // IPluginV2Ext Methods - nvinfer1::DataType getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept override; - - // IPluginV2 Methods - char const* getPluginType() const noexcept override; - char const* getPluginVersion() const noexcept override; - int getNbOutputs() const noexcept override; - int initialize() noexcept override; - void terminate() noexcept override; - size_t getSerializationSize() const noexcept override; - void serialize(void* buffer) const noexcept override; - void destroy() noexcept override; - -private: - const std::string mLayerName; - int mSideStreamId; - int mNbInputs; - nvinfer1::DataType mType; - nvinfer1::pluginInternal::SideStream* mSideStreamPtr; -}; - -class CudaStreamPluginCreator : public BaseCreator -{ -public: - CudaStreamPluginCreator(); - - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; - - nvinfer1::IPluginV2* deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept override; - -private: - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; -}; - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/cumsumLastDimPlugin/CMakeLists.txt b/cpp/tensorrt_llm/plugins/cumsumLastDimPlugin/CMakeLists.txt deleted file mode 100644 index ea25de075f34..000000000000 --- a/cpp/tensorrt_llm/plugins/cumsumLastDimPlugin/CMakeLists.txt +++ /dev/null @@ -1,22 +0,0 @@ -# -# SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. -# - -file(GLOB SRCS *.cpp) -set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS}) -set(PLUGIN_SOURCES - ${PLUGIN_SOURCES} - PARENT_SCOPE) diff --git a/cpp/tensorrt_llm/plugins/cumsumLastDimPlugin/cumsumLastDimPlugin.cpp b/cpp/tensorrt_llm/plugins/cumsumLastDimPlugin/cumsumLastDimPlugin.cpp deleted file mode 100644 index 927a42ebac2f..000000000000 --- a/cpp/tensorrt_llm/plugins/cumsumLastDimPlugin/cumsumLastDimPlugin.cpp +++ /dev/null @@ -1,299 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-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. - */ - -#include "cumsumLastDimPlugin.h" -#include "tensorrt_llm/common/assert.h" - -using namespace nvinfer1; -using namespace tensorrt_llm::kernels; -using namespace tensorrt_llm::common; -using tensorrt_llm::plugins::CumsumLastDimPluginCreator; -using tensorrt_llm::plugins::CumsumLastDimPlugin; - -static char const* CUMSUM_LAST_DIM_PLUGIN_VERSION{"1"}; -static char const* CUMSUM_LAST_DIM_PLUGIN_NAME{"CumsumLastDim"}; -PluginFieldCollection CumsumLastDimPluginCreator::mFC{}; -std::vector CumsumLastDimPluginCreator::mPluginAttributes; - -static constexpr SizeType32 LENGTH_LIMIT_FOR_BLOCKSCAN = 4096; - -CumsumLastDimPlugin::CumsumLastDimPlugin(SizeType32 inputLength, nvinfer1::DataType type, size_t temp_storage_bytes) - : mInputLength(inputLength) - , mTempStorageBytes(temp_storage_bytes) - , mType(type) -{ - TLLM_CHECK_WITH_INFO((mType == DataType::kBF16) || (mType == DataType::kFLOAT) || (mType == DataType::kHALF) - || (mType == DataType::kINT32), - "Only support int, float, half, and bfloat16."); - if (mTempStorageBytes == 0) - { - mTempStorageBytes = getWorkspaceSizeNeeded(inputLength, type); - } -} - -// Parameterized constructor -CumsumLastDimPlugin::CumsumLastDimPlugin(void const* data, size_t length) -{ - char const *d = reinterpret_cast(data), *a = d; - read(d, mInputLength); - read(d, mTempStorageBytes); - read(d, mType); - TLLM_CHECK(d == a + length); - TLLM_CHECK_WITH_INFO((mType == DataType::kBF16) || (mType == DataType::kFLOAT) || (mType == DataType::kHALF) - || (mType == DataType::kINT32), - "Only support int, float, half, and bfloat16."); -} - -// IPluginV2DynamicExt Methods -nvinfer1::IPluginV2DynamicExt* CumsumLastDimPlugin::clone() const noexcept -{ - auto* plugin = new CumsumLastDimPlugin(mInputLength, mType, mTempStorageBytes); - plugin->setPluginNamespace(mNamespace.c_str()); - return plugin; -} - -// Outputs -// output_tensor: [batch_size, inputLength] -nvinfer1::DimsExprs CumsumLastDimPlugin::getOutputDimensions( - int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept -{ - TLLM_CHECK_WITH_INFO(outputIndex == 0, "Only one output."); - return inputs[getInputTensorIdx()]; -} - -bool CumsumLastDimPlugin::supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept -{ - return (inOut[pos].type == mType) && (inOut[pos].format == TensorFormat::kLINEAR); -} - -void CumsumLastDimPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept -{ -} - -size_t CumsumLastDimPlugin::getWorkspaceSizeNeeded(SizeType32 inputLength, nvinfer1::DataType type) -{ - size_t tempStorageBytes{0}; - if (inputLength < LENGTH_LIMIT_FOR_BLOCKSCAN) // last dim unknown or small, use BlockScan - { - tempStorageBytes = 0; - } - else if (type == DataType::kINT32) - { - tempStorageBytes = invokeComputeCumsumLastDimWorkspaceSize(inputLength); - } - else if (type == DataType::kHALF) - { - tempStorageBytes = invokeComputeCumsumLastDimWorkspaceSize(inputLength); - } - else if (type == DataType::kFLOAT) - { - tempStorageBytes = invokeComputeCumsumLastDimWorkspaceSize(inputLength); - } -#ifdef ENABLE_BF16 - else if (type == DataType::kBF16) - { - tempStorageBytes = invokeComputeCumsumLastDimWorkspaceSize<__nv_bfloat16>(inputLength); - } -#endif - return tempStorageBytes; -} - -size_t CumsumLastDimPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - return mTempStorageBytes; -} - -template -int CumsumLastDimPlugin::enqueueImpl(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) -{ - // inputs - // 0. input_tensor [batch_size, inputLength] - // outputs - // 0. output_tensor [batch_size, inputLength] - auto const batchSize = inputDesc[getInputTensorIdx()].dims.d[0]; - auto const inputLength = inputDesc[getInputTensorIdx()].dims.d[1]; - /* - Two cases where we should use BlockScan: - 1. inputLength is small - 2. batchSize is large (since DeviceScan causes kernel launch per row) - */ - void* wp = inputLength < LENGTH_LIMIT_FOR_BLOCKSCAN || batchSize > 2 ? nullptr : workspace; - invokeCumsumLastDim( - batchSize, inputLength, inputs[getInputTensorIdx()], outputs[0], wp, mTempStorageBytes, stream); - - sync_check_cuda_error(stream); - return 0; -} - -int CumsumLastDimPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) noexcept -{ - if (mType == DataType::kINT32) - { - return enqueueImpl(inputDesc, outputDesc, inputs, outputs, workspace, stream); - } - else if (mType == DataType::kHALF) - { - return enqueueImpl(inputDesc, outputDesc, inputs, outputs, workspace, stream); - } - else if (mType == DataType::kFLOAT) - { - return enqueueImpl(inputDesc, outputDesc, inputs, outputs, workspace, stream); - } -#ifdef ENABLE_BF16 - else if (mType == DataType::kBF16) - { - return enqueueImpl<__nv_bfloat16>(inputDesc, outputDesc, inputs, outputs, workspace, stream); - } -#endif - return 0; -} - -// IPluginV2Ext Methods -nvinfer1::DataType CumsumLastDimPlugin::getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept -{ - TLLM_CHECK_WITH_INFO(index == 0, "Only one output."); - return inputTypes[getInputTensorIdx()]; -} - -// IPluginV2 Methods - -char const* CumsumLastDimPlugin::getPluginType() const noexcept -{ - return CUMSUM_LAST_DIM_PLUGIN_NAME; -} - -char const* CumsumLastDimPlugin::getPluginVersion() const noexcept -{ - return CUMSUM_LAST_DIM_PLUGIN_VERSION; -} - -int CumsumLastDimPlugin::getNbOutputs() const noexcept -{ - return 1; -} - -int CumsumLastDimPlugin::initialize() noexcept -{ - return 0; -} - -void CumsumLastDimPlugin::terminate() noexcept {} - -size_t CumsumLastDimPlugin::getSerializationSize() const noexcept -{ - return sizeof(mInputLength) + sizeof(mTempStorageBytes) + sizeof(mType); -} - -void CumsumLastDimPlugin::serialize(void* buffer) const noexcept -{ - char *d = static_cast(buffer), *a = d; - write(d, mInputLength); - write(d, mTempStorageBytes); - write(d, mType); - TLLM_CHECK(d == a + getSerializationSize()); -} - -void CumsumLastDimPlugin::destroy() noexcept -{ - delete this; -} - -/////////////// - -CumsumLastDimPluginCreator::CumsumLastDimPluginCreator() -{ - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - mPluginAttributes.emplace_back(PluginField("input_length", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("type_id", nullptr, PluginFieldType::kINT32)); - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -char const* CumsumLastDimPluginCreator::getPluginName() const noexcept -{ - return CUMSUM_LAST_DIM_PLUGIN_NAME; -} - -char const* CumsumLastDimPluginCreator::getPluginVersion() const noexcept -{ - return CUMSUM_LAST_DIM_PLUGIN_VERSION; -} - -PluginFieldCollection const* CumsumLastDimPluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -IPluginV2* CumsumLastDimPluginCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept -{ - PluginField const* fields = fc->fields; - int inputLength{}; - nvinfer1::DataType type{}; - // Read configurations from each fields - for (int i = 0; i < fc->nbFields; ++i) - { - char const* attrName = fields[i].name; - if (!strcmp(attrName, "input_length")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - inputLength = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "type_id")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - type = static_cast(*(static_cast(fields[i].data))); - } - } - try - { - auto* obj = new CumsumLastDimPlugin(inputLength, type); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -IPluginV2* CumsumLastDimPluginCreator::deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept -{ - // This object will be deleted when the network is destroyed, which will - // call CumsumLastDimPlugin::destroy() - try - { - auto* obj = new CumsumLastDimPlugin(serialData, serialLength); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} diff --git a/cpp/tensorrt_llm/plugins/cumsumLastDimPlugin/cumsumLastDimPlugin.h b/cpp/tensorrt_llm/plugins/cumsumLastDimPlugin/cumsumLastDimPlugin.h deleted file mode 100644 index 3cbf4e2356dd..000000000000 --- a/cpp/tensorrt_llm/plugins/cumsumLastDimPlugin/cumsumLastDimPlugin.h +++ /dev/null @@ -1,102 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-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. - */ - -#ifndef TRT_CUMSUM_LAST_DIM_PLUGIN_H -#define TRT_CUMSUM_LAST_DIM_PLUGIN_H - -#include "tensorrt_llm/kernels/cumsumLastDim.h" -#include "tensorrt_llm/plugins/common/plugin.h" -#include - -namespace tensorrt_llm::plugins -{ -class CumsumLastDimPlugin : public BasePlugin -{ -public: - using SizeType32 = tensorrt_llm::kernels::SizeType32; - - CumsumLastDimPlugin(SizeType32 inputLength, nvinfer1::DataType type, size_t tempStorageBytes = 0); - CumsumLastDimPlugin(void const* data, size_t length); - ~CumsumLastDimPlugin() override = default; - // IPluginV2DynamicExt Methods - nvinfer1::IPluginV2DynamicExt* clone() const noexcept override; - nvinfer1::DimsExprs getOutputDimensions(int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; - bool supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept override; - void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept override; - size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept override; - int enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - template - int enqueueImpl(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream); - size_t getWorkspaceSizeNeeded(SizeType32 inputLength, nvinfer1::DataType type); - - // IPluginV2Ext Methods - nvinfer1::DataType getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept override; - - // IPluginV2 Methods - char const* getPluginType() const noexcept override; - char const* getPluginVersion() const noexcept override; - int getNbOutputs() const noexcept override; - int initialize() noexcept override; - void terminate() noexcept override; - size_t getSerializationSize() const noexcept override; - void serialize(void* buffer) const noexcept override; - void destroy() noexcept override; - -private: - using IndexType = std::int32_t; - - IndexType getInputTensorIdx() const - { - return 0; - }; - -private: - SizeType32 mInputLength; - size_t mTempStorageBytes; - nvinfer1::DataType mType; -}; - -class CumsumLastDimPluginCreator : public BaseCreator -{ -public: - CumsumLastDimPluginCreator(); - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; - - nvinfer1::IPluginV2* deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept override; - -private: - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; -}; - -} // namespace tensorrt_llm::plugins - -#endif diff --git a/cpp/tensorrt_llm/plugins/doraPlugin/CMakeLists.txt b/cpp/tensorrt_llm/plugins/doraPlugin/CMakeLists.txt deleted file mode 100644 index 86876224fccd..000000000000 --- a/cpp/tensorrt_llm/plugins/doraPlugin/CMakeLists.txt +++ /dev/null @@ -1,21 +0,0 @@ -# -# SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. -# -file(GLOB SRCS *.cpp) -set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS}) -set(PLUGIN_SOURCES - ${PLUGIN_SOURCES} - PARENT_SCOPE) diff --git a/cpp/tensorrt_llm/plugins/doraPlugin/doraPlugin.cpp b/cpp/tensorrt_llm/plugins/doraPlugin/doraPlugin.cpp deleted file mode 100644 index 7c980f079ceb..000000000000 --- a/cpp/tensorrt_llm/plugins/doraPlugin/doraPlugin.cpp +++ /dev/null @@ -1,392 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-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. - */ - -#include "doraPlugin.h" - -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/cudaUtils.h" -#include "tensorrt_llm/common/memoryUtils.h" -#include "tensorrt_llm/runtime/iBuffer.h" - -#include - -using namespace nvinfer1; -using namespace tensorrt_llm::common; -using tensorrt_llm::plugins::DoraPlugin; -using tensorrt_llm::plugins::DoraPluginCreator; - -static char const* DORA_PLUGIN_VERSION{"1"}; -static char const* DORA_PLUGIN_NAME{"Dora"}; -PluginFieldCollection DoraPluginCreator::mFC{}; -std::vector DoraPluginCreator::mPluginAttributes; - -DoraPlugin::DoraPlugin(std::vector const& outHiddenSizes, nvinfer1::DataType type, bool removeInputPadding) - : mType(type) - , mRemoveInputPadding(removeInputPadding) - , mDoraImpl(outHiddenSizes, type) -{ - mOutHiddenSizes.resize(outHiddenSizes.size()); - mOutHiddenSizes.assign(outHiddenSizes.cbegin(), outHiddenSizes.cend()); - init(); -} - -void DoraPlugin::init() -{ - // initialize data to serialize - mDataToSerialize.clear(); - mDataToSerialize.emplace_back( - "out_hidden_sizes", mOutHiddenSizes.data(), PluginFieldType::kINT32, mOutHiddenSizes.size()); - mDataToSerialize.emplace_back("type", &mType, PluginFieldType::kINT32, 1); - mDataToSerialize.emplace_back("remove_input_padding", &mRemoveInputPadding, PluginFieldType::kINT8, 1); - mFieldsToSerialize.nbFields = static_cast(mDataToSerialize.size()); - mFieldsToSerialize.fields = mDataToSerialize.data(); -} - -// IPluginV3 methods -nvinfer1::IPluginCapability* DoraPlugin::getCapabilityInterface(nvinfer1::PluginCapabilityType type) noexcept -{ - switch (type) - { - case PluginCapabilityType::kBUILD: return static_cast(this); - case PluginCapabilityType::kRUNTIME: return static_cast(this); - case PluginCapabilityType::kCORE: return static_cast(this); - } - return nullptr; -} - -nvinfer1::IPluginV3* DoraPlugin::clone() noexcept -{ - std::unique_ptr plugin{std::make_unique(mOutHiddenSizes, mType, mRemoveInputPadding)}; - plugin->setPluginNamespace(mNamespace.c_str()); - return plugin.release(); -} - -// IPluginV3OneCore methods -char const* DoraPlugin::getPluginName() const noexcept -{ - return DORA_PLUGIN_NAME; -} - -char const* DoraPlugin::getPluginVersion() const noexcept -{ - return DORA_PLUGIN_VERSION; -} - -// IPluginV3OneBuild methods -int32_t DoraPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int32_t nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int32_t nbOutputs) noexcept -{ - return 0; -} - -int32_t DoraPlugin::getOutputDataTypes( - DataType* outputTypes, int32_t nbOutputs, DataType const* inputTypes, int32_t nbInputs) const noexcept -{ - try - { - TLLM_CHECK(nbOutputs == 1); - TLLM_CHECK(nbInputs == 2 + static_cast(mOutHiddenSizes.size()) + (mRemoveInputPadding ? 1 : 0)); - TLLM_CHECK(inputTypes[IdxEntry::kINPUT_TENSOR] == mType); - // output has the same dtype as the input, the plugin just applies scaling - outputTypes[0] = inputTypes[IdxEntry::kINPUT_TENSOR]; - } - catch (std::exception const& e) - { - caughtError(e); - } - return 0; -} - -int32_t DoraPlugin::getOutputShapes(DimsExprs const* inputs, int32_t nbInputs, DimsExprs const* shapeInputs, - int32_t nbShapeInputs, DimsExprs* outputs, int32_t nbOutputs, IExprBuilder& exprBuilder) noexcept -{ - try - { - TLLM_CHECK(nbOutputs == 1); - TLLM_CHECK(nbShapeInputs == 0); - TLLM_CHECK(nbInputs == 2 + static_cast(mOutHiddenSizes.size()) + (mRemoveInputPadding ? 1 : 0)); - - auto const inputTensorDims = inputs[IdxEntry::kINPUT_TENSOR]; - TLLM_CHECK(inputTensorDims.nbDims == (mRemoveInputPadding ? 2 : 3)); - - auto const lastDim = inputTensorDims.d[inputTensorDims.nbDims - 1]; - TLLM_CHECK(lastDim->isConstant()); - TLLM_CHECK(lastDim->getConstantValue() == std::accumulate(mOutHiddenSizes.cbegin(), mOutHiddenSizes.cend(), 0)); - - outputs[0].nbDims = inputTensorDims.nbDims; - for (auto dim = 0; dim < inputTensorDims.nbDims; ++dim) - { - outputs[0].d[dim] = inputTensorDims.d[dim]; - } - } - catch (std::exception const& e) - { - caughtError(e); - } - return 0; -} - -bool DoraPlugin::supportsFormatCombination( - int32_t pos, nvinfer1::DynamicPluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept -{ - auto const numModules = static_cast(mOutHiddenSizes.size()); - if (nbInputs != 2 + numModules + (mRemoveInputPadding ? 1 : 0)) - { - return false; - } - - bool const isInput = pos < nbInputs; - if (pos == IdxEntry::kHOST_REQUEST_TYPES) - { - return (inOut[pos].desc.type == nvinfer1::DataType::kINT32); - } - // optional host_context_lens after lora pointers - else if (pos == IdxEntry::kLORA_WEIGHTS_PTRS_START + numModules and isInput) - { - return (inOut[pos].desc.type == nvinfer1::DataType::kINT32 and mRemoveInputPadding); - } - // lora weight pointers - else if (pos >= IdxEntry::kLORA_WEIGHTS_PTRS_START and pos < IdxEntry::kLORA_WEIGHTS_PTRS_START + numModules) - { - return (inOut[pos].desc.type == nvinfer1::DataType::kINT64); - } - else if (pos != 0 and isInput) - { - TLLM_LOG_WARNING("%s: got an unexpected input at position %d", __PRETTY_FUNCTION__, pos); - return false; - } - - return (inOut[pos].desc.type == mType) and (inOut[pos].desc.format == TensorFormat::kLINEAR); -} - -int32_t DoraPlugin::getNbOutputs() const noexcept -{ - return 1; -} - -size_t DoraPlugin::getWorkspaceSize(nvinfer1::DynamicPluginTensorDesc const* inputs, int32_t nbInputs, - nvinfer1::DynamicPluginTensorDesc const* outputs, int32_t nbOutputs) const noexcept -{ - auto const inputTensorMax = inputs[IdxEntry::kINPUT_TENSOR].max; - auto const maxNumTokens = mRemoveInputPadding ? inputTensorMax.d[0] : inputTensorMax.d[0] * inputTensorMax.d[1]; - auto const size = mDoraImpl.getWorkspaceSize(maxNumTokens); - return size; -} - -int32_t DoraPlugin::getValidTactics(int32_t* tactics, int32_t nbTactics) noexcept -{ - return 0; -} - -int32_t DoraPlugin::getNbTactics() noexcept -{ - return 0; -} - -char const* DoraPlugin::getTimingCacheID() noexcept -{ - return nullptr; -} - -int32_t DoraPlugin::getFormatCombinationLimit() noexcept -{ - return 1; -} - -char const* DoraPlugin::getMetadataString() noexcept -{ - return nullptr; -} - -// IPluginV3OneRuntime methods -int32_t DoraPlugin::setTactic(int32_t tactic) noexcept -{ - return 0; -} - -int32_t DoraPlugin::onShapeChange(nvinfer1::PluginTensorDesc const* in, int32_t nbInputs, - nvinfer1::PluginTensorDesc const* out, int32_t nbOutputs) noexcept -{ - return 0; -} - -int32_t DoraPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept -{ - if (isBuilding()) - { - return 0; - } - - auto const numModules = static_cast(mOutHiddenSizes.size()); - auto const numReqs = inputDesc[IdxEntry::kHOST_REQUEST_TYPES].dims.d[0]; - - auto const inputTensorDesc = inputDesc[IdxEntry::kINPUT_TENSOR]; - auto const numTokens - = mRemoveInputPadding ? inputTensorDesc.dims.d[0] : inputTensorDesc.dims.d[0] * inputTensorDesc.dims.d[1]; - auto const seqLen = mRemoveInputPadding ? 0 : inputTensorDesc.dims.d[1]; - - void const* inputTensor = inputs[IdxEntry::kINPUT_TENSOR]; - auto const* hostRequestTypes = static_cast(inputs[IdxEntry::kHOST_REQUEST_TYPES]); - void const* const* loraWeightsPtrs = &inputs[IdxEntry::kLORA_WEIGHTS_PTRS_START]; - - int32_t const* hostContextLengths = mRemoveInputPadding - ? static_cast(inputs[IdxEntry::kLORA_WEIGHTS_PTRS_START + numModules]) - : nullptr; - - mExpandDoraWeightPtrs.clear(); - mExpandDoraWeightPtrs.reserve(numModules * numTokens); - - bool hasAnyDora = false; - - for (auto moduleIdx = 0; moduleIdx < numModules; moduleIdx++) - { - auto const loraWeightModulePtrs = static_cast(loraWeightsPtrs[moduleIdx]); - - int idx = 0; - for (int reqId = 0; reqId < numReqs; reqId++) - { - // loraWeightModulePtrs has 3 pointers for each module: A,B, and an optional DoRA magnitude - // the current DoRA plugin does not apply LoRA, so A and B are ignored. - RequestType const reqType = static_cast(hostRequestTypes[reqId]); - auto const* modulePtr = reinterpret_cast(loraWeightModulePtrs[reqId * 3 + 2]); - hasAnyDora = hasAnyDora or modulePtr != nullptr; - - if (reqType == RequestType::kGENERATION) - { - mExpandDoraWeightPtrs.push_back(modulePtr); - idx += 1; - } - else - { - int contextLen = (mRemoveInputPadding ? hostContextLengths[reqId] : seqLen); - - for (int contextId = 0; contextId < contextLen; contextId++) - { - mExpandDoraWeightPtrs.push_back(modulePtr); - idx += 1; - } - } - } - if (idx != numTokens) - { - TLLM_LOG_ERROR("LoraParams and input dims don't match, lora tokens %d input tokens %d", idx, numTokens); - return -1; - } - } - - if (hasAnyDora) - { - mDoraImpl.run(numTokens, inputTensor, mExpandDoraWeightPtrs.data(), outputs, workspace, stream); - } - else - { - // skip dora scaling if all requests are pure-lora - auto const inputRank = inputTensorDesc.dims.nbDims; - auto const numel - = std::accumulate(inputTensorDesc.dims.d, inputTensorDesc.dims.d + inputRank, 1, std::multiplies()); - auto const elemSize = tensorrt_llm::common::getDTypeSize(mType); - tensorrt_llm::common::cudaAutoCpy((int8_t*) outputs[0], (int8_t*) inputTensor, numel * elemSize, stream); - } - - sync_check_cuda_error(stream); - return 0; -} - -nvinfer1::IPluginV3* DoraPlugin::attachToContext(nvinfer1::IPluginResourceContext* context) noexcept -{ - return clone(); -} - -nvinfer1::PluginFieldCollection const* DoraPlugin::getFieldsToSerialize() noexcept -{ - return &mFieldsToSerialize; -} - -DoraPluginCreator::DoraPluginCreator() -{ - mPluginAttributes.clear(); - mPluginAttributes.emplace_back("num_modules", nullptr, PluginFieldType::kINT32, 1); - mPluginAttributes.emplace_back("type", nullptr, PluginFieldType::kINT32, 1); - mPluginAttributes.emplace_back("remove_input_padding", nullptr, PluginFieldType::kINT8, 1); - mFC.nbFields = static_cast(mPluginAttributes.size()); - mFC.fields = mPluginAttributes.data(); -} - -char const* DoraPluginCreator::getPluginName() const noexcept -{ - return DORA_PLUGIN_NAME; -} - -char const* DoraPluginCreator::getPluginVersion() const noexcept -{ - return DORA_PLUGIN_VERSION; -} - -nvinfer1::PluginFieldCollection const* DoraPluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -nvinfer1::IPluginV3* DoraPluginCreator::createPlugin( - char const* name, nvinfer1::PluginFieldCollection const* fc, nvinfer1::TensorRTPhase phase) noexcept -{ - PluginField const* fields = fc->fields; - nvinfer1::DataType type{}; - bool removeInputPadding{}; - std::vector outHiddenSizes; - - // Read configurations from each field - for (int i = 0; i < fc->nbFields; ++i) - { - auto const field = fields[i]; - char const* attrName = fields[i].name; - if (!strcmp(attrName, "type")) - { - TLLM_CHECK(field.type == PluginFieldType::kINT32 and field.length == 1); - type = *static_cast(field.data); - } - else if (!strcmp(attrName, "remove_input_padding")) - { - TLLM_CHECK(field.type == PluginFieldType::kINT8 and field.length == 1); - removeInputPadding = *static_cast(field.data); - } - else if (!strcmp(attrName, "out_hidden_sizes")) - { - TLLM_CHECK(field.type == PluginFieldType::kINT32); - auto const* outHiddenSizesPtr = static_cast(field.data); - outHiddenSizes.resize(field.length); - outHiddenSizes.assign(outHiddenSizesPtr, outHiddenSizesPtr + field.length); - } - else - { - TLLM_LOG_WARNING("%s: got an unexpected attribute: %s", __PRETTY_FUNCTION__, attrName); - } - } - - try - { - auto* obj = new DoraPlugin(outHiddenSizes, type, removeInputPadding); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} diff --git a/cpp/tensorrt_llm/plugins/doraPlugin/doraPlugin.h b/cpp/tensorrt_llm/plugins/doraPlugin/doraPlugin.h deleted file mode 100644 index dfee11fdc90e..000000000000 --- a/cpp/tensorrt_llm/plugins/doraPlugin/doraPlugin.h +++ /dev/null @@ -1,114 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. - */ -#include "tensorrt_llm/kernels/lora/dora.h" -#include "tensorrt_llm/plugins/common/plugin.h" - -namespace tensorrt_llm::plugins -{ - -class DoraPlugin : public BasePluginV3 -{ -public: - DoraPlugin() = delete; - DoraPlugin(std::vector const& outHiddenSizes, nvinfer1::DataType type, bool removeInputPadding); - DoraPlugin(DoraPlugin const& p) = default; - - // IPluginV3 methods - nvinfer1::IPluginCapability* getCapabilityInterface(nvinfer1::PluginCapabilityType type) noexcept override; - nvinfer1::IPluginV3* clone() noexcept override; - - // IPluginV3OneCore methods - char const* getPluginName() const noexcept override; - char const* getPluginVersion() const noexcept override; - - // IPluginV3OneBuild methods - int32_t configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int32_t nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int32_t nbOutputs) noexcept override; - int32_t getOutputDataTypes(nvinfer1::DataType* outputTypes, int32_t nbOutputs, nvinfer1::DataType const* inputTypes, - int32_t nbInputs) const noexcept override; - int32_t getOutputShapes(nvinfer1::DimsExprs const* inputs, int32_t nbInputs, nvinfer1::DimsExprs const* shapeInputs, - int32_t nbShapeInputs, nvinfer1::DimsExprs* outputs, int32_t nbOutputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; - bool supportsFormatCombination(int32_t pos, nvinfer1::DynamicPluginTensorDesc const* inOut, int32_t nbInputs, - int32_t nbOutputs) noexcept override; - int32_t getNbOutputs() const noexcept override; - size_t getWorkspaceSize(nvinfer1::DynamicPluginTensorDesc const* inputs, int32_t nbInputs, - nvinfer1::DynamicPluginTensorDesc const* outputs, int32_t nbOutputs) const noexcept override; - int32_t getValidTactics(int32_t* tactics, int32_t nbTactics) noexcept override; - int32_t getNbTactics() noexcept override; - char const* getTimingCacheID() noexcept override; - int32_t getFormatCombinationLimit() noexcept override; - char const* getMetadataString() noexcept override; - - // IPluginV3OneRuntime methods - int32_t setTactic(int32_t tactic) noexcept override; - int32_t onShapeChange(nvinfer1::PluginTensorDesc const* in, int32_t nbInputs, nvinfer1::PluginTensorDesc const* out, - int32_t nbOutputs) noexcept override; - int32_t enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) noexcept override; // fixed - nvinfer1::IPluginV3* attachToContext(nvinfer1::IPluginResourceContext* context) noexcept override; - nvinfer1::PluginFieldCollection const* getFieldsToSerialize() noexcept override; - -private: - void init(); - - std::vector mDataToSerialize; - nvinfer1::PluginFieldCollection mFieldsToSerialize; - - enum IdxEntry - { - kINPUT_TENSOR = 0, - kHOST_REQUEST_TYPES = 1, - kLORA_WEIGHTS_PTRS_START = 2 - }; - - // TODO(oargov) this is shared with the LoRA plugin, put it somewhere else - enum class RequestType : int32_t - { - kCONTEXT = 0, - kGENERATION = 1 - }; - - std::vector mOutHiddenSizes; - nvinfer1::DataType mType; - bool mRemoveInputPadding; - tensorrt_llm::kernels::DoraImpl mDoraImpl; - - std::vector mExpandDoraWeightPtrs{}; -}; - -class DoraPluginCreator : public BaseCreatorV3 -{ -public: - DoraPluginCreator(); - - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV3* createPlugin( - char const* name, nvinfer1::PluginFieldCollection const* fc, nvinfer1::TensorRTPhase phase) noexcept override; - -private: - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; -}; - -}; // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/eaglePlugin/CMakeLists.txt b/cpp/tensorrt_llm/plugins/eaglePlugin/CMakeLists.txt deleted file mode 100644 index b6bd0439cc0c..000000000000 --- a/cpp/tensorrt_llm/plugins/eaglePlugin/CMakeLists.txt +++ /dev/null @@ -1,21 +0,0 @@ -# -# SPDX-FileCopyrightText: Copyright (c) 1993-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. -# -file(GLOB SRCS *.cpp) -set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS}) -set(PLUGIN_SOURCES - ${PLUGIN_SOURCES} - PARENT_SCOPE) diff --git a/cpp/tensorrt_llm/plugins/eaglePlugin/eagleDecodeDraftTokensPlugin.cpp b/cpp/tensorrt_llm/plugins/eaglePlugin/eagleDecodeDraftTokensPlugin.cpp deleted file mode 100644 index 899c93855b9f..000000000000 --- a/cpp/tensorrt_llm/plugins/eaglePlugin/eagleDecodeDraftTokensPlugin.cpp +++ /dev/null @@ -1,945 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-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. - */ -#include "eagleDecodeDraftTokensPlugin.h" - -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/dataType.h" -#include "tensorrt_llm/common/memoryUtils.h" -#include "tensorrt_llm/kernels/samplingTopKKernels.h" -#include "tensorrt_llm/kernels/speculativeDecoding/eagleDecodingKernels.h" -#include "tensorrt_llm/kernels/speculativeDecoding/medusaDecodingKernels.h" -#include "tensorrt_llm/runtime/common.h" -#include "tensorrt_llm/runtime/iTensor.h" - -using namespace nvinfer1; -using tensorrt_llm::plugins::EagleDecodeDraftTokensPluginCreator; -using tensorrt_llm::plugins::EagleDecodeDraftTokensPlugin; -using namespace tensorrt_llm::kernels; -using namespace tensorrt_llm::kernels::speculative_decoding; -using namespace tensorrt_llm::runtime; -namespace tc = tensorrt_llm::common; - -static char const* EAGLE_DECODE_DRAFT_TOKENS_PLUGIN_VERSION{"1"}; -static char const* EAGLE_DECODE_DRAFT_TOKENS_PLUGIN_NAME{"EagleDecodeDraftTokens"}; -PluginFieldCollection EagleDecodeDraftTokensPluginCreator::mFC{}; -std::vector EagleDecodeDraftTokensPluginCreator::mPluginAttributes; - -EagleDecodeDraftTokensPlugin::EagleDecodeDraftTokensPlugin( - nvinfer1::DataType type, int32_t layerIdx, int32_t numEagleLayers, bool topKSampling) - : mDtype(type) - , mLayerIdx(layerIdx) - , mNumEagleLayers(numEagleLayers) - , mTopKSampling(topKSampling) -{ - TLLM_CHECK_WITH_INFO(mTopKSampling, "Multinomial sampling is not supported yet."); -} - -// Parameterized constructor -EagleDecodeDraftTokensPlugin::EagleDecodeDraftTokensPlugin(void const* data, size_t length) -{ - char const *d = reinterpret_cast(data), *a = d; - read(d, mDtype); - read(d, mLayerIdx); - read(d, mNumEagleLayers); - read(d, mTopKSampling); - TLLM_CHECK_WITH_INFO(d == a + length, - "Expected length (%d) != real length (%d). This is often " - "caused by using different TensorRT LLM version to build " - "engine and run engine.", - static_cast(length), static_cast(d - a)); -} - -// IPluginV2DynamicExt Methods -nvinfer1::IPluginV2DynamicExt* EagleDecodeDraftTokensPlugin::clone() const noexcept -{ - auto* plugin = new EagleDecodeDraftTokensPlugin(*this); - plugin->setPluginNamespace(mNamespace.c_str()); - return plugin; -} - -nvinfer1::DimsExprs EagleDecodeDraftTokensPlugin::getOutputDimensions( - int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept -{ - TLLM_CHECK(outputIndex < getNbOutputs()); - TLLM_CHECK(nbInputs == 12); - auto const batchSizeExpr = inputs[getIdx(InputIdxEntry::PATHS)].d[0]; - auto const maxDecodingTokensExpr = inputs[getIdx(InputIdxEntry::PATHS)].d[1]; - auto const maxPathLengthExpr = inputs[getIdx(InputIdxEntry::PATHS)].d[2]; - auto const maxDecodingDraftTokensExpr - = exprBuilder.operation(DimensionOperation::kSUB, *maxDecodingTokensExpr, *exprBuilder.constant(1)); - - auto const numEagleLayersExpr - = exprBuilder.operation(DimensionOperation::kSUB, *maxPathLengthExpr, *exprBuilder.constant(1)); - auto const maxDecodingDraftTokensSquareExpr - = exprBuilder.operation(DimensionOperation::kPROD, *maxDecodingDraftTokensExpr, - *maxDecodingDraftTokensExpr); // maxDecodingDraftTokensExpr * maxDecodingDraftTokensExpr - - nvinfer1::DimsExprs ret; - if (outputIndex == getIdx(OutputIdxEntry::OUTPUT_DRAFT_TOKEN_IDS)) - { - // output_draft_token_ids: [batch_size, max_decoding_draft_tokens] - ret.nbDims = 2; - ret.d[0] = batchSizeExpr; - ret.d[1] = maxDecodingDraftTokensExpr; - } - else if (outputIndex == getIdx(OutputIdxEntry::OUTPUT_DRAFT_LENS)) - { - // output_draft_lens: [batch_size] - ret.nbDims = 1; - ret.d[0] = batchSizeExpr; - } - else if (outputIndex == getIdx(OutputIdxEntry::OUTPUT_PATHS)) - { - // output_path: [batch_size, max_decoding_tokens, max_path_len] - ret.nbDims = 3; - ret.d[0] = batchSizeExpr; - ret.d[1] = maxDecodingTokensExpr; - ret.d[2] = maxPathLengthExpr; - } - else if (outputIndex == getIdx(OutputIdxEntry::OUTPUT_CURRENT_SCORES)) - { - // output_current_scores: [batch_size, max_decoding_draft_tokens] - ret.nbDims = 2; - ret.d[0] = batchSizeExpr; - ret.d[1] = maxDecodingDraftTokensExpr; - } - else if (outputIndex == getIdx(OutputIdxEntry::OUTPUT_NEXT_EXPAND_INDICES)) - { - // output_next_expand_index - ret.nbDims = 2; - ret.d[0] = batchSizeExpr; - ret.d[1] = maxDecodingDraftTokensExpr; - } - else if (outputIndex == getIdx(OutputIdxEntry::OUTPUT_ALL_LAYERS_SCORES)) - { - // output_all_layers_scores: - // [batch_size, num_eagle_layers, max_decoding_draft_tokens x max_decoding_draft_tokens] - ret.nbDims = 3; - ret.d[0] = batchSizeExpr; - ret.d[1] = numEagleLayersExpr; - ret.d[2] = maxDecodingDraftTokensSquareExpr; - } - else if (outputIndex == getIdx(OutputIdxEntry::OUTPUT_ALL_LAYERS_DRAFT_TOKEN_IDS)) - { - // output_all_layers_draft_token_ids: - // [batch_size, num_eagle_layers, max_decoding_draft_tokens x max_decoding_draft_tokens] - ret.nbDims = 3; - ret.d[0] = batchSizeExpr; - ret.d[1] = numEagleLayersExpr; - ret.d[2] = maxDecodingDraftTokensSquareExpr; - } - else if (outputIndex == getIdx(OutputIdxEntry::OUTPUT_ALL_LAYERS_DRAFT_TOKEN_IDS_PREDECESSOR)) - { - // output_all_layers_draft_token_ids_predecessor - // [batch_size, num_eagle_layers, max_decoding_draft_tokens x max_decoding_draft_tokens] - ret.nbDims = 3; - ret.d[0] = batchSizeExpr; - ret.d[1] = numEagleLayersExpr; - ret.d[2] = maxDecodingDraftTokensSquareExpr; - } - else - { - TLLM_CHECK_WITH_INFO( - false, "Wrong outputIndex %d in EagleDecodeDraftTokensPlugin::getOutputDimensions", outputIndex); - } - return ret; -} - -bool EagleDecodeDraftTokensPlugin::supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept -{ - TLLM_CHECK(nbInputs == 12 && nbOutputs == getNbOutputs()); - TLLM_CHECK(pos < nbInputs + nbOutputs); - - if (pos == getIdx(InputIdxEntry::LOGITS)) - { - // input: logits - // output: output_all_layers_scores - return (inOut[pos].type == mDtype) && (inOut[pos].format == TensorFormat::kLINEAR); - } - else if (pos == getIdx(InputIdxEntry::INPUT_ALL_LAYERS_SCORES) || pos == getIdx(InputIdxEntry::INPUT_PREV_SCORES) - || pos == nbInputs + getIdx(OutputIdxEntry::OUTPUT_ALL_LAYERS_SCORES) - || pos == nbInputs + getIdx(OutputIdxEntry::OUTPUT_CURRENT_SCORES)) - { - // input: rand_sample, input_all_layers_scores, input_prev_scores - // output: output_all_layers_scores, output_current_scores - return (inOut[pos].type == nvinfer1::DataType::kFLOAT) && (inOut[pos].format == TensorFormat::kLINEAR); - } - else - { - // input: path, num_valid_logits, use_dynamic_tree, dynamic_tree_max_topK, input_draft_token_ids, - // input_draft_lens, input_current_expand_index, input_all_layers_draft_token_ids - // output: output_draft_token_ids, output_draft_lens, output_path, output_next_expand_index - // output_all_layers_draft_token_ids, output_all_alyers_draft_token_predecessor - return (inOut[pos].type == nvinfer1::DataType::kINT32) && (inOut[pos].format == TensorFormat::kLINEAR); - } -} - -void EagleDecodeDraftTokensPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept -{ -} - -template -size_t EagleDecodeDraftTokensPlugin::getWorkspaceSizeType(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - size_t workspaceSize{0}; - auto const numInputLogits = inputs[getIdx(InputIdxEntry::LOGITS)].dims.d[0]; - auto const batchSize = inputs[getIdx(InputIdxEntry::PATHS)].dims.d[0]; - auto const vocabSizePadded = inputs[getIdx(InputIdxEntry::LOGITS)].dims.d[1]; - auto const maxDecodingTokens = inputs[getIdx(InputIdxEntry::PATHS)].dims.d[1]; - auto const maxDecodingDraftTokens = maxDecodingTokens - 1; - auto const maxTopK = maxDecodingDraftTokens; - auto const mNumEagleLayers = inputs[getIdx(InputIdxEntry::INPUT_ALL_LAYERS_SCORES)].dims.d[1]; - - // Greedy sampling - if (mTopKSampling) - { - // 0. The first topK sampling workspace - auto const draftTokenSamplingWorkspaceSize - = getTopKWorkspaceSize(numInputLogits, /* maxTokensPerStep */ 1, /* maxTopK */ maxTopK, vocabSizePadded); - - // 1. The first TopKs [numInputLogits] - auto const topKsSize = numInputLogits * sizeof(SizeType32); - - // 2. Topks offset [batchSize] - // Each request will have different number of logits that need to be sampled - // This tensor will record the start offset of the topK for each request - auto const topKOffsetSize = batchSize * sizeof(SizeType32); - - // 3. Logits ptrs [numInputLogits] - auto const logitsPtrsSize = numInputLogits * sizeof(T*); - - // 4. The first topK sampling's output ids ptrs [numInputLogits][maxDecodingDraftTokens] - auto const firstTopKOutputIdsPtrsSize = numInputLogits * sizeof(TokenIdType*); - - // 5. The first topK sampling's output ids (temporary buffer) [numInputLogits * maxDecodingDraftTokens] - auto const firstTopKOutputIdsSize = numInputLogits * maxDecodingDraftTokens * sizeof(TokenIdType); - - // 6. Number of successors for each nodes, extract from the paths and layerId - // [batchSize * maxDecodingTokens] - auto const numSuccessorsForEachNodeSize = batchSize * maxDecodingTokens * sizeof(SizeType32); - - // 7. Flag whether to do decoding or not. SamplingTopK is done for numInputLogits tokens. - // But only sum(numValidLogitsPerRequest[:]) of them are valid. - // [batchSize * maxDecodingTokens] - auto const skipDecodeSize = numInputLogits * sizeof(bool); - - // 8. The first topK sampling's logprobs [batchSize * maxDecodingDraftTokens] - auto const firstTopKOutputLogProbsSize = numInputLogits * maxDecodingDraftTokens * sizeof(float); - - // 9. Eagle-2, the second topK sampling workspace - // Sampling from [batchSize, maxTopK * maxTopK] to [batchSize, maxTopK] - auto const secondTopKSamplingWorkspaceSize = getTopKWorkspaceSize( - batchSize, /* maxTokensPerStep */ 1, /* maxTopK */ maxTopK, maxTopK * maxTopK); - - // 10. Eagle-2, the outputIds of the second topK sampling, shape [batchSize, maxDecodingTokens] - auto const secondTopKOutputIdsSize = batchSize * maxDecodingTokens * sizeof(TokenIdType); - // 11. Eagle-2, the outputIdsPtr of the second topK sampling, shape [batchSize] - auto const secondTopKOutputIdsPtrSize = batchSize * sizeof(TokenIdType*); - // 12. Eagle-2, the inputScoresPtrs of the second topK sampling, shape [batchSize] - auto const secondTopKInputScoresPtrsSize = batchSize * sizeof(float*); - // 13. Eagle-2, the outpuLogProbs of the second topK samplig, shape [batchSize, maxDecodingDraftTokens] - auto const secondTopKOutputLogProbsSize = batchSize * maxDecodingDraftTokens * sizeof(float); - - // 14. Eagle-2, the input scores pointers of the third topK sampling, shape [batchSize] - // Each points to a vocabSize = '(mNumEagleLayers - 1) * dynamicTreeMaxTopK * dynamicTreeMaxTopK + - // dynamicTreeMaxTopK' - auto const thirdTopKInputScoresPtrsSize = batchSize * sizeof(float*); - // 15. Eagle-2, the output of the third topK sampling, shape [batchSize, maxDecodingDraftTokens] - auto const thirdTopKOutputIdsSize = batchSize * maxDecodingDraftTokens * sizeof(TokenIdType); - // 16. Eagle-2, the output pointers of the third topK sampling, shape [batchSize] - auto const thirdTopKOutputIdsPtrsSize = batchSize * sizeof(TokenIdType*); - // 17. Eagle-2, the workspace of the third topK sampling - // Sampling from [batchSize, '(mNumEagleLayers - 1) * dynamicTreeMaxTopK * dynamicTreeMaxTopK + - // dynamicTreeMaxTopK'] to [batchSize, maxDecodingDraftTokens] We over-set the vocabsize here. - auto const thridTopKSamplingWorkspaceSize = getTopKWorkspaceSize(batchSize, /* maxTokensPerStep */ 1, - /* maxTopK */ maxDecodingDraftTokens, mNumEagleLayers * maxDecodingDraftTokens * maxDecodingDraftTokens); - - // 18. Eagle-2, the topKs for each request in the third topK sampling - // The real topK value is min(maxDecodingDraftTokens, totalNumDraftTokensForAllLayers) - auto const thirdTopKsSize = batchSize * sizeof(SizeType32); - - SizeType32 constexpr NUM_BUFFERS{19}; - size_t workspaces[NUM_BUFFERS]; - workspaces[0] = draftTokenSamplingWorkspaceSize; - workspaces[1] = topKsSize; - workspaces[2] = topKOffsetSize; - workspaces[3] = logitsPtrsSize; - workspaces[4] = firstTopKOutputIdsPtrsSize; - workspaces[5] = firstTopKOutputIdsSize; - workspaces[6] = numSuccessorsForEachNodeSize; - workspaces[7] = skipDecodeSize; - workspaces[8] = firstTopKOutputLogProbsSize; - workspaces[9] = secondTopKSamplingWorkspaceSize; - workspaces[10] = secondTopKOutputIdsSize; - workspaces[11] = secondTopKOutputIdsPtrSize; - workspaces[12] = secondTopKInputScoresPtrsSize; - workspaces[13] = secondTopKOutputLogProbsSize; - workspaces[14] = thirdTopKInputScoresPtrsSize; - workspaces[15] = thirdTopKOutputIdsSize; - workspaces[16] = thirdTopKOutputIdsPtrsSize; - workspaces[17] = thridTopKSamplingWorkspaceSize; - workspaces[18] = thirdTopKsSize; - workspaceSize = tc::calculateTotalWorkspaceSize(workspaces, NUM_BUFFERS); - } - else - { - // TODO fill me - // Multinomial sampling - TLLM_CHECK_WITH_INFO(false, "Multinomial sampling is not supported yet."); - } - - return workspaceSize; -} - -size_t EagleDecodeDraftTokensPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - auto const logitsType = inputs[getIdx(InputIdxEntry::LOGITS)].type; - if (logitsType == nvinfer1::DataType::kFLOAT) - { - return getWorkspaceSizeType(inputs, nbInputs, outputs, nbOutputs); - } - else if (logitsType == nvinfer1::DataType::kHALF) - { - return getWorkspaceSizeType<__half>(inputs, nbInputs, outputs, nbOutputs); - } - else - { - TLLM_CHECK_WITH_INFO(false, "Unsupported logits type"); - } - return 0; -} - -template -void EagleDecodeDraftTokensPlugin::doTopKSampling(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) noexcept -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - // We allocate many buffers with 'numInputLogits' size, but the input logits will include some padding logits. - // So only 'batchSize' or 'numValidLogits' size will be actually used. - auto const numInputLogits = inputDesc[getIdx(InputIdxEntry::LOGITS)].dims.d[0]; - auto const vocabSizePadded = inputDesc[getIdx(InputIdxEntry::LOGITS)].dims.d[1]; - auto const batchSize = inputDesc[getIdx(InputIdxEntry::PATHS)].dims.d[0]; - auto const maxDecodingTokens = inputDesc[getIdx(InputIdxEntry::PATHS)].dims.d[1]; - auto const maxPathLen = inputDesc[getIdx(InputIdxEntry::PATHS)].dims.d[2]; - auto const maxDecodingDraftTokens = maxDecodingTokens - 1; - auto const maxTopK = maxDecodingDraftTokens; - - ////////////////////////////////////////// Get plugin inputs ////////////////////////////////////////// - // Plugin inputs - // Input logits for sampling, shape: [numInputLogits, vocabSizePadded] - auto pluginInputLogits = static_cast(inputs[getIdx(InputIdxEntry::LOGITS)]); - // Input paths, shape: [batchSize, maxDecodingTokens, maxPathLen] - auto pluginInputPaths = static_cast(inputs[getIdx(InputIdxEntry::PATHS)]); - auto numValidLogits = static_cast(inputs[getIdx(InputIdxEntry::NUM_VALID_LOGITS)]); - // For Eagle-2 - // Whether to use dynamic tree (i.e., Eagle-2) - auto useDynamicTree = *(static_cast(inputs[getIdx(InputIdxEntry::USE_DYNAMIC_TREE)])); - // The max topK for dynamic tree. All the requests have the same expand topK. - // In Eagle-2, dynamicTreeMaxTopK is equal to maxNonLeavesPerLayer in the internal EagleNets. - auto dynamicTreeMaxTopK = *(static_cast(inputs[getIdx(InputIdxEntry::DYNAMIC_TREE_MAX_TOPK)])); - // All layer's draft tokenIds, shape: [batchSize, maxDecodingDraftTokens] - auto pluginInputDraftTokenIds - = reinterpret_cast(inputs[getIdx(InputIdxEntry::INPUT_DRAFT_TOKEN_IDS)]); - // The number of all layer's draft tokenIds, shape: [batchSize] - auto pluginInputDraftLens = reinterpret_cast(inputs[getIdx(InputIdxEntry::INPUT_DRAFT_LENS)]); - // The previous EagleNet's scores, shape: [batchSize, maxDecodingDraftTokens] - auto pluginInputPrevScores = static_cast(inputs[getIdx(InputIdxEntry::INPUT_PREV_SCORES)]); - // The indices of the nodes that will be expand in this layer, shape: [batchSize, maxDecodingDraftTokens] - // The index is related to the final output tree, which has max_decoding_draft_tokens draft tokens. - auto pluginInputCurrentExpandIndices - = reinterpret_cast(inputs[getIdx(InputIdxEntry::INPUT_CURRENT_EXPAND_INDICES)]); - // The scores from all previous EagleNets, - // shape: [batchSize, mNumEagleLayers, maxDecodingDraftTokens x maxDecodingDraftTokens] - auto pluginInputAllLayersScores = static_cast(inputs[getIdx(InputIdxEntry::INPUT_ALL_LAYERS_SCORES)]); - // The draft tokens from all previous EagleNets, - // shape: [batchSize, mNumEagleLayers, maxDecodingDraftTokens x maxDecodingDraftTokens] - auto pluginInputAllLayersDraftTokenIds - = reinterpret_cast(inputs[getIdx(InputIdxEntry::INPUT_ALL_LAYERS_DRAFT_TOKEN_IDS)]); - // The predecessor of all the draft tokens, - // shape: [batchSize, mNumEagleLayers, maxDecodingDraftTokens x maxDecodingDraftTokens] - auto pluginInputAllLayersDraftTokenIdsPredecessor = reinterpret_cast( - inputs[getIdx(InputIdxEntry::INPUT_ALL_LAYERS_DRAFT_TOKEN_IDS_PREDECESSOR)]); - - ////////////////////////////////////////// Get plugin outputs ////////////////////////////////////////// - // Plugin outputs - // All layer's draft tokenIds, shape: [batchSize, maxDecodingDraftTokens] - auto pluginOutputDraftTokenIds - = reinterpret_cast(outputs[getIdx(OutputIdxEntry::OUTPUT_DRAFT_TOKEN_IDS)]); - // The number of all layer's draft tokenIds, shape: [batchSize] - auto pluginOutputDraftLens = reinterpret_cast(outputs[getIdx(OutputIdxEntry::OUTPUT_DRAFT_LENS)]); - // For Eagle-2 - // Updated paths base on this layer's sampling result, shape: [batchSize, maxDecodingTokens, maxPathLen] - auto pluginOutputPaths = reinterpret_cast(outputs[getIdx(OutputIdxEntry::OUTPUT_PATHS)]); - // This layer's scores, which will be used in next layers [batchSize, maxDecodingDraftTokens] - auto pluginOutputCurrentScores = static_cast(outputs[getIdx(OutputIdxEntry::OUTPUT_CURRENT_SCORES)]); - // The indices of the nodes that will be expand in next layer, shape: [batchSize, maxDecodingDraftTokens] - // The index is related to the final output tree, which has max_decoding_draft_tokens draft tokens. - auto pluginOutputNextExpandIndices - = reinterpret_cast(outputs[getIdx(OutputIdxEntry::OUTPUT_NEXT_EXPAND_INDICES)]); - // Updated scores, shape: [batchSize, mNumEagleLayers, maxDecodingDraftTokens x maxDecodingDraftTokens] - auto pluginOutputAllLayersScores = static_cast(outputs[getIdx(OutputIdxEntry::OUTPUT_ALL_LAYERS_SCORES)]); - // Updated draft tokens, shape: [batchSize, mNumEagleLayers, maxDecodingDraftTokens x maxDecodingDraftTokens] - auto pluginOutputAllLayersDraftTokenIds - = reinterpret_cast(outputs[getIdx(OutputIdxEntry::OUTPUT_ALL_LAYERS_DRAFT_TOKEN_IDS)]); - // Update the predecessor of the draft tokens, shape: [batchSize, mNumEagleLayers, maxDecodingDraftTokens x - // maxDecodingDraftTokens] - auto pluginOutputAllLayersDraftTokenIdsPredecessor - = reinterpret_cast(outputs[getIdx(OutputIdxEntry::OUTPUT_ALL_LAYERS_DRAFT_TOKEN_IDS_PREDECESSOR)]); - - ////////////////////////////////////////// Get workspaces ////////////////////////////////////////// - int8_t* workspaceBytePtr = reinterpret_cast(workspace); - size_t offset{0}; - // Workspace 0: Sampling workspace. - // Treat numInputLogits as batchSize - auto const samplingWorkspaceSize - = getTopKWorkspaceSize(numInputLogits, /* maxTokensPerStep */ 1, /* maxTopK */ maxTopK, vocabSizePadded); - void* workspaceSampling - = reinterpret_cast(tc::nextWorkspacePtr(workspaceBytePtr, offset, samplingWorkspaceSize)); - - // Workspace 1: Topks tensor: shape [numInputLogits] - SizeType32* topKs = reinterpret_cast( - tc::nextWorkspacePtr(workspaceBytePtr, offset, numInputLogits * sizeof(SizeType32))); - - // Workspace 2: topKOffset tensor: shape: [batchSize], number of nodes that have successors for each requests - SizeType32* topKOffset - = reinterpret_cast(tc::nextWorkspacePtr(workspaceBytePtr, offset, batchSize * sizeof(SizeType32))); - - // Workspace 3: logits pointers tensor: shape: [numInputLogits] - T const** logitsPtrs - = reinterpret_cast(tc::nextWorkspacePtr(workspaceBytePtr, offset, numInputLogits * sizeof(T*))); - - // Workspace 4: outputIds pointers tensor: shape [numInputLogits], each points to a [maxDecodingDraftTokens] buffer - TokenIdType** firstTopKOutputIdsPtrs = reinterpret_cast( - tc::nextWorkspacePtr(workspaceBytePtr, offset, numInputLogits * sizeof(TokenIdType*))); - - // Workspace 5: outputIds tensor: flatten outputIds, shape [numInputLogits * maxDecodingDraftTokens] - TokenIdType* firstTopKOutputIdsFlatten = reinterpret_cast( - tc::nextWorkspacePtr(workspaceBytePtr, offset, numInputLogits * maxDecodingDraftTokens * sizeof(TokenIdType))); - - // Workspace 6: number of successors for each nodes tensor: shape [batchSize * maxDecodingTokens] - SizeType32* numSuccessorsForEachNode = reinterpret_cast( - tc::nextWorkspacePtr(workspaceBytePtr, offset, batchSize * maxDecodingTokens * sizeof(SizeType32))); - - // Workspace 7: skip decoding mask [numInputLogits] - bool* skipDecode - = reinterpret_cast(tc::nextWorkspacePtr(workspaceBytePtr, offset, numInputLogits * sizeof(bool))); - - // In Eagle-1, we do not need to return logProbs - float* firstTopKOutputLogProbs = nullptr; - if (useDynamicTree) - { - // Workspace 8. The output logProbs of the first topK sampling. - // Which will be updated with the previous layer's scores (i.e., pluginInputPrevScores), and will be treat as - // the input of the second topK sampling. For mLayerIdx == 0, shape: [numInputLogits(batchSize), - // maxDecodingDraftTokens] For mLayerIdx > 0, shape: [numInputLogits(batchSize * dynamicTreeMaxTopK), - // maxDecodingDraftTokens] - firstTopKOutputLogProbs = reinterpret_cast( - tc::nextWorkspacePtr(workspaceBytePtr, offset, numInputLogits * maxDecodingDraftTokens * sizeof(float))); - } - - SizeType32 const secondTopKVocabSize = dynamicTreeMaxTopK * maxDecodingDraftTokens; - // Workspace 9: Sampling from [batchSize, dynamicTreeMaxTopK * maxDecodingDraftTokens] to [batchSize, - // dynamicTreeMaxTopK] - auto const secondTopKSamplingWorkspaceSize - = getTopKWorkspaceSize(batchSize, /* maxTokensPerStep */ 1, /* maxTopK */ maxTopK, secondTopKVocabSize); - void* workspaceScoresSampling - = reinterpret_cast(tc::nextWorkspacePtr(workspaceBytePtr, offset, secondTopKSamplingWorkspaceSize)); - - // Workspace 10: the second (scores) sampling's outputIds, shape: [batchSize, maxDecodingDraftTokens] - TokenIdType* secondTopKOutputIdsFlatten = reinterpret_cast( - tc::nextWorkspacePtr(workspaceBytePtr, offset, batchSize * maxDecodingDraftTokens * sizeof(TokenIdType))); - - // Workspace 11: the second (scores) sampling's outputIdsPtrs - TokenIdType** secondTopKOutputIdsPtrs = reinterpret_cast( - tc::nextWorkspacePtr(workspaceBytePtr, offset, batchSize * sizeof(TokenIdType*))); - - // Workspace 12: input scores pointers - float** secondTopKInputScoresPtrs - = reinterpret_cast(tc::nextWorkspacePtr(workspaceBytePtr, offset, batchSize * sizeof(float*))); - - // Workspace 13: the second sampling's outputLogProbs - float* secondTopKOutputLogProbs = reinterpret_cast( - tc::nextWorkspacePtr(workspaceBytePtr, offset, batchSize * maxDecodingDraftTokens * sizeof(float))); - - // Workspace 14: The input scores pointers of the third topK sampling, shape [batchSize] - float** thirdTopKInputScoresPtrs - = reinterpret_cast(tc::nextWorkspacePtr(workspaceBytePtr, offset, batchSize * sizeof(float*))); - - // Workspace 15: The output of the third topK sampling, shape [batchSize, maxDecodingDraftTokens] - TokenIdType* thirdTopKOutputIds = reinterpret_cast( - tc::nextWorkspacePtr(workspaceBytePtr, offset, batchSize * maxDecodingDraftTokens * sizeof(TokenIdType))); - - // Workspace 16: The output pointers of the third topK sampling, shape [batchSize] - TokenIdType** thirdTopKOutputIdsPtrs = reinterpret_cast( - tc::nextWorkspacePtr(workspaceBytePtr, offset, batchSize * sizeof(TokenIdType*))); - - // The number of draft tokens among all layers - long const totalNumDraftTokensForAllLayers - = (mNumEagleLayers - 1) * dynamicTreeMaxTopK * dynamicTreeMaxTopK + dynamicTreeMaxTopK; - - auto const thridTopKSamplingWorkspaceSize = getTopKWorkspaceSize( - batchSize, /* maxTokensPerStep */ 1, /* maxTopK */ maxDecodingDraftTokens, totalNumDraftTokensForAllLayers); - // Workspace 17: The workspace of the third topK sampling - void* workspaceThirdTopKSampling - = reinterpret_cast(tc::nextWorkspacePtr(workspaceBytePtr, offset, thridTopKSamplingWorkspaceSize)); - - // Workspace 18. Eagle-2, the topKs for each request in the third topK sampling, shape [batchSize] - // The real topK value is min(maxDecodingDraftTokens, totalNumDraftTokensForAllLayers) - SizeType32* thirdTopKs - = reinterpret_cast(tc::nextWorkspacePtr(workspaceBytePtr, offset, batchSize * sizeof(SizeType32))); - - ////////////////////////////////////////// Main logic ////////////////////////////////////////// - // Fill logitsPtrs from plugin input logits - // And fill firstTopKOutputIdsPtrs from firstTopKOutputIdsFlatten - invokeAssembleDraftLogitsOffsets(logitsPtrs, pluginInputLogits, firstTopKOutputIdsPtrs, firstTopKOutputIdsFlatten, - skipDecode, numValidLogits, numInputLogits, batchSize, maxDecodingDraftTokens, vocabSizePadded, stream); - sync_check_cuda_error(stream); - - if (useDynamicTree) - { - // For Eagle-2, the topK value between different requests are the same, all set to 'dynamicTreeMaxTopK'. - invokeSetTopKsFromDyanmicTreeMaxTopK( - mLayerIdx, batchSize, numInputLogits, topKs, topKOffset, dynamicTreeMaxTopK, numValidLogits, stream); - sync_check_cuda_error(stream); - - // Do softmax for the input logits - // We set the 'batchSize' and 'maxBatchSize' to 'numInputLogits', while 'numInputLogits' logits may contain - // some padding logits, which do not need to be calculated. - // We use 'skipDecode' list to skip these padding logits. This could avoid redundant calculations. - BiasSoftmaxParams biasSoftmaxParams; - biasSoftmaxParams.logits = const_cast(pluginInputLogits); - biasSoftmaxParams.logitsPtrs = nullptr; - biasSoftmaxParams.probs = const_cast(pluginInputLogits); - biasSoftmaxParams.maxBeamWidth = 1; - biasSoftmaxParams.batchSlots = nullptr; - biasSoftmaxParams.batchSize = numInputLogits; - biasSoftmaxParams.maxBatchSize = numInputLogits; - biasSoftmaxParams.vocabSize = vocabSizePadded; - biasSoftmaxParams.vocabSizePadded = vocabSizePadded; - biasSoftmaxParams.skipSoftMax = false; - biasSoftmaxParams.batchSlotsLogits = false; - biasSoftmaxParams.skipDecode = skipDecode; - biasSoftmaxParams.checkParams(); - - invokeAddBiasSoftMax(biasSoftmaxParams, stream); - sync_check_cuda_error(stream); - } - else - { - // For Eagle-1, extract topK value from input path. - invokeExtractTopKsFromPath(pluginInputPaths, topKs, topKOffset, numSuccessorsForEachNode, mLayerIdx, batchSize, - maxDecodingTokens, maxPathLen, stream); - sync_check_cuda_error(stream); - } - - TopKSamplingKernelParams params{}; - params.logProbsPtrs = logitsPtrs; // [numInputLogits][vocabSizePadded] - params.outputIdsPtrs = firstTopKOutputIdsPtrs; // [numInputLogits][maxDecodingDraftTokens] - params.workspace = workspaceSampling; - params.maxTopK = maxTopK; - params.topKs = topKs; // [numInputLogits] - params.batchSize = numInputLogits; - params.maxBatchSize = numInputLogits; - params.maxTokensPerStep = 1; - params.vocabSizePadded = vocabSizePadded; - params.returnAllSelectedTokens = true; - params.strictTopPBoundary = false; - params.skipDecode = skipDecode; - params.outputLogProbs = firstTopKOutputLogProbs; // [numInputLogits * maxDecodingDraftTokens] - params.logitsHasProbs = true; - - invokeBatchTopKSampling(params, stream); - sync_check_cuda_error(stream); - - if (useDynamicTree) - { - // When mLayerIdx == 0, we do not need to update scores. - // We take the outputLogProbs of the first topK sampling as the scores directly. - if (mLayerIdx != 0) - { - // Update firstTopKOutputLogProbs with pluginInputPrevScores, which is the scores from the previous layer - invokeUpdateScores(batchSize, dynamicTreeMaxTopK, maxDecodingDraftTokens, firstTopKOutputLogProbs, - pluginInputPrevScores, stream); - sync_check_cuda_error(stream); - - // Do the second top-dynamicTreeMaxTopK sampling among this dynamicTreeMaxTopK x dynamicTreeMaxTopK draft - // tokens. Through the second topK sampling, we obtain the dynamicTreeMaxTopK output draft tokens of this - // layer. - - // Although theoretically we only need to select 'dynamicTreeMaxTopK' draft tokens from 'dynamicTreeMaxTopK - // * dynamicTreeMaxTopK' draft tokens, we over-set vocabSize here. This is because when we write the scores - // into firstTopKOutputLogProbs, we store it in the form of [batchSize * dynamicTreeMaxTopK, - // maxDecodingDraftTokens]. For each request, these 'dynamicTreeMaxTopK * dynamicTreeMaxTopK' scores are not - // saved continuously, but in the format of [dynamicTreeMaxTopK, maxDecodingDraftTokens]. For unused - // positions, we set '-inf' to ensure that they will not be sampled. Examples: For a request, - // dynamicTreeMaxTopK == 3, the scores in its buffer ([dynamicTreeMaxTopK, maxDecodingDraftTokens]) are as - // follow: - // [[1.1, 2.2, 3.3, -inf, -inf, ...], - // [4.4, 5.5, 6.6, -inf, -inf, ...], - // [7.7, 8.8, 9.9, -inf, -inf, ...]] - - // Prepare the input of the second topK sampling. - invokeAssembleSecondTopKSamplingInputs(batchSize, dynamicTreeMaxTopK, maxDecodingDraftTokens, - firstTopKOutputLogProbs, secondTopKInputScoresPtrs, secondTopKOutputIdsFlatten, secondTopKOutputIdsPtrs, - stream); - sync_check_cuda_error(stream); - - TopKSamplingKernelParams params{}; - params.logProbsPtrs = secondTopKInputScoresPtrs; - params.outputIdsPtrs = secondTopKOutputIdsPtrs; - params.workspace = workspaceScoresSampling; - params.maxTopK = maxTopK; // Same to maxDecodingTokens - params.topKs = topKs; // [batchSize], all set to dynamicTreeMaxTopK - params.batchSize = batchSize; - params.maxBatchSize = batchSize; - params.maxTokensPerStep = 1; - params.vocabSizePadded = secondTopKVocabSize; - params.returnAllSelectedTokens = true; - params.strictTopPBoundary = false; - - invokeBatchTopKSampling(params, stream); - sync_check_cuda_error(stream); - } - - // Copy this layer's scores and draft tokensId: - // 1) Copy this layer's scores to pluginOutputAllLayersScores - // 2) Copy dynamicTreeMaxTopK (or dynamicTreeMaxTopK * dynamicTreeMaxTopK) draft tokens to - // pluginOutputAllLayersDraftTokenIds 3) Set the predecessors of these draft tokens and save to - // pluginOutputAllLayersDraftTokenIdsPredecessor, - // which will be used to reconstruct the final output tree at the last layer - invokeCopyScoresAndDraftTokenIds(mLayerIdx, mNumEagleLayers, maxDecodingDraftTokens, batchSize, - dynamicTreeMaxTopK, - pluginInputCurrentExpandIndices, // The indices of the nodes that expand in this layer (i.e., the input - // logits). The index is related to the final tree. - pluginInputAllLayersScores, pluginInputAllLayersDraftTokenIds, pluginInputAllLayersDraftTokenIdsPredecessor, - pluginOutputAllLayersScores, pluginOutputAllLayersDraftTokenIds, - pluginOutputAllLayersDraftTokenIdsPredecessor, - firstTopKOutputLogProbs, // This layer's scores - firstTopKOutputIdsFlatten, // This layer's draft tokens - stream); - sync_check_cuda_error(stream); - - // Update Path - // For mLayerIdx == 0, the output of the first topK sampling are the output draft tokens of this layers. The - // update logic is simple. For mLayerIdx > 0, the output of the second topK sampling are the output draft tokens - // of this layers. 'secondTopKOutputIdsPtrs' contains the top-dynamicTreeMaxTopK selected from the second topK - // sampling. 'pluginOutputNextExpandIndices' record the selected the top-dynamicTreeMaxTopK draft token's Id of - // this layer, - // which will be used in the next layer to compute the predecessors. - // The last layer will completely reconstruct the paths, so there is no need to update the paths here. - if (mLayerIdx != mNumEagleLayers - 1) - { - invokeUpdatePath(mLayerIdx, batchSize, dynamicTreeMaxTopK, maxDecodingTokens, maxPathLen, pluginInputPaths, - pluginOutputPaths, - secondTopKOutputIdsPtrs, // if mLayerIdx == 0, secondTopKOutputIdsPtrs == nullptr, and it's useless - // during update paths - pluginOutputNextExpandIndices, stream); - sync_check_cuda_error(stream); - } - - if (mLayerIdx != 0) - { - // We will extract the real draft tokenIds and scores from 'firstTopKOutputIdsFlatten' and - // 'secondTopKInputScoresPtrs' according to the 'secondTopKOutputIdsPtrs'. And store them into - // 'secondTopKOutputIdsPtrs' and 'secondTopKOutputLogProbs' (reuse these buffers). - // secondTopKInputScoresPtrs: shape [batchSize * dynamicTreeMaxTopK, maxDecodingDraftTokens] - // The original scores, which were used to do the second TopK sampling - // secondTopKOutputIdsPtrs: shape [batchSize], each points to a [maxDecodingDraftTokens] buffer - // The output of the second TopK sampling, which are the indices of the top-dynamicTreeMaxTopK among - // 'dynamicTreeMaxTopK * dynamicTreeMaxTopK'. We need to figure out what these top-dynamicTreeMaxTopK - // draft tokens' real tokenIds. - // firstTopKOutputIdsFlatten: shape [batchSize * dynamicTreeMaxTopK, maxDecodingDraftTokens] - // The value are related to the vocabSize, which is the real tokenIds. - invokeExtractScoresAndRealDraftTokensIds(batchSize, dynamicTreeMaxTopK, maxDecodingDraftTokens, - secondTopKInputScoresPtrs, secondTopKOutputIdsPtrs, firstTopKOutputIdsFlatten, secondTopKOutputLogProbs, - stream); - sync_check_cuda_error(stream); - } - - // Copy this layer's output draft tokens and scores. - // This layer's output scores is next layer's previous scores. - // if mLayerIdx == 0, directly use the first topK's outputIds / logProbs as this layer's output draft tokens / - // scores if mLayerIdx > 0, we use the second topK's outputIds / logProbs, - // which is updated with the real draft tokenIds / logprobs in 'invokeExtractScoresAndRealDraftTokensIds' - invokeUpdateDraftTokensAndLensAndCurScores(mLayerIdx, batchSize, dynamicTreeMaxTopK, maxDecodingDraftTokens, - mLayerIdx == 0 ? firstTopKOutputIdsPtrs : secondTopKOutputIdsPtrs, pluginInputDraftTokenIds, - pluginInputDraftLens, pluginOutputDraftTokenIds, pluginOutputDraftLens, - mLayerIdx == 0 ? firstTopKOutputLogProbs : secondTopKOutputLogProbs, pluginOutputCurrentScores, stream); - sync_check_cuda_error(stream); - - if (mLayerIdx == mNumEagleLayers - 1) - { - // The maximum number of nodes on the final tree (exclude the root node) - auto const maxNodesOnFinalTree = std::min(maxDecodingDraftTokens, totalNumDraftTokensForAllLayers); - - // When reach the last EagleNet, we need to do the third sampling, which take all layers' draft tokens and - // scores as input, and then select top-maxDecodingDraftTokens draft tokens among them. We need to - // reconstruct the path/tree after the third topK sampling. - invokeAssembleThridTopKSamplingInputs(batchSize, maxDecodingDraftTokens, mNumEagleLayers, - maxNodesOnFinalTree, thirdTopKs, pluginOutputAllLayersScores, thirdTopKInputScoresPtrs, - thirdTopKOutputIds, thirdTopKOutputIdsPtrs, stream); - sync_check_cuda_error(stream); - - // 1) Do topK sampling among all previous draft tokens - TopKSamplingKernelParams params{}; - params.logProbsPtrs = thirdTopKInputScoresPtrs; - params.outputIdsPtrs = thirdTopKOutputIdsPtrs; - params.workspace = workspaceThirdTopKSampling; - params.topKs = thirdTopKs; // All set to 'maxNodesOnFinalTree' - params.maxTopK = maxDecodingDraftTokens; // We set maxTopK to 'maxDecodingDraftTokens' to align the - // outputIdsPtrs offsets when written back. - params.batchSize = batchSize; - params.maxBatchSize = batchSize; - params.maxTokensPerStep = 1; - params.vocabSizePadded = totalNumDraftTokensForAllLayers; - params.returnAllSelectedTokens = true; - params.strictTopPBoundary = false; // Make sure to select topK tokens. - - invokeBatchTopKSampling(params, stream); - sync_check_cuda_error(stream); - - // 2) Reconstruct the Path - invokeReconstructFinalPath(batchSize, dynamicTreeMaxTopK, maxDecodingDraftTokens, maxDecodingTokens, - maxPathLen, mNumEagleLayers, maxNodesOnFinalTree, thirdTopKOutputIdsPtrs, - pluginOutputAllLayersDraftTokenIdsPredecessor, pluginOutputPaths, stream); - sync_check_cuda_error(stream); - - // 3) Copy this layer's outputIds to outputDraftTokenIds - invokeCopyFinalDraftTokens(batchSize, maxDecodingDraftTokens, mNumEagleLayers, maxNodesOnFinalTree, - thirdTopKOutputIdsPtrs, pluginOutputAllLayersDraftTokenIds, pluginOutputDraftTokenIds, - pluginOutputDraftLens, stream); - sync_check_cuda_error(stream); - } - } - else - { - // Eagle-1: Copy output token id from outputIdsPtrs to the plugin output buffer - invokeCopyOutputTokensIds(firstTopKOutputIdsPtrs, topKs, topKOffset, pluginInputDraftTokenIds, - pluginInputDraftLens, numValidLogits, pluginOutputDraftTokenIds, pluginOutputDraftLens, mLayerIdx, - batchSize, maxDecodingDraftTokens, pluginInputPaths, pluginOutputPaths, maxPathLen, stream); - sync_check_cuda_error(stream); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void EagleDecodeDraftTokensPlugin::enqueueType(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) noexcept -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - // TODO split batch into greedy and non-greedy and execute both paths - if (mTopKSampling) - { - doTopKSampling(inputDesc, outputDesc, inputs, outputs, workspace, stream); - } - else - { - // TODO fill me - TLLM_CHECK_WITH_INFO(false, "Multinomial sampling is not supported yet"); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -int EagleDecodeDraftTokensPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) noexcept -{ - auto const logitsType = inputDesc[getIdx(InputIdxEntry::LOGITS)].type; - if (logitsType == nvinfer1::DataType::kFLOAT) - { - enqueueType(inputDesc, outputDesc, inputs, outputs, workspace, stream); - } - else if (logitsType == nvinfer1::DataType::kHALF) - { - enqueueType<__half>(inputDesc, outputDesc, inputs, outputs, workspace, stream); - } - else - { - TLLM_CHECK_WITH_INFO(false, "Unsupported logits type"); - } - - return 0; -} - -// IPluginV2Ext Methods -nvinfer1::DataType EagleDecodeDraftTokensPlugin::getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept -{ - TLLM_CHECK(index < getNbOutputs()); - TLLM_CHECK(index < getNbOutputs()); - if (index == getIdx(OutputIdxEntry::OUTPUT_ALL_LAYERS_SCORES) - || index == getIdx(OutputIdxEntry::OUTPUT_CURRENT_SCORES)) - { - // Only output_prev_socres are float - return inputTypes[getIdx(InputIdxEntry::INPUT_ALL_LAYERS_SCORES)]; - } - else - { - // output_draft_token_ids, output_draft_lens, output_paths, output_next_expand_index, - // output_all_layers_draft_token_ids, output_all_layers_draft_token_ids_predecessor - // are all int32 type, same as path - return inputTypes[getIdx(InputIdxEntry::PATHS)]; - } -} - -// IPluginV2 Methods - -char const* EagleDecodeDraftTokensPlugin::getPluginType() const noexcept -{ - return EAGLE_DECODE_DRAFT_TOKENS_PLUGIN_NAME; -} - -char const* EagleDecodeDraftTokensPlugin::getPluginVersion() const noexcept -{ - return EAGLE_DECODE_DRAFT_TOKENS_PLUGIN_VERSION; -} - -int EagleDecodeDraftTokensPlugin::getNbOutputs() const noexcept -{ - return 8; -} - -int EagleDecodeDraftTokensPlugin::initialize() noexcept -{ - return 0; -} - -void EagleDecodeDraftTokensPlugin::terminate() noexcept {} - -size_t EagleDecodeDraftTokensPlugin::getSerializationSize() const noexcept -{ - return sizeof(mDtype) + sizeof(mLayerIdx) + sizeof(mNumEagleLayers) + sizeof(mTopKSampling); -} - -void EagleDecodeDraftTokensPlugin::serialize(void* buffer) const noexcept -{ - char *d = static_cast(buffer), *a = d; - write(d, mDtype); - write(d, mLayerIdx); - write(d, mNumEagleLayers); - write(d, mTopKSampling); - TLLM_CHECK(d == a + getSerializationSize()); -} - -void EagleDecodeDraftTokensPlugin::destroy() noexcept -{ - // This gets called when the network containing plugin is destroyed - delete this; -} - -/////////////// - -EagleDecodeDraftTokensPluginCreator::EagleDecodeDraftTokensPluginCreator() -{ - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - mPluginAttributes.emplace_back(PluginField("type_id", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("layer_idx", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("num_eagle_layers", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("top_k_sampling", nullptr, PluginFieldType::kINT32)); - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -char const* EagleDecodeDraftTokensPluginCreator::getPluginName() const noexcept -{ - return EAGLE_DECODE_DRAFT_TOKENS_PLUGIN_NAME; -} - -char const* EagleDecodeDraftTokensPluginCreator::getPluginVersion() const noexcept -{ - return EAGLE_DECODE_DRAFT_TOKENS_PLUGIN_VERSION; -} - -PluginFieldCollection const* EagleDecodeDraftTokensPluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -IPluginV2* EagleDecodeDraftTokensPluginCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept -{ - PluginField const* fields = fc->fields; - int32_t layerIdx{}; - int32_t numEagleLayers{}; - nvinfer1::DataType type{}; - bool topKSampling{}; - // Read configurations from each fields - for (int i = 0; i < fc->nbFields; ++i) - { - char const* attrName = fields[i].name; - if (!strcmp(attrName, "layer_idx")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - layerIdx = *static_cast(fields[i].data); - } - else if (!strcmp(attrName, "num_eagle_layers")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - numEagleLayers = *static_cast(fields[i].data); - } - else if (!strcmp(attrName, "type_id")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - type = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "top_k_sampling")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - topKSampling = static_cast(*static_cast(fields[i].data)); - } - } - - try - { - auto* obj = new EagleDecodeDraftTokensPlugin(type, layerIdx, numEagleLayers, topKSampling); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -IPluginV2* EagleDecodeDraftTokensPluginCreator::deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept -{ - // This object will be deleted when the network is destroyed, which will - // call EagleDecodeDraftTokensPlugin::destroy() - try - { - auto* obj = new EagleDecodeDraftTokensPlugin(serialData, serialLength); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} diff --git a/cpp/tensorrt_llm/plugins/eaglePlugin/eagleDecodeDraftTokensPlugin.h b/cpp/tensorrt_llm/plugins/eaglePlugin/eagleDecodeDraftTokensPlugin.h deleted file mode 100644 index 8c144a1bc073..000000000000 --- a/cpp/tensorrt_llm/plugins/eaglePlugin/eagleDecodeDraftTokensPlugin.h +++ /dev/null @@ -1,174 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-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. - */ -#pragma once - -#include "tensorrt_llm/plugins/common/plugin.h" -#include -#include -#include -#include - -namespace tensorrt_llm::plugins -{ - -class EagleDecodeDraftTokensPlugin : public BasePlugin -{ -public: - EagleDecodeDraftTokensPlugin(nvinfer1::DataType type, int32_t layerIdx, int32_t numEagleLayers, bool topKSampling); - - EagleDecodeDraftTokensPlugin(void const* data, size_t length); - - ~EagleDecodeDraftTokensPlugin() override = default; - - // IPluginV2DynamicExt Methods - nvinfer1::IPluginV2DynamicExt* clone() const noexcept override; - nvinfer1::DimsExprs getOutputDimensions(int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; - bool supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept override; - void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept override; - size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept override; - int enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - - // IPluginV2Ext Methods - nvinfer1::DataType getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept override; - - // IPluginV2 Methods - char const* getPluginType() const noexcept override; - char const* getPluginVersion() const noexcept override; - int getNbOutputs() const noexcept override; - int initialize() noexcept override; - void terminate() noexcept override; - size_t getSerializationSize() const noexcept override; - void serialize(void* buffer) const noexcept override; - void destroy() noexcept override; - -private: - enum class InputIdxEntry : int32_t - { - // 12 inputs - // [num_input_logits, vocab_size_padded] - LOGITS = 0, - // [batch_size, max_decoding_tokens, max_path_len] - PATHS, - // [1] - NUM_VALID_LOGITS, - // [1] - USE_DYNAMIC_TREE, - // [1] - DYNAMIC_TREE_MAX_TOPK, - - // [batch_size, max_decoding_draft_tokens] - INPUT_DRAFT_TOKEN_IDS, - // [batch_size] - INPUT_DRAFT_LENS, - - // [batch_size, max_decoding_draft_tokens] - INPUT_PREV_SCORES, - - // [batch_size, max_decoding_draft_tokens] - INPUT_CURRENT_EXPAND_INDICES, - - // [batch_size, num_eagle_layers, max_decoding_draft_tokens x max_decoding_draft_tokens] - INPUT_ALL_LAYERS_SCORES, - // [batch_size, num_eagle_layers, max_decoding_draft_tokens x max_decoding_draft_tokens] - INPUT_ALL_LAYERS_DRAFT_TOKEN_IDS, - // [batch_size, num_eagle_layers, max_decoding_draft_tokens x max_decoding_draft_tokens] - INPUT_ALL_LAYERS_DRAFT_TOKEN_IDS_PREDECESSOR - }; - - enum class OutputIdxEntry : int32_t - { - // 8 outputs - // [batch_size, max_decoding_draft_tokens] - OUTPUT_DRAFT_TOKEN_IDS = 0, - // [batch_size] - OUTPUT_DRAFT_LENS, - - // [batch_size, max_decoding_tokens, max_path_len] - OUTPUT_PATHS, - - // [batch_size, max_decoding_draft_tokens] - OUTPUT_CURRENT_SCORES, - - // [batch_size, max_decoding_draft_tokens] - OUTPUT_NEXT_EXPAND_INDICES, - - // [batch_size, num_eagle_layers, max_decoding_draft_tokens x max_decoding_draft_tokens] - OUTPUT_ALL_LAYERS_SCORES, - // [batch_size, num_eagle_layers, max_decoding_draft_tokens x max_decoding_draft_tokens] - OUTPUT_ALL_LAYERS_DRAFT_TOKEN_IDS, - // [batch_size, num_eagle_layers, max_decoding_draft_tokens x max_decoding_draft_tokens] - OUTPUT_ALL_LAYERS_DRAFT_TOKEN_IDS_PREDECESSOR - }; - - int32_t getIdx(InputIdxEntry idx) const - { - return static_cast(idx); - } - - int32_t getIdx(OutputIdxEntry idx) const - { - return static_cast(idx); - } - -private: - template - size_t getWorkspaceSizeType(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept; - - template - void enqueueType(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept; - - template - void doTopKSampling(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept; - -private: - nvinfer1::DataType mDtype; // Logit datatype - int32_t mLayerIdx{-1}; // Index of eagle layer - int32_t mNumEagleLayers{-1}; // Number of eagle layers - bool mTopKSampling; // Use TopK sampling or multinomial sampling -}; - -class EagleDecodeDraftTokensPluginCreator : public BaseCreator -{ -public: - EagleDecodeDraftTokensPluginCreator(); - - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; - - nvinfer1::IPluginV2* deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept override; - -private: - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; -}; - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/eaglePlugin/eaglePrepareDrafterInputsPlugin.cpp b/cpp/tensorrt_llm/plugins/eaglePlugin/eaglePrepareDrafterInputsPlugin.cpp deleted file mode 100644 index 2cd8c695e296..000000000000 --- a/cpp/tensorrt_llm/plugins/eaglePlugin/eaglePrepareDrafterInputsPlugin.cpp +++ /dev/null @@ -1,548 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-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. - */ -#include "eaglePrepareDrafterInputsPlugin.h" - -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/dataType.h" -#include "tensorrt_llm/common/memoryUtils.h" -#include "tensorrt_llm/kernels/speculativeDecoding/eagleDecodingKernels.h" -#include "tensorrt_llm/runtime/common.h" -#include "tensorrt_llm/runtime/iTensor.h" - -using namespace nvinfer1; -using tensorrt_llm::plugins::EaglePrepareDrafterInputsPluginCreator; -using tensorrt_llm::plugins::EaglePrepareDrafterInputsPlugin; -using namespace tensorrt_llm::kernels; -using namespace tensorrt_llm::kernels::speculative_decoding; -using namespace tensorrt_llm::runtime; -namespace tc = tensorrt_llm::common; - -static char const* EAGLE_PREPARE_DRAFTER_INPUTS_PLUGIN_VERSION{"1"}; -static char const* EAGLE_PREPARE_DRAFTER_INPUTS_PLUGIN_NAME{"EaglePrepareDrafterInputs"}; -PluginFieldCollection EaglePrepareDrafterInputsPluginCreator::mFC{}; -std::vector EaglePrepareDrafterInputsPluginCreator::mPluginAttributes; - -EaglePrepareDrafterInputsPlugin::EaglePrepareDrafterInputsPlugin( - int32_t layerIdx, int32_t numLayers, int32_t maxNonLeavesPerLayer) - : mLayerIdx(layerIdx) - , mNumLayers(numLayers) - , mMaxNonLeavesPerLayer(maxNonLeavesPerLayer) -{ -} - -void EaglePrepareDrafterInputsPlugin::initFieldsToSerialize() -{ - mDataToSerialize.clear(); - mDataToSerialize.emplace_back(PluginField("layer_idx", &mLayerIdx, PluginFieldType::kINT32, 1)); - mDataToSerialize.emplace_back(PluginField("num_layers", &mNumLayers, PluginFieldType::kINT32, 1)); - mDataToSerialize.emplace_back( - PluginField("max_non_leaves_per_layer", &mMaxNonLeavesPerLayer, PluginFieldType::kINT32, 1)); - mFCToSerialize.nbFields = mDataToSerialize.size(); - mFCToSerialize.fields = mDataToSerialize.data(); -} - -nvinfer1::IPluginCapability* EaglePrepareDrafterInputsPlugin::getCapabilityInterface( - nvinfer1::PluginCapabilityType type) noexcept -{ - try - { - if (type == nvinfer1::PluginCapabilityType::kBUILD) - { - return static_cast(this); - } - if (type == nvinfer1::PluginCapabilityType::kRUNTIME) - { - return static_cast(this); - } - TLLM_CHECK(type == nvinfer1::PluginCapabilityType::kCORE); - return static_cast(this); - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -// IPluginV3 methods -nvinfer1::IPluginV3* EaglePrepareDrafterInputsPlugin::clone() noexcept -{ - auto clone = std::make_unique(*this); - clone->initFieldsToSerialize(); - return clone.release(); -} - -// IPluginV3OneCore methods -char const* EaglePrepareDrafterInputsPlugin::getPluginName() const noexcept -{ - return EAGLE_PREPARE_DRAFTER_INPUTS_PLUGIN_NAME; -} - -char const* EaglePrepareDrafterInputsPlugin::getPluginVersion() const noexcept -{ - return EAGLE_PREPARE_DRAFTER_INPUTS_PLUGIN_VERSION; -} - -char const* EaglePrepareDrafterInputsPlugin::getPluginNamespace() const noexcept -{ - return tensorrt_llm::plugins::api::kDefaultNamespace; -} - -// IPluginV3OneBuild methods -int32_t EaglePrepareDrafterInputsPlugin::getNbOutputs() const noexcept -{ - return 11; -} - -int32_t EaglePrepareDrafterInputsPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int32_t nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int32_t nbOutputs) noexcept -{ - return 0; -} - -bool EaglePrepareDrafterInputsPlugin::supportsFormatCombination( - int32_t pos, nvinfer1::DynamicPluginTensorDesc const* inOut, int32_t nbInputs, int32_t nbOutputs) noexcept -{ - return (inOut[pos].desc.type == nvinfer1::DataType::kINT32) && (inOut[pos].desc.format == TensorFormat::kLINEAR); -} - -int32_t EaglePrepareDrafterInputsPlugin::getOutputDataTypes(nvinfer1::DataType* outputTypes, int32_t nbOutputs, - nvinfer1::DataType const* inputTypes, int32_t nbInputs) const noexcept -{ - outputTypes[0] = nvinfer1::DataType::kINT32; - outputTypes[1] = nvinfer1::DataType::kINT32; - outputTypes[2] = nvinfer1::DataType::kINT32; - outputTypes[3] = nvinfer1::DataType::kINT32; - outputTypes[4] = nvinfer1::DataType::kINT32; - outputTypes[5] = nvinfer1::DataType::kINT32; - outputTypes[6] = nvinfer1::DataType::kINT32; - outputTypes[7] = nvinfer1::DataType::kINT32; - outputTypes[8] = nvinfer1::DataType::kINT32; - outputTypes[9] = nvinfer1::DataType::kINT32; - outputTypes[10] = nvinfer1::DataType::kINT32; - outputTypes[11] = nvinfer1::DataType::kINT32; - return 0; -} - -int32_t EaglePrepareDrafterInputsPlugin::getOutputShapes(nvinfer1::DimsExprs const* inputs, int32_t nbInputs, - nvinfer1::DimsExprs const* shapeInputs, int32_t nbShapeInputs, nvinfer1::DimsExprs* outputs, int32_t nbOutputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept -{ - TLLM_CHECK(nbOutputs == 11); - TLLM_CHECK(nbInputs == 15); - TLLM_CHECK(nbShapeInputs == 0); - auto const numTokens = inputs[getIdx(InputIdxEntry::INPUT_IDS)].d[0]; - auto const batchSizeExpr = inputs[getIdx(InputIdxEntry::PREV_DRAFT_PATHS)].d[0]; - auto const numGenRequestsExpr = inputs[getIdx(InputIdxEntry::SPEC_DECODING_GENERATION_LENGTHS)].d[0]; - auto const numInputGenTokensExpr = inputs[getIdx(InputIdxEntry::INPUT_GEN_TOKENS)].d[0]; - auto const maxDecodingLenExpr = inputs[getIdx(InputIdxEntry::PREV_DRAFT_PATHS)].d[1]; - auto const maxPathLenExpr = inputs[getIdx(InputIdxEntry::PREV_DRAFT_PATHS)].d[2]; - - for (SizeType32 outputIndex = 0; outputIndex < nbOutputs; ++outputIndex) - { - if (outputIndex == getIdx(OutputIdxEntry::SEQUENCE_LENGTHS) - || outputIndex == getIdx(OutputIdxEntry::CONTEXT_LENGTHS) - || outputIndex == getIdx(OutputIdxEntry::SPEC_DECODING_GENERATION_LENGTHS)) - { - outputs[outputIndex] = inputs[getIdx(InputIdxEntry::SEQUENCE_LENGTHS)]; - } - else if (outputIndex == getIdx(OutputIdxEntry::SPEC_DECODING_PACKED_MASK)) - { - outputs[outputIndex].nbDims = 3; - outputs[outputIndex].d[0] = batchSizeExpr; - outputs[outputIndex].d[1] = maxDecodingLenExpr; - outputs[outputIndex].d[2] - = exprBuilder.operation(DimensionOperation::kCEIL_DIV, *maxDecodingLenExpr, *exprBuilder.constant(32)); - } - else if (outputIndex == getIdx(OutputIdxEntry::SPEC_DECODING_POSITION_OFFSETS)) - { - outputs[outputIndex].nbDims = 2; - outputs[outputIndex].d[0] = batchSizeExpr; - outputs[outputIndex].d[1] = maxDecodingLenExpr; - } - else if (outputIndex == getIdx(OutputIdxEntry::OUTPUT_IDS) - || outputIndex == getIdx(OutputIdxEntry::HIDDEN_STATES_INDICES) - || (mLayerIdx == 0 && outputIndex == getIdx(OutputIdxEntry::POSITION_IDS))) - { - if (mLayerIdx == 0) - { - // We have at most numGenRequests * (mNumLayers + 1) accepted tokens per step for gen requests and - // input_ids - numGenTokens tokens for context requests. - auto numOutputGenTokensExpr = exprBuilder.operation( - DimensionOperation::kPROD, *numGenRequestsExpr, *exprBuilder.constant(mNumLayers + 1)); - auto numInputCtxTokensExpr - = exprBuilder.operation(DimensionOperation::kSUB, *numTokens, *numInputGenTokensExpr); - outputs[outputIndex].nbDims = 1; - outputs[outputIndex].d[0] = exprBuilder.operation(DimensionOperation::kMAX, *exprBuilder.constant(1), - *exprBuilder.operation(DimensionOperation::kSUM, *numOutputGenTokensExpr, *numInputCtxTokensExpr)); - } - else - { - // At most we have mMaxNonLeavesPerLayer non-leaves at this layer. - // And in total we pass all non-leaves + all their preceding nodes. - // batchSize * mMaxNonLeavesPerLayer * layerIdx - outputs[outputIndex].nbDims = 1; - outputs[outputIndex].d[0] = exprBuilder.operation(DimensionOperation::kPROD, - *exprBuilder.operation(DimensionOperation::kPROD, *exprBuilder.constant(mLayerIdx), - *exprBuilder.constant(mMaxNonLeavesPerLayer)), - *batchSizeExpr); - } - } - else if (mLayerIdx > 0 && outputIndex == getIdx(OutputIdxEntry::POSITION_IDS)) - { - outputs[outputIndex].nbDims = 1; - outputs[outputIndex].d[0] = batchSizeExpr; - } - else if (outputIndex == getIdx(OutputIdxEntry::LAST_TOKEN_INDICES)) - { - outputs[outputIndex].nbDims = 1; - outputs[outputIndex].d[0] = exprBuilder.operation( - DimensionOperation::kPROD, *exprBuilder.constant(mMaxNonLeavesPerLayer), *batchSizeExpr); - } - else if (outputIndex == getIdx(OutputIdxEntry::NUM_LAST_TOKEN_INDICES)) - { - outputs[outputIndex].nbDims = 1; - outputs[outputIndex].d[0] = exprBuilder.constant(1); - } - else if (outputIndex == getIdx(OutputIdxEntry::HIDDEN_SIZE_BATCH_LEVEL_STARTS)) - { - // batchSize * (maxPathLen - 1) + 1 - outputs[outputIndex].nbDims = 1; - outputs[outputIndex].d[0] = exprBuilder.operation(DimensionOperation::kSUM, *exprBuilder.constant(1), - *exprBuilder.operation(DimensionOperation::kPROD, *batchSizeExpr, - *exprBuilder.operation(DimensionOperation::kSUB, *maxPathLenExpr, *exprBuilder.constant(1)))); - } - } - return 0; -} - -int32_t EaglePrepareDrafterInputsPlugin::onShapeChange(nvinfer1::PluginTensorDesc const* in, int32_t nbInputs, - nvinfer1::PluginTensorDesc const* out, int32_t nbOutputs) noexcept -{ - return 0; -} - -nvinfer1::IPluginV3* EaglePrepareDrafterInputsPlugin::attachToContext( - nvinfer1::IPluginResourceContext* context) noexcept -{ - return clone(); -} - -PluginFieldCollection const* EaglePrepareDrafterInputsPlugin::getFieldsToSerialize() noexcept -{ - return &mFCToSerialize; -} - -size_t EaglePrepareDrafterInputsPlugin::getWorkspaceSize(nvinfer1::DynamicPluginTensorDesc const* inputs, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - size_t workspaceSize{0}; - - auto const batchSize = inputs[getIdx(InputIdxEntry::NEXT_DRAFT_PATHS)].max.d[0]; - auto const maxDecodingTokens = inputs[getIdx(InputIdxEntry::NEXT_DRAFT_PATHS)].max.d[1]; - - if (mLayerIdx > 0) - { - SizeType32 constexpr NUM_BUFFERS{9}; - size_t workspaces[NUM_BUFFERS]; - workspaces[0] = batchSize * maxDecodingTokens * sizeof(int8_t); // isLeafMask - workspaces[1] = batchSize * maxDecodingTokens * sizeof(SizeType32); // selectedDraftIndices - workspaces[2] = batchSize * maxDecodingTokens * sizeof(SizeType32); // selectedDraftPosOffsets - workspaces[3] = batchSize * sizeof(SizeType32); // numSelectedDraftIndices - workspaces[4] = batchSize * maxDecodingTokens * maxDecodingTokens * sizeof(int8_t); // selectedMasks - workspaces[5] = (batchSize + 1) * sizeof(SizeType32); // cumSumGenerationLengths - workspaces[6] = batchSize * maxDecodingTokens * sizeof(SizeType32); // nonLeavesInLevelOffsets - workspaces[7] = batchSize * maxDecodingTokens * sizeof(SizeType32); // parentNonLeafInLevelOffset - workspaces[8] = 1 * sizeof(SizeType32); // maxGenerationLength - workspaceSize = tc::calculateTotalWorkspaceSize(workspaces, NUM_BUFFERS); - } - - return workspaceSize; -} - -void EaglePrepareDrafterInputsPlugin::prepareCtxEagleNetData(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) noexcept -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - auto const batchSize = inputDesc[getIdx(InputIdxEntry::SEQUENCE_LENGTHS)].dims.d[0]; - - auto const numTokens = inputDesc[getIdx(InputIdxEntry::INPUT_IDS)].dims.d[0]; - auto const numGenRequests = inputDesc[getIdx(InputIdxEntry::SPEC_DECODING_GENERATION_LENGTHS)].dims.d[0]; - auto const numInputGenTokens = inputDesc[getIdx(InputIdxEntry::INPUT_GEN_TOKENS)].dims.d[0]; - - auto const maxPathLen = inputDesc[getIdx(InputIdxEntry::ACCEPTED_TOKENS)].dims.d[1]; - auto const maxDecodingTokens = inputDesc[getIdx(InputIdxEntry::NEXT_DRAFT_PATHS)].dims.d[1]; - - auto eagleNetSequenceLengths = reinterpret_cast(outputs[getIdx(OutputIdxEntry::SEQUENCE_LENGTHS)]); - auto eagleNetContextLengths = reinterpret_cast(outputs[getIdx(OutputIdxEntry::CONTEXT_LENGTHS)]); - auto outputIds = reinterpret_cast(outputs[getIdx(OutputIdxEntry::OUTPUT_IDS)]); - auto positionIds = reinterpret_cast(outputs[getIdx(OutputIdxEntry::POSITION_IDS)]); - auto hiddenStatesIndices = reinterpret_cast(outputs[getIdx(OutputIdxEntry::HIDDEN_STATES_INDICES)]); - auto lastTokenIndices = reinterpret_cast(outputs[getIdx(OutputIdxEntry::LAST_TOKEN_INDICES)]); - auto numLastTokenIndices = reinterpret_cast(outputs[getIdx(OutputIdxEntry::NUM_LAST_TOKEN_INDICES)]); - auto hiddenSizeBatchLevelStarts - = reinterpret_cast(outputs[getIdx(OutputIdxEntry::HIDDEN_SIZE_BATCH_LEVEL_STARTS)]); - - auto inputIds = reinterpret_cast(inputs[getIdx(InputIdxEntry::INPUT_IDS)]); - auto chunkedContextNextTokens - = reinterpret_cast(inputs[getIdx(InputIdxEntry::CHUNKED_CONTEXT_NEXT_TOKENS)]); - auto baseNetSequenceLengths = reinterpret_cast(inputs[getIdx(InputIdxEntry::SEQUENCE_LENGTHS)]); - auto baseNetContextLengths = reinterpret_cast(inputs[getIdx(InputIdxEntry::CONTEXT_LENGTHS)]); - auto acceptedTokens = reinterpret_cast(inputs[getIdx(InputIdxEntry::ACCEPTED_TOKENS)]); - auto acceptedLens = reinterpret_cast(inputs[getIdx(InputIdxEntry::ACCEPTED_LENS)]); - auto prevDraftLens = reinterpret_cast(inputs[getIdx(InputIdxEntry::PREV_DRAFT_LENS)]); - auto prevPaths = reinterpret_cast(inputs[getIdx(InputIdxEntry::PREV_DRAFT_PATHS)]); - auto bestPathIds = reinterpret_cast(inputs[getIdx(InputIdxEntry::ACCEPTED_PATHS)]); - - auto const numOutputTokens = (numTokens - numInputGenTokens) + (numGenRequests * (mNumLayers + 1)); - cudaMemsetAsync(positionIds, 0, numOutputTokens * sizeof(SizeType32), stream); - cudaMemsetAsync(hiddenStatesIndices, 0, numOutputTokens * sizeof(SizeType32), stream); - - invokePrepareCtxEagleNetInputs(eagleNetSequenceLengths, eagleNetContextLengths, outputIds, positionIds, - hiddenStatesIndices, lastTokenIndices, numLastTokenIndices, hiddenSizeBatchLevelStarts, inputIds, - chunkedContextNextTokens, baseNetSequenceLengths, baseNetContextLengths, acceptedTokens, acceptedLens, - prevDraftLens, prevPaths, bestPathIds, batchSize, maxPathLen, maxDecodingTokens, mMaxNonLeavesPerLayer, stream); - - sync_check_cuda_error(stream); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void EaglePrepareDrafterInputsPlugin::prepareGenEagleNetData(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) noexcept -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - auto const batchSize = inputDesc[getIdx(InputIdxEntry::SEQUENCE_LENGTHS)].dims.d[0]; - auto const maxDecodingTokens = inputDesc[getIdx(InputIdxEntry::NEXT_DRAFT_PATHS)].dims.d[1]; - auto const maxPathLen = inputDesc[getIdx(InputIdxEntry::NEXT_DRAFT_PATHS)].dims.d[2]; - - auto eagleNetSequenceLengths = reinterpret_cast(outputs[getIdx(OutputIdxEntry::SEQUENCE_LENGTHS)]); - auto eagleNetContextLengths = reinterpret_cast(outputs[getIdx(OutputIdxEntry::CONTEXT_LENGTHS)]); - auto outputIds = reinterpret_cast(outputs[getIdx(OutputIdxEntry::OUTPUT_IDS)]); - auto positionIds = reinterpret_cast(outputs[getIdx(OutputIdxEntry::POSITION_IDS)]); - auto specDecodingGenLengths - = reinterpret_cast(outputs[getIdx(OutputIdxEntry::SPEC_DECODING_GENERATION_LENGTHS)]); - auto specDecodingPositionOffsets - = reinterpret_cast(outputs[getIdx(OutputIdxEntry::SPEC_DECODING_POSITION_OFFSETS)]); - auto specDecodingPackedMasks - = reinterpret_cast(outputs[getIdx(OutputIdxEntry::SPEC_DECODING_PACKED_MASK)]); - auto hiddenStatesIndices = reinterpret_cast(outputs[getIdx(OutputIdxEntry::HIDDEN_STATES_INDICES)]); - auto lastTokenIndices = reinterpret_cast(outputs[getIdx(OutputIdxEntry::LAST_TOKEN_INDICES)]); - auto numLastTokenIndices = reinterpret_cast(outputs[getIdx(OutputIdxEntry::NUM_LAST_TOKEN_INDICES)]); - auto outputHiddenSizeBatchStartsPerLevel - = reinterpret_cast(outputs[getIdx(OutputIdxEntry::HIDDEN_SIZE_BATCH_LEVEL_STARTS)]); - - auto eagleNet0SequenceLengths - = reinterpret_cast(inputs[getIdx(InputIdxEntry::SEQUENCE_LENGTHS)]); - auto eagleNet0ContextLength = reinterpret_cast(inputs[getIdx(InputIdxEntry::CONTEXT_LENGTHS)]); - auto nextDraftPaths = reinterpret_cast(inputs[getIdx(InputIdxEntry::NEXT_DRAFT_PATHS)]); - auto nextDraftIds = reinterpret_cast(inputs[getIdx(InputIdxEntry::NEXT_DRAFT_TOKENS)]); - auto inputHiddenSizeBatchStartsPerLevel - = reinterpret_cast(inputs[getIdx(InputIdxEntry::HIDDEN_SIZE_BATCH_LEVEL_STARTS)]); - - int8_t* workspaceBytePtr = reinterpret_cast(workspace); - size_t offset{0}; - - int8_t* isLeafMask = reinterpret_cast( - tc::nextWorkspacePtr(workspaceBytePtr, offset, batchSize * maxDecodingTokens * sizeof(int8_t))); - TokenIdType* selectedDraftIndices = reinterpret_cast( - tc::nextWorkspacePtr(workspaceBytePtr, offset, batchSize * maxDecodingTokens * sizeof(TokenIdType))); - SizeType32* selectedDraftPosOffsets = reinterpret_cast( - tc::nextWorkspacePtr(workspaceBytePtr, offset, batchSize * maxDecodingTokens * sizeof(SizeType32))); - SizeType32* numSelectedDraftIndices - = reinterpret_cast(tc::nextWorkspacePtr(workspaceBytePtr, offset, batchSize * sizeof(SizeType32))); - bool* selectedMasks = reinterpret_cast(tc::nextWorkspacePtr( - workspaceBytePtr, offset, batchSize * maxDecodingTokens * maxDecodingTokens * sizeof(int8_t))); - SizeType32* cumSumGenerationLengths = reinterpret_cast( - tc::nextWorkspacePtr(workspaceBytePtr, offset, (batchSize + 1) * sizeof(SizeType32))); - SizeType32* nonLeavesInLevelOffsets = reinterpret_cast( - tc::nextWorkspacePtr(workspaceBytePtr, offset, batchSize * maxDecodingTokens * sizeof(SizeType32))); - SizeType32* parentNonLeafInLevelOffset = reinterpret_cast( - tc::nextWorkspacePtr(workspaceBytePtr, offset, batchSize * maxDecodingTokens * sizeof(SizeType32))); - SizeType32* maxGenerationLength - = reinterpret_cast(tc::nextWorkspacePtr(workspaceBytePtr, offset, 1 * sizeof(SizeType32))); - - cudaMemsetAsync(hiddenStatesIndices, 0, batchSize * mMaxNonLeavesPerLayer * mLayerIdx * sizeof(SizeType32), stream); - cudaMemsetAsync(selectedMasks, 0, batchSize * maxDecodingTokens * maxDecodingTokens * sizeof(int8_t), stream); - // Prefill mask setting all to leaves. - cudaMemsetAsync(isLeafMask, 1, batchSize * maxDecodingTokens * sizeof(int8_t), stream); - - PrepareGenEagleNetInputsParams params; - params.nextSequenceLengths = eagleNetSequenceLengths; - params.nextContextLengths = eagleNetContextLengths; - params.outputIds = outputIds; - params.positionIds = positionIds; - params.specDecodingGenLengths = specDecodingGenLengths; - params.specDecodingPositionOffsets = specDecodingPositionOffsets; - params.specDecodingPackedMasks = specDecodingPackedMasks; - params.hiddenStatesIndices = hiddenStatesIndices; - params.lastTokenIndices = lastTokenIndices; - params.numLastTokenIndices = numLastTokenIndices; - params.outputHiddenSizeBatchStartsPerLevel = outputHiddenSizeBatchStartsPerLevel; - - // tmp data - params.isLeafMask = isLeafMask; - params.selectedDraftIndices = selectedDraftIndices; - params.selectedDraftPosOffsets = selectedDraftPosOffsets; - params.numSelectedDraftIndices = numSelectedDraftIndices; - params.selectedMasks = selectedMasks; - params.cumSumGenerationLengths = cumSumGenerationLengths; - params.maxGenerationLength = maxGenerationLength; - params.nonLeavesInLevelOffsets = nonLeavesInLevelOffsets; - params.parentNonLeafInLevelOffset = parentNonLeafInLevelOffset; - - params.nextDraftIds = nextDraftIds; - params.eagleNet0SequenceLengths = eagleNet0SequenceLengths; - params.prevContextLengths = eagleNet0ContextLength; - params.nextPaths = nextDraftPaths; - params.inputHiddenSizeBatchStartsPerLevel = inputHiddenSizeBatchStartsPerLevel; - params.levelIdx = mLayerIdx; - params.batchSize = batchSize; - params.maxPathLen = maxPathLen; - params.maxDecodingTokens = maxDecodingTokens; - params.maxNonLeavesPerLayer = mMaxNonLeavesPerLayer; - params.stream = stream; - - params.checkParams(); - - invokePrepareGenEagleNetInputs(params); - - sync_check_cuda_error(stream); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -int EaglePrepareDrafterInputsPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) noexcept -{ - // First EagleNet instance (EagleNet0) is always chunked context attn, - // where we process either context tokens or newly accepted tokens and append them to EagleNet KV cache. - - // For all following EagleNetX (X > 0) instances there is need for masked spec decoding attn. - // Ideally with mask for context. - // Let's say we have prompt ABCD and two variants of tokens spec decoding tokens E and F - // predicted by EagleNet0. If we draw full attn mask, it becomes: - // |A|B|C|D|E|F - // E|1|1|1|1|1|0 - // F|1|1|1|1|0|1 - // - // In the next step we predict token G from ABCDE branch and token H from ABCDF branch -- like beam search. - // And we'd need spec decoding mask that includes kv cache: - // |A|B|C|D|E|F|G|H - // G|1|1|1|1|1|0|1|0 - // H|1|1|1|1|0|1|0|1 - // - // But TRT-LLM does not support such mask for now. We can only provide - // |G|H - // G|1|0 - // H|0|1 - // , which is wrong mask. - // - // For now we WAR this by passing EFGH for the EagleNet1 with right mask - // and using only G and H logits for sampling, but that's redundant compute: - // |E|F|G|H - // E|1|0|0|0 - // F|0|1|0|0 - // G|1|0|1|0 - // H|0|1|0|1 - - if (mLayerIdx == 0) - { - prepareCtxEagleNetData(inputDesc, outputDesc, inputs, outputs, workspace, stream); - } - else - { - prepareGenEagleNetData(inputDesc, outputDesc, inputs, outputs, workspace, stream); - } - - return 0; -} - -/////////////// - -EaglePrepareDrafterInputsPluginCreator::EaglePrepareDrafterInputsPluginCreator() -{ - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - mPluginAttributes.emplace_back(PluginField("layer_idx", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("num_layers", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("max_non_leaves_per_layer", nullptr, PluginFieldType::kINT32)); - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -char const* EaglePrepareDrafterInputsPluginCreator::getPluginName() const noexcept -{ - return EAGLE_PREPARE_DRAFTER_INPUTS_PLUGIN_NAME; -} - -char const* EaglePrepareDrafterInputsPluginCreator::getPluginVersion() const noexcept -{ - return EAGLE_PREPARE_DRAFTER_INPUTS_PLUGIN_VERSION; -} - -PluginFieldCollection const* EaglePrepareDrafterInputsPluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -nvinfer1::IPluginV3* EaglePrepareDrafterInputsPluginCreator::createPlugin( - char const* name, nvinfer1::PluginFieldCollection const* fc, nvinfer1::TensorRTPhase phase) noexcept -{ - try - { - int32_t layerIdx{0}; - int32_t numLayers{0}; - int32_t maxNonLeavesPerLayer{0}; - // Read configurations from each fields - for (int i = 0; i < fc->nbFields; ++i) - { - char const* attrName = fc->fields[i].name; - if (!strcmp(attrName, "layer_idx")) - { - TLLM_CHECK(fc->fields[i].type == PluginFieldType::kINT32); - layerIdx = *static_cast(fc->fields[i].data); - } - else if (!strcmp(attrName, "num_layers")) - { - TLLM_CHECK(fc->fields[i].type == PluginFieldType::kINT32); - numLayers = *static_cast(fc->fields[i].data); - } - else if (!strcmp(attrName, "max_non_leaves_per_layer")) - { - TLLM_CHECK(fc->fields[i].type == PluginFieldType::kINT32); - maxNonLeavesPerLayer = *static_cast(fc->fields[i].data); - } - } - return new EaglePrepareDrafterInputsPlugin(layerIdx, numLayers, maxNonLeavesPerLayer); - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -char const* EaglePrepareDrafterInputsPluginCreator::getPluginNamespace() const noexcept -{ - return tensorrt_llm::plugins::api::kDefaultNamespace; -} diff --git a/cpp/tensorrt_llm/plugins/eaglePlugin/eaglePrepareDrafterInputsPlugin.h b/cpp/tensorrt_llm/plugins/eaglePlugin/eaglePrepareDrafterInputsPlugin.h deleted file mode 100644 index 0059c46f6c8d..000000000000 --- a/cpp/tensorrt_llm/plugins/eaglePlugin/eaglePrepareDrafterInputsPlugin.h +++ /dev/null @@ -1,186 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-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. - */ -#pragma once - -#include "tensorrt_llm/plugins/common/plugin.h" -#include -#include -#include -#include - -namespace tensorrt_llm::plugins -{ - -class EaglePrepareDrafterInputsPlugin : public nvinfer1::IPluginV3, - public nvinfer1::IPluginV3OneCore, - public nvinfer1::IPluginV3OneBuild, - public nvinfer1::IPluginV3OneRuntime -{ -public: - EaglePrepareDrafterInputsPlugin(EaglePrepareDrafterInputsPlugin const& p) = default; - - EaglePrepareDrafterInputsPlugin(int32_t layerIdx, int32_t numLayers, int32_t maxNonLeavesPerLayer); - - nvinfer1::IPluginV3* clone() noexcept override; - - nvinfer1::IPluginCapability* getCapabilityInterface(nvinfer1::PluginCapabilityType type) noexcept override; - - void initFieldsToSerialize(); - - char const* getPluginName() const noexcept override; - char const* getPluginVersion() const noexcept override; - char const* getPluginNamespace() const noexcept override; - - int32_t getNbOutputs() const noexcept override; - - bool supportsFormatCombination( - int pos, nvinfer1::DynamicPluginTensorDesc const* inOut, int32_t nbInputs, int32_t nbOutputs) noexcept override; - int32_t configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int32_t nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int32_t nbOutputs) noexcept override; - - int32_t getOutputDataTypes(nvinfer1::DataType* outputTypes, int32_t nbOutputs, nvinfer1::DataType const* inputTypes, - int32_t nbInputs) const noexcept override; - - int32_t getOutputShapes(nvinfer1::DimsExprs const* inputs, int32_t nbInputs, nvinfer1::DimsExprs const* shapeInputs, - int32_t nbShapeInputs, nvinfer1::DimsExprs* outputs, int32_t nbOutputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; - - int32_t onShapeChange(nvinfer1::PluginTensorDesc const* in, int32_t nbInputs, nvinfer1::PluginTensorDesc const* out, - int32_t nbOutputs) noexcept override; - - nvinfer1::IPluginV3* attachToContext(nvinfer1::IPluginResourceContext* context) noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldsToSerialize() noexcept override; - - size_t getWorkspaceSize(nvinfer1::DynamicPluginTensorDesc const* inputs, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* outputs, int nbOutputs) const noexcept override; - int enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - -private: - enum class InputIdxEntry : int32_t - { - //! [batch_size] - SEQUENCE_LENGTHS = 0, - //! [batch_size] - CONTEXT_LENGTHS, - //! [num_tokens] - INPUT_IDS, - //! [batch_size] - CHUNKED_CONTEXT_NEXT_TOKENS, - //! [batch_size, max_path_len] - ACCEPTED_TOKENS, - //! [batch_size] - ACCEPTED_LENS, - //! [batch_size] - ACCEPTED_PATHS, - //! [batch_size, max_decoding_draft_tokens] - NEXT_DRAFT_TOKENS, - //! [batch_size] - NEXT_DRAFT_LENS, - //! [batch_size, max_decoding_tokens, max_path_len] - NEXT_DRAFT_PATHS, - //! [batch_size] - PREV_DRAFT_LENS, - //! [batch_size, max_decoding_tokens, max_path_len] - PREV_DRAFT_PATHS, - //! [(max_path_len - 1) * batch_size + 1] - HIDDEN_SIZE_BATCH_LEVEL_STARTS, - //! [num_gen_tokens] - INPUT_GEN_TOKENS, - //! [num_gen_requests] - SPEC_DECODING_GENERATION_LENGTHS, - }; - - enum class OutputIdxEntry : int32_t - { - //! [batch_size] - SEQUENCE_LENGTHS = 0, - //! [batch_size] - CONTEXT_LENGTHS, - //! [batch_size] - SPEC_DECODING_GENERATION_LENGTHS, - //! [batch_size, max_decoding_tokens] - SPEC_DECODING_POSITION_OFFSETS, - //! [batchSize, maxDecodingTokens, ceil(maxDecodingTokens / 32)] - SPEC_DECODING_PACKED_MASK, - //! [batchSize * mMaxNonLeavesPerLayer * layerIdx] for layerIdx > 0 - //! [num_tokens - numGenTokens + numGenRequests * (mNumLayers + 1)] for layerIdx == 0 - OUTPUT_IDS, - //! [batchSize] for layerIdx > 0 - //! [num_tokens - numGenTokens + numGenRequests * (mNumLayers + 1)] for layerIdx == 0 - POSITION_IDS, - //! [batchSize * mMaxNonLeavesPerLayer * layerIdx] for layerIdx > 0 - //! [num_tokens - numGenTokens + numGenRequests * (mNumLayers + 1)] for layerIdx == 0 - HIDDEN_STATES_INDICES, - //! [batchSize * mMaxNonLeavesPerLayer] - LAST_TOKEN_INDICES, - //! [1] - NUM_LAST_TOKEN_INDICES, - //! [(max_path_len - 1) * batch_size + 1] - HIDDEN_SIZE_BATCH_LEVEL_STARTS, - }; - - int32_t getIdx(InputIdxEntry idx) const - { - return static_cast(idx); - } - - int32_t getIdx(OutputIdxEntry idx) const - { - return static_cast(idx); - } - -private: - void prepareCtxEagleNetData(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) noexcept; - - void prepareGenEagleNetData(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) noexcept; - -private: - int32_t mLayerIdx{0}; - int32_t mNumLayers{0}; - int32_t mMaxNonLeavesPerLayer{0}; - std::vector mDataToSerialize; - nvinfer1::PluginFieldCollection mFCToSerialize; -}; - -class EaglePrepareDrafterInputsPluginCreator : public nvinfer1::IPluginCreatorV3One -{ -public: - EaglePrepareDrafterInputsPluginCreator(); - - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - char const* getPluginNamespace() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV3* createPlugin( - char const* name, nvinfer1::PluginFieldCollection const* fc, nvinfer1::TensorRTPhase phase) noexcept override; - -private: - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; -}; - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/eaglePlugin/eagleSampleAndAcceptDraftTokensPlugin.cpp b/cpp/tensorrt_llm/plugins/eaglePlugin/eagleSampleAndAcceptDraftTokensPlugin.cpp deleted file mode 100644 index 5fb30f583712..000000000000 --- a/cpp/tensorrt_llm/plugins/eaglePlugin/eagleSampleAndAcceptDraftTokensPlugin.cpp +++ /dev/null @@ -1,565 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-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. - */ -#include "eagleSampleAndAcceptDraftTokensPlugin.h" - -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/cudaUtils.h" -#include "tensorrt_llm/common/dataType.h" -#include "tensorrt_llm/common/memoryUtils.h" -#include "tensorrt_llm/kernels/samplingTopKKernels.h" -#include "tensorrt_llm/kernels/speculativeDecoding/common.h" -#include "tensorrt_llm/kernels/speculativeDecoding/eagleDecodingKernels.h" -#include "tensorrt_llm/kernels/speculativeDecoding/medusaDecodingKernels.h" -#include "tensorrt_llm/runtime/common.h" -#include "tensorrt_llm/runtime/iTensor.h" - -using namespace nvinfer1; -using tensorrt_llm::plugins::EagleSampleAndAcceptDraftTokensPluginCreator; -using tensorrt_llm::plugins::EagleSampleAndAcceptDraftTokensPlugin; -using namespace tensorrt_llm::kernels; -using namespace tensorrt_llm::kernels::speculative_decoding; -using namespace tensorrt_llm::runtime; -namespace tc = tensorrt_llm::common; - -static char const* EAGLE_SAMPLE_AND_ACCEPT_DRAFT_TOKENS_PLUGIN_VERSION{"1"}; -static char const* EAGLE_SAMPLE_AND_ACCEPT_DRAFT_TOKENS_PLUGIN_NAME{"EagleSampleAndAcceptDraftTokens"}; -PluginFieldCollection EagleSampleAndAcceptDraftTokensPluginCreator::mFC{}; -std::vector EagleSampleAndAcceptDraftTokensPluginCreator::mPluginAttributes; - -EagleSampleAndAcceptDraftTokensPlugin::EagleSampleAndAcceptDraftTokensPlugin(nvinfer1::DataType type) - : mDtype(type) -{ -} - -// Parameterized constructor -EagleSampleAndAcceptDraftTokensPlugin::EagleSampleAndAcceptDraftTokensPlugin(void const* data, size_t length) -{ - char const *d = reinterpret_cast(data), *a = d; - read(d, mDtype); - TLLM_CHECK_WITH_INFO(d == a + length, - "Expected length (%d) != real length (%d). This is often " - "caused by using different TensorRT LLM version to build " - "engine and run engine.", - (int) length, (int) (d - a)); -} - -// IPluginV2DynamicExt Methods -nvinfer1::IPluginV2DynamicExt* EagleSampleAndAcceptDraftTokensPlugin::clone() const noexcept -{ - auto* plugin = new EagleSampleAndAcceptDraftTokensPlugin(*this); - plugin->setPluginNamespace(mNamespace.c_str()); - return plugin; -} - -nvinfer1::DimsExprs EagleSampleAndAcceptDraftTokensPlugin::getOutputDimensions( - int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept -{ - TLLM_CHECK(nbInputs == 10); - TLLM_CHECK(outputIndex < 7); - auto const batchSizeExpr = inputs[getIdx(InputIdxEntry::PATHS)].d[0]; - auto const maxDecodingDraftTokensExpr = inputs[getIdx(InputIdxEntry::DRAFT_TOKEN_IDS)].d[1]; - auto const maxDecodingTokensExpr = inputs[getIdx(InputIdxEntry::PATHS)].d[1]; - auto const maxPathLenExpr = inputs[getIdx(InputIdxEntry::PATHS)].d[2]; - - nvinfer1::DimsExprs ret; - if (outputIndex == getIdx(OutputIdxEntry::ACCEPTED_TOKENS)) - { - ret.nbDims = 2; - ret.d[0] = batchSizeExpr; - ret.d[1] = maxPathLenExpr; - } - else if (outputIndex == getIdx(OutputIdxEntry::ACCEPTED_LENS)) - { - ret.nbDims = 1; - ret.d[0] = batchSizeExpr; - } - else if (outputIndex == getIdx(OutputIdxEntry::BEST_ACCEPTED_PATHS)) - { - ret.nbDims = 1; - ret.d[0] = batchSizeExpr; - } - else if (outputIndex == getIdx(OutputIdxEntry::NEXT_DRAFT_TOKEN_IDS)) - { - ret.nbDims = 2; - ret.d[0] = batchSizeExpr; - ret.d[1] = maxDecodingDraftTokensExpr; - } - else if (outputIndex == getIdx(OutputIdxEntry::NEXT_DRAFT_LENS)) - { - ret.nbDims = 1; - ret.d[0] = batchSizeExpr; - } - else if (outputIndex == getIdx(OutputIdxEntry::NEXT_DRAFT_PATHS)) - { - ret.nbDims = 3; - ret.d[0] = batchSizeExpr; - ret.d[1] = maxDecodingTokensExpr; - ret.d[2] = maxPathLenExpr; - } - else if (outputIndex == getIdx(OutputIdxEntry::HIDDEN_SIZE_BATCH_LEVEL_STARTS)) - { - ret.nbDims = 1; - ret.d[0] = exprBuilder.operation(DimensionOperation::kSUM, *exprBuilder.constant(1), - *exprBuilder.operation(DimensionOperation::kPROD, - *exprBuilder.operation(DimensionOperation::kSUB, *maxPathLenExpr, *exprBuilder.constant(1)), - *batchSizeExpr)); - } - return ret; -} - -bool EagleSampleAndAcceptDraftTokensPlugin::supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept -{ - if (pos == getIdx(InputIdxEntry::LOGITS)) // logits - { - return (inOut[pos].type == mDtype) && (inOut[pos].format == TensorFormat::kLINEAR); - } - else if (pos == getIdx(InputIdxEntry::TEMPERATURE) || pos == getIdx(InputIdxEntry::RAND_VALIDATION) - || pos == getIdx(InputIdxEntry::POSTERIOR_ALPHA) - || pos == getIdx(InputIdxEntry::POSTERIOR_THRESHOLD)) // temperature, rand_validation - { - return (inOut[pos].type == nvinfer1::DataType::kFLOAT) && (inOut[pos].format == TensorFormat::kLINEAR); - } - else // everything else - { - return (inOut[pos].type == nvinfer1::DataType::kINT32) && (inOut[pos].format == TensorFormat::kLINEAR); - } -} - -void EagleSampleAndAcceptDraftTokensPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept -{ -} - -template -size_t EagleSampleAndAcceptDraftTokensPlugin::getWorkspaceSizeType(nvinfer1::PluginTensorDesc const* inputs, - int nbInputs, nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - size_t workspaceSize{0}; - - auto const vocabSizePadded = inputs[getIdx(InputIdxEntry::LOGITS)].dims.d[1]; - auto const batchSize = inputs[getIdx(InputIdxEntry::PATHS)].dims.d[0]; - auto const maxDecodingTokens = inputs[getIdx(InputIdxEntry::PATHS)].dims.d[1]; - - // Greedy sampling - // Top1 sampling workspace - auto const greedySamplingWorkspaceSize - = getTopKWorkspaceSize(batchSize, maxDecodingTokens, /* maxTopK */ 1, vocabSizePadded); - - // Multinomial sampling - auto const typicalSamplingWorkspaceSize - = getTypicalAcceptanceWorkspaceSize(batchSize, maxDecodingTokens, vocabSizePadded); - - auto const primarySamplingWorkspaceSize = std::max(greedySamplingWorkspaceSize, typicalSamplingWorkspaceSize); - - // Target output ids - auto const targetOutputIdsSize = batchSize * maxDecodingTokens * sizeof(TokenIdType); - // Logits ptrs - auto const logitsPtrsSize = batchSize * maxDecodingTokens * sizeof(T*); - SizeType32 constexpr NUM_BUFFERS{4}; - size_t workspaces[NUM_BUFFERS]; - workspaces[0] = targetOutputIdsSize; - workspaces[1] = primarySamplingWorkspaceSize; - workspaces[2] = logitsPtrsSize; - workspaces[3] = batchSize * sizeof(SizeType32); - workspaceSize = tc::calculateTotalWorkspaceSize(workspaces, NUM_BUFFERS); - - return workspaceSize; -} - -size_t EagleSampleAndAcceptDraftTokensPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - auto const logitsType = inputs[getIdx(InputIdxEntry::LOGITS)].type; - if (logitsType == nvinfer1::DataType::kFLOAT) - { - return getWorkspaceSizeType(inputs, nbInputs, outputs, nbOutputs); - } - else if (logitsType == nvinfer1::DataType::kHALF) - { - return getWorkspaceSizeType<__half>(inputs, nbInputs, outputs, nbOutputs); - } - else - { - TLLM_CHECK_WITH_INFO(false, "Unsupported logits type"); - } - return 0; -} - -template -void EagleSampleAndAcceptDraftTokensPlugin::samplePrimeHeadTokens(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) noexcept -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - // auto const maxNumTokens = inputDesc[getIdx(InputIdxEntry::LOGITS)].dims.d[0]; - auto const vocabSizePadded = inputDesc[getIdx(InputIdxEntry::LOGITS)].dims.d[1]; - auto const batchSize = inputDesc[getIdx(InputIdxEntry::PATHS)].dims.d[0]; - auto const maxDecodingTokens = inputDesc[getIdx(InputIdxEntry::PATHS)].dims.d[1]; - - auto logits = static_cast(inputs[getIdx(InputIdxEntry::LOGITS)]); - auto prevDraftLens = reinterpret_cast(inputs[getIdx(InputIdxEntry::DRAFT_LENS)]); - - int8_t* workspaceBytePtr = reinterpret_cast(workspace); - size_t offset{0}; - - auto const samplingWorkspaceSize - = getTopKWorkspaceSize(batchSize, maxDecodingTokens, /* maxTopK */ 1, vocabSizePadded); - - TokenIdType* outputIds = reinterpret_cast( - tc::nextWorkspacePtr(workspaceBytePtr, offset, batchSize * maxDecodingTokens * sizeof(TokenIdType))); - void* workspaceSampling - = reinterpret_cast(tc::nextWorkspacePtr(workspaceBytePtr, offset, samplingWorkspaceSize)); - T const** logitsPtrs = reinterpret_cast( - tc::nextWorkspacePtr(workspaceBytePtr, offset, batchSize * maxDecodingTokens * sizeof(T*))); - SizeType32* decodingTokens - = reinterpret_cast(tc::nextWorkspacePtr(workspaceBytePtr, offset, batchSize * sizeof(SizeType32))); - - // Assemble pointers to logits - invokeAssembleTargetLogitsOffsets( - logitsPtrs, decodingTokens, logits, prevDraftLens, batchSize, maxDecodingTokens, vocabSizePadded, stream); - - sync_check_cuda_error(stream); - - TopKSamplingKernelParams params; - params.logProbsPtrs = logitsPtrs; - params.outputIds = outputIds; - params.workspace = workspaceSampling; - params.maxTopK = 1; - params.batchSize = batchSize; - params.maxBatchSize = batchSize; - params.tokensPerStep = decodingTokens; - params.maxTokensPerStep = maxDecodingTokens; - params.maxSeqLen = maxDecodingTokens; - params.vocabSizePadded = vocabSizePadded; - - invokeBatchTopKSampling(params, stream); - - sync_check_cuda_error(stream); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void EagleSampleAndAcceptDraftTokensPlugin::doTypicalAcceptance(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) noexcept -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - // auto const maxNumTokens = inputDesc[getIdx(InputIdxEntry::LOGITS)].dims.d[0]; - auto const vocabSizePadded = inputDesc[getIdx(InputIdxEntry::LOGITS)].dims.d[1]; - - auto const batchSize = inputDesc[getIdx(InputIdxEntry::PATHS)].dims.d[0]; - auto const maxDecodingTokens = inputDesc[getIdx(InputIdxEntry::PATHS)].dims.d[1]; - // auto const maxPathLen = inputDesc[getIdx(InputIdxEntry::PATHS)].dims.d[2]; - // auto const maxDraftPathLen = maxPathLen - 1; - - auto logits = static_cast(inputs[getIdx(InputIdxEntry::LOGITS)]); - auto prevDraftLens = reinterpret_cast(inputs[getIdx(InputIdxEntry::DRAFT_LENS)]); - - int8_t* workspaceBytePtr = reinterpret_cast(workspace); - size_t offset{0}; - - // Multinomial sampling - auto const primarySamplingWorkspaceSize - = getTypicalAcceptanceWorkspaceSize(batchSize, maxDecodingTokens, vocabSizePadded); - - TokenIdType* outputIds = reinterpret_cast( - tc::nextWorkspacePtr(workspaceBytePtr, offset, batchSize * maxDecodingTokens * sizeof(TokenIdType))); - void* workspaceSampling - = reinterpret_cast(tc::nextWorkspacePtr(workspaceBytePtr, offset, primarySamplingWorkspaceSize)); - T** logitsPtrs = reinterpret_cast( - tc::nextWorkspacePtr(workspaceBytePtr, offset, batchSize * maxDecodingTokens * sizeof(T*))); - SizeType32* decodingTokens - = reinterpret_cast(tc::nextWorkspacePtr(workspaceBytePtr, offset, batchSize * sizeof(SizeType32))); - - // Assemble pointers to logits - invokeAssembleTargetLogitsOffsets(const_cast(logitsPtrs), decodingTokens, logits, prevDraftLens, - batchSize, maxDecodingTokens, vocabSizePadded, stream); - - sync_check_cuda_error(stream); - - TypicalAcceptanceSampling params; - params.logitsPtrs = logitsPtrs; - params.generationLengths = decodingTokens; - params.temperatures = reinterpret_cast(inputs[getIdx(InputIdxEntry::TEMPERATURE)]); - params.posteriorThresholds = reinterpret_cast(inputs[getIdx(InputIdxEntry::POSTERIOR_THRESHOLD)]); - params.posteriorAlphas = reinterpret_cast(inputs[getIdx(InputIdxEntry::POSTERIOR_ALPHA)]); - params.outputIds = outputIds; - params.workspace = reinterpret_cast(workspaceSampling); - params.randomVals = reinterpret_cast(inputs[getIdx(InputIdxEntry::RAND_VALIDATION)]); - - params.batchSize = batchSize; - params.maxBatchSize = batchSize; - params.maxDecodingTokens = maxDecodingTokens; - params.vocabSize = vocabSizePadded; - - if (mSmCnt <= 0) - { - auto const deviceId = tensorrt_llm::common::getDevice(); - cudaDeviceProp prop{}; - TLLM_CUDA_CHECK(cudaGetDeviceProperties(&prop, deviceId)); - mSmCnt = prop.multiProcessorCount; - } - params.smCnt = mSmCnt; - - params.checkParams(); - - typicalAcceptanceSampling(params, stream); - - sync_check_cuda_error(stream); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void EagleSampleAndAcceptDraftTokensPlugin::acceptDraftTokens(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) noexcept -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - // auto const maxNumTokens = inputDesc[getIdx(InputIdxEntry::LOGITS)].dims.d[0]; - auto const vocabSizePadded = inputDesc[getIdx(InputIdxEntry::LOGITS)].dims.d[1]; - - auto const batchSize = inputDesc[getIdx(InputIdxEntry::PATHS)].dims.d[0]; - auto const maxDecodingTokens = inputDesc[getIdx(InputIdxEntry::PATHS)].dims.d[1]; - auto const maxPathLen = inputDesc[getIdx(InputIdxEntry::PATHS)].dims.d[2]; - auto const maxDraftPathLen = maxPathLen - 1; - - auto const useDynamicTree = *(reinterpret_cast(inputs[getIdx(InputIdxEntry::USE_DYNAMIC_TREE)])); - - int8_t* workspaceBytePtr = reinterpret_cast(workspace); - size_t offset{0}; - - // auto const samplingWorkspaceSize - // = getTopKWorkspaceSize(batchSize, maxDecodingTokens, /* maxTopK */ 1, vocabSizePadded); - - TokenIdType* outputIds = reinterpret_cast( - tc::nextWorkspacePtr(workspaceBytePtr, offset, batchSize * maxDecodingTokens * sizeof(TokenIdType))); - - AcceptDraftTokensByIdsWithPathsParams params; - params.outputIds = reinterpret_cast(outputs[getIdx(OutputIdxEntry::ACCEPTED_TOKENS)]); - params.draftIds = reinterpret_cast(inputs[getIdx(InputIdxEntry::DRAFT_TOKEN_IDS)]); - params.targetIds = outputIds; - params.acceptedLengths = reinterpret_cast(outputs[getIdx(OutputIdxEntry::ACCEPTED_LENS)]); - params.paths = reinterpret_cast(inputs[getIdx(InputIdxEntry::PATHS)]); - params.bestPathIds = reinterpret_cast(outputs[getIdx(OutputIdxEntry::BEST_ACCEPTED_PATHS)]); - params.batchSize = batchSize; - params.maxBatchSize = batchSize; - params.vocabSize = vocabSizePadded; - params.maxSeqLen = maxPathLen; - params.maxDraftPathLen = maxDraftPathLen; - params.maxDecodingTokens = maxDecodingTokens; - params.stream = stream; - - params.checkParams(); - - acceptDraftTokensByIdsWithPaths(params); - - if (useDynamicTree) - { - // For Eagle-2, after verification and acceptance, the original path becomes useless. - // All set to '-1' - cudaMemsetAsync(outputs[getIdx(OutputIdxEntry::NEXT_DRAFT_PATHS)], -1, - batchSize * maxDecodingTokens * maxPathLen * sizeof(SizeType32), stream); - } - else - { - // For Eagle-1 - // Copy input paths to the output - cudaMemcpyAsync(outputs[getIdx(OutputIdxEntry::NEXT_DRAFT_PATHS)], inputs[getIdx(InputIdxEntry::PATHS)], - batchSize * maxDecodingTokens * maxPathLen * sizeof(SizeType32), cudaMemcpyDeviceToDevice, stream); - } - - sync_check_cuda_error(stream); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void EagleSampleAndAcceptDraftTokensPlugin::enqueueType(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) noexcept -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - auto const greedySampling = reinterpret_cast(inputs[getIdx(InputIdxEntry::GREEDY_SAMPLING)])[0]; - // TODO split batch into greedy and non-greedy and execute both paths - if (greedySampling) - { - // Sample all main head tokens with Top-1. - samplePrimeHeadTokens(inputDesc, outputDesc, inputs, outputs, workspace, stream); - } - else - { - // Typical sampling for typical acceptance. - doTypicalAcceptance(inputDesc, outputDesc, inputs, outputs, workspace, stream); - } - - // Accept tokens based on token ids, write the best path and best token id. - acceptDraftTokens(inputDesc, outputDesc, inputs, outputs, workspace, stream); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -int EagleSampleAndAcceptDraftTokensPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) noexcept -{ - auto const logitsType = inputDesc[getIdx(InputIdxEntry::LOGITS)].type; - if (logitsType == nvinfer1::DataType::kFLOAT) - { - enqueueType(inputDesc, outputDesc, inputs, outputs, workspace, stream); - } - else if (logitsType == nvinfer1::DataType::kHALF) - { - enqueueType<__half>(inputDesc, outputDesc, inputs, outputs, workspace, stream); - } - else - { - TLLM_CHECK_WITH_INFO(false, "Unsupported logits type"); - } - - return 0; -} - -// IPluginV2Ext Methods -nvinfer1::DataType EagleSampleAndAcceptDraftTokensPlugin::getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept -{ - TLLM_CHECK(index < 7); - // input 1 is draft tokens now of int32 type. All outputs are int32_t as well. - return inputTypes[getIdx(InputIdxEntry::DRAFT_TOKEN_IDS)]; -} - -// IPluginV2 Methods - -char const* EagleSampleAndAcceptDraftTokensPlugin::getPluginType() const noexcept -{ - return EAGLE_SAMPLE_AND_ACCEPT_DRAFT_TOKENS_PLUGIN_NAME; -} - -char const* EagleSampleAndAcceptDraftTokensPlugin::getPluginVersion() const noexcept -{ - return EAGLE_SAMPLE_AND_ACCEPT_DRAFT_TOKENS_PLUGIN_VERSION; -} - -int EagleSampleAndAcceptDraftTokensPlugin::getNbOutputs() const noexcept -{ - return 7; -} - -int EagleSampleAndAcceptDraftTokensPlugin::initialize() noexcept -{ - return 0; -} - -void EagleSampleAndAcceptDraftTokensPlugin::terminate() noexcept {} - -size_t EagleSampleAndAcceptDraftTokensPlugin::getSerializationSize() const noexcept -{ - return sizeof(mDtype); -} - -void EagleSampleAndAcceptDraftTokensPlugin::serialize(void* buffer) const noexcept -{ - char *d = static_cast(buffer), *a = d; - write(d, mDtype); - TLLM_CHECK(d == a + getSerializationSize()); -} - -void EagleSampleAndAcceptDraftTokensPlugin::destroy() noexcept -{ - // This gets called when the network containing plugin is destroyed - delete this; -} - -/////////////// - -EagleSampleAndAcceptDraftTokensPluginCreator::EagleSampleAndAcceptDraftTokensPluginCreator() -{ - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - mPluginAttributes.emplace_back(PluginField("type_id", nullptr, PluginFieldType::kINT32)); - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -char const* EagleSampleAndAcceptDraftTokensPluginCreator::getPluginName() const noexcept -{ - return EAGLE_SAMPLE_AND_ACCEPT_DRAFT_TOKENS_PLUGIN_NAME; -} - -char const* EagleSampleAndAcceptDraftTokensPluginCreator::getPluginVersion() const noexcept -{ - return EAGLE_SAMPLE_AND_ACCEPT_DRAFT_TOKENS_PLUGIN_VERSION; -} - -PluginFieldCollection const* EagleSampleAndAcceptDraftTokensPluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -IPluginV2* EagleSampleAndAcceptDraftTokensPluginCreator::createPlugin( - char const* name, PluginFieldCollection const* fc) noexcept -{ - PluginField const* fields = fc->fields; - nvinfer1::DataType type{}; - // Read configurations from each fields - for (int i = 0; i < fc->nbFields; ++i) - { - char const* attrName = fields[i].name; - if (!strcmp(attrName, "type_id")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - type = static_cast(*(static_cast(fields[i].data))); - } - } - - try - { - auto* obj = new EagleSampleAndAcceptDraftTokensPlugin(type); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -IPluginV2* EagleSampleAndAcceptDraftTokensPluginCreator::deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept -{ - // This object will be deleted when the network is destroyed, which will - // call EagleSampleAndAcceptDraftTokensPlugin::destroy() - try - { - auto* obj = new EagleSampleAndAcceptDraftTokensPlugin(serialData, serialLength); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} diff --git a/cpp/tensorrt_llm/plugins/eaglePlugin/eagleSampleAndAcceptDraftTokensPlugin.h b/cpp/tensorrt_llm/plugins/eaglePlugin/eagleSampleAndAcceptDraftTokensPlugin.h deleted file mode 100644 index 3b14bab83170..000000000000 --- a/cpp/tensorrt_llm/plugins/eaglePlugin/eagleSampleAndAcceptDraftTokensPlugin.h +++ /dev/null @@ -1,167 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-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. - */ -#pragma once - -#include "tensorrt_llm/plugins/common/plugin.h" - -#include -#include -#include -#include -#include - -namespace tensorrt_llm::plugins -{ - -class EagleSampleAndAcceptDraftTokensPlugin : public BasePlugin -{ -public: - EagleSampleAndAcceptDraftTokensPlugin(nvinfer1::DataType type); - - EagleSampleAndAcceptDraftTokensPlugin(void const* data, size_t length); - - ~EagleSampleAndAcceptDraftTokensPlugin() override = default; - - // IPluginV2DynamicExt Methods - nvinfer1::IPluginV2DynamicExt* clone() const noexcept override; - nvinfer1::DimsExprs getOutputDimensions(int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; - bool supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept override; - void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept override; - size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept override; - int enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - - // IPluginV2Ext Methods - nvinfer1::DataType getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept override; - - // IPluginV2 Methods - char const* getPluginType() const noexcept override; - char const* getPluginVersion() const noexcept override; - int getNbOutputs() const noexcept override; - int initialize() noexcept override; - void terminate() noexcept override; - size_t getSerializationSize() const noexcept override; - void serialize(void* buffer) const noexcept override; - void destroy() noexcept override; - -private: - enum class InputIdxEntry : int32_t - { - //! [num_tokens, vocab_size_padded] - LOGITS = 0, - //! [batch_size, max_decoding_draft_tokens] - DRAFT_TOKEN_IDS, - //! [batch_size] - DRAFT_LENS, - //! [batch_size] - TEMPERATURE, - //! [batch_size, max_decoding_tokens] - RAND_VALIDATION, - //! [batch_size] - POSTERIOR_ALPHA, - //! [batch_size] - POSTERIOR_THRESHOLD, - //! [batch_size, max_decoding_tokens, max_path_len] - PATHS, - //! [1] - GREEDY_SAMPLING, - //! [1] - USE_DYNAMIC_TREE - }; - - enum class OutputIdxEntry : int32_t - { - //! [batch_size, max_path_len] - ACCEPTED_TOKENS = 0, - //! [batch_size] - ACCEPTED_LENS, - //! [batch_size] - BEST_ACCEPTED_PATHS, - //! [batch_size, max_decoding_draft_tokens] - NEXT_DRAFT_TOKEN_IDS, - //! [batch_size] - NEXT_DRAFT_LENS, - //! [batch_size, max_decoding_tokens, max_path_len] - NEXT_DRAFT_PATHS, - //! [max_draft_path_len * batch_size] - HIDDEN_SIZE_BATCH_LEVEL_STARTS, - }; - - int32_t getIdx(InputIdxEntry idx) const - { - return static_cast(idx); - } - - int32_t getIdx(OutputIdxEntry idx) const - { - return static_cast(idx); - } - -private: - template - size_t getWorkspaceSizeType(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept; - - template - void samplePrimeHeadTokens(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) noexcept; - - template - void doTypicalAcceptance(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept; - - template - void acceptDraftTokens(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept; - - template - void enqueueType(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept; - -private: - nvinfer1::DataType mDtype; - int32_t mSmCnt{0}; -}; - -class EagleSampleAndAcceptDraftTokensPluginCreator : public BaseCreator -{ -public: - EagleSampleAndAcceptDraftTokensPluginCreator(); - - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; - - nvinfer1::IPluginV2* deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept override; - -private: - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; -}; - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/exports.def b/cpp/tensorrt_llm/plugins/exports.def deleted file mode 100644 index 5d4ac9e3e793..000000000000 --- a/cpp/tensorrt_llm/plugins/exports.def +++ /dev/null @@ -1,19 +0,0 @@ -; SPDX-FileCopyrightText: Copyright (c) 1993-2023 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. - -LIBRARY nvinfer_plugin_tensorrt_llm -EXPORTS -getPluginRegistry -initLibNvInferPlugins diff --git a/cpp/tensorrt_llm/plugins/exports.map b/cpp/tensorrt_llm/plugins/exports.map deleted file mode 100644 index c6c949775079..000000000000 --- a/cpp/tensorrt_llm/plugins/exports.map +++ /dev/null @@ -1,34 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. - */ - -/* Hides all symbols except those specified in the global section */ -{ - global: - initTrtLlmPlugins; - setLoggerFinder; - getPluginCreators; - getCreators; - extern "C++" { - nvinfer1::IPluginCreator::*; - nvinfer1::IPluginV2Ext::*; - nvinfer1::IPluginV2IOExt::*; - nvinfer1::PluginRegistrar*; - tensorrt_llm::plugins::api::*; - tensorrt_llm::plugins::*; - }; - local: *; -}; diff --git a/cpp/tensorrt_llm/plugins/fp4GemmPlugin/CMakeLists.txt b/cpp/tensorrt_llm/plugins/fp4GemmPlugin/CMakeLists.txt deleted file mode 100644 index 86876224fccd..000000000000 --- a/cpp/tensorrt_llm/plugins/fp4GemmPlugin/CMakeLists.txt +++ /dev/null @@ -1,21 +0,0 @@ -# -# SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. -# -file(GLOB SRCS *.cpp) -set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS}) -set(PLUGIN_SOURCES - ${PLUGIN_SOURCES} - PARENT_SCOPE) diff --git a/cpp/tensorrt_llm/plugins/fp4GemmPlugin/fp4GemmPlugin.cpp b/cpp/tensorrt_llm/plugins/fp4GemmPlugin/fp4GemmPlugin.cpp deleted file mode 100644 index 05f06ae38feb..000000000000 --- a/cpp/tensorrt_llm/plugins/fp4GemmPlugin/fp4GemmPlugin.cpp +++ /dev/null @@ -1,434 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-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. - */ - -#include - -#include "fp4GemmPlugin.h" -#include "tensorrt_llm/common/assert.h" - -using namespace nvinfer1; -using namespace tensorrt_llm::common; -using tensorrt_llm::plugins::Fp4GemmPluginCreator; -using tensorrt_llm::plugins::Fp4GemmPlugin; -using tensorrt_llm::plugins::Fp4GemmPluginProfiler; -#if defined(USING_OSS_CUTLASS_FP4_GEMM) -using namespace tensorrt_llm::kernels::cutlass_kernels; -#else -using namespace tensorrt_llm::kernels::internal_cutlass_kernels; -#endif - -constexpr nvinfer1::DataType FP4_DTYPE = nvinfer1::DataType::kFP4; -constexpr nvinfer1::DataType FP8_DTYPE = nvinfer1::DataType::kFP8; - -static char const* FP4_GEMM_PLUGIN_VERSION{"1"}; -static char const* FP4_GEMM_PLUGIN_NAME{"Fp4Gemm"}; -PluginFieldCollection Fp4GemmPluginCreator::mFC{}; -std::vector Fp4GemmPluginCreator::mPluginAttributes; - -void Fp4GemmPluginProfiler::runTactic( - int m, int n, int k, Fp4GemmPluginProfiler::Config const& tactic, char* workspace, cudaStream_t const& stream) -{ - // Workspace size required by gemm runner - // NB: this function will throw exception when selected tactic exceeds SMEM, which is then - // caught by gemmPluginProfiler and it will register this tactic as invalid - size_t wsSizeRunner = mRunner->getWorkspaceSize(m, n, k, /* batch_count */ 1); - - // Workspace size required by profiling - size_t wsByteOffset = 0; - int8_t* wsBytePointer = reinterpret_cast(workspace); - void* aTmp = reinterpret_cast(nextWorkspacePtr(wsBytePointer, wsByteOffset, (m * k) / 2)); - void* bTmp = reinterpret_cast(nextWorkspacePtr(wsBytePointer, wsByteOffset, (n * k) / 2)); - void* dTmp = reinterpret_cast( - nextWorkspacePtr(wsBytePointer, wsByteOffset, m * n * (mType == nvinfer1::DataType::kFLOAT ? 4u : 2u))); - // SF M/N is padded along 128 and K is padded along 4. - int vector_size = 16; - int sf_round_m = ((m + 127) / 128) * 128; - int sf_round_n = ((n + 127) / 128) * 128; - int sf_round_k = ((k / vector_size + 3) / 4) * 4; - float* a_sf = reinterpret_cast(nextWorkspacePtr(wsBytePointer, wsByteOffset, sf_round_m * sf_round_k)); - float* b_sf = reinterpret_cast(nextWorkspacePtr(wsBytePointer, wsByteOffset, sf_round_n * sf_round_k)); - float* global_sf = reinterpret_cast(nextWorkspacePtr(wsBytePointer, wsByteOffset, sizeof(float))); - char* workspaceTmp = reinterpret_cast(nextWorkspacePtr(wsBytePointer, wsByteOffset, wsSizeRunner)); - - // Run profiling - mRunner->gemm(dTmp, aTmp, bTmp, a_sf, b_sf, global_sf, m, n, k, /* batch_count */ 1, tactic, workspaceTmp, - wsSizeRunner, stream); - sync_check_cuda_error(stream); -} - -void Fp4GemmPluginProfiler::computeTmpSize(size_t maxM, size_t n, size_t k) -{ - size_t vector_size = 16; - size_t sf_round_m = ((maxM + 127) / 128) * 128; - size_t sf_round_n = ((n + 127) / 128) * 128; - size_t sf_round_k = ((k / vector_size + 3) / 4) * 4; - std::vector workspaces = { - (size_t) (maxM * k / 2), // A - (size_t) (n * k / 2), // B - maxM * n * (mType == nvinfer1::DataType::kFLOAT ? 4u : 2u), // D - (size_t) (sf_round_m * sf_round_k), // A_SF - (size_t) (sf_round_n * sf_round_k), // B_SF - sizeof(float), // Global_SF - mRunner->getWorkspaceSize(maxM, n, k, /* batch_count */ 1) // workspace - }; - size_t bytes = calculateTotalWorkspaceSize(workspaces.data(), workspaces.size()); - setTmpWorkspaceSizeInBytes(bytes); -} - -std::vector Fp4GemmPluginProfiler::getTactics(int m, int n, int k) const -{ - return mRunner->getConfigs(); -} - -Fp4GemmPlugin::Fp4GemmPlugin( - int sfVecSize, nvinfer1::DataType OutputType, Fp4GemmPlugin::PluginProfilerPtr const& pluginProfiler) - : mPluginProfiler(pluginProfiler) - , mSfVecSize(sfVecSize) - , mOutputType(OutputType) -{ - init(OutputType); -} - -Fp4GemmPlugin::Fp4GemmPlugin(void const* data, size_t length, Fp4GemmPlugin::PluginProfilerPtr const& pluginProfiler) - : mPluginProfiler(pluginProfiler) -{ - char const *d = reinterpret_cast(data), *a = d; - read(d, mSfVecSize); - read(d, mOutputType); - read(d, mDims); - - init(mOutputType); - mPluginProfiler->deserialize(d, mDims, mGemmId); - - TLLM_CHECK(d == a + length); -} - -void Fp4GemmPlugin::init(nvinfer1::DataType type) -{ - TLLM_CHECK_WITH_INFO((getSMVersion() >= 100), "FP4 Gemm not supported before Blackwell"); - TLLM_CHECK_WITH_INFO( - (mOutputType == DataType::kBF16) || (mOutputType == DataType::kFLOAT) || (mOutputType == DataType::kHALF), - "Only support float, half, bfloat16, got %d.", (int) mOutputType); - mOutputType = type; - if (mOutputType == nvinfer1::DataType::kHALF) - { - mGemmRunner = std::make_shared>(); - } - else if (mOutputType == nvinfer1::DataType::kFLOAT) - { - mGemmRunner = std::make_shared>(); - } -#ifdef ENABLE_BF16 - else if (mOutputType == nvinfer1::DataType::kBF16) - { - mGemmRunner = std::make_shared>(); - } -#endif - - mGemmId = GemmIdCore(mDims.n, mDims.k, mOutputType); -} - -// IPluginV2DynamicExt Methods -nvinfer1::IPluginV2DynamicExt* Fp4GemmPlugin::clone() const noexcept -{ - auto* plugin = new Fp4GemmPlugin(*this); - return plugin; -} - -nvinfer1::DimsExprs Fp4GemmPlugin::getOutputDimensions( - int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept -{ - TLLM_CHECK_WITH_INFO(outputIndex == 0, "Only support one output"); - auto const& dimsInput = inputs[getInputTensorIdx()]; - auto const& dimsWeights = inputs[getWeightsTensorIdx()]; - TLLM_CHECK_WITH_INFO(dimsInput.nbDims >= 2 && dimsWeights.nbDims == 2, "Fp4GemmPlugin input dim=%d, weights dim=%d", - dimsInput.nbDims, dimsWeights.nbDims); - nvinfer1::DimsExprs ret; - if (outputIndex == 0) - { - ret.nbDims = dimsInput.nbDims; - for (int i = 0; i < dimsInput.nbDims - 1; ++i) - { - ret.d[i] = dimsInput.d[i]; - } - ret.d[dimsInput.nbDims - 1] = dimsWeights.d[0]; - } - else - { - TLLM_CHECK_WITH_INFO(outputIndex == 0, "output fp4 not supported now."); - ret.nbDims = 1; - auto vecCount = dimsInput.d[0]; - int numDim = dimsInput.nbDims; - for (int idx = 1; idx < numDim - 1; ++idx) - { - vecCount = exprBuilder.operation(nvinfer1::DimensionOperation::kPROD, *vecCount, *dimsInput.d[idx]); - } - auto constant128 = exprBuilder.constant(128); - auto alignedRowCount = exprBuilder.operation(nvinfer1::DimensionOperation::kCEIL_DIV, *vecCount, *constant128); - alignedRowCount = exprBuilder.operation(nvinfer1::DimensionOperation::kPROD, *alignedRowCount, *constant128); - auto constant4 = exprBuilder.constant(4); - auto constantSFSize = exprBuilder.constant(mSfVecSize); - auto sfColumn - = exprBuilder.operation(nvinfer1::DimensionOperation::kCEIL_DIV, *dimsInput.d[numDim - 1], *constantSFSize); - auto alignedColumnCount = exprBuilder.operation(nvinfer1::DimensionOperation::kCEIL_DIV, *sfColumn, *constant4); - alignedColumnCount - = exprBuilder.operation(nvinfer1::DimensionOperation::kPROD, *alignedColumnCount, *constant4); - auto totalSize - = exprBuilder.operation(nvinfer1::DimensionOperation::kPROD, *alignedColumnCount, *alignedRowCount); - ret.d[0] = totalSize; - } - return ret; -} - -bool Fp4GemmPlugin::supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept -{ - if (inOut[pos].format != TensorFormat::kLINEAR) - { - return false; - } - if (pos == getInputTensorIdx()) - { - return (inOut[pos].type == FP4_DTYPE); - } - else if (pos == getWeightsTensorIdx()) - { - return (inOut[pos].type == FP4_DTYPE); - } - else if (pos == getInputSFTensorIdx() || pos == getWeightsSFTensorIdx()) - { - return (inOut[pos].type == FP8_DTYPE); - } - else if (pos == getGlobalSFTensorIdx()) - { - return (inOut[pos].type == DataType::kFLOAT); - } - else if (pos == nbInputs) - { - // Output - return (inOut[pos].type == DataType::kFLOAT || inOut[pos].type == DataType::kBF16 - || inOut[pos].type == DataType::kHALF); - } - return false; -} - -void Fp4GemmPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept -{ - auto const minM = std::accumulate(in[0].min.d, in[0].min.d + in[0].min.nbDims - 1, 1, std::multiplies()); - auto const maxM = std::accumulate(in[0].max.d, in[0].max.d + in[0].max.nbDims - 1, 1, std::multiplies()); - - int const maxK = in[0].max.d[in[0].max.nbDims - 1]; - int const maxN = in[2].max.d[0]; - int const minK = in[0].min.d[in[0].min.nbDims - 1]; - int const minN = in[2].min.d[0]; - - TLLM_CHECK_WITH_INFO(minN == maxN, "Variable out channels is not allowed"); - TLLM_CHECK_WITH_INFO(minK == maxK, "Variable in channels is not allowed"); - - if (!mDims.isInitialized()) - { - mDims = {minM, maxM, maxN, maxK}; - } - mGemmId = {maxN, maxK, mOutputType}; - m_workspaceMaxSize = mGemmRunner->getWorkspaceSize(maxM, maxN, maxK, /* batch_count */ 1); -} - -size_t Fp4GemmPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - return m_workspaceMaxSize; -} - -int Fp4GemmPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept -{ - // inputs - // 0. input_tensor [num_tokens, dim] - // 1. input_block_scale [num_tokens, dim / SFVecSize] (padded) - // 2. weights_tensor [out_dim, dim] - // 3. weights_block_scale [out_dim, dim / SFVecSize] (padded) - // 4. alpha (global scaling factor) [1] - // outputs - // 0. output_tensor [num_tokens, out_dim] - int64_t m = 1; - for (int i = 0; i < inputDesc[getInputTensorIdx()].dims.nbDims - 1; ++i) - { - m *= inputDesc[getInputTensorIdx()].dims.d[i]; - } - int const n = inputDesc[getWeightsTensorIdx()].dims.d[0]; - int const k = inputDesc[getWeightsTensorIdx()].dims.d[1]; - TLLM_CHECK_WITH_INFO(k % 32 == 0, "K dim should be aligned to 16 Bytes"); - int N_align = mOutputType == nvinfer1::DataType::kFLOAT ? 4u : 8u; - TLLM_CHECK_WITH_INFO(n % N_align == 0, "N dim should be aligned to 16 Bytes"); - size_t const wsSize = mGemmRunner->getWorkspaceSize(m, n, k, /* batch_count */ 1); - auto const bestTactic = mPluginProfiler->getBestConfig(m, mGemmId); - TLLM_CHECK_WITH_INFO(bestTactic, "No valid FP4 GEMM tactic"); - if (m >= 1) - { - mGemmRunner->gemm(outputs[0], inputs[0], inputs[2], inputs[1], inputs[3], - reinterpret_cast(inputs[4]), m, n, k, /* batch_count */ 1, *bestTactic, - reinterpret_cast(workspace), wsSize, stream); - } - sync_check_cuda_error(stream); - return 0; -} - -// IPluginV2Ext Methods -nvinfer1::DataType Fp4GemmPlugin::getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept -{ - TLLM_CHECK_WITH_INFO(index == 0, "Only support one output"); - return mOutputType; -} - -// IPluginV2 Methods - -char const* Fp4GemmPlugin::getPluginType() const noexcept -{ - return FP4_GEMM_PLUGIN_NAME; -} - -char const* Fp4GemmPlugin::getPluginVersion() const noexcept -{ - return FP4_GEMM_PLUGIN_VERSION; -} - -int Fp4GemmPlugin::getNbOutputs() const noexcept -{ - return 1; -} - -int Fp4GemmPlugin::initialize() noexcept -{ - configGemm(); - return 0; -} - -void Fp4GemmPlugin::terminate() noexcept {} - -size_t Fp4GemmPlugin::getSerializationSize() const noexcept -{ - return sizeof(mSfVecSize) + // mSfVecSize - sizeof(nvinfer1::DataType) + // dtype - sizeof(mDims) + // Dimensions - mPluginProfiler->getSerializationSize(mGemmId); // selected tactics container size -} - -void Fp4GemmPlugin::serialize(void* buffer) const noexcept -{ - char *d = static_cast(buffer), *a = d; - write(d, mSfVecSize); - write(d, mOutputType); - write(d, mDims); - mPluginProfiler->serialize(d, mGemmId); - TLLM_CHECK(d == a + getSerializationSize()); -} - -void Fp4GemmPlugin::destroy() noexcept -{ - delete this; -} - -void Fp4GemmPlugin::configGemm() -{ - mPluginProfiler->profileTactics(mGemmRunner, mOutputType, mDims, mGemmId); -} - -/////////////// - -Fp4GemmPluginCreator::Fp4GemmPluginCreator() -{ - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - mPluginAttributes.emplace_back(PluginField("sv_vec_size", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("output_type_id", nullptr, PluginFieldType::kINT32)); - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -char const* Fp4GemmPluginCreator::getPluginName() const noexcept -{ - return FP4_GEMM_PLUGIN_NAME; -} - -char const* Fp4GemmPluginCreator::getPluginVersion() const noexcept -{ - return FP4_GEMM_PLUGIN_VERSION; -} - -PluginFieldCollection const* Fp4GemmPluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -IPluginV2* Fp4GemmPluginCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept -{ - PluginField const* fields = fc->fields; - TLLM_CHECK(fc->nbFields == 2); - int sf_vec_size{}; - nvinfer1::DataType output_type{}; - // Read configurations from each fields - for (int i = 0; i < fc->nbFields; ++i) - { - char const* attrName = fields[i].name; - if (!strcmp(attrName, "sf_vec_size")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - sf_vec_size = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "output_type_id")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - output_type = static_cast(*(static_cast(fields[i].data))); - } - } - try - { - // Fp4GemmPluginCreator is unique and shared for an engine generation - // Create plugin profiler with shared tactics map - auto pluginProfiler = mGemmPluginProfileManager.createGemmPluginProfiler(/* inference */ false); - auto* obj = new Fp4GemmPlugin(sf_vec_size, output_type, pluginProfiler); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -IPluginV2* Fp4GemmPluginCreator::deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept -{ - // This object will be deleted when the network is destroyed, which will - // call CumsumLastDimPlugin::destroy() - try - { - // Create plugin profiler with private tactics map which is read from the serialized engine - auto pluginProfiler = mGemmPluginProfileManager.createGemmPluginProfiler(/* inference */ true); - auto* obj = new Fp4GemmPlugin(serialData, serialLength, pluginProfiler); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} diff --git a/cpp/tensorrt_llm/plugins/fp4GemmPlugin/fp4GemmPlugin.h b/cpp/tensorrt_llm/plugins/fp4GemmPlugin/fp4GemmPlugin.h deleted file mode 100644 index 9947e849d84f..000000000000 --- a/cpp/tensorrt_llm/plugins/fp4GemmPlugin/fp4GemmPlugin.h +++ /dev/null @@ -1,162 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-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. - */ - -#include "tensorrt_llm/plugins/common/gemmPluginProfiler.h" -#include "tensorrt_llm/plugins/common/plugin.h" -#if defined(USING_OSS_CUTLASS_FP4_GEMM) -#include "tensorrt_llm/kernels/cutlass_kernels/include/fp4_gemm.h" -#else -#include "fp4_gemm.h" -#endif - -#include -#include -#include -#include -#include - -namespace tensorrt_llm::plugins -{ - -#if defined(USING_OSS_CUTLASS_FP4_GEMM) -using Fp4GemmRunnerPtr = std::shared_ptr; -#else -using Fp4GemmRunnerPtr - = std::shared_ptr; -#endif - -class Fp4GemmPluginProfiler : public GemmPluginProfiler -{ -public: - using Config = tensorrt_llm::cutlass_extensions::CutlassGemmConfig; - -protected: - void runTactic(int m, int n, int k, Config const& tactic, char* workspace, cudaStream_t const& stream) override; - - void computeTmpSize(size_t maxM, size_t n, size_t k) override; - - std::vector getTactics(int m, int n, int k) const override; - -private: - tensorrt_llm::common::QuantMode mQuantMode; -}; - -class Fp4GemmPlugin : public BasePlugin -{ -public: - using PluginProfilerPtr = std::shared_ptr; - - Fp4GemmPlugin() = delete; - - Fp4GemmPlugin(int sfVecSize, nvinfer1::DataType OutputType, PluginProfilerPtr const& pluginProfiler); - - Fp4GemmPlugin(void const* data, size_t length, PluginProfilerPtr const& pluginProfiler); - - ~Fp4GemmPlugin() override = default; - - // IPluginV2DynamicExt Methods - nvinfer1::IPluginV2DynamicExt* clone() const noexcept override; - nvinfer1::DimsExprs getOutputDimensions(int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; - bool supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept override; - void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept override; - size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept override; - int enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - - // IPluginV2Ext Methods - nvinfer1::DataType getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept override; - - // IPluginV2 Methods - char const* getPluginType() const noexcept override; - char const* getPluginVersion() const noexcept override; - int getNbOutputs() const noexcept override; - int initialize() noexcept override; - void terminate() noexcept override; - size_t getSerializationSize() const noexcept override; - void serialize(void* buffer) const noexcept override; - void destroy() noexcept override; - -private: - using IndexType = std::int32_t; - - IndexType getInputTensorIdx() const - { - return 0; - }; - - IndexType getInputSFTensorIdx() const - { - return 1; - }; - - IndexType getWeightsTensorIdx() const - { - return 2; - }; - - IndexType getWeightsSFTensorIdx() const - { - return 3; - }; - - IndexType getGlobalSFTensorIdx() const - { - return 4; - } - - void init(nvinfer1::DataType type); - void configGemm(); - - Fp4GemmRunnerPtr mGemmRunner; - PluginProfilerPtr mPluginProfiler; - - int mSfVecSize; - nvinfer1::DataType mOutputType; - size_t m_workspaceMaxSize; - GemmDims mDims{}; - GemmIdCore mGemmId{}; -}; - -class Fp4GemmPluginCreator : public BaseCreator -{ -public: - Fp4GemmPluginCreator(); - - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; - - nvinfer1::IPluginV2* deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept override; - -private: - GemmPluginProfilerManager mGemmPluginProfileManager; - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; -}; - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/fp8RowwiseGemmPlugin/CMakeLists.txt b/cpp/tensorrt_llm/plugins/fp8RowwiseGemmPlugin/CMakeLists.txt deleted file mode 100644 index 3b714a3928fb..000000000000 --- a/cpp/tensorrt_llm/plugins/fp8RowwiseGemmPlugin/CMakeLists.txt +++ /dev/null @@ -1,21 +0,0 @@ -# -# SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. -# -file(GLOB SRCS *.cpp *.cu) -set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS}) -set(PLUGIN_SOURCES - ${PLUGIN_SOURCES} - PARENT_SCOPE) diff --git a/cpp/tensorrt_llm/plugins/fp8RowwiseGemmPlugin/fp8RowwiseGemmPlugin.cpp b/cpp/tensorrt_llm/plugins/fp8RowwiseGemmPlugin/fp8RowwiseGemmPlugin.cpp deleted file mode 100644 index 84963df50a21..000000000000 --- a/cpp/tensorrt_llm/plugins/fp8RowwiseGemmPlugin/fp8RowwiseGemmPlugin.cpp +++ /dev/null @@ -1,422 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. - */ - -#include "fp8RowwiseGemmPlugin.h" -#include "cutlass_extensions/gemm_configs.h" - -#include -#include - -using namespace nvinfer1; -using namespace tensorrt_llm::common; -using namespace tensorrt_llm::kernels::cutlass_kernels; -using tensorrt_llm::plugins::Fp8RowwiseGemmPluginCreator; -using tensorrt_llm::plugins::Fp8RowwiseGemmPlugin; -using tensorrt_llm::plugins::Fp8RowwiseGemmPluginProfiler; - -static char const* FP8_ROWWISE_GEMM_PLUGIN_VERSION{"1"}; -static char const* FP8_ROWWISE_GEMM_PLUGIN_NAME{"Fp8RowwiseGemm"}; -PluginFieldCollection Fp8RowwiseGemmPluginCreator::mFC{}; -std::vector Fp8RowwiseGemmPluginCreator::mPluginAttributes; - -size_t Fp8RowwiseGemmPluginProfiler::getBytePerElement(nvinfer1::DataType type) -{ - size_t bpe; - if (type == nvinfer1::DataType::kHALF || type == nvinfer1::DataType::kBF16) - { - bpe = 2; - } - else if (type == nvinfer1::DataType::kINT8 || type == nvinfer1::DataType::kFP8) - { - bpe = 1; - } - else - { - TLLM_THROW("Not recognized/implemented"); - } - return bpe; -} - -void Fp8RowwiseGemmPluginProfiler::setQuantMode(tensorrt_llm::common::QuantMode const& quantMode) -{ - mQuantMode = quantMode; -} - -void Fp8RowwiseGemmPluginProfiler::runTactic(int m, int n, int k, Fp8RowwiseGemmPluginProfiler::Config const& tactic, - char* workspace, cudaStream_t const& stream) -{ - size_t bpeIn = getBytePerElement(nvinfer1::DataType::kFP8); - size_t bpeOut = getBytePerElement(mType); - - // Workspace size required by gemm runner - // NB: this function will throw exception when selected tactic exceeds SMEM, which is then - // caught by gemmPluginProfiler and it will register this tactic as invalid - size_t wsSizeRunner = mRunner->getWorkspaceSize(m, n, k); - - // Workspace size required by profiling - size_t wsByteOffset = 0; - int8_t* wsBytePointer = reinterpret_cast(workspace); - void* aTmp = reinterpret_cast(nextWorkspacePtr(wsBytePointer, wsByteOffset, m * k * bpeIn)); - void* bTmp = reinterpret_cast(nextWorkspacePtr(wsBytePointer, wsByteOffset, n * k * bpeIn)); - // void* cTmp = reinterpret_cast(nextWorkspacePtr(wsBytePointer, wsByteOffset, n * bpeOut)); - void* dTmp = reinterpret_cast(nextWorkspacePtr(wsBytePointer, wsByteOffset, m * n * bpeOut)); - float* scaleD0Tmp = reinterpret_cast(nextWorkspacePtr(wsBytePointer, wsByteOffset, m * sizeof(float))); - float* scaleD1Tmp = reinterpret_cast(nextWorkspacePtr(wsBytePointer, wsByteOffset, n * sizeof(float))); - char* workspaceTmp = reinterpret_cast(nextWorkspacePtr(wsBytePointer, wsByteOffset, wsSizeRunner)); - - // Run profiling - mRunner->gemm(dTmp, aTmp, bTmp, nullptr, mQuantMode, m, n, k, scaleD0Tmp, scaleD1Tmp, tactic, workspaceTmp, - wsSizeRunner, stream); - sync_check_cuda_error(stream); -} - -int Fp8RowwiseGemmPluginProfiler::getMaxProfileM() const -{ - // Max_num_tokens are not suggested to be set larger than 16k. - return 16384; -} - -void Fp8RowwiseGemmPluginProfiler::computeTmpSize(size_t maxM, size_t n, size_t k) -{ - std::vector workspaces = { - maxM * k * getBytePerElement(nvinfer1::DataType::kFP8), // A - n * k * getBytePerElement(nvinfer1::DataType::kFP8), // B - // n * getBytePerElement(mType), // C_bias - maxM * n * getBytePerElement(mType), // D - maxM * sizeof(float), // alphaRow - n * sizeof(float), // alphaCol - maxM * sizeof(float), // alphaOutput - mRunner->getWorkspaceSize(maxM, n, k) // workspace - }; - size_t bytes = calculateTotalWorkspaceSize(workspaces.data(), workspaces.size()); - setTmpWorkspaceSizeInBytes(bytes); -} - -std::vector Fp8RowwiseGemmPluginProfiler::getTactics(int m, int n, int k) const -{ - return mRunner->getConfigs(); -} - -Fp8RowwiseGemmPlugin::Fp8RowwiseGemmPlugin( - QuantMode quantMode, nvinfer1::DataType type, Fp8RowwiseGemmPlugin::PluginProfilerPtr const& pluginProfiler) - : mQuantMode(quantMode) - , mPluginProfiler(pluginProfiler) -{ - init(type); -} - -// Parameterized constructor -Fp8RowwiseGemmPlugin::Fp8RowwiseGemmPlugin( - void const* data, size_t length, Fp8RowwiseGemmPlugin::PluginProfilerPtr const& pluginProfiler) - : mPluginProfiler(pluginProfiler) -{ - char const *d = reinterpret_cast(data), *a = d; - nvinfer1::DataType type; - unsigned int quantMode; - read(d, quantMode); - read(d, type); - read(d, mDims); - - mQuantMode = QuantMode(quantMode); - - init(type); - - mPluginProfiler->deserialize(d, mDims, mGemmId); - - TLLM_CHECK(d == a + length); -} - -void Fp8RowwiseGemmPlugin::init(nvinfer1::DataType type) -{ - mType = type; - if (mType == nvinfer1::DataType::kHALF) - { - mGemmRunner = std::make_shared>(); - } -#ifdef ENABLE_BF16 - else if (mType == nvinfer1::DataType::kBF16) - { - mGemmRunner = std::make_shared>(); - } -#endif - else - { - TLLM_THROW("Fp8 Rowwise Gemm plugin doesn't support this type now"); - } - - mPluginProfiler->setQuantMode(mQuantMode); - - mGemmId = GemmIdCore(mDims.n, mDims.k, mType); -} - -// IPluginV2DynamicExt Methods -nvinfer1::IPluginV2DynamicExt* Fp8RowwiseGemmPlugin::clone() const noexcept -{ - auto* plugin = new Fp8RowwiseGemmPlugin(*this); - return plugin; -} - -nvinfer1::DimsExprs Fp8RowwiseGemmPlugin::getOutputDimensions( - int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept -{ - try - { - TLLM_CHECK(nbInputs == 4); - TLLM_CHECK(outputIndex == 0); - int const nbDimsA = inputs[0].nbDims; - TLLM_CHECK(nbDimsA >= 2); - DimsExprs ret; - ret.nbDims = nbDimsA; - for (int ii = 0; ii < nbDimsA - 1; ++ii) - { - ret.d[ii] = inputs[0].d[ii]; - } - ret.d[nbDimsA - 1] = inputs[1].d[0]; - return ret; - } - catch (std::exception const& e) - { - caughtError(e); - } - return DimsExprs{}; -} - -bool Fp8RowwiseGemmPlugin::supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept -{ - switch (pos) - { - case 0: - // activation - return inOut[pos].type == nvinfer1::DataType::kFP8 && inOut[pos].format == TensorFormat::kLINEAR; - case 1: - // weights - // Weights stored in checkpoint must have fp8 type - return inOut[pos].type == nvinfer1::DataType::kFP8 && inOut[pos].format == TensorFormat::kLINEAR; - case 2: - // scales channels - case 3: - // scales tokens - return inOut[pos].type == nvinfer1::DataType::kFLOAT && inOut[pos].format == TensorFormat::kLINEAR; - case 4: - // out - return inOut[pos].type == mType && inOut[pos].format == TensorFormat::kLINEAR; - default: - // All other format combinations are unsupported. - return false; - } -} - -void Fp8RowwiseGemmPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept -{ - auto const minM = std::accumulate(in[0].min.d, in[0].min.d + in[0].min.nbDims - 1, 1, std::multiplies()); - auto const maxM = std::accumulate(in[0].max.d, in[0].max.d + in[0].max.nbDims - 1, 1, std::multiplies()); - - int const maxK = in[0].max.d[in[0].max.nbDims - 1]; - int const maxN = in[1].max.d[0]; - int const minK = in[0].min.d[in[0].min.nbDims - 1]; - int const minN = in[1].min.d[0]; - - TLLM_CHECK_WITH_INFO(minN == maxN, "Variable out channels is not allowed"); - TLLM_CHECK_WITH_INFO(minK == maxK, "Variable in channels is not allowed"); - - if (!mDims.isInitialized()) - { - mDims = {minM, maxM, maxN, maxK}; - } - mGemmId = {maxN, maxK, mType}; - - mWorkspaceMaxSize = mGemmRunner->getWorkspaceSize(maxM, maxN, maxK); -} - -size_t Fp8RowwiseGemmPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - return mWorkspaceMaxSize; -} - -int Fp8RowwiseGemmPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) noexcept -{ - // inputs - // mat1 [M(*), K] - // mat2 [N, K] - // scale_tokens [M, 1] if has_per_token_scaling else [1, 1] - // scale_channels [1, N] if has_per_channel_scaling else [1, 1] - // outputs - // mat [M(*), N] - int m = 1; - for (int ii = 0; ii < inputDesc[0].dims.nbDims - 1; ++ii) - { - m *= inputDesc[0].dims.d[ii]; - } - int const n = inputDesc[1].dims.d[0]; - int const k = inputDesc[0].dims.d[inputDesc[0].dims.nbDims - 1]; - size_t const wsSize = mGemmRunner->getWorkspaceSize(m, n, k); - - auto const bestTactic = mPluginProfiler->getBestConfig(m, mGemmId); - TLLM_CHECK_WITH_INFO(bestTactic, "No valid GEMM tactic"); - mGemmRunner->gemm(outputs[0], inputs[0], inputs[1], nullptr, mQuantMode, m, n, k, - reinterpret_cast(inputs[2]), reinterpret_cast(inputs[3]), *bestTactic, - reinterpret_cast(workspace), wsSize, stream); - sync_check_cuda_error(stream); - - return 0; -} - -// IPluginV2Ext Methods -nvinfer1::DataType Fp8RowwiseGemmPlugin::getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept -{ - TLLM_CHECK(index == 0); - return mType; -} - -// IPluginV2 Methods - -char const* Fp8RowwiseGemmPlugin::getPluginType() const noexcept -{ - return FP8_ROWWISE_GEMM_PLUGIN_NAME; -} - -char const* Fp8RowwiseGemmPlugin::getPluginVersion() const noexcept -{ - return FP8_ROWWISE_GEMM_PLUGIN_VERSION; -} - -int Fp8RowwiseGemmPlugin::getNbOutputs() const noexcept -{ - return 1; -} - -int Fp8RowwiseGemmPlugin::initialize() noexcept -{ - configGemm(); // gemm profiler in action - return 0; -} - -void Fp8RowwiseGemmPlugin::terminate() noexcept {} - -size_t Fp8RowwiseGemmPlugin::getSerializationSize() const noexcept -{ - return sizeof(unsigned int) + // QuantMode - sizeof(nvinfer1::DataType) + // dtype - sizeof(mDims) + // Dimensions - mPluginProfiler->getSerializationSize(mGemmId); // selected tactics container size -} - -void Fp8RowwiseGemmPlugin::serialize(void* buffer) const noexcept -{ - char *d = static_cast(buffer), *a = d; - write(d, mQuantMode.value()); - write(d, mType); - write(d, mDims); - - mPluginProfiler->serialize(d, mGemmId); - TLLM_CHECK(d == a + getSerializationSize()); -} - -void Fp8RowwiseGemmPlugin::destroy() noexcept -{ - // This gets called when the network containing plugin is destroyed - delete this; -} - -void Fp8RowwiseGemmPlugin::configGemm() -{ - mPluginProfiler->profileTactics(mGemmRunner, mType, mDims, mGemmId); -} - -Fp8RowwiseGemmPluginCreator::Fp8RowwiseGemmPluginCreator() -{ - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - mPluginAttributes.emplace_back(PluginField("has_per_channel_scaling", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("has_per_token_scaling", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("type_id", nullptr, PluginFieldType::kINT32)); - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -char const* Fp8RowwiseGemmPluginCreator::getPluginName() const noexcept -{ - return FP8_ROWWISE_GEMM_PLUGIN_NAME; -} - -char const* Fp8RowwiseGemmPluginCreator::getPluginVersion() const noexcept -{ - return FP8_ROWWISE_GEMM_PLUGIN_VERSION; -} - -PluginFieldCollection const* Fp8RowwiseGemmPluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -IPluginV2* Fp8RowwiseGemmPluginCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept -{ - PluginField const* fields = fc->fields; - TLLM_CHECK(fc->nbFields == 3); - nvinfer1::DataType type{}; - // Read configurations from each fields - for (int i = 0; i < fc->nbFields; ++i) - { - char const* attrName = fields[i].name; - if (!strcmp(attrName, "type_id")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - type = static_cast(*(static_cast(fields[i].data))); - } - } - try - { - // Fp8RowwiseGemmPluginCreator is unique and shared for an engine generation - // Create plugin profiler with shared tactics map - auto pluginProfiler = mGemmPluginProfileManager.createGemmPluginProfiler(/* inference */ false); - QuantMode quantMode = QuantMode{}; - auto* obj = new Fp8RowwiseGemmPlugin(quantMode, type, pluginProfiler); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -IPluginV2* Fp8RowwiseGemmPluginCreator::deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept -{ - // This object will be deleted when the network is destroyed, which will - // call Fp8RowwiseGemmPlugin::destroy() - try - { - // Create plugin profiler with private tactics map which is read from the serialized engine - auto pluginProfiler = mGemmPluginProfileManager.createGemmPluginProfiler(/* inference */ true); - auto* obj = new Fp8RowwiseGemmPlugin(serialData, serialLength, pluginProfiler); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} diff --git a/cpp/tensorrt_llm/plugins/fp8RowwiseGemmPlugin/fp8RowwiseGemmPlugin.h b/cpp/tensorrt_llm/plugins/fp8RowwiseGemmPlugin/fp8RowwiseGemmPlugin.h deleted file mode 100644 index 36f22ad5885d..000000000000 --- a/cpp/tensorrt_llm/plugins/fp8RowwiseGemmPlugin/fp8RowwiseGemmPlugin.h +++ /dev/null @@ -1,140 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. - */ -#pragma once - -#include "tensorrt_llm/kernels/cutlass_kernels/fp8_rowwise_gemm/fp8_rowwise_gemm.h" -#include "tensorrt_llm/plugins/common/gemmPluginProfiler.h" -#include "tensorrt_llm/plugins/common/plugin.h" -#include -#include -#include -#include - -namespace tensorrt_llm::plugins -{ - -using Fp8RowwiseGemmRunnerPtr - = std::shared_ptr; - -class Fp8RowwiseGemmPluginProfiler : public GemmPluginProfiler - -{ -public: - using Config = tensorrt_llm::cutlass_extensions::CutlassGemmConfig; - - void setQuantMode(tensorrt_llm::common::QuantMode const& quantMode); - - virtual int getMaxProfileM() const override; - -protected: - void runTactic(int m, int n, int k, Config const& tactic, char* workspace, cudaStream_t const& stream) override; - - void computeTmpSize(size_t maxM, size_t n, size_t k) override; - - std::vector getTactics(int m, int n, int k) const override; - -private: - size_t getBytePerElement(nvinfer1::DataType type); - - tensorrt_llm::common::QuantMode mQuantMode; -}; - -class Fp8RowwiseGemmPlugin : public BasePlugin -{ -public: - using PluginProfilerPtr = std::shared_ptr; - - Fp8RowwiseGemmPlugin() = delete; - - Fp8RowwiseGemmPlugin( - tensorrt_llm::common::QuantMode quantMode, nvinfer1::DataType type, PluginProfilerPtr const& pluginProfiler); - - Fp8RowwiseGemmPlugin(void const* data, size_t length, PluginProfilerPtr const& profiler); - - ~Fp8RowwiseGemmPlugin() override = default; - - // IPluginV2DynamicExt Methods - nvinfer1::IPluginV2DynamicExt* clone() const noexcept override; - nvinfer1::DimsExprs getOutputDimensions(int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; - bool supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept override; - void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept override; - size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept override; - int enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - - // IPluginV2Ext Methods - nvinfer1::DataType getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept override; - - // IPluginV2 Methods - char const* getPluginType() const noexcept override; - char const* getPluginVersion() const noexcept override; - int getNbOutputs() const noexcept override; - int initialize() noexcept override; - void terminate() noexcept override; - size_t getSerializationSize() const noexcept override; - void serialize(void* buffer) const noexcept override; - void destroy() noexcept override; - -private: - void init(nvinfer1::DataType type); - - void configGemm(); - -private: - const std::string mLayerName; - - Fp8RowwiseGemmRunnerPtr mGemmRunner; - tensorrt_llm::common::QuantMode mQuantMode; // not configurable yet - size_t mWorkspaceMaxSize; - - GemmDims mDims{}; - GemmIdCore mGemmId{}; - - PluginProfilerPtr mPluginProfiler; - - nvinfer1::DataType mType; -}; - -class Fp8RowwiseGemmPluginCreator : public BaseCreator -{ -public: - Fp8RowwiseGemmPluginCreator(); - - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; - - nvinfer1::IPluginV2* deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept override; - -private: - GemmPluginProfilerManager mGemmPluginProfileManager; - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; -}; - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/fusedLayernormPlugin/CMakeLists.txt b/cpp/tensorrt_llm/plugins/fusedLayernormPlugin/CMakeLists.txt deleted file mode 100755 index 7cc985b60b7a..000000000000 --- a/cpp/tensorrt_llm/plugins/fusedLayernormPlugin/CMakeLists.txt +++ /dev/null @@ -1,21 +0,0 @@ -# -# SPDX-FileCopyrightText: Copyright (c) 1993-2023 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. -# -file(GLOB SRCS *.cpp) -set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS}) -set(PLUGIN_SOURCES - ${PLUGIN_SOURCES} - PARENT_SCOPE) diff --git a/cpp/tensorrt_llm/plugins/fusedLayernormPlugin/fusedLayernormPlugin.cpp b/cpp/tensorrt_llm/plugins/fusedLayernormPlugin/fusedLayernormPlugin.cpp deleted file mode 100644 index 541afdadc4c8..000000000000 --- a/cpp/tensorrt_llm/plugins/fusedLayernormPlugin/fusedLayernormPlugin.cpp +++ /dev/null @@ -1,388 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. - */ -#include "fusedLayernormPlugin.h" -#include "pluginUtils.h" -#include "tensorrt_llm/common/assert.h" - -using namespace nvinfer1; -using namespace tensorrt_llm::kernels; -using namespace tensorrt_llm::common; -using tensorrt_llm::plugins::FusedLayernormPluginCreator; -using tensorrt_llm::plugins::FusedLayernormPlugin; - -static char const* FUSED_LAYERNORM_PLUGIN_VERSION{"1"}; -static char const* FUSED_LAYERNORM_PLUGIN_NAME{"FusedLayernorm"}; -PluginFieldCollection FusedLayernormPluginCreator::mFC{}; -std::vector FusedLayernormPluginCreator::mPluginAttributes; - -FusedLayernormPlugin::FusedLayernormPlugin(float eps, bool needFP32Output, bool needQuantize, nvinfer1::DataType type) - : mEps(eps) - , mNeedFP32Output(needFP32Output) - , mNeedQuantize(needQuantize) - , mType(type) -{ -} - -// Parameterized constructor -FusedLayernormPlugin::FusedLayernormPlugin(void const* data, size_t length) -{ - char const *d = reinterpret_cast(data), *a = d; - read(d, mEps); - read(d, mNeedFP32Output); - read(d, mNeedQuantize); - read(d, mType); - TLLM_CHECK_WITH_INFO(d == a + length, - "Expected length (%d) != real length (%d). This is often " - "caused by using different TensorRT LLM version to build " - "engine and run engine.", - (int) length, (int) (d - a)); -} - -// IPluginV2DynamicExt Methods -nvinfer1::IPluginV2DynamicExt* FusedLayernormPlugin::clone() const noexcept -{ - auto* plugin = new FusedLayernormPlugin(mEps, mNeedFP32Output, mNeedQuantize, mType); - plugin->setPluginNamespace(mNamespace.c_str()); - return plugin; -} - -nvinfer1::DimsExprs FusedLayernormPlugin::getOutputDimensions( - int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept -{ - // Dim should be the same as input hidden states - if (!mNeedQuantize) - { - return inputs[0]; - } - - if (outputIndex == 1) // un-normed output fp16 - { - return inputs[0]; - } - if (outputIndex == 0) // quantized normed output - { - // Quantized output with int64_t data type (16 FP4 values per element). - DimsExprs ret; - ret.nbDims = inputs[0].nbDims; - for (int di = 0; di < ret.nbDims; ++di) - { - ret.d[di] = inputs[0].d[di]; - } - return ret; - } - - // Scaling Factors. - try - { - TLLM_CHECK(outputIndex == 2); - - DimsExprs ret; - ret.nbDims = inputs[0].nbDims; - for (int di = 0; di < ret.nbDims; ++di) - { - ret.d[di] = inputs[0].d[di]; - } - // Sequence dimension or token dimension. - // Pad to multiple of 128. - auto dimM - = exprBuilder.operation(DimensionOperation::kCEIL_DIV, *ret.d[ret.nbDims - 2], *exprBuilder.constant(128)); - ret.d[ret.nbDims - 2] = exprBuilder.operation(DimensionOperation::kPROD, *dimM, *exprBuilder.constant(128)); - // Hidden size dimension. - ret.d[ret.nbDims - 1] - = exprBuilder.operation(DimensionOperation::kCEIL_DIV, *ret.d[ret.nbDims - 1], *exprBuilder.constant(16)); - return ret; - } - catch (std::exception const& e) - { - caughtError(e); - } - return DimsExprs{}; -} - -bool FusedLayernormPlugin::supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept -{ - int const totalPoses = 5 + 2 * static_cast(mNeedQuantize); - TLLM_CHECK(0 <= pos && pos < totalPoses); - TLLM_CHECK(nbInputs == 3 + static_cast(mNeedQuantize)); - if (pos < nbInputs) - { - switch (pos) - { - case 0: - case 1: - case 2: return (inOut[pos].type == mType) && (inOut[pos].format == TensorFormat::kLINEAR); - case 3: return (inOut[pos].type == nvinfer1::DataType::kFLOAT); - } - } - if (pos == nbInputs) // Normed output - { - if (mNeedQuantize) - { - // fp4 quantized output -- fp4 padded tp int64 - return (inOut[pos].type == nvinfer1::DataType::kFP4) && (inOut[pos].format == TensorFormat::kLINEAR); - } - return (inOut[pos].type == mType) && (inOut[pos].format == TensorFormat::kLINEAR); - } - else if (pos == nbInputs + 1) // Un-normed output - { - return (inOut[pos].type == mType) && (inOut[pos].format == TensorFormat::kLINEAR); - } - // fp4 act_per_block_scale -- fp8 padded to int32 - return (inOut[pos].type == nvinfer1::DataType::kFP8) && (inOut[pos].format == TensorFormat::kLINEAR); -} - -void FusedLayernormPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept -{ -} - -size_t FusedLayernormPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - return sizeof(WarpSpecializedCounters); -} - -int FusedLayernormPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) noexcept -{ - // inputs - // input [M(*), N] - // residual [M(*), N] - // weight [N, ] - // scale [1, ] - if needQuantize - // outputs - // output [M(*), N] - fp4 padded to int64 / fp16 - // un-normed output [M(*), N] - fp16 - // act_per_block_scale - fp8 padded to int32 - if needQuantize - -#define SETUP_PARAM \ - Param param; \ - int64_t m64 = 1; \ - for (int i = 0; i < inputDesc[0].dims.nbDims - 1; ++i) \ - { \ - m64 *= inputDesc[0].dims.d[i]; \ - } \ - int const m = TLLM_INT32_CAST(m64); \ - int const n = TLLM_INT32_CAST(inputDesc[2].dims.d[0]); \ - param.m = m; \ - param.n = n; \ - param.layernorm_eps = mEps; \ - param.input = const_cast(reinterpret_cast(inputs[0])); \ - param.residual = const_cast(reinterpret_cast(inputs[1])); \ - param.gamma = const_cast(reinterpret_cast(inputs[2])); \ - if (mNeedQuantize) \ - { \ - param.sf_scale = const_cast(reinterpret_cast(inputs[3])); \ - } \ - param.counters = reinterpret_cast(workspace); \ - param.stream = stream; \ - param.normed_output = reinterpret_cast(outputs[0]); \ - param.output = reinterpret_cast(outputs[1]); \ - param.sf_out = reinterpret_cast(outputs[2]); - -#define CLEANUP_AND_INVOKE \ - TLLM_CUDA_CHECK(cudaMemsetAsync(workspace, 0, sizeof(WarpSpecializedCounters), stream)); \ - invokeWSLayerNorm(param, true, num_sms); - - int num_sms = tensorrt_llm::common::getMultiProcessorCount(); - - if (mType == DataType::kHALF) - { - using Input = half; - using Param = WarpSpecializedParam>; - SETUP_PARAM - CLEANUP_AND_INVOKE - } -#ifdef ENABLE_BF16 - else if (mType == DataType::kBF16) - { - using Input = __nv_bfloat16; - using Param = WarpSpecializedParam>; - SETUP_PARAM - CLEANUP_AND_INVOKE - } -#endif - else - { - TLLM_LOG_ERROR("Unsupported data type"); - return 1; - } - return 0; -} - -// IPluginV2Ext Methods -nvinfer1::DataType FusedLayernormPlugin::getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept -{ - // assert((mNeedFP32Output && index < 3) || (!mNeedFP32Output && index < 2)); - assert((mNeedQuantize && index < 3) || (!mNeedQuantize && index < 2)); - if (index == 0) - { - // Output 0 quantized output of layernorm - fp4 padded to int64 - if (mNeedQuantize) - { - return nvinfer1::DataType::kFP4; - } - return mType; - } - else if (index == 1) - { - // Output 1 un-normed output - return mType; - } - // Output 2 act_per_block_scale - fp8 padded to int32 - return nvinfer1::DataType::kFP8; -} - -// IPluginV2 Methods - -char const* FusedLayernormPlugin::getPluginType() const noexcept -{ - return FUSED_LAYERNORM_PLUGIN_NAME; -} - -char const* FusedLayernormPlugin::getPluginVersion() const noexcept -{ - return FUSED_LAYERNORM_PLUGIN_VERSION; -} - -int FusedLayernormPlugin::getNbOutputs() const noexcept -{ - return 2 + static_cast(mNeedQuantize); -} - -int FusedLayernormPlugin::initialize() noexcept -{ - return 0; -} - -void FusedLayernormPlugin::terminate() noexcept {} - -size_t FusedLayernormPlugin::getSerializationSize() const noexcept -{ - return sizeof(mEps) + sizeof(mNeedFP32Output) + sizeof(mNeedQuantize) + sizeof(mType); -} - -void FusedLayernormPlugin::serialize(void* buffer) const noexcept -{ - char *d = static_cast(buffer), *a = d; - write(d, mEps); - write(d, mNeedFP32Output); - write(d, mNeedQuantize); - write(d, mType); - TLLM_CHECK(d == a + getSerializationSize()); -} - -void FusedLayernormPlugin::destroy() noexcept -{ - // This gets called when the network containing plugin is destroyed - delete this; -} - -/////////////// - -FusedLayernormPluginCreator::FusedLayernormPluginCreator() -{ - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - mPluginAttributes.emplace_back(PluginField("eps", nullptr, PluginFieldType::kFLOAT32)); - mPluginAttributes.emplace_back(PluginField("need_fp32_output", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("need_quantize", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("type_id", nullptr, PluginFieldType::kINT32)); - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -char const* FusedLayernormPluginCreator::getPluginName() const noexcept -{ - return FUSED_LAYERNORM_PLUGIN_NAME; -} - -char const* FusedLayernormPluginCreator::getPluginVersion() const noexcept -{ - return FUSED_LAYERNORM_PLUGIN_VERSION; -} - -PluginFieldCollection const* FusedLayernormPluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -IPluginV2* FusedLayernormPluginCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept -{ - PluginField const* fields = fc->fields; - float eps{}; - nvinfer1::DataType type{}; - bool needFP32Output{}; - bool needQuantize{}; - // Read configurations from each fields - for (int i = 0; i < fc->nbFields; ++i) - { - char const* attrName = fields[i].name; - if (!strcmp(attrName, "eps")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kFLOAT32); - eps = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "type_id")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - type = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "need_fp32_output")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - needFP32Output = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "need_quantize")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - needQuantize = static_cast(*(static_cast(fields[i].data))); - } - } - try - { - auto* obj = new FusedLayernormPlugin(eps, needFP32Output, needQuantize, type); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -IPluginV2* FusedLayernormPluginCreator::deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept -{ - // This object will be deleted when the network is destroyed, which will - // call FusedLayernormPlugin::destroy() - try - { - auto* obj = new FusedLayernormPlugin(serialData, serialLength); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} diff --git a/cpp/tensorrt_llm/plugins/fusedLayernormPlugin/fusedLayernormPlugin.h b/cpp/tensorrt_llm/plugins/fusedLayernormPlugin/fusedLayernormPlugin.h deleted file mode 100755 index c6c899950fdc..000000000000 --- a/cpp/tensorrt_llm/plugins/fusedLayernormPlugin/fusedLayernormPlugin.h +++ /dev/null @@ -1,98 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. - */ -#pragma once - -#include "tensorrt_llm/kernels/fusedLayernormKernels/layernorm_param.h" -#include "tensorrt_llm/kernels/fusedLayernormKernels/ws_layernorm.h" -#include "tensorrt_llm/plugins/common/plugin.h" -#include -#include -#include -#include - -namespace tensorrt_llm::plugins -{ - -class FusedLayernormPlugin : public BasePlugin -{ -public: - FusedLayernormPlugin() = delete; - - FusedLayernormPlugin(float eps, bool needFP32Output, bool needQuantize, nvinfer1::DataType type); - - FusedLayernormPlugin(void const* data, size_t length); - - ~FusedLayernormPlugin() override = default; - - // IPluginV2DynamicExt Methods - nvinfer1::IPluginV2DynamicExt* clone() const noexcept override; - nvinfer1::DimsExprs getOutputDimensions(int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; - bool supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept override; - void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept override; - size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept override; - int enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - - // IPluginV2Ext Methods - nvinfer1::DataType getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept override; - - // IPluginV2 Methods - char const* getPluginType() const noexcept override; - char const* getPluginVersion() const noexcept override; - int getNbOutputs() const noexcept override; - int initialize() noexcept override; - void terminate() noexcept override; - size_t getSerializationSize() const noexcept override; - void serialize(void* buffer) const noexcept override; - void destroy() noexcept override; - -private: - float mEps; - bool mNeedFP32Output; - bool mNeedQuantize; - nvinfer1::DataType mType; - - const std::string mLayerName; -}; - -class FusedLayernormPluginCreator : public BaseCreator -{ -public: - FusedLayernormPluginCreator(); - - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; - - nvinfer1::IPluginV2* deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept override; - -private: - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; -}; - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/gemmAllReducePlugin/CMakeLists.txt b/cpp/tensorrt_llm/plugins/gemmAllReducePlugin/CMakeLists.txt deleted file mode 100644 index 1d1fa98f4132..000000000000 --- a/cpp/tensorrt_llm/plugins/gemmAllReducePlugin/CMakeLists.txt +++ /dev/null @@ -1,19 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 1993-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. -file(GLOB SRCS *.cpp) -set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS}) -set(PLUGIN_SOURCES - ${PLUGIN_SOURCES} - PARENT_SCOPE) diff --git a/cpp/tensorrt_llm/plugins/gemmAllReducePlugin/gemmAllReducePlugin.cpp b/cpp/tensorrt_llm/plugins/gemmAllReducePlugin/gemmAllReducePlugin.cpp deleted file mode 100644 index 08ee2af55406..000000000000 --- a/cpp/tensorrt_llm/plugins/gemmAllReducePlugin/gemmAllReducePlugin.cpp +++ /dev/null @@ -1,721 +0,0 @@ -/* - * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved. - * - * 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. - */ -#include "gemmAllReducePlugin.h" -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/kernels/cutlass_kernels/cutlass_type_conversion.h" -#include "tensorrt_llm/plugins/common/pluginUtils.h" - -#include - -static char const* GEMM_ALLREDUCE_PLUGIN_VERSION = "1"; -static char const* GEMM_ALLREDUCE_PLUGIN_NAME = "GemmAllReduce"; -template -using CutlassType = ::tensorrt_llm::kernels::cutlass_kernels::CutlassType; - -namespace tensorrt_llm::plugins -{ -template -static std::pair makeEntry() -{ - return {std::make_tuple(ElementA, ElementB, ElementD), - [&]() - { - using GemmTraits - = cutlass_kernels::GemmTypes::type, typename CutlassType::type, - typename CutlassType::type, // C, unused - typename CutlassType::type, - std::conditional_t, // SFA - std::conditional_t, // SFB - cutlass::layout::RowMajor, cutlass::layout::ColumnMajor, - cutlass::layout::RowMajor, // C, unused - cutlass::layout::RowMajor>; - return new cutlass_kernels::GemmAllReduceImplRunner(); - }}; -} - -template -static std::map getTypedInstantiators() -{ - return std::map({makeEntry(), - makeEntry(), - makeEntry(), - makeEntry(), - makeEntry(), - makeEntry()}); -} - -//////////////////////////////////////////////////////////// -// GemmAllReducePlugin Methods -//////////////////////////////////////////////////////////// -GemmAllReducePlugin::GemmAllReducePlugin(GemmAllReducePluginOptions const& options) - : mOptions(options) - , mGemmId(GemmIdCore(options.maxProblemShape.n, options.maxProblemShape.k, options.typeD)) - , mProfiler(mGemmPluginProfileManager.createGemmPluginProfiler(/*inference=*/options.deserialize)) -{ - // construct mapping of input/output pos to argument - int argIdx = 0; - // inputs - mArgMap[argIdx++] = TensorArg::IN_ACTIVATION; - mArgMap[argIdx++] = TensorArg::IN_WEIGHT; - if (mOptions.hasSFA) - { - mArgMap[argIdx++] = TensorArg::IN_ACTIVATION_SF; - } - if (mOptions.hasSFB) - { - mArgMap[argIdx++] = TensorArg::IN_WEIGHT_SF; - } - if (mOptions.alphaIsPtr) - { - mArgMap[argIdx++] = TensorArg::IN_ALPHA; - } - mNbInputs = argIdx; - // outputs - mArgMap[argIdx++] = TensorArg::OUT_D_UC; - mArgMap[argIdx++] = TensorArg::OUT_D_MC; - mArgMap[argIdx++] = TensorArg::OUT_D_IPC; - mNbOutputs = argIdx - mNbInputs; - - // Create mapping of argument to tensor pos - for (auto const& pair : mArgMap) - { - mArgInvMap[pair.second] = pair.first; - } - - // Use map instead of huge switch case - mTypedInstantiators = getTypedInstantiators(); - - auto key = std::make_tuple(mOptions.typeA, mOptions.typeB, mOptions.typeD); - - TLLM_CHECK_WITH_INFO(mTypedInstantiators.count(key) > 0, "No cutlass gemm for impl."); - mGemm = std::shared_ptr(mTypedInstantiators[key]()); -} - -void GemmAllReducePlugin::allocatePersistentWorkspace() -{ - TLLM_CHECK(mOptions.maxProblemShape.isInitialized()); - - mWorkspaceKey = "gemm_allreduce_workspace_m" + std::to_string(mOptions.maxProblemShape.maxM); - - cutlass_kernels::GemmAllReduceImplInterface::LaunchConfig smallest_tile_config - = mGemm->getSupportedLaunchConfigs()[0]; - cutlass_kernels::GemmAllReduceImplInterface::ProblemArgs args; - args.argProblemShape(mOptions.maxProblemShape.maxM, mOptions.maxProblemShape.n, mOptions.maxProblemShape.k, 1) - .argRanks(mRank, mOptions.group) - .argLaunchConfig(smallest_tile_config); - - TLLM_CHECK(mWorkspace == nullptr); - - // Wrap persistent workspace in IPluginResource type - // so that clone() can be called to allocate memory - GemmAllReducePersistentWorkspace unallocated_resource(mGemm->getPersistentWorkspace(args)); - - // Register and allocate workspace - mWorkspace = static_cast( - getPluginRegistry()->acquirePluginResource(mWorkspaceKey.c_str(), &unallocated_resource)); - TLLM_CHECK(mWorkspace != nullptr); -} - -LaunchConfig GemmAllReducePlugin::getStaticHeuristicLaunchConfig(int M) const -{ - using namespace tensorrt_llm::cutlass_extensions; - // This is only applicable when we swap and transpose A & B. - // When M is small we want to select tile that best fits it to maximize MMA efficiency. - auto filterByM = [&](std::vector candidateConfigs) - { - std::vector result; - if (M <= 16) - { - std::copy_if(candidateConfigs.begin(), candidateConfigs.end(), std::back_inserter(result), - [](const LaunchConfig& config) - { return config.tile_shape == TileShape::TileShape_128x16x128 and config.transposed; }); - } - else if (M <= 32) - { - std::copy_if(candidateConfigs.begin(), candidateConfigs.end(), std::back_inserter(result), - [](const LaunchConfig& config) - { return config.tile_shape == TileShape::TileShape_128x32x128 and config.transposed; }); - } - else if (M <= 64) - { - std::copy_if(candidateConfigs.begin(), candidateConfigs.end(), std::back_inserter(result), - [](const LaunchConfig& config) - { return config.tile_shape == TileShape::TileShape_128x64x128 and config.transposed; }); - } - else - { - std::copy_if(candidateConfigs.begin(), candidateConfigs.end(), std::back_inserter(result), - [](const LaunchConfig& config) - { return config.tile_shape == TileShape::TileShape_128x128x128 and config.transposed; }); - } - // If result empty then use any. - if (result.empty()) - { - result = candidateConfigs; - } - return result; - }; - - auto bestLaunchConfigs = mGemm->getSupportedLaunchConfigs(); - bestLaunchConfigs = filterByM(bestLaunchConfigs); - TLLM_CHECK(!bestLaunchConfigs.empty()); - // Return first one, because who knows which is best. - return bestLaunchConfigs.front(); -} - -static GemmAllReducePluginOptions deserializeOptions(void const*& data, size_t length) -{ - char const* begin = reinterpret_cast(data); - char const*& end = reinterpret_cast(data); - GemmAllReducePluginOptions options; - options.deserialize = true; - - read(end, options.typeA); - read(end, options.typeB); - read(end, options.typeD); - read(end, options.transA); - read(end, options.transB); - read(end, options.alpha); - read(end, options.maxProblemShape); - read(end, options.groupSize); - for (int i = 0; i < options.groupSize; ++i) - { - int rank = -1; - read(end, rank); - options.group.insert(rank); - } - read(end, options.hasSFA); - read(end, options.hasSFB); - read(end, options.alphaIsPtr); - - TLLM_CHECK_WITH_INFO(end == begin + length, - "Expected length (%d) != real length (%d). This is often " - "caused by using different TensorRT LLM version to build " - "engine and run engine.", - (int) length, (int) (end - begin)); - - return options; -} - -GemmAllReducePlugin::GemmAllReducePlugin(void const* data, size_t length) - : GemmAllReducePlugin(deserializeOptions(std::ref(data), length)) -{ - if (mProfiler->useProfiler()) - { - mProfiler->deserializeFromOwnFile(mGemmId, mOptions.maxProblemShape); - } -} - -////////////////////////////////// -// IPluginV2DynamicExt Methods -////////////////////////////////// -IPluginV2DynamicExt* GemmAllReducePlugin::clone() const noexcept -{ - return new GemmAllReducePlugin(*this); -} - -DimsExprs GemmAllReducePlugin::getOutputDimensions( - int outputIndex, DimsExprs const* inputs, int nbInputs, IExprBuilder& exprBuilder) noexcept -{ - try - { - TLLM_CHECK(nbInputs == mNbInputs); // number of input tensors - TLLM_CHECK(inputs[0].nbDims == inputs[1].nbDims); - TLLM_CHECK(outputIndex < getNbOutputs()); - - // List of pointers to D on each rank - if ((nbInputs + outputIndex) == TensorArg::OUT_D_IPC) - { - DimsExprs out_dims; - out_dims.nbDims = 1; - out_dims.d[0] = exprBuilder.constant(mOptions.groupSize); - return out_dims; - } - - TLLM_CHECK(mOptions.transA == false); - TLLM_CHECK(mOptions.transB == true); - - int const nbDimsA = inputs[0].nbDims; // number of dims - int const nbDimsB = inputs[1].nbDims; - - DimsExprs out_dims; - // subtract 2 -> K from each input - out_dims.nbDims = nbDimsA + nbDimsB - 2; - - if (mOptions.transA) - { - for (int i = 1; i < nbDimsA; ++i) - { - out_dims.d[i - 1] = inputs[0].d[i]; - } - } - else - { - for (int i = 0; i < nbDimsA - 1; ++i) - { - out_dims.d[i] = inputs[0].d[i]; - } - } - if (mOptions.transB) - { - for (int i = 0; i < nbDimsB - 1; ++i) - { - out_dims.d[nbDimsA - 1 + i] = inputs[1].d[i]; - } - } - else - { - for (int i = 1; i < nbDimsB; ++i) - { - out_dims.d[nbDimsA - 2 + i] = inputs[1].d[i]; - } - } - return out_dims; - } - catch (std::exception const& e) - { - caughtError(e); - } - return DimsExprs{}; -} - -bool GemmAllReducePlugin::supportsFormatCombination( - int32_t pos, PluginTensorDesc const* inOut, int32_t nbInputs, int32_t nbOutputs) noexcept -{ - // inOut[0] -> activation - // inOut[1] -> weight - // inOut[1+hasInputSF] -> activation_sf - // inOut[1+hasInputSF*2] -> weight_sf - // inOut[2+hasInputSF*2] -> output[0] = D_uc - // inOut[3+hasInputSF*2] -> output[1] = D_mc - - TLLM_CHECK_WITH_INFO(pos < mNbInputs + mNbOutputs, "Unexpected pos: %d", pos); - auto const& desc = inOut[pos]; - - TLLM_CHECK_WITH_INFO(mArgMap.count(pos) > 0, "pos %d not found in mArgMap.", pos); - TensorArg arg = mArgMap[pos]; - - auto typeExists = [&](DataType dtype, auto idx) -> bool - { - for (const auto& [key, value] : mTypedInstantiators) - { - // key format: - if (std::get(key) == dtype) - { - return true; - } - } - return false; - }; - - switch (arg) - { - case TensorArg::IN_ACTIVATION: return typeExists(desc.type, std::integral_constant{}); - case TensorArg::IN_WEIGHT: return typeExists(desc.type, std::integral_constant{}); - case TensorArg::IN_ACTIVATION_SF: - case TensorArg::IN_WEIGHT_SF: - // Assumed SF for only FP4 at the moment - return desc.type == DataType::kFP8; - case TensorArg::IN_ALPHA: return desc.type == DataType::kFLOAT; - case TensorArg::OUT_D_UC: - case TensorArg::OUT_D_MC: - case TensorArg::OUT_D_IPC: return typeExists(desc.type, std::integral_constant{}); - default: return false; - } -} - -void GemmAllReducePlugin::configurePlugin( - DynamicPluginTensorDesc const* in, int32_t nbInputs, DynamicPluginTensorDesc const* out, int32_t nbOutputs) noexcept -{ - // Get problem shape - int const nbDimsA = in[0].max.nbDims; - int const minM = utils::computeMDimension(mOptions.transA, in[0].min); - int const maxM = utils::computeMDimension(mOptions.transA, in[0].max); - int const N = utils::computeNDimension(mOptions.transB, in[1].max); - int const K = mOptions.transA ? in[0].max.d[0] : in[0].max.d[nbDimsA - 1]; - - TLLM_CHECK_WITH_INFO(out[0].desc.type == mOptions.typeD, "Output type mismatch."); - - // Ensure call from execution phase does - // not override call from build phase - if (!mOptions.maxProblemShape.isInitialized()) - { - mOptions.maxProblemShape = {minM, maxM, N, K}; - mGemmId = {N, K, mOptions.typeD}; - } - - // Build phase doesn't have COMM_SESSION (i.e built on single rank) - // so do not allocate persistent workspace - if (!isBuilding()) - { - auto getTPRank = [&]() - { - int rank = COMM_SESSION.getRank(); - auto it = std::find(mOptions.group.begin(), mOptions.group.end(), rank); - TLLM_CHECK_WITH_INFO(it != mOptions.group.end(), - "Incorrect group specified - rank " + std::to_string(rank) + " not found in group"); - return std::distance(mOptions.group.begin(), it); - }; - - mRank = getTPRank(); - - if (mWorkspace == nullptr) - { - allocatePersistentWorkspace(); - } - } -} - -size_t GemmAllReducePlugin::getWorkspaceSize( - PluginTensorDesc const* inputs, int32_t nbInputs, PluginTensorDesc const* outputs, int32_t nbOutputs) const noexcept -{ - return 0; -} - -int GemmAllReducePlugin::enqueue(PluginTensorDesc const* inputDesc, PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept -{ - // inputs[0] -> [M(*), K] - // inputs[1] -> [K, N] - // outputs[0] -> [M(*), N] unicast ptr - // outputs[1] -> [M(*), N] multicast ptr - auto const nbDimsA = inputDesc[0].dims.nbDims; - auto const M = utils::computeMDimension(mOptions.transA, inputDesc[0].dims); - auto const N = utils::computeNDimension(mOptions.transB, inputDesc[1].dims); - auto const K = mOptions.transA ? inputDesc[0].dims.d[0] : inputDesc[0].dims.d[nbDimsA - 1]; - - TLLM_CHECK_WITH_INFO(M <= mOptions.maxProblemShape.maxM, "GemmAllReducePlugin M > maxM."); - TLLM_CHECK_WITH_INFO(M > 0, "GemmAllReducePlugin M is 0."); - TLLM_CHECK_WITH_INFO(N > 0, "GemmAllReducePlugin N is 0."); - TLLM_CHECK_WITH_INFO(K > 0, "GemmAllReducePlugin K is 0."); - TLLM_CHECK_WITH_INFO(mWorkspace != nullptr, "GemmAllReducePlugin workspace is null."); - - LaunchConfig bestLaunchConfig; - if (mProfiler->useProfiler()) - { - bestLaunchConfig = mProfiler->getBestConfig(M, mGemmId).value(); - } - else - { - bestLaunchConfig = getStaticHeuristicLaunchConfig(M); - } - - void const* activation = inputs[mArgInvMap[TensorArg::IN_ACTIVATION]]; - void const* weight = inputs[mArgInvMap[TensorArg::IN_WEIGHT]]; - void* D_out_uc = outputs[mArgInvMap[TensorArg::OUT_D_UC] - mNbInputs]; - void* D_out_mc = outputs[mArgInvMap[TensorArg::OUT_D_MC] - mNbInputs]; - void* D_out_ipc = outputs[mArgInvMap[TensorArg::OUT_D_IPC] - mNbInputs]; - - TLLM_CHECK_WITH_INFO(activation != nullptr, "GemmAllReducePlugin activation is NULL"); - TLLM_CHECK_WITH_INFO(weight != nullptr, "GemmAllReducePlugin weight is NULL"); - TLLM_CHECK_WITH_INFO(D_out_uc != nullptr, "GemmAllReducePlugin out_uc is NULL"); - TLLM_CHECK_WITH_INFO(D_out_mc != nullptr, "GemmAllReducePlugin out_mc is NULL"); - TLLM_CHECK_WITH_INFO(D_out_ipc != nullptr, "GemmAllReducePlugin out_ipc is NULL"); - - cutlass_kernels::GemmAllReduceImplInterface::ProblemArgs args; - args.argProblemShape(M, N, K, 1) - .argA(activation) - .argB(weight) - .argC(nullptr) - .argD(D_out_uc, D_out_mc, (void**) D_out_ipc) - .argRanks(mRank, mOptions.group) - .argBeta(0.f) // no bias - .argLaunchConfig(bestLaunchConfig) - .argWorkspace(mWorkspace->mWorkspace.get()); - // tensor for scaling input A - if (mOptions.hasSFA) - { - void const* activation_sf = inputs[mArgInvMap[TensorArg::IN_ACTIVATION_SF]]; - TLLM_CHECK_WITH_INFO(activation_sf != nullptr, "GemmAllReducePlugin activation_sf is NULL"); - args.argAScale(activation_sf); - } - // tensor for scaling input B - if (mOptions.hasSFB) - { - void const* weight_sf = inputs[mArgInvMap[TensorArg::IN_WEIGHT_SF]]; - TLLM_CHECK_WITH_INFO(weight_sf != nullptr, "GemmAllReducePlugin weight_sf is NULL"); - args.argBScale(weight_sf); - } - // tensor for scaling output D - if (mOptions.alphaIsPtr) - { - void const* alpha_vec = inputs[mArgInvMap[TensorArg::IN_ALPHA]]; - TLLM_CHECK_WITH_INFO(alpha_vec != nullptr, "GemmAllReducePlugin alpha_vec is NULL"); - args.argAlphaPtr(reinterpret_cast(alpha_vec)); - } - else - { - args.argAlpha(mOptions.alpha); - } - - mGemm->run(args, stream); - - return 0; -} - -////////////////////////////////// -// IPluginV2Ext Methods -////////////////////////////////// -DataType GemmAllReducePlugin::getOutputDataType(int index, DataType const* inputTypes, int nbInputs) const noexcept -{ - TLLM_CHECK_WITH_INFO(index < getNbOutputs(), "Output index out of bounds: %d", index); - return mOptions.typeD; -} - -////////////////////////////////// -// IPluginV2 Methods -////////////////////////////////// -char const* GemmAllReducePlugin::getPluginType() const noexcept -{ - return GEMM_ALLREDUCE_PLUGIN_NAME; -} - -char const* GemmAllReducePlugin::getPluginVersion() const noexcept -{ - return GEMM_ALLREDUCE_PLUGIN_VERSION; -} - -int GemmAllReducePlugin::getNbOutputs() const noexcept -{ - return mNbOutputs; -} - -int GemmAllReducePlugin::initialize() noexcept -{ - if (isBuilding() && mProfiler->useProfiler()) - { - // TODO (xsimmons): interfaces between GemmPluginProfiler and Plugin - // needs to be relooked at - current interface implicitly assigns runner to profiler - // object in profileTactics() - assert(mOptions.maxProblemShape.isInitialized()); - mProfiler->profileTactics(mGemm, mOptions.typeD, mOptions.maxProblemShape, mGemmId); - } - return 0; -} - -void GemmAllReducePlugin::terminate() noexcept -{ - if (isBuilding()) // need this otherwise getComm will crash during build phase - { - return; - } - - // free mWorkspace - if (mWorkspace) - { - getPluginRegistry()->releasePluginResource(mWorkspaceKey.c_str()); - mWorkspace = nullptr; - } -} - -size_t GemmAllReducePlugin::getSerializationSize() const noexcept -{ - // cannot use sizeof(GemmAllReducePluginOptions) - // becaused need packed attribute which doesn't work on enum - // without making the enum also packed - size_t size = 0; - size += sizeof(mOptions.typeA); - size += sizeof(mOptions.typeB); - size += sizeof(mOptions.typeD); - size += sizeof(mOptions.transA); - size += sizeof(mOptions.transB); - size += sizeof(mOptions.alpha); - size += sizeof(mOptions.maxProblemShape); - size += sizeof(mOptions.groupSize); - size += mOptions.group.size() * sizeof(int); - size += sizeof(mOptions.hasSFA); - size += sizeof(mOptions.hasSFB); - size += sizeof(mOptions.alphaIsPtr); - return size; -} - -void GemmAllReducePlugin::serialize(void* buffer) const noexcept -{ - char* begin = reinterpret_cast(buffer); - char* end = reinterpret_cast(buffer); - - write(end, mOptions.typeA); - write(end, mOptions.typeB); - write(end, mOptions.typeD); - write(end, mOptions.transA); - write(end, mOptions.transB); - write(end, mOptions.alpha); - write(end, mOptions.maxProblemShape); - write(end, mOptions.groupSize); - for (auto const& rank : mOptions.group) - { - write(end, rank); - } - write(end, mOptions.hasSFA); - write(end, mOptions.hasSFB); - write(end, mOptions.alphaIsPtr); - TLLM_CHECK(end == begin + getSerializationSize()); - - // Profiler MNK->kernel mappings need to be deterministic and consistent across ranks - // to ensure correct functionality (unlike standalone GEMMs). - // Since by default each rank will generate and serialize its own profiler mapping - // this can lead to different mappings between ranks which will result in fatal - // error. Therefore only generate and use profiler mapping for single rank. - if (mProfiler->useProfiler() && COMM_SESSION.getRank() == 0) - { - mProfiler->serializeToOwnFile(mGemmId); - } -} - -void GemmAllReducePlugin::destroy() noexcept -{ - delete this; -} - -//////////////////////////////////////////////////////////// -// GemmAllReducePluginCreator Methods -//////////////////////////////////////////////////////////// -PluginFieldCollection GemmAllReducePluginCreator::mFC; -std::vector GemmAllReducePluginCreator::mPluginAttributes; - -GemmAllReducePluginCreator::GemmAllReducePluginCreator() -{ - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - mPluginAttributes.emplace_back("type_a", nullptr, PluginFieldType::kINT32, 1); - mPluginAttributes.emplace_back("type_b", nullptr, PluginFieldType::kINT32, 1); - mPluginAttributes.emplace_back("type_d", nullptr, PluginFieldType::kINT32, 1); - mPluginAttributes.emplace_back("transa", nullptr, PluginFieldType::kINT32, 1); - mPluginAttributes.emplace_back("transb", nullptr, PluginFieldType::kINT32, 1); - mPluginAttributes.emplace_back("alpha", nullptr, PluginFieldType::kFLOAT32, 1); - mPluginAttributes.emplace_back("group", nullptr, PluginFieldType::kINT32, 1); - mPluginAttributes.emplace_back("has_sfa", nullptr, PluginFieldType::kINT8, 1); - mPluginAttributes.emplace_back("has_sfb", nullptr, PluginFieldType::kINT8, 1); - mPluginAttributes.emplace_back("alpha_is_ptr", nullptr, PluginFieldType::kINT8, 1); - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -char const* GemmAllReducePluginCreator::getPluginName() const noexcept -{ - return GEMM_ALLREDUCE_PLUGIN_NAME; -} - -char const* GemmAllReducePluginCreator::getPluginVersion() const noexcept -{ - return GEMM_ALLREDUCE_PLUGIN_VERSION; -} - -PluginFieldCollection const* GemmAllReducePluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -IPluginV2* GemmAllReducePluginCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept -{ - PluginField const* fields = fc->fields; - GemmAllReducePluginOptions options; - options.deserialize = false; - - // Read configurations from each fields - for (int i = 0; i < fc->nbFields; ++i) - { - char const* attrName = fields[i].name; - if (!strcmp(attrName, "type_a")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - options.typeA = *static_cast(fields[i].data); - } - else if (!strcmp(attrName, "type_b")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - options.typeB = *static_cast(fields[i].data); - } - else if (!strcmp(attrName, "type_d")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - options.typeD = *static_cast(fields[i].data); - } - else if (!strcmp(attrName, "transa")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - options.transA = *static_cast(fields[i].data); - } - else if (!strcmp(attrName, "transb")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - options.transB = *static_cast(fields[i].data); - } - else if (!strcmp(attrName, "alpha")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kFLOAT32); - options.alpha = *static_cast(fields[i].data); - } - else if (!strcmp(attrName, "group")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - auto const* ranks = static_cast(fields[i].data); - for (int j = 0; j < fields[i].length; ++j) - { - options.group.insert(ranks[j]); - } - options.groupSize = options.group.size(); - } - else if (!strcmp(attrName, "has_sfa")) // passed in as input tensor - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT8); - options.hasSFA = *static_cast(fields[i].data); - } - else if (!strcmp(attrName, "has_sfb")) // passed in as input tensor - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT8); - options.hasSFB = *static_cast(fields[i].data); - } - else if (!strcmp(attrName, "alpha_is_ptr")) // passed in as input tensor - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT8); - options.alphaIsPtr = *static_cast(fields[i].data); - } - } - - try - { - // GemmAllReducePluginCreator is unique and shared for an engine generation - auto* obj = new GemmAllReducePlugin(options); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - return nullptr; - } -} - -IPluginV2* GemmAllReducePluginCreator::deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept -{ - // This object will be deleted when the network is destroyed, which will - // call GemmAllReducePlugin::destroy() - try - { - auto* obj = new GemmAllReducePlugin(serialData, serialLength); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/gemmAllReducePlugin/gemmAllReducePlugin.h b/cpp/tensorrt_llm/plugins/gemmAllReducePlugin/gemmAllReducePlugin.h deleted file mode 100644 index 457926246002..000000000000 --- a/cpp/tensorrt_llm/plugins/gemmAllReducePlugin/gemmAllReducePlugin.h +++ /dev/null @@ -1,189 +0,0 @@ -/* - * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved. - * - * 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. - */ -#pragma once - -#if defined(USING_OSS_CUTLASS_ALLREDUCE_GEMM) -#include "tensorrt_llm/kernels/cutlass_kernels/include/allreduce_gemm_runner.h" -#else -#include "allreduce_gemm_runner.h" -#endif - -#include "gemmAllReducePluginProfiler.h" -#include "gemmAllReducePluginResource.h" -#include "tensorrt_llm/plugins/common/plugin.h" -#include "tensorrt_llm/runtime/utils/mpiUtils.h" - -using namespace nvinfer1; - -using nvinfer1::DataType; -#if defined(USING_OSS_CUTLASS_ALLREDUCE_GEMM) -namespace cutlass_kernels = ::tensorrt_llm::kernels::opened_cutlass_kernels; -#else -namespace cutlass_kernels = ::tensorrt_llm::kernels::cutlass_kernels; -#endif - -using LaunchConfig = typename cutlass_kernels::GemmAllReduceImplInterface::LaunchConfig; - -namespace tensorrt_llm::plugins -{ -struct GemmAllReducePluginOptions -{ - // Don't need to specify problem shape, this - // is specified in configurePlugin - DataType typeA; - DataType typeB; - DataType typeD; - int transA; - int transB; - float alpha; - // ranks participating in collective - std::set group; - int groupSize; - // Set in configurePlugin during build phase - GemmDims maxProblemShape; - bool deserialize; // used for profiler instantiation - int8_t hasSFA = 0; - int8_t hasSFB = 0; - int8_t alphaIsPtr = 0; -}; - -class GemmAllReducePlugin : public BasePlugin -{ - friend class GemmAllReducePluginCreator; - -public: - ~GemmAllReducePlugin() override = default; - - ////////////////////////////////// - // IPluginV2DynamicExt Methods - ////////////////////////////////// - IPluginV2DynamicExt* clone() const noexcept override; - - DimsExprs getOutputDimensions( - int outputIndex, DimsExprs const* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept override; - - // inOut[0] -> activation - // inOut[1] -> weight - // inOut[2] -> result - bool supportsFormatCombination( - int32_t pos, PluginTensorDesc const* inOut, int32_t nbInputs, int32_t nbOutputs) noexcept override; - - // in[0] -> activation - // in[1] -> weight - // no bias needed - void configurePlugin(DynamicPluginTensorDesc const* in, int32_t nbInputs, DynamicPluginTensorDesc const* out, - int32_t nbOutputs) noexcept override; - - size_t getWorkspaceSize(PluginTensorDesc const* inputs, int32_t nbInputs, PluginTensorDesc const* outputs, - int32_t nbOutputs) const noexcept override; - - // in[0] -> activation - // in[1] -> weight - // out[0] -> result_uc - // out[1] -> result_mc - int enqueue(PluginTensorDesc const* inputDesc, PluginTensorDesc const* outputDesc, void const* const* inputs, - void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - - ////////////////////////////////// - // IPluginV2Ext Methods - ////////////////////////////////// - DataType getOutputDataType(int index, DataType const* inputTypes, int nbInputs) const noexcept override; - - ////////////////////////////////// - // IPluginV2 Methods - ////////////////////////////////// - char const* getPluginType() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - int getNbOutputs() const noexcept override; - - int initialize() noexcept override; - - void terminate() noexcept override; - - size_t getSerializationSize() const noexcept override; - - void serialize(void* buffer) const noexcept override; - - void destroy() noexcept override; - -private: - explicit GemmAllReducePlugin(GemmAllReducePluginOptions const& options); - // Parameterized constructor - explicit GemmAllReducePlugin(void const* data, size_t length); - - void allocatePersistentWorkspace(); - - LaunchConfig getStaticHeuristicLaunchConfig(int M) const; - - // Params that are initialized during constructor - using KeyType = std::tuple; - using ValueType = std::function; - GemmAllReducePluginOptions mOptions; - int mRank = 0; - - enum TensorArg - { - IN_ACTIVATION, - IN_ACTIVATION_SF, - IN_WEIGHT, - IN_WEIGHT_SF, - IN_ALPHA, - OUT_D_UC, - OUT_D_MC, - OUT_D_IPC - }; - - std::unordered_map mArgMap; - std::unordered_map mArgInvMap; - int mNbInputs = 0; - int mNbOutputs = 0; - - std::map mTypedInstantiators; - std::string mWorkspaceKey; - std::shared_ptr mGemm; - // Params that are initialized during configurePlugin() - GemmAllReducePersistentWorkspace* mWorkspace = nullptr; - - // Used for selecting best GEMM for given problem shapes - GemmIdCore mGemmId{}; - GemmPluginProfilerManager mGemmPluginProfileManager; - std::shared_ptr mProfiler; -}; - -class GemmAllReducePluginCreator : public BaseCreator -{ -public: - GemmAllReducePluginCreator(); - - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; - - nvinfer1::IPluginV2* deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept override; - -private: - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; -}; - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/gemmAllReducePlugin/gemmAllReducePluginProfiler.cpp b/cpp/tensorrt_llm/plugins/gemmAllReducePlugin/gemmAllReducePluginProfiler.cpp deleted file mode 100644 index a6f7ca2615df..000000000000 --- a/cpp/tensorrt_llm/plugins/gemmAllReducePlugin/gemmAllReducePluginProfiler.cpp +++ /dev/null @@ -1,140 +0,0 @@ -/* - * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved. - * - * 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. - */ -#include "gemmAllReducePlugin.h" -#include "tensorrt_llm/common/dataType.h" -#include "tensorrt_llm/kernels/cutlass_kernels/cutlass_type_conversion.h" -#include "tensorrt_llm/plugins/common/pluginUtils.h" - -namespace tc = tensorrt_llm::common; - -namespace tensorrt_llm::plugins -{ -void GemmAllReducePluginProfiler::serializeToOwnFile(GemmIdCore gemmId) -{ - std::vector file_buf(getSerializationSize(gemmId)); - char* begin = file_buf.data(); - char* end = file_buf.data(); - serialize(end, gemmId); - assert(end == begin + file_buf.size()); - - auto fileName = getCacheFileName(gemmId); - std::ofstream file(fileName, std::ios::binary); - TLLM_CHECK(file.is_open()); - file.write(begin, file_buf.size()); - file.flush(); - file.close(); -} - -void GemmAllReducePluginProfiler::deserializeFromOwnFile(GemmIdCore gemmId, GemmDims problemShape) -{ - auto fileName = getCacheFileName(gemmId); - std::ifstream file(fileName, std::ios::binary); - TLLM_CHECK(file.is_open()); - file.seekg(0, std::ios::end); - std::streamsize size = file.tellg(); - TLLM_CHECK(size > 0); - file.seekg(0, std::ios::beg); - - std::vector file_buf(size); - file.read(file_buf.data(), size); - file.close(); - - char const* begin = const_cast(file_buf.data()); - char const* end = begin; - deserialize(end, problemShape, gemmId); - assert(end == begin + size); -} - -bool GemmAllReducePluginProfiler::useProfiler() -{ - // char const* envDir = getenv("GEMM_AR_PLUGIN_PROFILE_DIR"); - // return envDir != nullptr; - // TODO(xsimmons): currently the profiler does not add any perf gain - // due to static heuristics being sufficient. We can re-enable this - // when we need more configurations. - return false; -} - -std::string GemmAllReducePluginProfiler::getCacheFileName(GemmIdCore gemmId) -{ - std::stringstream fileName; - char const* envDir = getenv("GEMM_AR_PLUGIN_PROFILE_DIR"); - std::string directory = envDir ? std::string(envDir) : "/tmp/"; - fileName << directory + "/gemm-AR"; - fileName << "-n" << std::to_string(gemmId.n); - fileName << "-k" << std::to_string(gemmId.k); - fileName << "-" << tc::getDtypeString(gemmId.dtype); - fileName << ".prof_cache"; - return fileName.str(); -} - -void GemmAllReducePluginProfiler::runTactic(int m, int n, int k, - cutlass_kernels::GemmAllReduceImplInterface::LaunchConfig const& tactic, char* workspace, - cudaStream_t const& stream) -{ - const size_t dtype_size = tc::getDTypeSize(mType); - char* inputA = workspace; - char* inputB = inputA + m * k * dtype_size; - char* outputD = inputB + n * k * dtype_size; - char* inputSFA = outputD + m * n * dtype_size; - char* inputSFB = inputSFA + m * k * dtype_size; - std::set tpGroup = {0}; - - // Run on single-GPU - cutlass_kernels::GemmAllReduceImplInterface::ProblemArgs args; - args.argProblemShape(m, n, k, 1) - .argA((void*) inputA) - .argB((void*) inputB) - .argD((void*) outputD, /*output_mc=*/nullptr) - .argAScale((void*) inputSFA) - .argBScale((void*) inputSFB) - .argRanks(0, tpGroup) - .argAlpha(1.f) - .argBeta(0.f) // no bias - .argLaunchConfig(tactic); - - TLLM_CHECK(mRunner != nullptr); - mRunner->run(args, stream); -} - -void GemmAllReducePluginProfiler::computeTmpSize(size_t maxM, size_t n, size_t k) -{ - TLLM_CHECK(maxM != 0); - TLLM_CHECK(n != 0); - TLLM_CHECK(k != 0); - // mType refers to the output data type - // WARNING: This code assumes that the output precision is >= to input precision - const size_t dtype_size = tc::getDTypeSize(mType); - size_t bytes = 0; - bytes += maxM * k * dtype_size; // A - bytes += n * k * dtype_size; // B - // No C - // Note that D is typically IPC, however, when tuning GEMM we need it to run on single GPU - bytes += maxM * n * dtype_size; // D - // scale tensors for A & B - will at most be same size as A/B - bytes += maxM * k * dtype_size; // A - bytes += n * k * dtype_size; // B - - setTmpWorkspaceSizeInBytes(bytes); -} - -std::vector GemmAllReducePluginProfiler::getTactics( - int m, int n, int k) const -{ - TLLM_CHECK(mRunner != nullptr); - return mRunner->getSupportedLaunchConfigs(); -} -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/gemmAllReducePlugin/gemmAllReducePluginProfiler.h b/cpp/tensorrt_llm/plugins/gemmAllReducePlugin/gemmAllReducePluginProfiler.h deleted file mode 100644 index faacbb3b8c0f..000000000000 --- a/cpp/tensorrt_llm/plugins/gemmAllReducePlugin/gemmAllReducePluginProfiler.h +++ /dev/null @@ -1,68 +0,0 @@ -/* - * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved. - * - * 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. - */ -#pragma once - -#if defined(USING_OSS_CUTLASS_ALLREDUCE_GEMM) -#include "tensorrt_llm/kernels/cutlass_kernels/include/allreduce_gemm_runner.h" -#else -#include "allreduce_gemm_runner.h" -#endif -#include "tensorrt_llm/plugins/common/gemmPluginProfiler.h" -#include "tensorrt_llm/plugins/common/plugin.h" - -namespace tensorrt_llm::plugins -{ -/* - * Used for tuning to find best GEMM configs for different problem shapes. - * WARNING: Tuning GEMM+AR kernel may not be fully representable of real - * multi-GPU workloads as tuning only runs on single-GPU. - * IMPORTANT: TRT-LLM does not support deterministic tuning across ranks. - * Because of this, we have to serialize/deserialize our own configuration file. - */ - -#if defined(USING_OSS_CUTLASS_ALLREDUCE_GEMM) -namespace cutlass_kernels = ::tensorrt_llm::kernels::opened_cutlass_kernels; -#else -namespace cutlass_kernels = ::tensorrt_llm::kernels::cutlass_kernels; -#endif -class GemmAllReducePluginProfiler - : public GemmPluginProfiler, GemmIdCore, GemmIdCoreHash> -{ -public: - void serializeToOwnFile(GemmIdCore gemmId); - - void deserializeFromOwnFile(GemmIdCore gemmId, GemmDims problemShape); - - bool useProfiler(); - -protected: - //////////////////////////////////// - // GemmPluginProfiler methods - //////////////////////////////////// - void runTactic(int m, int n, int k, cutlass_kernels::GemmAllReduceImplInterface::LaunchConfig const& tactic, - char* workspace, cudaStream_t const& stream) override; - - void computeTmpSize(size_t maxM, size_t n, size_t k) override; - - std::vector getTactics( - int m, int n, int k) const override; - -private: - static std::string getCacheFileName(GemmIdCore gemmId); -}; - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/gemmAllReducePlugin/gemmAllReducePluginResource.h b/cpp/tensorrt_llm/plugins/gemmAllReducePlugin/gemmAllReducePluginResource.h deleted file mode 100644 index 8136bd363bd7..000000000000 --- a/cpp/tensorrt_llm/plugins/gemmAllReducePlugin/gemmAllReducePluginResource.h +++ /dev/null @@ -1,82 +0,0 @@ -/* - * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved. - * - * 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. - */ -#pragma once - -#include "NvInferPlugin.h" - -#if defined(USING_OSS_CUTLASS_ALLREDUCE_GEMM) -#include "tensorrt_llm/kernels/cutlass_kernels/include/allreduce_gemm_runner.h" -#else -#include "allreduce_gemm_runner.h" -#endif -#include "tensorrt_llm/common/cudaUtils.h" -#include "tensorrt_llm/plugins/common/plugin.h" - -using namespace nvinfer1; - -namespace tensorrt_llm::plugins -{ - -#if defined(USING_OSS_CUTLASS_ALLREDUCE_GEMM) -namespace cutlass_kernels = ::tensorrt_llm::kernels::opened_cutlass_kernels; -#else -namespace cutlass_kernels = ::tensorrt_llm::kernels::cutlass_kernels; -#endif -class GemmAllReducePersistentWorkspace : public IPluginResource -{ -public: - GemmAllReducePersistentWorkspace(std::shared_ptr workspace) - : mWorkspace(workspace) - { - } - - ////////////////////////////////// - // IPluginResource Methods - ////////////////////////////////// - IPluginResource* clone() noexcept override - { - auto copy = new GemmAllReducePersistentWorkspace(mWorkspace); - // Resource initialization (if any) may be skipped for non-cloned objects - // since only clones will be registered by TensorRT. - try - { - copy->mWorkspace->allocate(); - return copy; - } - catch (std::exception const& e) - { - TLLM_LOG_ERROR(e.what()); - return nullptr; - } - } - - int32_t release() noexcept override - { - try - { - return mWorkspace->free(); - } - catch (std::exception const& e) - { - TLLM_LOG_ERROR(e.what()); - return -1; - } - } - - std::shared_ptr mWorkspace; -}; - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/gemmPlugin/CMakeLists.txt b/cpp/tensorrt_llm/plugins/gemmPlugin/CMakeLists.txt deleted file mode 100644 index 86876224fccd..000000000000 --- a/cpp/tensorrt_llm/plugins/gemmPlugin/CMakeLists.txt +++ /dev/null @@ -1,21 +0,0 @@ -# -# SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. -# -file(GLOB SRCS *.cpp) -set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS}) -set(PLUGIN_SOURCES - ${PLUGIN_SOURCES} - PARENT_SCOPE) diff --git a/cpp/tensorrt_llm/plugins/gemmPlugin/gemmPlugin.cpp b/cpp/tensorrt_llm/plugins/gemmPlugin/gemmPlugin.cpp deleted file mode 100644 index 9e06ad01d10f..000000000000 --- a/cpp/tensorrt_llm/plugins/gemmPlugin/gemmPlugin.cpp +++ /dev/null @@ -1,614 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. - */ - -#include "gemmPlugin.h" - -#include "gemmPluginProfiler.h" -#include "plugin.h" -#include "pluginUtils.h" -#include "tensorrt_llm/kernels/weightOnlyBatchedGemv/cudaCoreGemm.h" -#include "tensorrt_llm/runtime/utils/debugUtils.h" - -#include - -#include - -using namespace nvinfer1; -using namespace tensorrt_llm::common; -using tensorrt_llm::plugins::GemmDims; -using tensorrt_llm::plugins::GemmPluginCreator; -using tensorrt_llm::plugins::GemmPlugin; -using tensorrt_llm::plugins::CublasLtGemmPluginProfiler; -using tensorrt_llm::plugins::CublasGemmWrapperPtr; -using tensorrt_llm::plugins::read; -using tensorrt_llm::plugins::write; - -static char const* GEMM_PLUGIN_VERSION{"1"}; -static char const* GEMM_PLUGIN_NAME{"Gemm"}; -PluginFieldCollection GemmPluginCreator::mFC{}; -std::vector GemmPluginCreator::mPluginAttributes; - -void getProblemParams(cublasOperation_t& transa, cublasOperation_t& transb, int& m, int& n, int& k, int& lda, int& ldb, - int& ldc, bool transA, bool transB, int M, int N, int K, int padLda, int padLdb, int padLdc) -{ - transa = transB ? CUBLAS_OP_T : CUBLAS_OP_N; - transb = transA ? CUBLAS_OP_T : CUBLAS_OP_N; - m = N; - n = M; - k = K; - lda = transB ? K + padLdb : N + padLdb; - ldb = transA ? M + padLda : K + padLda; - ldc = N + padLdc; -} - -void runGemm(int const M, int const N, int const K, bool const transA, bool const transB, int const padLda, - int const padLdb, int const padLdc, nvinfer1::DataType const type, CublasGemmWrapperPtr const& cublasWrapperPtr, - void const* act, void const* weight, float const alpha, void* output, - std::optional const& heuristic, void* workspace, cudaStream_t stream) -{ - if (M == 0 || N == 0 || K == 0) - return; - - cublasWrapperPtr->setStream(stream); - cublasWrapperPtr->setWorkspace(workspace); - - cublasOperation_t transa, transb; - int m, n, k; - int lda, ldb, ldc; - getProblemParams(transa, transb, m, n, k, lda, ldb, ldc, transA, transB, M, N, K, padLda, padLdb, padLdc); - - cublasWrapperPtr->createDescriptors(transa, transb, m, n, k, lda, ldb, ldc); - cublasWrapperPtr->Gemm(transa, transb, m, n, k, weight, lda, act, ldb, output, ldc, alpha, 0.0f, heuristic); - cublasWrapperPtr->destroyDescriptors(); -} - -void CublasLtGemmPluginProfiler::runTactic( - int m, int n, int k, CublasLtGemmPluginProfiler::Config const& tactic, char* workspace, cudaStream_t const& stream) -{ - size_t dataSize = sizeof(half); - if (mType == nvinfer1::DataType::kFLOAT) - { - dataSize = sizeof(float); - } - - void* actPtr = reinterpret_cast(workspace); - void* weightPtr = reinterpret_cast( - nextWorkspacePtrWithAlignment(reinterpret_cast(actPtr), m * k * dataSize, ALIGNMENT)); - void* outputPtr = reinterpret_cast( - nextWorkspacePtrWithAlignment(reinterpret_cast(weightPtr), n * k * dataSize, ALIGNMENT)); - char* workspacePtr = reinterpret_cast( - nextWorkspacePtrWithAlignment(reinterpret_cast(outputPtr), m * (n + mPadLdc) * dataSize, ALIGNMENT)); - runGemm(m, n, k, mTransA, mTransB, mPadLda, mPadLdb, mPadLdc, mType, mRunner, actPtr, weightPtr, 1.0f, outputPtr, - {tactic}, workspacePtr, stream); -} - -bool CublasLtGemmPluginProfiler::checkTactic(int m, int n, int k, Config const& tactic) const -{ - cublasOperation_t transa, transb; - int M = m, N = n, K = k; - int lda, ldb, ldc; - getProblemParams(transa, transb, m, n, k, lda, ldb, ldc, mTransA, mTransB, M, N, K, mPadLda, mPadLdb, mPadLdc); - - mRunner->createDescriptors(transa, transb, m, n, k, lda, ldb, ldc); - - auto const checkResult = mRunner->checkTactic(transa, transb, m, n, k, lda, ldb, ldc, tactic.algo); - - mRunner->destroyDescriptors(); - - return checkResult; -} - -void CublasLtGemmPluginProfiler::computeTmpSize(size_t maxM, size_t n, size_t k) -{ - size_t dataSize = getDTypeSize(mType); - size_t outputDataSize = getDTypeSize(mOutputType); - - std::vector workspaces = { - maxM * k * dataSize, // A - n * k * dataSize, // B - maxM * (n + mPadLdc) * outputDataSize, // C - CUBLAS_WORKSPACE_SIZE // workspace - }; - size_t bytes = calculateTotalWorkspaceSize(workspaces.data(), workspaces.size(), ALIGNMENT); - setTmpWorkspaceSizeInBytes(bytes); -} - -std::vector CublasLtGemmPluginProfiler::getTactics(int M, int N, int K) const -{ - cublasOperation_t transa, transb; - int m, n, k; - int lda, ldb, ldc; - getProblemParams(transa, transb, m, n, k, lda, ldb, ldc, mTransA, mTransB, M, N, K, mPadLda, mPadLdb, mPadLdc); - - mRunner->createDescriptors(transa, transb, m, n, k, lda, ldb, ldc); - auto const heruistics = mRunner->getTactics(transa, transb, m, n, k, lda, ldb, ldc); - mRunner->destroyDescriptors(); - - return heruistics; -} - -GemmPlugin::GemmPlugin(int transA, int transB, int padLda, int padLdb, int padLdc, nvinfer1::DataType type, bool useFp8, - float alpha, GemmPlugin::PluginProfilerPtr const& pluginProfiler) - : mTransA(transA) - , mTransB(transB) - , mPadLda(padLda) - , mPadLdb(padLdb) - , mPadLdc(padLdc) - , mType(type) - , mOutputType(type) - , mUseFp8(useFp8) - , mAlpha(alpha) - , mPluginProfiler(pluginProfiler) -{ - init(); -} - -// Parameterized constructor -GemmPlugin::GemmPlugin(void const* data, size_t length, GemmPlugin::PluginProfilerPtr const& pluginProfiler) - : mPluginProfiler(pluginProfiler) -{ - char const *d = reinterpret_cast(data), *a = d; - read(d, mTransA); - read(d, mTransB); - read(d, mPadLda); - read(d, mPadLdb); - read(d, mPadLdc); - read(d, mType); - read(d, mUseFp8); - read(d, mAlpha); - read(d, mDims); - read(d, mOutputType); - - init(); - - mPluginProfiler->deserialize(d, mDims, mGemmId); - - TLLM_CHECK_WITH_INFO(d == a + length, - "Expected length (%d) != real length (%d). This is often " - "caused by using different TensorRT LLM version to build " - "engine and run engine.", - (int) length, (int) (d - a)); -} - -thread_local CublasGemmWrapperPtr GemmPlugin::mCublasWrapper = nullptr; - -void GemmPlugin::init() -{ - auto cublasHandle = getCublasHandle(); - auto cublasLtHandle = getCublasLtHandle(); - mCublasWrapper = std::make_shared(cublasHandle, cublasLtHandle, nullptr, nullptr); - - mPluginProfiler->setTranspose(mTransA, mTransB); - mPluginProfiler->setOutputType(mOutputType); - mPluginProfiler->setPadLd(mPadLda, mPadLdb, mPadLdc); - - mGemmId = GemmIdCublas(mDims.n, mDims.k, mType, mTransA, mTransB, mOutputType); - - mArch = tensorrt_llm::common::getSMVersion(); -} - -void GemmPlugin::setGemmConfig() -{ - if (mType == nvinfer1::DataType::kHALF) - { - mCublasWrapper->setFP16GemmConfig(trtToCublasDtype(mOutputType)); - } - else if (mType == nvinfer1::DataType::kFLOAT) - { - mCublasWrapper->setFP32GemmConfig(); - } -#ifdef ENABLE_BF16 - else if (mType == nvinfer1::DataType::kBF16) - { - mCublasWrapper->setBF16GemmConfig(trtToCublasDtype(mOutputType)); - } -#endif - -#ifdef ENABLE_FP8 - if (mUseFp8) - { - mCublasWrapper->setFP8GemmConfig(trtToCublasDtype(mOutputType)); - } -#endif -} - -void GemmPlugin::configGemm() -{ - if (!mDims.isInitialized()) - { - return; - } - - setGemmConfig(); - - mPluginProfiler->profileTactics(mCublasWrapper, mType, mDims, mGemmId); -} - -// IPluginV2DynamicExt Methods -nvinfer1::IPluginV2DynamicExt* GemmPlugin::clone() const noexcept -{ - auto* plugin = new GemmPlugin(*this); - return plugin; -} - -nvinfer1::DimsExprs GemmPlugin::getOutputDimensions( - int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept -{ - try - { - TLLM_CHECK(nbInputs == 2); - TLLM_CHECK(outputIndex == 0); - int const nbDimsA = inputs[0].nbDims; - int const nbDimsB = inputs[1].nbDims; - DimsExprs ret; - ret.nbDims = nbDimsA + nbDimsB - 2; - - if (mTransA) - { - for (int i = 1; i < nbDimsA; ++i) - { - ret.d[i - 1] = inputs[0].d[i]; - } - } - else - { - for (int i = 0; i < nbDimsA - 1; ++i) - { - ret.d[i] = inputs[0].d[i]; - } - } - if (mTransB) - { - for (int i = 0; i < nbDimsB - 1; ++i) - { - ret.d[nbDimsA - 1 + i] = exprBuilder.constant(inputs[1].d[i]->getConstantValue() + mPadLdc); - } - } - else - { - for (int i = 1; i < nbDimsB; ++i) - { - ret.d[nbDimsA - 2 + i] = exprBuilder.constant(inputs[1].d[i]->getConstantValue() + mPadLdc); - } - } - return ret; - } - catch (std::exception const& e) - { - caughtError(e); - } - return DimsExprs{}; -} - -bool GemmPlugin::supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept -{ - auto const& desc = inOut[pos]; - if (desc.format != TensorFormat::kLINEAR) - { - return false; - } - - if (pos < nbInputs) - { - // If use FP8, act/weight dtype should be kFP8 - if (mUseFp8) - { - return desc.type == nvinfer1::DataType::kFP8; - } - else - { - return desc.type == mType; - } - } - - return desc.type == mType || desc.type == nvinfer1::DataType::kFLOAT; -} - -void GemmPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept -{ - auto const nbDimsA = in[0].max.nbDims; - - auto const minM = utils::computeMDimension(mTransA, in[0].min); - auto const maxM = utils::computeMDimension(mTransA, in[0].max); - auto const N = utils::computeNDimension(mTransB, in[1].max); - auto const K = static_cast(mTransA ? in[0].max.d[0] : in[0].max.d[nbDimsA - 1]); - - if (!mDims.isInitialized()) - { - mDims = {minM, maxM, N, K}; - } - mGemmId.n = N; - mGemmId.k = K; - - mOutputType = out[0].desc.type; -} - -size_t GemmPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - return CUBLAS_WORKSPACE_SIZE; -} - -int GemmPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept -{ - // inputs - // mat1 [M, K] (mTransA = False) - // mat2 [K, N] (mTransB = False) - // outputs - // mat [M, N] - if (mCublasWrapper == nullptr) - { - auto cublasHandle = getCublasHandle(); - auto cublasLtHandle = getCublasLtHandle(); - mCublasWrapper = std::make_shared(cublasHandle, cublasLtHandle, nullptr, nullptr); - } - setGemmConfig(); - - int const nbDimsA = inputDesc[0].dims.nbDims; - int const padM = mTransA ? mPadLda : 0; - int const padN = mTransB ? 0 : mPadLdb; - int const padK = mTransA ? 0 : mPadLda; - auto const M = utils::computeMDimension(mTransA, inputDesc[0].dims) - padM; - auto const N = utils::computeNDimension(mTransB, inputDesc[1].dims) - padN; - int const K = static_cast( - mTransA ? inputDesc[0].dims.d[0] - padK : inputDesc[0].dims.d[nbDimsA - 1] - padK); - - bool noPadDim = padM == 0 && padN == 0 && padK == 0 && mPadLdc == 0; - bool cudaKernelSupportType = mType == nvinfer1::DataType::kHALF || mType == nvinfer1::DataType::kFLOAT - || mType == nvinfer1::DataType::kBF16; - - // skip computation for a TRT empty tensor - if (M == 0) - { - return 0; - } - - std::string mnkStr = "MNK={" + std::to_string(M) + ", " + std::to_string(N) + ", " + std::to_string(K) + "}"; - { - std::string const activationStr = "GEMM layer's activation before GEMM with " + mnkStr; - TLLM_CHECK_DEBUG_WITH_INFO( - tensorrt_llm::runtime::utils::tensorHasInvalid(M, K, mType, inputs[0], stream, activationStr) == false, - "Found invalid number (NaN or Inf) in " + activationStr); - } - - bool cudaKernelFinished = false; - bool isArch90or100 = mArch >= 90 && mArch < 120; - // TODO: sub tensor matmul is not supported in fp8 gemm cuda kernel - if (!isArch90or100 && M <= 4 && N <= 128000 && mUseFp8 && noPadDim && cudaKernelSupportType) - { - tensorrt_llm::kernels::cuda_core_gemm::Params params(reinterpret_cast(inputs[0]), - reinterpret_cast(inputs[1]), mAlpha, reinterpret_cast(outputs[0]), M, N, K, - CUDA_R_8F_E4M3, trtToCublasDtype(mOutputType)); - cudaKernelFinished = tensorrt_llm::kernels::cuda_core_gemm::cudaCoreGemmDispatcher(params, stream); - } - else if (!isArch90or100 && ((mArch < 90 && M <= 6) || (isArch90or100 && M <= 2)) && N <= 128000 && !mUseFp8 - && noPadDim && cudaKernelSupportType) - { - tensorrt_llm::kernels::cuda_core_gemm::Params params(reinterpret_cast(inputs[0]), - reinterpret_cast(inputs[1]), mAlpha, reinterpret_cast(outputs[0]), M, N, K, - trtToCublasDtype(mType), trtToCublasDtype(mOutputType)); - cudaKernelFinished = tensorrt_llm::kernels::cuda_core_gemm::cudaCoreGemmDispatcher(params, stream); - } - - if (!cudaKernelFinished) - { - auto bestTactic = mPluginProfiler->getBestConfig(M, mGemmId); - runGemm(M, N, K, mTransA, mTransB, mPadLda, mPadLdb, mPadLdc, mType, mCublasWrapper, inputs[0], inputs[1], - mAlpha, outputs[0], bestTactic, workspace, stream); - } - - { - std::string const outputStr = "GEMM layer's output after GEMM with " + mnkStr; - TLLM_CHECK_DEBUG_WITH_INFO( - tensorrt_llm::runtime::utils::tensorHasInvalid(M, N + mPadLdc, mType, outputs[0], stream, outputStr) - == false, - "Found invalid number (NaN or Inf) in " + outputStr); - } - return 0; -} - -// IPluginV2Ext Methods -nvinfer1::DataType GemmPlugin::getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept -{ - TLLM_CHECK(index == 0); - return mType; -} - -// IPluginV2 Methods - -char const* GemmPlugin::getPluginType() const noexcept -{ - return GEMM_PLUGIN_NAME; -} - -char const* GemmPlugin::getPluginVersion() const noexcept -{ - return GEMM_PLUGIN_VERSION; -} - -int GemmPlugin::getNbOutputs() const noexcept -{ - return 1; -} - -int GemmPlugin::initialize() noexcept -{ - configGemm(); - return 0; -} - -void GemmPlugin::destroy() noexcept -{ - delete this; -} - -size_t GemmPlugin::getSerializationSize() const noexcept -{ - return sizeof(mTransA) + sizeof(mTransB) + sizeof(mPadLda) + sizeof(mPadLdb) + sizeof(mPadLdc) + sizeof(mType) - + sizeof(mDims) + sizeof(mUseFp8) + sizeof(mAlpha) + mPluginProfiler->getSerializationSize(mGemmId) - + sizeof(mOutputType); // selected tactics container size -} - -void GemmPlugin::serialize(void* buffer) const noexcept -{ - char *d = static_cast(buffer), *a = d; - write(d, mTransA); - write(d, mTransB); - write(d, mPadLda); - write(d, mPadLdb); - write(d, mPadLdc); - write(d, mType); - write(d, mUseFp8); - write(d, mAlpha); - write(d, mDims); - write(d, mOutputType); - mPluginProfiler->serialize(d, mGemmId); - - TLLM_CHECK(d == a + getSerializationSize()); -} - -void GemmPlugin::terminate() noexcept {} - -/////////////// - -GemmPluginCreator::GemmPluginCreator() -{ - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - mPluginAttributes.emplace_back(PluginField("transA", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("transB", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("padLda", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("padLdb", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("padLdc", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("type_id", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("use_fp8", nullptr, PluginFieldType::kINT32)); - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -char const* GemmPluginCreator::getPluginName() const noexcept -{ - return GEMM_PLUGIN_NAME; -} - -char const* GemmPluginCreator::getPluginVersion() const noexcept -{ - return GEMM_PLUGIN_VERSION; -} - -PluginFieldCollection const* GemmPluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -IPluginV2* GemmPluginCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept -{ - PluginField const* fields = fc->fields; - int transA{}; - int transB{}; - int padLda{}; - int padLdb{}; - int padLdc{}; - nvinfer1::DataType type{}; - int useFp8{}; - float alpha = 1.F; - // Read configurations from each fields - for (int i = 0; i < fc->nbFields; ++i) - { - char const* attrName = fields[i].name; - if (!strcmp(attrName, "transa")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - transA = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "transb")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - transB = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "pad_lda")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - padLda = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "pad_ldb")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - padLdb = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "pad_ldc")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - padLdc = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "type_id")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - type = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "use_fp8")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - useFp8 = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "alpha")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kFLOAT32); - alpha = static_cast(*(static_cast(fields[i].data))); - } - } - try - { - // GemmPluginCreator is unique and shared for an engine generation - // Create plugin profiler with shared tactics map - // FIXME enable tactic profiler - auto pluginProfiler = gemmPluginProfileManager.createGemmPluginProfiler(/* inference */ false, /* skip */ true); - auto* obj = new GemmPlugin(transA, transB, padLda, padLdb, padLdc, type, useFp8, alpha, pluginProfiler); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -IPluginV2* GemmPluginCreator::deserializePlugin(char const* name, void const* serialData, size_t serialLength) noexcept -{ - // This object will be deleted when the network is destroyed, which will - // call GemmPlugin::destroy() - try - { - // GemmPluginCreator is unique and shared for an engine generation - // Create plugin profiler with shared tactics map - // FIXME enable tactic profiler - auto pluginProfiler = gemmPluginProfileManager.createGemmPluginProfiler(/* inference */ true, /* skip */ true); - auto* obj = new GemmPlugin(serialData, serialLength, pluginProfiler); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} diff --git a/cpp/tensorrt_llm/plugins/gemmPlugin/gemmPlugin.h b/cpp/tensorrt_llm/plugins/gemmPlugin/gemmPlugin.h deleted file mode 100644 index 1ba553c23d4b..000000000000 --- a/cpp/tensorrt_llm/plugins/gemmPlugin/gemmPlugin.h +++ /dev/null @@ -1,169 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. - */ -#ifndef TRT_GEMM_PLUGIN_H -#define TRT_GEMM_PLUGIN_H - -#include "tensorrt_llm/common/cublasMMWrapper.h" -#include "tensorrt_llm/plugins/common/gemmPluginProfiler.h" -#include "tensorrt_llm/plugins/common/plugin.h" - -#include -#include - -namespace tensorrt_llm::plugins -{ - -using CublasGemmWrapper = tensorrt_llm::common::CublasMMWrapper; -using CublasGemmWrapperPtr = std::shared_ptr; - -class CublasLtGemmPluginProfiler - : public GemmPluginProfiler -{ -public: - using Config = cublasLtMatmulHeuristicResult_t; - - void setTranspose(bool transposeA, bool transposeB) - { - mTransA = transposeA; - mTransB = transposeB; - } - - void setPadLd(int padLda, int padLdb, int padLdc) - { - mPadLda = padLda; - mPadLdb = padLdb; - mPadLdc = padLdc; - } - - void setOutputType(nvinfer1::DataType type) - { - mOutputType = type; - } - -protected: - void runTactic(int m, int n, int k, Config const& tactic, char* workspace, cudaStream_t const& stream) override; - - void computeTmpSize(size_t maxM, size_t n, size_t k) override; - - bool checkTactic(int m, int n, int k, Config const& tactic) const override; - - std::vector getTactics(int m, int n, int k) const override; - -private: - bool mTransA; - bool mTransB; - int mPadLda; - int mPadLdb; - int mPadLdc; - nvinfer1::DataType mOutputType; - - static constexpr size_t ALIGNMENT = 256; -}; - -class GemmPlugin : public BasePlugin -{ -public: - using PluginProfilerPtr = std::shared_ptr; - - GemmPlugin() = delete; - - GemmPlugin(int transA, int transB, int padLda, int padLdb, int padLdc, nvinfer1::DataType type, bool useFp8, - float alpha, PluginProfilerPtr const& profiler); - - GemmPlugin(void const* data, size_t length, PluginProfilerPtr const& profiler); - - ~GemmPlugin() override = default; - - // IPluginV2DynamicExt Methods - nvinfer1::IPluginV2DynamicExt* clone() const noexcept override; - nvinfer1::DimsExprs getOutputDimensions(int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; - bool supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept override; - void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept override; - size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept override; - int enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - - // IPluginV2Ext Methods - nvinfer1::DataType getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept override; - - // IPluginV2 Methods - char const* getPluginType() const noexcept override; - char const* getPluginVersion() const noexcept override; - int getNbOutputs() const noexcept override; - int initialize() noexcept override; - void terminate() noexcept override; - size_t getSerializationSize() const noexcept override; - void serialize(void* buffer) const noexcept override; - void destroy() noexcept override; - -private: - void init(); - void configGemm(); - void setGemmConfig(); - -private: - const std::string mLayerName; - - int mTransA; - int mTransB; - int mPadLda; - int mPadLdb; - int mPadLdc; - int mArch; - nvinfer1::DataType mType; - nvinfer1::DataType mOutputType; - - static thread_local CublasGemmWrapperPtr mCublasWrapper; - - GemmDims mDims{}; - GemmIdCublas mGemmId{}; - bool mUseFp8{false}; - float mAlpha{1.f}; - - PluginProfilerPtr mPluginProfiler; -}; - -class GemmPluginCreator : public BaseCreator -{ -public: - GemmPluginCreator(); - - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; - - nvinfer1::IPluginV2* deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept override; - -private: - GemmPluginProfilerManager gemmPluginProfileManager; - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; -}; - -} // namespace tensorrt_llm::plugins - -#endif // TRT_GEMM_PLUGIN_H diff --git a/cpp/tensorrt_llm/plugins/gemmSwigluPlugin/CMakeLists.txt b/cpp/tensorrt_llm/plugins/gemmSwigluPlugin/CMakeLists.txt deleted file mode 100644 index 3b714a3928fb..000000000000 --- a/cpp/tensorrt_llm/plugins/gemmSwigluPlugin/CMakeLists.txt +++ /dev/null @@ -1,21 +0,0 @@ -# -# SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. -# -file(GLOB SRCS *.cpp *.cu) -set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS}) -set(PLUGIN_SOURCES - ${PLUGIN_SOURCES} - PARENT_SCOPE) diff --git a/cpp/tensorrt_llm/plugins/gemmSwigluPlugin/gemmSwigluPlugin.cpp b/cpp/tensorrt_llm/plugins/gemmSwigluPlugin/gemmSwigluPlugin.cpp deleted file mode 100644 index ed964ace695f..000000000000 --- a/cpp/tensorrt_llm/plugins/gemmSwigluPlugin/gemmSwigluPlugin.cpp +++ /dev/null @@ -1,446 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. - */ - -#include "gemmSwigluPlugin.h" -#include "cutlass_extensions/gemm_configs.h" - -#include -#include - -using namespace nvinfer1; -using namespace tensorrt_llm::common; -using namespace tensorrt_llm::kernels::cutlass_kernels; -using tensorrt_llm::plugins::GemmSwigluPluginCreator; -using tensorrt_llm::plugins::GemmSwigluPlugin; -using tensorrt_llm::plugins::GemmSwigluPluginProfiler; -using tensorrt_llm::plugins::read; -using tensorrt_llm::plugins::write; - -static char const* GEMM_SWIGLU_PLUGIN_VERSION{"1"}; -static char const* GEMM_SWIGLU_PLUGIN_NAME{"GemmSwiglu"}; -PluginFieldCollection GemmSwigluPluginCreator::mFC{}; -std::vector GemmSwigluPluginCreator::mPluginAttributes; - -size_t GemmSwigluPluginProfiler::getBytePerElement(nvinfer1::DataType type) -{ - size_t bpe; - if (type == nvinfer1::DataType::kHALF || type == nvinfer1::DataType::kBF16) - { - bpe = 2; - } - else if (type == nvinfer1::DataType::kINT8 || type == nvinfer1::DataType::kFP8) - { - bpe = 1; - } - else - { - TLLM_THROW("Not recognized/implemented"); - } - return bpe; -} - -void GemmSwigluPluginProfiler::setQuantMode(tensorrt_llm::common::QuantMode const& quantMode) -{ - mQuantMode = quantMode; -} - -void GemmSwigluPluginProfiler::runTactic( - int m, int n, int k, GemmSwigluPluginProfiler::Config const& tactic, char* workspace, cudaStream_t const& stream) -{ - size_t bpe = getBytePerElement(mType); - - // Workspace size required by gemm runner - // NB: this function will throw exception when selected tactic exceeds SMEM, which is then - // caught by gemmPluginProfiler and it will register this tactic as invalid - size_t wsSizeRunner = mRunner->getWorkspaceSize(m, n, k); - - // Workspace size required by profiling - size_t wsByteOffset = 0; - int8_t* wsBytePointer = reinterpret_cast(workspace); - void* aTmp = reinterpret_cast(nextWorkspacePtr(wsBytePointer, wsByteOffset, m * k * bpe)); - void* bTmp = reinterpret_cast(nextWorkspacePtr(wsBytePointer, wsByteOffset, n * k * bpe)); - void* cTmp = reinterpret_cast(nextWorkspacePtr(wsBytePointer, wsByteOffset, 1 * n * bpe)); - void* dTmp = reinterpret_cast(nextWorkspacePtr(wsBytePointer, wsByteOffset, m * (n / 2) * bpe)); - char* workspaceTmp = reinterpret_cast(nextWorkspacePtr(wsBytePointer, wsByteOffset, wsSizeRunner)); - - // Run profiling - mRunner->gemm( - dTmp, aTmp, bTmp, cTmp, mQuantMode, m, n, k, 1.0, 1.0, 1.0, tactic, workspaceTmp, wsSizeRunner, stream); -} - -int GemmSwigluPluginProfiler::getMaxProfileM() const -{ - return 32768; -} - -void GemmSwigluPluginProfiler::computeTmpSize(size_t maxM, size_t n, size_t k) -{ - std::vector workspaces = { - maxM * k * getBytePerElement(mType), // A - n * k * getBytePerElement(mType), // B - 1 * n * getBytePerElement(mType), // C_bias - maxM * (n / 2) * getBytePerElement(mType), // D - mRunner->getWorkspaceSize(maxM, n, k) // workspace - }; - size_t bytes = calculateTotalWorkspaceSize(workspaces.data(), workspaces.size()); - setTmpWorkspaceSizeInBytes(bytes); -} - -std::vector GemmSwigluPluginProfiler::getTactics(int m, int n, int k) const -{ - return mRunner->getConfigs(); -} - -GemmSwigluPlugin::GemmSwigluPlugin(QuantMode quantMode, nvinfer1::DataType type, bool hasBias, float scale_d0, - float scale_d1, float scale_output, GemmSwigluPlugin::PluginProfilerPtr const& pluginProfiler) - : mQuantMode(quantMode) - , mPluginProfiler(pluginProfiler) - , mHasBias(hasBias) - , mScaleD0(scale_d0) - , mScaleD1(scale_d1) - , mScaleOutput(scale_output) -{ - init(type); -} - -// Parameterized constructor -GemmSwigluPlugin::GemmSwigluPlugin( - void const* data, size_t length, GemmSwigluPlugin::PluginProfilerPtr const& pluginProfiler) - : mPluginProfiler(pluginProfiler) -{ - char const *d = reinterpret_cast(data), *a = d; - nvinfer1::DataType type; - unsigned int quantMode; - read(d, quantMode); - read(d, type); - read(d, mHasBias); - read(d, mScaleD0); - read(d, mScaleD1); - read(d, mScaleOutput); - read(d, mDims); - - mQuantMode = QuantMode(quantMode); - - init(type); - - mPluginProfiler->deserialize(d, mDims, mGemmId); - - TLLM_CHECK(d == a + length); -} - -void GemmSwigluPlugin::init(nvinfer1::DataType type) -{ - mType = type; - if (mType == nvinfer1::DataType::kFP8) - { - mGemmRunner = std::make_shared>(); - } - else - { - TLLM_THROW("Gemm Swiglu plugin only supports fp8 now"); - } - - mPluginProfiler->setQuantMode(mQuantMode); - - mGemmId = GemmIdCore(mDims.n, mDims.k, mType); -} - -// IPluginV2DynamicExt Methods -nvinfer1::IPluginV2DynamicExt* GemmSwigluPlugin::clone() const noexcept -{ - auto* plugin = new GemmSwigluPlugin(*this); - return plugin; -} - -nvinfer1::DimsExprs GemmSwigluPlugin::getOutputDimensions( - int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept -{ - try - { - TLLM_CHECK(nbInputs == 3); - TLLM_CHECK(outputIndex == 0); - int const nbDimsA = inputs[0].nbDims; - TLLM_CHECK(nbDimsA >= 2); - DimsExprs ret; - ret.nbDims = nbDimsA; - for (int ii = 0; ii < nbDimsA - 1; ++ii) - { - ret.d[ii] = inputs[0].d[ii]; - } - ret.d[nbDimsA - 1] = exprBuilder.constant(inputs[1].d[1]->getConstantValue() / 2); - return ret; - } - catch (std::exception const& e) - { - caughtError(e); - } - return DimsExprs{}; -} - -bool GemmSwigluPlugin::supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept -{ - switch (pos) - { - case 0: - // activation - return inOut[pos].type == mType && inOut[pos].format == TensorFormat::kLINEAR; - case 1: - // weights - return inOut[pos].type == mType && inOut[pos].format == TensorFormat::kLINEAR; - case 2: - // bias - return inOut[pos].type == mType && inOut[pos].format == TensorFormat::kLINEAR; - case 3: - // out - return inOut[pos].type == mType && inOut[pos].format == TensorFormat::kLINEAR; - default: - // Never should be here - TLLM_CHECK(false); - return false; - } -} - -void GemmSwigluPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept -{ - auto const minM = std::accumulate(in[0].min.d, in[0].min.d + in[0].min.nbDims - 1, 1, std::multiplies()); - auto const maxM = std::accumulate(in[0].max.d, in[0].max.d + in[0].max.nbDims - 1, 1, std::multiplies()); - - int const maxK = in[0].max.d[in[0].max.nbDims - 1]; - int const maxN = in[1].max.d[1]; - int const minK = in[0].min.d[in[0].min.nbDims - 1]; - int const minN = in[1].min.d[1]; - - TLLM_CHECK_WITH_INFO(minN == maxN, "Variable out channels is not allowed"); - TLLM_CHECK_WITH_INFO(minK == maxK, "Variable in channels is not allowed"); - - if (!mDims.isInitialized()) - { - mDims = {minM, maxM, maxN, maxK}; - } - mGemmId = {maxN, maxK, mType}; - - mWorkspaceMaxSize = mGemmRunner->getWorkspaceSize(maxM, maxN, maxK); -} - -size_t GemmSwigluPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - return mWorkspaceMaxSize; -} - -int GemmSwigluPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept -{ - // inputs - // mat1 [M(*), K] - // mat2 [K, N] - // bias [1, N] - // outputs - // mat [M(*), N / 2] - int m = 1; - for (int ii = 0; ii < inputDesc[0].dims.nbDims - 1; ++ii) - { - m *= inputDesc[0].dims.d[ii]; - } - int const n = inputDesc[1].dims.d[1]; - int const k = inputDesc[0].dims.d[inputDesc[0].dims.nbDims - 1]; - size_t const wsSize = mGemmRunner->getWorkspaceSize(m, n, k); - - auto const bestTactic = mPluginProfiler->getBestConfig(m, mGemmId); - TLLM_CHECK_WITH_INFO(bestTactic, "No valid GEMM tactic"); - mGemmRunner->gemm(outputs[0], inputs[0], inputs[1], inputs[2], mQuantMode, m, n, k, mScaleD0, mScaleD1, - mScaleOutput, *bestTactic, reinterpret_cast(workspace), wsSize, stream); - - return 0; -} - -// IPluginV2Ext Methods -nvinfer1::DataType GemmSwigluPlugin::getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept -{ - TLLM_CHECK(index == 0); - return mType; -} - -// IPluginV2 Methods - -char const* GemmSwigluPlugin::getPluginType() const noexcept -{ - return GEMM_SWIGLU_PLUGIN_NAME; -} - -char const* GemmSwigluPlugin::getPluginVersion() const noexcept -{ - return GEMM_SWIGLU_PLUGIN_VERSION; -} - -int GemmSwigluPlugin::getNbOutputs() const noexcept -{ - return 1; -} - -int GemmSwigluPlugin::initialize() noexcept -{ - configGemm(); // gemm profiler in action - return 0; -} - -void GemmSwigluPlugin::terminate() noexcept {} - -size_t GemmSwigluPlugin::getSerializationSize() const noexcept -{ - return sizeof(unsigned int) + // QuantMode - sizeof(nvinfer1::DataType) + // dtype - sizeof(bool) + // hasBias - sizeof(float) * 3 + // scales - sizeof(mDims) + // Dimensions - mPluginProfiler->getSerializationSize(mGemmId); // selected tactics container size -} - -void GemmSwigluPlugin::serialize(void* buffer) const noexcept -{ - char *d = static_cast(buffer), *a = d; - write(d, mQuantMode.value()); - write(d, mType); - write(d, mHasBias); - write(d, mScaleD0); - write(d, mScaleD1); - write(d, mScaleOutput); - write(d, mDims); - - mPluginProfiler->serialize(d, mGemmId); - TLLM_CHECK(d == a + getSerializationSize()); -} - -void GemmSwigluPlugin::destroy() noexcept -{ - // This gets called when the network containing plugin is destroyed - delete this; -} - -void GemmSwigluPlugin::configGemm() -{ - mPluginProfiler->profileTactics(mGemmRunner, mType, mDims, mGemmId); -} - -/////////////// - -GemmSwigluPluginCreator::GemmSwigluPluginCreator() -{ - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - mPluginAttributes.emplace_back(PluginField("type_id", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("has_bias", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("scale_d0", nullptr, PluginFieldType::kFLOAT32)); - mPluginAttributes.emplace_back(PluginField("scale_d1", nullptr, PluginFieldType::kFLOAT32)); - mPluginAttributes.emplace_back(PluginField("scale_output", nullptr, PluginFieldType::kFLOAT32)); - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -char const* GemmSwigluPluginCreator::getPluginName() const noexcept -{ - return GEMM_SWIGLU_PLUGIN_NAME; -} - -char const* GemmSwigluPluginCreator::getPluginVersion() const noexcept -{ - return GEMM_SWIGLU_PLUGIN_VERSION; -} - -PluginFieldCollection const* GemmSwigluPluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -IPluginV2* GemmSwigluPluginCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept -{ - PluginField const* fields = fc->fields; - TLLM_CHECK(fc->nbFields == 5); - nvinfer1::DataType type{}; - bool hasBias{}; - float scale_d0{}; - float scale_d1{}; - float scale_output{}; - // Read configurations from each fields - for (int i = 0; i < fc->nbFields; ++i) - { - char const* attrName = fields[i].name; - if (!strcmp(attrName, "type_id")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - type = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "has_bias")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT8); - hasBias = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "scale_d0")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kFLOAT32); - scale_d0 = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "scale_d1")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kFLOAT32); - scale_d1 = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "scale_output")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kFLOAT32); - scale_output = static_cast(*(static_cast(fields[i].data))); - } - } - try - { - // GemmSwigluPluginCreator is unique and shared for an engine generation - // Create plugin profiler with shared tactics map - auto pluginProfiler = mGemmPluginProfileManager.createGemmPluginProfiler(/* inference */ false); - QuantMode quantMode = QuantMode{}; - auto* obj = new GemmSwigluPlugin(quantMode, type, hasBias, scale_d0, scale_d1, scale_output, pluginProfiler); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -IPluginV2* GemmSwigluPluginCreator::deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept -{ - // This object will be deleted when the network is destroyed, which will - // call GemmSwigluPlugin::destroy() - try - { - // Create plugin profiler with private tactics map which is read from the serialized engine - auto pluginProfiler = mGemmPluginProfileManager.createGemmPluginProfiler(/* inference */ true); - auto* obj = new GemmSwigluPlugin(serialData, serialLength, pluginProfiler); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} diff --git a/cpp/tensorrt_llm/plugins/gemmSwigluPlugin/gemmSwigluPlugin.cu b/cpp/tensorrt_llm/plugins/gemmSwigluPlugin/gemmSwigluPlugin.cu deleted file mode 100644 index 339c432b1113..000000000000 --- a/cpp/tensorrt_llm/plugins/gemmSwigluPlugin/gemmSwigluPlugin.cu +++ /dev/null @@ -1,41 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2023 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. - */ - -#include "gemmSwigluPlugin.h" - -#include "cutlass/util/reference/device/tensor_fill.h" -#include "cutlass_extensions/gemm_configs.h" - -using namespace nvinfer1; -using namespace tensorrt_llm::common; -using namespace tensorrt_llm::kernels::cutlass_kernels; -using tensorrt_llm::plugins::GemmSwigluPluginCreator; -using tensorrt_llm::plugins::GemmSwigluPlugin; -using tensorrt_llm::plugins::GemmSwigluPluginProfiler; -using tensorrt_llm::plugins::read; -using tensorrt_llm::plugins::write; - -void GemmSwigluPluginProfiler::initTmpData(int m, int n, int k, char* workspace, size_t size, cudaStream_t stream) -{ - size_t bpe = getBytePerElement(mType); - - if (mType == nvinfer1::DataType::kFP8) - { - cutlass::reference::device::BlockFillRandomUniform(reinterpret_cast(workspace), - m * k + n * k + 1 * n, 42, cutlass::float_e4m3_t{128}, -cutlass::float_e4m3_t{128}, -1, 0, stream); - } -} diff --git a/cpp/tensorrt_llm/plugins/gemmSwigluPlugin/gemmSwigluPlugin.h b/cpp/tensorrt_llm/plugins/gemmSwigluPlugin/gemmSwigluPlugin.h deleted file mode 100644 index 766e59aad258..000000000000 --- a/cpp/tensorrt_llm/plugins/gemmSwigluPlugin/gemmSwigluPlugin.h +++ /dev/null @@ -1,150 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. - */ -#pragma once - -#include "tensorrt_llm/kernels/cutlass_kernels/fused_gated_gemm/fused_gated_gemm.h" -#include "tensorrt_llm/plugins/common/gemmPluginProfiler.h" -#include "tensorrt_llm/plugins/common/plugin.h" -#include -#include -#include -#include - -namespace tensorrt_llm::plugins -{ - -using GemmSwigluRunnerPtr - = std::shared_ptr; - -class GemmSwigluPluginProfiler : public GemmPluginProfiler - -{ -public: - using Config = tensorrt_llm::cutlass_extensions::CutlassGemmConfig; - - void setQuantMode(tensorrt_llm::common::QuantMode const& quantMode); - - virtual int getMaxProfileM() const override; - -protected: - void runTactic(int m, int n, int k, Config const& tactic, char* workspace, cudaStream_t const& stream) override; - - void computeTmpSize(size_t maxM, size_t n, size_t k) override; - - // TODO(anchengc) implement checkTactic - // bool checkTactic(int m, int n, int k, const Config& tactic) const override; - - std::vector getTactics(int m, int n, int k) const override; - - void initTmpData(int m, int n, int k, char* workspace, size_t size, cudaStream_t stream) override; - -private: - size_t getBytePerElement(nvinfer1::DataType type); - - tensorrt_llm::common::QuantMode mQuantMode; -}; - -class GemmSwigluPlugin : public BasePlugin -{ -public: - using PluginProfilerPtr = std::shared_ptr; - - GemmSwigluPlugin() = delete; - - GemmSwigluPlugin(tensorrt_llm::common::QuantMode quantMode, nvinfer1::DataType type, bool hasBias, float scale_d0, - float scale_d1, float scale_output, PluginProfilerPtr const& pluginProfiler); - - GemmSwigluPlugin(void const* data, size_t length, PluginProfilerPtr const& profiler); - - ~GemmSwigluPlugin() override = default; - - // IPluginV2DynamicExt Methods - nvinfer1::IPluginV2DynamicExt* clone() const noexcept override; - nvinfer1::DimsExprs getOutputDimensions(int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; - bool supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept override; - void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept override; - size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept override; - int enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - - // IPluginV2Ext Methods - nvinfer1::DataType getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept override; - - // IPluginV2 Methods - char const* getPluginType() const noexcept override; - char const* getPluginVersion() const noexcept override; - int getNbOutputs() const noexcept override; - int initialize() noexcept override; - void terminate() noexcept override; - size_t getSerializationSize() const noexcept override; - void serialize(void* buffer) const noexcept override; - void destroy() noexcept override; - -private: - void init(nvinfer1::DataType type); - - void configGemm(); - // void setGemmConfig(); - -private: - const std::string mLayerName; - - GemmSwigluRunnerPtr mGemmRunner; - tensorrt_llm::common::QuantMode mQuantMode; // not configurable yet - size_t mWorkspaceMaxSize; - - GemmDims mDims{}; - GemmIdCore mGemmId{}; - - PluginProfilerPtr mPluginProfiler; - - nvinfer1::DataType mType; - bool mHasBias; - float mScaleD0; - float mScaleD1; - float mScaleOutput; -}; - -class GemmSwigluPluginCreator : public BaseCreator -{ -public: - GemmSwigluPluginCreator(); - - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; - - nvinfer1::IPluginV2* deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept override; - -private: - GemmPluginProfilerManager mGemmPluginProfileManager; - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; -}; - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/gptAttentionCommon/CMakeLists.txt b/cpp/tensorrt_llm/plugins/gptAttentionCommon/CMakeLists.txt deleted file mode 100644 index 86876224fccd..000000000000 --- a/cpp/tensorrt_llm/plugins/gptAttentionCommon/CMakeLists.txt +++ /dev/null @@ -1,21 +0,0 @@ -# -# SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. -# -file(GLOB SRCS *.cpp) -set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS}) -set(PLUGIN_SOURCES - ${PLUGIN_SOURCES} - PARENT_SCOPE) diff --git a/cpp/tensorrt_llm/plugins/gptAttentionCommon/gptAttentionCommon.cpp b/cpp/tensorrt_llm/plugins/gptAttentionCommon/gptAttentionCommon.cpp deleted file mode 100644 index 717ab3083e5f..000000000000 --- a/cpp/tensorrt_llm/plugins/gptAttentionCommon/gptAttentionCommon.cpp +++ /dev/null @@ -1,380 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. - */ -#include "gptAttentionCommon.h" -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderXQARunner.h" -#include "tensorrt_llm/kernels/gptKernels.h" -#include -#include - -using namespace nvinfer1; -using namespace tensorrt_llm::kernels; -namespace tc = tensorrt_llm::common; -using tensorrt_llm::plugins::GPTAttentionPluginCreatorCommon; -using tensorrt_llm::plugins::GPTAttentionPluginCommon; - -GPTAttentionPluginCommon::GPTAttentionPluginCommon(int layer_idx, int num_heads, int vision_start, int vision_length, - int num_kv_heads, int num_kv_heads_origin, int head_size, int unidirectional, float q_scaling, - float attn_logit_softcapping_scale, tensorrt_llm::kernels::PositionEmbeddingType position_embedding_type, - int rotary_embedding_dim, // for RoPE. Use 0 for non-RoPE - float rotary_embedding_base, tensorrt_llm::kernels::RotaryScalingType rotary_embedding_scale_type, - float rotary_embedding_scale, float rotary_embedding_short_m_scale, float rotary_embedding_long_m_scale, - int rotary_embedding_max_positions, int rotary_embedding_original_max_positions, int tp_size, - int tp_rank, // for ALiBi - bool unfuse_qkv_gemm, // for AutoPP - bool use_logn_scaling, // for LognScaling - tensorrt_llm::kernels::ContextFMHAType context_fmha_type, int kv_cache_quant_mode, bool remove_input_padding, - tensorrt_llm::kernels::AttentionMaskType mask_type, tensorrt_llm::kernels::BlockSparseParams block_sparse_params, - bool paged_kv_cache, int tokens_per_block, nvinfer1::DataType type, int32_t max_context_length, - bool qkv_bias_enabled, bool cross_attention, int max_distance, bool pos_shift_enabled, bool dense_context_fmha, - bool use_paged_context_fmha, bool use_fp8_context_fmha, bool has_full_attention_mask, bool use_cache, - bool is_spec_decoding_enabled, bool spec_decoding_is_generation_length_variable, - int32_t spec_decoding_max_generation_length, bool is_mla_enabled, int q_lora_rank, int kv_lora_rank, - int qk_nope_head_dim, int qk_rope_head_dim, int v_head_dim, bool fuse_fp4_quant, bool skip_attn, int cp_size, - int cp_rank, std::set cp_group) - : mResource{DecoderXQARunner::getResourceGlobal()} -{ - mLayerIdx = layer_idx; - mNumHeads = num_heads; - mVisionStart = vision_start; - mVisionLength = vision_length; - mNumKVHeads = num_kv_heads; - mNumKVHeadsOrigin = num_kv_heads_origin; - mHeadSize = head_size; - mUnidirectional = unidirectional; - mQScaling = q_scaling; - mAttnLogitSoftcappingScale = attn_logit_softcapping_scale; - mRotaryEmbeddingDim = rotary_embedding_dim; - mRotaryEmbeddingBase = rotary_embedding_base; - mRotaryEmbeddingScaleType = rotary_embedding_scale_type; - mRotaryEmbeddingScale = rotary_embedding_scale; - mRotaryEmbeddingShortMscale = rotary_embedding_short_m_scale; - mRotaryEmbeddingLongMscale = rotary_embedding_long_m_scale; - mRotaryEmbeddingMaxPositions = rotary_embedding_max_positions; - mRotaryEmbeddingOriginalMaxPositions = rotary_embedding_original_max_positions; - mPositionEmbeddingType = position_embedding_type; - mEnableContextFMHA = context_fmha_type != ContextFMHAType::DISABLED; - mFMHAForceFP32Acc = type == nvinfer1::DataType::kBF16; - mMaskType = mask_type; - mBlockSparseParams = block_sparse_params; - mType = type; - mMultiBlockMode = true; - mEnableXQA = true; - mKVCacheQuantMode = tc::QuantMode(kv_cache_quant_mode); - mRemovePadding = remove_input_padding; - mPagedKVCache = paged_kv_cache; - mTokensPerBlock = tokens_per_block; - mTpSize = tp_size; - mTpRank = tp_rank; - mUnfuseQkvGemm = unfuse_qkv_gemm; - mUseLognScaling = use_logn_scaling; - mMaxContextLength = max_context_length; - mQKVBiasEnabled = qkv_bias_enabled; - mCrossAttention = cross_attention; - mMaxDistance = max_distance; - mPosShiftEnabled = pos_shift_enabled; - mDenseContextFMHA = dense_context_fmha; - mPagedContextFMHA = use_paged_context_fmha; - mFP8ContextFMHA = use_fp8_context_fmha; - mFP8AttenOutput = use_fp8_context_fmha; - mHasFullAttentionMask = has_full_attention_mask; - mUseKVCache = use_cache; - mIsSpecDecodingEnabled = is_spec_decoding_enabled; - mSpecDecodingIsGenerationLengthVariable = spec_decoding_is_generation_length_variable; - mSpecDecodingMaxGenerationLength = spec_decoding_max_generation_length; - mIsMLAEnabled = is_mla_enabled; - mMLAParams = {q_lora_rank, kv_lora_rank, qk_nope_head_dim, qk_rope_head_dim, v_head_dim}; - mCpSize = cp_size; - mCpRank = cp_rank; - mCpGroup = std::move(cp_group); - mFuseFp4Quant = fuse_fp4_quant; - mSkipAttn = skip_attn; -} - -// Parameterized constructor -GPTAttentionPluginCommon::GPTAttentionPluginCommon(void const* data, size_t length) - : mResource{DecoderXQARunner::getResourceGlobal()} -{ - char const *d = reinterpret_cast(data), *a = d; - unsigned int kvCacheQuantMode; - - read(d, mLayerIdx); - read(d, mNumHeads); - read(d, mVisionStart); - read(d, mVisionLength); - read(d, mNumKVHeads); - read(d, mNumKVHeadsOrigin); - read(d, mHeadSize); - read(d, mUnidirectional); - read(d, mQScaling); - read(d, mAttnLogitSoftcappingScale); - read(d, mPositionEmbeddingType); - read(d, mRotaryEmbeddingDim); - read(d, mRotaryEmbeddingBase); - read(d, mRotaryEmbeddingScaleType); - read(d, mRotaryEmbeddingScale); - read(d, mRotaryEmbeddingShortMscale); - read(d, mRotaryEmbeddingLongMscale); - read(d, mRotaryEmbeddingMaxPositions); - read(d, mRotaryEmbeddingOriginalMaxPositions); - read(d, mTpSize); - read(d, mTpRank); - read(d, mUnfuseQkvGemm); - read(d, mUseLognScaling); - read(d, mEnableContextFMHA); - read(d, mFMHAForceFP32Acc); - read(d, mMultiBlockMode); - read(d, mEnableXQA); - read(d, kvCacheQuantMode); - read(d, mRemovePadding); - read(d, mMaskType); - read(d, mBlockSparseParams); - read(d, mPagedKVCache); - read(d, mTokensPerBlock); - read(d, mType); - read(d, mMaxContextLength); - read(d, mQKVBiasEnabled); - read(d, mCrossAttention); - read(d, mMaxDistance); - read(d, mPosShiftEnabled); - read(d, mDenseContextFMHA); - read(d, mPagedContextFMHA); - read(d, mFP8ContextFMHA); - read(d, mFP8AttenOutput); - read(d, mHasFullAttentionMask); - read(d, mUseKVCache); - read(d, mIsSpecDecodingEnabled); - read(d, mUseSpecDecoding); - read(d, mSpecDecodingIsGenerationLengthVariable); - read(d, mSpecDecodingMaxGenerationLength); - read(d, mIsMLAEnabled); - read(d, mMLAParams); - read(d, mNbMultiBlockSemaphores); - read(d, mFuseFp4Quant); - read(d, mSkipAttn); - read(d, mCpSize); - read(d, mCpRank); - - mKVCacheQuantMode = tc::QuantMode(kvCacheQuantMode); - - uint32_t decoderXQARunnerResourceSerializedSize; - read(d, decoderXQARunnerResourceSerializedSize); - mResource->merge(DecoderXQARunnerResource(d, decoderXQARunnerResourceSerializedSize), /*initialize=*/true); - d += decoderXQARunnerResourceSerializedSize; - - mCpGroup.clear(); - int32_t groupItem = 0; - while (d != a + length) - { - read(d, groupItem); - mCpGroup.insert(groupItem); - } - TLLM_CHECK_WITH_INFO(d == a + length, - "Expected length (%d) != real length (%d). This is often " - "caused by using different TensorRT LLM version to build " - "engine and run engine.", - (int) length, (int) (d - a)); - TLLM_CHECK_WITH_INFO((smVersion() >= 80) || (mType != nvinfer1::DataType::kBF16), - "Unsupported data type, pre SM 80 GPUs do not support bfloat16"); -} - -int GPTAttentionPluginCommon::initialize() noexcept -{ - return AttentionOp::initialize(); -} - -void GPTAttentionPluginCommon::destroy() noexcept -{ - delete this; -} - -size_t GPTAttentionPluginCommon::getCommonSerializationSize() const noexcept -{ - return sizeof(mLayerIdx) + sizeof(mNumHeads) + +sizeof(mVisionStart) + sizeof(mVisionLength) + sizeof(mNumKVHeads) - + sizeof(mNumKVHeadsOrigin) + sizeof(mHeadSize) + sizeof(mUnidirectional) + sizeof(mQScaling) - + sizeof(mAttnLogitSoftcappingScale) + sizeof(mPositionEmbeddingType) + sizeof(mRotaryEmbeddingDim) - + sizeof(mRotaryEmbeddingBase) + sizeof(mRotaryEmbeddingScaleType) + sizeof(mRotaryEmbeddingScale) - + sizeof(mRotaryEmbeddingShortMscale) + sizeof(mRotaryEmbeddingLongMscale) - + sizeof(mRotaryEmbeddingMaxPositions) + sizeof(mRotaryEmbeddingOriginalMaxPositions) + sizeof(mTpSize) - + sizeof(mTpRank) + sizeof(mEnableContextFMHA) + sizeof(mFMHAForceFP32Acc) + sizeof(mMultiBlockMode) - + sizeof(mEnableXQA) + sizeof(unsigned int) // mKVCacheQuantMode - + sizeof(mRemovePadding) + sizeof(mMaskType) + sizeof(mBlockSparseParams) + sizeof(mPagedKVCache) - + sizeof(mTokensPerBlock) + sizeof(mType) + sizeof(mMaxContextLength) + sizeof(mQKVBiasEnabled) - + sizeof(mCrossAttention) + sizeof(mMaxDistance) + sizeof(mPosShiftEnabled) + sizeof(mDenseContextFMHA) - + sizeof(mPagedContextFMHA) + sizeof(mFP8ContextFMHA) + sizeof(mFP8AttenOutput) + sizeof(mHasFullAttentionMask) - + sizeof(mUseKVCache) + sizeof(mUnfuseQkvGemm) + sizeof(mUseLognScaling) + sizeof(mIsSpecDecodingEnabled) - + sizeof(mUseSpecDecoding) + sizeof(mSpecDecodingIsGenerationLengthVariable) - + sizeof(mSpecDecodingMaxGenerationLength) + sizeof(mNbMultiBlockSemaphores) + sizeof(mIsMLAEnabled) - + sizeof(mMLAParams) + sizeof(mFuseFp4Quant) + sizeof(mSkipAttn) - + sizeof(uint32_t) // size of DecoderXQARunnerResource buffer. - + sizeof(mCpSize) + sizeof(mCpRank) + sizeof(int32_t) * mCpGroup.size() + mResource->getSerializationSize(); -} - -void GPTAttentionPluginCommon::serializeCommon(void* buffer) const noexcept -{ - char *d = static_cast(buffer), *a = d; - write(d, mLayerIdx); - write(d, mNumHeads); - write(d, mVisionStart); - write(d, mVisionLength); - write(d, mNumKVHeads); - write(d, mNumKVHeadsOrigin); - write(d, mHeadSize); - write(d, mUnidirectional); - write(d, mQScaling); - write(d, mAttnLogitSoftcappingScale); - write(d, mPositionEmbeddingType); - write(d, mRotaryEmbeddingDim); - write(d, mRotaryEmbeddingBase); - write(d, mRotaryEmbeddingScaleType); - write(d, mRotaryEmbeddingScale); - write(d, mRotaryEmbeddingShortMscale); - write(d, mRotaryEmbeddingLongMscale); - write(d, mRotaryEmbeddingMaxPositions); - write(d, mRotaryEmbeddingOriginalMaxPositions); - write(d, mTpSize); - write(d, mTpRank); - write(d, mUnfuseQkvGemm); - write(d, mUseLognScaling); - write(d, mEnableContextFMHA); - write(d, mFMHAForceFP32Acc); - write(d, mMultiBlockMode); - write(d, mEnableXQA); - write(d, mKVCacheQuantMode.value()); - write(d, mRemovePadding); - write(d, mMaskType); - write(d, mBlockSparseParams); - write(d, mPagedKVCache); - write(d, mTokensPerBlock); - write(d, mType); - write(d, mMaxContextLength); - write(d, mQKVBiasEnabled); - write(d, mCrossAttention); - write(d, mMaxDistance); - write(d, mPosShiftEnabled); - write(d, mDenseContextFMHA); - write(d, mPagedContextFMHA); - write(d, mFP8ContextFMHA); - write(d, mFP8AttenOutput); - write(d, mHasFullAttentionMask); - write(d, mUseKVCache); - write(d, mIsSpecDecodingEnabled); - write(d, mUseSpecDecoding); - write(d, mSpecDecodingIsGenerationLengthVariable); - write(d, mSpecDecodingMaxGenerationLength); - write(d, mIsMLAEnabled); - write(d, mMLAParams); - write(d, mNbMultiBlockSemaphores); - write(d, mFuseFp4Quant); - write(d, mSkipAttn); - write(d, mCpSize); - write(d, mCpRank); - - // An uint32_t that specifies the size of the serialized buffer, followed by the actual content. - uint32_t decoderXQARunnerResourceSerializedSize = mResource->getSerializationSize(); - write(d, decoderXQARunnerResourceSerializedSize); - mResource->serialize(d, decoderXQARunnerResourceSerializedSize); - d += decoderXQARunnerResourceSerializedSize; - - for (auto it = mCpGroup.begin(); it != mCpGroup.end(); ++it) - { - write(d, *it); - } - TLLM_CHECK(d == a + getCommonSerializationSize()); -} - -void GPTAttentionPluginCommon::terminate() noexcept -{ - // Do nothing, destroy will always be called, so release the resources there. -} - -/////////////// - -GPTAttentionPluginCreatorCommon::GPTAttentionPluginCreatorCommon() -{ - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - mPluginAttributes.emplace_back(PluginField("layer_idx", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("num_heads", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("vision_start", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("vision_length", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("num_kv_heads", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("num_kv_heads_origin", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("layer_idx_in_cache_pool", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("head_size", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("unidirectional", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("q_scaling", nullptr, PluginFieldType::kFLOAT32)); - mPluginAttributes.emplace_back(PluginField("attn_logit_softcapping_scale", nullptr, PluginFieldType::kFLOAT32)); - mPluginAttributes.emplace_back(PluginField("position_embedding_type", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("rotary_embedding_dim", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("rotary_embedding_base", nullptr, PluginFieldType::kFLOAT32)); - mPluginAttributes.emplace_back(PluginField("rotary_embedding_scale_type", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("rotary_embedding_scale", nullptr, PluginFieldType::kFLOAT32)); - mPluginAttributes.emplace_back(PluginField("rotary_embedding_short_m_scale", nullptr, PluginFieldType::kFLOAT32)); - mPluginAttributes.emplace_back(PluginField("rotary_embedding_long_m_scale", nullptr, PluginFieldType::kFLOAT32)); - mPluginAttributes.emplace_back(PluginField("rotary_embedding_max_positions", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back( - PluginField("rotary_embedding_original_max_positions", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("tp_size", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("tp_rank", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("unfuse_qkv_gemm", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("use_logn_scaling", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("context_fmha_type", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("kv_cache_quant_mode", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("remove_input_padding", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("mask_type", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("block_sparse_block_size", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("block_sparse_homo_head_pattern", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("block_sparse_num_local_blocks", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("block_sparse_vertical_stride", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("paged_kv_cache", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("tokens_per_block", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("type_id", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("max_context_length", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("qkv_bias_enabled", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("do_cross_attention", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("max_distance", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("pos_shift_enabled", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("dense_context_fmha", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("use_paged_context_fmha", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("use_fp8_context_fmha", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("has_full_attention_mask", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("use_cache", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("is_spec_decoding_enabled", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back( - PluginField("spec_decoding_is_generation_length_variable", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back( - PluginField("spec_decoding_max_generation_length", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("is_mla_enabled", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("q_lora_rank", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("kv_lora_rank", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("qk_nope_head_dim", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("qk_rope_head_dim", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("v_head_dim", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("fuse_fp4_quant", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("skip_attn", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("cp_size", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("cp_rank", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("cp_group", nullptr, PluginFieldType::kINT32)); - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -PluginFieldCollection const* GPTAttentionPluginCreatorCommon::getFieldNames() noexcept -{ - return &mFC; -} diff --git a/cpp/tensorrt_llm/plugins/gptAttentionCommon/gptAttentionCommon.h b/cpp/tensorrt_llm/plugins/gptAttentionCommon/gptAttentionCommon.h deleted file mode 100644 index dd87d67aab9e..000000000000 --- a/cpp/tensorrt_llm/plugins/gptAttentionCommon/gptAttentionCommon.h +++ /dev/null @@ -1,112 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. - */ -#pragma once - -#include "tensorrt_llm/common/attentionOp.h" -#include "tensorrt_llm/common/cublasMMWrapper.h" -#include "tensorrt_llm/common/quantization.h" -#include "tensorrt_llm/kernels/gptKernels.h" -#include "tensorrt_llm/plugins/common/plugin.h" -#include -#include -#include -#include - -namespace tensorrt_llm::kernels -{ -class DecoderXQARunnerResource; -} - -namespace tensorrt_llm::plugins -{ - -class GPTAttentionPluginCommon : public BasePlugin, public tensorrt_llm::common::op::AttentionOp -{ -public: - GPTAttentionPluginCommon() = delete; - - GPTAttentionPluginCommon(int layer_idx, int num_heads, int vision_start, int vision_length, int num_kv_heads, - int num_kv_heads_origin, int head_size, int unidirectional, float q_scaling, float attn_logit_softcapping_scale, - tensorrt_llm::kernels::PositionEmbeddingType position_embedding_type, - int rotary_embedding_dim, // for RoPE. Use 0 for non-RoPE - float rotary_embedding_base, tensorrt_llm::kernels::RotaryScalingType rotary_embedding_scale_type, - float rotary_embedding_scale, float rotary_embedding_short_m_scale, float rotary_embedding_long_m_scale, - int rotary_embedding_max_positions, int rotary_embedding_original_max_positions, int tp_size, - int tp_rank, // for ALiBi - bool unfuse_qkv_gemm, // for AutoPP - bool use_logn_scaling, // for LognScaling - tensorrt_llm::kernels::ContextFMHAType context_fmha_type, int kv_cache_quant_mode, bool remove_input_padding, - tensorrt_llm::kernels::AttentionMaskType mask_type, - tensorrt_llm::kernels::BlockSparseParams block_sparse_params, bool paged_kv_cache, int tokens_per_block, - nvinfer1::DataType type, int32_t max_context_length, bool qkv_bias_enabled, bool cross_attention = false, - int max_distance = 0, bool pos_shift_enabled = false, bool dense_context_fmha = false, - bool use_paged_context_fmha = true, bool use_fp8_context_fmha = true, bool has_full_attention_mask = false, - bool use_cache = true, bool is_spec_decoding_enabled = false, - bool spec_decoding_is_generation_length_variable = false, int32_t spec_decoding_max_generation_length = 1, - bool is_mla_enabled = false, int q_lora_rank = 0, int kv_lora_rank = 0, int qk_nope_head_dim = 0, - int qk_rope_head_dim = 0, int v_head_dim = 0, bool fuse_fp4_quant = false, bool skip_attn = false, - int cp_size = 1, int cp_rank = 0, std::set cp_group = {}); - - GPTAttentionPluginCommon(void const* data, size_t length); - - ~GPTAttentionPluginCommon() override = default; - - template - int enqueueImpl(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream); - - //! This is called on every trt Engine creation - int initialize() noexcept override; - //! This is called on every trt Engine destroy - void terminate() noexcept override; - - //! This is called on every trt ExecutionContext creation by TRT - //! Note TRT does not call the initialize on cloned plugin, so clone internally should do initialization. - template - T* cloneImpl() const noexcept; - - //! This is called on evert trt Engine or ExecutionContext destroy. - //! None-cloned plugins will call terminate and then call destroy, while the cloned plugins will call destroy only - //! So plugin should put the resource release inside destroy. - void destroy() noexcept override; - - size_t getCommonSerializationSize() const noexcept; - void serializeCommon(void* buffer) const noexcept; - -protected: - std::string const mLayerName; - -private: - std::shared_ptr mResource; -}; - -class GPTAttentionPluginCreatorCommon : public BaseCreator -{ -public: - GPTAttentionPluginCreatorCommon(); - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - template - T* deserializePluginImpl(char const* name, void const* serialData, size_t serialLength) noexcept; - -protected: - std::vector mPluginAttributes; - nvinfer1::PluginFieldCollection mFC{}; -}; - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/gptAttentionCommon/gptAttentionCommonImpl.h b/cpp/tensorrt_llm/plugins/gptAttentionCommon/gptAttentionCommonImpl.h deleted file mode 100644 index 51462cee6f40..000000000000 --- a/cpp/tensorrt_llm/plugins/gptAttentionCommon/gptAttentionCommonImpl.h +++ /dev/null @@ -1,54 +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. - */ - -#pragma once - -#include "gptAttentionCommon.h" - -namespace tensorrt_llm::plugins -{ -template -T* GPTAttentionPluginCommon::cloneImpl() const noexcept -{ - static_assert(std::is_base_of_v); - auto* plugin = new T(static_cast(*this)); - plugin->setPluginNamespace(mNamespace.c_str()); - - // Cloned plugins should be in initialized state with correct resources ready to be enqueued. - plugin->initialize(); - return plugin; -} - -template -T* GPTAttentionPluginCreatorCommon::deserializePluginImpl( - char const* name, void const* serialData, size_t serialLength) noexcept -{ - // This object will be deleted when the network is destroyed, which will - // call GPTAttentionPluginCommon::destroy() - try - { - auto* obj = new T(serialData, serialLength); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/gptAttentionPlugin/CMakeLists.txt b/cpp/tensorrt_llm/plugins/gptAttentionPlugin/CMakeLists.txt deleted file mode 100644 index 86876224fccd..000000000000 --- a/cpp/tensorrt_llm/plugins/gptAttentionPlugin/CMakeLists.txt +++ /dev/null @@ -1,21 +0,0 @@ -# -# SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. -# -file(GLOB SRCS *.cpp) -set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS}) -set(PLUGIN_SOURCES - ${PLUGIN_SOURCES} - PARENT_SCOPE) diff --git a/cpp/tensorrt_llm/plugins/gptAttentionPlugin/gptAttentionPlugin.cpp b/cpp/tensorrt_llm/plugins/gptAttentionPlugin/gptAttentionPlugin.cpp deleted file mode 100644 index 6f8c41c94131..000000000000 --- a/cpp/tensorrt_llm/plugins/gptAttentionPlugin/gptAttentionPlugin.cpp +++ /dev/null @@ -1,1387 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. - */ - -#include "gptAttentionPlugin.h" - -#include "tensorrt_llm/batch_manager/contextProgress.h" -#include "tensorrt_llm/common/logger.h" -#include "tensorrt_llm/kernels/decoderMaskedMultiheadAttention.h" -#include "tensorrt_llm/kernels/gptKernels.h" -#include "tensorrt_llm/kernels/unfusedAttentionKernels.h" -#include "tensorrt_llm/plugins/common/checkMacrosPlugin.h" -#include "tensorrt_llm/plugins/common/plugin.h" -#include "tensorrt_llm/plugins/gptAttentionCommon/gptAttentionCommon.h" -#include "tensorrt_llm/plugins/gptAttentionCommon/gptAttentionCommonImpl.h" -#include "tensorrt_llm/runtime/common.h" -#include "tensorrt_llm/runtime/iBuffer.h" -#include "tensorrt_llm/runtime/utils/debugUtils.h" - -#include -#include -#include -#include -#include - -using namespace nvinfer1; -using namespace tensorrt_llm::kernels; -using namespace tensorrt_llm::common; -using tensorrt_llm::plugins::GPTAttentionPluginCreator; -using tensorrt_llm::plugins::GPTAttentionPlugin; - -static char const* GPT_ATTENTION_PLUGIN_VERSION{"1"}; -static char const* GPT_ATTENTION_PLUGIN_NAME{"GPTAttention"}; - -GPTAttentionPlugin::GPTAttentionPlugin(int layer_idx, int num_heads, int vision_start, int vision_length, - int num_kv_heads, int num_kv_heads_origin, int head_size, int unidirectional, float q_scaling, - float attn_logit_softcapping_scale, tensorrt_llm::kernels::PositionEmbeddingType position_embedding_type, - int rotary_embedding_dim, // for RoPE. 0 for non-RoPE - float rotary_embedding_base, tensorrt_llm::kernels::RotaryScalingType rotary_embedding_scale_type, - float rotary_embedding_scale, float rotary_embedding_short_m_scale, - float rotary_embedding_long_m_scale, // magnitude scaling factors for Phi-3 long RoPE - int rotary_embedding_max_positions, int rotary_embedding_original_max_positions, int tp_size, - int tp_rank, // for ALiBi - bool unfuse_qkv_gemm, // for AutoPP - bool use_logn_scaling, // for LognScaling - tensorrt_llm::kernels::ContextFMHAType context_fmha_type, int kv_cache_quant_mode, bool remove_input_padding, - tensorrt_llm::kernels::AttentionMaskType mask_type, tensorrt_llm::kernels::BlockSparseParams block_sparse_params, - bool paged_kv_cache, int tokens_per_block, nvinfer1::DataType type, int32_t max_context_length, - bool qkv_bias_enabled, bool cross_attention, int max_distance, bool pos_shift_enabled, bool dense_context_fmha, - bool use_paged_context_fmha, bool use_fp8_context_fmha, bool has_full_attention_mask, bool use_cache, - bool is_spec_decoding_enabled, bool spec_decoding_is_generation_length_variable, - int spec_decoding_max_generation_length, bool is_mla_enabled, int q_lora_rank, int kv_lora_rank, - int qk_nope_head_dim, int qk_rope_head_dim, int v_head_dim, bool fuse_fp4_quant, bool skip_attn, int cp_size, - int cp_rank, std::set cp_group) - : GPTAttentionPluginCommon(layer_idx, num_heads, vision_start, vision_length, num_kv_heads, num_kv_heads_origin, - head_size, unidirectional, q_scaling, attn_logit_softcapping_scale, position_embedding_type, - rotary_embedding_dim, rotary_embedding_base, rotary_embedding_scale_type, rotary_embedding_scale, - rotary_embedding_short_m_scale, rotary_embedding_long_m_scale, rotary_embedding_max_positions, - rotary_embedding_original_max_positions, tp_size, tp_rank, unfuse_qkv_gemm, use_logn_scaling, context_fmha_type, - kv_cache_quant_mode, remove_input_padding, mask_type, block_sparse_params, paged_kv_cache, tokens_per_block, - type, max_context_length, qkv_bias_enabled, cross_attention, max_distance, pos_shift_enabled, - dense_context_fmha, use_paged_context_fmha, use_fp8_context_fmha, has_full_attention_mask, use_cache, - is_spec_decoding_enabled, spec_decoding_is_generation_length_variable, spec_decoding_max_generation_length, - is_mla_enabled, q_lora_rank, kv_lora_rank, qk_nope_head_dim, qk_rope_head_dim, v_head_dim, fuse_fp4_quant, - skip_attn, cp_size, cp_rank, cp_group) -{ - TLLM_CHECK_WITH_INFO( - !is_mla_enabled, "GPTAttentionPlugin no longer supports MLA. Please use the PyTorch workflow instead."); - initEntryIdx(); -} - -GPTAttentionPlugin::GPTAttentionPlugin(void const* data, size_t length) - : GPTAttentionPluginCommon(data, length) -{ - initEntryIdx(); -} - -std::string GPTAttentionPlugin::toString(IdxEntry const& entry) const -{ -#define TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(name) \ - case IdxEntry::name: return #name - - switch (entry) - { - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(QKV_TENSOR); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(K_TENSOR); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(V_TENSOR); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(ATTENTION_MASK); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(ATTENTION_PACKED_MASK); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(SEQUENCE_LENGTH); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(HOST_PAST_KEY_VALUE_LENGTHS); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(HOST_MAX_ATTENTION_WINDOW); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(HOST_SINK_TOKEN_LENGTH); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(CONTEXT_LENGTHS); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(CACHE_INDIR); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(REQUEST_TYPES); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(KV_CACHE_BLOCK_OFFSETS); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(HOST_KV_CACHE_BLOCK_OFFSETS); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(HOST_KV_CACHE_POOL_POINTERS); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(HOST_KV_CACHE_POOL_MAPPING); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(PAST_KEY_VALUE); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(KV_CACHE_QUANTIZATION_SCALE); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(KV_CACHE_DEQUANTIZATION_SCALE); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(ATTENTION_OUTPUT_QUANTIZATION_SCALE); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(ATTENTION_OUTPUT_SF_SCALE); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(ROTARY_INV_FREQ); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(ROTARY_COS_SIN); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(ALIBI_SLOPES); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(RELATIVE_ATTENTION_BIAS); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(CROSS_KV); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(CROSS_KV_LENGTH); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(ENCODER_INPUT_LENGTH); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(HOST_CONTEXT_LENGTH); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(QKV_BIAS_TENSOR); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(SPEC_DECODING_GENERATION_LENGTHS); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(SPEC_DECODING_PACKED_MASK); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(SPEC_DECODING_POSITION_OFFSETS); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(SPEC_DECODING_USE); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(LONG_ROPE_ROTARY_INV_FREQ); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(LONG_ROPE_ROTARY_COS_SIN); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(MROPE_ROTARY_COS_SIN); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(MROPE_POSITION_DELTAS); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(HOST_RUNTIME_PERF_KNOBS); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(HOST_CONTEXT_PROGRESS); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(MLA_Q_B_PROJ_TENSOR); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(MLA_KV_B_PROJ_TENSOR); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(MLA_K_B_PROJ_TRANS_TENSOR); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(SKIP_ATTN); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(LOGN_SCALING); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(ENUM_SIZE); - } -#undef TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING - TLLM_LOG_TRACE(common::fmtstr("Missing string description for IdxEntry enum %lu.\n", static_cast(entry))); - return ""; -} - -bool GPTAttentionPlugin::isEntryUsed(IdxEntry const& entry) const -{ - switch (entry) - { - case IdxEntry::QKV_TENSOR: return true; - case IdxEntry::K_TENSOR: return mUnfuseQkvGemm; - case IdxEntry::V_TENSOR: return mUnfuseQkvGemm; - case IdxEntry::ATTENTION_MASK: return useFullCustomMask(); - case IdxEntry::ATTENTION_PACKED_MASK: return useCustomMask(); - case IdxEntry::SEQUENCE_LENGTH: return useKVCache(); - case IdxEntry::HOST_PAST_KEY_VALUE_LENGTHS: return useKVCache(); - case IdxEntry::HOST_MAX_ATTENTION_WINDOW: return true; - case IdxEntry::HOST_SINK_TOKEN_LENGTH: return true; - case IdxEntry::CONTEXT_LENGTHS: return true; - case IdxEntry::CACHE_INDIR: return useKVCache(); - case IdxEntry::REQUEST_TYPES: return true; - case IdxEntry::KV_CACHE_BLOCK_OFFSETS: return useKVCache() && mPagedKVCache; - case IdxEntry::HOST_KV_CACHE_BLOCK_OFFSETS: return useKVCache() && mPagedKVCache; - case IdxEntry::HOST_KV_CACHE_POOL_POINTERS: return useKVCache() && mPagedKVCache; - case IdxEntry::HOST_KV_CACHE_POOL_MAPPING: return useKVCache() && mPagedKVCache; - case IdxEntry::PAST_KEY_VALUE: return useKVCache() && !mPagedKVCache; - case IdxEntry::KV_CACHE_QUANTIZATION_SCALE: return useKVCache() && mKVCacheQuantMode.hasKvCacheQuant(); - case IdxEntry::KV_CACHE_DEQUANTIZATION_SCALE: return useKVCache() && mKVCacheQuantMode.hasKvCacheQuant(); - case IdxEntry::ATTENTION_OUTPUT_QUANTIZATION_SCALE: return mFP8ContextFMHA; - case IdxEntry::ATTENTION_OUTPUT_SF_SCALE: return mFuseFp4Quant; - case IdxEntry::ROTARY_INV_FREQ: return isRoPE(); - case IdxEntry::ROTARY_COS_SIN: return isRoPE(); - case IdxEntry::ALIBI_SLOPES: return isALiBi(); - case IdxEntry::RELATIVE_ATTENTION_BIAS: return isRelativePosition(); - case IdxEntry::CROSS_KV: return isCrossAttention(); - case IdxEntry::CROSS_KV_LENGTH: return isCrossAttention(); - case IdxEntry::LOGN_SCALING: return isLognScaling(); - case IdxEntry::ENCODER_INPUT_LENGTH: return isCrossAttention(); - case IdxEntry::HOST_CONTEXT_LENGTH: return mRemovePadding; - case IdxEntry::QKV_BIAS_TENSOR: return mQKVBiasEnabled; - case IdxEntry::SPEC_DECODING_GENERATION_LENGTHS: return mIsSpecDecodingEnabled; - case IdxEntry::SPEC_DECODING_PACKED_MASK: return mIsSpecDecodingEnabled; - case IdxEntry::SPEC_DECODING_POSITION_OFFSETS: return mIsSpecDecodingEnabled; - case IdxEntry::SPEC_DECODING_USE: return mIsSpecDecodingEnabled; - case IdxEntry::LONG_ROPE_ROTARY_INV_FREQ: return isLongRoPE(); - case IdxEntry::LONG_ROPE_ROTARY_COS_SIN: return isLongRoPE(); - case IdxEntry::MROPE_ROTARY_COS_SIN: return isMRoPE(); - case IdxEntry::MROPE_POSITION_DELTAS: return isMRoPE(); - case IdxEntry::HOST_RUNTIME_PERF_KNOBS: return true; - case IdxEntry::HOST_CONTEXT_PROGRESS: return true; - case IdxEntry::MLA_Q_B_PROJ_TENSOR: return mIsMLAEnabled; - case IdxEntry::MLA_KV_B_PROJ_TENSOR: return mIsMLAEnabled; - case IdxEntry::MLA_K_B_PROJ_TRANS_TENSOR: return mIsMLAEnabled; - case IdxEntry::SKIP_ATTN: return mSkipAttn; - default: return false; - } -} - -void GPTAttentionPlugin::initEntryIdx() -{ - mEntryIdx.resize(static_cast(IdxEntry::ENUM_SIZE)); - size_t entryIdx = 0; - for (size_t i = 0; i < static_cast(IdxEntry::ENUM_SIZE); i++) - { - mEntryIdx[i] = entryIdx; - entryIdx += isEntryUsed(static_cast(i)); - } -} - -GPTAttentionPlugin::IndexType GPTAttentionPlugin::getIdx(IdxEntry const& entry) const -{ - TLLM_CHECK_WITH_INFO( - isEntryUsed(entry), common::fmtstr("getIdx() should not be used with entry %s.\n", toString(entry).data())); - return mEntryIdx[static_cast(entry)]; -} - -// IPluginV2DynamicExt Methods -GPTAttentionPlugin* GPTAttentionPlugin::clone() const noexcept -{ - return dynamic_cast(this->cloneImpl()); -} - -static int getPackedTensorHiddenDimIndex(bool removePadding) -{ - return removePadding ? 1 : 2; -} - -// NOTE: generation input length might be larger than one in the spec decoding mode. -int GPTAttentionPlugin::getGenerationInputSequenceLength( - nvinfer1::PluginTensorDesc const* inputDesc, int32_t localNbSeq, int32_t localNbTokens) const -{ - if (mRemovePadding) - { - // Speculative decoding mode might need variable generation input sequence length. - if (mIsSpecDecodingEnabled && mUseSpecDecoding) - { - TLLM_CHECK_WITH_INFO(mCpSize <= 1, "Context Parallel does not support speculative decoding mode for now"); - // SPEC_DECODING_POSITION_OFFSETS: [batch_size, max_generation_input_length]. - return inputDesc[getIdx(IdxEntry::SPEC_DECODING_POSITION_OFFSETS)].dims.d[1]; - } - else - { - if (mCpSize > 1) - { - // Given that localNbTokens == (beamSize * localNbSeq + mCpSize - 1) / mCpSize, but when mCpSize - 1 > - // localNbSeq, there are multiple choices for beamSize. Assume beamSize == 1 here. - TLLM_CHECK_WITH_INFO(localNbTokens == (localNbSeq + mCpSize - 1) / mCpSize, - "Context Parallel does not support beamSize > 1 for non-speculative decoding mode, " - "localNbTokens=%d, localNbSeq=%d", - localNbTokens, localNbSeq); - return 1; - } - // [num_tokens, local_hidden_size] where num_tokens = batch_size * generation_input_length - TLLM_CHECK_WITH_INFO(localNbTokens % localNbSeq == 0, - "seq_len should be same for all generation requests, localNbTokens=%d, localNbSeq=%d", localNbTokens, - localNbSeq); - return localNbTokens / localNbSeq; - } - } - else - { - // We don't have IFB without mRemovePadding, so just take it out from inputDesc - // [batch_size, seq_len, local_hidden_size] - return inputDesc[getIdx(IdxEntry::QKV_TENSOR)].dims.d[1]; - } -} - -// outputs -// output_tensor [batch_size, seq_len, local_hidden_size] or [num_tokens, local_hidden_size] -// present_key_value_pool (optional if mPagedKVCache is false) [batch_size, 2, local_num_kv_heads, max_seq_len, -// head_size] -nvinfer1::DimsExprs GPTAttentionPlugin::getOutputDimensions( - int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept -{ - if (mFuseFp4Quant) - { - TLLM_CHECK(outputIndex == 0 || outputIndex == 1 || (!mPagedKVCache && useKVCache() && outputIndex == 2)); - // Compute the output dimension for FP4 quantized tensor. Consistent with QuantizeToFP4Plugin. - if (outputIndex == 0) - { - auto ret = inputs[getIdx(IdxEntry::QKV_TENSOR)]; - return ret; - } - // Compute the output dimension for output scaling factor tensor. Consistent with QuantizeToFP4Plugin. - if (outputIndex == 1) - { - auto ret = inputs[getIdx(IdxEntry::QKV_TENSOR)]; - // Sequence dimension or token dimension. - // Pad to multiple of 128. - auto dimM = exprBuilder.operation(DimensionOperation::kCEIL_DIV, - *ret.d[getPackedTensorHiddenDimIndex(mRemovePadding) - 1], *exprBuilder.constant(128)); - ret.d[getPackedTensorHiddenDimIndex(mRemovePadding) - 1] - = exprBuilder.operation(DimensionOperation::kPROD, *dimM, *exprBuilder.constant(128)); - // Hidden size dimension. - // Div (rounding up) by 16 since 16 elements share one SF and SF padded to k%4==0. - ret.d[getPackedTensorHiddenDimIndex(mRemovePadding)] = exprBuilder.operation(DimensionOperation::kCEIL_DIV, - *ret.d[getPackedTensorHiddenDimIndex(mRemovePadding)], *exprBuilder.constant(16)); - return ret; - } - } - else - { - TLLM_CHECK(outputIndex == 0 || (!mPagedKVCache && useKVCache() && outputIndex == 1)); - if (outputIndex == 0) - { - auto ret = inputs[getIdx(IdxEntry::QKV_TENSOR)]; - // In MLA, the output dim is v_head_dim - auto const head_size = mHeadSize; - ret.d[getPackedTensorHiddenDimIndex(mRemovePadding)] = exprBuilder.operation( - DimensionOperation::kPROD, *exprBuilder.constant(head_size), *exprBuilder.constant(mNumHeads)); - return ret; - } - } - return inputs[getIdx(IdxEntry::PAST_KEY_VALUE)]; -} - -bool GPTAttentionPlugin::supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept -{ - bool result = false; - int posCaseLine = -1; - if (pos == getIdx(IdxEntry::CONTEXT_LENGTHS) || pos == getIdx(IdxEntry::REQUEST_TYPES) - || pos == getIdx(IdxEntry::HOST_MAX_ATTENTION_WINDOW) || pos == getIdx(IdxEntry::HOST_SINK_TOKEN_LENGTH) - || (isEntryUsed(IdxEntry::SPEC_DECODING_PACKED_MASK) && pos == getIdx(IdxEntry::SPEC_DECODING_PACKED_MASK)) - || (isEntryUsed(IdxEntry::SPEC_DECODING_POSITION_OFFSETS) - && pos == getIdx(IdxEntry::SPEC_DECODING_POSITION_OFFSETS)) - || (isEntryUsed(IdxEntry::SPEC_DECODING_GENERATION_LENGTHS) - && pos == getIdx(IdxEntry::SPEC_DECODING_GENERATION_LENGTHS)) - || (isEntryUsed(IdxEntry::SPEC_DECODING_USE) && pos == getIdx(IdxEntry::SPEC_DECODING_USE))) - { - posCaseLine = __LINE__; - result = inOut[pos].type == nvinfer1::DataType::kINT32; - } - else if (isMRoPE() && (pos == getIdx(IdxEntry::MROPE_ROTARY_COS_SIN))) - { - return inOut[pos].type == nvinfer1::DataType::kFLOAT; - } - else if (isMRoPE() && (pos == getIdx(IdxEntry::MROPE_POSITION_DELTAS))) - { - return inOut[pos].type == nvinfer1::DataType::kINT32; - } - else if (pos == getIdx(IdxEntry::HOST_RUNTIME_PERF_KNOBS) || pos == getIdx(IdxEntry::HOST_CONTEXT_PROGRESS)) - { - posCaseLine = __LINE__; - result = inOut[pos].type == nvinfer1::DataType::kINT64; - } - else if (useKVCache() - && (pos == getIdx(IdxEntry::SEQUENCE_LENGTH) || pos == getIdx(IdxEntry::HOST_PAST_KEY_VALUE_LENGTHS) - || pos == getIdx(IdxEntry::CACHE_INDIR))) - { - posCaseLine = __LINE__; - result = inOut[pos].type == nvinfer1::DataType::kINT32; - } - else if (isRoPE() && (pos == getIdx(IdxEntry::ROTARY_INV_FREQ) || pos == getIdx(IdxEntry::ROTARY_COS_SIN))) - { - posCaseLine = __LINE__; - result = inOut[pos].type == nvinfer1::DataType::kFLOAT; - } - else if (isLongRoPE() - && (pos == getIdx(IdxEntry::LONG_ROPE_ROTARY_INV_FREQ) || pos == getIdx(IdxEntry::LONG_ROPE_ROTARY_COS_SIN))) - { - posCaseLine = __LINE__; - result = inOut[pos].type == nvinfer1::DataType::kFLOAT; - } - else if (useKVCache() && mKVCacheQuantMode.hasKvCacheQuant() - && (pos == getIdx(IdxEntry::KV_CACHE_DEQUANTIZATION_SCALE) - || pos == getIdx(IdxEntry::KV_CACHE_QUANTIZATION_SCALE))) - { - // kv_scale for mType->int8/fp8 and int8/fp8->mType conversion - posCaseLine = __LINE__; - result = inOut[pos].type == nvinfer1::DataType::kFLOAT && inOut[pos].format == TensorFormat::kLINEAR; - } - else if (mFP8ContextFMHA && pos == getIdx(IdxEntry::ATTENTION_OUTPUT_QUANTIZATION_SCALE)) - { - posCaseLine = __LINE__; - result = inOut[pos].type == nvinfer1::DataType::kFLOAT && inOut[pos].format == TensorFormat::kLINEAR; - } - else if (mFuseFp4Quant && pos == getIdx(IdxEntry::ATTENTION_OUTPUT_SF_SCALE)) - { - posCaseLine = __LINE__; - result = inOut[pos].type == nvinfer1::DataType::kFLOAT && inOut[pos].format == TensorFormat::kLINEAR; - } - else if (useFullCustomMask() && pos == getIdx(IdxEntry::ATTENTION_MASK)) - { - posCaseLine = __LINE__; - result = inOut[pos].type == nvinfer1::DataType::kBOOL && inOut[pos].format == TensorFormat::kLINEAR; - } - else if (useCustomMask() && pos == getIdx(IdxEntry::ATTENTION_PACKED_MASK)) - { - posCaseLine = __LINE__; - result = inOut[pos].type == nvinfer1::DataType::kINT32 && inOut[pos].format == TensorFormat::kLINEAR; - } - else if (useKVCache() && mPagedKVCache - && (pos == getIdx(IdxEntry::KV_CACHE_BLOCK_OFFSETS) || pos == getIdx(IdxEntry::HOST_KV_CACHE_BLOCK_OFFSETS))) - { - // kv cache block offsets - posCaseLine = __LINE__; - result = inOut[pos].type == nvinfer1::DataType::kINT32 && inOut[pos].format == TensorFormat::kLINEAR; - } - else if (useKVCache() && mPagedKVCache && (pos == getIdx(IdxEntry::HOST_KV_CACHE_POOL_POINTERS))) - { - // kv cache pool pointers - posCaseLine = __LINE__; - result = inOut[pos].type == nvinfer1::DataType::kINT64 && inOut[pos].format == TensorFormat::kLINEAR; - } - else if (useKVCache() && mPagedKVCache && (pos == getIdx(IdxEntry::HOST_KV_CACHE_POOL_MAPPING))) - { - // kv cache pool mapping - posCaseLine = __LINE__; - result = inOut[pos].type == nvinfer1::DataType::kINT32 && inOut[pos].format == TensorFormat::kLINEAR; - } - else if (useKVCache() && mKVCacheQuantMode.hasInt8KvCache() - && (!mPagedKVCache && (pos == getIdx(IdxEntry::PAST_KEY_VALUE) || pos == nbInputs + 1))) - { - // If use Int8 K/V cache we require I/O KV values to int8 - posCaseLine = __LINE__; - result = (inOut[pos].type == nvinfer1::DataType::kINT8) && (inOut[pos].format == TensorFormat::kLINEAR); - } - else if (useKVCache() && mKVCacheQuantMode.hasFp8KvCache() - && (!mPagedKVCache && (pos == getIdx(IdxEntry::PAST_KEY_VALUE) || pos == nbInputs + 1))) - { - // If use FP8 K/V cache we require I/O KV values to FP8 - posCaseLine = __LINE__; - result = (inOut[pos].type == nvinfer1::DataType::kFP8) && (inOut[pos].format == TensorFormat::kLINEAR); - } - else if (mRemovePadding && (pos == getIdx(IdxEntry::HOST_CONTEXT_LENGTH))) - { - posCaseLine = __LINE__; - result = inOut[pos].type == nvinfer1::DataType::kINT32 && inOut[pos].format == TensorFormat::kLINEAR; - } - else if (mCrossAttention - && (pos == getIdx(IdxEntry::CROSS_KV_LENGTH) || pos == getIdx(IdxEntry::ENCODER_INPUT_LENGTH))) - { - posCaseLine = __LINE__; - result = inOut[pos].type == nvinfer1::DataType::kINT32; - } - else if (isLognScaling() && pos == getIdx(IdxEntry::LOGN_SCALING)) - { - return inOut[pos].type == nvinfer1::DataType::kFLOAT; - } - else if (pos == nbInputs && mFuseFp4Quant) - { - // Set dtype for output FP4 quantized tensor. - posCaseLine = __LINE__; - result = (inOut[pos].type == nvinfer1::DataType::kFP4) && (inOut[pos].format == TensorFormat::kLINEAR); - } - else if (pos == nbInputs + 1 && mFuseFp4Quant) - { - // Set dtype for output scaling factor tensor. Use kINT32 as storage type (same as QuantizeToFP4Plugin). - posCaseLine = __LINE__; - result = (inOut[pos].type == nvinfer1::DataType::kFP8) && (inOut[pos].format == TensorFormat::kLINEAR); - } - else if (pos == nbInputs && mFP8ContextFMHA) - { - // Output tensor now supports fp8 data type. - posCaseLine = __LINE__; - result = (inOut[pos].type == nvinfer1::DataType::kFP8) && (inOut[pos].format == TensorFormat::kLINEAR); - } - else if (mSkipAttn && pos == getIdx(IdxEntry::SKIP_ATTN)) - { - posCaseLine = __LINE__; - result = inOut[pos].type == nvinfer1::DataType::kBOOL && inOut[pos].format == TensorFormat::kLINEAR; - } - else - { - posCaseLine = __LINE__; - result = (inOut[pos].type == mType) && (inOut[pos].format == TensorFormat::kLINEAR); - } - TLLM_LOG_DEBUG( - "%s: pos: %d, result: %d, posCaseLine: %d", __PRETTY_FUNCTION__, pos, static_cast(result), posCaseLine); - return result; -} - -template -void GPTAttentionPlugin::configurePluginImpl(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept -{ - TLLM_CHECK(mHeadSize > 0); - - int beamWidth = -1; - if (!isCrossAttention() && useKVCache()) - { - // desc_val == -1 means beam_width is not static, we should look at min/max/opt. - // - // In prepareEnqueueGeneration, we'll prepare for all cases where beam_width doesn't exceed max. - // TODO: pass min AND max to prepareEnqueueGeneration instead of max only. - int desc_val = in[getIdx(IdxEntry::CACHE_INDIR)].desc.dims.d[1]; - int max_val = in[getIdx(IdxEntry::CACHE_INDIR)].max.d[1]; - beamWidth = desc_val == -1 ? max_val : desc_val; - } - else - { - beamWidth = 1; - } - TLLM_CHECK(beamWidth != -1); - - // Commonly, cyclic_attention_window_size, and max_attention_window_size will be the same - // unless each layer has different attention window sizes. - // the kv_cache capacity. - int max_encoder_context_len = isCrossAttention() ? in[getIdx(IdxEntry::CROSS_KV_LENGTH)].desc.dims.d[0] : 0; - int const max_attention_window_size = isCrossAttention() - ? max_encoder_context_len - : (useKVCache() ? in[getIdx(IdxEntry::CACHE_INDIR)].desc.dims.d[2] : 0); - int const cyclic_attention_window_size = max_attention_window_size; - - int const num_requests = 256; - int const sink_token_length = 0; - - EnqueueGenerationParams enqueueParams; - enqueueParams.max_attention_window_size = max_attention_window_size; - enqueueParams.cyclic_attention_window_size = cyclic_attention_window_size; - enqueueParams.max_cyclic_attention_window_size = cyclic_attention_window_size; - enqueueParams.sink_token_length = sink_token_length; - enqueueParams.beam_width = beamWidth; - enqueueParams.num_requests = num_requests; - - prepareEnqueueGeneration(enqueueParams); - - // Always reserve SemaphoreArray (for multi-block mode) as MMHA may enable multi-block mode when shared memory is - // not enough. - auto const& ctxLenTensor = in[getIdx(IdxEntry::CONTEXT_LENGTHS)]; - TLLM_CHECK_DEBUG(ctxLenTensor.max.nbDims == 1); - int32_t const max_batch_beam = in[getIdx(IdxEntry::CONTEXT_LENGTHS)].max.d[0]; - reserveSemaphoreArray(mNumHeads * max_batch_beam); -} - -template -void GPTAttentionPlugin::configurePluginDispatchKVCacheType(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept -{ - if (mPagedKVCache) - { - configurePluginImpl(in, nbInputs, out, nbOutputs); - } - else - { - configurePluginImpl(in, nbInputs, out, nbOutputs); - } -} - -void GPTAttentionPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept -{ - if (mType == nvinfer1::DataType::kHALF) - { - configurePluginDispatchKVCacheType(in, nbInputs, out, nbOutputs); - } - else if (mType == nvinfer1::DataType::kFLOAT) - { - configurePluginDispatchKVCacheType(in, nbInputs, out, nbOutputs); - } -#ifdef ENABLE_BF16 - else if (mType == nvinfer1::DataType::kBF16) - { - configurePluginDispatchKVCacheType<__nv_bfloat16>(in, nbInputs, out, nbOutputs); - } -#endif -} - -size_t GPTAttentionPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - int const max_context_length = mMaxContextLength; - int const cross_kv_length = isCrossAttention() ? inputs[getIdx(IdxEntry::CROSS_KV_LENGTH)].dims.d[0] : 0; - int const max_num_seq = inputs[getIdx(IdxEntry::CONTEXT_LENGTHS)].dims.d[0]; - auto const type = inputs[getIdx(IdxEntry::QKV_TENSOR)].type; - int const max_kv_cache_length - = isCrossAttention() ? cross_kv_length : (useKVCache() ? inputs[getIdx(IdxEntry::CACHE_INDIR)].dims.d[2] : 0); - int const max_num_tokens - = mRemovePadding ? inputs[getIdx(IdxEntry::QKV_TENSOR)].dims.d[0] : max_num_seq * max_context_length; - int const max_blocks_per_sequence - = (useKVCache() && mPagedKVCache) ? inputs[getIdx(IdxEntry::KV_CACHE_BLOCK_OFFSETS)].dims.d[3] : 0; - - size_t const context_workspace_size - = getWorkspaceSizeForContext(type, max_num_seq, max_context_length, cross_kv_length, max_num_tokens); - - size_t const generation_workspace_size = getWorkspaceSizeForGeneration( - type, max_num_seq, max_kv_cache_length, max_num_tokens, max_blocks_per_sequence); - - size_t attention_input_workspace_size = 0; - - if (mUnfuseQkvGemm) - { - int const local_hidden_units_q - = inputs[getIdx(IdxEntry::QKV_TENSOR)].dims.d[getPackedTensorHiddenDimIndex(mRemovePadding)]; - int const local_hidden_units_kv - = inputs[getIdx(IdxEntry::K_TENSOR)].dims.d[getPackedTensorHiddenDimIndex(mRemovePadding)]; - size_t const size = tensorrt_llm::runtime::BufferDataType(type).getSize(); - size_t const attention_input_size = size * max_num_tokens * (local_hidden_units_q + 2 * local_hidden_units_kv); - size_t workspaces[1]; - workspaces[0] = attention_input_size; - attention_input_workspace_size = tensorrt_llm::common::calculateTotalWorkspaceSize(workspaces, 1); - } - - return std::max(context_workspace_size, generation_workspace_size) + attention_input_workspace_size; -} - -static size_t getStride(nvinfer1::Dims const& dims, int n) -{ - TLLM_CHECK(n >= 0 && n < dims.nbDims); - return std::accumulate(dims.d + n + 1, dims.d + dims.nbDims, 1, std::multiplies{}); -} - -template -int GPTAttentionPlugin::enqueueImpl(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) -{ - TLLM_LOG_TRACE("Attention plugin start at layer %d", mLayerIdx); - - using runtime::RequestType; - - int32_t const nbSeq = inputDesc[getIdx(IdxEntry::CONTEXT_LENGTHS)].dims.d[0]; - RequestType const* reqTypes = static_cast(inputs[getIdx(IdxEntry::REQUEST_TYPES)]); - - int32_t nbContextRequests = 0; - int32_t contextTokenIdxEnd = 0; - int32_t contextTokenIdxEndForCp = 0; - // count context requests - for (int32_t seqIdx = 0; seqIdx < nbSeq; seqIdx++) - { - if (reqTypes[seqIdx] != RequestType::kCONTEXT) - { - break; - } - ++nbContextRequests; - contextTokenIdxEnd += mRemovePadding - ? static_cast(inputs[getIdx(IdxEntry::HOST_CONTEXT_LENGTH)])[seqIdx] - : inputDesc[getIdx(IdxEntry::QKV_TENSOR)].dims.d[1]; - contextTokenIdxEndForCp += mRemovePadding - ? (static_cast(inputs[getIdx(IdxEntry::HOST_CONTEXT_LENGTH)])[seqIdx] + mCpSize - 1) - / mCpSize - : (inputDesc[getIdx(IdxEntry::QKV_TENSOR)].dims.d[1] + mCpSize - 1) / mCpSize; - } - - for (int32_t seqIdx = nbContextRequests; seqIdx < nbSeq; seqIdx++) - { - TLLM_CHECK(reqTypes[seqIdx] == RequestType::kGENERATION); - } - - // mixed requests require mRemovePadding and mPagedKVCache - if (nbContextRequests != 0 && nbContextRequests != nbSeq) - { - TLLM_CHECK(mRemovePadding && mPagedKVCache); - } - - if (nbContextRequests > 0) - { - auto seqIdxBeg = 0; - auto tokenIdxBeg = 0; - auto localNbTokens = contextTokenIdxEnd; - enqueueSome(seqIdxBeg, nbContextRequests, tokenIdxBeg, localNbTokens, - inputDesc, outputDesc, inputs, outputs, workspace, stream); - } - - if (auto nbGenerationSeq = nbSeq - nbContextRequests; nbGenerationSeq > 0) - { - auto seqIdxBeg = nbContextRequests; - auto tokenIdxBeg = mCpSize > 1 ? contextTokenIdxEndForCp : contextTokenIdxEnd; - // if mRemovePadding is true, we may have IFB, and need to remove context tokens. - // if mRemovePadding is false, it is only generation requests, so just multiply batch_beam and seq_len (May not - // 1 for Parallel Decoding) - auto localNbTokens = mRemovePadding - ? inputDesc[getIdx(IdxEntry::QKV_TENSOR)].dims.d[0] - tokenIdxBeg - : inputDesc[getIdx(IdxEntry::QKV_TENSOR)].dims.d[0] * inputDesc[getIdx(IdxEntry::QKV_TENSOR)].dims.d[1]; - enqueueSome(seqIdxBeg, nbGenerationSeq, tokenIdxBeg, localNbTokens, inputDesc, - outputDesc, inputs, outputs, workspace, stream); - } - - sync_check_cuda_error(stream); - TLLM_LOG_TRACE("Attention plugin stop at layer %d", mLayerIdx); - - return 0; -} - -template -int GPTAttentionPlugin::enqueueSome(int32_t seqIdxBeg, int32_t localNbSeq, int32_t tokenIdxBeg, int32_t localNbTokens, - nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) -{ - // relative_attention_bias [head_num, max_seq_len, max_seq_len] (optional in relative position) - // or [head_num, num_buckets] (optional in implicit relative attention) - // cross_kv [batch_size, seq_len, 2 * local_hidden_size] or [num_tokens, 2 * local_hidden_size] - // when enable remove_input_padding (optional in cross attention mode) - // cross_kv_length [int] max encoder input context length (optional in cross attention mode) - // encoder_input_lengths [batch_size] raw sequence lengths (optional in cross attention mode) - - using runtime::RequestType; - - auto const* const reqTypeInBatchPtr - = static_cast(inputs[getIdx(IdxEntry::REQUEST_TYPES)]) + seqIdxBeg; - bool const is_context = (reqTypeInBatchPtr[0] == RequestType::kCONTEXT); - - T const* attention_input = static_cast(inputs[getIdx(IdxEntry::QKV_TENSOR)]) - + inputDesc[getIdx(IdxEntry::QKV_TENSOR)].dims.d[getPackedTensorHiddenDimIndex(mRemovePadding)] - * size_t(tokenIdxBeg); - - bool changeSpecDecodingMode = false; - if (mIsSpecDecodingEnabled) - { - bool useSpecDecoding - = static_cast(reinterpret_cast(inputs[getIdx(IdxEntry::SPEC_DECODING_USE)])[0]); - changeSpecDecodingMode = mUseSpecDecoding != useSpecDecoding; - mUseSpecDecoding = useSpecDecoding; - } - - [[maybe_unused]] MlaParams mla_params; - - T const* qkv_bias = nullptr; - if (mQKVBiasEnabled) - { - qkv_bias = reinterpret_cast(inputs[getIdx(IdxEntry::QKV_BIAS_TENSOR)]); - } - - // Note we still need context length during generation for MMHA optimization. - int32_t const max_context_q_len = [&]() - { - if (!mRemovePadding) - { - return static_cast(inputDesc[getIdx(IdxEntry::QKV_TENSOR)].dims.d[1]); - } - auto const host_context_lengths - = static_cast(inputs[getIdx(IdxEntry::HOST_CONTEXT_LENGTH)]) + seqIdxBeg; - return *std::max_element(host_context_lengths, host_context_lengths + localNbSeq); - }(); - - // Rotary inv_freq, cos_sin cache to avoid re-computing. - float const* rotary_inv_freq = nullptr; - float2 const* rotary_cos_sin = nullptr; - - bool const useLongRoPECache = isLongRoPE() && max_context_q_len > mRotaryEmbeddingOriginalMaxPositions; - if (isRoPE()) - { - auto inputName = useLongRoPECache ? IdxEntry::LONG_ROPE_ROTARY_INV_FREQ : IdxEntry::ROTARY_INV_FREQ; - rotary_inv_freq = reinterpret_cast(inputs[getIdx(inputName)]); - } - if (isRoPE()) - { - auto inputName = useLongRoPECache ? IdxEntry::LONG_ROPE_ROTARY_COS_SIN : IdxEntry::ROTARY_COS_SIN; - rotary_cos_sin = reinterpret_cast(inputs[getIdx(inputName)]); - } - - auto const mrope_rotary_cos_sin - = isMRoPE() ? reinterpret_cast(inputs[getIdx(IdxEntry::MROPE_ROTARY_COS_SIN)]) : nullptr; - - auto const mrope_position_deltas - = isMRoPE() ? reinterpret_cast(inputs[getIdx(IdxEntry::MROPE_POSITION_DELTAS)]) : nullptr; - - if (mUnfuseQkvGemm) - { - int const max_seqlen = inputDesc[getIdx(IdxEntry::QKV_TENSOR)].dims.d[mRemovePadding ? 0 : 1]; - int const batch_size = mRemovePadding ? 1 : inputDesc[getIdx(IdxEntry::QKV_TENSOR)].dims.d[0]; - - T const* attention_input_q = static_cast(inputs[getIdx(IdxEntry::QKV_TENSOR)]); - T const* attention_input_k = static_cast(inputs[getIdx(IdxEntry::K_TENSOR)]); - T const* attention_input_v = static_cast(inputs[getIdx(IdxEntry::V_TENSOR)]); - size_t const hidden_units_q - = inputDesc[getIdx(IdxEntry::QKV_TENSOR)].dims.d[getPackedTensorHiddenDimIndex(mRemovePadding)]; - size_t const hidden_units_kv - = inputDesc[getIdx(IdxEntry::K_TENSOR)].dims.d[getPackedTensorHiddenDimIndex(mRemovePadding)]; - size_t const hidden_units = hidden_units_q + 2 * hidden_units_kv; - size_t const size_qkv = sizeof(T) * hidden_units; - size_t const size_q = sizeof(T) * hidden_units_q; - size_t const size_kv = sizeof(T) * hidden_units_kv; - size_t const total_size = size_qkv * batch_size * max_seqlen; - int8_t* workspace_byte_ptr = reinterpret_cast(workspace); - size_t offset = 0; - T* attention_input_qkv = reinterpret_cast(nextWorkspacePtr(workspace_byte_ptr, offset, total_size)); - workspace = reinterpret_cast(workspace_byte_ptr + offset); - - cudaMemcpy2DAsync(attention_input_qkv, size_qkv, attention_input_q, size_q, size_q, batch_size * max_seqlen, - cudaMemcpyDeviceToDevice, stream); - cudaMemcpy2DAsync(attention_input_qkv + hidden_units_q, size_qkv, attention_input_k, size_kv, size_kv, - batch_size * max_seqlen, cudaMemcpyDeviceToDevice, stream); - cudaMemcpy2DAsync(attention_input_qkv + hidden_units_q + hidden_units_kv, size_qkv, attention_input_v, size_kv, - size_kv, batch_size * max_seqlen, cudaMemcpyDeviceToDevice, stream); - - attention_input = attention_input_qkv + hidden_units * tokenIdxBeg; - } - - int const* context_q_lengths = reinterpret_cast(inputs[getIdx(IdxEntry::CONTEXT_LENGTHS)]) + seqIdxBeg; - int const* sequence_kv_length = useKVCache() - ? static_cast(inputs[getIdx(IdxEntry::SEQUENCE_LENGTH)]) + seqIdxBeg - : context_q_lengths; - - int max_encoder_context_len = isCrossAttention() ? inputDesc[getIdx(IdxEntry::CROSS_KV_LENGTH)].dims.d[0] : 0; - // for enc-dec model, since decoder_input_ids could be longer than 1, - // such model has an encoder context (for cross attn) and an decoder context (for self attn) - // clarify 3 lens: - // -- max_context_q_len: len of decoder input. No "max" concept, it's what it is given. - // Also called (decoder_)input_seq_length, normally 1 for encoder-decoder start token - // -- max_seq_len: max allowed len of decoder output, i.e. final results - // -- max_encoder_context_len: len of encoder input (in cross attn). Also called encoder_input_seq_length - - int const beamWidth - = isCrossAttention() ? 1 : (useKVCache() ? inputDesc[getIdx(IdxEntry::CACHE_INDIR)].dims.d[1] : 1); - - // Commonly, cyclic_attention_window_size, and max_attention_window_size will be the same - // unless each layer has different attention window sizes. - // the kv_cache capacity. - int const max_attention_window_size = isCrossAttention() - ? max_encoder_context_len - : (useKVCache() ? inputDesc[getIdx(IdxEntry::CACHE_INDIR)].dims.d[2] : 0); - // The cyclic_attention_window_size will determine the cyclic kv cache position of new tokens. - // Note that this cyclic_attention_window_size might be smaller than the actual kv cache capactity. - int const* cyclic_attention_window_sizes - = reinterpret_cast(inputs[getIdx(IdxEntry::HOST_MAX_ATTENTION_WINDOW)]); - int const cyclic_attention_window_size - = isCrossAttention() ? max_encoder_context_len : cyclic_attention_window_sizes[mLayerIdx]; - int const sink_token_length = reinterpret_cast(inputs[getIdx(IdxEntry::HOST_SINK_TOKEN_LENGTH)])[0]; - int const num_attn_layer = inputDesc[getIdx(IdxEntry::HOST_MAX_ATTENTION_WINDOW)].dims.d[0]; - int const max_cyclic_attention_window_size = isCrossAttention() - ? max_encoder_context_len - : *std::max_element(cyclic_attention_window_sizes, cyclic_attention_window_sizes + num_attn_layer); - bool const can_use_one_more_block = beamWidth > 1; - - float const* kv_scale_orig_quant = nullptr; - float const* kv_scale_quant_orig = nullptr; - if (useKVCache() && mKVCacheQuantMode.hasKvCacheQuant()) - { - assert(inputDesc[getIdx(IdxEntry::KV_CACHE_QUANTIZATION_SCALE)].type == nvinfer1::DataType::kFLOAT); - assert(inputDesc[getIdx(IdxEntry::KV_CACHE_DEQUANTIZATION_SCALE)].type == nvinfer1::DataType::kFLOAT); - kv_scale_orig_quant = reinterpret_cast(inputs[getIdx(IdxEntry::KV_CACHE_QUANTIZATION_SCALE)]); - kv_scale_quant_orig = reinterpret_cast(inputs[getIdx(IdxEntry::KV_CACHE_DEQUANTIZATION_SCALE)]); - } - - float const* attention_output_orig_quant = nullptr; - if (mFP8ContextFMHA) - { - assert(inputDesc[getIdx(IdxEntry::ATTENTION_OUTPUT_QUANTIZATION_SCALE)].type == nvinfer1::DataType::kFLOAT); - attention_output_orig_quant - = reinterpret_cast(inputs[getIdx(IdxEntry::ATTENTION_OUTPUT_QUANTIZATION_SCALE)]); - } - float const* attention_output_sf_scale = nullptr; - if (mFuseFp4Quant) - { - assert(inputDesc[getIdx(IdxEntry::ATTENTION_OUTPUT_SF_SCALE)].type == nvinfer1::DataType::kFLOAT); - attention_output_sf_scale = reinterpret_cast(inputs[getIdx(IdxEntry::ATTENTION_OUTPUT_SF_SCALE)]); - } - uint32_t const* attention_packed_mask = nullptr; - if (useCustomMask()) - { - assert(inputDesc[getIdx(IdxEntry::ATTENTION_PACKED_MASK)].type == nvinfer1::DataType::kINT32); - attention_packed_mask = reinterpret_cast(inputs[getIdx(IdxEntry::ATTENTION_PACKED_MASK)]); - } - bool const* attention_mask = nullptr; - int attention_mask_stride = 0; - if (useFullCustomMask()) - { - attention_mask_stride = static_cast(inputDesc[getIdx(IdxEntry::ATTENTION_MASK)].dims.d[1]); - attention_mask = reinterpret_cast(inputs[getIdx(IdxEntry::ATTENTION_MASK)]) - + attention_mask_stride * static_cast(tokenIdxBeg); - } - - int max_blocks_per_sequence = 0; - kernels::KVBlockArray::DataType* block_offsets = nullptr; - void* host_primary_pool_pointer = nullptr; - void* host_secondary_pool_pointer = nullptr; - if (useKVCache() && mPagedKVCache) - { - auto const& kvCacheBlockOffsetsShape = inputDesc[getIdx(IdxEntry::KV_CACHE_BLOCK_OFFSETS)].dims; - max_blocks_per_sequence = kvCacheBlockOffsetsShape.d[kvCacheBlockOffsetsShape.nbDims - 1]; - - std::int32_t const* host_pool_mapping - = static_cast(inputs[getIdx(IdxEntry::HOST_KV_CACHE_POOL_MAPPING)]); - - int32_t const layerToPool = host_pool_mapping[mLayerIdx * 2]; - int32_t const layerIdxInCachePool = host_pool_mapping[mLayerIdx * 2 + 1]; - TLLM_LOG_TRACE("Layer%d: LayerCachePoolLocator{.indexOfPool=%d, .layerIdxInCachePool=%d}", mLayerIdx, - layerToPool, layerIdxInCachePool); - auto const seqStride = getStride(kvCacheBlockOffsetsShape, 1); - auto const poolStride = getStride(kvCacheBlockOffsetsShape, 0); - auto const seqOffset = seqIdxBeg * seqStride; - auto const poolOffset = layerToPool * poolStride; - - block_offsets - = reinterpret_cast(inputs[getIdx(IdxEntry::KV_CACHE_BLOCK_OFFSETS)]) - + poolOffset + seqOffset; - - auto const* const typed_host_pool_pointers - = static_cast(inputs[getIdx(IdxEntry::HOST_KV_CACHE_POOL_POINTERS)]); - - auto const cacheElemSize = (mKVCacheQuantMode.hasKvCacheQuant() ? 1 : sizeof(T)); - - auto const kv_cache_head_num = (mNumKVHeads + mCpSize - 1) / mCpSize; - auto const blockSize = mTokensPerBlock * kv_cache_head_num * mHeadSize; - auto const bytesPerBlock = blockSize * cacheElemSize; - auto const layerOffset = layerIdxInCachePool * 2 * bytesPerBlock; - - host_primary_pool_pointer = reinterpret_cast(typed_host_pool_pointers[layerToPool * 2] + layerOffset); - host_secondary_pool_pointer - = reinterpret_cast(typed_host_pool_pointers[layerToPool * 2 + 1] + layerOffset); - } - - // The index of kv cache tensor in outputs. If fuse FP4 quant, an additional scaling factor output is added before - // the kv cache tensor. - int const kvCacheIdxInOutputs = mFuseFp4Quant ? 2 : 1; - // The number of elements per storage type. For FP4 output, storage type is uint8_t. - int const numEltsPerStorageType = mFuseFp4Quant ? 2 : 1; - - AttentionOutT* context_buf_ = static_cast(outputs[0]) - + outputDesc[0].dims.d[getPackedTensorHiddenDimIndex(mRemovePadding)] * tokenIdxBeg / numEltsPerStorageType; - - __nv_fp8_e4m3* context_buf_sf_ = nullptr; - if (mFuseFp4Quant) - { - // The output address for FP4 scaling factor. - context_buf_sf_ = static_cast<__nv_fp8_e4m3*>(outputs[1]); - } - - void* key_value_cache = nullptr; - if (useKVCache() && !mPagedKVCache) - { - auto const cacheElemSize = (mKVCacheQuantMode.hasKvCacheQuant() ? 1 : sizeof(T)); - key_value_cache = static_cast(outputs[kvCacheIdxInOutputs]) - + cacheElemSize * getStride(outputDesc[kvCacheIdxInOutputs].dims, 0) * seqIdxBeg; - void const* past_key_value_cache = inputs[getIdx(IdxEntry::PAST_KEY_VALUE)]; - if (past_key_value_cache != outputs[kvCacheIdxInOutputs]) - { - auto shape = outputDesc[kvCacheIdxInOutputs].dims; - auto const size - = cacheElemSize * std::accumulate(shape.d, shape.d + shape.nbDims, 1, std::multiplies{}); - cudaMemcpyAsync(outputs[kvCacheIdxInOutputs], past_key_value_cache, size, cudaMemcpyDeviceToDevice, stream); - } - } - - T const* alibi_slopes = isALiBi() ? static_cast(inputs[getIdx(IdxEntry::ALIBI_SLOPES)]) : nullptr; - - int const* spec_decoding_packed_mask = nullptr; - int const* spec_decoding_position_offsets = nullptr; - int const* spec_decoding_generation_lengths = nullptr; - int num_decoding_draft_tokens = 0; - if (mIsSpecDecodingEnabled && mUseSpecDecoding) - { - // Second dimension of spec_decoding_position_offsets is num_decoding_draft_tokens + 1. - // [batch_size, num_decoding_draft_tokens + 1] - num_decoding_draft_tokens = inputDesc[getIdx(IdxEntry::SPEC_DECODING_POSITION_OFFSETS)].dims.d[1] - 1; - if (num_decoding_draft_tokens > 0) - { - // spec_decoding_* tensors are not filled for context requests. Hence, always strting from 0th index - int32_t constexpr genSeqIdx = 0; - spec_decoding_packed_mask = static_cast(inputs[getIdx(IdxEntry::SPEC_DECODING_PACKED_MASK)]) - + genSeqIdx * getStride(inputDesc[getIdx(IdxEntry::SPEC_DECODING_PACKED_MASK)].dims, 0); - // Packed as [num_tokens, packed_mask_size] - // Use seqIdxBeg * (num_decoding_draft_tokens + 1) here as only generation tokens have the packed_mask - // buffer. - // TODO: support variable sequence length based on generationTokenIdxBeg. - spec_decoding_packed_mask = static_cast(inputs[getIdx(IdxEntry::SPEC_DECODING_PACKED_MASK)]) - + genSeqIdx * (num_decoding_draft_tokens + 1) - * getStride(inputDesc[getIdx(IdxEntry::SPEC_DECODING_PACKED_MASK)].dims, 0); - spec_decoding_position_offsets - = static_cast(inputs[getIdx(IdxEntry::SPEC_DECODING_POSITION_OFFSETS)]) - + genSeqIdx * getStride(inputDesc[getIdx(IdxEntry::SPEC_DECODING_POSITION_OFFSETS)].dims, 0); - spec_decoding_generation_lengths - = static_cast(inputs[getIdx(IdxEntry::SPEC_DECODING_GENERATION_LENGTHS)]) + genSeqIdx; - } - } - - int32_t const* host_past_kv_len_list = useKVCache() - ? static_cast(inputs[getIdx(IdxEntry::HOST_PAST_KEY_VALUE_LENGTHS)]) + seqIdxBeg - : nullptr; - int32_t const max_context_kv_len = useKVCache() - ? *std::max_element(host_past_kv_len_list, host_past_kv_len_list + localNbSeq) - : max_context_q_len; - - int const* host_context_lengths - = mRemovePadding ? reinterpret_cast(inputs[getIdx(IdxEntry::HOST_CONTEXT_LENGTH)]) : nullptr; - - int64_t const* runtime_perf_knobs = static_cast(inputs[getIdx(IdxEntry::HOST_RUNTIME_PERF_KNOBS)]); - - EnqueueParams common_enqueue_params; - common_enqueue_params.attention_input = attention_input; - common_enqueue_params.qkv_bias = qkv_bias; - common_enqueue_params.attention_mask = attention_mask; - common_enqueue_params.rotary_inv_freq = rotary_inv_freq; - common_enqueue_params.rotary_cos_sin = rotary_cos_sin; - common_enqueue_params.max_attention_window_size = max_attention_window_size; - common_enqueue_params.cyclic_attention_window_size = cyclic_attention_window_size; - common_enqueue_params.max_cyclic_attention_window_size = max_cyclic_attention_window_size; - common_enqueue_params.can_use_one_more_block = can_use_one_more_block; - common_enqueue_params.sink_token_length = sink_token_length; - common_enqueue_params.kv_scale_orig_quant = kv_scale_orig_quant; - common_enqueue_params.kv_scale_quant_orig = kv_scale_quant_orig; - common_enqueue_params.attention_output_orig_quant = attention_output_orig_quant; - common_enqueue_params.attention_output_sf_scale = attention_output_sf_scale; - common_enqueue_params.alibi_slopes = alibi_slopes; - common_enqueue_params.context_buf = context_buf_; - common_enqueue_params.context_buf_sf = context_buf_sf_; - common_enqueue_params.key_value_cache = key_value_cache; - common_enqueue_params.block_offsets = block_offsets; - common_enqueue_params.host_primary_pool_pointer = host_primary_pool_pointer; - common_enqueue_params.host_secondary_pool_pointer = host_secondary_pool_pointer; - common_enqueue_params.num_tokens = localNbTokens; - common_enqueue_params.max_blocks_per_sequence = max_blocks_per_sequence; - common_enqueue_params.sequence_lengths = sequence_kv_length; - common_enqueue_params.context_lengths = context_q_lengths; - common_enqueue_params.host_context_lengths = host_context_lengths; - common_enqueue_params.workspace = workspace; - common_enqueue_params.runtime_perf_knobs = runtime_perf_knobs; - - if (isRelativePosition()) - { - common_enqueue_params.relative_attention_bias - = static_cast(inputs[getIdx(IdxEntry::RELATIVE_ATTENTION_BIAS)]); - common_enqueue_params.relative_attention_bias_stride - = inputDesc[getIdx(IdxEntry::RELATIVE_ATTENTION_BIAS)].dims.d[1]; // max_seq_len or num_buckets - } - if (isLognScaling()) - { - common_enqueue_params.logn_scaling_ptr = static_cast(inputs[getIdx(IdxEntry::LOGN_SCALING)]); - } - if (isCrossAttention()) - { - common_enqueue_params.encoder_input_lengths - = reinterpret_cast(inputs[getIdx(IdxEntry::ENCODER_INPUT_LENGTH)]) + seqIdxBeg; - } - - if (is_context) // context stage - { - int const batch_size = localNbSeq; - int const request_batch_size = batch_size; - // num of total tokens (without paddings when remove paddings). - int num_encoder_tokens = 0; - if (isCrossAttention()) - { - if (!mRemovePadding) - { - num_encoder_tokens = request_batch_size * max_encoder_context_len; - } - else - { - num_encoder_tokens = inputDesc[getIdx(IdxEntry::CROSS_KV)].dims.d[0]; - } - } - - common_enqueue_params.input_seq_length = max_context_q_len; - common_enqueue_params.max_past_kv_length = max_context_kv_len; - EnqueueContextParams enqueue_params{common_enqueue_params}; - enqueue_params.attention_packed_mask = attention_packed_mask; - enqueue_params.batch_size = batch_size; - enqueue_params.mrope_rotary_cos_sin = mrope_rotary_cos_sin; - enqueue_params.total_kv_len = enqueue_params.num_tokens; - - if (isCrossAttention()) - { - enqueue_params.cross_kv = static_cast(inputs[getIdx(IdxEntry::CROSS_KV)]); - enqueue_params.cross_kv_length = max_encoder_context_len; - enqueue_params.num_encoder_tokens = num_encoder_tokens; - } - - enqueueContext(enqueue_params, stream); - - { - std::string const afterContexStr = "ctx attention at layer " + std::to_string(mLayerIdx); - TLLM_LOG_TRACE("GPTAttentionPlugin - %s", afterContexStr.c_str()); - - auto progress = static_cast( - inputs[getIdx(IdxEntry::HOST_CONTEXT_PROGRESS)])[0]; - if (progress != nullptr) - { - progress->recordEvent(mLayerIdx, stream); - } - - if (!mFuseFp4Quant) - { - TLLM_CHECK_DEBUG_WITH_INFO( - tensorrt_llm::runtime::utils::tensorHasInvalid(localNbTokens, - outputDesc[0].dims.d[getPackedTensorHiddenDimIndex(mRemovePadding)], - mFP8ContextFMHA ? nvinfer1::DataType::kFP8 : mType, context_buf_, stream, afterContexStr) - == false, - "Found invalid number (NaN or Inf) in " + afterContexStr); - } - } - } - else // generation stage; max_context_q_len == input_seq_len == 1 - { - TLLM_CHECK_WITH_INFO(useKVCache(), "KV-cache-less is only supported for context"); - int batch_beam = localNbSeq; - TLLM_CHECK(batch_beam % beamWidth == 0); - int32_t const num_requests = batch_beam / beamWidth; - - int const* cache_indir - = beamWidth == 1 ? nullptr : reinterpret_cast(inputs[getIdx(IdxEntry::CACHE_INDIR)]); - - // Medusa: the max input sequence length if variable sequence length is needed. - int const input_seq_length = getGenerationInputSequenceLength(inputDesc, localNbSeq, localNbTokens); - int const max_past_kv_length = isCrossAttention() ? max_encoder_context_len : max_context_kv_len; - auto qkvDims = inputDesc[getIdx(IdxEntry::QKV_TENSOR)].dims; - TLLM_CHECK_WITH_INFO(input_seq_length == 1 || (mIsSpecDecodingEnabled && mUseSpecDecoding), - "Only speculative decoding mode supports input length > 1 in the generation phase, input_seq_length=%d, " - "mIsSpecDecodingEnabled=%s, nDims=%d, (" FMT_DIM ", " FMT_DIM ", " FMT_DIM ")", - input_seq_length, mIsSpecDecodingEnabled ? "true" : "false", qkvDims.nbDims, qkvDims.d[0], qkvDims.d[1], - qkvDims.d[2]); - TLLM_CHECK_WITH_INFO( - input_seq_length == num_decoding_draft_tokens + 1, "The generation input length is not expected."); - common_enqueue_params.input_seq_length = input_seq_length; - common_enqueue_params.max_past_kv_length = max_past_kv_length; - EnqueueGenerationParams enqueue_params{common_enqueue_params}; - enqueue_params.beam_width = beamWidth; - enqueue_params.attention_mask_stride = attention_mask_stride; - enqueue_params.num_requests = num_requests; - enqueue_params.cache_indir = cache_indir; - enqueue_params.semaphores = multiBlockSemaphores(); - enqueue_params.host_past_key_value_lengths = host_past_kv_len_list; - enqueue_params.mrope_position_deltas = mrope_position_deltas; - if (mIsSpecDecodingEnabled && mUseSpecDecoding) - { - enqueue_params.spec_decoding_packed_mask = spec_decoding_packed_mask; - enqueue_params.spec_decoding_position_offsets = spec_decoding_position_offsets; - enqueue_params.spec_decoding_generation_lengths = spec_decoding_generation_lengths; - enqueue_params.spec_decoding_is_generation_length_variable = mSpecDecodingIsGenerationLengthVariable; - enqueue_params.spec_decoding_max_generation_length = mSpecDecodingMaxGenerationLength; - } - if (mFuseFp4Quant) - { - enqueue_params.start_token_idx_sf = tokenIdxBeg; - } - - if (changeSpecDecodingMode) - { - // mUseSpecDecoding is changed, need to re-prepare the DecoderXQARunner - prepareEnqueueGeneration(enqueue_params); - } - - enqueueGeneration(enqueue_params, stream); - - { - std::string const afterGenStr = "gen attention at layer " + std::to_string(mLayerIdx); - { - TLLM_CHECK_DEBUG_WITH_INFO( - tensorrt_llm::runtime::utils::tensorHasInvalid(localNbTokens, - outputDesc[0].dims.d[getPackedTensorHiddenDimIndex(mRemovePadding)], - mFP8ContextFMHA ? nvinfer1::DataType::kFP8 : mType, context_buf_, stream, afterGenStr) - == false, - "Found invalid number (NaN or Inf) in " + afterGenStr); - } - } - } - - return 0; -} - -template -int GPTAttentionPlugin::enqueueDispatchKVCacheType(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) -{ - if (mPagedKVCache) - { - return enqueueImpl(inputDesc, outputDesc, inputs, outputs, workspace, stream); - } - else - { - return enqueueImpl(inputDesc, outputDesc, inputs, outputs, workspace, stream); - } - return 0; -} - -int GPTAttentionPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) noexcept -{ - if (isBuilding()) - { - return 0; - } - if (mSkipAttn) - { - bool const* SKIP_ATTN = reinterpret_cast(inputs[getIdx(IdxEntry::SKIP_ATTN)]); - if (SKIP_ATTN[0]) - { - return 0; - } - } - - if (mType == nvinfer1::DataType::kHALF) - { - if (mFuseFp4Quant) - { - return enqueueDispatchKVCacheType(inputDesc, outputDesc, inputs, outputs, workspace, stream); - } -#ifdef ENABLE_FP8 - if (mFP8ContextFMHA) - { - return enqueueDispatchKVCacheType( - inputDesc, outputDesc, inputs, outputs, workspace, stream); - } -#endif - return enqueueDispatchKVCacheType(inputDesc, outputDesc, inputs, outputs, workspace, stream); - } - else if (mType == nvinfer1::DataType::kFLOAT) - { - return enqueueDispatchKVCacheType(inputDesc, outputDesc, inputs, outputs, workspace, stream); - } -#ifdef ENABLE_BF16 - else if (mType == nvinfer1::DataType::kBF16) - { - if (mFuseFp4Quant) - { - return enqueueDispatchKVCacheType<__nv_bfloat16, uint8_t>( - inputDesc, outputDesc, inputs, outputs, workspace, stream); - } -#ifdef ENABLE_FP8 - if (mFP8ContextFMHA) - { - return enqueueDispatchKVCacheType<__nv_bfloat16, __nv_fp8_e4m3>( - inputDesc, outputDesc, inputs, outputs, workspace, stream); - } -#endif - return enqueueDispatchKVCacheType<__nv_bfloat16>(inputDesc, outputDesc, inputs, outputs, workspace, stream); - } -#endif - return 0; -} - -// IPluginV2Ext Methods -nvinfer1::DataType GPTAttentionPlugin::getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept -{ - if (mFuseFp4Quant) - { - TLLM_CHECK(index == 0 || index == 1 || (!mPagedKVCache && useKVCache() && index == 2)); - } - else - { - TLLM_CHECK(index == 0 || (!mPagedKVCache && useKVCache() && index == 1)); - } - if (index == 0) - { - if (mFuseFp4Quant) - { - return nvinfer1::DataType::kFP4; - } - return mFP8ContextFMHA && mEnableContextFMHA ? nvinfer1::DataType::kFP8 - : inputTypes[getIdx(IdxEntry::QKV_TENSOR)]; - } - if (mFuseFp4Quant && index == 1) - { - return nvinfer1::DataType::kFP8; - } - return inputTypes[getIdx(IdxEntry::PAST_KEY_VALUE)]; -} - -// IPluginV2 Methods - -char const* GPTAttentionPlugin::getPluginType() const noexcept -{ - return GPT_ATTENTION_PLUGIN_NAME; -} - -char const* GPTAttentionPlugin::getPluginVersion() const noexcept -{ - return GPT_ATTENTION_PLUGIN_VERSION; -} - -int GPTAttentionPlugin::getNbOutputs() const noexcept -{ - int nbOutputs = mFuseFp4Quant ? 2 : 1; - if (!mPagedKVCache && useKVCache()) - { - nbOutputs += 1; - } - return nbOutputs; -} - -size_t GPTAttentionPlugin::getSerializationSize() const noexcept -{ - return GPTAttentionPluginCommon::getCommonSerializationSize(); -} - -void GPTAttentionPlugin::serialize(void* buffer) const noexcept -{ - GPTAttentionPluginCommon::serializeCommon(buffer); -} - -/////////////// - -GPTAttentionPluginCreator::GPTAttentionPluginCreator() - : GPTAttentionPluginCreatorCommon() -{ -} - -char const* GPTAttentionPluginCreator::getPluginName() const noexcept -{ - return GPT_ATTENTION_PLUGIN_NAME; -} - -char const* GPTAttentionPluginCreator::getPluginVersion() const noexcept -{ - return GPT_ATTENTION_PLUGIN_VERSION; -} - -PluginFieldCollection const* GPTAttentionPluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -IPluginV2* GPTAttentionPluginCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept -{ - PluginFieldParser p{fc->nbFields, fc->fields}; - - try - { - auto* obj = new GPTAttentionPlugin(p.getScalar("layer_idx").value(), - p.getScalar("num_heads").value(), p.getScalar("vision_start").value(), - p.getScalar("vision_length").value(), p.getScalar("num_kv_heads").value(), - p.getScalar("num_kv_heads_origin").value(), p.getScalar("head_size").value(), - p.getScalar("unidirectional").value(), p.getScalar("q_scaling").value(), - p.getScalar("attn_logit_softcapping_scale").value(), - static_cast(p.getScalar("position_embedding_type").value()), - p.getScalar("rotary_embedding_dim").value(), p.getScalar("rotary_embedding_base").value(), - static_cast(p.getScalar("rotary_embedding_scale_type").value()), - p.getScalar("rotary_embedding_scale").value(), - p.getScalar("rotary_embedding_short_m_scale").value(), - p.getScalar("rotary_embedding_long_m_scale").value(), - p.getScalar("rotary_embedding_max_positions").value(), - p.getScalar("rotary_embedding_original_max_positions").value(), - static_cast(p.getScalar("tp_size").value()), - static_cast(p.getScalar("tp_rank").value()), - static_cast(p.getScalar("unfuse_qkv_gemm").value()), - static_cast(p.getScalar("use_logn_scaling").value()), - static_cast(p.getScalar("context_fmha_type").value()), - p.getScalar("kv_cache_quant_mode").value(), - static_cast(p.getScalar("remove_input_padding").value()), - static_cast(p.getScalar("mask_type").value()), - BlockSparseParams{p.getScalar("block_sparse_block_size").value(), - static_cast(p.getScalar("block_sparse_homo_head_pattern").value()), - p.getScalar("block_sparse_num_local_blocks").value(), - p.getScalar("block_sparse_vertical_stride").value()}, - static_cast(p.getScalar("paged_kv_cache").value()), - p.getScalar("tokens_per_block").value(), - static_cast(p.getScalar("type_id").value()), - p.getScalar("max_context_length").value(), - static_cast(p.getScalar("qkv_bias_enabled").value()), - static_cast(p.getScalar("do_cross_attention").value()), - static_cast(p.getScalar("max_distance").value()), - static_cast(p.getScalar("pos_shift_enabled").value()), - static_cast(p.getScalar("dense_context_fmha").value()), - static_cast(p.getScalar("use_paged_context_fmha").value()), - static_cast(p.getScalar("use_fp8_context_fmha").value()), - static_cast(p.getScalar("has_full_attention_mask").value()), - static_cast(p.getScalar("use_cache").value()), - static_cast(p.getScalar("is_spec_decoding_enabled").value()), - static_cast(p.getScalar("spec_decoding_is_generation_length_variable").value()), - p.getScalar("spec_decoding_max_generation_length").value(), - static_cast(p.getScalar("is_mla_enabled").value()), - static_cast(p.getScalar("q_lora_rank").value()), - static_cast(p.getScalar("kv_lora_rank").value()), - static_cast(p.getScalar("qk_nope_head_dim").value()), - static_cast(p.getScalar("qk_rope_head_dim").value()), - static_cast(p.getScalar("v_head_dim").value()), - static_cast(p.getScalar("fuse_fp4_quant").value()), - static_cast(p.getScalar("skip_attn").value()), - static_cast(p.getScalar("cp_size").value()), - static_cast(p.getScalar("cp_rank").value()), - static_cast>(p.getSet("cp_group").value())); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -IPluginV2* GPTAttentionPluginCreator::deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept -{ - // This object will be deleted when the network is destroyed, which will - // call GPTAttentionPlugin::destroy() - try - { - auto* obj = new GPTAttentionPlugin(serialData, serialLength); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} diff --git a/cpp/tensorrt_llm/plugins/gptAttentionPlugin/gptAttentionPlugin.h b/cpp/tensorrt_llm/plugins/gptAttentionPlugin/gptAttentionPlugin.h deleted file mode 100644 index 3e34703c6221..000000000000 --- a/cpp/tensorrt_llm/plugins/gptAttentionPlugin/gptAttentionPlugin.h +++ /dev/null @@ -1,259 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. - */ -#pragma once - -#include "checkMacrosPlugin.h" -#include "tensorrt_llm/common/cublasMMWrapper.h" -#include "tensorrt_llm/common/logger.h" -#include "tensorrt_llm/common/quantization.h" -#include "tensorrt_llm/common/stringUtils.h" -#include "tensorrt_llm/kernels/contextFusedMultiHeadAttention/fmhaRunner.h" -#include "tensorrt_llm/kernels/contextFusedMultiHeadAttention/fused_multihead_attention_common.h" -#include "tensorrt_llm/kernels/gptKernels.h" -#include "tensorrt_llm/plugins/common/plugin.h" -#include "tensorrt_llm/plugins/gptAttentionCommon/gptAttentionCommon.h" -#include -#include -#include -#include -#include -#include - -namespace tensorrt_llm::plugins -{ -// batch_size = num_ctx_requests + num_gen_requests * beam_width -// num_ctx_requests = number of context requests (single sequence per request). -// num_gen_requests = number of generation requests (beam_width sequences per request). -// Context sequences have to appear first, generation sequences after - -// inputs (see GPTAttentionPlugin::isEntryUsed for when each tensor is actually used) -// 0. input_tensor [batch_size, seq_len, local_hidden_size + 2 * local_num_kv_heads * head_size] or -// [num_tokens, local_hidden_size + 2 * local_num_kv_heads * head_size] when -// enable_remove_input_padding -// 1. sequence_length [batch_size] (optional) -// 2. host_past_key_value_lengths [batch_size] (int32) (optional) -// 3. host_max_attention_window_sizes [num_layers] (int32) -// 4. host_sink_token_length [1] (int32) -// 5. context_lengths [batch_size] -// 6. cache_indir [num_gen_requests, beam_width, memory_max_len] (required in beamsearch) (optional) -// 7. host_request_types [batch_size] int32. 0: context; 1: generation: 2: none. When not in inflight-batching -// mode, -// all elements must be identical. -// 8. past_key_value_pool [batch_size, 2, local_num_kv_heads, max_seq_len, head_size] or -// block_offsets [batch_size, 2, max_blocks_per_seq] if paged kv cache (optional) -// 8.1 host_pool_pointers [2] if paged kv cache (optional) -// 9. kv_cache_quantization_scale [1] (optional) -// 10. kv_cache_dequantization_scale [1] (optional) -// 11. attention_output_quantization_scale [1] (on device, optional) -// 12. attention_mask [num_tokens, kv_seqlen] (on device, bool, optional) -// 13. attention_packed_mask [num_tokens, kv_seqlen / 32] (on device, uint32_t, optional) -// - pack masks by encoding multiple mask positions into a single 32-bit unsigned integer. -// - see kernels/contextMultiHeadAttention/fmhaPackedMask.cpp for more details. -// 14. rotary_inv_freq [head_size / 2] or [head_size] (longrope type) (float) (on device, optional) -// 15. rotary_cos_sin [max_num_embedding_positions, 2] (float) (on device, optional) -// 16. alibi_slopes [num_heads] (optional for ALiBi position embedding) -// 17. relative_attention_bias [num_heads] (optional for ALiBi position embedding) -// 18. host_context_lengths [batch_size] int32. (optional, required when remove_input_padding is true) -// 19. qkv_bias (optional) [local_hidden_size * 3] -// 20. spec_decoding_generation_lengths (optional, required when medusa is enabled) (int32_t) [batch_size] -// 21. spec_decoding_packed_mask (optional, required when medusa is enabled) (int32_t) [num_tokens, packed_mask_dim] -// packed_mask_dim = divUp(max_num_spec_decoding_tokens + 1, 32) -// 22. spec_decoding_position_offsets (optional, required when medusa is enabled) (int32_t) [batch_size, -// max_num_spec_decoding_tokens + 1] -// 23. spec_decoding_use (optional, bool) [1]: If it is set as true, enable speculative decoding -// 24. long_rope_rotary_inv_freq [head / 2] (float) (on device, optional) -// 25. long_rope_rotary_cos_sin [max_num_embedding_positions, 2] (float) (on device, optional) -// 26. host_runtime_perf_knobs (int64) -// 27. host_context_progress (void*) -// 28. position_id_tensor(MLA) [total_tokens], used for rope embedding in MLA -// 29. q_a_proj_tensor(MLA) [hidden_dim, c_q_dim + c_k_dim + ropd_dim], used to proj compacted QKV -// 30. q_a_layernorm_tensor(MLA) [c_q_dim], rmsnorm weight for compacted q -// 31. q_b_proj_tensor(MLA) [c_q_dim, head_num * head_size], weight for companted q to q in context -// 32. kv_a_proj_with_mqa_tensor(MLA) [c_q_dim, head_num * (c_k_dim + rope_dim)], weight for companted q to kdim in -// generation -// 33. kv_a_layernorm_tensor(MLA) [c_k_dim], rmsnorm weight for compacted kv -// 34. kv_b_proj_tensor(MLA) [c_k_dim, head_num * 2 * (head_size - rope_dim)], weight for compacted kv to kv in -// context -// 35. skip_attn (optional, bool) [1]: If it is set as true, skip the atteniton plugin and return -// directly. -// -// outputs -// output_tensor [batch_size, seq_len, local_hidden_size] -// present_key_value_pool (optional if not paged kv cache) [batch_size, 2, local_num_kv_heads, max_seq_len, -// head_size] - -class GPTAttentionPlugin : public GPTAttentionPluginCommon -{ -public: - GPTAttentionPlugin(int layer_idx, int num_heads, int vision_start, int vision_length, int num_kv_heads, - int num_kv_heads_origin, int head_size, int unidirectional, float q_scaling, float attn_logit_softcapping_scale, - tensorrt_llm::kernels::PositionEmbeddingType position_embedding_type, - int rotary_embedding_dim, // for RoPE. 0 for non-RoPE - float rotary_embedding_base, tensorrt_llm::kernels::RotaryScalingType rotary_embedding_scale_type, - float rotary_embedding_scale, float rotary_embedding_short_m_scale, float rotary_embedding_long_m_scale, - int rotary_embedding_max_positions, int rotary_embedding_original_max_positions, int tp_size, - int tp_rank, // for ALiBi - bool unfuse_qkv_gemm, // for AutoPP - bool use_logn_scaling, // for LognScaling - tensorrt_llm::kernels::ContextFMHAType context_fmha_type, int kv_cache_quant_mode, bool remove_input_padding, - tensorrt_llm::kernels::AttentionMaskType mask_type, - tensorrt_llm::kernels::BlockSparseParams block_sparse_params, bool paged_kv_cache, int tokens_per_block, - nvinfer1::DataType type, int32_t max_context_length, bool qkv_bias_enabled, bool cross_attention = false, - int max_distance = 0, bool pos_shift_enabled = false, bool dense_context_fmha = false, - bool use_paged_context_fmha = true, bool use_fp8_context_fmha = true, bool has_full_attention_mask = false, - bool use_cache = true, bool is_spec_decoding_enabled = false, - bool spec_decoding_is_generation_length_variable = false, int spec_decoding_max_generation_length = 1, - bool is_mla_enabled = false, int q_lora_rank = 0, int kv_lora_rank = 0, int qk_nope_head_dim = 0, - int qk_rope_head_dim = 0, int v_head_dim = 0, bool fuse_fp4_quant = false, bool skip_attn = false, - int cp_size = 1, int cp_rank = 0, std::set cp_group = {}); - - GPTAttentionPlugin(void const* data, size_t length); - - ~GPTAttentionPlugin() override = default; - - // IPluginV2DynamicExt Methods - nvinfer1::DimsExprs getOutputDimensions(int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; - - bool supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept override; - size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept override; - int enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - - template - int enqueueImpl(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream); - - template - int enqueueDispatchKVCacheType(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream); - - template - void configurePluginImpl(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept; - template - void configurePluginDispatchKVCacheType(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept; - void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept override; - - // IPluginV2Ext Methods - nvinfer1::DataType getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept override; - - // IPluginV2 Methods - char const* getPluginType() const noexcept override; - char const* getPluginVersion() const noexcept override; - int getNbOutputs() const noexcept override; - - //! This is called on every trt ExecutionContext creation by TRT - //! Note TRT does not call the initialize on cloned plugin, so clone internally should do initialization. - GPTAttentionPlugin* clone() const noexcept override; - - size_t getSerializationSize() const noexcept override; - void serialize(void* buffer) const noexcept override; - -private: - template - int enqueueSome(int32_t seqIdxBeg, int32_t localNbSeq, int32_t tokenIdxBeg, int32_t localNbTokens, - nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream); - - using IndexType = std::int32_t; - - std::vector mEntryIdx; - enum class IdxEntry : size_t - { - QKV_TENSOR, - K_TENSOR, - V_TENSOR, - ATTENTION_MASK, - ATTENTION_PACKED_MASK, - SEQUENCE_LENGTH, - HOST_PAST_KEY_VALUE_LENGTHS, - HOST_MAX_ATTENTION_WINDOW, - HOST_SINK_TOKEN_LENGTH, - CONTEXT_LENGTHS, - CACHE_INDIR, - REQUEST_TYPES, - KV_CACHE_BLOCK_OFFSETS, - HOST_KV_CACHE_BLOCK_OFFSETS, - HOST_KV_CACHE_POOL_POINTERS, - HOST_KV_CACHE_POOL_MAPPING, - PAST_KEY_VALUE, - KV_CACHE_QUANTIZATION_SCALE, - KV_CACHE_DEQUANTIZATION_SCALE, - ATTENTION_OUTPUT_QUANTIZATION_SCALE, - ATTENTION_OUTPUT_SF_SCALE, - ROTARY_INV_FREQ, - ROTARY_COS_SIN, - ALIBI_SLOPES, - RELATIVE_ATTENTION_BIAS, - CROSS_KV, - CROSS_KV_LENGTH, - ENCODER_INPUT_LENGTH, - HOST_CONTEXT_LENGTH, - QKV_BIAS_TENSOR, - SPEC_DECODING_GENERATION_LENGTHS, - SPEC_DECODING_PACKED_MASK, - SPEC_DECODING_POSITION_OFFSETS, - SPEC_DECODING_USE, - LONG_ROPE_ROTARY_INV_FREQ, - LONG_ROPE_ROTARY_COS_SIN, - MROPE_ROTARY_COS_SIN, - MROPE_POSITION_DELTAS, - HOST_RUNTIME_PERF_KNOBS, - HOST_CONTEXT_PROGRESS, - MLA_Q_B_PROJ_TENSOR, - MLA_KV_B_PROJ_TENSOR, - MLA_K_B_PROJ_TRANS_TENSOR, - SKIP_ATTN, - LOGN_SCALING, - ENUM_SIZE, // Used to count the number of IdxEntry, must put in last - }; - - std::string toString(IdxEntry const& entry) const; - bool isEntryUsed(IdxEntry const& entry) const; - void initEntryIdx(); - IndexType getIdx(IdxEntry const& entry) const; - - // Get generation input sequence length (might be larger than 1 in the speculative decoding mode). - int getGenerationInputSequenceLength( - nvinfer1::PluginTensorDesc const* inputDesc, int32_t localNbSeq, int32_t localNbTokens) const; -}; - -class GPTAttentionPluginCreator : public GPTAttentionPluginCreatorCommon -{ -public: - GPTAttentionPluginCreator(); - - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; - - nvinfer1::IPluginV2* deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept override; -}; - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/identityPlugin/CMakeLists.txt b/cpp/tensorrt_llm/plugins/identityPlugin/CMakeLists.txt deleted file mode 100644 index 86876224fccd..000000000000 --- a/cpp/tensorrt_llm/plugins/identityPlugin/CMakeLists.txt +++ /dev/null @@ -1,21 +0,0 @@ -# -# SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. -# -file(GLOB SRCS *.cpp) -set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS}) -set(PLUGIN_SOURCES - ${PLUGIN_SOURCES} - PARENT_SCOPE) diff --git a/cpp/tensorrt_llm/plugins/identityPlugin/identityPlugin.cpp b/cpp/tensorrt_llm/plugins/identityPlugin/identityPlugin.cpp deleted file mode 100644 index 109010e7a933..000000000000 --- a/cpp/tensorrt_llm/plugins/identityPlugin/identityPlugin.cpp +++ /dev/null @@ -1,199 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. - */ -#include "identityPlugin.h" -#include "tensorrt_llm/runtime/iBuffer.h" -#include "tensorrt_llm/runtime/iTensor.h" - -using namespace nvinfer1; -using tensorrt_llm::plugins::IdentityPluginCreator; -using tensorrt_llm::plugins::IdentityPlugin; - -static char const* IDENTITY_PLUGIN_VERSION{"1"}; -static char const* IDENTITY_PLUGIN_NAME{"Identity"}; -PluginFieldCollection IdentityPluginCreator::mFC{}; -std::vector IdentityPluginCreator::mPluginAttributes; - -IdentityPlugin::IdentityPlugin() {} - -// Parameterized constructor -IdentityPlugin::IdentityPlugin(void const* data, size_t length) -{ - char const *d = reinterpret_cast(data), *a = d; - TLLM_CHECK_WITH_INFO(d == a + length, - "Expected length (%d) != real length (%d). This is often " - "caused by using different TensorRT LLM version to build " - "engine and run engine.", - (int) length, (int) (d - a)); -} - -// IPluginV2DynamicExt Methods -nvinfer1::IPluginV2DynamicExt* IdentityPlugin::clone() const noexcept -{ - auto* plugin = new IdentityPlugin(*this); - plugin->setPluginNamespace(mNamespace.c_str()); - return plugin; -} - -nvinfer1::DimsExprs IdentityPlugin::getOutputDimensions( - int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept -{ - return inputs[outputIndex]; -} - -bool IdentityPlugin::supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept -{ - assert(0 <= pos && pos < 2); - PluginTensorDesc const& input = inOut[0]; - PluginTensorDesc const& output = inOut[1]; - switch (pos) - { - case 0: return input.format == nvinfer1::TensorFormat::kLINEAR; - case 1: return output.type == input.type && output.format == nvinfer1::TensorFormat::kLINEAR; - } - return false; -} - -void IdentityPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept -{ -} - -size_t IdentityPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - return 0; -} - -int IdentityPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept -{ - size_t count = 1; - for (int i = 0; i < inputDesc[0].dims.nbDims; ++i) - { - count *= inputDesc[0].dims.d[i]; - } - count *= tensorrt_llm::runtime::BufferDataType(inputDesc[0].type).getSize(); - - cudaMemcpyAsync(outputs[0], inputs[0], count, cudaMemcpyDeviceToDevice, stream); - - sync_check_cuda_error(stream); - return 0; -} - -// IPluginV2Ext Methods -nvinfer1::DataType IdentityPlugin::getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept -{ - assert(index == 0); - return inputTypes[0]; -} - -// IPluginV2 Methods - -char const* IdentityPlugin::getPluginType() const noexcept -{ - return IDENTITY_PLUGIN_NAME; -} - -char const* IdentityPlugin::getPluginVersion() const noexcept -{ - return IDENTITY_PLUGIN_VERSION; -} - -int IdentityPlugin::getNbOutputs() const noexcept -{ - return 1; -} - -int IdentityPlugin::initialize() noexcept -{ - return 0; -} - -void IdentityPlugin::terminate() noexcept {} - -size_t IdentityPlugin::getSerializationSize() const noexcept -{ - return 0; -} - -void IdentityPlugin::serialize(void* buffer) const noexcept {} - -void IdentityPlugin::destroy() noexcept -{ - // This gets called when the network containing plugin is destroyed - delete this; -} - -/////////////// - -IdentityPluginCreator::IdentityPluginCreator() -{ - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -char const* IdentityPluginCreator::getPluginName() const noexcept -{ - return IDENTITY_PLUGIN_NAME; -} - -char const* IdentityPluginCreator::getPluginVersion() const noexcept -{ - return IDENTITY_PLUGIN_VERSION; -} - -PluginFieldCollection const* IdentityPluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -IPluginV2* IdentityPluginCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept -{ - try - { - auto* obj = new IdentityPlugin(); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -IPluginV2* IdentityPluginCreator::deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept -{ - // This object will be deleted when the network is destroyed, which will - // call IdentityPlugin::destroy() - try - { - auto* obj = new IdentityPlugin(serialData, serialLength); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} diff --git a/cpp/tensorrt_llm/plugins/identityPlugin/identityPlugin.h b/cpp/tensorrt_llm/plugins/identityPlugin/identityPlugin.h deleted file mode 100644 index 9ab10601ae59..000000000000 --- a/cpp/tensorrt_llm/plugins/identityPlugin/identityPlugin.h +++ /dev/null @@ -1,89 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. - */ -#pragma once - -#include "tensorrt_llm/plugins/common/plugin.h" -#include -#include -#include -#include - -namespace tensorrt_llm::plugins -{ - -class IdentityPlugin : public BasePlugin -{ -public: - IdentityPlugin(); - - IdentityPlugin(void const* data, size_t length); - - ~IdentityPlugin() override = default; - - // IPluginV2DynamicExt Methods - nvinfer1::IPluginV2DynamicExt* clone() const noexcept override; - nvinfer1::DimsExprs getOutputDimensions(int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; - bool supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept override; - void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept override; - size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept override; - int enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - - // IPluginV2Ext Methods - nvinfer1::DataType getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept override; - - // IPluginV2 Methods - char const* getPluginType() const noexcept override; - char const* getPluginVersion() const noexcept override; - int getNbOutputs() const noexcept override; - int initialize() noexcept override; - void terminate() noexcept override; - size_t getSerializationSize() const noexcept override; - void serialize(void* buffer) const noexcept override; - void destroy() noexcept override; - -private: - const std::string mLayerName; -}; - -class IdentityPluginCreator : public BaseCreator -{ -public: - IdentityPluginCreator(); - - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; - - nvinfer1::IPluginV2* deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept override; - -private: - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; -}; - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/layernormQuantizationPlugin/CMakeLists.txt b/cpp/tensorrt_llm/plugins/layernormQuantizationPlugin/CMakeLists.txt deleted file mode 100755 index 86876224fccd..000000000000 --- a/cpp/tensorrt_llm/plugins/layernormQuantizationPlugin/CMakeLists.txt +++ /dev/null @@ -1,21 +0,0 @@ -# -# SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. -# -file(GLOB SRCS *.cpp) -set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS}) -set(PLUGIN_SOURCES - ${PLUGIN_SOURCES} - PARENT_SCOPE) diff --git a/cpp/tensorrt_llm/plugins/layernormQuantizationPlugin/layernormQuantizationPlugin.cpp b/cpp/tensorrt_llm/plugins/layernormQuantizationPlugin/layernormQuantizationPlugin.cpp deleted file mode 100644 index 02a40a00c919..000000000000 --- a/cpp/tensorrt_llm/plugins/layernormQuantizationPlugin/layernormQuantizationPlugin.cpp +++ /dev/null @@ -1,472 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. - */ -#include "layernormQuantizationPlugin.h" -#include "pluginUtils.h" -#include "tensorrt_llm/kernels/layernormKernels.h" - -using namespace nvinfer1; -using namespace tensorrt_llm::kernels; -using namespace tensorrt_llm::common; -using tensorrt_llm::plugins::LayernormQuantizationPluginCreator; -using tensorrt_llm::plugins::LayernormQuantizationPlugin; - -static char const* LAYERNORM_QUANTIZATION_PLUGIN_VERSION{"1"}; -static char const* LAYERNORM_QUANTIZATION_PLUGIN_NAME{"LayernormQuantization"}; -PluginFieldCollection LayernormQuantizationPluginCreator::mFC{}; -std::vector LayernormQuantizationPluginCreator::mPluginAttributes; - -LayernormQuantizationPlugin::LayernormQuantizationPlugin(float eps, bool useDiffOfSquares, - bool dynamicActivationScaling, bool sumPerToken, bool clampValEnabled, tensorrt_llm::common::QuantMode quantMode, - nvinfer1::DataType type, nvinfer1::DataType outputType) - : mEps(eps) - , mUseDiffOfSquares(useDiffOfSquares) - , mDynActScaling(dynamicActivationScaling) - , mType(type) - , mOutputType(outputType) - , mClampValEnabled(clampValEnabled) - , mQuantMode(quantMode) - , mSumPerToken(sumPerToken) -{ - TLLM_CHECK_WITH_INFO(mOutputType == nvinfer1::DataType::kINT8 || mOutputType == nvinfer1::DataType::kFP8, - "Only int8 or fp8 output type is allowed."); - // Check if the quant mode is valid. - TLLM_CHECK_WITH_INFO(mQuantMode.hasPerTokenScaling(), "The quant mode is not valid."); -} - -// Parameterized constructor -LayernormQuantizationPlugin::LayernormQuantizationPlugin(void const* data, size_t length) -{ - char const *d = reinterpret_cast(data), *a = d; - read(d, mEps); - read(d, mUseDiffOfSquares); - read(d, mDynActScaling); - read(d, mSumPerToken); - read(d, mClampValEnabled); - read(d, mQuantMode); - read(d, mType); - read(d, mOutputType); - TLLM_CHECK_WITH_INFO(d == a + length, - "Expected length (%d) != real length (%d). This is often " - "caused by using different TensorRT LLM version to build " - "engine and run engine.", - (int) length, (int) (d - a)); -} - -// IPluginV2DynamicExt Methods -nvinfer1::IPluginV2DynamicExt* LayernormQuantizationPlugin::clone() const noexcept -{ - auto* plugin = new LayernormQuantizationPlugin( - mEps, mUseDiffOfSquares, mDynActScaling, mSumPerToken, mClampValEnabled, mQuantMode, mType, mOutputType); - plugin->setPluginNamespace(mNamespace.c_str()); - return plugin; -} - -nvinfer1::DimsExprs LayernormQuantizationPlugin::getOutputDimensions( - int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept -{ - if (outputIndex == 0) - { - // Quantized output - return inputs[outputIndex]; - } - - // Dynamic scaling or per-token sum if enabled - try - { - if (outputIndex == 1) - { - TLLM_CHECK(mDynActScaling); - } - else if (outputIndex == 2) - { - TLLM_CHECK(mSumPerToken); - } - else - { - TLLM_CHECK(false); - } - DimsExprs ret; - ret.nbDims = inputs[0].nbDims; - for (int di = 0; di < ret.nbDims - 1; ++di) - { - ret.d[di] = inputs[0].d[di]; - } - ret.d[ret.nbDims - 1] = exprBuilder.constant(1); - return ret; - } - catch (std::exception const& e) - { - caughtError(e); - } - return DimsExprs{}; -} - -bool LayernormQuantizationPlugin::supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept -{ - int const totalPoses - = 6 + static_cast(mClampValEnabled) + static_cast(mDynActScaling) + static_cast(mSumPerToken); - TLLM_CHECK(0 <= pos && pos < totalPoses); - TLLM_CHECK(nbInputs == 4 + static_cast(mClampValEnabled)); - if (pos < nbInputs) - { - if (pos < 3) - { - // activatation, weight, bias - return (inOut[pos].type == mType) && (inOut[pos].format == TensorFormat::kLINEAR); - } - else if (pos == 3) - { - // scale - return (inOut[pos].type == nvinfer1::DataType::kFLOAT) && (inOut[pos].format == TensorFormat::kLINEAR); - } - else if (pos == 4 && mClampValEnabled) - { - // clamp_max_v - return (inOut[pos].type == nvinfer1::DataType::kFLOAT) && (inOut[pos].format == TensorFormat::kLINEAR); - } - } - else - { - auto const output_pos = pos - nbInputs; - if (output_pos == 0) - { - // Quantized output - return (inOut[pos].type == mOutputType) && (inOut[pos].format == TensorFormat::kLINEAR); - } - else if (output_pos == 1 && mDynActScaling) - { - // Dynamic scaling if enabled - return (inOut[pos].type == nvinfer1::DataType::kFLOAT) && (inOut[pos].format == TensorFormat::kLINEAR); - } - else if (output_pos == 2 && static_cast(mClampValEnabled)) - { - // Clamp value - return (inOut[pos].type == nvinfer1::DataType::kFLOAT) && (inOut[pos].format == TensorFormat::kLINEAR); - } - } - - // We should never reach this point - TLLM_CHECK_WITH_INFO(false, "The input/output is not supported."); - return false; -} - -void LayernormQuantizationPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept -{ -} - -size_t LayernormQuantizationPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - return 0; -} - -int LayernormQuantizationPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) noexcept -{ - // inputs - // input [M(*), N] - // weight [N, ] - // bias [N, ] - // scale_to_int [1] - // clamp_max_v [2], contains min val, and max val (optional) - // outputs - // output [M(*), N] Normalized activations, potentially with quantization applied. - // dynamic_scaling [M(*), 1] (Optional) Per-token scales if quantization is enabled. - // token_sums [M(*), 1] (Optional) Per-token sums of all the channels (before quantization). - - int64_t m64 = 1; - for (int i = 0; i < inputDesc[0].dims.nbDims - 1; ++i) - { - m64 *= inputDesc[0].dims.d[i]; - } - int const m = TLLM_INT32_CAST(m64); - int const n = TLLM_INT32_CAST(inputDesc[1].dims.d[0]); - - void const* input = inputs[0]; - void const* weight = inputs[1]; - void const* bias = inputs[2]; - void const* scale = inputs[3]; - void const* clampValPtr = mClampValEnabled ? inputs[4] : nullptr; - void* output = outputs[0]; - void* dynamic_scale = mDynActScaling ? outputs[1] : nullptr; - void* sum_per_token = mSumPerToken ? outputs[2] : nullptr; - - if (inputDesc[0].type == DataType::kFLOAT && mOutputType == DataType::kINT8) - { - dispatchDataType(nullptr, input, weight, bias, mEps, m, n, stream, mUseDiffOfSquares, - clampValPtr, scale, dynamic_scale, sum_per_token, output); - } -#ifdef ENABLE_FP8 - else if (inputDesc[0].type == DataType::kFLOAT && mOutputType == DataType::kFP8) - { - dispatchDataType(nullptr, input, weight, bias, mEps, m, n, stream, mUseDiffOfSquares, - clampValPtr, scale, dynamic_scale, sum_per_token, output); - } -#endif // ENABLE_FP8 - else if (inputDesc[0].type == DataType::kHALF && mOutputType == DataType::kINT8) - { - dispatchDataType(nullptr, input, weight, bias, mEps, m, n, stream, mUseDiffOfSquares, clampValPtr, - scale, dynamic_scale, sum_per_token, output); - } -#ifdef ENABLE_FP8 - else if (inputDesc[0].type == DataType::kHALF && mOutputType == DataType::kFP8) - { - dispatchDataType(nullptr, input, weight, bias, mEps, m, n, stream, mUseDiffOfSquares, - clampValPtr, scale, dynamic_scale, sum_per_token, output); - } -#endif // ENABLE_FP8 -#ifdef ENABLE_BF16 - else if (inputDesc[0].type == DataType::kBF16 && mOutputType == DataType::kINT8) - { - dispatchDataType<__nv_bfloat16, int8_t>(nullptr, input, weight, bias, mEps, m, n, stream, mUseDiffOfSquares, - clampValPtr, scale, dynamic_scale, sum_per_token, output); - } -#ifdef ENABLE_FP8 - else if (inputDesc[0].type == DataType::kBF16 && mOutputType == DataType::kFP8) - { - dispatchDataType<__nv_bfloat16, __nv_fp8_e4m3>(nullptr, input, weight, bias, mEps, m, n, stream, - mUseDiffOfSquares, clampValPtr, scale, dynamic_scale, sum_per_token, output); - } -#endif // ENABLE_FP8 -#endif // ENABLE_BF16 - sync_check_cuda_error(stream); - return 0; -} - -template -void LayernormQuantizationPlugin::dispatchDataType(void* out, void const* input, void const* gamma, void const* beta, - float const eps, int const tokens, int const hidden_dim, cudaStream_t stream, bool use_diff_of_squares, - void const* clampValPtr, void const* scale, void* dynamic_scale, void* sum_per_token, - void* normed_output_quant) noexcept -{ - // inputs - // activation [dim0(*), dim1] - // clamp_value [2], contains min val, and max val (optional) - // outputs - // quant [dim0(*), dim1] - // scale_tokens [dim0(*), 1] - - invokeGeneralLayerNorm(reinterpret_cast(out), reinterpret_cast(input), - reinterpret_cast(gamma), reinterpret_cast(beta), eps, tokens, hidden_dim, mQuantMode, - stream, use_diff_of_squares, reinterpret_cast(clampValPtr), reinterpret_cast(scale), - reinterpret_cast(dynamic_scale), reinterpret_cast(sum_per_token), - reinterpret_cast(normed_output_quant)); -} - -// IPluginV2Ext Methods -nvinfer1::DataType LayernormQuantizationPlugin::getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept -{ - assert(index <= 2); - - if (index == 0) - { - // Output 0 quantized output of layer norm - return mOutputType; - } - else if (index == 1) - { - assert(mDynActScaling); - // Output 1 dynamic act scaling - return nvinfer1::DataType::kFLOAT; - } - else if (index == 2) - { - assert(mDynActScaling && mSumPerToken); - // Output 2 per-token sums - return nvinfer1::DataType::kFLOAT; - } - - // We should never reach this point - TLLM_CHECK_WITH_INFO(false, "The output index is not supported."); - return nvinfer1::DataType::kFLOAT; -} - -// IPluginV2 Methods - -char const* LayernormQuantizationPlugin::getPluginType() const noexcept -{ - return LAYERNORM_QUANTIZATION_PLUGIN_NAME; -} - -char const* LayernormQuantizationPlugin::getPluginVersion() const noexcept -{ - return LAYERNORM_QUANTIZATION_PLUGIN_VERSION; -} - -int LayernormQuantizationPlugin::getNbOutputs() const noexcept -{ - return 1 + static_cast(mDynActScaling) + static_cast(mSumPerToken); -} - -int LayernormQuantizationPlugin::initialize() noexcept -{ - return 0; -} - -void LayernormQuantizationPlugin::terminate() noexcept {} - -size_t LayernormQuantizationPlugin::getSerializationSize() const noexcept -{ - return sizeof(mEps) + sizeof(mUseDiffOfSquares) + sizeof(mDynActScaling) + sizeof(mSumPerToken) - + sizeof(mClampValEnabled) + sizeof(mQuantMode) + sizeof(mType) + sizeof(mOutputType); -} - -void LayernormQuantizationPlugin::serialize(void* buffer) const noexcept -{ - char *d = static_cast(buffer), *a = d; - write(d, mEps); - write(d, mUseDiffOfSquares); - write(d, mDynActScaling); - write(d, mSumPerToken); - write(d, mClampValEnabled); - write(d, mQuantMode); - write(d, mType); - write(d, mOutputType); - TLLM_CHECK(d == a + getSerializationSize()); -} - -void LayernormQuantizationPlugin::destroy() noexcept -{ - // This gets called when the network containing plugin is destroyed - delete this; -} - -/////////////// - -LayernormQuantizationPluginCreator::LayernormQuantizationPluginCreator() -{ - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - mPluginAttributes.emplace_back(PluginField("eps", nullptr, PluginFieldType::kFLOAT32)); - mPluginAttributes.emplace_back(PluginField("use_diff_of_squares", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("dyn_act_scaling", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("sum_per_token", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("clamp_val_enabled", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("quant_mode", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("type_id", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("out_type_id", nullptr, PluginFieldType::kINT32)); - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -char const* LayernormQuantizationPluginCreator::getPluginName() const noexcept -{ - return LAYERNORM_QUANTIZATION_PLUGIN_NAME; -} - -char const* LayernormQuantizationPluginCreator::getPluginVersion() const noexcept -{ - return LAYERNORM_QUANTIZATION_PLUGIN_VERSION; -} - -PluginFieldCollection const* LayernormQuantizationPluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -IPluginV2* LayernormQuantizationPluginCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept -{ - PluginField const* fields = fc->fields; - tensorrt_llm::common::QuantMode quantMode{}; - float eps{}; - nvinfer1::DataType type{}; - nvinfer1::DataType outputType{}; - bool useDiffOfSquares{}; - bool dynamicActivationScaling{}; - bool sumPerToken{}; - bool clampValEnabled{}; - - // Read configurations from each fields - for (int i = 0; i < fc->nbFields; ++i) - { - char const* attrName = fields[i].name; - if (!strcmp(attrName, "eps")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kFLOAT32); - eps = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "type_id")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - type = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "dyn_act_scaling")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - dynamicActivationScaling = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "use_diff_of_squares")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - useDiffOfSquares = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "sum_per_token")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - sumPerToken = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "clamp_val_enabled")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - clampValEnabled = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "quant_mode")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - quantMode = QuantMode(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "out_type_id")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - outputType = static_cast(*(static_cast(fields[i].data))); - } - } - try - { - auto* obj = new LayernormQuantizationPlugin( - eps, useDiffOfSquares, dynamicActivationScaling, sumPerToken, clampValEnabled, quantMode, type, outputType); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -IPluginV2* LayernormQuantizationPluginCreator::deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept -{ - // This object will be deleted when the network is destroyed, which will - // call LayernormQuantizationPlugin::destroy() - try - { - auto* obj = new LayernormQuantizationPlugin(serialData, serialLength); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} diff --git a/cpp/tensorrt_llm/plugins/layernormQuantizationPlugin/layernormQuantizationPlugin.h b/cpp/tensorrt_llm/plugins/layernormQuantizationPlugin/layernormQuantizationPlugin.h deleted file mode 100644 index 5cf3fa7e022f..000000000000 --- a/cpp/tensorrt_llm/plugins/layernormQuantizationPlugin/layernormQuantizationPlugin.h +++ /dev/null @@ -1,110 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. - */ -#pragma once - -#include "tensorrt_llm/common/quantization.h" -#include "tensorrt_llm/plugins/common/plugin.h" -#include -#include -#include -#include - -namespace tensorrt_llm::plugins -{ - -class LayernormQuantizationPlugin : public BasePlugin -{ -public: - LayernormQuantizationPlugin(float eps, bool useDiffOfSquares, bool dynamicActivationScaling, bool sumPerToken, - bool clampValEnabled, tensorrt_llm::common::QuantMode quantMode, nvinfer1::DataType type, - nvinfer1::DataType outputType); - - LayernormQuantizationPlugin(void const* data, size_t length); - - ~LayernormQuantizationPlugin() override = default; - - // IPluginV2DynamicExt Methods - nvinfer1::IPluginV2DynamicExt* clone() const noexcept override; - nvinfer1::DimsExprs getOutputDimensions(int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; - bool supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept override; - void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept override; - size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept override; - int enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - - template - void dispatchDataType(void* out, void const* input, void const* gamma, void const* beta, float const eps, - int const tokens, int const hidden_dim, cudaStream_t stream, bool use_diff_of_squares, void const* clampValPtr, - void const* scale, void* dynamic_scale, void* sum_per_token, void* normed_output_quant) noexcept; - - // IPluginV2Ext Methods - nvinfer1::DataType getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept override; - - // IPluginV2 Methods - char const* getPluginType() const noexcept override; - char const* getPluginVersion() const noexcept override; - int getNbOutputs() const noexcept override; - int initialize() noexcept override; - void terminate() noexcept override; - size_t getSerializationSize() const noexcept override; - void serialize(void* buffer) const noexcept override; - void destroy() noexcept override; - -private: - float mEps; - bool mUseDiffOfSquares; - bool mDynActScaling; - nvinfer1::DataType mType; - - const std::string mLayerName; - // The quantized output data type - nvinfer1::DataType mOutputType; - // Do we clamp the input tensor? - bool mClampValEnabled; - // The quantization mode - tensorrt_llm::common::QuantMode mQuantMode; - // Should we output the sum of channels per-token? (Used by QServe GEMM) - bool mSumPerToken; -}; - -class LayernormQuantizationPluginCreator : public BaseCreator -{ -public: - LayernormQuantizationPluginCreator(); - - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; - - nvinfer1::IPluginV2* deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept override; - -private: - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; -}; - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/lookupPlugin/CMakeLists.txt b/cpp/tensorrt_llm/plugins/lookupPlugin/CMakeLists.txt deleted file mode 100644 index 86876224fccd..000000000000 --- a/cpp/tensorrt_llm/plugins/lookupPlugin/CMakeLists.txt +++ /dev/null @@ -1,21 +0,0 @@ -# -# SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. -# -file(GLOB SRCS *.cpp) -set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS}) -set(PLUGIN_SOURCES - ${PLUGIN_SOURCES} - PARENT_SCOPE) diff --git a/cpp/tensorrt_llm/plugins/lookupPlugin/lookupPlugin.cpp b/cpp/tensorrt_llm/plugins/lookupPlugin/lookupPlugin.cpp deleted file mode 100644 index e4d26f9e5ec6..000000000000 --- a/cpp/tensorrt_llm/plugins/lookupPlugin/lookupPlugin.cpp +++ /dev/null @@ -1,341 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. - */ - -#include - -#include "lookupPlugin.h" -#include "tensorrt_llm/kernels/lookupKernels.h" -#include "tensorrt_llm/plugins/common/plugin.h" - -using namespace nvinfer1; -using namespace tensorrt_llm::kernels; -using namespace tensorrt_llm::common; -using tensorrt_llm::plugins::LookupPluginCreator; -using tensorrt_llm::plugins::LookupPlugin; - -static char const* LOOKUP_PLUGIN_VERSION{"1"}; -static char const* LOOKUP_PLUGIN_NAME{"Lookup"}; -PluginFieldCollection LookupPluginCreator::mFC{}; -std::vector LookupPluginCreator::mPluginAttributes; - -LookupPlugin::LookupPlugin(nvinfer1::DataType type, int rank) - : mType(type) - , mRank(rank) -{ - mArch = tensorrt_llm::common::getSMVersion(); -} - -// Parameterized constructor -LookupPlugin::LookupPlugin(void const* data, size_t length) -{ - mArch = tensorrt_llm::common::getSMVersion(); - char const *d = reinterpret_cast(data), *a = d; - read(d, mType); - read(d, mRank); - TLLM_CHECK_WITH_INFO(d == a + length, - "Expected length (%d) != real length (%d). This is often " - "caused by using different TensorRT LLM version to build " - "engine and run engine.", - (int) length, (int) (d - a)); -} - -// IPluginV2DynamicExt Methods -nvinfer1::IPluginV2DynamicExt* LookupPlugin::clone() const noexcept -{ - auto* plugin = new LookupPlugin(*this); - plugin->setPluginNamespace(mNamespace.c_str()); - plugin->initialize(); - return plugin; -} - -nvinfer1::DimsExprs LookupPlugin::getOutputDimensions( - int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept -{ - try - { - TLLM_CHECK(nbInputs == 2 || nbInputs == 3); - TLLM_CHECK(outputIndex == 0); - DimsExprs ret; - int const nbDimsInput = inputs[0].nbDims; - int const nbDimsWeight = inputs[1].nbDims; - ret.nbDims = nbDimsInput + 1; - - for (int i = 0; i < nbDimsInput; ++i) - { - ret.d[i] = inputs[0].d[i]; - } - ret.d[nbDimsInput] = inputs[1].d[nbDimsWeight - 1]; - - return ret; - } - catch (std::exception const& e) - { - caughtError(e); - } - return DimsExprs{}; -} - -bool LookupPlugin::supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept -{ - bool res = false; - if (nbInputs == 2) - { - switch (pos) - { - case 0: res = ((inOut[0].type == DataType::kINT32) && (inOut[0].format == TensorFormat::kLINEAR)); break; - case 1: res = ((inOut[1].type == mType) && (inOut[1].format == TensorFormat::kLINEAR)); break; - case 2: res = ((inOut[2].type == mType) && (inOut[2].format == TensorFormat::kLINEAR)); break; - default: // should NOT be here! - res = false; - } - } - else - { - TLLM_CHECK_WITH_INFO(mArch == 90, "int8 weight only lookupPlugin is only supported in SM 90 now."); - switch (pos) - { - case 0: res = ((inOut[0].type == DataType::kINT32) && (inOut[0].format == TensorFormat::kLINEAR)); break; - case 1: - res = ((inOut[1].type == DataType::kINT8 || inOut[1].type == mType) - && (inOut[1].format == TensorFormat::kLINEAR)); - break; - case 2: res = ((inOut[2].type == mType) && (inOut[2].format == TensorFormat::kLINEAR)); break; - case 3: res = ((inOut[3].type == mType) && (inOut[3].format == TensorFormat::kLINEAR)); break; - default: // should NOT be here! - res = false; - } - } - return res; -} - -void LookupPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept -{ - mNbInputs = nbInputs; -} - -size_t LookupPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - return 0; -} - -int LookupPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept -{ - // inputs - // input [tokenNum] - // weight [localVocabSize, hidden] - // per_token_scales [localVocabSize], optional - // outputs - // embedding [tokenNum, hidden] - - int64_t tokenNum = 1; - for (int i = 0; i < inputDesc[0].dims.nbDims; ++i) - { - tokenNum *= inputDesc[0].dims.d[i]; - } - - int const localVocabSize = inputDesc[1].dims.d[0]; - int const hidden = inputDesc[1].dims.d[inputDesc[1].dims.nbDims - 1]; - int const* input = reinterpret_cast(inputs[0]); - - int offset = mRank * localVocabSize; - - if (mNbInputs == 3) - { - int8_t const* weight = reinterpret_cast(inputs[1]); - if (mType == DataType::kHALF) - { - half const* per_token_scales = reinterpret_cast(inputs[2]); - half* output = reinterpret_cast(outputs[0]); - invokeLookUp( - output, input, weight, tokenNum, offset, localVocabSize, hidden, per_token_scales, stream); - } - else if (mType == DataType::kFLOAT) - { - float const* per_token_scales = reinterpret_cast(inputs[2]); - float* output = reinterpret_cast(outputs[0]); - invokeLookUp( - output, input, weight, tokenNum, offset, localVocabSize, hidden, per_token_scales, stream); - } - else if (mType == DataType::kBF16) - { - __nv_bfloat16 const* per_token_scales = reinterpret_cast<__nv_bfloat16 const*>(inputs[2]); - __nv_bfloat16* output = reinterpret_cast<__nv_bfloat16*>(outputs[0]); - invokeLookUp<__nv_bfloat16, int8_t, int>( - output, input, weight, tokenNum, offset, localVocabSize, hidden, per_token_scales, stream); - } - } - else - { - if (mType == DataType::kHALF) - { - half const* weight = reinterpret_cast(inputs[1]); - half* output = reinterpret_cast(outputs[0]); - invokeLookUp( - output, input, weight, tokenNum, offset, localVocabSize, hidden, nullptr, stream); - } - else if (mType == DataType::kFLOAT) - { - float const* weight = reinterpret_cast(inputs[1]); - float* output = reinterpret_cast(outputs[0]); - invokeLookUp( - output, input, weight, tokenNum, offset, localVocabSize, hidden, nullptr, stream); - } - else if (mType == DataType::kBF16) - { - __nv_bfloat16 const* weight = reinterpret_cast<__nv_bfloat16 const*>(inputs[1]); - __nv_bfloat16* output = reinterpret_cast<__nv_bfloat16*>(outputs[0]); - invokeLookUp<__nv_bfloat16, __nv_bfloat16, int>( - output, input, weight, tokenNum, offset, localVocabSize, hidden, nullptr, stream); - } - } - sync_check_cuda_error(stream); - - return 0; -} - -// IPluginV2Ext Methods -nvinfer1::DataType LookupPlugin::getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept -{ - TLLM_CHECK(index == 0); - return mType; -} - -// IPluginV2 Methods - -char const* LookupPlugin::getPluginType() const noexcept -{ - return LOOKUP_PLUGIN_NAME; -} - -char const* LookupPlugin::getPluginVersion() const noexcept -{ - return LOOKUP_PLUGIN_VERSION; -} - -int LookupPlugin::getNbOutputs() const noexcept -{ - return 1; -} - -int LookupPlugin::initialize() noexcept -{ - return 0; -} - -void LookupPlugin::destroy() noexcept -{ - delete this; -} - -size_t LookupPlugin::getSerializationSize() const noexcept -{ - return sizeof(mType) + sizeof(mRank); -} - -void LookupPlugin::serialize(void* buffer) const noexcept -{ - char *d = static_cast(buffer), *a = d; - write(d, mType); - write(d, mRank); - - TLLM_CHECK(d == a + getSerializationSize()); -} - -void LookupPlugin::terminate() noexcept {} - -/////////////// - -LookupPluginCreator::LookupPluginCreator() -{ - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - mPluginAttributes.emplace_back(PluginField("type_id", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("rank", nullptr, PluginFieldType::kINT32)); - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -char const* LookupPluginCreator::getPluginName() const noexcept -{ - return LOOKUP_PLUGIN_NAME; -} - -char const* LookupPluginCreator::getPluginVersion() const noexcept -{ - return LOOKUP_PLUGIN_VERSION; -} - -PluginFieldCollection const* LookupPluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -IPluginV2* LookupPluginCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept -{ - PluginField const* fields = fc->fields; - nvinfer1::DataType type{}; - int rank{}; - // Read configurations from each fields - for (int i = 0; i < fc->nbFields; ++i) - { - char const* attrName = fields[i].name; - if (!strcmp(attrName, "type_id")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - type = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "rank")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - rank = static_cast(*(static_cast(fields[i].data))); - } - } - try - { - auto* obj = new LookupPlugin(type, rank); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -IPluginV2* LookupPluginCreator::deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept -{ - // This object will be deleted when the network is destroyed, which will - // call LookupPlugin::destroy() - try - { - auto* obj = new LookupPlugin(serialData, serialLength); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} diff --git a/cpp/tensorrt_llm/plugins/lookupPlugin/lookupPlugin.h b/cpp/tensorrt_llm/plugins/lookupPlugin/lookupPlugin.h deleted file mode 100644 index 4dddaa1d8bdc..000000000000 --- a/cpp/tensorrt_llm/plugins/lookupPlugin/lookupPlugin.h +++ /dev/null @@ -1,96 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. - */ -#pragma once - -#include "tensorrt_llm/plugins/common/plugin.h" -#include -#include -#include -#include - -namespace tensorrt_llm::plugins -{ - -class LookupPlugin : public BasePlugin -{ -public: - LookupPlugin() = delete; - - LookupPlugin(nvinfer1::DataType type, int rank); - - LookupPlugin(void const* data, size_t length); - - ~LookupPlugin() override = default; - - // IPluginV2DynamicExt Methods - nvinfer1::IPluginV2DynamicExt* clone() const noexcept override; - nvinfer1::DimsExprs getOutputDimensions(int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; - bool supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept override; - void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept override; - size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept override; - int enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - - // IPluginV2Ext Methods - nvinfer1::DataType getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept override; - - // IPluginV2 Methods - char const* getPluginType() const noexcept override; - char const* getPluginVersion() const noexcept override; - int getNbOutputs() const noexcept override; - int initialize() noexcept override; - void terminate() noexcept override; - size_t getSerializationSize() const noexcept override; - void serialize(void* buffer) const noexcept override; - void destroy() noexcept override; - -private: - const std::string mLayerName; - - nvinfer1::DataType mType; - int mRank; - int mNbInputs = 0; - int mArch; -}; - -class LookupPluginCreator : public BaseCreator -{ -public: - LookupPluginCreator(); - - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; - - nvinfer1::IPluginV2* deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept override; - -private: - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; -}; - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/loraPlugin/CMakeLists.txt b/cpp/tensorrt_llm/plugins/loraPlugin/CMakeLists.txt deleted file mode 100644 index 86876224fccd..000000000000 --- a/cpp/tensorrt_llm/plugins/loraPlugin/CMakeLists.txt +++ /dev/null @@ -1,21 +0,0 @@ -# -# SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. -# -file(GLOB SRCS *.cpp) -set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS}) -set(PLUGIN_SOURCES - ${PLUGIN_SOURCES} - PARENT_SCOPE) diff --git a/cpp/tensorrt_llm/plugins/loraPlugin/loraPlugin.cpp b/cpp/tensorrt_llm/plugins/loraPlugin/loraPlugin.cpp deleted file mode 100644 index 7a7d925a74f6..000000000000 --- a/cpp/tensorrt_llm/plugins/loraPlugin/loraPlugin.cpp +++ /dev/null @@ -1,525 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-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. - */ - -#include "loraPlugin.h" - -#include "pluginUtils.h" -#include "tensorrt_llm/common/assert.h" - -#include - -using namespace nvinfer1; -using namespace tensorrt_llm::common; -using tensorrt_llm::plugins::LoraPluginCreator; -using tensorrt_llm::plugins::LoraPlugin; -using tensorrt_llm::plugins::read; -using tensorrt_llm::plugins::write; - -static char const* LORA_PLUGIN_VERSION{"1"}; -static char const* LORA_PLUGIN_NAME{"Lora"}; -PluginFieldCollection LoraPluginCreator::mFC{}; -std::vector LoraPluginCreator::mPluginAttributes; - -LoraPlugin::LoraPlugin(int in_hidden_size, std::vector out_hidden_sizes, int transA, int transB, - int num_lora_modules, nvinfer1::DataType type, LoraPlugin::PluginProfilerPtr const& pluginProfiler, - bool remove_input_padding, int max_low_rank, int weight_index) - : mTransA(transA) - , mTransB(transB) - , mType(type) - , mRemoveInputPadding(remove_input_padding) - , mNumLoraModules(num_lora_modules) - , mInHiddenSize(in_hidden_size) - , mMaxLowRank(max_low_rank) - , mWeightIndex(weight_index) - , mPluginProfiler(pluginProfiler) -{ - TLLM_LOG_DEBUG("%s", __PRETTY_FUNCTION__); - mOutHiddenSizes.resize(mNumLoraModules); - mOutHiddenSizes.assign(out_hidden_sizes.begin(), out_hidden_sizes.end()); - init(); -} - -// Parameterized constructor -LoraPlugin::LoraPlugin(void const* data, size_t length, LoraPlugin::PluginProfilerPtr const& pluginProfiler) - : mPluginProfiler(pluginProfiler) -{ - TLLM_LOG_DEBUG("%s", __PRETTY_FUNCTION__); - char const *d = reinterpret_cast(data), *a = d; - read(d, mInHiddenSize); - read(d, mTransA); - read(d, mTransB); - read(d, mNumLoraModules); - read(d, mType); - read(d, mRemoveInputPadding); - read(d, mMaxLowRank); - read(d, mWeightIndex); - mOutHiddenSizes.resize(mNumLoraModules); - for (int i = 0; i < mNumLoraModules; i++) - { - read(d, mOutHiddenSizes[i]); - } - init(); - - mPluginProfiler->deserialize(d, mDims, mGemmId); - - TLLM_CHECK_WITH_INFO(d == a + length, - "Expected length (%d) != real length (%d). This is often " - "caused by using different TensorRT LLM version to build " - "engine and run engine.", - (int) length, (int) (d - a)); -} - -void LoraPlugin::init() -{ - TLLM_LOG_DEBUG("%s", __PRETTY_FUNCTION__); - - auto cublasHandle = getCublasHandle(); - auto cublasLtHandle = getCublasLtHandle(); - auto cublasWraper = std::make_shared(cublasHandle, cublasLtHandle, nullptr, nullptr); - - mLoraImpl = std::make_shared( - mInHiddenSize, mOutHiddenSizes, mTransA, mTransB, mNumLoraModules, mType, mMaxLowRank, cublasWraper); - - mPluginProfiler->setTranspose(mTransA, mTransB); - mGemmId = GemmIdCublas(mDims.n, mDims.k, mType, mTransA, mTransB, mType); -} - -// IPluginV2DynamicExt Methods -nvinfer1::IPluginV2DynamicExt* LoraPlugin::clone() const noexcept -{ - TLLM_LOG_DEBUG("%s", __PRETTY_FUNCTION__); - auto* plugin = new LoraPlugin(*this); - return plugin; -} - -nvinfer1::DimsExprs LoraPlugin::getOutputDimensions( - int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept -{ - TLLM_LOG_DEBUG("%s", __PRETTY_FUNCTION__); - try - { - TLLM_CHECK(outputIndex < mNumLoraModules); - int const nbDimsA = inputs[getInputTensorIdx()].nbDims; - DimsExprs ret; - ret.nbDims = nbDimsA; - - for (int i = 0; i < ret.nbDims; ++i) - { - ret.d[0] = 0; - } - - if (mTransA) - { - for (int i = 1; i < nbDimsA; ++i) - { - ret.d[i - 1] = inputs[getInputTensorIdx()].d[i]; - } - } - else - { - for (int i = 0; i < nbDimsA - 1; ++i) - { - ret.d[i] = inputs[getInputTensorIdx()].d[i]; - } - } - - auto const* outHiddenSize = exprBuilder.constant(mOutHiddenSizes.at(outputIndex)); - TLLM_CHECK(outHiddenSize != nullptr); - ret.d[ret.nbDims - 1] = outHiddenSize; - return ret; - } - catch (std::exception const& e) - { - caughtError(e); - } - return DimsExprs{}; -} - -bool LoraPlugin::supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept -{ - TLLM_LOG_DEBUG("%s", __PRETTY_FUNCTION__); - if (pos == getHostRequestTypesIdx()) - { - return inOut[pos].type == nvinfer1::DataType::kINT32; - } - else if (pos >= getLoraRanksIdx() && pos < getLoraRanksIdx() + mNumLoraModules) - { - return inOut[pos].type == nvinfer1::DataType::kINT32; - } - else if (pos >= getLoraWeightsPtrsIdx() && pos < getLoraWeightsPtrsIdx() + mNumLoraModules) - { - return inOut[pos].type == nvinfer1::DataType::kINT64; - } - else if (mRemoveInputPadding && pos == getHostContextLengthsIdx()) - { - return inOut[pos].type == nvinfer1::DataType::kINT32; - } - else - { - return (inOut[pos].type == mType) && (inOut[pos].format == TensorFormat::kLINEAR); - } -} - -void LoraPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept -{ - TLLM_LOG_DEBUG("%s", __PRETTY_FUNCTION__); - - auto const input = in[getInputTensorIdx()]; - - int const nbDimsA = input.max.nbDims; - - auto const minM = utils::computeMDimension(mTransA, input.min); - auto const maxM = utils::computeMDimension(mTransA, input.max); - auto const N = utils::computeNDimension(mTransB, in[getHostRequestTypesIdx()].max); - auto const K = static_cast(mTransA ? input.max.d[0] : input.max.d[nbDimsA - 1]); - - if (!mDims.isInitialized()) - { - mDims = {minM, maxM, N, K}; - } - mGemmId.n = N; - mGemmId.k = K; - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -size_t LoraPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - TLLM_LOG_DEBUG("%s", __PRETTY_FUNCTION__); - - int const nbReq = inputs[getLoraRanksIdx()].dims.d[0]; - auto const type = inputs[getInputTensorIdx()].type; - auto const numTokens = getNumTokens(inputs); - return mLoraImpl->getWorkspaceSize(numTokens, nbReq, type); -} - -int64_t LoraPlugin::getNumTokens(nvinfer1::PluginTensorDesc const* input_tensors) const -{ - int ndim = input_tensors[getInputTensorIdx()].dims.nbDims; - TLLM_CHECK_WITH_INFO( - 3 == ndim || 2 == ndim, "hidden_state dimension should be either 2 [numTokens, hidden], or 3 [b, s, hidden]"); - int64_t num_tokens = input_tensors[getInputTensorIdx()].dims.d[0]; - if (ndim == 3) - { - num_tokens *= input_tensors[getInputTensorIdx()].dims.d[1]; - } - return num_tokens; -} - -int LoraPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - if (isBuilding()) - { - return 0; - } - - auto const numReqs = inputDesc[getLoraRanksIdx()].dims.d[0]; - void const* input = inputs[getInputTensorIdx()]; - int const seqLen = mRemoveInputPadding ? 0 : inputDesc[getInputTensorIdx()].dims.d[1]; - int32_t const* reqTypes = static_cast(inputs[getHostRequestTypesIdx()]); - void const* const* loraRanks = &inputs[getLoraRanksIdx()]; - void const* const* loraWeightPtrs = &inputs[getLoraWeightsPtrsIdx()]; - int32_t const* hostContextLengths - = mRemoveInputPadding ? static_cast(inputs[getHostContextLengthsIdx()]) : nullptr; - - int numTokens = getNumTokens(inputDesc); - mExpandLoraWeightPtrs.clear(); - mExpandLoraRanks.clear(); - mExpandLoraWeightPtrs.reserve(mNumLoraModules * numTokens * 2); - mExpandLoraRanks.reserve(mNumLoraModules * numTokens); - - for (int loraModuleIdx = 0; loraModuleIdx < mNumLoraModules; loraModuleIdx++) - { - auto const loraWeightModulePtrs = static_cast(loraWeightPtrs[loraModuleIdx]); - auto const loraRankModule = static_cast(loraRanks[loraModuleIdx]); - - int idx = 0; - for (int reqId = 0; reqId < numReqs; reqId++) - { - // loraWeightModulePtrs has 3 pointers for each module: A,B, and an optional DoRA magnitude - // the current LoRA plugin does not apply DoRA scaling, so the magnitude is ignored - RequestType const reqType = static_cast(reqTypes[reqId]); - if (reqType == RequestType::kGENERATION) - { - mExpandLoraWeightPtrs.push_back(reinterpret_cast(loraWeightModulePtrs[reqId * 3])); - mExpandLoraWeightPtrs.push_back(reinterpret_cast(loraWeightModulePtrs[reqId * 3 + 1])); - mExpandLoraRanks.push_back(loraRankModule[reqId]); - idx += 1; - } - else - { - int contextLen = (mRemoveInputPadding ? hostContextLengths[reqId] : seqLen); - - for (int contextId = 0; contextId < contextLen; contextId++) - { - mExpandLoraWeightPtrs.push_back(reinterpret_cast(loraWeightModulePtrs[reqId * 3])); - mExpandLoraWeightPtrs.push_back(reinterpret_cast(loraWeightModulePtrs[reqId * 3 + 1])); - mExpandLoraRanks.push_back(loraRankModule[reqId]); - idx += 1; - } - } - } - - // In 1st generation phase cross attention qkv lora, cross qkv is skipped by passing an empty encoder_output - // (passing 0 to dim) getNumTokens() will get in cross qkv_lora. Skipping the check for this case. - if (numTokens > 0) - { - TLLM_CHECK_WITH_INFO(idx == numTokens, - fmtstr("LoraParams and input dims don't match, lora tokens %d input tokens %d", idx, numTokens)); - } - } - - // only used for unified gemm - auto bestTactic = mPluginProfiler->getBestConfig(numTokens, mGemmId); - mLoraImpl->setBestTactic(bestTactic); - mLoraImpl->run(numTokens, numReqs, input, mExpandLoraRanks.data(), mExpandLoraWeightPtrs.data(), mWeightIndex, - outputs, workspace, stream); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); - return 0; -} - -// IPluginV2Ext Methods -nvinfer1::DataType LoraPlugin::getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept -{ - TLLM_LOG_DEBUG("%s", __PRETTY_FUNCTION__); - TLLM_CHECK(index < mNumLoraModules); - return mType; -} - -// IPluginV2 Methods - -char const* LoraPlugin::getPluginType() const noexcept -{ - TLLM_LOG_DEBUG("%s", __PRETTY_FUNCTION__); - return LORA_PLUGIN_NAME; -} - -char const* LoraPlugin::getPluginVersion() const noexcept -{ - TLLM_LOG_DEBUG("%s", __PRETTY_FUNCTION__); - return LORA_PLUGIN_VERSION; -} - -int LoraPlugin::getNbOutputs() const noexcept -{ - TLLM_LOG_DEBUG("%s", __PRETTY_FUNCTION__); - return mNumLoraModules; -} - -int LoraPlugin::initialize() noexcept -{ - TLLM_LOG_DEBUG("%s", __PRETTY_FUNCTION__); - if (!mDims.isInitialized()) - { - return 0; - } - - mLoraImpl->setGemmConfig(); - - mPluginProfiler->profileTactics(mLoraImpl->getCublasWrapper(), mType, mDims, mGemmId); - return 0; -} - -void LoraPlugin::destroy() noexcept -{ - TLLM_LOG_DEBUG("%s", __PRETTY_FUNCTION__); - delete this; -} - -size_t LoraPlugin::getSerializationSize() const noexcept -{ - TLLM_LOG_DEBUG("%s", __PRETTY_FUNCTION__); - return sizeof(mInHiddenSize) + sizeof(mTransA) + sizeof(mTransB) + sizeof(mNumLoraModules) + sizeof(mType) - + mPluginProfiler->getSerializationSize(mGemmId) + sizeof(mRemoveInputPadding) + sizeof(mMaxLowRank) - + sizeof(mWeightIndex) + sizeof(int) * mNumLoraModules; // selected tactics container size -} - -void LoraPlugin::serialize(void* buffer) const noexcept -{ - TLLM_LOG_DEBUG("%s", __PRETTY_FUNCTION__); - char *d = static_cast(buffer), *a = d; - write(d, mInHiddenSize); - write(d, mTransA); - write(d, mTransB); - write(d, mNumLoraModules); - write(d, mType); - write(d, mRemoveInputPadding); - write(d, mMaxLowRank); - write(d, mWeightIndex); - for (int i = 0; i < mNumLoraModules; i++) - { - write(d, mOutHiddenSizes.at(i)); - } - mPluginProfiler->serialize(d, mGemmId); - TLLM_CHECK(d == a + getSerializationSize()); -} - -void LoraPlugin::terminate() noexcept {} - -/////////////// - -LoraPluginCreator::LoraPluginCreator() -{ - TLLM_LOG_DEBUG("%s", __PRETTY_FUNCTION__); - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - mPluginAttributes.emplace_back(PluginField("transA", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("transB", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("num_lora_modules", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("type_id", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("weight_index", nullptr, PluginFieldType::kINT32)); - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -char const* LoraPluginCreator::getPluginName() const noexcept -{ - TLLM_LOG_DEBUG("%s", __PRETTY_FUNCTION__); - return LORA_PLUGIN_NAME; -} - -char const* LoraPluginCreator::getPluginVersion() const noexcept -{ - TLLM_LOG_DEBUG("%s", __PRETTY_FUNCTION__); - return LORA_PLUGIN_VERSION; -} - -PluginFieldCollection const* LoraPluginCreator::getFieldNames() noexcept -{ - TLLM_LOG_DEBUG("%s", __PRETTY_FUNCTION__); - return &mFC; -} - -IPluginV2* LoraPluginCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept -{ - TLLM_LOG_DEBUG("%s", __PRETTY_FUNCTION__); - - PluginField const* fields = fc->fields; - nvinfer1::DataType type{}; - int num_lora_modules{}; - int in_hidden_size{}; - int transA{}; - int transB{}; - bool remove_input_padding{}; - int max_low_rank{}; - int weight_index{}; - // Read configurations from each fields - for (int i = 0; i < fc->nbFields; ++i) - { - char const* attrName = fields[i].name; - if (!strcmp(attrName, "in_hidden_size")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - in_hidden_size = *(static_cast(fields[i].data)); - } - else if (!strcmp(attrName, "transa")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - transA = *(static_cast(fields[i].data)); - } - else if (!strcmp(attrName, "transb")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - transB = *(static_cast(fields[i].data)); - } - else if (!strcmp(attrName, "type_id")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - type = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "remove_input_padding")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT8); - remove_input_padding = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "max_low_rank")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - max_low_rank = *(static_cast(fields[i].data)); - } - else if (!strcmp(attrName, "num_lora_modules")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - num_lora_modules = *(static_cast(fields[i].data)); - } - else if (!strcmp(attrName, "weight_index")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - weight_index = *(static_cast(fields[i].data)); - } - } - std::vector out_hidden_sizes; - out_hidden_sizes.resize(num_lora_modules); - for (int i = 0; i < fc->nbFields; ++i) - { - char const* attrName = fields[i].name; - for (int j = 0; j < num_lora_modules; j++) - { - if (!strcmp(attrName, fmtstr("out_hidden_size_%d", j).c_str())) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - out_hidden_sizes.at(j) = *(static_cast(fields[i].data)); - } - } - } - try - { - // LoraPluginCreator is unique and shared for an engine generation - // Create plugin profiler with shared tactics map - // FIXME enable tactic profiler - auto pluginProfiler = gemmPluginProfileManager.createGemmPluginProfiler(/* inference */ false, /* skip */ true); - auto* obj = new LoraPlugin(in_hidden_size, out_hidden_sizes, transA, transB, num_lora_modules, type, - pluginProfiler, remove_input_padding, max_low_rank, weight_index); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -IPluginV2* LoraPluginCreator::deserializePlugin(char const* name, void const* serialData, size_t serialLength) noexcept -{ - TLLM_LOG_DEBUG("%s", __PRETTY_FUNCTION__); - // This object will be deleted when the network is destroyed, which will - // call LoraPlugin::destroy() - try - { - // LoraPluginCreator is unique and shared for an engine generation - // Create plugin profiler with shared tactics map - // FIXME enable tactic profiler - auto pluginProfiler = gemmPluginProfileManager.createGemmPluginProfiler(/* inference */ true, /* skip */ true); - auto* obj = new LoraPlugin(serialData, serialLength, pluginProfiler); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} diff --git a/cpp/tensorrt_llm/plugins/loraPlugin/loraPlugin.h b/cpp/tensorrt_llm/plugins/loraPlugin/loraPlugin.h deleted file mode 100644 index 7795f7b7c76d..000000000000 --- a/cpp/tensorrt_llm/plugins/loraPlugin/loraPlugin.h +++ /dev/null @@ -1,159 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. - */ -#ifndef TRT_LORA_PLUGIN_H -#define TRT_LORA_PLUGIN_H -#include "tensorrt_llm/kernels/lora/lora.h" -#include "tensorrt_llm/plugins/common/gemmPluginProfiler.h" -#include "tensorrt_llm/plugins/common/plugin.h" -#include "tensorrt_llm/plugins/gemmPlugin/gemmPlugin.h" -#include -#include -#include - -namespace tensorrt_llm::plugins -{ - -class LoraPlugin : public BasePlugin -{ -public: - using PluginProfilerPtr = std::shared_ptr; - using ImplPtr = std::shared_ptr; - using Config = cublasLtMatmulHeuristicResult_t; - - LoraPlugin() = delete; - - LoraPlugin(int in_hidden_size, std::vector out_hidden_sizes, int transA, int transB, int num_lora_modules, - nvinfer1::DataType type, PluginProfilerPtr const& profiler, bool remove_input_padding, int max_low_rank, - int weight_index); - - LoraPlugin(void const* data, size_t length, PluginProfilerPtr const& profiler); - - ~LoraPlugin() override = default; - - // IPluginV2DynamicExt Methods - nvinfer1::IPluginV2DynamicExt* clone() const noexcept override; - nvinfer1::DimsExprs getOutputDimensions(int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; - bool supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept override; - void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept override; - size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept override; - int enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - - // IPluginV2Ext Methods - nvinfer1::DataType getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept override; - - // IPluginV2 Methods - char const* getPluginType() const noexcept override; - char const* getPluginVersion() const noexcept override; - int getNbOutputs() const noexcept override; - int initialize() noexcept override; - void terminate() noexcept override; - size_t getSerializationSize() const noexcept override; - void serialize(void* buffer) const noexcept override; - void destroy() noexcept override; - -private: - int64_t getNumTokens(nvinfer1::PluginTensorDesc const* input_tensors) const; - void init(); - - using IndexType = std::int32_t; - - IndexType getInputTensorIdx() const - { - return 0; - } - - IndexType getHostRequestTypesIdx() const - { - return 1; - } - - IndexType getLoraRanksIdx() const - { - return 2; - } - - IndexType getLoraWeightsPtrsIdx() const - { - return 2 + mNumLoraModules; - } - - IndexType getHostContextLengthsIdx() const - { - TLLM_CHECK(mRemoveInputPadding); - return 2 + mNumLoraModules + mNumLoraModules; - } - - enum class RequestType : int32_t - { - kCONTEXT = 0, - kGENERATION = 1 - }; - -private: - const std::string mLayerName; - - std::vector mOutHiddenSizes; - int mTransA; - int mTransB; - nvinfer1::DataType mType; - bool mRemoveInputPadding; - int mNumLoraModules; - int mInHiddenSize; - int mMaxLowRank; - int mWeightIndex; - - std::vector mExpandLoraWeightPtrs{}; - std::vector mExpandLoraRanks{}; - - GemmDims mDims{}; - GemmIdCublas mGemmId{}; - - PluginProfilerPtr mPluginProfiler; - ImplPtr mLoraImpl; -}; - -class LoraPluginCreator : public BaseCreator -{ -public: - LoraPluginCreator(); - - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; - - nvinfer1::IPluginV2* deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept override; - -private: - GemmPluginProfilerManager gemmPluginProfileManager; - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; -}; - -} // namespace tensorrt_llm::plugins - -#endif // TRT_LORA_PLUGIN_H diff --git a/cpp/tensorrt_llm/plugins/lowLatencyGemmPlugin/CMakeLists.txt b/cpp/tensorrt_llm/plugins/lowLatencyGemmPlugin/CMakeLists.txt deleted file mode 100644 index b6bd0439cc0c..000000000000 --- a/cpp/tensorrt_llm/plugins/lowLatencyGemmPlugin/CMakeLists.txt +++ /dev/null @@ -1,21 +0,0 @@ -# -# SPDX-FileCopyrightText: Copyright (c) 1993-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. -# -file(GLOB SRCS *.cpp) -set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS}) -set(PLUGIN_SOURCES - ${PLUGIN_SOURCES} - PARENT_SCOPE) diff --git a/cpp/tensorrt_llm/plugins/lowLatencyGemmPlugin/lowLatencyGemmPlugin.cpp b/cpp/tensorrt_llm/plugins/lowLatencyGemmPlugin/lowLatencyGemmPlugin.cpp deleted file mode 100644 index 6165d6210f29..000000000000 --- a/cpp/tensorrt_llm/plugins/lowLatencyGemmPlugin/lowLatencyGemmPlugin.cpp +++ /dev/null @@ -1,425 +0,0 @@ - -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-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. - */ - -#include "lowLatencyGemmPlugin.h" -#include "low_latency_gemm.h" -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/cudaFp8Utils.h" -#include "tensorrt_llm/common/logger.h" -#include -#include -#include -#include -#include -#include -#include -#include -#include - -using namespace nvinfer1; -using namespace tensorrt_llm::common; -using namespace tensorrt_llm::kernels::internal_cutlass_kernels; -using tensorrt_llm::plugins::LowLatencyGemmPluginCreator; -using tensorrt_llm::plugins::LowLatencyGemmPlugin; -using tensorrt_llm::plugins::LowLatencyGemmPluginProfiler; -using tensorrt_llm::plugins::read; -using tensorrt_llm::plugins::write; - -static char const* LOW_LATENCY_GEMM_PLUGIN_VERSION{"1"}; -static char const* LOW_LATENCY_GEMM_PLUGIN_NAME{"LowLatencyGemm"}; - -PluginFieldCollection LowLatencyGemmPluginCreator::mFC{}; -std::vector LowLatencyGemmPluginCreator::mPluginAttributes; - -using FP8Type = __nv_fp8_e4m3; - -static std::optional getFloatEnv(char const* name) -{ - char const* const env = std::getenv(name); - if (env == nullptr) - { - return std::nullopt; - } - try - { - float value = std::stof(env); - return {value}; - } - catch (std::invalid_argument const& e) - { - return std::nullopt; - } - catch (std::out_of_range const& e) - { - return std::nullopt; - } -}; - -void LowLatencyGemmPluginProfiler::runTactic(int m, int n, int k, LowLatencyGemmPluginProfiler::Config const& tactic, - char* workspace, cudaStream_t const& stream) -{ - - float default_pdl_overlap_ratio = 0.5; - float default_prefetch_ratio = -1.0; - FP8Type* aTmp = reinterpret_cast(workspace); - FP8Type* bTmp - = reinterpret_cast(nextWorkspacePtr(reinterpret_cast(aTmp), m * k * sizeof(FP8Type))); - void* cTmp = reinterpret_cast(nextWorkspacePtr(reinterpret_cast(bTmp), n * k * sizeof(FP8Type))); - size_t workspaceSize = mRunner->getWorkspaceSize(m, n, k); - char* workspaceTmp = reinterpret_cast(nextWorkspacePtr( - reinterpret_cast(cTmp), m * n * (mType == nvinfer1::DataType::kFLOAT ? sizeof(float) : sizeof(half)))); - mRunner->gemm(aTmp, bTmp, 1.0f, 0.0f, nullptr, cTmp, m, n, k, default_pdl_overlap_ratio, default_prefetch_ratio, - tactic, workspaceTmp, workspaceSize, stream); -} - -void LowLatencyGemmPluginProfiler::computeTmpSize(size_t maxM, size_t n, size_t k) -{ - - std::vector workspaces = {maxM * k * sizeof(FP8Type), n * k * sizeof(FP8Type), - maxM * n * (mType == nvinfer1::DataType::kFLOAT ? sizeof(float) : sizeof(half)), - mRunner->getWorkspaceSize(maxM, n, k)}; - - size_t bytes = calculateTotalWorkspaceSize(workspaces.data(), workspaces.size()); - setTmpWorkspaceSizeInBytes(bytes); -} - -std::vector LowLatencyGemmPluginProfiler::getTactics(int m, int n, int k) const -{ - return mRunner->getConfigs(); -} - -LowLatencyGemmPlugin::LowLatencyGemmPlugin( - nvinfer1::DataType type, float alpha, PluginProfilerPtr const& pluginProfiler) - : mPluginProfiler(pluginProfiler) - , mAplha(alpha) -{ - init(type); -} - -LowLatencyGemmPlugin::LowLatencyGemmPlugin(void const* data, size_t length, PluginProfilerPtr const& pluginProfiler) - : mPluginProfiler(pluginProfiler) -{ - - char const *d = reinterpret_cast(data), *a = d; - nvinfer1::DataType type; - read(d, type); - read(d, mAplha); - read(d, mDims); - init(type); - mPluginProfiler->deserialize(d, mDims, mGemmId); - TLLM_CHECK_WITH_INFO(d == a + length, - "Expected length (%d) != real length (%d). This is often " - "caused by using different TensorRT LLM version to build " - "engine and run engine.", - (int) length, (int) (d - a)); -} - -void LowLatencyGemmPlugin::init(nvinfer1::DataType type) -{ - - mType = type; - - if (mType == nvinfer1::DataType::kFLOAT) - { - m_lowLatencyGemmRunner = std::make_shared>(); - } - else if (mType == nvinfer1::DataType::kHALF) - { - m_lowLatencyGemmRunner = std::make_shared>(); - } -#ifdef ENABLE_BF16 - - else if (mType == nvinfer1::DataType::kBF16) - { - m_lowLatencyGemmRunner = std::make_shared>(); - } -#endif - else - { - TLLM_THROW("Unsupported data type"); - } - mGemmId = GemmIdCore(mDims.n, mDims.k, mType); -} - -nvinfer1::DimsExprs LowLatencyGemmPlugin::getOutputDimensions( - int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept -{ - try - { - TLLM_CHECK(nbInputs == 2); - TLLM_CHECK(outputIndex == 0); - int const nbDimsA = inputs[0].nbDims; - TLLM_CHECK(nbDimsA >= 2); - DimsExprs ret; - ret.nbDims = nbDimsA; - for (int ii = 0; ii < nbDimsA - 1; ++ii) - { - ret.d[ii] = inputs[0].d[ii]; - } - // input[1] , weights [n,k] - ret.d[nbDimsA - 1] = inputs[1].d[0]; - return ret; - } - catch (std::exception const& e) - { - caughtError(e); - } - return DimsExprs{}; -} - -bool LowLatencyGemmPlugin::supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept -{ - switch (pos) - { - case 0: - // activation - return inOut[pos].type == nvinfer1::DataType::kFP8 && inOut[pos].format == TensorFormat::kLINEAR; - case 1: - // weights - // Weights stored in checkpoint must have fp8 type - return inOut[pos].type == nvinfer1::DataType::kFP8 && inOut[pos].format == TensorFormat::kLINEAR; - case 2: - // out - return inOut[pos].type == mType && inOut[pos].format == TensorFormat::kLINEAR; - default: - // Never should be here - assert(false); - return false; - } -} - -void LowLatencyGemmPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept -{ - auto const minM = std::accumulate(in[0].min.d, in[0].min.d + in[0].min.nbDims - 1, 1, std::multiplies()); - auto const maxM = std::accumulate(in[0].max.d, in[0].max.d + in[0].max.nbDims - 1, 1, std::multiplies()); - - int const maxK = in[0].max.d[in[0].max.nbDims - 1]; - int const maxN = in[1].max.d[0]; - int const minK = in[0].min.d[in[0].min.nbDims - 1]; - int const minN = in[1].min.d[0]; - - TLLM_CHECK_WITH_INFO(minN == maxN, "Variable out channels is not allowed"); - TLLM_CHECK_WITH_INFO(minK == maxK, "Variable in channels is not allowed"); - - if (!mDims.isInitialized()) - { - mDims = {minM, maxM, maxN, maxK}; - } - mGemmId = {maxN, maxK, mType}; - - m_workspaceMaxSize = m_lowLatencyGemmRunner->getWorkspaceSize(maxM, maxN, maxK); -} - -size_t LowLatencyGemmPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - return m_workspaceMaxSize; -} - -int LowLatencyGemmPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) noexcept -{ - - // input0 activation [M,K] - // input1 weights [N,K] - // output0 [M,N] - - int64_t m64 = 1; - for (int ii = 0; ii < inputDesc[0].dims.nbDims - 1; ++ii) - { - m64 *= inputDesc[0].dims.d[ii]; - } - int const m = TLLM_INT32_CAST(m64); - int const n = TLLM_INT32_CAST(inputDesc[1].dims.d[0]); - int const k = TLLM_INT32_CAST(inputDesc[0].dims.d[inputDesc[0].dims.nbDims - 1]); - int const wsSize = m_lowLatencyGemmRunner->getWorkspaceSize(m, n, k); - auto const& bestTactic = mPluginProfiler->getBestConfig(m, mGemmId); - TLLM_CHECK_WITH_INFO(bestTactic, "No valid Low Latency GEMM tactic"); - - auto env_pdl_overlap_ratio = getFloatEnv("TRTLLM_PDL_OVERLAP_RATIO"); - auto env_prefetch_ratio = getFloatEnv("TRTLLM_PREFETCH_RATIO"); - auto valid_ratio = [](std::optional& env_val, float default_val) - { - if (env_val.has_value()) - { - TLLM_CHECK_WITH_INFO(env_val.value() <= 1.0f, "Valid ratio should be less than or equal to 1.0"); - return env_val.value(); - } - return default_val; - }; - float pdl_overlap_ratio = valid_ratio(env_pdl_overlap_ratio, /*default_val=*/0.5); - float prefetch_ratio = valid_ratio(env_prefetch_ratio, /*default_val=*/-1.0); - m_lowLatencyGemmRunner->gemm(const_cast(reinterpret_cast(inputs[0])), - const_cast(reinterpret_cast(inputs[1])), mAplha, 0.0F, nullptr, outputs[0], m, n, k, - pdl_overlap_ratio, prefetch_ratio, *bestTactic, reinterpret_cast(workspace), wsSize, stream); - - return 0; -} - -nvinfer1::DataType LowLatencyGemmPlugin::getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept -{ - TLLM_CHECK(index == 0); - return mType; -} - -// IPluginV2 Methods - -char const* LowLatencyGemmPlugin::getPluginType() const noexcept -{ - return LOW_LATENCY_GEMM_PLUGIN_NAME; -} - -char const* LowLatencyGemmPlugin::getPluginVersion() const noexcept -{ - return LOW_LATENCY_GEMM_PLUGIN_VERSION; -} - -int LowLatencyGemmPlugin::getNbOutputs() const noexcept -{ - return 1; -} - -int LowLatencyGemmPlugin::initialize() noexcept -{ - configGemm(); - return 0; -} - -void LowLatencyGemmPlugin::terminate() noexcept {} - -nvinfer1::IPluginV2DynamicExt* LowLatencyGemmPlugin::clone() const noexcept -{ - auto* plugin = new LowLatencyGemmPlugin(*this); - return plugin; -} - -size_t LowLatencyGemmPlugin::getSerializationSize() const noexcept -{ - return sizeof(nvinfer1::DataType) + // dtype - sizeof(float) * 1 + // alpha - sizeof(mDims) + mPluginProfiler->getSerializationSize(mGemmId); -} - -void LowLatencyGemmPlugin::serialize(void* buffer) const noexcept -{ - char *d = static_cast(buffer), *a = d; - write(d, mType); - write(d, mAplha); - write(d, mDims); - mPluginProfiler->serialize(d, mGemmId); - TLLM_CHECK(d == a + getSerializationSize()); -} - -void LowLatencyGemmPlugin::destroy() noexcept -{ - // This gets called when the network containing plugin is destroyed - delete this; -} - -void LowLatencyGemmPlugin::configGemm() -{ - mPluginProfiler->profileTactics(m_lowLatencyGemmRunner, mType, mDims, mGemmId); -} - -LowLatencyGemmPluginCreator::LowLatencyGemmPluginCreator() -{ - - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - mPluginAttributes.emplace_back(PluginField("alpha", nullptr, PluginFieldType::kFLOAT32)); - mPluginAttributes.emplace_back(PluginField("type_id", nullptr, PluginFieldType::kINT32)); - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -char const* LowLatencyGemmPluginCreator::getPluginName() const noexcept -{ - return LOW_LATENCY_GEMM_PLUGIN_NAME; -} - -char const* LowLatencyGemmPluginCreator::getPluginVersion() const noexcept -{ - return LOW_LATENCY_GEMM_PLUGIN_VERSION; -} - -PluginFieldCollection const* LowLatencyGemmPluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -IPluginV2* LowLatencyGemmPluginCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept -{ - PluginField const* fields = fc->fields; - float alpha{}; - nvinfer1::DataType type{}; - for (int i = 0; i < fc->nbFields; i++) - { - char const* attrName = fields[i].name; - if (!strcmp(attrName, "alpha")) - { - - TLLM_CHECK(fields[i].type == PluginFieldType::kFLOAT32); - alpha = *(static_cast(fields[i].data)); - } - else if (!strcmp(attrName, "type_id")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - type = static_cast(*(static_cast(fields[i].data))); - } - } - - try - { - - // - // GemmPluginCreator is unique and shared for an engine generation - // Create plugin profiler with shared tactics map - - auto pluginProfiler = gemmPluginProfileManager.createGemmPluginProfiler(/*inference=*/false); - auto* obj = new LowLatencyGemmPlugin(type, alpha, pluginProfiler); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -IPluginV2* LowLatencyGemmPluginCreator::deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept -{ - try - { - auto pluginProfiler = gemmPluginProfileManager.createGemmPluginProfiler(/*inference=*/true); - auto* obj = new LowLatencyGemmPlugin(serialData, serialLength, pluginProfiler); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} diff --git a/cpp/tensorrt_llm/plugins/lowLatencyGemmPlugin/lowLatencyGemmPlugin.h b/cpp/tensorrt_llm/plugins/lowLatencyGemmPlugin/lowLatencyGemmPlugin.h deleted file mode 100644 index 98b8f4807174..000000000000 --- a/cpp/tensorrt_llm/plugins/lowLatencyGemmPlugin/lowLatencyGemmPlugin.h +++ /dev/null @@ -1,135 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-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. - */ - -#pragma once - -#include "low_latency_gemm.h" - -#include "tensorrt_llm/plugins/common/gemmPluginProfiler.h" -#include "tensorrt_llm/plugins/common/plugin.h" -#include -#include -#include -#include -#include -#include - -namespace tensorrt_llm::plugins -{ - -using LowLatencyGemmRunnerPtr - = std::shared_ptr; - -class LowLatencyGemmPluginProfiler - : public GemmPluginProfiler< - tensorrt_llm::kernels::internal_cutlass_kernels::CutlassLowLatencyFp8GemmRunnerInterface::ConfigType, - LowLatencyGemmRunnerPtr, GemmIdCore, GemmIdCoreHash> -{ - -public: - using Config = tensorrt_llm::kernels::internal_cutlass_kernels::CutlassLowLatencyFp8GemmRunnerInterface::ConfigType; - -protected: - void runTactic(int m, int n, int k, Config const& tactic, char* workspace, cudaStream_t const& stream) override; - - void computeTmpSize(size_t maxM, size_t n, size_t k) override; - - std::vector getTactics(int m, int n, int k) const override; -}; - -class LowLatencyGemmPlugin : public BasePlugin -{ - -public: - using PluginProfilerPtr = std::shared_ptr; - - LowLatencyGemmPlugin() = delete; - - LowLatencyGemmPlugin(nvinfer1::DataType type, float alpha, PluginProfilerPtr const& pluginProfiler); - - LowLatencyGemmPlugin(void const* data, size_t length, PluginProfilerPtr const& pluginProfiler); - ~LowLatencyGemmPlugin() override = default; - - // IPluginV2DynamicExt Methods - nvinfer1::IPluginV2DynamicExt* clone() const noexcept override; - nvinfer1::DimsExprs getOutputDimensions(int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; - bool supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept override; - void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept override; - size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept override; - int enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - - // IPluginV2Ext Methods - nvinfer1::DataType getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept override; - - // IPluginV2 Methods - char const* getPluginType() const noexcept override; - char const* getPluginVersion() const noexcept override; - int getNbOutputs() const noexcept override; - int initialize() noexcept override; - void terminate() noexcept override; - size_t getSerializationSize() const noexcept override; - void serialize(void* buffer) const noexcept override; - void destroy() noexcept override; - -private: - void init(nvinfer1::DataType type); - void configGemm(); - -private: - std::string const mLayerName; - - LowLatencyGemmRunnerPtr m_lowLatencyGemmRunner; - size_t m_workspaceMaxSize; - - GemmDims mDims{}; - GemmIdCore mGemmId{}; - - PluginProfilerPtr mPluginProfiler; - - nvinfer1::DataType mType; - float mAplha{1.0F}; -}; - -class LowLatencyGemmPluginCreator : public BaseCreator -{ -public: - LowLatencyGemmPluginCreator(); - - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; - - nvinfer1::IPluginV2* deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept override; - -private: - GemmPluginProfilerManager gemmPluginProfileManager; - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; -}; - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/lowLatencyGemmSwigluPlugin/CMakeLists.txt b/cpp/tensorrt_llm/plugins/lowLatencyGemmSwigluPlugin/CMakeLists.txt deleted file mode 100644 index b6bd0439cc0c..000000000000 --- a/cpp/tensorrt_llm/plugins/lowLatencyGemmSwigluPlugin/CMakeLists.txt +++ /dev/null @@ -1,21 +0,0 @@ -# -# SPDX-FileCopyrightText: Copyright (c) 1993-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. -# -file(GLOB SRCS *.cpp) -set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS}) -set(PLUGIN_SOURCES - ${PLUGIN_SOURCES} - PARENT_SCOPE) diff --git a/cpp/tensorrt_llm/plugins/lowLatencyGemmSwigluPlugin/lowLatencyGemmSwigluPlugin.cpp b/cpp/tensorrt_llm/plugins/lowLatencyGemmSwigluPlugin/lowLatencyGemmSwigluPlugin.cpp deleted file mode 100644 index a1aa11c2f165..000000000000 --- a/cpp/tensorrt_llm/plugins/lowLatencyGemmSwigluPlugin/lowLatencyGemmSwigluPlugin.cpp +++ /dev/null @@ -1,468 +0,0 @@ - -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-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. - */ - -#include "lowLatencyGemmSwigluPlugin.h" -#include "low_latency_gemm_swiglu.h" -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/cudaFp8Utils.h" -#include "tensorrt_llm/common/logger.h" -#include -#include -#include -#include -#include -#include -#include -#include -#include - -using namespace nvinfer1; -using namespace tensorrt_llm::common; -using namespace tensorrt_llm::kernels::internal_cutlass_kernels; -using tensorrt_llm::plugins::LowLatencyGemmSwigluPluginCreator; -using tensorrt_llm::plugins::LowLatencyGemmSwigluPlugin; -using tensorrt_llm::plugins::LowLatencyGemmSwigluPluginProfiler; -using tensorrt_llm::plugins::read; -using tensorrt_llm::plugins::write; - -static char const* LOW_LATENCY_GEMM_SWIGLU_PLUGIN_VERSION{"1"}; -static char const* LOW_LATENCY_GEMM_SWIGLU_PLUGIN_NAME{"LowLatencyGemmSwiglu"}; - -PluginFieldCollection LowLatencyGemmSwigluPluginCreator::mFC{}; -std::vector LowLatencyGemmSwigluPluginCreator::mPluginAttributes; - -using FP8Type = __nv_fp8_e4m3; - -static std::optional getFloatEnv(char const* name) -{ - char const* const env = std::getenv(name); - if (env == nullptr) - { - return std::nullopt; - } - try - { - float value = std::stof(env); - return {value}; - } - catch (std::invalid_argument const& e) - { - return std::nullopt; - } - catch (std::out_of_range const& e) - { - return std::nullopt; - } -}; - -static size_t getBytePerElement(nvinfer1::DataType type) -{ - size_t bpe; - if (type == nvinfer1::DataType::kFLOAT) - { - bpe = 4; - } - else if (type == nvinfer1::DataType::kHALF || type == nvinfer1::DataType::kBF16) - { - bpe = 2; - } - else if (type == nvinfer1::DataType::kINT8 || type == nvinfer1::DataType::kFP8) - { - bpe = 1; - } - else - { - TLLM_THROW("Not recognized/implemented"); - } - return bpe; -} - -void LowLatencyGemmSwigluPluginProfiler::runTactic(int m, int n, int k, - LowLatencyGemmSwigluPluginProfiler::Config const& tactic, char* workspace, cudaStream_t const& stream) -{ - - float default_pdl_overlap_ratio = 0.5; - float default_prefetch_ratio = -1.0; - FP8Type* aTmp = reinterpret_cast(workspace); - FP8Type* bTmp - = reinterpret_cast(nextWorkspacePtr(reinterpret_cast(aTmp), m * k * sizeof(FP8Type))); - void* dTmp = reinterpret_cast(nextWorkspacePtr(reinterpret_cast(bTmp), n * k * sizeof(FP8Type))); - size_t workspaceSize = mRunner->getWorkspaceSize(m, n, k); - char* workspaceTmp = reinterpret_cast( - nextWorkspacePtr(reinterpret_cast(dTmp), (n / 2 * m * getBytePerElement(mType)))); - mRunner->gemm(aTmp, bTmp, 1.0f, 0.0f, 1.0f, 1.0f, nullptr, dTmp, m, n, k, default_pdl_overlap_ratio, - default_prefetch_ratio, tactic, workspaceTmp, workspaceSize, stream); -} - -int LowLatencyGemmSwigluPluginProfiler::getMaxProfileM() const -{ - return 32768; -} - -void LowLatencyGemmSwigluPluginProfiler::computeTmpSize(size_t maxM, size_t n, size_t k) -{ - - std::vector workspaces = {maxM * k * sizeof(FP8Type), // A - n * k * sizeof(FP8Type), // B - maxM * (n / 2) * getBytePerElement(mType), // D - mRunner->getWorkspaceSize(maxM, n, k)}; // workspace - - size_t bytes = calculateTotalWorkspaceSize(workspaces.data(), workspaces.size()); - setTmpWorkspaceSizeInBytes(bytes); -} - -std::vector LowLatencyGemmSwigluPluginProfiler::getTactics( - int m, int n, int k) const -{ - return mRunner->getConfigs(); -} - -LowLatencyGemmSwigluPlugin::LowLatencyGemmSwigluPlugin(nvinfer1::DataType type, float scale_output, float scale_d0, - float scale_d1, PluginProfilerPtr const& pluginProfiler) - : mPluginProfiler(pluginProfiler) - , mScaleOutput(scale_output) - , mScaleD0(scale_d0) - , mScaleD1(scale_d1) -{ - init(type); -} - -LowLatencyGemmSwigluPlugin::LowLatencyGemmSwigluPlugin( - void const* data, size_t length, PluginProfilerPtr const& pluginProfiler) - : mPluginProfiler(pluginProfiler) -{ - - char const *d = reinterpret_cast(data), *a = d; - nvinfer1::DataType type; - read(d, type); - read(d, mScaleOutput); - read(d, mScaleD0); - read(d, mScaleD1); - read(d, mDims); - - init(type); - mPluginProfiler->deserialize(d, mDims, mGemmId); - TLLM_CHECK_WITH_INFO(d == a + length, - "Expected length (%d) != real length (%d). This is often " - "caused by using different TensorRT LLM version to build " - "engine and run engine.", - (int) length, (int) (d - a)); -} - -void LowLatencyGemmSwigluPlugin::init(nvinfer1::DataType type) -{ - - mType = type; - - if (mType == nvinfer1::DataType::kFP8) - { - mLowLatencyGemmSwigluRunner = std::make_shared>(); - } - else - { - TLLM_THROW("Unsupported data type"); - } - mGemmId = GemmIdCore(mDims.n, mDims.k, mType); -} - -// IPluginV2DynamicExt Methods -nvinfer1::IPluginV2DynamicExt* LowLatencyGemmSwigluPlugin::clone() const noexcept -{ - auto* plugin = new LowLatencyGemmSwigluPlugin(*this); - return plugin; -} - -nvinfer1::DimsExprs LowLatencyGemmSwigluPlugin::getOutputDimensions( - int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept -{ - try - { - TLLM_CHECK(nbInputs == 2); - TLLM_CHECK(outputIndex == 0); - int const nbDimsA = inputs[0].nbDims; - TLLM_CHECK(nbDimsA >= 2); - DimsExprs ret; - ret.nbDims = nbDimsA; - for (int ii = 0; ii < nbDimsA - 1; ++ii) - { - ret.d[ii] = inputs[0].d[ii]; - } - ret.d[nbDimsA - 1] = exprBuilder.constant(inputs[1].d[1]->getConstantValue() / 2); - return ret; - } - catch (std::exception const& e) - { - caughtError(e); - } - return DimsExprs{}; -} - -bool LowLatencyGemmSwigluPlugin::supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept -{ - switch (pos) - { - case 0: - // activation - return inOut[pos].type == nvinfer1::DataType::kFP8 && inOut[pos].format == TensorFormat::kLINEAR; - case 1: - // weights - // Weights stored in checkpoint must have fp8 type - return inOut[pos].type == nvinfer1::DataType::kFP8 && inOut[pos].format == TensorFormat::kLINEAR; - case 2: - // out - return inOut[pos].type == mType && inOut[pos].format == TensorFormat::kLINEAR; - default: - // Never should be here - TLLM_CHECK(false); - return false; - } -} - -void LowLatencyGemmSwigluPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept -{ - auto const minM = std::accumulate(in[0].min.d, in[0].min.d + in[0].min.nbDims - 1, 1, std::multiplies()); - auto const maxM = std::accumulate(in[0].max.d, in[0].max.d + in[0].max.nbDims - 1, 1, std::multiplies()); - - int const maxK = in[0].max.d[in[0].max.nbDims - 1]; - int const maxN = in[1].max.d[1]; - int const minK = in[0].min.d[in[0].min.nbDims - 1]; - int const minN = in[1].min.d[1]; - - TLLM_CHECK_WITH_INFO(minN == maxN, "Variable out channels is not allowed"); - TLLM_CHECK_WITH_INFO(minK == maxK, "Variable in channels is not allowed"); - - if (!mDims.isInitialized()) - { - mDims = {minM, maxM, maxN, maxK}; - } - mGemmId = {maxN, maxK, mType}; - - mWorkspaceMaxSize = mLowLatencyGemmSwigluRunner->getWorkspaceSize(maxM, maxN, maxK); -} - -size_t LowLatencyGemmSwigluPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - return mWorkspaceMaxSize; -} - -int LowLatencyGemmSwigluPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) noexcept -{ - - // input0 activation [M,K] row-major - // input1 weights [K, N] col-major - // output0 [M,N / 2] row-major - - int64_t m64 = 1; - for (int ii = 0; ii < inputDesc[0].dims.nbDims - 1; ++ii) - { - m64 *= inputDesc[0].dims.d[ii]; - } - int const m = TLLM_INT32_CAST(m64); - int const n = TLLM_INT32_CAST(inputDesc[1].dims.d[1]); - int const k = TLLM_INT32_CAST(inputDesc[0].dims.d[inputDesc[0].dims.nbDims - 1]); - int const wsSize = mLowLatencyGemmSwigluRunner->getWorkspaceSize(m, n, k); - auto const& bestTactic = mPluginProfiler->getBestConfig(m, mGemmId); - TLLM_CHECK_WITH_INFO(bestTactic, "No valid Low Latency GEMM SWIGLU tactic"); - - auto env_pdl_overlap_ratio = getFloatEnv("TRTLLM_PDL_OVERLAP_RATIO"); - auto env_prefetch_ratio = getFloatEnv("TRTLLM_PREFETCH_RATIO"); - auto valid_ratio = [](std::optional& env_val, float default_val) - { - if (env_val.has_value()) - { - TLLM_CHECK_WITH_INFO(env_val.value() <= 1.0f, "Valid ratio should be less than or equal to 1.0"); - return env_val.value(); - } - return default_val; - }; - float pdl_overlap_ratio = valid_ratio(env_pdl_overlap_ratio, /*default_val=*/0.5); - float prefetch_ratio = valid_ratio(env_prefetch_ratio, /*default_val=*/-1.0); - mLowLatencyGemmSwigluRunner->gemm(const_cast(reinterpret_cast(inputs[0])), - const_cast(reinterpret_cast(inputs[1])), mScaleOutput, 0.0F, mScaleD0, mScaleD1, - nullptr, outputs[0], m, n, k, pdl_overlap_ratio, prefetch_ratio, *bestTactic, - reinterpret_cast(workspace), wsSize, stream); - - return 0; -} - -// IPluginV2Ext Methods -nvinfer1::DataType LowLatencyGemmSwigluPlugin::getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept -{ - TLLM_CHECK(index == 0); - return mType; -} - -// IPluginV2 Methods - -char const* LowLatencyGemmSwigluPlugin::getPluginType() const noexcept -{ - return LOW_LATENCY_GEMM_SWIGLU_PLUGIN_NAME; -} - -char const* LowLatencyGemmSwigluPlugin::getPluginVersion() const noexcept -{ - return LOW_LATENCY_GEMM_SWIGLU_PLUGIN_VERSION; -} - -int LowLatencyGemmSwigluPlugin::getNbOutputs() const noexcept -{ - return 1; -} - -int LowLatencyGemmSwigluPlugin::initialize() noexcept -{ - configGemm(); - return 0; -} - -void LowLatencyGemmSwigluPlugin::terminate() noexcept {} - -size_t LowLatencyGemmSwigluPlugin::getSerializationSize() const noexcept -{ - return sizeof(nvinfer1::DataType) + // dtype - sizeof(float) * 3 + // scales - sizeof(mDims) + mPluginProfiler->getSerializationSize(mGemmId); -} - -void LowLatencyGemmSwigluPlugin::serialize(void* buffer) const noexcept -{ - char *d = static_cast(buffer), *a = d; - write(d, mType); - write(d, mScaleOutput); - write(d, mScaleD0); - write(d, mScaleD1); - write(d, mDims); - mPluginProfiler->serialize(d, mGemmId); - TLLM_CHECK(d == a + getSerializationSize()); -} - -void LowLatencyGemmSwigluPlugin::destroy() noexcept -{ - // This gets called when the network containing plugin is destroyed - delete this; -} - -void LowLatencyGemmSwigluPlugin::configGemm() -{ - mPluginProfiler->profileTactics(mLowLatencyGemmSwigluRunner, mType, mDims, mGemmId); -} - -////////////////////////////////////////////////////////////////////////// - -LowLatencyGemmSwigluPluginCreator::LowLatencyGemmSwigluPluginCreator() -{ - - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - mPluginAttributes.emplace_back(PluginField("type_id", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("scale_output", nullptr, PluginFieldType::kFLOAT32)); - mPluginAttributes.emplace_back(PluginField("scale_d0", nullptr, PluginFieldType::kFLOAT32)); - mPluginAttributes.emplace_back(PluginField("scale_d1", nullptr, PluginFieldType::kFLOAT32)); - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -char const* LowLatencyGemmSwigluPluginCreator::getPluginName() const noexcept -{ - return LOW_LATENCY_GEMM_SWIGLU_PLUGIN_NAME; -} - -char const* LowLatencyGemmSwigluPluginCreator::getPluginVersion() const noexcept -{ - return LOW_LATENCY_GEMM_SWIGLU_PLUGIN_VERSION; -} - -PluginFieldCollection const* LowLatencyGemmSwigluPluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -IPluginV2* LowLatencyGemmSwigluPluginCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept -{ - PluginField const* fields = fc->fields; - TLLM_CHECK(fc->nbFields == 4); - nvinfer1::DataType type{}; - float scale_output{}; - float scale_d0{}; - float scale_d1{}; - for (int i = 0; i < fc->nbFields; i++) - { - char const* attrName = fields[i].name; - if (!strcmp(attrName, "type_id")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - type = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "scale_output")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kFLOAT32); - scale_output = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "scale_d0")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kFLOAT32); - scale_d0 = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "scale_d1")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kFLOAT32); - scale_d1 = static_cast(*(static_cast(fields[i].data))); - } - } - - try - { - - // - // LowLatencyGemmSwigluPluginCreator is unique and shared for an engine generation - // Create plugin profiler with shared tactics map - auto pluginProfiler = gemmPluginProfileManager.createGemmPluginProfiler(/*inference=*/false); - auto* obj = new LowLatencyGemmSwigluPlugin(type, scale_output, scale_d0, scale_d1, pluginProfiler); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -IPluginV2* LowLatencyGemmSwigluPluginCreator::deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept -{ - try - { - auto pluginProfiler = gemmPluginProfileManager.createGemmPluginProfiler(/*inference=*/true); - auto* obj = new LowLatencyGemmSwigluPlugin(serialData, serialLength, pluginProfiler); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} diff --git a/cpp/tensorrt_llm/plugins/lowLatencyGemmSwigluPlugin/lowLatencyGemmSwigluPlugin.h b/cpp/tensorrt_llm/plugins/lowLatencyGemmSwigluPlugin/lowLatencyGemmSwigluPlugin.h deleted file mode 100644 index 3f73324e7740..000000000000 --- a/cpp/tensorrt_llm/plugins/lowLatencyGemmSwigluPlugin/lowLatencyGemmSwigluPlugin.h +++ /dev/null @@ -1,140 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-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. - */ - -#pragma once - -#include "low_latency_gemm_swiglu.h" - -#include "tensorrt_llm/plugins/common/gemmPluginProfiler.h" -#include "tensorrt_llm/plugins/common/plugin.h" -#include -#include -#include -#include -#include -#include - -namespace tensorrt_llm::plugins -{ -using LowLatencyGemmSwigluRunnerPtr - = std::shared_ptr; - -class LowLatencyGemmSwigluPluginProfiler - : public GemmPluginProfiler< - tensorrt_llm::kernels::internal_cutlass_kernels::CutlassLowLatencyFp8GemmSwigluRunnerInterface::ConfigType, - LowLatencyGemmSwigluRunnerPtr, GemmIdCore, GemmIdCoreHash> -{ - -public: - using Config - = tensorrt_llm::kernels::internal_cutlass_kernels::CutlassLowLatencyFp8GemmSwigluRunnerInterface::ConfigType; - - virtual int getMaxProfileM() const override; - -protected: - void runTactic(int m, int n, int k, Config const& tactic, char* workspace, cudaStream_t const& stream) override; - - void computeTmpSize(size_t maxM, size_t n, size_t k) override; - - std::vector getTactics(int m, int n, int k) const override; -}; - -class LowLatencyGemmSwigluPlugin : public BasePlugin -{ - -public: - using PluginProfilerPtr = std::shared_ptr; - - LowLatencyGemmSwigluPlugin() = delete; - - LowLatencyGemmSwigluPlugin(nvinfer1::DataType type, float scale_output, float scale_d0, float scale_d1, - PluginProfilerPtr const& pluginProfiler); - - LowLatencyGemmSwigluPlugin(void const* data, size_t length, PluginProfilerPtr const& pluginProfiler); - ~LowLatencyGemmSwigluPlugin() override = default; - - // IPluginV2DynamicExt Methods - nvinfer1::IPluginV2DynamicExt* clone() const noexcept override; - nvinfer1::DimsExprs getOutputDimensions(int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; - bool supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept override; - void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept override; - size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept override; - int enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - - // IPluginV2Ext Methods - nvinfer1::DataType getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept override; - - // IPluginV2 Methods - char const* getPluginType() const noexcept override; - char const* getPluginVersion() const noexcept override; - int getNbOutputs() const noexcept override; - int initialize() noexcept override; - void terminate() noexcept override; - size_t getSerializationSize() const noexcept override; - void serialize(void* buffer) const noexcept override; - void destroy() noexcept override; - -private: - void init(nvinfer1::DataType type); - void configGemm(); - -private: - std::string const mLayerName; - - LowLatencyGemmSwigluRunnerPtr mLowLatencyGemmSwigluRunner; - size_t mWorkspaceMaxSize; - - GemmDims mDims{}; - GemmIdCore mGemmId{}; - - PluginProfilerPtr mPluginProfiler; - - nvinfer1::DataType mType; - float mScaleOutput; - float mScaleD0; - float mScaleD1; -}; - -class LowLatencyGemmSwigluPluginCreator : public BaseCreator -{ -public: - LowLatencyGemmSwigluPluginCreator(); - - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; - - nvinfer1::IPluginV2* deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept override; - -private: - GemmPluginProfilerManager gemmPluginProfileManager; - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; -}; - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/lruPlugin/CMakeLists.txt b/cpp/tensorrt_llm/plugins/lruPlugin/CMakeLists.txt deleted file mode 100644 index 86876224fccd..000000000000 --- a/cpp/tensorrt_llm/plugins/lruPlugin/CMakeLists.txt +++ /dev/null @@ -1,21 +0,0 @@ -# -# SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. -# -file(GLOB SRCS *.cpp) -set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS}) -set(PLUGIN_SOURCES - ${PLUGIN_SOURCES} - PARENT_SCOPE) diff --git a/cpp/tensorrt_llm/plugins/lruPlugin/lruPlugin.cpp b/cpp/tensorrt_llm/plugins/lruPlugin/lruPlugin.cpp deleted file mode 100644 index 9d86b8cb8acd..000000000000 --- a/cpp/tensorrt_llm/plugins/lruPlugin/lruPlugin.cpp +++ /dev/null @@ -1,431 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-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. - */ - -#include "lruPlugin.h" -#include "tensorrt_llm/common/assert.h" - -using namespace nvinfer1; -using namespace tensorrt_llm::kernels; -using namespace tensorrt_llm::common; -using tensorrt_llm::plugins::lruPluginCreator; -using tensorrt_llm::plugins::lruPlugin; - -static char const* LRU_PLUGIN_VERSION{"1"}; -static char const* LRU_PLUGIN_NAME{"LRU"}; -PluginFieldCollection lruPluginCreator::mFC{}; -std::vector lruPluginCreator::mPluginAttributes; - -lruPlugin::lruPlugin(int dim, int block_size, nvinfer1::DataType type, bool removePadding, bool pagedState, - bool yEnabled, bool yBiasEnabled, bool fuseGateEnabled, bool gateBiasEnabled) - : mDim(dim) - , mBlockSize(block_size) - , mType(type) - , mRemovePadding(removePadding) - , mPagedState(pagedState) - , mYEnabled(yEnabled) - , mYBiasEnabled(yBiasEnabled) - , mFuseGateEnabled(fuseGateEnabled) - , mGateBiasEnabled(gateBiasEnabled) -{ - TLLM_CHECK_WITH_INFO((mType == DataType::kBF16) || (mType == DataType::kFLOAT) || (mType == DataType::kHALF), - "Only support float, half, and bfloat16."); -} - -// Parameterized constructor -lruPlugin::lruPlugin(void const* data, size_t length) -{ - char const *d = reinterpret_cast(data), *a = d; - read(d, mDim); - read(d, mBlockSize); - read(d, mType); - read(d, mRemovePadding); - read(d, mPagedState); - read(d, mYEnabled); - read(d, mYBiasEnabled); - read(d, mFuseGateEnabled); - read(d, mGateBiasEnabled); - TLLM_CHECK(d == a + length); - TLLM_CHECK_WITH_INFO((mType == DataType::kBF16) || (mType == DataType::kFLOAT) || (mType == DataType::kHALF), - "Only support float, half, and bfloat16."); -} - -// IPluginV2DynamicExt Methods -nvinfer1::IPluginV2DynamicExt* lruPlugin::clone() const noexcept -{ - auto* plugin = new lruPlugin(mDim, mBlockSize, mType, mRemovePadding, mPagedState, mYEnabled, mYBiasEnabled, - mFuseGateEnabled, mGateBiasEnabled); - plugin->setPluginNamespace(mNamespace.c_str()); - return plugin; -} - -// Outputs -// output_tensor: [batch_size, seq_len, dim] or [num_tokens, dim] for remove_input_padding -// state: [batch_size, dim] -nvinfer1::DimsExprs lruPlugin::getOutputDimensions( - int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept -{ - if (outputIndex == 0) - { - return inputs[getXIdx()]; - } - return inputs[getStateIdx()]; -} - -bool lruPlugin::supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept -{ - if (pos == getHostRequestTypesIdx() || pos == getLastTokenIdsIdx() || (mPagedState && pos == getSlotMappingIdx())) - { - return inOut[pos].type == nvinfer1::DataType::kINT32; - } - else if (mPagedState && pos == getStateIdx()) - { - return inOut[pos].type == nvinfer1::DataType::kINT64; - } - else if (pos == getStateIdx() || pos == (nbInputs + 1)) - { - // Use float for both input and output state - return (inOut[pos].type == nvinfer1::DataType::kFLOAT) && (inOut[pos].format == TensorFormat::kLINEAR); - } - else - { - return (inOut[pos].type == mType) && (inOut[pos].format == TensorFormat::kLINEAR); - } -} - -void lruPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept -{ -} - -size_t lruPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - return 0; -} - -void lruPlugin::setLruParams(lruParams& params, const size_t batch, const size_t dim, const size_t block_size, - const size_t maxSeqLen, void* statePtr, void const* x, void const* gate, void const* gate_bias, void const* gate_x, - void const* gate_x_bias, void const* gate_a, void const* gate_a_bias, void const* y, void const* y_bias, - void const* A, int const* lastTokenIds, int const* slotMapping, void* out, bool removePadding) -{ - // Reset the parameters - memset(¶ms, 0, sizeof(params)); - - params.batch = batch; - params.width = dim; - params.block_size = block_size; - params.max_seqlen = maxSeqLen; - params.remove_padding = removePadding; - - // Set the pointers and strides. - params.A_ptr = const_cast(A); - params.x_ptr = const_cast(x); - params.y_ptr = const_cast(y); - params.y_bias_ptr = const_cast(y_bias); - params.gate_ptr = const_cast(gate); - params.gate_bias_ptr = const_cast(gate_bias); - params.gate_x_ptr = const_cast(gate_x); - params.gate_x_bias_ptr = const_cast(gate_x_bias); - params.gate_a_ptr = const_cast(gate_a); - params.gate_a_bias_ptr = const_cast(gate_a_bias); - params.state_ptr = statePtr; - params.out_ptr = out; - params.last_token_ids_ptr = lastTokenIds; - params.slot_mapping_ptr = slotMapping; -} - -template -int lruPlugin::enqueueImpl(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) -{ - // inputs - // 0. x [batch_size, seq_len, dim] or [num_tokens, dim] for remove_input_padding - // 1. A [dim] - // 2. state [batch_size, dim] or host [1] containing only pointer for paged_state - // 3. host_request_types [batch_size] int32. 0: context; 1: generation; 2: none. - // 4. last_token_ids [batch_size] int32 - // 5. state_slot_mapping [batch_size] int32, optional for paged state - // 6. y [batch_size, seq_len, dim] or [num_tokens, dim] for remove_input_padding - // 7. y_bias [dim] - // 8. gate [batch_size, seq_len, 2 * dim] or [num_tokens, 2 * dim] for remove_input_padding - // 9. gate_bias [2 * dim] - // 10. gate_x [batch_size, seq_len, dim] or [num_tokens, dim] for remove_input_padding - // 11. gate_a [batch_size, seq_len, dim] or [num_tokens, dim] for remove_input_padding - // 12. gate_x_bias [2 * dim] - // 13. gate_a_bias [2 * dim] - // outputs - // 0. output_tensor [batch_size, seq_len, dim] or [num_tokens, dim] for remove_input_padding - // 1. state [batch_size, dim] - auto const batch_size = inputDesc[getHostRequestTypesIdx()].dims.d[0]; - int max_seq_len; - if (mRemovePadding) - { - max_seq_len = -1; - } - else - { - max_seq_len = inputDesc[getXIdx()].dims.d[1]; - } - - // only support context or generation, not for both of them - RequestType const* reqTypes = static_cast(inputs[getHostRequestTypesIdx()]); - - lruParams lru_params; - - int const* slotMapping = mPagedState ? static_cast(inputs[getSlotMappingIdx()]) : nullptr; - void const* y = mYEnabled ? inputs[getYIdx()] : nullptr; - void const* y_bias = mYBiasEnabled ? inputs[getYBiasIdx()] : nullptr; - void const* gate = mFuseGateEnabled ? inputs[getGateIdx()] : nullptr; - void const* gate_bias = (mFuseGateEnabled && mGateBiasEnabled) ? inputs[getGateBiasIdx()] : nullptr; - void const* gate_x = mFuseGateEnabled ? nullptr : inputs[getGateXIdx()]; - void const* gate_a = mFuseGateEnabled ? nullptr : inputs[getGateAIdx()]; - void const* gate_x_bias = (!mFuseGateEnabled && mGateBiasEnabled) ? inputs[getGateXBiasIdx()] : nullptr; - void const* gate_a_bias = (!mFuseGateEnabled && mGateBiasEnabled) ? inputs[getGateABiasIdx()] : nullptr; - - void* statePtr = mPagedState ? *reinterpret_cast(const_cast(inputs[getStateIdx()])) : outputs[1]; - - setLruParams(lru_params, batch_size, mDim, mBlockSize, max_seq_len, statePtr, inputs[getXIdx()], gate, gate_bias, - gate_x, gate_x_bias, gate_a, gate_a_bias, y, y_bias, inputs[getAIdx()], - static_cast(inputs[getLastTokenIdsIdx()]), slotMapping, outputs[0], mRemovePadding); - - if (reqTypes[0] == RequestType::kCONTEXT) - { - invokeRGLRU(lru_params, stream); - } - else if (reqTypes[0] == RequestType::kGENERATION) - { - invokeRGLRUUpdate(lru_params, stream); - } - sync_check_cuda_error(stream); - return 0; -} - -int lruPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept -{ - if (isBuilding()) - { - return 0; - } - if (mType == DataType::kHALF) - { - return enqueueImpl(inputDesc, outputDesc, inputs, outputs, workspace, stream); - } - else if (mType == DataType::kFLOAT) - { - return enqueueImpl(inputDesc, outputDesc, inputs, outputs, workspace, stream); - } -#ifdef ENABLE_BF16 - else if (mType == DataType::kBF16) - { - return enqueueImpl<__nv_bfloat16>(inputDesc, outputDesc, inputs, outputs, workspace, stream); - } -#endif - return 0; -} - -// IPluginV2Ext Methods -nvinfer1::DataType lruPlugin::getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept -{ - if (index == 0) - { - return inputTypes[getXIdx()]; - } - else - { - return inputTypes[getStateIdx()]; - } -} - -// IPluginV2 Methods - -char const* lruPlugin::getPluginType() const noexcept -{ - return LRU_PLUGIN_NAME; -} - -char const* lruPlugin::getPluginVersion() const noexcept -{ - return LRU_PLUGIN_VERSION; -} - -int lruPlugin::getNbOutputs() const noexcept -{ - return mPagedState ? 1 : 2; -} - -int lruPlugin::initialize() noexcept -{ - return 0; -} - -void lruPlugin::terminate() noexcept {} - -size_t lruPlugin::getSerializationSize() const noexcept -{ - return sizeof(mDim) + sizeof(mBlockSize) + sizeof(mType) + sizeof(mRemovePadding) + sizeof(mPagedState) - + sizeof(mYEnabled) + sizeof(mYBiasEnabled) + sizeof(mFuseGateEnabled) + sizeof(mGateBiasEnabled); -} - -void lruPlugin::serialize(void* buffer) const noexcept -{ - char *d = static_cast(buffer), *a = d; - write(d, mDim); - write(d, mBlockSize); - write(d, mType); - write(d, mRemovePadding); - write(d, mPagedState); - write(d, mYEnabled); - write(d, mYBiasEnabled); - write(d, mFuseGateEnabled); - write(d, mGateBiasEnabled); - TLLM_CHECK(d == a + getSerializationSize()); -} - -void lruPlugin::destroy() noexcept -{ - delete this; -} - -/////////////// - -lruPluginCreator::lruPluginCreator() -{ - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - mPluginAttributes.emplace_back(PluginField("dim", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("block_size", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("type_id", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("remove_input_padding", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("paged_state", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("y_enabled", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("y_bias_enabled", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("fuse_gate_enabled", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("gate_bias_enabled", nullptr, PluginFieldType::kINT8)); - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -char const* lruPluginCreator::getPluginName() const noexcept -{ - return LRU_PLUGIN_NAME; -} - -char const* lruPluginCreator::getPluginVersion() const noexcept -{ - return LRU_PLUGIN_VERSION; -} - -PluginFieldCollection const* lruPluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -IPluginV2* lruPluginCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept -{ - PluginField const* fields = fc->fields; - int dim{}; - int block_size{}; - bool removePadding{}; - bool pagedState{}; - bool yEnabled{}; - bool yBiasEnabled{}; - bool fuseGateEnabled{}; - bool gateBiasEnabled{}; - nvinfer1::DataType type{}; - // Read configurations from each fields - for (int i = 0; i < fc->nbFields; ++i) - { - char const* attrName = fields[i].name; - if (!strcmp(attrName, "dim")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - dim = static_cast(*(static_cast(fields[i].data))); - } - if (!strcmp(attrName, "block_size")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - block_size = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "type_id")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - type = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "remove_input_padding")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT8); - removePadding = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "paged_state")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT8); - pagedState = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "y_enabled")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT8); - yEnabled = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "y_bias_enabled")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT8); - yBiasEnabled = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "fuse_gate_enabled")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT8); - fuseGateEnabled = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "gate_bias_enabled")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT8); - gateBiasEnabled = static_cast(*(static_cast(fields[i].data))); - } - } - try - { - auto* obj = new lruPlugin( - dim, block_size, type, removePadding, pagedState, yEnabled, yBiasEnabled, fuseGateEnabled, gateBiasEnabled); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -IPluginV2* lruPluginCreator::deserializePlugin(char const* name, void const* serialData, size_t serialLength) noexcept -{ - // This object will be deleted when the network is destroyed, which will - // call lruPlugin::destroy() - try - { - auto* obj = new lruPlugin(serialData, serialLength); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} diff --git a/cpp/tensorrt_llm/plugins/lruPlugin/lruPlugin.h b/cpp/tensorrt_llm/plugins/lruPlugin/lruPlugin.h deleted file mode 100644 index ee4e0b989b34..000000000000 --- a/cpp/tensorrt_llm/plugins/lruPlugin/lruPlugin.h +++ /dev/null @@ -1,239 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-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. - */ - -#ifndef TRT_LRU_PLUGIN_H -#define TRT_LRU_PLUGIN_H -#include "tensorrt_llm/kernels/lruKernel.h" -#include "tensorrt_llm/plugins/common/plugin.h" -#include - -namespace tensorrt_llm::plugins -{ -// batch_size = num_ctx_requests or num_gen_requests -// num_ctx_requests = number of context requests (single sequence per request). -// num_gen_requests = number of generation requests (single sequences per request). -// can not support beam search - -// inputs -// 0. x [batch_size, seq_len, dim] or [num_tokens, dim] for remove_input_padding -// 1. A [dim] -// 2. state [batch_size, dim] or host [1] containing only pointer for paged_state -// 3. host_request_types [batch_size] int32. 0: context; 1: generation; 2: none. -// 4. last_token_ids [batch_size] int32 -// 5. state_slot_mapping [batch_size] int32, optional for paged state -// 6. y [batch_size, seq_len, dim] or [num_tokens, dim] for remove_input_padding -// 7. y_bias [dim] -// 8. gate [batch_size, seq_len, 2 * dim] or [num_tokens, 2 * dim] for remove_input_padding -// 9. gate_bias [2 * dim] -// 10. gate_x [batch_size, seq_len, dim] or [num_tokens, dim] for remove_input_padding -// 11. gate_a [batch_size, seq_len, dim] or [num_tokens, dim] for remove_input_padding -// 12. gate_x_bias [2 * dim] -// 13. gate_a_bias [2 * dim] -// outputs -// 0. output_tensor [batch_size, seq_len, dim] or [num_tokens, dim] for remove_input_padding -// 1. state [batch_size, dim] - -class lruPlugin : public BasePlugin -{ -public: - lruPlugin(int dim, int block_size, nvinfer1::DataType type, bool removePadding, bool pagedState, bool yEnabled, - bool yBiasEnabled, bool fuseGateEnabled, bool gateBiasEnabled); - - lruPlugin(void const* data, size_t length); - - ~lruPlugin() override = default; - - // IPluginV2DynamicExt Methods - nvinfer1::IPluginV2DynamicExt* clone() const noexcept override; - nvinfer1::DimsExprs getOutputDimensions(int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; - bool supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept override; - void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept override; - size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept override; - int enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - template - int enqueueImpl(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream); - - // IPluginV2Ext Methods - nvinfer1::DataType getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept override; - - // IPluginV2 Methods - char const* getPluginType() const noexcept override; - char const* getPluginVersion() const noexcept override; - int getNbOutputs() const noexcept override; - int initialize() noexcept override; - void terminate() noexcept override; - size_t getSerializationSize() const noexcept override; - void serialize(void* buffer) const noexcept override; - void destroy() noexcept override; - - enum class RequestType : int32_t - { - kCONTEXT = 0, - kGENERATION = 1 - }; - -private: - using IndexType = std::int32_t; - - IndexType getXIdx() const - { - return 0; - }; - - IndexType getAIdx() const - { - return 1; - }; - - IndexType getStateIdx() const - { - return 2; - }; - - IndexType getHostRequestTypesIdx() const - { - return 3; - }; - - IndexType getLastTokenIdsIdx() const - { - return 4; - }; - - IndexType getSlotMappingIdx() const - { - if (mPagedState) - return 5; - else - return 4; - }; - - IndexType getYIdx() const - { - if (mYEnabled) - return getSlotMappingIdx() + 1; - else - return getSlotMappingIdx(); - }; - - IndexType getYBiasIdx() const - { - if (mYBiasEnabled) - return getYIdx() + 1; - else - return getYIdx(); - }; - - IndexType getGateIdx() const - { - if (mFuseGateEnabled) - return getYBiasIdx() + 1; - else - return getYBiasIdx(); - }; - - IndexType getGateBiasIdx() const - { - if (mFuseGateEnabled && mGateBiasEnabled) - return getGateIdx() + 1; - else - return getGateIdx(); - }; - - IndexType getGateXIdx() const - { - if (mFuseGateEnabled) - return getGateBiasIdx(); - else - return getGateBiasIdx() + 1; - }; - - IndexType getGateAIdx() const - { - if (mFuseGateEnabled) - return getGateXIdx(); - else - return getGateXIdx() + 1; - }; - - IndexType getGateXBiasIdx() const - { - if (!mFuseGateEnabled && mGateBiasEnabled) - return getGateAIdx() + 1; - else - return getGateAIdx(); - }; - - IndexType getGateABiasIdx() const - { - if (!mFuseGateEnabled && mGateBiasEnabled) - return getGateXBiasIdx() + 1; - else - return getGateXBiasIdx(); - }; - - static void setLruParams(tensorrt_llm::kernels::lruParams& params, - // sizes - const size_t batch, const size_t dim, const size_t block_size, const size_t maxSeqLen, - // device pointers - void* statePtr, void const* x, void const* gate, void const* gate_bias, void const* gate_x, - void const* gate_x_bias, void const* gate_a, void const* gate_a_bias, void const* y, void const* y_bias, - void const* A, int const* lastTokenIds, int const* slotMapping, void* out, bool removePadding); - -private: - int mDim; - int mBlockSize; - nvinfer1::DataType mType; - bool mRemovePadding = false; - bool mPagedState = false; - bool mYEnabled = false; - bool mYBiasEnabled = false; - bool mFuseGateEnabled = false; - bool mGateBiasEnabled = false; -}; - -class lruPluginCreator : public BaseCreator -{ -public: - lruPluginCreator(); - - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; - - nvinfer1::IPluginV2* deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept override; - -private: - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; -}; - -} // namespace tensorrt_llm::plugins - -#endif // TRT_LRU_PLUGIN_H diff --git a/cpp/tensorrt_llm/plugins/mambaConv1dPlugin/CMakeLists.txt b/cpp/tensorrt_llm/plugins/mambaConv1dPlugin/CMakeLists.txt deleted file mode 100644 index 86876224fccd..000000000000 --- a/cpp/tensorrt_llm/plugins/mambaConv1dPlugin/CMakeLists.txt +++ /dev/null @@ -1,21 +0,0 @@ -# -# SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. -# -file(GLOB SRCS *.cpp) -set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS}) -set(PLUGIN_SOURCES - ${PLUGIN_SOURCES} - PARENT_SCOPE) diff --git a/cpp/tensorrt_llm/plugins/mambaConv1dPlugin/mambaConv1dPlugin.cpp b/cpp/tensorrt_llm/plugins/mambaConv1dPlugin/mambaConv1dPlugin.cpp deleted file mode 100644 index 16754248b84d..000000000000 --- a/cpp/tensorrt_llm/plugins/mambaConv1dPlugin/mambaConv1dPlugin.cpp +++ /dev/null @@ -1,404 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-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. - */ - -#include "mambaConv1dPlugin.h" -#include "tensorrt_llm/common/assert.h" -#include - -using namespace nvinfer1; -using namespace tensorrt_llm::kernels; -using namespace tensorrt_llm::common; -using tensorrt_llm::plugins::MambaConv1dPluginCreator; -using tensorrt_llm::plugins::MambaConv1dPlugin; - -static char const* MAMBA_CONV1D_PLUGIN_VERSION{"1"}; -static char const* MAMBA_CONV1D_PLUGIN_NAME{"MambaConv1d"}; - -PluginFieldCollection MambaConv1dPluginCreator::mFC{}; -std::vector MambaConv1dPluginCreator::mPluginAttributes; - -MambaConv1dPlugin::MambaConv1dPlugin(int dim, int dconv, int preStride, int postStride, nvinfer1::DataType type, - bool removePadding, bool pagedState, bool applySilu) - : mDim(dim) - , mDConv(dconv) - , mPreStride(preStride) - , mPostStride(postStride) - , mType(type) - , mRemovePadding(removePadding) - , mPagedState(pagedState) - , mApplySilu(applySilu) -{ - TLLM_CHECK_WITH_INFO((mType == DataType::kBF16) || (mType == DataType::kFLOAT) || (mType == DataType::kHALF), - "Only support float, half, and bfloat16."); -} - -// Parameterized constructor -MambaConv1dPlugin::MambaConv1dPlugin(void const* data, size_t length) -{ - char const *d = reinterpret_cast(data), *a = d; - read(d, mDim); - read(d, mDConv); - read(d, mPreStride); - read(d, mPostStride); - read(d, mType); - read(d, mRemovePadding); - read(d, mPagedState); - read(d, mApplySilu); - TLLM_CHECK(d == a + length); - TLLM_CHECK_WITH_INFO((mType == DataType::kBF16) || (mType == DataType::kFLOAT) || (mType == DataType::kHALF), - "Only support float, half, and bfloat16."); -} - -// IPluginV2DynamicExt Methods -nvinfer1::IPluginV2DynamicExt* MambaConv1dPlugin::clone() const noexcept -{ - auto* plugin - = new MambaConv1dPlugin(mDim, mDConv, mPreStride, mPostStride, mType, mRemovePadding, mPagedState, mApplySilu); - plugin->setPluginNamespace(mNamespace.c_str()); - return plugin; -} - -// Outputs -// output_tensor: [batch_size, seq_len, dim] or [num_tokens, dim] for remove_input_padding -// state: [batch_size, dconv - 1, dim] -nvinfer1::DimsExprs MambaConv1dPlugin::getOutputDimensions( - int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept -{ - if (outputIndex == 0) - { - auto ret = inputs[getInputTensorIdx()]; - ret.d[mRemovePadding ? 1 : 2] = exprBuilder.constant(mDim); - return ret; - } - return inputs[getConvStateIdx()]; -} - -bool MambaConv1dPlugin::supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept -{ - if (pos == getHostRequestTypesIdx() || pos == getLastTokenIdsIdx() - || (mRemovePadding && pos == getHostContextLengthIdx()) || (mPagedState && pos == getSlotMappingIdx())) - { - return inOut[pos].type == nvinfer1::DataType::kINT32; - } - else if (mPagedState && pos == getConvStateIdx()) - { - return inOut[pos].type == nvinfer1::DataType::kINT64; - } - else - { - return (inOut[pos].type == mType) && (inOut[pos].format == TensorFormat::kLINEAR); - } -} - -void MambaConv1dPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept -{ -} - -size_t MambaConv1dPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - return 0; -} - -void MambaConv1dPlugin::setMambaConv1dParams(tensorrt_llm::kernels::MambaConv1dParamsBase& params, const size_t batch, - const size_t dim, const size_t maxSeqLen, const size_t dconv, const size_t preStride, const size_t postStride, - void const* inPtr, void const* stateInPtr, void* stateOutPtr, void const* convWeight, void const* convBias, - void* outPtr, int const* lastTokenIds, int const* stateSlotMapping, bool removePadding, bool applySilu) -{ - // Reset the parameters - memset(¶ms, 0, sizeof(params)); - - params.batch = batch; - params.dim = dim; - params.max_seqlen = maxSeqLen; - params.dconv = dconv; - params.pre_stride = preStride; - params.post_stride = postStride; - - params.remove_padding = removePadding; - params.apply_silu = applySilu; - - // Set the pointers and strides. - params.in_ptr = const_cast(inPtr); - params.state_in_ptr = const_cast(stateInPtr); - params.state_out_ptr = stateOutPtr; - params.weight_ptr = const_cast(convWeight); - params.bias_ptr = const_cast(convBias); - params.out_ptr = outPtr; - params.last_token_ids_ptr = lastTokenIds; - params.state_slot_mapping_ptr = stateSlotMapping; -} - -template -int MambaConv1dPlugin::enqueueImpl(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) -{ - // inputs - // 0. input_tensor [batch_size, seq_len, dim] or [num_tokens, dim] for remove_input_padding - // 1. conv_state [batch_size, dconv - 1, dim] or host [1] containing only pointer for paged_state - // 2. weight [dim, 1, dconv] - // 3. bias [dim] - // 4. host_request_types [batch_size] int32. 0: context; 1: generation; 2: none. - // 5. last_token_ids [batch_size] int32 - // 6. host_context_lengths [batch_size] int32, optional for remove_input_padding - // 7. state_slot_mapping [batch_size] int32, optional - // outputs - // 0. output_tensor [batch_size, seq_len, dim] or [num_tokens, dim] for remove_input_padding - // 1. conv_state [batch_size, dconv - 1, dim] - auto const batchSize = inputDesc[getHostRequestTypesIdx()].dims.d[0]; - int maxSeqLen; - if (mRemovePadding) - { - int const* host_context_length = static_cast(inputs[getHostContextLengthIdx()]); - maxSeqLen = *std::max_element(host_context_length, host_context_length + batchSize); - } - else - { - maxSeqLen = inputDesc[getInputTensorIdx()].dims.d[1]; - } - - // only support context or generation, not for both of them - RequestType const* reqTypes = static_cast(inputs[getHostRequestTypesIdx()]); - - MambaConv1dParamsBase mambaConv1dParams; - - int const* slotMapping = mPagedState ? static_cast(inputs[getSlotMappingIdx()]) : nullptr; - void* stateInPtr = mPagedState ? *reinterpret_cast(const_cast(inputs[getConvStateIdx()])) - : const_cast(inputs[getConvStateIdx()]); - void* stateOutPtr - = mPagedState ? *reinterpret_cast(const_cast(inputs[getConvStateIdx()])) : outputs[1]; - - setMambaConv1dParams(mambaConv1dParams, batchSize, mDim, maxSeqLen, mDConv, mPreStride, mPostStride, - inputs[getInputTensorIdx()], stateInPtr, stateOutPtr, inputs[getWeightIdx()], inputs[getBiasIdx()], outputs[0], - static_cast(inputs[getLastTokenIdsIdx()]), slotMapping, mRemovePadding, mApplySilu); - - if (reqTypes[0] == RequestType::kCONTEXT) - { - invokeMambaConv1dContext(mambaConv1dParams, stream); - } - else if (reqTypes[0] == RequestType::kGENERATION) - { - invokeMambaConv1dGeneration(mambaConv1dParams, stream); - } - sync_check_cuda_error(stream); - return 0; -} - -int MambaConv1dPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) noexcept -{ - if (isBuilding()) - { - return 0; - } - if (mType == DataType::kHALF) - { - return enqueueImpl(inputDesc, outputDesc, inputs, outputs, workspace, stream); - } - else if (mType == DataType::kFLOAT) - { - return enqueueImpl(inputDesc, outputDesc, inputs, outputs, workspace, stream); - } -#ifdef ENABLE_BF16 - else if (mType == DataType::kBF16) - { - return enqueueImpl<__nv_bfloat16>(inputDesc, outputDesc, inputs, outputs, workspace, stream); - } -#endif - return 0; -} - -// IPluginV2Ext Methods -nvinfer1::DataType MambaConv1dPlugin::getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept -{ - return inputTypes[getInputTensorIdx()]; -} - -// IPluginV2 Methods - -char const* MambaConv1dPlugin::getPluginType() const noexcept -{ - return MAMBA_CONV1D_PLUGIN_NAME; -} - -char const* MambaConv1dPlugin::getPluginVersion() const noexcept -{ - return MAMBA_CONV1D_PLUGIN_VERSION; -} - -int MambaConv1dPlugin::getNbOutputs() const noexcept -{ - return 2; -} - -int MambaConv1dPlugin::initialize() noexcept -{ - return 0; -} - -void MambaConv1dPlugin::terminate() noexcept {} - -size_t MambaConv1dPlugin::getSerializationSize() const noexcept -{ - return sizeof(mDim) + sizeof(mDConv) + sizeof(mPreStride) + sizeof(mPostStride) + sizeof(mType) - + sizeof(mRemovePadding) + sizeof(mPagedState) + sizeof(mApplySilu); -} - -void MambaConv1dPlugin::serialize(void* buffer) const noexcept -{ - char *d = static_cast(buffer), *a = d; - write(d, mDim); - write(d, mDConv); - write(d, mPreStride); - write(d, mPostStride); - write(d, mType); - write(d, mRemovePadding); - write(d, mPagedState); - write(d, mApplySilu); - TLLM_CHECK(d == a + getSerializationSize()); -} - -void MambaConv1dPlugin::destroy() noexcept -{ - delete this; -} - -/////////////// - -MambaConv1dPluginCreator::MambaConv1dPluginCreator() -{ - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - mPluginAttributes.emplace_back(PluginField("dim", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("dconv", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("pre_stride", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("post_stride", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("type_id", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("remove_input_padding", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("paged_state", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("apply_silu", nullptr, PluginFieldType::kINT8)); - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -char const* MambaConv1dPluginCreator::getPluginName() const noexcept -{ - return MAMBA_CONV1D_PLUGIN_NAME; -} - -char const* MambaConv1dPluginCreator::getPluginVersion() const noexcept -{ - return MAMBA_CONV1D_PLUGIN_VERSION; -} - -PluginFieldCollection const* MambaConv1dPluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -IPluginV2* MambaConv1dPluginCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept -{ - PluginField const* fields = fc->fields; - int dim{}; - int dconv{}; - int pre_stride{}; - int post_stride{}; - bool removePadding{}; - bool pagedState{}; - bool applySilu{}; - nvinfer1::DataType type{}; - // Read configurations from each fields - for (int i = 0; i < fc->nbFields; ++i) - { - char const* attrName = fields[i].name; - if (!strcmp(attrName, "dim")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - dim = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "dconv")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - dconv = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "pre_stride")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - pre_stride = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "post_stride")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - post_stride = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "type_id")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - type = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "remove_input_padding")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT8); - removePadding = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "paged_state")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT8); - pagedState = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "apply_silu")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT8); - applySilu = static_cast(*(static_cast(fields[i].data))); - } - } - try - { - auto* obj - = new MambaConv1dPlugin(dim, dconv, pre_stride, post_stride, type, removePadding, pagedState, applySilu); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -IPluginV2* MambaConv1dPluginCreator::deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept -{ - // This object will be deleted when the network is destroyed, which will - // call MambaConv1dPlugin::destroy() - try - { - auto* obj = new MambaConv1dPlugin(serialData, serialLength); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} diff --git a/cpp/tensorrt_llm/plugins/mambaConv1dPlugin/mambaConv1dPlugin.h b/cpp/tensorrt_llm/plugins/mambaConv1dPlugin/mambaConv1dPlugin.h deleted file mode 100644 index d351b1cdc237..000000000000 --- a/cpp/tensorrt_llm/plugins/mambaConv1dPlugin/mambaConv1dPlugin.h +++ /dev/null @@ -1,176 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-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. - */ - -#ifndef TRT_MAMBA_CONV1D_PLUGIN_H -#define TRT_MAMBA_CONV1D_PLUGIN_H -#include "tensorrt_llm/kernels/mambaConv1dKernels.h" -#include "tensorrt_llm/plugins/common/plugin.h" -#include - -namespace tensorrt_llm::plugins -{ -// batch_size = num_ctx_requests or num_gen_requests -// num_ctx_requests = number of context requests (single sequence per request). -// num_gen_requests = number of generation requests (single sequences per request). -// can not support beam search - -// inputs -// 0. input_tensor [batch_size, seq_len, dim] or [num_tokens, dim] for remove_input_padding -// 1. conv_state [batch_size, dconv - 1, dim] or host [1] containing only pointer for paged_state -// 2. weight [1, dconv, dim] -// 3. bias [dim] -// 4. host_request_types [batch_size] int32. 0: context; 1: generation; 2: none. -// 5. last_token_ids [batch_size] int32 -// 6. host_context_lengths [batch_size] int32, optional for remove_input_padding -// 7. state_slot_mapping [batch_size] int32, optional -// outputs -// 0. output_tensor [batch_size, seq_len, dim] or [num_tokens, dim] for remove_input_padding -// 1. conv_state [batch_size, dconv - 1, dim] - -class MambaConv1dPlugin : public BasePlugin -{ -public: - MambaConv1dPlugin(int dim, int dconv, int preStride, int postStride, nvinfer1::DataType type, bool removePadding, - bool pagedState, bool applySilu); - - MambaConv1dPlugin(void const* data, size_t length); - - ~MambaConv1dPlugin() override = default; - - // IPluginV2DynamicExt Methods - nvinfer1::IPluginV2DynamicExt* clone() const noexcept override; - nvinfer1::DimsExprs getOutputDimensions(int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; - bool supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept override; - void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept override; - size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept override; - int enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - template - int enqueueImpl(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream); - - // IPluginV2Ext Methods - nvinfer1::DataType getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept override; - - // IPluginV2 Methods - char const* getPluginType() const noexcept override; - char const* getPluginVersion() const noexcept override; - int getNbOutputs() const noexcept override; - int initialize() noexcept override; - void terminate() noexcept override; - size_t getSerializationSize() const noexcept override; - void serialize(void* buffer) const noexcept override; - void destroy() noexcept override; - - enum class RequestType : int32_t - { - kCONTEXT = 0, - kGENERATION = 1 - }; - -private: - using IndexType = std::int32_t; - - IndexType getInputTensorIdx() const - { - return 0; - }; - - IndexType getConvStateIdx() const - { - return 1; - }; - - IndexType getWeightIdx() const - { - return 2; - }; - - IndexType getBiasIdx() const - { - return 3; - }; - - IndexType getHostRequestTypesIdx() const - { - return 4; - }; - - IndexType getLastTokenIdsIdx() const - { - return 5; - }; - - IndexType getHostContextLengthIdx() const - { - return 6; - }; - - IndexType getSlotMappingIdx() const - { - // if not remove input padding, host_context_length is not used, so the index is 6 - return mRemovePadding ? 7 : 6; - }; - - void setMambaConv1dParams(tensorrt_llm::kernels::MambaConv1dParamsBase& params, - // sizes - const size_t batch, const size_t dim, const size_t maxSeqLen, const size_t dconv, const size_t preStride, - const size_t postStride, - // device pointers - void const* inPtr, void const* stateInPtr, void* stateOutPtr, void const* convWeight, void const* convBias, - void* outPtr, int const* lastTokenIds, int const* stateSlotMapping, bool removePadding, bool applySilu); - -private: - int mDim; - int mDConv; - int mPreStride; - int mPostStride; - nvinfer1::DataType mType; - bool mRemovePadding = false; - bool mPagedState = false; - bool mApplySilu = true; -}; - -class MambaConv1dPluginCreator : public BaseCreator -{ -public: - MambaConv1dPluginCreator(); - - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; - - nvinfer1::IPluginV2* deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept override; - -private: - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; -}; - -} // namespace tensorrt_llm::plugins - -#endif // TRT_MAMBA_CONV1D_PLUGIN_H diff --git a/cpp/tensorrt_llm/plugins/mixtureOfExperts/CMakeLists.txt b/cpp/tensorrt_llm/plugins/mixtureOfExperts/CMakeLists.txt deleted file mode 100644 index 7cc985b60b7a..000000000000 --- a/cpp/tensorrt_llm/plugins/mixtureOfExperts/CMakeLists.txt +++ /dev/null @@ -1,21 +0,0 @@ -# -# SPDX-FileCopyrightText: Copyright (c) 1993-2023 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. -# -file(GLOB SRCS *.cpp) -set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS}) -set(PLUGIN_SOURCES - ${PLUGIN_SOURCES} - PARENT_SCOPE) diff --git a/cpp/tensorrt_llm/plugins/mixtureOfExperts/mixtureOfExpertsPlugin.cpp b/cpp/tensorrt_llm/plugins/mixtureOfExperts/mixtureOfExpertsPlugin.cpp deleted file mode 100644 index ccce34850730..000000000000 --- a/cpp/tensorrt_llm/plugins/mixtureOfExperts/mixtureOfExpertsPlugin.cpp +++ /dev/null @@ -1,1314 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2023 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. - */ -#include "tensorrt_llm/plugins/mixtureOfExperts/mixtureOfExpertsPlugin.h" -#include "tensorrt_llm/common/cudaBf16Wrapper.h" -#include "tensorrt_llm/common/dataType.h" -#include "tensorrt_llm/common/envUtils.h" -#include "tensorrt_llm/common/quantization.h" -#include "tensorrt_llm/runtime/iBuffer.h" -#include "tensorrt_llm/runtime/utils/debugUtils.h" -#include - -using namespace nvinfer1; -using namespace tensorrt_llm::common; -using namespace tensorrt_llm::plugins; -using tensorrt_llm::common::QuantMode; -using tensorrt_llm::common::nextWorkspacePtr; -using tensorrt_llm::common::calculateTotalWorkspaceSize; -using tensorrt_llm::plugins::MixtureOfExpertsPluginCreator; -using tensorrt_llm::plugins::MixtureOfExpertsPlugin; -using tensorrt_llm::plugins::read; -using tensorrt_llm::plugins::write; - -using LoraImpl = tensorrt_llm::kernels::LoraImpl; -using LoraParams = tensorrt_llm::kernels::LoraParams; - -static char const* MIXTURE_OF_EXPERTS_PLUGIN_VERSION{"1"}; -static char const* MIXTURE_OF_EXPERTS_PLUGIN_NAME{"MixtureOfExperts"}; -nvinfer1::PluginFieldCollection MixtureOfExpertsPluginCreator::mFC{}; -std::vector MixtureOfExpertsPluginCreator::mPluginAttributes; - -MixtureOfExpertsPlugin::MixtureOfExpertsPlugin(bool remove_input_padding, int number_of_experts, int experts_per_token, - int expert_hidden_size, int expert_inter_size, int groupwise_quant_algo, int group_size, - ActivationType activation_type, nvinfer1::DataType type, nvinfer1::DataType weight_type, - nvinfer1::DataType output_type, QuantMode quant_mode, bool use_final_scales, bool use_bias, int tp_size, - int tp_rank, int ep_size, int ep_rank, bool force_determinism, int side_stream_id, - MixtureOfExpertsPluginProfilerPtr gemm_profiler_ptr, bool use_lora, nvinfer1::DataType lora_type, - LoraPluginProfilerPtr lora_profiler, int max_low_rank) - : mNumExperts(number_of_experts) - , mExpertsPerToken(experts_per_token) - , mExpertHiddenSize(expert_hidden_size) - , mExpertInterSize(expert_inter_size) - , mGroupwiseQuantAlgo(groupwise_quant_algo) - , mGroupSize(group_size) - , mActivationType(activation_type) - , mType(type) - , mWeightType(weight_type) - , mOutputType(output_type) - , mQuantMode(quant_mode) - , mUseFinalScales(use_final_scales) - , mUseBias(use_bias) - , mParallelismConfig(MOEParallelismConfig{tp_size, tp_rank, ep_size, ep_rank}) - , mUseDeterministicKernels(force_determinism) - , mSideStreamId(side_stream_id) - , mGemmProfiler(std::move(gemm_profiler_ptr)) - , mUseLora(use_lora) - , mLoraType(lora_type) - , mMaxLowRank(max_low_rank) - , mRemoveInputPadding(remove_input_padding) - , mLoraProfiler(std::move(lora_profiler)) -{ - init(); -} - -tensorrt_llm::plugins::MixtureOfExpertsPlugin::MixtureOfExpertsPlugin(MixtureOfExpertsPlugin const& other) - : mMOERunner() - , mNumExperts(other.mNumExperts) - , mExpertsPerToken(other.mExpertsPerToken) - , mExpertHiddenSize(other.mExpertHiddenSize) - , mExpertInterSize(other.mExpertInterSize) - , mGroupwiseQuantAlgo(other.mGroupwiseQuantAlgo) - , mGroupSize(other.mGroupSize) - , mActivationType(other.mActivationType) - , mType(other.mType) - , mWeightType(other.mWeightType) - , mOutputType(other.mOutputType) - , mQuantMode(other.mQuantMode) - , mUseFinalScales(other.mUseFinalScales) - , mUseBias(other.mUseBias) - , mParallelismConfig(other.mParallelismConfig) - , mDims(other.mDims) - , mUseDeterministicKernels(other.mUseDeterministicKernels) - , mSideStreamId(other.mSideStreamId) - , mGemmId1(other.mGemmId1) - , mGemmId2(other.mGemmId2) - , mGemmProfiler(other.mGemmProfiler) - , mUseLora(other.mUseLora) - , mLoraType(other.mLoraType) - , mMaxLowRank(other.mMaxLowRank) - , mRemoveInputPadding(other.mRemoveInputPadding) - , mLoraImpl1(other.mLoraImpl1) - , mLoraImpl2(other.mLoraImpl2) - , mLoraGemmId1(other.mLoraGemmId1) - , mLoraGemmId2(other.mLoraGemmId2) - , mLoraProfiler(other.mLoraProfiler) - , mLayerName(other.mLayerName) - , mNamespace(other.mNamespace) -{ - init(); -} - -size_t MixtureOfExpertsPlugin::getSerializationSize() const noexcept -{ - size_t size = sizeof(mRemoveInputPadding) + sizeof(mNumExperts) + sizeof(mExpertsPerToken) - + sizeof(mExpertHiddenSize) + sizeof(mExpertInterSize) + sizeof(mGroupwiseQuantAlgo) + sizeof(mGroupSize) - + sizeof(mActivationType) + sizeof(mType) + sizeof(mWeightType) + sizeof(mOutputType) - + sizeof(QuantMode::BaseType) + sizeof(mUseFinalScales) + sizeof(mUseBias) + sizeof(mParallelismConfig) - + sizeof(mDims) + sizeof(mUseDeterministicKernels) + sizeof(mSideStreamId) - + mGemmProfiler->getSerializationSize(mGemmId1) + mGemmProfiler->getSerializationSize(mGemmId2) - + sizeof(mUseLora) + sizeof(mLoraType) + sizeof(mMaxLowRank); - - if (hasLora()) - { - size += mLoraProfiler->getSerializationSize(mLoraGemmId1); - size += mLoraProfiler->getSerializationSize(mLoraGemmId2); - } - - return size; -} - -MixtureOfExpertsPlugin::MixtureOfExpertsPlugin(void const* data, size_t length, - MixtureOfExpertsPluginProfilerPtr gemm_profiler_ptr, LoraPluginProfilerPtr lora_profiler) - : mGemmProfiler(gemm_profiler_ptr) - , mLoraProfiler(lora_profiler) -{ - char const* d = reinterpret_cast(data); - char const* a = d; - read(d, mRemoveInputPadding); - read(d, mNumExperts); - read(d, mExpertsPerToken); - read(d, mExpertHiddenSize); - read(d, mExpertInterSize); - read(d, mGroupwiseQuantAlgo); - read(d, mGroupSize); - read(d, mActivationType); - read(d, mType); - read(d, mWeightType); - read(d, mOutputType); - QuantMode::BaseType quant_mode; - read(d, quant_mode); - mQuantMode = QuantMode{quant_mode}; - read(d, mUseFinalScales); - read(d, mUseBias); - read(d, mParallelismConfig); - read(d, mDims); - read(d, mUseDeterministicKernels); - read(d, mSideStreamId); - read(d, mUseLora); - read(d, mLoraType); - read(d, mMaxLowRank); - - // Call init before deserialising the profiler to initialize mGemmId - init(); - mGemmProfiler->deserialize(d, mDims, mGemmId1); - mGemmProfiler->deserialize(d, mDims, mGemmId2); - - if (hasLora()) - { - mLoraProfiler->deserialize(d, mDims, mLoraGemmId1); - mLoraProfiler->deserialize(d, mDims, mLoraGemmId2); - } - - TLLM_CHECK_WITH_INFO(d == a + length, - "Expected length (%d) != real length (%d). This is often " - "caused by using different TensorRT LLM version to build " - "engine and run engine.", - (int) length, (int) (d - a)); -} - -void MixtureOfExpertsPlugin::serialize(void* buffer) const noexcept -{ - char* d = static_cast(buffer); - char* a = d; - - write(d, mRemoveInputPadding); - write(d, mNumExperts); - write(d, mExpertsPerToken); - write(d, mExpertHiddenSize); - write(d, mExpertInterSize); - write(d, mGroupwiseQuantAlgo); - write(d, mGroupSize); - write(d, mActivationType); - write(d, mType); - write(d, mWeightType); - write(d, mOutputType); - write(d, mQuantMode.value()); - write(d, mUseFinalScales); - write(d, mUseBias); - write(d, mParallelismConfig); - write(d, mDims); - write(d, mUseDeterministicKernels); - write(d, mSideStreamId); - write(d, mUseLora); - write(d, mLoraType); - write(d, mMaxLowRank); - - mGemmProfiler->serialize(d, mGemmId1); - mGemmProfiler->serialize(d, mGemmId2); - - if (hasLora()) - { - mLoraProfiler->serialize(d, mLoraGemmId1); - mLoraProfiler->serialize(d, mLoraGemmId2); - } - - TLLM_CHECK(d == a + getSerializationSize()); -} - -template -std::unique_ptr switch_output_type(nvinfer1::DataType output_type) -{ - switch (output_type) - { - case nvinfer1::DataType::kFP4: - case nvinfer1::DataType::kFP8: - // TODO We need an atomic FP8 reduction for the finalize fusions - TLLM_THROW("Outputting %d directly is not currently supported", static_cast(output_type)); - // return std::make_unique>(); - case nvinfer1::DataType::kHALF: - if constexpr (NeedQuant) - { - return std::make_unique>(); - } - else - { - return std::make_unique>(); - } -#ifdef ENABLE_BF16 - case nvinfer1::DataType::kBF16: - if constexpr (NeedQuant) - { - return std::make_unique>(); - } - else - { - return std::make_unique>(); - } -#endif - default: TLLM_THROW("Invalid output type %d", static_cast(output_type)); - } -}; - -void MixtureOfExpertsPlugin::init() -{ - TLLM_CHECK_WITH_INFO(mType == DataType::kFP8 || mType == DataType::kFP4 || mOutputType == mType, - "MOE plugin only supports a different output type for FP4/FP8"); - TLLM_CHECK_WITH_INFO(mType != DataType::kFP8 || tensorrt_llm::common::getSMVersion() >= 89, - "MoE FP8 is not supported for architectures less than SM89"); - TLLM_CHECK_WITH_INFO(mType != DataType::kFP4 || (tensorrt_llm::common::getSMVersion() >= 100), - "MoE FP4 is only supported on architecture SM100 or later"); - - TLLM_CHECK_WITH_INFO(!hasLora() || mLoraType == mOutputType, "The LoraType need to keep same with moe OutputType."); - - if (mWeightType == nvinfer1::DataType::kINT8 && mQuantMode.hasInt4Weights()) - { - mWeightType = DataType::kINT4; - } - - if (mType == DataType::kHALF && mWeightType == DataType::kHALF) - { - mMOERunner = std::make_unique>(); - } - else if (mType == DataType::kFLOAT && mWeightType == DataType::kFLOAT) - { - mMOERunner = std::make_unique>(); - } - else if (mType == DataType::kHALF && mWeightType == DataType::kINT8) - { - mMOERunner = std::make_unique>(); - } - else if (mType == DataType::kHALF && mWeightType == DataType::kINT4) - { - mMOERunner = std::make_unique>(); - } -#ifdef ENABLE_FP8 - else if (mType == DataType::kFP8 && mWeightType == DataType::kINT4 && mOutputType == DataType::kHALF) - { - mMOERunner = std::make_unique>(); - } -#endif -#ifdef ENABLE_BF16 - else if (mType == DataType::kBF16 && mWeightType == DataType::kBF16) - { - mMOERunner = std::make_unique>(); - } - else if (mType == DataType::kBF16 && mWeightType == DataType::kINT8) - { - mMOERunner = std::make_unique>(); - } - else if (mType == DataType::kBF16 && mWeightType == DataType::kINT4) - { - mMOERunner = std::make_unique>(); - } -#ifdef ENABLE_FP8 - else if (mType == DataType::kFP8 && mWeightType == DataType::kINT4 && mOutputType == DataType::kBF16) - { - mMOERunner = std::make_unique< - kernels::CutlassMoeFCRunner<__nv_fp8_e4m3, cutlass::uint4b_t, __nv_bfloat16, __nv_bfloat16>>(); - } -#endif -#endif - -#ifdef ENABLE_FP8 - if (mType == DataType::kFP8 && mWeightType == DataType::kFP8) - { - mMOERunner = switch_output_type<__nv_fp8_e4m3>(mOutputType); - } -#endif -#ifdef ENABLE_FP4 - if (mType == DataType::kFP4 && mWeightType == DataType::kFP4) - { - mMOERunner = switch_output_type<__nv_fp4_e2m1, true>(mOutputType); - } -#endif - - if (!mMOERunner) - { - TLLM_THROW( - "Could not construct the mixture of experts plugin with the requested input combination Activation: %d " - "Weight: %d Output: %d", - static_cast(mType), static_cast(mWeightType), static_cast(mOutputType)); - } - - // Finalize fusion should be disabled if Lora is used. - mMOERunner->use_fused_finalize_ - = (mExpertsPerToken < 3 || !mUseDeterministicKernels) && !getEnvMOEDisableFinalizeFusion() && !hasLora(); - - mGemmId1 = GemmIDMoe{1, mNumExperts, mExpertsPerToken, mParallelismConfig, mExpertHiddenSize, mExpertInterSize, - mGroupSize, mActivationType, mType, mWeightType, mQuantMode, !mMOERunner->use_fused_finalize_}; - mGemmId2 = GemmIDMoe{2, mNumExperts, mExpertsPerToken, mParallelismConfig, mExpertHiddenSize, mExpertInterSize, - mGroupSize, mActivationType, mType, mWeightType, mQuantMode, !mMOERunner->use_fused_finalize_}; - mGemmProfiler->setMaxProfileM(16384 * mNumExperts / mExpertsPerToken); - - if (hasLora()) - { - auto cublasHandle = getCublasHandle(); - auto cublasLtHandle = getCublasLtHandle(); - auto cublasWrapper = std::make_shared(cublasHandle, cublasLtHandle, nullptr, nullptr); - mLoraGemmId1 = GemmIdCublas(mExpertInterSize, mExpertHiddenSize, mLoraType, false, true, mLoraType); - mLoraGemmId2 = GemmIdCublas(mExpertHiddenSize, mExpertInterSize, mLoraType, false, true, mLoraType); - std::vector loraOutSizes1 = {static_cast(mExpertInterSize)}; - mLoraImpl1 = std::make_shared( - mExpertHiddenSize, loraOutSizes1, false, true, 1, mLoraType, mMaxLowRank, cublasWrapper); - std::vector loraOutSizes2 = {static_cast(mExpertHiddenSize)}; - mLoraImpl2 = std::make_shared( - mExpertInterSize, loraOutSizes2, false, true, 1, mLoraType, mMaxLowRank, cublasWrapper); - - TLLM_CUDA_CHECK(cudaEventCreate(&mMemcpyEvent)); - } - mSideStreamPtr = nullptr; - mDebugStallMain = tensorrt_llm::runtime::utils::stallStream("TLLM_DEBUG_MOE_STALL_MAIN"); - mDebugStallSide = tensorrt_llm::runtime::utils::stallStream("TLLM_DEBUG_MOE_STALL_SIDE"); -} - -// IPluginV2DynamicExt Methods -nvinfer1::IPluginV2DynamicExt* MixtureOfExpertsPlugin::clone() const noexcept -{ - auto* plugin = new MixtureOfExpertsPlugin(*this); - plugin->setPluginNamespace(mNamespace.c_str()); - return plugin; -} - -nvinfer1::DimsExprs MixtureOfExpertsPlugin::getOutputDimensions( - int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept -{ - assert(outputIndex == getOutputTensorIndex() || outputIndex == getOutputDummyTensorIndex()); - return inputs[getInputTensorIndex()]; -} - -bool MixtureOfExpertsPlugin::supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept -{ - TLLM_CHECK(0 <= pos && pos < getNbInputs() + getNbOutputs()); - TLLM_CHECK_WITH_INFO( - nbInputs == getNbInputs(), "Required input to plugin is missing. Expected %d Got %d", getNbInputs(), nbInputs); - TLLM_CHECK_WITH_INFO(nbOutputs == getNbOutputs(), "Required output to plugin is missing. Expected %d Got %d", - getNbOutputs(), nbOutputs); - - if (inOut[pos].format != TensorFormat::kLINEAR) - { - return false; - } - - if (pos == getExpertWeights1Index() || pos == getExpertWeights2Index()) - { - if (mGroupwiseQuantAlgo == 0) - { - auto normalized_weight_type - = mWeightType == nvinfer1::DataType::kINT4 ? nvinfer1::DataType::kINT8 : mWeightType; - return inOut[pos].type == normalized_weight_type; - } - else - { - return inOut[pos].type == mOutputType; - } - } - else if (pos == getTokenSelectedExpertsIndex()) - { - return inOut[pos].type == DataType::kINT32; - } - else if (pos == getTokenFinalScalesIndex()) - { - return inOut[pos].type == DataType::kFLOAT; - } - else if (pos == getExpertBias1Index() || pos == getExpertBias2Index()) - { - return inOut[pos].type == mOutputType; - } - else if (pos == nbInputs + getOutputTensorIndex()) - { - return inOut[pos].type == mOutputType; - } - else if (useSideStream() && pos == nbInputs + getOutputDummyTensorIndex()) - { - return inOut[pos].type == inOut[getInputDummyTensorIndex()].type; - } - else if (useSideStream() && pos == getInputDummyTensorIndex()) - { - return true; - } - else if (hasExpertFp8QuantScales() && getExpertFP8Dequant1Index() <= pos && pos <= getExpertFP8QuantFinalIndex()) - { - return inOut[pos].type == DataType::kFLOAT; - } - else if (hasExpertIntQuantScales() && getExpertIntQuantScale1Index() <= pos - && pos <= getExpertIntQuantScale2Index()) - { - return inOut[pos].type == mOutputType; - } - else if (hasFP4QuantScales() && getFP4GlobalActSF1Index() <= pos && pos <= getFP4GlobalSF2Index()) - { - if (pos == getFP4WeightSF1Index() || pos == getFP4WeightSF2Index()) - return inOut[pos].type == nvinfer1::DataType::kFP8; - else - return inOut[pos].type == nvinfer1::DataType::kFLOAT; - } - else if (hasLora() && hasExpertFp8QuantScales() && pos == getInputFP8DequantIndex()) - { - return inOut[pos].type == nvinfer1::DataType::kFLOAT; - } - else if (hasExpertWeightQuantZeros() && getExpertIntQuantZeros1Index() <= pos - && pos <= getExpertIntQuantZeros2Index()) - { - return inOut[pos].type == mOutputType; - } - else if (hasExpertPrequantScales() && getExpertPrequantScales1Index() <= pos - && pos <= getExpertPrequantScales2Index()) - { - return inOut[pos].type == mOutputType; - } - else if (hasGroupwiseFp8Alpha() && getExpertFp8Alpha1Index() <= pos && pos <= getExpertFp8Alpha2Index()) - { - return inOut[pos].type == DataType::kFLOAT; - } - else if (hasLora() && pos == getHostRequestTypeIndex()) - { - return inOut[pos].type == nvinfer1::DataType::kINT32; - } - else if (hasLora() && (pos == getLoraFC1RanksIndex() || pos == getLoraFC2RanksIndex())) - { - return inOut[pos].type == nvinfer1::DataType::kINT32; - } - else if (hasGatedLoraWeightsAndRanks() && pos == getLoraGatedRanksIndex()) - { - return inOut[pos].type == nvinfer1::DataType::kINT32; - } - else if (hasLora() && (pos == getLoraFC1WeightPtrsIndex() || pos == getLoraFC2WeightPtrsIndex())) - { - return inOut[pos].type == nvinfer1::DataType::kINT64; - } - else if (hasGatedLoraWeightsAndRanks() && pos == getLoraGatedWeightPtrsIndex()) - { - return inOut[pos].type == nvinfer1::DataType::kINT64; - } - else if (hasLora() && mRemoveInputPadding && pos == getHostContextLengthIndex()) - { - return inOut[pos].type == nvinfer1::DataType::kINT32; - } - else if ((hasFP4QuantScales() || hasGroupwiseFp8Alpha()) && pos == getInputTensorIndex()) - { - return inOut[pos].type == mOutputType; - } - else - { - return inOut[pos].type == mType; - } - - return false; -} - -void MixtureOfExpertsPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept -{ - TLLM_CHECK_WITH_INFO( - nbInputs == getNbInputs(), "Required input to plugin is missing. Expected %d Got %d", getNbInputs(), nbInputs); - TLLM_CHECK_WITH_INFO(nbOutputs == getNbOutputs(), "Required output to plugin is missing. Expected %d Got %d", - getNbOutputs(), nbOutputs); - - auto in_tensor = in[getInputTensorIndex()]; - - auto const minM - = std::accumulate(in_tensor.min.d, in_tensor.min.d + in_tensor.min.nbDims - 1, 1, std::multiplies()); - auto const maxM - = std::accumulate(in_tensor.max.d, in_tensor.max.d + in_tensor.max.nbDims - 1, 1, std::multiplies()); - - auto weights_1 = in[getExpertWeights1Index()]; - auto weights_2 = in[getExpertWeights2Index()]; - int inner_dim_idx = getGemmShapeInnerDimIndex(); - int const maxK = weights_1.max.d[inner_dim_idx]; - int const maxN = weights_2.max.d[inner_dim_idx]; - int const minK = weights_1.min.d[inner_dim_idx]; - int const minN = weights_2.min.d[inner_dim_idx]; - - TLLM_CHECK_WITH_INFO(minN == maxN, "Variable out channels is not allowed"); - TLLM_CHECK_WITH_INFO(minK == maxK, "Variable in channels is not allowed"); - TLLM_CHECK_WITH_INFO(maxK == mExpertHiddenSize && maxN == mExpertInterSize, - "Configured tensor sizes %dx%d does not match constructor param size %ldx%ld", maxK, maxN, mExpertHiddenSize, - mExpertInterSize); - - if (!mDims.isInitialized()) - { - mDims = {minM, maxM, maxN, maxK}; - } - - mGemmId1 = GemmIDMoe{1, mNumExperts, mExpertsPerToken, mParallelismConfig, mExpertHiddenSize, mExpertInterSize, - mGroupSize, mActivationType, mType, mWeightType, mQuantMode, !mMOERunner->use_fused_finalize_}; - mGemmId2 = GemmIDMoe{2, mNumExperts, mExpertsPerToken, mParallelismConfig, mExpertHiddenSize, mExpertInterSize, - mGroupSize, mActivationType, mType, mWeightType, mQuantMode, !mMOERunner->use_fused_finalize_}; - - if (hasLora()) - { - auto const N = utils::computeNDimension(true, in[getHostRequestTypeIndex()].max); - mLoraGemmId1 = GemmIdCublas(N, mExpertHiddenSize, mLoraType, false, true, mLoraType); - mLoraGemmId2 = GemmIdCublas(N, mExpertInterSize, mLoraType, false, true, mLoraType); - } -} - -auto MixtureOfExpertsPlugin::setupWorkspace(void* base_ptr, int64_t num_tokens, int num_reqs) const -> WorkspaceInfo -{ - size_t moe_workspace_size - = mMOERunner->getWorkspaceSize(num_tokens, mExpertHiddenSize, mExpertInterSize, mNumExperts, mExpertsPerToken, - mActivationType, mParallelismConfig, hasLora(), /*use_deepseek_fp8_block_scale=*/false, - /*min_latency_mode=*/false, hasExpertPrequantScales()); - - // Permutation map - size_t src_to_dest_map_size = mExpertsPerToken * num_tokens * sizeof(int); - - size_t lora_workspace_size = 0; - if (hasLora()) - { - int64_t num_reqs_lora = std::min(num_tokens * mExpertsPerToken, static_cast(num_reqs * mNumExperts)); - lora_workspace_size - = std::max(mLoraImpl1->getWorkspaceSize(num_tokens * mExpertsPerToken, num_reqs_lora, mLoraType), - mLoraImpl2->getWorkspaceSize(num_tokens * mExpertsPerToken, num_reqs_lora, mLoraType)); - } - - std::vector workspaces{ - moe_workspace_size, - src_to_dest_map_size, - lora_workspace_size, - }; - - WorkspaceInfo info{}; - info.size = calculateTotalWorkspaceSize(workspaces.data(), workspaces.size()); - - if (base_ptr) - { - info.workspace = base_ptr; - info.src_to_dest_map = nextWorkspacePtr((int8_t*) info.workspace, moe_workspace_size); - info.lora_workspace = nextWorkspacePtr((int8_t*) info.src_to_dest_map, src_to_dest_map_size); - } - - return info; -} - -int64_t MixtureOfExpertsPlugin::getNumTokens(nvinfer1::PluginTensorDesc const* input_tensors) const -{ - int ndim = input_tensors[getInputTensorIndex()].dims.nbDims; - TLLM_CHECK_WITH_INFO( - 3 == ndim || 2 == ndim, "hidden_state dimension should be either 2 [b*s, hidden], or 3 [b, s, hidden]"); - int64_t num_tokens = input_tensors[getInputTensorIndex()].dims.d[0]; - if (ndim == 3) - { - num_tokens *= input_tensors[getInputTensorIndex()].dims.d[1]; - } - return num_tokens; -} - -size_t MixtureOfExpertsPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - TLLM_CHECK_WITH_INFO( - nbInputs == getNbInputs(), "Required input to plugin is missing. Expected %d Got %d", getNbInputs(), nbInputs); - TLLM_CHECK_WITH_INFO(nbOutputs == getNbOutputs(), "Required output to plugin is missing. Expected %d Got %d", - getNbOutputs(), nbOutputs); - - if (useSideStream()) - { - return 0; - } - int const num_tokens = getNumTokens(inputs); - int const num_lora_reqs = getNumLoraRequests(inputs); - return setupWorkspace(nullptr, num_tokens, num_lora_reqs).size; -} - -MOEParallelismConfig MixtureOfExpertsPlugin::getParallelismConfig() const -{ - return mParallelismConfig; -} - -QuantParams tensorrt_llm::plugins::MixtureOfExpertsPlugin::getQuantParams(nvinfer1::PluginTensorDesc const* inputDesc, - void const* const* inputs, int scale_1_idx, int scale_2_idx, int scale_3_idx, int scale_4_idx, int scale_5_idx, - int scale_6_idx, int scale_7_idx, int scale_8_idx) const -{ - void const* scale_1 = scale_1_idx >= 0 ? inputs[scale_1_idx] : nullptr; - void const* scale_2 = scale_2_idx >= 0 ? inputs[scale_2_idx] : nullptr; - void const* scale_3 = scale_3_idx >= 0 ? inputs[scale_3_idx] : nullptr; - void const* scale_4 = scale_4_idx >= 0 ? inputs[scale_4_idx] : nullptr; - void const* scale_5 = scale_5_idx >= 0 ? inputs[scale_5_idx] : nullptr; - void const* scale_6 = scale_6_idx >= 0 ? inputs[scale_6_idx] : nullptr; - void const* scale_7 = scale_7_idx >= 0 ? inputs[scale_7_idx] : nullptr; - void const* scale_8 = scale_8_idx >= 0 ? inputs[scale_8_idx] : nullptr; - nvinfer1::PluginTensorDesc const* desc_1 = scale_1_idx >= 0 ? &inputDesc[scale_1_idx] : nullptr; - nvinfer1::PluginTensorDesc const* desc_2 = scale_2_idx >= 0 ? &inputDesc[scale_2_idx] : nullptr; - nvinfer1::PluginTensorDesc const* desc_3 = scale_3_idx >= 0 ? &inputDesc[scale_3_idx] : nullptr; - nvinfer1::PluginTensorDesc const* desc_4 = scale_4_idx >= 0 ? &inputDesc[scale_4_idx] : nullptr; - nvinfer1::PluginTensorDesc const* desc_5 = scale_5_idx >= 0 ? &inputDesc[scale_5_idx] : nullptr; - nvinfer1::PluginTensorDesc const* desc_6 = scale_6_idx >= 0 ? &inputDesc[scale_6_idx] : nullptr; - auto const gated_inter_size = isGatedActivation(mActivationType) ? mExpertInterSize * 2 : mExpertInterSize; - auto const experts_per_node = mNumExperts / mParallelismConfig.ep_size; - if (hasExpertIntQuantScales()) - { - TLLM_CHECK(scale_1 && scale_2); - if (!hasGroupwiseIntQuantScales()) - { - TLLM_CHECK(!scale_3 && !scale_4 && !scale_5 && !scale_6); - TLLM_CHECK(desc_1->dims.nbDims == 2); - TLLM_CHECK(desc_2->dims.nbDims == 2); - TLLM_CHECK_WITH_INFO( - desc_1->dims.d[0] == experts_per_node, "Incorrect number of experts in int quant scale"); - TLLM_CHECK(desc_1->dims.d[1] == gated_inter_size); - TLLM_CHECK_WITH_INFO( - desc_2->dims.d[0] == experts_per_node, "Incorrect number of experts in int quant scale"); - TLLM_CHECK(desc_2->dims.d[1] == mExpertHiddenSize); - return QuantParams::Int(scale_1, scale_2); - } - else - { - TLLM_CHECK(desc_1->dims.nbDims == 3); - TLLM_CHECK(desc_2->dims.nbDims == 3); - TLLM_CHECK((scale_3 && scale_4) || !hasExpertPrequantScales()); - TLLM_CHECK((scale_5 && scale_6) || !hasExpertWeightQuantZeros()); - TLLM_CHECK((scale_7 && scale_8) || !hasGroupwiseFp8Alpha()); - return QuantParams::GroupWise(mGroupSize, scale_1, scale_2, scale_3, scale_4, scale_5, scale_6, - static_cast(scale_7), static_cast(scale_8)); - } - } - else if (hasExpertFp8QuantScales()) - { - TLLM_CHECK(scale_1 && scale_2 && scale_3); - TLLM_CHECK(scale_4 || !hasExpertFp8FinalQuantScales()); - TLLM_CHECK((scale_5 != nullptr) == hasLora()); - TLLM_CHECK(!scale_6); - TLLM_CHECK(desc_1->dims.nbDims == 2); - TLLM_CHECK(desc_2->dims.nbDims == 1); - TLLM_CHECK(desc_3->dims.nbDims == 2); - TLLM_CHECK_WITH_INFO( - desc_1->dims.d[0] == experts_per_node && desc_1->dims.d[1] == 1, "Incorrect shape for weight FP8 scale"); - TLLM_CHECK(desc_2->dims.d[0] == 1); - TLLM_CHECK_WITH_INFO( - desc_3->dims.d[0] == experts_per_node && desc_3->dims.d[1] == 1, "Incorrect shape for weight FP8 scale"); - return QuantParams::FP8(static_cast(scale_1), static_cast(scale_2), - static_cast(scale_3), static_cast(scale_4), static_cast(scale_5)); - } - else if (hasFP4QuantScales()) - { - TLLM_CHECK(scale_1 && scale_2 && scale_3 && scale_4 && scale_5 && scale_6); - TLLM_CHECK(desc_1->dims.nbDims == 1); - TLLM_CHECK(desc_2->dims.nbDims == 3); - TLLM_CHECK(desc_3->dims.nbDims == 1); - TLLM_CHECK(desc_4->dims.nbDims == 1); - TLLM_CHECK(desc_5->dims.nbDims == 3); - TLLM_CHECK(desc_6->dims.nbDims == 1); - TLLM_CHECK(desc_1->dims.d[0] == 1); - TLLM_CHECK_WITH_INFO(desc_2->dims.d[0] == experts_per_node && desc_2->dims.d[1] == gated_inter_size - && desc_2->dims.d[2] - == mExpertHiddenSize / TmaWarpSpecializedGroupedGemmInput::NVFP4BlockScaleVectorSize, - "Incorrect shape for FP4 scale"); - TLLM_CHECK_WITH_INFO(desc_3->dims.d[0] == experts_per_node, "Incorrect shape for FP4 scale"); - TLLM_CHECK(desc_4->dims.d[0] == 1); - TLLM_CHECK_WITH_INFO(desc_5->dims.d[0] == experts_per_node && desc_5->dims.d[1] == mExpertHiddenSize - && desc_5->dims.d[2] - == mExpertInterSize / TmaWarpSpecializedGroupedGemmInput::NVFP4BlockScaleVectorSize, - "Incorrect shape for FP4 scale"); - TLLM_CHECK_WITH_INFO(desc_6->dims.d[0] == experts_per_node, "Incorrect shape for FP4 scale"); - return QuantParams::FP4(static_cast(scale_1), - static_cast(scale_2), - static_cast(scale_3), static_cast(scale_4), - static_cast(scale_5), - static_cast(scale_6)); - } - return {}; -} - -int MixtureOfExpertsPlugin::getNumLoraRequests(nvinfer1::PluginTensorDesc const* input_tensors) const -{ - if (!hasLora()) - return 0; - int num_reqs = input_tensors[getLoraFC1RanksIndex()].dims.d[0]; - return num_reqs; -} - -LoraParams MixtureOfExpertsPlugin::getLoraParams( - nvinfer1::PluginTensorDesc const* inputDesc, void const* const* inputs, void* workspace) -{ - TLLM_CHECK(hasLora()); - - int const num_reqs = getNumLoraRequests(inputDesc); - int64_t const num_tokens = getNumTokens(inputDesc); - bool is_gated_actiation = isGatedActivation(mActivationType); - - mLoraExpandFC1WeightPtrs.clear(); - mLoraExpandFC2WeightPtrs.clear(); - mLoraExpandFC1Ranks.clear(); - mLoraExpandFC2Ranks.clear(); - - mLoraExpandFC1WeightPtrs.reserve(num_tokens * 2); - mLoraExpandFC2WeightPtrs.reserve(num_tokens * 2); - mLoraExpandFC1Ranks.reserve(num_tokens); - mLoraExpandFC2Ranks.reserve(num_tokens); - - if (is_gated_actiation) - { - mLoraExpandGatedWeightPtrs.clear(); - mLoraExpandGatedRanks.clear(); - mLoraExpandGatedWeightPtrs.reserve(num_tokens * 2); - mLoraExpandGatedRanks.reserve(num_tokens); - } - - int const seq_len = mRemoveInputPadding ? 0 : inputDesc[getInputTensorIndex()].dims.d[1]; - int32_t const* req_types = static_cast(inputs[getHostRequestTypeIndex()]); - int32_t const* host_context_lens - = mRemoveInputPadding ? static_cast(inputs[getHostContextLengthIndex()]) : nullptr; - - auto const fc1_lora_weight_ptrs = static_cast(inputs[getLoraFC1WeightPtrsIndex()]); - auto const fc1_lora_ranks = static_cast(inputs[getLoraFC1RanksIndex()]); - - auto const fc2_lora_weight_ptrs = static_cast(inputs[getLoraFC2WeightPtrsIndex()]); - auto const fc2_lora_ranks = static_cast(inputs[getLoraFC2RanksIndex()]); - - auto const gated_lora_weight_ptrs - = is_gated_actiation ? static_cast(inputs[getLoraGatedWeightPtrsIndex()]) : nullptr; - auto const gated_lora_ranks - = is_gated_actiation ? static_cast(inputs[getLoraGatedRanksIndex()]) : nullptr; - - int idx = 0; - for (int req_id = 0; req_id < num_reqs; req_id++) - { - RequestType const reqType = static_cast(req_types[req_id]); - if (reqType == RequestType::kGENERATION) - { - // lora_weight_ptrs has 3 pointers for each module: A,B, and an optional DoRA magnitude - // the current LoRA implementation does not apply DoRA scaling, so the magnitude is ignored - mLoraExpandFC1WeightPtrs.push_back(fc1_lora_weight_ptrs[req_id * 3]); - mLoraExpandFC1WeightPtrs.push_back(fc1_lora_weight_ptrs[req_id * 3 + 1]); - mLoraExpandFC1Ranks.push_back(fc1_lora_ranks[req_id]); - - mLoraExpandFC2WeightPtrs.push_back(fc2_lora_weight_ptrs[req_id * 3]); - mLoraExpandFC2WeightPtrs.push_back(fc2_lora_weight_ptrs[req_id * 3 + 1]); - mLoraExpandFC2Ranks.push_back(fc2_lora_ranks[req_id]); - - if (is_gated_actiation) - { - mLoraExpandGatedWeightPtrs.push_back(gated_lora_weight_ptrs[req_id * 3]); - mLoraExpandGatedWeightPtrs.push_back(gated_lora_weight_ptrs[req_id * 3 + 1]); - mLoraExpandGatedRanks.push_back(gated_lora_ranks[req_id]); - } - - idx += 1; - } - else - { - int context_len = (mRemoveInputPadding ? host_context_lens[req_id] : seq_len); - - for (int context_id = 0; context_id < context_len; context_id++) - { - mLoraExpandFC1WeightPtrs.push_back(fc1_lora_weight_ptrs[req_id * 3]); - mLoraExpandFC1WeightPtrs.push_back(fc1_lora_weight_ptrs[req_id * 3 + 1]); - mLoraExpandFC1Ranks.push_back(fc1_lora_ranks[req_id]); - - mLoraExpandFC2WeightPtrs.push_back(fc2_lora_weight_ptrs[req_id * 3]); - mLoraExpandFC2WeightPtrs.push_back(fc2_lora_weight_ptrs[req_id * 3 + 1]); - mLoraExpandFC2Ranks.push_back(fc2_lora_ranks[req_id]); - - if (is_gated_actiation) - { - mLoraExpandGatedWeightPtrs.push_back(gated_lora_weight_ptrs[req_id * 3]); - mLoraExpandGatedWeightPtrs.push_back(gated_lora_weight_ptrs[req_id * 3 + 1]); - mLoraExpandGatedRanks.push_back(gated_lora_ranks[req_id]); - } - } - idx += context_len; - } - } - - TLLM_CHECK_WITH_INFO(idx == num_tokens, fmtstr("idx %d num_tokens %ld", idx, num_tokens)); - - return LoraParams(num_reqs, mLoraExpandFC1Ranks.data(), mLoraExpandFC1WeightPtrs.data(), mLoraExpandFC2Ranks.data(), - mLoraExpandFC2WeightPtrs.data(), mLoraImpl1, mLoraImpl2, workspace, &mMemcpyEvent, mLoraExpandGatedRanks.data(), - mLoraExpandGatedWeightPtrs.data()); -} - -int MixtureOfExpertsPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace_ptr, - cudaStream_t stream) noexcept -{ - if (isBuilding()) - { - return 0; - } - - int64_t const num_tokens = getNumTokens(inputDesc); - int64_t const num_reqs = getNumLoraRequests(inputDesc); - - if (useSideStream()) - { - // Prepare the side stream - if (!mSideStreamPtr) - { - auto const resource_name = nvinfer1::pluginInternal::SideStream::getResourceKey(mSideStreamId); - nvinfer1::pluginInternal::SideStream side_stream{}; - mSideStreamPtr = reinterpret_cast( - getPluginRegistry()->acquirePluginResource(resource_name.c_str(), &side_stream)); - } - // Debug the code with the main stream stalled (only executed when the environment variable - // TLLM_DEBUG_MOE_STALL_MAIN is set and has a positive value) - mSideStreamPtr->stallMainStream("TLLM_DEBUG_MOE_STALL_MAIN", stream, mDebugStallMain); - // The side stream waits for the inputs managed by the main stream to be ready - mSideStreamPtr->waitMainStreamOnSideStream(stream); - // Provide data dependency for the shared experts running after this plugin by copying inputs on the main stream - size_t count = 1; - for (int i = 0; i < inputDesc[getInputDummyTensorIndex()].dims.nbDims; ++i) - { - count *= inputDesc[getInputDummyTensorIndex()].dims.d[i]; - } - count *= tensorrt_llm::runtime::BufferDataType(inputDesc[getInputDummyTensorIndex()].type).getSize(); - TLLM_CUDA_CHECK(cudaMemcpyAsync(outputs[getOutputDummyTensorIndex()], inputs[getInputDummyTensorIndex()], count, - cudaMemcpyDeviceToDevice, stream)); - // Switch from the main stream to the side stream - stream = mSideStreamPtr->getStream(); - // The workspace is managed by the side stream (otherwise, the lifetime of workspace may be incorrect) - auto const workspace_size = setupWorkspace(nullptr, num_tokens, num_reqs).size; - workspace_ptr = mSideStreamPtr->getWorkspacePtr(workspace_size); - } - auto workspace = setupWorkspace(workspace_ptr, num_tokens, num_reqs); - - auto w1_desc = inputDesc[getExpertWeights1Index()]; - auto w2_desc = inputDesc[getExpertWeights2Index()]; - TLLM_CHECK(w1_desc.dims.nbDims == 3); - auto const experts_per_node = mNumExperts / mParallelismConfig.ep_size; - TLLM_CHECK(w1_desc.dims.d[0] == experts_per_node); - TLLM_CHECK(w2_desc.dims.nbDims == 3); - TLLM_CHECK(w2_desc.dims.d[0] == experts_per_node); - - auto [inner_packed_elements, outer_packed_elements] = getWeightPackedElements(); - int inner_dim_idx = getGemmShapeInnerDimIndex(); - int outer_dim_idx = getGemmShapeOuterDimIndex(); - TLLM_CHECK(w1_desc.dims.d[inner_dim_idx] * inner_packed_elements == mExpertHiddenSize); - if (isGatedActivation(mActivationType)) - { - TLLM_CHECK(w1_desc.dims.d[outer_dim_idx] * outer_packed_elements == mExpertInterSize * 2); - } - else - { - TLLM_CHECK(w1_desc.dims.d[outer_dim_idx] * outer_packed_elements == mExpertInterSize); - } - - TLLM_CHECK(w2_desc.dims.d[inner_dim_idx] * inner_packed_elements == mExpertInterSize); - TLLM_CHECK(w2_desc.dims.d[outer_dim_idx] * outer_packed_elements == mExpertHiddenSize); - - QuantParams quant_params{}; - if (hasExpertIntQuantScales()) - { - if (mGroupSize > 0) - { - quant_params = getQuantParams(inputDesc, inputs, getExpertIntQuantScale1Index(), - getExpertIntQuantScale2Index(), hasExpertPrequantScales() ? getExpertPrequantScales1Index() : -1, - hasExpertPrequantScales() ? getExpertPrequantScales2Index() : -1, - hasExpertWeightQuantZeros() ? getExpertIntQuantZeros1Index() : -1, - hasExpertWeightQuantZeros() ? getExpertIntQuantZeros2Index() : -1, - hasGroupwiseFp8Alpha() ? getExpertFp8Alpha1Index() : -1, - hasGroupwiseFp8Alpha() ? getExpertFp8Alpha2Index() : -1); - } - else - { - quant_params - = getQuantParams(inputDesc, inputs, getExpertIntQuantScale1Index(), getExpertIntQuantScale2Index()); - } - } - else if (hasExpertFp8QuantScales()) - { - quant_params = getQuantParams(inputDesc, inputs, // - getExpertFP8Dequant1Index(), // - getExpertFP8Quant2Index(), // - getExpertFP8Dequant2Index(), // - hasExpertFp8FinalQuantScales() ? getExpertFP8QuantFinalIndex() : -1, - hasLora() ? getInputFP8DequantIndex() : -1); - } - else if (hasFP4QuantScales()) - { - quant_params = getQuantParams(inputDesc, inputs, // - getFP4GlobalActSF1Index(), // - getFP4WeightSF1Index(), // - getFP4GlobalSF1Index(), // - getFP4GlobalActSF2Index(), // - getFP4WeightSF2Index(), // - getFP4GlobalSF2Index() // - ); - } - - LoraParams lora_params{}; - - if (hasLora()) - { - lora_params = getLoraParams(inputDesc, inputs, workspace.lora_workspace); - auto lora_gemm1 = mLoraProfiler->getBestConfig(num_tokens, mLoraGemmId1); - auto lora_gemm2 = mLoraProfiler->getBestConfig(num_tokens, mLoraGemmId2); - - mLoraImpl1->setBestTactic(lora_gemm1); - mLoraImpl2->setBestTactic(lora_gemm2); - } - - std::optional gemm1; - std::optional gemm2; - if (common::getEnvForceDeterministicMOE()) - { - gemm1 = mMOERunner->getTactics(MoeGemmId::GEMM_1)[0]; - gemm2 = mMOERunner->getTactics(MoeGemmId::GEMM_2)[0]; - } - else - { - gemm1 = mGemmProfiler->getBestConfig(num_tokens, mGemmId1); - gemm2 = mGemmProfiler->getBestConfig(num_tokens, mGemmId2); - } - - MoeMinLatencyParams min_latency_params{}; - mMOERunner->setTactic(gemm1, gemm2); -#ifdef USING_OSS_CUTLASS_MOE_GEMM - mMOERunner->runMoe(inputs[getInputTensorIndex()], nullptr, true, - static_cast(inputs[getTokenSelectedExpertsIndex()]), - hasFinalScales() ? static_cast(inputs[getTokenFinalScalesIndex()]) : nullptr, - inputs[getExpertWeights1Index()], hasBias() ? inputs[getExpertBias1Index()] : nullptr, - ActivationParams(mActivationType), inputs[getExpertWeights2Index()], - hasBias() ? inputs[getExpertBias2Index()] : nullptr, quant_params, num_tokens, num_tokens, mExpertHiddenSize, - mExpertHiddenSize /*TRT does not support padding, safe to assume padded/unpadded hidden sizes are the same*/, - mExpertInterSize, mNumExperts, mExpertsPerToken, static_cast(workspace.workspace), - // Outputs - outputs[getOutputTensorIndex()], static_cast(workspace.src_to_dest_map), mParallelismConfig, - /*enable_alltoall=*/false, hasLora(), lora_params, /*use_deepseek_fp8_block_scale=*/false, - /*min_latency_mode=*/false, min_latency_params, stream); -#else - mMOERunner->runMoe(inputs[getInputTensorIndex()], nullptr, true, - static_cast(inputs[getTokenSelectedExpertsIndex()]), - hasFinalScales() ? static_cast(inputs[getTokenFinalScalesIndex()]) : nullptr, - inputs[getExpertWeights1Index()], hasBias() ? inputs[getExpertBias1Index()] : nullptr, - ActivationParams(mActivationType), inputs[getExpertWeights2Index()], - hasBias() ? inputs[getExpertBias2Index()] : nullptr, quant_params, num_tokens, num_tokens, mExpertHiddenSize, - mExpertInterSize, mNumExperts, mExpertsPerToken, static_cast(workspace.workspace), - // Outputs - outputs[getOutputTensorIndex()], static_cast(workspace.src_to_dest_map), mParallelismConfig, hasLora(), - lora_params, /*use_deepseek_fp8_block_scale=*/false, - /*min_latency_mode=*/false, min_latency_params, stream); -#endif - - if (useSideStream()) - { - // Debug the code with the side stream stalled (only executed when the environment variable - // TLLM_DEBUG_MOE_STALL_SIDE is set and has a positive value) - mSideStreamPtr->stallSideStream("TLLM_DEBUG_MOE_STALL_SIDE", mDebugStallSide); - } - - return 0; -} - -// IPluginV2Ext Methods -nvinfer1::DataType MixtureOfExpertsPlugin::getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept -{ - TLLM_CHECK(index == getOutputTensorIndex() || index == getOutputDummyTensorIndex()); - if (useSideStream() && index == getOutputDummyTensorIndex()) - { - return inputTypes[getInputDummyTensorIndex()]; - } - return mOutputType; -} - -// IPluginV2 Methods -char const* MixtureOfExpertsPlugin::getPluginType() const noexcept -{ - return MIXTURE_OF_EXPERTS_PLUGIN_NAME; -} - -char const* MixtureOfExpertsPlugin::getPluginVersion() const noexcept -{ - return MIXTURE_OF_EXPERTS_PLUGIN_VERSION; -} - -int MixtureOfExpertsPlugin::initialize() noexcept -{ - mGemmProfiler->setGemmToProfile(kernels::GemmProfilerBackend::GemmToProfile::GEMM_1); - mGemmProfiler->profileTactics(this, mType, mDims, mGemmId1); - mGemmProfiler->setGemmToProfile(kernels::GemmProfilerBackend::GemmToProfile::GEMM_2); - mGemmProfiler->profileTactics(this, mType, mDims, mGemmId2); - - if (hasLora()) - { - mLoraImpl1->setGemmConfig(); - mLoraImpl2->setGemmConfig(); - - mLoraProfiler->profileTactics(mLoraImpl1->getCublasWrapper(), mType, mDims, mLoraGemmId1); - mLoraProfiler->profileTactics(mLoraImpl2->getCublasWrapper(), mType, mDims, mLoraGemmId2); - } - return 0; -} - -void MixtureOfExpertsPlugin::terminate() noexcept -{ - if (mSideStreamPtr) - { - auto const resource_name = nvinfer1::pluginInternal::SideStream::getResourceKey(mSideStreamId); - getPluginRegistry()->releasePluginResource(resource_name.c_str()); - mSideStreamPtr = nullptr; - } -} - -void MixtureOfExpertsPlugin::destroy() noexcept -{ - if (hasLora()) - { - TLLM_CUDA_CHECK(cudaEventDestroy(mMemcpyEvent)); - } - // This gets called when the network containing plugin is destroyed - delete this; -} - -void MixtureOfExpertsPlugin::setPluginNamespace(char const* libNamespace) noexcept -{ - mNamespace = libNamespace; -} - -char const* MixtureOfExpertsPlugin::getPluginNamespace() const noexcept -{ - return mNamespace.c_str(); -} - -/////////////// - -char const* MixtureOfExpertsPluginCreator::getPluginName() const noexcept -{ - return MIXTURE_OF_EXPERTS_PLUGIN_NAME; -} - -char const* MixtureOfExpertsPluginCreator::getPluginVersion() const noexcept -{ - return MIXTURE_OF_EXPERTS_PLUGIN_VERSION; -} - -nvinfer1::PluginFieldCollection const* MixtureOfExpertsPluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -MixtureOfExpertsPluginCreator::MixtureOfExpertsPluginCreator() -{ - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - mPluginAttributes.emplace_back(nvinfer1::PluginField("remove_input_padding", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(nvinfer1::PluginField("number_of_experts", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(nvinfer1::PluginField("experts_per_token", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(nvinfer1::PluginField("expert_hidden_size", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(nvinfer1::PluginField("expert_inter_size", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(nvinfer1::PluginField("groupwise_quant_algo", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(nvinfer1::PluginField("group_size", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(nvinfer1::PluginField("activation_type", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(nvinfer1::PluginField("type_id", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(nvinfer1::PluginField("weight_type_id", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(nvinfer1::PluginField("quant_mode", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(nvinfer1::PluginField("use_final_scales", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(nvinfer1::PluginField("use_bias", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(nvinfer1::PluginField("tp_size", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(nvinfer1::PluginField("tp_rank", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(nvinfer1::PluginField("ep_size", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(nvinfer1::PluginField("ep_rank", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(nvinfer1::PluginField("side_stream_id", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(nvinfer1::PluginField("use_lora", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(nvinfer1::PluginField("lora_type_id", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(nvinfer1::PluginField("max_low_rank", nullptr, PluginFieldType::kINT32)); - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -IPluginV2* MixtureOfExpertsPluginCreator::createPlugin( - char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept -{ - nvinfer1::PluginField const* fields = fc->fields; - int mRemoveInputPadding{}; - int mNumExperts{}; - int mExpertsPerToken{}; - int mExpertHiddenSize{}; - int mExpertInterSize{}; - int mGroupwiseQuantAlgo{}; - int mGroupSize{}; - int mActivationType{}; - int mType{}; - int mWeightType{}; - int mOutputType{INT_MAX}; - int mQuantMode{}; - int mUseFinalScales{1}; // Default to true - int mUseBias{0}; - int mTPSize{}; - int mTPRank{}; - int mEPSize{}; - int mEPRank{}; - int mRequiresDeterminism{0}; - int mSideStreamId{0}; - int mUseLora{}; - int mLoraType{INT_MAX}; - int mMaxLowRank{0}; - - // Read configurations from each fields - struct MapPair - { - char const* key; - int& field; - bool optional = false; - bool set = false; - }; - - std::array input_map{ - MapPair{"remove_input_padding", std::ref(mRemoveInputPadding)}, - MapPair{"number_of_experts", std::ref(mNumExperts)}, - MapPair{"experts_per_token", std::ref(mExpertsPerToken)}, - MapPair{"expert_hidden_size", std::ref(mExpertHiddenSize)}, - MapPair{"expert_inter_size", std::ref(mExpertInterSize)}, - MapPair{"groupwise_quant_algo", std::ref(mGroupwiseQuantAlgo)}, - MapPair{"group_size", std::ref(mGroupSize)}, - MapPair{"activation_type", std::ref(mActivationType)}, - MapPair{"type_id", std::ref(mType)}, - MapPair{"weight_type_id", std::ref(mWeightType)}, - MapPair{"quant_mode", std::ref(mQuantMode)}, - MapPair{"tp_size", std::ref(mTPSize)}, - MapPair{"tp_rank", std::ref(mTPRank)}, - MapPair{"ep_size", std::ref(mEPSize)}, - MapPair{"ep_rank", std::ref(mEPRank)}, - MapPair{"use_lora", std::ref(mUseLora)}, - MapPair{"use_final_scales", std::ref(mUseFinalScales)}, - - // Optional - MapPair{"use_bias", std::ref(mUseBias), true}, - MapPair{"output_type_id", std::ref(mOutputType), true}, - MapPair{"force_determinism", std::ref(mRequiresDeterminism), true}, - MapPair{"side_stream_id", std::ref(mSideStreamId), true}, - MapPair{"lora_type_id", std::ref(mLoraType), true}, - MapPair{"max_low_rank", std::ref(mMaxLowRank), true}, - }; - for (int i = 0; i < fc->nbFields; ++i) - { - char const* attrName = fields[i].name; - for (auto& item : input_map) - { - if (!strcmp(item.key, attrName)) - { - TLLM_CHECK(fields[i].type == nvinfer1::PluginFieldType::kINT32); - TLLM_CHECK_WITH_INFO(!item.set, "Parameter %s was set twice", item.key); - item.field = static_cast(*(static_cast(fields[i].data))); - item.set = true; - } - } - } - - for (auto& item : input_map) - { - TLLM_CHECK_WITH_INFO(item.set || item.optional, "Parameter %s is required but not set", item.key); - } - - // Output type is optional, if not set it to the same as mType - if (mOutputType == INT_MAX) - { - mOutputType = mType; - } - - if (mUseLora) - { - TLLM_CHECK_WITH_INFO(mLoraType != INT_MAX && mMaxLowRank != 0, - "MoE fuse lora, lora_type_id and max_low_rank are required but not set"); - } - - try - { - auto gemmProfiler = moePluginProfiler.createGemmPluginProfiler(/* inference */ false); - auto loraProfiler = loraPluginProfileManager.createGemmPluginProfiler(/* inference */ false, /* skip */ true); - auto* obj = new MixtureOfExpertsPlugin( - // Constructor parameters - mRemoveInputPadding, mNumExperts, mExpertsPerToken, mExpertHiddenSize, mExpertInterSize, - mGroupwiseQuantAlgo, mGroupSize, static_cast(mActivationType), - static_cast(mType), static_cast(mWeightType), - static_cast(mOutputType), QuantMode(mQuantMode), mUseFinalScales != 0, mUseBias != 0, - mTPSize, mTPRank, mEPSize, mEPRank, mRequiresDeterminism != 0, mSideStreamId, gemmProfiler, mUseLora != 0, - static_cast(mLoraType), loraProfiler, mMaxLowRank); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -IPluginV2* MixtureOfExpertsPluginCreator::deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept -{ - // This object will be deleted when the network is destroyed, which will - // call MixtureOfExpertsPlugin::destroy() - try - { - auto gemmProfiler = moePluginProfiler.createGemmPluginProfiler(/* inference */ true); - auto loraProfiler = loraPluginProfileManager.createGemmPluginProfiler(/* inference */ false, /* skip */ true); - - auto* obj = new MixtureOfExpertsPlugin( - // Constructor parameters - serialData, serialLength, gemmProfiler, loraProfiler); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -void MixtureOfExpertsPluginCreator::setPluginNamespace(char const* libNamespace) noexcept -{ - mNamespace = libNamespace; -} - -char const* MixtureOfExpertsPluginCreator::getPluginNamespace() const noexcept -{ - return mNamespace.c_str(); -} - -void MixtureOfExpertsGemmProfiler::computeTmpSize(size_t maxM, size_t n, size_t k) -{ - checkInit(); - size_t bytes = backend.getWorkspaceSize(maxM); - this->setTmpWorkspaceSizeInBytes(bytes); -} - -void MixtureOfExpertsGemmProfiler::runTactic(int m, int n, int k, MixtureOfExpertsGemmProfiler::Config const& tactic, - char* workspace_ptr_char, cudaStream_t const& stream) -{ - checkInit(); - backend.runProfiler(m, tactic, workspace_ptr_char, /*expert_weights*/ nullptr, stream); -} - -auto MixtureOfExpertsGemmProfiler::getTactics(int m, int n, int k) const -> std::vector -{ - assert(mRunner); - return mRunner->mMOERunner->getTactics(backend.mGemmToProfile); -} - -void MixtureOfExpertsGemmProfiler::initTmpData( - int m, int n, int k, char* workspace, size_t ws_size, cudaStream_t stream) -{ - checkInit(); - backend.prepare(m, workspace, /*expert_weights*/ nullptr, stream); -} - -void MixtureOfExpertsGemmProfiler::checkInit() -{ - assert(mRunner); - if (init_backend) - { - return; - } - init_backend = true; - auto& plugin = *mRunner; -#ifdef USING_OSS_CUTLASS_MOE_GEMM - backend.init(*plugin.mMOERunner, backend.mGemmToProfile, plugin.mType, plugin.mWeightType, plugin.mOutputType, - plugin.mNumExperts, plugin.mExpertsPerToken, plugin.mExpertHiddenSize, - plugin.mExpertHiddenSize /*TRT backend does not support unpadded hidden size*/, plugin.mExpertInterSize, - plugin.mGroupSize, plugin.mActivationType, plugin.hasBias(), plugin.hasLora(), /*min_latency_mode=*/false, - /*need_weights=*/true, plugin.getParallelismConfig(), /*enable_alltoall=*/false); -#else - backend.init(*plugin.mMOERunner, backend.mGemmToProfile, plugin.mType, plugin.mWeightType, plugin.mOutputType, - plugin.mNumExperts, plugin.mExpertsPerToken, plugin.mExpertHiddenSize, plugin.mExpertInterSize, - plugin.mGroupSize, plugin.mActivationType, plugin.hasBias(), plugin.hasLora(), /*min_latency_mode=*/false, - /*need_weights=*/true, plugin.getParallelismConfig()); -#endif -} diff --git a/cpp/tensorrt_llm/plugins/mixtureOfExperts/mixtureOfExpertsPlugin.h b/cpp/tensorrt_llm/plugins/mixtureOfExperts/mixtureOfExpertsPlugin.h deleted file mode 100644 index feb1f10cdc70..000000000000 --- a/cpp/tensorrt_llm/plugins/mixtureOfExperts/mixtureOfExpertsPlugin.h +++ /dev/null @@ -1,637 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2023 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. - */ -#ifndef TRT_MIXTURE_OF_EXPERTS_PLUGIN_H -#define TRT_MIXTURE_OF_EXPERTS_PLUGIN_H - -#include "NvInferPlugin.h" -#include "tensorrt_llm/kernels/cutlass_kernels/include/cutlass_kernel_selector.h" -#if defined(USING_OSS_CUTLASS_MOE_GEMM) -#include "tensorrt_llm/kernels/cutlass_kernels/include/moe_kernels.h" -#else -#include "moe_kernels.h" -#endif -#include "tensorrt_llm/common/cudaUtils.h" -#include "tensorrt_llm/common/quantization.h" -#include "tensorrt_llm/kernels/lora/lora.h" -#include "tensorrt_llm/plugins/common/gemmPluginProfiler.h" -#include "tensorrt_llm/plugins/common/plugin.h" -#include "tensorrt_llm/plugins/cudaStreamPlugin/cudaStreamPlugin.h" -#include "tensorrt_llm/plugins/gemmPlugin/gemmPlugin.h" -#include "tensorrt_llm/runtime/cudaStream.h" -#include -#include -#include -#include - -namespace tensorrt_llm::plugins -{ -namespace kernels = CUTLASS_MOE_GEMM_KERNELS_NAMESPACE; -using MoeMinLatencyParams = CUTLASS_MOE_GEMM_KERNELS_NAMESPACE::MoeMinLatencyParams; -using MOEParallelismConfig = CUTLASS_MOE_GEMM_KERNELS_NAMESPACE::MOEParallelismConfig; -using QuantParams = CUTLASS_MOE_GEMM_KERNELS_NAMESPACE::QuantParams; -using MoeGemmId = CUTLASS_MOE_GEMM_NAMESPACE::MoeGemmId; -using ActivationType = CUTLASS_MOE_GEMM_NAMESPACE::ActivationType; -using ActivationParams = CUTLASS_MOE_GEMM_KERNELS_NAMESPACE::ActivationParams; -using TmaWarpSpecializedGroupedGemmInput = CUTLASS_MOE_GEMM_NAMESPACE::TmaWarpSpecializedGroupedGemmInput; -using CUTLASS_MOE_GEMM_NAMESPACE::isGatedActivation; - -class MixtureOfExpertsGemmProfiler; -using MixtureOfExpertsPluginProfilerPtr = std::shared_ptr; -using GroupwiseQuantAlgo = tensorrt_llm::common::GroupwiseQuantAlgo; - -struct GemmIDMoe -{ - int gemm_idx; - int num_experts{}; - int experts_per_token{}; - kernels::MOEParallelismConfig parallelism_config{}; - int64_t hidden{}; - int64_t inter{}; - int64_t group_size{}; - ActivationType actfn{}; - nvinfer1::DataType dtype{}; - nvinfer1::DataType wdtype{}; - tensorrt_llm::common::QuantMode quant_mode; - bool determinism_mode = false; - - bool operator==(GemmIDMoe const& id) const - { - return id.gemm_idx == gemm_idx && id.num_experts == num_experts && id.experts_per_token == experts_per_token - && id.parallelism_config == parallelism_config && id.hidden == hidden && id.inter == inter - && id.group_size == group_size && id.actfn == actfn && id.dtype == dtype && id.wdtype == wdtype - && id.quant_mode == quant_mode && id.determinism_mode == determinism_mode; - } - - friend std::ostream& operator<<(std::ostream& out, GemmIDMoe const& id) - { - out << "gemm idx, experts, experts_per_token, parallelism_config, hidden, inter, group_size, actfn, dtype, " - "weight " - "type, parallelism mode, determinism mode=" - - << id.gemm_idx << "," << id.num_experts << "," << id.experts_per_token << "," << id.parallelism_config - << "," << id.hidden << "," << id.inter << "," << id.group_size << "," << static_cast(id.actfn) << "," - << static_cast(id.dtype) << "," << static_cast(id.wdtype) << "," << id.quant_mode.value() << "," - << id.determinism_mode; - return out; - } -}; - -// Hash of GemmIDMoe -struct GemmIDMoeHash -{ - std::size_t operator()(GemmIDMoe const& id) const - { - size_t hash = std::hash{}(id.gemm_idx); - hash ^= std::hash{}(id.num_experts); - hash ^= std::hash{}(id.experts_per_token); - hash ^= std::hash{}(id.parallelism_config.tp_size); - hash ^= std::hash{}(id.parallelism_config.ep_size); - hash ^= std::hash{}(id.parallelism_config.tp_rank); - hash ^= std::hash{}(id.parallelism_config.ep_rank); - hash ^= std::hash{}(id.hidden); - hash ^= std::hash{}(id.inter); - hash ^= std::hash{}(id.group_size); - hash ^= std::hash{}(static_cast(id.actfn)); - hash ^= std::hash{}(static_cast(id.dtype)); - hash ^= std::hash{}(static_cast(id.wdtype)); - hash ^= std::hash{}(static_cast(id.quant_mode.value())); - return hash; - } -}; - -class MixtureOfExpertsPlugin : public nvinfer1::IPluginV2DynamicExt -{ -public: - using LoraPluginProfilerPtr = std::shared_ptr; - using LoraImplPtr = std::shared_ptr; - MixtureOfExpertsPlugin() = delete; - MixtureOfExpertsPlugin(bool remove_input_padding, int number_of_experts, int experts_per_token, - int expert_hidden_size, int expert_inter_size, int groupwise_quant_algo, int group_size, - ActivationType activation_type, nvinfer1::DataType type, nvinfer1::DataType weight_type, - nvinfer1::DataType output_type, tensorrt_llm::common::QuantMode quant_mode, bool use_final_scales, - bool use_bias, int tp_size, int tp_rank, int ep_size, int ep_rank, bool force_determinism, int side_stream_id, - MixtureOfExpertsPluginProfilerPtr gemm_profiler_ptr, bool use_lora, nvinfer1::DataType lora_type, - LoraPluginProfilerPtr lora_profiler, int max_low_rank); - MixtureOfExpertsPlugin(void const* data, size_t length, MixtureOfExpertsPluginProfilerPtr gemm_profiler_ptr, - LoraPluginProfilerPtr lora_profiler); - MixtureOfExpertsPlugin(MixtureOfExpertsPlugin const&); - - void init(); - - ~MixtureOfExpertsPlugin() override = default; - - // IPluginV2DynamicExt Methods - nvinfer1::IPluginV2DynamicExt* clone() const noexcept override; - nvinfer1::DimsExprs getOutputDimensions(int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; - bool supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept override; - void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept override; - size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept override; - int enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - - // IPluginV2Ext Methods - nvinfer1::DataType getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept override; - - // IPluginV2 Methods - char const* getPluginType() const noexcept override; - char const* getPluginVersion() const noexcept override; - - int getNbOutputs() const noexcept override - { - return 1 + useSideStream(); - } - - int initialize() noexcept override; - void terminate() noexcept override; - size_t getSerializationSize() const noexcept override; - void serialize(void* buffer) const noexcept override; - void destroy() noexcept override; - void setPluginNamespace(char const* pluginNamespace) noexcept override; - char const* getPluginNamespace() const noexcept override; - -private: - friend class MixtureOfExpertsGemmProfiler; - std::unique_ptr mMOERunner{}; - int mNumExperts{}; - int mExpertsPerToken{}; - int64_t mExpertHiddenSize{}; - int64_t mExpertInterSize{}; - int64_t mGroupwiseQuantAlgo{}; - int64_t mGroupSize{}; - ActivationType mActivationType; - nvinfer1::DataType mType{}; - nvinfer1::DataType mWeightType{}; - nvinfer1::DataType mOutputType{}; - tensorrt_llm::common::QuantMode mQuantMode; - bool mUseFinalScales{}; - bool mUseBias{}; - MOEParallelismConfig mParallelismConfig{}; - - GemmDims mDims{}; - bool mUseDeterministicKernels = false; - int mSideStreamId = 0; - - int mDebugStallMain = 0; - int mDebugStallSide = 0; - - GemmIDMoe mGemmId1{}; - GemmIDMoe mGemmId2{}; - - MixtureOfExpertsPluginProfilerPtr mGemmProfiler; - - // lora related - bool mUseLora{}; - nvinfer1::DataType mLoraType{}; - int mMaxLowRank{}; - bool mRemoveInputPadding{}; - - LoraImplPtr mLoraImpl1; - LoraImplPtr mLoraImpl2; - - GemmIdCublas mLoraGemmId1{}; - GemmIdCublas mLoraGemmId2{}; - LoraPluginProfilerPtr mLoraProfiler; - - std::vector mLoraExpandFC1WeightPtrs{}; - std::vector mLoraExpandFC2WeightPtrs{}; - std::vector mLoraExpandGatedWeightPtrs{}; - std::vector mLoraExpandFC1Ranks{}; - std::vector mLoraExpandFC2Ranks{}; - std::vector mLoraExpandGatedRanks{}; - - cudaEvent_t mMemcpyEvent; - nvinfer1::pluginInternal::SideStream* mSideStreamPtr; - - // The below are not serialised - std::string const mLayerName{}; - std::string mNamespace{}; - - struct WorkspaceInfo - { - void* workspace{}; - void* src_to_dest_map{}; - void* lora_workspace{}; - size_t size{}; - }; - - int64_t getNumTokens(nvinfer1::PluginTensorDesc const* input_tensor) const; - WorkspaceInfo setupWorkspace(void* base_ptr, int64_t num_tokens, int num_reqs = 0) const; - - MOEParallelismConfig getParallelismConfig() const; - QuantParams getQuantParams(nvinfer1::PluginTensorDesc const* inputDesc, void const* const* inputs, - int scale_1_idx = -1, int scale_2_idx = -1, int scale_3_idx = -1, int scale_4_idx = -1, int scale_5_idx = -1, - int scale_6_idx = -1, int scale_7_idx = -1, int scale_8_idx = -1) const; - - int getNumLoraRequests(nvinfer1::PluginTensorDesc const* input_tensor) const; - tensorrt_llm::kernels::LoraParams getLoraParams( - nvinfer1::PluginTensorDesc const* inputDesc, void const* const* inputs, void* workspace); - - enum class RequestType : int32_t - { - kCONTEXT = 0, - kGENERATION = 1 - }; - - using IndexType = std::int32_t; - - // Inputs - constexpr static IndexType getInputTensorIndex() - { - return 0; - } - - constexpr static IndexType getExpertWeights1Index() - { - return getInputTensorIndex() + 1; - } - - constexpr static IndexType getExpertWeights2Index() - { - return getExpertWeights1Index() + 1; - } - - constexpr static IndexType getTokenSelectedExpertsIndex() - { - return getExpertWeights2Index() + 1; - } - - // Conditional inputs, we only allocate a new index if actually used - bool hasBias() const - { - return mUseBias; - } - - bool hasFinalScales() const - { - return mUseFinalScales; - } - - bool hasExpertIntQuantScales() const - { - return mQuantMode.hasInt4Weights() || mQuantMode.hasInt8Weights(); - } - - bool hasExpertFp8QuantScales() const - { - return mQuantMode.hasFp8Qdq(); - } - - bool hasExpertFp8FinalQuantScales() const - { - return hasExpertFp8QuantScales() && mOutputType == nvinfer1::DataType::kFP8; - } - - bool hasFP4QuantScales() const - { - return mQuantMode.hasNvfp4(); - } - - bool hasGroupwiseIntQuantScales() const - { - return mGroupwiseQuantAlgo > 0; - } - - bool hasExpertWeightQuantZeros() const - { - return mGroupwiseQuantAlgo & GroupwiseQuantAlgo::ZERO; - } - - bool hasExpertPrequantScales() const - { - return mGroupwiseQuantAlgo & GroupwiseQuantAlgo::PRE_QUANT_SCALE; - } - - bool hasGroupwiseFp8Alpha() const - { - return mGroupwiseQuantAlgo & GroupwiseQuantAlgo::FP8_ALPHA; - } - - bool useSideStream() const - { - return mSideStreamId > 0; - } - - bool hasLora() const - { - return mUseLora; - } - - bool hasGatedLoraWeightsAndRanks() const - { - return mUseLora && isGatedActivation(mActivationType); - } - - IndexType getTokenFinalScalesIndex() const - { - return getTokenSelectedExpertsIndex() + hasFinalScales(); - } - - IndexType getExpertBias1Index() const - { - return getTokenFinalScalesIndex() + hasBias(); - } - - IndexType getExpertBias2Index() const - { - return getExpertBias1Index() + hasBias(); - } - - /* - * Weight-Only int quant scales - */ - IndexType getExpertIntQuantScale1Index() const - { - return getExpertBias2Index() + hasExpertIntQuantScales(); - } - - IndexType getExpertIntQuantScale2Index() const - { - return getExpertIntQuantScale1Index() + hasExpertIntQuantScales(); - } - - /* - * FP8 Quant Scales - */ - IndexType getExpertFP8Dequant1Index() const - { - return getExpertIntQuantScale2Index() + hasExpertFp8QuantScales(); - } - - IndexType getExpertFP8Quant2Index() const - { - return getExpertFP8Dequant1Index() + hasExpertFp8QuantScales(); - } - - IndexType getExpertFP8Dequant2Index() const - { - return getExpertFP8Quant2Index() + hasExpertFp8QuantScales(); - } - - IndexType getExpertFP8QuantFinalIndex() const - { - return getExpertFP8Dequant2Index() + hasExpertFp8FinalQuantScales(); - } - - IndexType getInputFP8DequantIndex() const - { - return getExpertFP8QuantFinalIndex() + (hasExpertFp8QuantScales() && hasLora()); - } - - /* - * FP4 Quant Scales - */ - IndexType getFP4GlobalActSF1Index() const - { - return getInputFP8DequantIndex() + hasFP4QuantScales(); - } - - IndexType getFP4WeightSF1Index() const - { - return getFP4GlobalActSF1Index() + hasFP4QuantScales(); - } - - IndexType getFP4GlobalSF1Index() const - { - return getFP4WeightSF1Index() + hasFP4QuantScales(); - } - - IndexType getFP4GlobalActSF2Index() const - { - return getFP4GlobalSF1Index() + hasFP4QuantScales(); - } - - IndexType getFP4WeightSF2Index() const - { - return getFP4GlobalActSF2Index() + hasFP4QuantScales(); - } - - IndexType getFP4GlobalSF2Index() const - { - return getFP4WeightSF2Index() + hasFP4QuantScales(); - } - - /* - * Groupwise Params - */ - IndexType getExpertPrequantScales1Index() const - { - return getFP4GlobalSF2Index() + hasExpertPrequantScales(); - } - - IndexType getExpertPrequantScales2Index() const - { - return getExpertPrequantScales1Index() + hasExpertPrequantScales(); - } - - IndexType getExpertIntQuantZeros1Index() const - { - return getExpertPrequantScales2Index() + hasExpertWeightQuantZeros(); - } - - IndexType getExpertIntQuantZeros2Index() const - { - return getExpertIntQuantZeros1Index() + hasExpertWeightQuantZeros(); - } - - IndexType getExpertFp8Alpha1Index() const - { - return getExpertIntQuantZeros2Index() + hasGroupwiseFp8Alpha(); - } - - IndexType getExpertFp8Alpha2Index() const - { - return getExpertFp8Alpha1Index() + hasGroupwiseFp8Alpha(); - } - - /* - * LoRA params - */ - IndexType getLoraFC1WeightPtrsIndex() const - { - return getExpertFp8Alpha2Index() + hasLora(); - } - - IndexType getLoraFC1RanksIndex() const - { - return getLoraFC1WeightPtrsIndex() + hasLora(); - } - - IndexType getLoraFC2WeightPtrsIndex() const - { - return getLoraFC1RanksIndex() + hasLora(); - } - - IndexType getLoraFC2RanksIndex() const - { - return getLoraFC2WeightPtrsIndex() + hasLora(); - } - - IndexType getLoraGatedWeightPtrsIndex() const - { - return getLoraFC2RanksIndex() + hasGatedLoraWeightsAndRanks(); - } - - IndexType getLoraGatedRanksIndex() const - { - return getLoraGatedWeightPtrsIndex() + hasGatedLoraWeightsAndRanks(); - } - - IndexType getHostRequestTypeIndex() const - { - return getLoraGatedRanksIndex() + hasLora(); - } - - IndexType getHostContextLengthIndex() const - { - return getHostRequestTypeIndex() + (mRemoveInputPadding && hasLora()); - } - - IndexType getInputDummyTensorIndex() const - { - return getHostContextLengthIndex() + useSideStream(); - } - - IndexType getNbInputs() const - { - return getInputDummyTensorIndex() + 1; - } - - // Outputs - constexpr static IndexType getOutputTensorIndex() - { - return 0; - } - - IndexType getOutputDummyTensorIndex() const - { - return getOutputTensorIndex() + useSideStream(); - } - - /** - * Get the index of the expert shape tuple that represents the inner dimension - */ - int getGemmShapeInnerDimIndex() const - { - // In weight only mode the shape is transposed - return hasExpertIntQuantScales() ? 1 : 2; - } - - /** - * Get the index of the expert shape tuple that represents the outer dimension - */ - int getGemmShapeOuterDimIndex() const - { - // In weight only mode the shape is transposed - return hasExpertIntQuantScales() ? 2 : 1; - } - - /** - * Get quantization dimension scaling factor - */ - std::pair getWeightPackedElements() const - { - if (mGroupwiseQuantAlgo == 0) - { - return {1, mQuantMode.hasInt4Weights() ? 2 : 1}; - } - else - { - return {1, 4}; - } - } -}; - -class MixtureOfExpertsGemmProfiler - : public tensorrt_llm::plugins::GemmPluginProfiler -{ -public: - MixtureOfExpertsGemmProfiler() - { - // NOTE: Do not access mPlugin here, since we are called from the constructor before all fields are init - } - - void setGemmToProfile(kernels::GemmProfilerBackend::GemmToProfile gemm_to_profile) - { - // Just set the backend directly. This will just be reused in checkInit(). - backend.mGemmToProfile = gemm_to_profile; - // We need to set the backend to reinitialise itself with the new GEMM - init_backend = false; - } - - void setMaxProfileM(int maxProfileM) - { - mMaxProfileM = maxProfileM; - } - - virtual int getMaxProfileM() const override - { - return mMaxProfileM; - } - -protected: - using Config = tensorrt_llm::cutlass_extensions::CutlassGemmConfig; - void runTactic(int m, int n, int k, Config const& tactic, char* workspace, cudaStream_t const& stream) override; - void computeTmpSize(size_t maxM, size_t n, size_t k) override; - std::vector getTactics(int m, int n, int k) const override; - void initTmpData(int maxM, int n, int k, char* workspace, size_t size, cudaStream_t stream) override; - - void checkInit(); - - bool init_backend = false; - kernels::GemmProfilerBackend backend{}; - -private: - int mMaxProfileM = 0; -}; - -class MixtureOfExpertsPluginCreator : public nvinfer1::IPluginCreator -{ -public: - MixtureOfExpertsPluginCreator(); - - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; - - nvinfer1::IPluginV2* deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept override; - - void setPluginNamespace(char const* pluginNamespace) noexcept override; - - char const* getPluginNamespace() const noexcept override; - -private: - GemmPluginProfilerManager moePluginProfiler; - GemmPluginProfilerManager loraPluginProfileManager; - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; - std::string mNamespace; -}; - -} // namespace tensorrt_llm::plugins - -#endif // TRT_MIXTURE_OF_EXPERTS_PLUGIN_H diff --git a/cpp/tensorrt_llm/plugins/ncclPlugin/CMakeLists.txt b/cpp/tensorrt_llm/plugins/ncclPlugin/CMakeLists.txt deleted file mode 100644 index 86876224fccd..000000000000 --- a/cpp/tensorrt_llm/plugins/ncclPlugin/CMakeLists.txt +++ /dev/null @@ -1,21 +0,0 @@ -# -# SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. -# -file(GLOB SRCS *.cpp) -set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS}) -set(PLUGIN_SOURCES - ${PLUGIN_SOURCES} - PARENT_SCOPE) diff --git a/cpp/tensorrt_llm/plugins/ncclPlugin/allgatherPlugin.cpp b/cpp/tensorrt_llm/plugins/ncclPlugin/allgatherPlugin.cpp deleted file mode 100644 index 4825dd51bbab..000000000000 --- a/cpp/tensorrt_llm/plugins/ncclPlugin/allgatherPlugin.cpp +++ /dev/null @@ -1,253 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. - */ -#include "allgatherPlugin.h" -#include "tensorrt_llm/runtime/utils/mpiUtils.h" - -#include - -using namespace nvinfer1; -using tensorrt_llm::plugins::AllgatherPluginCreator; -using tensorrt_llm::plugins::AllgatherPlugin; - -static char const* ALLGATHER_PLUGIN_VERSION{"1"}; -static char const* ALLGATHER_PLUGIN_NAME{"AllGather"}; -PluginFieldCollection AllgatherPluginCreator::mFC{}; -std::vector AllgatherPluginCreator::mPluginAttributes; - -AllgatherPlugin::AllgatherPlugin(std::set group, nvinfer1::DataType type) - : mGroup(std::move(group)) - , mType(type) -{ -} - -// Parameterized constructor -AllgatherPlugin::AllgatherPlugin(void const* data, size_t length) -{ - char const *d = reinterpret_cast(data), *a = d; - read(d, mType); - mGroup.clear(); - int groupItem = 0; - while (d != a + length) - { - read(d, groupItem); - mGroup.insert(groupItem); - } - TLLM_CHECK_WITH_INFO(d == a + length, - "Expected length (%d) != real length (%d). This is often " - "caused by using different TensorRT LLM version to build " - "engine and run engine.", - (int) length, (int) (d - a)); -} - -// IPluginV2DynamicExt Methods -nvinfer1::IPluginV2DynamicExt* AllgatherPlugin::clone() const noexcept -{ - auto* plugin = new AllgatherPlugin(*this); - plugin->setPluginNamespace(mNamespace.c_str()); - return plugin; -} - -nvinfer1::DimsExprs AllgatherPlugin::getOutputDimensions( - int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept -{ - auto ret = inputs[0]; - auto groupSize = exprBuilder.constant(mGroup.size()); - ret.d[0] = exprBuilder.operation(DimensionOperation::kPROD, *ret.d[0], *groupSize); - return ret; -} - -bool AllgatherPlugin::supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept -{ - - return (inOut[pos].type == mType) && (inOut[pos].format == TensorFormat::kLINEAR); -} - -void AllgatherPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept -{ -} - -size_t AllgatherPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - return 0; -} - -int AllgatherPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept -{ - if (isBuilding()) - { - return 0; - } - size_t size = 1; - for (int i = 0; i < inputDesc[0].dims.nbDims; ++i) - { - size *= inputDesc[0].dims.d[i]; - } - - TLLM_CHECK_WITH_INFO(mNcclComm.get() != nullptr, "mNcclComm should be initialized before used"); - NCCLCHECK(ncclAllGather(inputs[0], outputs[0], size, (*getDtypeMap())[inputDesc[0].type], *mNcclComm, stream)); - - return 0; -} - -// IPluginV2Ext Methods -nvinfer1::DataType AllgatherPlugin::getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept -{ - assert(index == 0); - return inputTypes[0]; -} - -// IPluginV2 Methods - -char const* AllgatherPlugin::getPluginType() const noexcept -{ - return ALLGATHER_PLUGIN_NAME; -} - -char const* AllgatherPlugin::getPluginVersion() const noexcept -{ - return ALLGATHER_PLUGIN_VERSION; -} - -int AllgatherPlugin::getNbOutputs() const noexcept -{ - return 1; -} - -int AllgatherPlugin::initialize() noexcept -{ - if (isBuilding()) - { - return 0; - } - TLLM_LOG_TRACE("%s start for rank %d", __PRETTY_FUNCTION__, COMM_SESSION.getRank()); - mNcclComm = getComm(mGroup); - TLLM_LOG_TRACE("%s stop for rank %d", __PRETTY_FUNCTION__, COMM_SESSION.getRank()); - return 0; -} - -void AllgatherPlugin::terminate() noexcept {} - -size_t AllgatherPlugin::getSerializationSize() const noexcept -{ - return sizeof(int) * mGroup.size() + sizeof(mType); -} - -void AllgatherPlugin::serialize(void* buffer) const noexcept -{ - char *d = static_cast(buffer), *a = d; - write(d, mType); - for (auto it = mGroup.begin(); it != mGroup.end(); ++it) - { - write(d, *it); - } - TLLM_CHECK(d == a + getSerializationSize()); -} - -void AllgatherPlugin::destroy() noexcept -{ - // This gets called when the network containing plugin is destroyed - delete this; -} - -/////////////// - -AllgatherPluginCreator::AllgatherPluginCreator() -{ - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - mPluginAttributes.emplace_back(PluginField("group", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("type_id", nullptr, PluginFieldType::kINT32)); - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -char const* AllgatherPluginCreator::getPluginName() const noexcept -{ - return ALLGATHER_PLUGIN_NAME; -} - -char const* AllgatherPluginCreator::getPluginVersion() const noexcept -{ - return ALLGATHER_PLUGIN_VERSION; -} - -PluginFieldCollection const* AllgatherPluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -IPluginV2* AllgatherPluginCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept -{ - PluginField const* fields = fc->fields; - std::set group; - nvinfer1::DataType type{}; - // Read configurations from each fields - for (int i = 0; i < fc->nbFields; ++i) - { - char const* attrName = fields[i].name; - if (!strcmp(attrName, "group")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - auto const* r = static_cast(fields[i].data); - for (int j = 0; j < fields[i].length; ++j) - { - group.insert(*r); - ++r; - } - } - else if (!strcmp(attrName, "type_id")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - type = static_cast(*(static_cast(fields[i].data))); - } - } - - try - { - auto* obj = new AllgatherPlugin(group, type); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -IPluginV2* AllgatherPluginCreator::deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept -{ - // This object will be deleted when the network is destroyed, which will - // call AllgatherPlugin::destroy() - try - { - auto* obj = new AllgatherPlugin(serialData, serialLength); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} diff --git a/cpp/tensorrt_llm/plugins/ncclPlugin/allgatherPlugin.h b/cpp/tensorrt_llm/plugins/ncclPlugin/allgatherPlugin.h deleted file mode 100644 index 3d7810e6bd49..000000000000 --- a/cpp/tensorrt_llm/plugins/ncclPlugin/allgatherPlugin.h +++ /dev/null @@ -1,92 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. - */ -#pragma once - -#include "tensorrt_llm/plugins/common/plugin.h" -#include -#include -#include -#include - -namespace tensorrt_llm::plugins -{ - -class AllgatherPlugin : public BasePlugin -{ -public: - AllgatherPlugin(std::set group, nvinfer1::DataType type); - - AllgatherPlugin(void const* data, size_t length); - - ~AllgatherPlugin() override = default; - - // IPluginV2DynamicExt Methods - nvinfer1::IPluginV2DynamicExt* clone() const noexcept override; - nvinfer1::DimsExprs getOutputDimensions(int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; - bool supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept override; - void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept override; - size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept override; - int enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - - // IPluginV2Ext Methods - nvinfer1::DataType getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept override; - - // IPluginV2 Methods - char const* getPluginType() const noexcept override; - char const* getPluginVersion() const noexcept override; - int getNbOutputs() const noexcept override; - int initialize() noexcept override; - void terminate() noexcept override; - size_t getSerializationSize() const noexcept override; - void serialize(void* buffer) const noexcept override; - void destroy() noexcept override; - -private: - const std::string mLayerName; - std::set mGroup; - nvinfer1::DataType mType; - std::shared_ptr mNcclComm; -}; - -class AllgatherPluginCreator : public BaseCreator -{ -public: - AllgatherPluginCreator(); - - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; - - nvinfer1::IPluginV2* deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept override; - -private: - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; -}; - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/ncclPlugin/allreducePlugin.cpp b/cpp/tensorrt_llm/plugins/ncclPlugin/allreducePlugin.cpp deleted file mode 100644 index 24d9aff418f7..000000000000 --- a/cpp/tensorrt_llm/plugins/ncclPlugin/allreducePlugin.cpp +++ /dev/null @@ -1,986 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. - */ -#include "allreducePlugin.h" - -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/customAllReduceUtils.h" -#include "tensorrt_llm/common/dataType.h" -#include "tensorrt_llm/common/nvmlWrapper.h" -#include "tensorrt_llm/kernels/customAllReduceKernels.h" -#include "tensorrt_llm/kernels/userbuffers/ub_interface.h" -#include "tensorrt_llm/runtime/utils/mpiUtils.h" - -#include - -#include - -using namespace nvinfer1; -using tensorrt_llm::plugins::AllreducePluginCreator; -using tensorrt_llm::plugins::AllreducePlugin; -using tensorrt_llm::kernels::AllReduceFusionOp; -using tensorrt_llm::kernels::AllReduceStrategyType; -using tensorrt_llm::kernels::AllReduceStrategyConfig; -using tensorrt_llm::mpi::MpiTag; - -static char const* ALLREDUCE_PLUGIN_VERSION{"1"}; -static char const* ALLREDUCE_PLUGIN_NAME{"AllReduce"}; -PluginFieldCollection AllreducePluginCreator::mFC{}; -std::vector AllreducePluginCreator::mPluginAttributes; - -AllreducePlugin::AllreducePlugin(std::set group, nvinfer1::DataType type, AllReduceStrategyType strategy, - AllReduceStrategyConfig config, AllReduceFusionOp op, int32_t counter, float eps, int8_t affine, int8_t bias, - int8_t scale) - : mGroup(std::move(group)) - , mType(type) - , mStrategy(strategy) - , mConfig(config) - , mOp(op) - , mEps(eps) - , mAffine(affine) - , mBias(bias) - , mScale(scale) -{ - check(); -} - -// Parameterized constructor -AllreducePlugin::AllreducePlugin(void const* data, size_t length) -{ - char const *d = reinterpret_cast(data), *a = d; - read(d, mType); - read(d, mStrategy); - read(d, mConfig); - read(d, mOp); - read(d, mEps); - read(d, mAffine); - read(d, mBias); - read(d, mScale); - mGroup.clear(); - int groupItem = 0; - while (d != a + length) - { - read(d, groupItem); - mGroup.insert(groupItem); - } - TLLM_CHECK_WITH_INFO(d == a + length, - "Expected length (%d) != real length (%d). This is often " - "caused by using different TensorRT LLM version to build " - "engine and run engine.", - (int) length, (int) (d - a)); - check(); -} - -void AllreducePlugin::check() noexcept -{ - if (mStrategy != AllReduceStrategyType::UB) - { - TLLM_CHECK(mOp != AllReduceFusionOp::LAST_PROCESS_FOR_UB); - } -} - -// IPluginV2DynamicExt Methods -nvinfer1::IPluginV2DynamicExt* AllreducePlugin::clone() const noexcept -{ - auto* plugin = new AllreducePlugin(*this); - plugin->setPluginNamespace(mNamespace.c_str()); - return plugin; -} - -nvinfer1::DimsExprs AllreducePlugin::getOutputDimensions( - int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept -{ - if (mOp == AllReduceFusionOp::RESIDUAL_RMS_NORM_QUANT_NVFP4 && mStrategy == AllReduceStrategyType::UB && mScale) - { - if (outputIndex == 0) - { - DimsExprs ret; - ret.nbDims = inputs[0].nbDims; - for (int di = 0; di < ret.nbDims; ++di) - { - ret.d[di] = inputs[0].d[di]; - } - return ret; - } - else if (outputIndex == 2) - { - DimsExprs ret; - ret.nbDims = inputs[0].nbDims; - for (int di = 0; di < ret.nbDims; ++di) - { - ret.d[di] = inputs[0].d[di]; - } - auto dimM = exprBuilder.operation( - DimensionOperation::kCEIL_DIV, *ret.d[ret.nbDims - 2], *exprBuilder.constant(128)); - ret.d[ret.nbDims - 2] = exprBuilder.operation(DimensionOperation::kPROD, *dimM, *exprBuilder.constant(128)); - ret.d[ret.nbDims - 1] = exprBuilder.operation( - DimensionOperation::kCEIL_DIV, *ret.d[ret.nbDims - 1], *exprBuilder.constant(16)); - return ret; - } - } - return inputs[0]; -} - -bool AllreducePlugin::supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept -{ - int base_inputs = 0; - switch (mStrategy) - { - case AllReduceStrategyType::NCCL: - case AllReduceStrategyType::UB: - case AllReduceStrategyType::NCCL_SYMMETRIC: base_inputs = 1; break; - default: base_inputs = 2; break; - } - int fusion_op_extra_inputs = 0; - int scale_idx = 0; - if (mOp != AllReduceFusionOp::NONE) - { - ++fusion_op_extra_inputs; - if (mAffine) - { - if (mOp == AllReduceFusionOp::RESIDUAL_RMS_PREPOST_NORM) - ++fusion_op_extra_inputs; - ++fusion_op_extra_inputs; - } - if (mBias) - { - ++fusion_op_extra_inputs; - } - if (mScale) - { - scale_idx = base_inputs + fusion_op_extra_inputs; - ++fusion_op_extra_inputs; - } - } - - TLLM_CHECK(nbInputs == (base_inputs + fusion_op_extra_inputs)); - - if (pos == 1) - { - switch (mStrategy) - { - case AllReduceStrategyType::NCCL: - case AllReduceStrategyType::UB: - case AllReduceStrategyType::NCCL_SYMMETRIC: break; - default: return (inOut[pos].type == nvinfer1::DataType::kINT64) && (inOut[pos].format == TensorFormat::kLINEAR); - } - } - if (mStrategy == AllReduceStrategyType::UB) - { - if (mScale && pos == scale_idx) - { - return (inOut[pos].type == nvinfer1::DataType::kFLOAT) && (inOut[pos].format == TensorFormat::kLINEAR); - } - if (mOp == AllReduceFusionOp::RESIDUAL_RMS_NORM_QUANT_NVFP4) - { - if (pos == nbInputs) - { - return (inOut[pos].type == nvinfer1::DataType::kFP4) && (inOut[pos].format == TensorFormat::kLINEAR); - } - if (pos == (nbInputs + 2)) - { - return (inOut[pos].type == nvinfer1::DataType::kFP8) && (inOut[pos].format == TensorFormat::kLINEAR); - } - } - if (mOp == AllReduceFusionOp::RESIDUAL_RMS_NORM_QUANT_FP8) - { - if (pos == nbInputs) - { - return (inOut[pos].type == nvinfer1::DataType::kFP8) && (inOut[pos].format == TensorFormat::kLINEAR); - } - } - } - return (inOut[pos].type == mType) && (inOut[pos].format == TensorFormat::kLINEAR); -} - -void AllreducePlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept -{ -} - -size_t AllreducePlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - return 0; -} - -AllReduceStrategyType AllreducePlugin::selectImplementation( - size_t messageSize, int worldSize, nvinfer1::DataType type) noexcept -{ - bool const isAuto = (mStrategy == AllReduceStrategyType::AUTO); - - bool const forceDeterministic = common::getEnvForceDeterministicAllReduce(); - if (!mIsP2PSupported) - { - if (!isAuto) - { - TLLM_LOG_INFO("Since Peer to Peer not supported, fallback to AllReduceStrategy: NCCL_SYMMETRIC"); - } - else if (forceDeterministic) - { - TLLM_LOG_WARNING( - "Since Peer to Peer not supported, fallback to AllReduceStrategy: NCCL_SYMMETRIC. NCCL_SYMMETRIC might " - "produce " - "non-deterministic results."); - } - return AllReduceStrategyType::NCCL_SYMMETRIC; - } - - if (isAuto && !mIsNVLINKSupported && !forceDeterministic) - { - return AllReduceStrategyType::NCCL_SYMMETRIC; - } - - auto const maxWorkspaceSize = utils::customAllReduceUtils::getMaxRequiredWorkspaceSize(worldSize); - - AllReduceStrategyType strat = AllReduceStrategyType::NCCL_SYMMETRIC; - auto const messageSizeBytes = messageSize * common::getDTypeSize(type); - - if (messageSizeBytes <= maxWorkspaceSize) - { - // In some instances, the two-shot strategy has exhibited significant performance issues. - // As a temporary measure, we have disabled the two-shot strategy. - // TODO: remove this WAR after https://nvbugspro.nvidia.com/bug/4718747 is fixed. - if (!isAuto) - { - strat = mStrategy; - } - else if (forceDeterministic) - { - strat = AllReduceStrategyType::ONESHOT; - } - else if (worldSize <= 2) - { - strat = AllReduceStrategyType::ONESHOT; - } - else if (worldSize <= 4) - { - if (messageSizeBytes < 1 * 1000 * 1000) - { - strat = AllReduceStrategyType::ONESHOT; - } - else - { - strat = AllReduceStrategyType::NCCL_SYMMETRIC; - } - } - else - { - if (messageSizeBytes < 500 * 1000) - { - strat = AllReduceStrategyType::ONESHOT; - } - else - { - strat = AllReduceStrategyType::NCCL_SYMMETRIC; - } - } - - if (!kernels::configurationSupported(strat, messageSize, worldSize, type)) - { - if (!isAuto) - { - TLLM_LOG_WARNING("Since not aligned, fallback to AllReduceStrategy: NCCL_SYMMETRIC"); - } - else if (forceDeterministic) - { - TLLM_LOG_WARNING( - "Since not aligned, fallback to AllReduceStrategy: NCCL_SYMMETRIC. NCCL_SYMMETRIC might produce " - "non-deterministic results."); - } - strat = AllReduceStrategyType::NCCL_SYMMETRIC; - } - } - else - { - if (!isAuto) - { - TLLM_LOG_WARNING("Since messageSize > maxWorkspace, fallback to AllReduceStrategy: NCCL_SYMMETRIC"); - } - else if (forceDeterministic) - { - TLLM_LOG_WARNING( - "Since messageSize > maxWorkspace, fallback to AllReduceStrategy: NCCL_SYMMETRIC. NCCL_SYMMETRIC might " - "produce " - "non-deterministic results."); - } - strat = AllReduceStrategyType::NCCL_SYMMETRIC; - } - - return strat; -} - -int AllreducePlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept -{ - if (isBuilding()) - { - return 0; - } - size_t size = 1; - for (int i = 0; i < inputDesc[0].dims.nbDims; ++i) - { - size *= inputDesc[0].dims.d[i]; - } - - kernels::AllReduceStrategyType runtimeStrategy; - - static char* forceNcclAllReduceStrategyChar = std::getenv("FORCE_NCCL_ALL_REDUCE_STRATEGY"); - bool forceNcclAllReduceStrategy = (forceNcclAllReduceStrategyChar != nullptr); - if (forceNcclAllReduceStrategy || mStrategy == AllReduceStrategyType::NCCL) - { - runtimeStrategy = AllReduceStrategyType::NCCL; - } - else if (mStrategy == AllReduceStrategyType::NCCL_SYMMETRIC) - { - runtimeStrategy = AllReduceStrategyType::NCCL_SYMMETRIC; - } - else if (mStrategy == AllReduceStrategyType::UB) - { - runtimeStrategy = AllReduceStrategyType::UB; - } - else - { - runtimeStrategy = selectImplementation(size, mGroup.size(), mType); - } - - // Log runtime strategy - auto const rank = COMM_SESSION.getRank(); - switch (runtimeStrategy) - { - case AllReduceStrategyType::NCCL: - { - TLLM_LOG_DEBUG("AllReducePlugin strategy for rank %d: NCCL", rank); - break; - } - case AllReduceStrategyType::NCCL_SYMMETRIC: - { - TLLM_LOG_DEBUG("AllReducePlugin strategy for rank %d: NCCL_SYMMETRIC", rank); - break; - } - case AllReduceStrategyType::ONESHOT: - { - TLLM_LOG_DEBUG("AllReducePlugin strategy for rank %d: ONESHOT", rank); - break; - } - case AllReduceStrategyType::TWOSHOT: - { - TLLM_LOG_DEBUG("AllReducePlugin strategy for rank %d: TWOSHOT", rank); - break; - } - case AllReduceStrategyType::UB: - { - TLLM_LOG_DEBUG("AllReducePlugin strategy for rank %d: UB", rank); - break; - } - default: break; - } - - if (runtimeStrategy == AllReduceStrategyType::NCCL || runtimeStrategy == AllReduceStrategyType::NCCL_SYMMETRIC) - { - if (mOp == AllReduceFusionOp::RESIDUAL_RMS_NORM || mOp == AllReduceFusionOp::RESIDUAL_RMS_PREPOST_NORM) - { - NCCLCHECK(ncclAllReduce(inputs[0], outputs[1], size, (*getDtypeMap())[mType], ncclSum, *mNcclComm, stream)); - tensorrt_llm::kernels::AllReduceParams params; - int fusion_ptr_idx = 0; - if (mStrategy == AllReduceStrategyType::NCCL || mStrategy == AllReduceStrategyType::NCCL_SYMMETRIC) - { - fusion_ptr_idx = 1; - } - else - { - fusion_ptr_idx = 2; - } - params.fusion_params.bias_buffer = mBias ? inputs[fusion_ptr_idx++] : nullptr; - params.fusion_params.residual_buffer = inputs[fusion_ptr_idx++]; - params.fusion_params.weight_buffer = mAffine ? inputs[fusion_ptr_idx++] : nullptr; - if (mOp == AllReduceFusionOp::RESIDUAL_RMS_PREPOST_NORM) - { - params.fusion_params.weight_buffer_pre_residual_norm = mAffine ? inputs[fusion_ptr_idx++] : nullptr; - } - params.local_output_buffer_ptr = outputs[0]; - params.elts_total = size; - params.fusion_params.hidden_size = inputDesc[0].dims.d[inputDesc[0].dims.nbDims - 1]; - params.fusion_params.eps = mEps; - params.fusion_params.intermediate_buffer = outputs[1]; - TLLM_LOG_DEBUG("residualRmsNorm called"); - tensorrt_llm::kernels::residualRmsNorm(params, mType, stream, mOp); - } - else - { - NCCLCHECK(ncclAllReduce(inputs[0], outputs[0], size, (*getDtypeMap())[mType], ncclSum, *mNcclComm, stream)); - } - } - else if (runtimeStrategy == AllReduceStrategyType::UB) - { - TLLM_CHECK(!mBias); - - size_t dtype_size = tensorrt_llm::common::getDTypeSize(mType); - int hidden_size = inputDesc[0].dims.d[inputDesc[0].dims.nbDims - 1]; - - TLLM_CHECK_WITH_INFO(tensorrt_llm::runtime::ub::ub_is_initialized(), "UserBuffer has not been initialized!"); - auto ub_buffer0 = tensorrt_llm::runtime::ub::ub_get(0); - auto ub_buffer1 = tensorrt_llm::runtime::ub::ub_get(1); - TLLM_CHECK(inputs[0] == ub_buffer0.addr); - auto ub_comm = tensorrt_llm::runtime::ub::ub_comm(); - if (mOp == AllReduceFusionOp::RESIDUAL_RMS_NORM_QUANT_FP8) - { - TLLM_CHECK(mAffine); - TLLM_CHECK(mScale); - TLLM_CHECK(outputs[0] == ub_buffer1.addr); - void* residual = const_cast(inputs[1]); - void* gamma = const_cast(inputs[2]); - float* scale = const_cast(reinterpret_cast(inputs[3])); - tensorrt_llm::kernels::ub::allreduce2_userbuff_inplace_rmsnorm_quant_launcher(ub_buffer0.handle, 0, - ub_buffer1.handle, 0, size, hidden_size, nullptr, gamma, mEps, scale, residual, outputs[1], mType, - ub_comm, stream); - } - else if (mOp == AllReduceFusionOp::RESIDUAL_RMS_NORM_QUANT_NVFP4) - { - auto ub_buffer2 = tensorrt_llm::runtime::ub::ub_get(2); - TLLM_CHECK(mAffine); - TLLM_CHECK(mScale); - TLLM_CHECK(outputs[0] == ub_buffer1.addr); - TLLM_CHECK(outputs[2] == ub_buffer2.addr); - void* residual = const_cast(inputs[1]); - void* gamma = const_cast(inputs[2]); - float* scale = const_cast(reinterpret_cast(inputs[3])); - tensorrt_llm::kernels::ub::allreduce2_userbuff_inplace_rmsnorm_quant_fp4_launcher(ub_buffer0.handle, 0, - ub_buffer1.handle, 0, ub_buffer2.handle, 0, size, hidden_size, nullptr, gamma, mEps, scale, residual, - outputs[1], mType, ub_comm, stream); - } - else if (mOp == AllReduceFusionOp::LAST_PROCESS_FOR_UB) - { - TLLM_CHECK(outputs[1] == ub_buffer1.addr); - void* residual = const_cast(inputs[1]); - tensorrt_llm::kernels::ub::allreduce2_userbuff_inplace_launcher( - ub_buffer0.handle, 0, size, mType, ub_comm, stream); - tensorrt_llm::kernels::ub::allgather2_userbuff_residual_launcher( - ub_buffer1.handle, 0, size, hidden_size, residual, mType, ub_comm, stream); - TLLM_CUDA_CHECK( - cudaMemcpyAsync(outputs[0], ub_buffer0.addr, size * dtype_size, cudaMemcpyDeviceToDevice, stream)); - } - else if (mOp == AllReduceFusionOp::NONE) - { - tensorrt_llm::kernels::ub::allreduce2_userbuff_inplace_launcher( - ub_buffer0.handle, 0, size, mType, ub_comm, stream); - TLLM_CUDA_CHECK( - cudaMemcpyAsync(outputs[0], ub_buffer0.addr, size * dtype_size, cudaMemcpyDeviceToDevice, stream)); - } - else - { - TLLM_CHECK_WITH_INFO(false, "Unsupported UB allreduce fusion op"); - } - } - else - { - auto const tpSize = mGroup.size(); - int tpRank = 0; - for (auto const& currentRank : mGroup) - { - if (rank == currentRank) - break; - ++tpRank; - } - - int token_num = size / inputDesc[0].dims.d[inputDesc[0].dims.nbDims - 1]; - int hidden_size = inputDesc[0].dims.d[inputDesc[0].dims.nbDims - 1]; - auto params = tensorrt_llm::kernels::AllReduceParams::deserialize( - reinterpret_cast(const_cast(inputs[1])), tpSize, tpRank, mType, token_num, hidden_size, - mOp); - - params.local_output_buffer_ptr = outputs[0]; - params.local_input_buffer_ptr = inputs[0]; - params.elts_total = size; - - int fusion_ptr_idx = 2; - params.fusion_params.bias_buffer = mBias ? inputs[fusion_ptr_idx++] : nullptr; - params.fusion_params.residual_buffer = inputs[fusion_ptr_idx++]; - params.fusion_params.weight_buffer = mAffine ? inputs[fusion_ptr_idx++] : nullptr; - if (mOp == AllReduceFusionOp::RESIDUAL_RMS_PREPOST_NORM) - params.fusion_params.weight_buffer_pre_residual_norm = mAffine ? inputs[fusion_ptr_idx++] : nullptr; - params.fusion_params.hidden_size = hidden_size; - params.fusion_params.eps = mEps; - params.fusion_params.intermediate_buffer = outputs[1]; - if (mOp == AllReduceFusionOp::RESIDUAL_RMS_NORM) - { - for (size_t i = 0; i < tpSize; ++i) - { - params.fusion_params.lamport_peer_comm_buffer_ptrs[i] - = reinterpret_cast(const_cast(inputs[1]))[tpSize * 4 + i]; - params.fusion_params.lamport_peer_comm_buffer_ptrs[i + tensorrt_llm::kernels::MAX_RANKS_PER_NODE] - = reinterpret_cast(const_cast(inputs[1]))[tpSize * 5 + i]; - params.fusion_params.lamport_peer_comm_buffer_ptrs[i + tensorrt_llm::kernels::MAX_RANKS_PER_NODE * 2] - = reinterpret_cast(const_cast(inputs[1]))[tpSize * 6 + i]; - } - } - TLLM_LOG_DEBUG("customAllReduce called"); - tensorrt_llm::kernels::customAllReduce(params, mType, runtimeStrategy, mConfig, mOp, stream); - } - - return 0; -} - -// IPluginV2Ext Methods -nvinfer1::DataType AllreducePlugin::getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept -{ - TLLM_CHECK(index < getNbOutputs()); - if (mOp == AllReduceFusionOp::RESIDUAL_RMS_NORM_QUANT_NVFP4) - { - if (index == 0) - { - return nvinfer1::DataType::kFP4; - } - else if (index == 2) - { - return nvinfer1::DataType::kFP8; - } - } - if (mOp == AllReduceFusionOp::RESIDUAL_RMS_NORM_QUANT_FP8) - { - if (index == 0) - { - return nvinfer1::DataType::kFP8; - } - } - return inputTypes[0]; -} - -// IPluginV2 Methods - -char const* AllreducePlugin::getPluginType() const noexcept -{ - return ALLREDUCE_PLUGIN_NAME; -} - -char const* AllreducePlugin::getPluginVersion() const noexcept -{ - return ALLREDUCE_PLUGIN_VERSION; -} - -int AllreducePlugin::getNbOutputs() const noexcept -{ - if (mOp == AllReduceFusionOp::NONE) - { - return 1; - } - else if (mOp == AllReduceFusionOp::RESIDUAL_RMS_NORM_QUANT_NVFP4) - { - return 3; - } - else - { - return 2; - } -} - -bool AllreducePlugin::isCustomAllReduceSupported(int ranks_per_node) const noexcept -{ - constexpr bool isCudaVersionSupported = -#if defined(CUDART_VERSION) && CUDART_VERSION >= 11020 - true; -#else - false; -#endif - - return isCudaVersionSupported && (ranks_per_node % 2 == 0) - && (static_cast(ranks_per_node) <= kernels::MAX_RANKS_PER_NODE) && (ranks_per_node > 0); -} - -using tensorrt_llm::common::NvmlManager; -using tensorrt_llm::common::NVMLWrapper; - -std::set getLocalGroup(std::set const& group) -{ - auto const myRank = COMM_SESSION.getRank(); - auto const myLocalRank = LOCAL_COMM_SESSION.getRank(); - auto const localSize = LOCAL_COMM_SESSION.getSize(); - - std::vector ranks(localSize, 0); - std::vector localRanks(localSize, 0); - if (group.size() >= static_cast(localSize)) - { - LOCAL_COMM_SESSION.allgather(&myRank, ranks.data(), 1, tensorrt_llm::mpi::MpiType::kINT32); - LOCAL_COMM_SESSION.allgather(&myLocalRank, localRanks.data(), 1, tensorrt_llm::mpi::MpiType::kINT32); - } - else - { - if (myRank == *group.begin()) - { - ranks.clear(); - int rank; - ranks.push_back(myRank); - for (auto it = std::next(std::begin(group), 1); it != group.end(); ++it) - { - COMM_SESSION.recvValue(rank, *it, MpiTag::kDefault); - ranks.push_back(rank); - } - for (auto it = std::next(std::begin(group), 1); it != group.end(); ++it) - { - COMM_SESSION.send(ranks.data(), localSize, tensorrt_llm::mpi::MpiType::kINT32, *it, MpiTag::kDefault); - } - - localRanks.clear(); - localRanks.push_back(myLocalRank); - for (auto it = std::next(std::begin(group), 1); it != group.end(); ++it) - { - COMM_SESSION.recvValue(rank, *it, MpiTag::kDefault); - localRanks.push_back(rank); - } - for (auto it = std::next(std::begin(group), 1); it != group.end(); ++it) - { - COMM_SESSION.send( - localRanks.data(), localSize, tensorrt_llm::mpi::MpiType::kINT32, *it, MpiTag::kDefault); - } - } - else - { - COMM_SESSION.sendValue(myRank, *group.begin(), MpiTag::kDefault); - COMM_SESSION.recv( - ranks.data(), localSize, tensorrt_llm::mpi::MpiType::kINT32, *group.begin(), MpiTag::kDefault); - - COMM_SESSION.sendValue(myLocalRank, *group.begin(), MpiTag::kDefault); - COMM_SESSION.recv( - localRanks.data(), localSize, tensorrt_llm::mpi::MpiType::kINT32, *group.begin(), MpiTag::kDefault); - } - } - - std::set localGroup; - for (size_t i = 0; i < ranks.size(); ++i) - { - auto rank = ranks[i]; - if (group.find(rank) != group.end()) - { - localGroup.insert(localRanks[i]); - } - } - return localGroup; -} - -void AllreducePlugin::initGroupTopology() noexcept -{ - static std::map, std::tuple> cache; - if (cache.find(mGroup) != cache.end()) - { - auto [isNVLINKSupported, isP2PSupported] = cache[mGroup]; - mIsNVLINKSupported = isNVLINKSupported; - mIsP2PSupported = isP2PSupported; - return; - } - setGroupTopology(); - cache[mGroup] = {mIsNVLINKSupported, mIsP2PSupported}; -} - -void AllreducePlugin::setGroupTopology() noexcept -{ - auto const rank = COMM_SESSION.getRank(); - TLLM_LOG_INFO("Detecting local TP group for rank %d", rank); - std::set localGroup = getLocalGroup(mGroup); - if (mGroup.size() != localGroup.size()) - { - mIsP2PSupported = false; - mIsNVLINKSupported = false; - TLLM_LOG_INFO("Found inter-node TP group for rank %d", rank); - return; - } - TLLM_LOG_INFO("TP group is intra-node for rank %d", rank); - - NvmlManager nvmlManager; - auto const& nvml = nvmlManager.sharedWrapper(); - std::unordered_set visitedDevice; - mIsP2PSupported = true; - mIsNVLINKSupported = true; - - // Use cudaDeviceCanAccessPeer to determine whether p2p is supported, - // and use nvml to determine whether there are nvlink links between ranks. - for (int firstDeviceId : localGroup) - { - for (int secondDeviceId : localGroup) - { - if (firstDeviceId == secondDeviceId || visitedDevice.find(secondDeviceId) != visitedDevice.end()) - { - continue; - } - - int canAccessPeer = 0; - TLLM_CUDA_CHECK(cudaDeviceCanAccessPeer(&canAccessPeer, firstDeviceId, secondDeviceId)); - - if (!canAccessPeer) - { - mIsP2PSupported = false; - mIsNVLINKSupported = false; - - return; - } - - nvmlDevice_t firstDevice; - NVML_CHECK(nvml->nvmlDeviceGetHandleByIndex(firstDeviceId, &firstDevice)); - - bool isNVLINK = false; - - for (unsigned int link = 0; link < NVML_NVLINK_MAX_LINKS; link++) - { - nvmlPciInfo_t remotePciInfo; - if (nvml->nvmlDeviceGetNvLinkRemotePciInfo(firstDevice, link, &remotePciInfo) != NVML_SUCCESS) - { - continue; - } - - nvmlDevice_t remoteDevice; - auto const result = nvml->nvmlDeviceGetHandleByPciBusId(remotePciInfo.busId, &remoteDevice); - - if (result == NVML_SUCCESS) - { - // Two GPUs are connected directly through nvlink - unsigned int remoteDeviceId; - NVML_CHECK(nvml->nvmlDeviceGetIndex(remoteDevice, &remoteDeviceId)); - - if (remoteDeviceId == static_cast(secondDeviceId)) - { - isNVLINK = true; - } - } - else if (result == NVML_ERROR_NOT_FOUND) - { - // Maybe Two GPUs are connected via nvswitch, - // now remotePciInfo represents the pci information of nvswitch, - // determine whether nvlink is supported by whether two GPUs are connected to the same nvswitch. - nvmlDevice_t secondDevice; - NVML_CHECK(nvml->nvmlDeviceGetHandleByIndex(secondDeviceId, &secondDevice)); - - for (unsigned int secondLink = 0; secondLink < NVML_NVLINK_MAX_LINKS; secondLink++) - { - nvmlPciInfo_t secondRemotePciInfo; - if (nvml->nvmlDeviceGetNvLinkRemotePciInfo(secondDevice, secondLink, &secondRemotePciInfo) - != NVML_SUCCESS) - { - continue; - } - - if (strcmp(remotePciInfo.busId, secondRemotePciInfo.busId) == 0) - { - isNVLINK = true; - break; - } - } - } - else - { - NVML_CHECK(result); - } - - if (isNVLINK) - { - break; - } - } - - mIsNVLINKSupported &= isNVLINK; - } - visitedDevice.insert(firstDeviceId); - } -} - -int AllreducePlugin::initialize() noexcept -{ - if (isBuilding()) - { - return 0; - } - - TLLM_LOG_TRACE("%s start for rank %d", __PRETTY_FUNCTION__, COMM_SESSION.getRank()); - mNcclComm = getComm(mGroup); - if (mStrategy != AllReduceStrategyType::NCCL) - { - initGroupTopology(); - } - - TLLM_LOG_TRACE("%s stop for rank %d", __PRETTY_FUNCTION__, COMM_SESSION.getRank()); - return 0; -} - -void AllreducePlugin::terminate() noexcept {} - -size_t AllreducePlugin::getSerializationSize() const noexcept -{ - return sizeof(int) * mGroup.size() + sizeof(mType) + sizeof(mStrategy) + sizeof(mConfig) + sizeof(mOp) - + sizeof(mEps) + sizeof(mAffine) + sizeof(mBias) + sizeof(mScale); -} - -void AllreducePlugin::serialize(void* buffer) const noexcept -{ - char *d = static_cast(buffer), *a = d; - write(d, mType); - write(d, mStrategy); - write(d, mConfig); - write(d, mOp); - write(d, mEps); - write(d, mAffine); - write(d, mBias); - write(d, mScale); - for (auto it = mGroup.begin(); it != mGroup.end(); ++it) - { - write(d, *it); - } - TLLM_CHECK(d == a + getSerializationSize()); -} - -void AllreducePlugin::destroy() noexcept -{ - // This gets called when the network containing plugin is destroyed - delete this; -} - -/////////////// - -AllreducePluginCreator::AllreducePluginCreator() -{ - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - mPluginAttributes.emplace_back(PluginField("group", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("type_id", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("strategy", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("config", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("fusion_op", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("counter", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("eps", nullptr, PluginFieldType::kFLOAT32)); - mPluginAttributes.emplace_back(PluginField("affine", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("bias", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("scale", nullptr, PluginFieldType::kINT8)); - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -char const* AllreducePluginCreator::getPluginName() const noexcept -{ - return ALLREDUCE_PLUGIN_NAME; -} - -char const* AllreducePluginCreator::getPluginVersion() const noexcept -{ - return ALLREDUCE_PLUGIN_VERSION; -} - -PluginFieldCollection const* AllreducePluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -IPluginV2* AllreducePluginCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept -{ - PluginField const* fields = fc->fields; - std::set group; - nvinfer1::DataType type{}; - AllReduceStrategyType strategy{}; - AllReduceStrategyConfig config{}; - AllReduceFusionOp fusion_op{}; - int32_t counter{}; - float eps{}; - int8_t affine{}; - int8_t bias{}; - int8_t scale{}; - // Read configurations from each fields - for (int i = 0; i < fc->nbFields; ++i) - { - char const* attrName = fields[i].name; - if (!strcmp(attrName, "group")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - auto const* r = static_cast(fields[i].data); - for (int j = 0; j < fields[i].length; ++j) - { - group.insert(*r); - ++r; - } - } - else if (!strcmp(attrName, "type_id")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - type = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "strategy")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT8); - strategy = static_cast(*static_cast(fields[i].data)); - } - else if (!strcmp(attrName, "config")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT8); - config = static_cast(*static_cast(fields[i].data)); - } - else if (!strcmp(attrName, "fusion_op")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT8); - fusion_op = static_cast(*static_cast(fields[i].data)); - } - else if (!strcmp(attrName, "counter")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - counter = *static_cast(fields[i].data); - } - else if (!strcmp(attrName, "eps")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kFLOAT32); - eps = *static_cast(fields[i].data); - } - else if (!strcmp(attrName, "affine")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT8); - affine = *static_cast(fields[i].data); - } - else if (!strcmp(attrName, "bias")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT8); - bias = *static_cast(fields[i].data); - } - else if (!strcmp(attrName, "scale")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT8); - scale = *static_cast(fields[i].data); - } - } - try - { - auto* obj = new AllreducePlugin(group, type, strategy, config, fusion_op, counter, eps, affine, bias, scale); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -IPluginV2* AllreducePluginCreator::deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept -{ - // This object will be deleted when the network is destroyed, which will - // call AllreducePlugin::destroy() - try - { - auto* obj = new AllreducePlugin(serialData, serialLength); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} diff --git a/cpp/tensorrt_llm/plugins/ncclPlugin/allreducePlugin.h b/cpp/tensorrt_llm/plugins/ncclPlugin/allreducePlugin.h deleted file mode 100644 index 881fbf3b89a5..000000000000 --- a/cpp/tensorrt_llm/plugins/ncclPlugin/allreducePlugin.h +++ /dev/null @@ -1,114 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. - */ -#pragma once - -#include "tensorrt_llm/kernels/customAllReduceKernels.h" -#include "tensorrt_llm/plugins/common/plugin.h" - -#include -#include -#include -#include -#include - -namespace tensorrt_llm::plugins -{ -namespace tk = ::tensorrt_llm::kernels; - -class AllreducePlugin : public BasePlugin -{ -public: - AllreducePlugin(std::set group, nvinfer1::DataType type, tk::AllReduceStrategyType strategy, - tk::AllReduceStrategyConfig config, tk::AllReduceFusionOp op, int32_t counter, float eps, int8_t affine, - int8_t bias, int8_t scale); - - AllreducePlugin(void const* data, size_t length); - - ~AllreducePlugin() override = default; - - // IPluginV2DynamicExt Methods - nvinfer1::IPluginV2DynamicExt* clone() const noexcept override; - nvinfer1::DimsExprs getOutputDimensions(int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; - bool supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept override; - void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept override; - size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept override; - int enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - - // IPluginV2Ext Methods - nvinfer1::DataType getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept override; - - // IPluginV2 Methods - char const* getPluginType() const noexcept override; - char const* getPluginVersion() const noexcept override; - int getNbOutputs() const noexcept override; - int initialize() noexcept override; - void terminate() noexcept override; - size_t getSerializationSize() const noexcept override; - void serialize(void* buffer) const noexcept override; - void destroy() noexcept override; - -private: - bool isCustomAllReduceSupported(int ranks_per_node) const noexcept; - void initGroupTopology() noexcept; - void setGroupTopology() noexcept; - tk::AllReduceStrategyType selectImplementation(size_t messageSize, int worldSize, nvinfer1::DataType type) noexcept; - void check() noexcept; - -private: - std::string const mLayerName; - std::set mGroup; - bool mIsNVLINKSupported; - bool mIsP2PSupported; - nvinfer1::DataType mType; - tk::AllReduceStrategyType mStrategy; - tk::AllReduceStrategyConfig mConfig; - tk::AllReduceFusionOp mOp; - float mEps; - std::shared_ptr mNcclComm; - int8_t mAffine; - int8_t mBias; - int8_t mScale; -}; - -class AllreducePluginCreator : public BaseCreator -{ -public: - AllreducePluginCreator(); - - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; - - nvinfer1::IPluginV2* deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept override; - -private: - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; -}; - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/ncclPlugin/recvPlugin.cpp b/cpp/tensorrt_llm/plugins/ncclPlugin/recvPlugin.cpp deleted file mode 100644 index 089ed31175b2..000000000000 --- a/cpp/tensorrt_llm/plugins/ncclPlugin/recvPlugin.cpp +++ /dev/null @@ -1,252 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. - */ -#include "recvPlugin.h" - -#include "tensorrt_llm/common/logger.h" -#include "tensorrt_llm/runtime/utils/mpiUtils.h" - -#include - -using namespace nvinfer1; -using tensorrt_llm::plugins::RecvPluginCreator; -using tensorrt_llm::plugins::RecvPlugin; -using tensorrt_llm::mpi::MpiTag; - -static char const* RECV_PLUGIN_VERSION{"1"}; -static char const* RECV_PLUGIN_NAME{"Recv"}; -PluginFieldCollection RecvPluginCreator::mFC{}; -std::vector RecvPluginCreator::mPluginAttributes; - -RecvPlugin::RecvPlugin(int srcRank, nvinfer1::DataType type) - : mSrcRank(srcRank) - , mType(type) -{ -} - -// Parameterized constructor -RecvPlugin::RecvPlugin(void const* data, size_t length) -{ - char const *d = reinterpret_cast(data), *a = d; - read(d, mType); - read(d, mSrcRank); - TLLM_CHECK_WITH_INFO(d == a + length, - "Expected length (%d) != real length (%d). This is often " - "caused by using different TensorRT LLM version to build " - "engine and run engine.", - (int) length, (int) (d - a)); -} - -// IPluginV2DynamicExt Methods -nvinfer1::IPluginV2DynamicExt* RecvPlugin::clone() const noexcept -{ - auto* plugin = new RecvPlugin(*this); - plugin->setPluginNamespace(mNamespace.c_str()); - return plugin; -} - -nvinfer1::DimsExprs RecvPlugin::getOutputDimensions( - int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept -{ - return inputs[0]; -} - -bool RecvPlugin::supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept -{ - return (inOut[pos].type == mType) && (inOut[pos].format == TensorFormat::kLINEAR); -} - -void RecvPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept -{ -} - -size_t RecvPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - return 0; -} - -int RecvPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept -{ - if (isBuilding()) - { - return 0; - } - size_t size = 1; - for (int i = 0; i < inputDesc[0].dims.nbDims; ++i) - { - size *= inputDesc[0].dims.d[i]; - } - TLLM_LOG_DEBUG("start ncclRecv with size %d", size); - NCCLCHECK(ncclRecv(outputs[0], size, (*getDtypeMap())[inputDesc[0].type], 0, mComm, stream)); - TLLM_LOG_DEBUG("end ncclRecv with size %d", size); - - return 0; -} - -// IPluginV2Ext Methods -nvinfer1::DataType RecvPlugin::getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept -{ - assert(index == 0); - return inputTypes[0]; -} - -// IPluginV2 Methods - -char const* RecvPlugin::getPluginType() const noexcept -{ - return RECV_PLUGIN_NAME; -} - -char const* RecvPlugin::getPluginVersion() const noexcept -{ - return RECV_PLUGIN_VERSION; -} - -int RecvPlugin::getNbOutputs() const noexcept -{ - return 1; -} - -int RecvPlugin::initialize() noexcept -{ - if (isBuilding()) - { - return 0; - } - ncclUniqueId id; - COMM_SESSION.recvValue(id, mSrcRank, MpiTag::kDefault); -// Need static connection initialization for accurate KV cache size estimation -#if defined(_WIN32) - if (getenv("NCCL_RUNTIME_CONNECT") == nullptr) - _putenv_s("NCCL_RUNTIME_CONNECT", "0"); -#else - setenv("NCCL_RUNTIME_CONNECT", "0", 0); -#endif // _WIN32 - NCCLCHECK(ncclCommInitRank(&mComm, 2, id, 1)); - return 0; -} - -void RecvPlugin::terminate() noexcept -{ - if (isBuilding()) - { - return; - } - NCCLCHECK(ncclCommDestroy(mComm)); -} - -size_t RecvPlugin::getSerializationSize() const noexcept -{ - return sizeof(mSrcRank) + sizeof(mType); -} - -void RecvPlugin::serialize(void* buffer) const noexcept -{ - char *d = static_cast(buffer), *a = d; - write(d, mType); - write(d, mSrcRank); - TLLM_CHECK(d == a + getSerializationSize()); -} - -void RecvPlugin::destroy() noexcept -{ - // This gets called when the network containing plugin is destroyed - delete this; -} - -/////////////// - -RecvPluginCreator::RecvPluginCreator() -{ - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - mPluginAttributes.emplace_back(PluginField("src_rank", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("type_id", nullptr, PluginFieldType::kINT32)); - - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -char const* RecvPluginCreator::getPluginName() const noexcept -{ - return RECV_PLUGIN_NAME; -} - -char const* RecvPluginCreator::getPluginVersion() const noexcept -{ - return RECV_PLUGIN_VERSION; -} - -PluginFieldCollection const* RecvPluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -IPluginV2* RecvPluginCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept -{ - PluginField const* fields = fc->fields; - int srcRank{}; - nvinfer1::DataType type{}; - // Read configurations from each fields - for (int i = 0; i < fc->nbFields; ++i) - { - char const* attrName = fields[i].name; - if (!strcmp(attrName, "src_rank")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - srcRank = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "type_id")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - type = static_cast(*(static_cast(fields[i].data))); - } - } - - try - { - auto* obj = new RecvPlugin(srcRank, type); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -IPluginV2* RecvPluginCreator::deserializePlugin(char const* name, void const* serialData, size_t serialLength) noexcept -{ - // This object will be deleted when the network is destroyed, which will - // call RecvPlugin::destroy() - try - { - auto* obj = new RecvPlugin(serialData, serialLength); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} diff --git a/cpp/tensorrt_llm/plugins/ncclPlugin/recvPlugin.h b/cpp/tensorrt_llm/plugins/ncclPlugin/recvPlugin.h deleted file mode 100644 index 5c8eedfb5218..000000000000 --- a/cpp/tensorrt_llm/plugins/ncclPlugin/recvPlugin.h +++ /dev/null @@ -1,90 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. - */ -#pragma once - -#include "tensorrt_llm/plugins/common/plugin.h" -#include -#include -#include - -namespace tensorrt_llm::plugins -{ - -class RecvPlugin : public BasePlugin -{ -public: - RecvPlugin(int srcRank, nvinfer1::DataType type); - - RecvPlugin(void const* data, size_t length); - - ~RecvPlugin() override = default; - - // IPluginV2DynamicExt Methods - nvinfer1::IPluginV2DynamicExt* clone() const noexcept override; - nvinfer1::DimsExprs getOutputDimensions(int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; - bool supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept override; - void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept override; - size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept override; - int enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - - // IPluginV2Ext Methods - nvinfer1::DataType getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept override; - - // IPluginV2 Methods - char const* getPluginType() const noexcept override; - char const* getPluginVersion() const noexcept override; - int getNbOutputs() const noexcept override; - int initialize() noexcept override; - void terminate() noexcept override; - size_t getSerializationSize() const noexcept override; - void serialize(void* buffer) const noexcept override; - void destroy() noexcept override; - -private: - ncclComm_t mComm; // TODO: Remove this - int mSrcRank; - nvinfer1::DataType mType; -}; - -class RecvPluginCreator : public BaseCreator -{ -public: - RecvPluginCreator(); - - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; - - nvinfer1::IPluginV2* deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept override; - -private: - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; -}; - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/ncclPlugin/reduceScatterPlugin.cpp b/cpp/tensorrt_llm/plugins/ncclPlugin/reduceScatterPlugin.cpp deleted file mode 100644 index fe17c44fc418..000000000000 --- a/cpp/tensorrt_llm/plugins/ncclPlugin/reduceScatterPlugin.cpp +++ /dev/null @@ -1,252 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. - */ -#include "reduceScatterPlugin.h" - -#include -#include - -using namespace nvinfer1; -using tensorrt_llm::plugins::ReduceScatterPluginCreator; -using tensorrt_llm::plugins::ReduceScatterPlugin; - -static char const* REDUCE_SCATTER_PLUGIN_VERSION{"1"}; -static char const* REDUCE_SCATTER_PLUGIN_NAME{"ReduceScatter"}; -PluginFieldCollection ReduceScatterPluginCreator::mFC{}; -std::vector ReduceScatterPluginCreator::mPluginAttributes; - -ReduceScatterPlugin::ReduceScatterPlugin(std::set group, nvinfer1::DataType type) - : mGroup(std::move(group)) - , mType(type) -{ -} - -// Parameterized constructor -ReduceScatterPlugin::ReduceScatterPlugin(void const* data, size_t length) -{ - char const *d = reinterpret_cast(data), *a = d; - read(d, mType); - mGroup.clear(); - int groupItem = 0; - while (d != a + length) - { - read(d, groupItem); - mGroup.insert(groupItem); - } - TLLM_CHECK_WITH_INFO(d == a + length, - "Expected length (%d) != real length (%d). This is often " - "caused by using different TensorRT LLM version to build " - "engine and run engine.", - (int) length, (int) (d - a)); -} - -// IPluginV2DynamicExt Methods -nvinfer1::IPluginV2DynamicExt* ReduceScatterPlugin::clone() const noexcept -{ - auto* plugin = new ReduceScatterPlugin(*this); - plugin->setPluginNamespace(mNamespace.c_str()); - return plugin; -} - -nvinfer1::DimsExprs ReduceScatterPlugin::getOutputDimensions( - int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept -{ - auto output = inputs[0]; - output.d[0] - = exprBuilder.operation(DimensionOperation::kFLOOR_DIV, *output.d[0], *exprBuilder.constant(mGroup.size())); - return output; -} - -bool ReduceScatterPlugin::supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept -{ - return (inOut[pos].type == mType) && (inOut[pos].format == TensorFormat::kLINEAR); -} - -void ReduceScatterPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept -{ -} - -size_t ReduceScatterPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - return 0; -} - -int ReduceScatterPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) noexcept -{ - if (isBuilding()) - { - return 0; - } - size_t size = 1; - for (int i = 0; i < outputDesc[0].dims.nbDims; ++i) - { - size *= outputDesc[0].dims.d[i]; - } - - TLLM_CHECK_WITH_INFO(mNcclComm.get() != nullptr, "mNcclComm should be initialized before used"); - NCCLCHECK(ncclReduceScatter( - inputs[0], outputs[0], size, (*getDtypeMap())[inputDesc[0].type], ncclSum, *mNcclComm, stream)); - - return 0; -} - -// IPluginV2Ext Methods -nvinfer1::DataType ReduceScatterPlugin::getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept -{ - assert(index == 0); - return inputTypes[0]; -} - -// IPluginV2 Methods - -char const* ReduceScatterPlugin::getPluginType() const noexcept -{ - return REDUCE_SCATTER_PLUGIN_NAME; -} - -char const* ReduceScatterPlugin::getPluginVersion() const noexcept -{ - return REDUCE_SCATTER_PLUGIN_VERSION; -} - -int ReduceScatterPlugin::getNbOutputs() const noexcept -{ - return 1; -} - -int ReduceScatterPlugin::initialize() noexcept -{ - if (isBuilding()) - { - return 0; - } - mNcclComm = getComm(mGroup); - return 0; -} - -void ReduceScatterPlugin::terminate() noexcept {} - -size_t ReduceScatterPlugin::getSerializationSize() const noexcept -{ - return sizeof(int) * mGroup.size() + sizeof(mType); -} - -void ReduceScatterPlugin::serialize(void* buffer) const noexcept -{ - char *d = static_cast(buffer), *a = d; - write(d, mType); - for (auto it = mGroup.begin(); it != mGroup.end(); ++it) - { - write(d, *it); - } - TLLM_CHECK(d == a + getSerializationSize()); -} - -void ReduceScatterPlugin::destroy() noexcept -{ - // This gets called when the network containing plugin is destroyed - delete this; -} - -/////////////// - -ReduceScatterPluginCreator::ReduceScatterPluginCreator() -{ - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - mPluginAttributes.emplace_back(PluginField("group", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("type_id", nullptr, PluginFieldType::kINT32)); - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -char const* ReduceScatterPluginCreator::getPluginName() const noexcept -{ - return REDUCE_SCATTER_PLUGIN_NAME; -} - -char const* ReduceScatterPluginCreator::getPluginVersion() const noexcept -{ - return REDUCE_SCATTER_PLUGIN_VERSION; -} - -PluginFieldCollection const* ReduceScatterPluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -IPluginV2* ReduceScatterPluginCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept -{ - PluginField const* fields = fc->fields; - std::set group; - nvinfer1::DataType type{}; - // Read configurations from each fields - for (int i = 0; i < fc->nbFields; ++i) - { - char const* attrName = fields[i].name; - if (!strcmp(attrName, "group")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - auto const* r = static_cast(fields[i].data); - for (int j = 0; j < fields[i].length; ++j) - { - group.insert(*r); - ++r; - } - } - else if (!strcmp(attrName, "type_id")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - type = static_cast(*(static_cast(fields[i].data))); - } - } - - try - { - auto* obj = new ReduceScatterPlugin(group, type); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -IPluginV2* ReduceScatterPluginCreator::deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept -{ - // This object will be deleted when the network is destroyed, which will - // call ReduceScatterPlugin::destroy() - try - { - auto* obj = new ReduceScatterPlugin(serialData, serialLength); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} diff --git a/cpp/tensorrt_llm/plugins/ncclPlugin/reduceScatterPlugin.h b/cpp/tensorrt_llm/plugins/ncclPlugin/reduceScatterPlugin.h deleted file mode 100644 index c630b57a2b98..000000000000 --- a/cpp/tensorrt_llm/plugins/ncclPlugin/reduceScatterPlugin.h +++ /dev/null @@ -1,91 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. - */ -#pragma once - -#include "tensorrt_llm/plugins/common/plugin.h" -#include -#include -#include - -namespace tensorrt_llm::plugins -{ - -class ReduceScatterPlugin : public BasePlugin -{ -public: - ReduceScatterPlugin(std::set group, nvinfer1::DataType type); - - ReduceScatterPlugin(void const* data, size_t length); - - ~ReduceScatterPlugin() override = default; - - // IPluginV2DynamicExt Methods - nvinfer1::IPluginV2DynamicExt* clone() const noexcept override; - nvinfer1::DimsExprs getOutputDimensions(int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; - bool supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept override; - void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept override; - size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept override; - int enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - - // IPluginV2Ext Methods - nvinfer1::DataType getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept override; - - // IPluginV2 Methods - char const* getPluginType() const noexcept override; - char const* getPluginVersion() const noexcept override; - int getNbOutputs() const noexcept override; - int initialize() noexcept override; - void terminate() noexcept override; - size_t getSerializationSize() const noexcept override; - void serialize(void* buffer) const noexcept override; - void destroy() noexcept override; - -private: - const std::string mLayerName; - std::set mGroup; - nvinfer1::DataType mType; - std::shared_ptr mNcclComm; -}; - -class ReduceScatterPluginCreator : public BaseCreator -{ -public: - ReduceScatterPluginCreator(); - - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; - - nvinfer1::IPluginV2* deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept override; - -private: - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; -}; - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/ncclPlugin/sendPlugin.cpp b/cpp/tensorrt_llm/plugins/ncclPlugin/sendPlugin.cpp deleted file mode 100644 index 81d66aa8211e..000000000000 --- a/cpp/tensorrt_llm/plugins/ncclPlugin/sendPlugin.cpp +++ /dev/null @@ -1,255 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. - */ -#include "sendPlugin.h" - -#include "tensorrt_llm/common/logger.h" -#include "tensorrt_llm/runtime/utils/mpiUtils.h" - -#include -#include - -using namespace nvinfer1; -using tensorrt_llm::plugins::SendPluginCreator; -using tensorrt_llm::plugins::SendPlugin; -using tensorrt_llm::mpi::MpiTag; - -static char const* SEND_PLUGIN_VERSION{"1"}; -static char const* SEND_PLUGIN_NAME{"Send"}; -PluginFieldCollection SendPluginCreator::mFC{}; -std::vector SendPluginCreator::mPluginAttributes; - -SendPlugin::SendPlugin(int tgtRank, nvinfer1::DataType type) - : mTgtRank(tgtRank) - , mType(type) -{ -} - -// Parameterized constructor -SendPlugin::SendPlugin(void const* data, size_t length) -{ - char const *d = reinterpret_cast(data), *a = d; - read(d, mType); - read(d, mTgtRank); - TLLM_CHECK_WITH_INFO(d == a + length, - "Expected length (%d) != real length (%d). This is often " - "caused by using different TensorRT LLM version to build " - "engine and run engine.", - (int) length, (int) (d - a)); -} - -// IPluginV2DynamicExt Methods -nvinfer1::IPluginV2DynamicExt* SendPlugin::clone() const noexcept -{ - auto* plugin = new SendPlugin(*this); - plugin->setPluginNamespace(mNamespace.c_str()); - return plugin; -} - -nvinfer1::DimsExprs SendPlugin::getOutputDimensions( - int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept -{ - return inputs[0]; -} - -bool SendPlugin::supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept -{ - return (inOut[pos].type == mType) && (inOut[pos].format == TensorFormat::kLINEAR); -} - -void SendPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept -{ -} - -size_t SendPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - return 0; -} - -int SendPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept -{ - if (isBuilding()) - { - return 0; - } - size_t size = 1; - for (int i = 0; i < inputDesc[0].dims.nbDims; ++i) - { - size *= inputDesc[0].dims.d[i]; - } - - TLLM_LOG_DEBUG("start ncclSend with size %d", size); - NCCLCHECK(ncclSend(inputs[0], size, (*getDtypeMap())[inputDesc[0].type], 1, mComm, stream)); - TLLM_LOG_DEBUG("end ncclSend with size %d", size); - return 0; -} - -// IPluginV2Ext Methods -nvinfer1::DataType SendPlugin::getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept -{ - assert(index == 0); - return inputTypes[0]; -} - -// IPluginV2 Methods - -char const* SendPlugin::getPluginType() const noexcept -{ - return SEND_PLUGIN_NAME; -} - -char const* SendPlugin::getPluginVersion() const noexcept -{ - return SEND_PLUGIN_VERSION; -} - -int SendPlugin::getNbOutputs() const noexcept -{ - return 1; -} - -int SendPlugin::initialize() noexcept -{ - if (isBuilding()) - { - return 0; - } - - ncclUniqueId id; - ncclGetUniqueId(&id); - COMM_SESSION.sendValue(id, mTgtRank, MpiTag::kDefault); -// Need static connection initialization for accurate KV cache size estimation -#if defined(_WIN32) - if (getenv("NCCL_RUNTIME_CONNECT") == nullptr) - _putenv_s("NCCL_RUNTIME_CONNECT", "0"); -#else - setenv("NCCL_RUNTIME_CONNECT", "0", 0); -#endif // _WIN32 - NCCLCHECK(ncclCommInitRank(&mComm, 2, id, 0)); - return 0; -} - -void SendPlugin::terminate() noexcept -{ - if (isBuilding()) - { - return; - } - NCCLCHECK(ncclCommDestroy(mComm)); -} - -size_t SendPlugin::getSerializationSize() const noexcept -{ - return sizeof(mTgtRank) + sizeof(mType); -} - -void SendPlugin::serialize(void* buffer) const noexcept -{ - char *d = static_cast(buffer), *a = d; - write(d, mType); - write(d, mTgtRank); - TLLM_CHECK(d == a + getSerializationSize()); -} - -void SendPlugin::destroy() noexcept -{ - // This gets called when the network containing plugin is destroyed - delete this; -} - -/////////////// - -SendPluginCreator::SendPluginCreator() -{ - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - mPluginAttributes.emplace_back(PluginField("tgt_rank", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("type_id", nullptr, PluginFieldType::kINT32)); - - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -char const* SendPluginCreator::getPluginName() const noexcept -{ - return SEND_PLUGIN_NAME; -} - -char const* SendPluginCreator::getPluginVersion() const noexcept -{ - return SEND_PLUGIN_VERSION; -} - -PluginFieldCollection const* SendPluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -IPluginV2* SendPluginCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept -{ - PluginField const* fields = fc->fields; - int tgtRank{}; - nvinfer1::DataType type{}; - // Read configurations from each fields - for (int i = 0; i < fc->nbFields; ++i) - { - char const* attrName = fields[i].name; - if (!strcmp(attrName, "tgt_rank")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - tgtRank = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "type_id")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - type = static_cast(*(static_cast(fields[i].data))); - } - } - - try - { - auto* obj = new SendPlugin(tgtRank, type); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -IPluginV2* SendPluginCreator::deserializePlugin(char const* name, void const* serialData, size_t serialLength) noexcept -{ - // This object will be deleted when the network is destroyed, which will - // call SendPlugin::destroy() - try - { - auto* obj = new SendPlugin(serialData, serialLength); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} diff --git a/cpp/tensorrt_llm/plugins/ncclPlugin/sendPlugin.h b/cpp/tensorrt_llm/plugins/ncclPlugin/sendPlugin.h deleted file mode 100644 index 0d36b0ebff28..000000000000 --- a/cpp/tensorrt_llm/plugins/ncclPlugin/sendPlugin.h +++ /dev/null @@ -1,89 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. - */ -#pragma once - -#include "tensorrt_llm/plugins/common/plugin.h" -#include -#include - -namespace tensorrt_llm::plugins -{ - -class SendPlugin : public BasePlugin -{ -public: - SendPlugin(int tgtRank, nvinfer1::DataType type); - - SendPlugin(void const* data, size_t length); - - ~SendPlugin() override = default; - - // IPluginV2DynamicExt Methods - nvinfer1::IPluginV2DynamicExt* clone() const noexcept override; - nvinfer1::DimsExprs getOutputDimensions(int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; - bool supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept override; - void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept override; - size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept override; - int enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - - // IPluginV2Ext Methods - nvinfer1::DataType getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept override; - - // IPluginV2 Methods - char const* getPluginType() const noexcept override; - char const* getPluginVersion() const noexcept override; - int getNbOutputs() const noexcept override; - int initialize() noexcept override; - void terminate() noexcept override; - size_t getSerializationSize() const noexcept override; - void serialize(void* buffer) const noexcept override; - void destroy() noexcept override; - -private: - ncclComm_t mComm; // TODO: Remove this - int mTgtRank; - nvinfer1::DataType mType; -}; - -class SendPluginCreator : public BaseCreator -{ -public: - SendPluginCreator(); - - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; - - nvinfer1::IPluginV2* deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept override; - -private: - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; -}; - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/qserveGemmPlugin/CMakeLists.txt b/cpp/tensorrt_llm/plugins/qserveGemmPlugin/CMakeLists.txt deleted file mode 100755 index 86876224fccd..000000000000 --- a/cpp/tensorrt_llm/plugins/qserveGemmPlugin/CMakeLists.txt +++ /dev/null @@ -1,21 +0,0 @@ -# -# SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. -# -file(GLOB SRCS *.cpp) -set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS}) -set(PLUGIN_SOURCES - ${PLUGIN_SOURCES} - PARENT_SCOPE) diff --git a/cpp/tensorrt_llm/plugins/qserveGemmPlugin/qserveGemmPlugin.cpp b/cpp/tensorrt_llm/plugins/qserveGemmPlugin/qserveGemmPlugin.cpp deleted file mode 100644 index 166f1cc32cbe..000000000000 --- a/cpp/tensorrt_llm/plugins/qserveGemmPlugin/qserveGemmPlugin.cpp +++ /dev/null @@ -1,416 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. - */ - -#include "qserveGemmPlugin.h" -#include "tensorrt_llm/kernels/qserveGemm.h" -#include -#include - -using namespace nvinfer1; -using namespace tensorrt_llm::common; -using tensorrt_llm::plugins::QServeGemmPluginCreator; -using tensorrt_llm::plugins::QServeGemmPlugin; -using tensorrt_llm::plugins::read; -using tensorrt_llm::plugins::write; -using namespace tensorrt_llm::kernels::qserve; - -static char const* QSERVE_GEMM_PLUGIN_VERSION{"1"}; -static char const* QSERVE_GEMM_PLUGIN_NAME{"QServeGemm"}; - -PluginFieldCollection QServeGemmPluginCreator::mFC{}; -std::vector QServeGemmPluginCreator::mPluginAttributes; - -namespace tensorrt_llm::plugins -{ - -QServeGemmPlugin::QServeGemmPlugin( - // QuantMode quantMode, - nvinfer1::DataType dtype, int groupSize) -{ - init(dtype, groupSize); -} - -QServeGemmPlugin::QServeGemmPlugin(void const* data, size_t length) -{ - char const *d = reinterpret_cast(data), *a = d; - - nvinfer1::DataType type; - unsigned int quantMode; - int groupSize; - - read(d, quantMode); - read(d, type); - read(d, groupSize); - - read(d, mDims); - - // mQuantMode = QuantMode(quantMode); - - init(type, groupSize); - - TLLM_CHECK_WITH_INFO(d == a + length, - "Expected length (%d) != real length (%d). This is often " - "caused by using different TensorRT LLM version to build " - "engine and run engine.", - (int) length, (int) (d - a)); -} - -void QServeGemmPlugin::init(nvinfer1::DataType dtype, int groupSize) -{ - if (groupSize <= 0) - groupSize = -1; // Per-channel - mGroupSize = groupSize; - mType = dtype; - mRunner = std::make_shared(); -} - -// IPluginV2DynamicExt Methods -nvinfer1::IPluginV2DynamicExt* QServeGemmPlugin::clone() const noexcept -{ - auto* plugin = new QServeGemmPlugin(*this); - return plugin; -} - -nvinfer1::DimsExprs QServeGemmPlugin::getOutputDimensions( - int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept -{ - try - { - TLLM_CHECK(nbInputs == 6); - TLLM_CHECK(outputIndex == 0); - int const nbDimsA = inputs[0].nbDims; - TLLM_CHECK(nbDimsA >= 2); - DimsExprs ret; - ret.nbDims = nbDimsA; - for (int ii = 0; ii < nbDimsA - 1; ++ii) - { - ret.d[ii] = inputs[0].d[ii]; - } - ret.d[nbDimsA - 1] = inputs[1].d[0]; - return ret; - } - catch (std::exception const& e) - { - caughtError(e); - } - return DimsExprs{}; -} - -bool QServeGemmPlugin::supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept -{ - if (mGroupSize != -1) - { // Per-group - switch (pos) - { - case 0: - // activation - return inOut[pos].type == nvinfer1::DataType::kINT8 && inOut[pos].format == TensorFormat::kLINEAR; - case 1: - // uint4 weights packed in int8 - return inOut[pos].type == nvinfer1::DataType::kINT8 && inOut[pos].format == TensorFormat::kLINEAR; - case 2: - // int8 weight s2_zeros - return inOut[pos].type == nvinfer1::DataType::kINT8 && inOut[pos].format == TensorFormat::kLINEAR; - case 3: - // int8 weight s2_scales - return inOut[pos].type == nvinfer1::DataType::kINT8 && inOut[pos].format == TensorFormat::kLINEAR; - case 4: - // fp16 weight s1_scales - return inOut[pos].type == nvinfer1::DataType::kHALF && inOut[pos].format == TensorFormat::kLINEAR; - case 5: - // fp16 activation scales - return inOut[pos].type == nvinfer1::DataType::kHALF && inOut[pos].format == TensorFormat::kLINEAR; - case 6: - // fp16 output activation - return inOut[pos].type == nvinfer1::DataType::kHALF && inOut[pos].format == TensorFormat::kLINEAR; - default: return false; - } - } - - else - { // Per-channel - switch (pos) - { - case 0: - // activation - return inOut[pos].type == nvinfer1::DataType::kINT8 && inOut[pos].format == TensorFormat::kLINEAR; - case 1: - // uint4 weights packed in int8 - return inOut[pos].type == nvinfer1::DataType::kINT8 && inOut[pos].format == TensorFormat::kLINEAR; - case 2: - // fp16 s1_scales - return inOut[pos].type == nvinfer1::DataType::kHALF && inOut[pos].format == TensorFormat::kLINEAR; - case 3: - // fp16 s1_szeros - return inOut[pos].type == nvinfer1::DataType::kHALF && inOut[pos].format == TensorFormat::kLINEAR; - case 4: - // fp16 act_sums - return inOut[pos].type == nvinfer1::DataType::kHALF && inOut[pos].format == TensorFormat::kLINEAR; - case 5: - // fp16 act_scales - return inOut[pos].type == nvinfer1::DataType::kHALF && inOut[pos].format == TensorFormat::kLINEAR; - case 6: - // fp16 output activation - return inOut[pos].type == nvinfer1::DataType::kHALF && inOut[pos].format == TensorFormat::kLINEAR; - default: return false; - } - } -} - -void QServeGemmPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept -{ - auto const minM = std::accumulate(in[0].min.d, in[0].min.d + in[0].min.nbDims - 1, 1, std::multiplies()); - auto const maxM = std::accumulate(in[0].max.d, in[0].max.d + in[0].max.nbDims - 1, 1, std::multiplies()); - - int const maxK = in[0].max.d[in[0].max.nbDims - 1]; - int const maxN = in[1].max.d[0]; - int const minK = in[0].min.d[in[0].min.nbDims - 1]; - int const minN = in[1].min.d[0]; - - TLLM_CHECK_WITH_INFO(minN == maxN, "Variable out channels is not allowed"); - TLLM_CHECK_WITH_INFO(minK == maxK, "Variable in channels is not allowed"); - - if (!mDims.isInitialized()) - { - mDims = {minM, maxM, maxN, maxK}; - } - m_workspaceMaxSize = mRunner->getWorkspaceSize(maxM, maxN, maxK); -} - -size_t QServeGemmPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - return m_workspaceMaxSize; -} - -int QServeGemmPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept -{ - // inputs - - // Per group: - // activation [M, K] int8_t Quantized sint8 activations - // weights [N, K/2] int8_t Quantized uint4 weights (packed as int8_t) - // s2_zeros [K/group_size, N] int8_t Level-2 sint8 scaled zeros of weights - // s2_scales [K/group_size, N] int8_t Level-2 sint8 scales of weights - // s1_scales [N] half Level-1 fp16 scales of weights - // act_scales [M] half Scales of activations - - // Per channel: - // activation [M, K] int8_t Quantized sint8 activations - // weights [N, K/2] int8_t Quantized uint4 weights (packed as int8_t) - // s1_scales [N] half Level-1 scales of weights - // s1_szeros [N] half Level-1 scaled zeros of weights - // act_sums [M] half Per-token sums of activations - // act_scales [M] half Scales of activations - - // outputs - // mat [M(*), N] half - - int64_t m64 = 1; - for (int ii = 0; ii < inputDesc[0].dims.nbDims - 1; ++ii) - { - m64 *= inputDesc[0].dims.d[ii]; - } - int const m = TLLM_INT32_CAST(m64); - int const n = TLLM_INT32_CAST(inputDesc[1].dims.d[0]); - int const k = TLLM_INT32_CAST(inputDesc[0].dims.d[inputDesc[0].dims.nbDims - 1]); - - // TODO: Implement optimized kernels if (m <= 4) - - if (mGroupSize != -1) - { - ParamsPerGroup params = {reinterpret_cast(inputs[0]), // A - reinterpret_cast(inputs[1]), // B - reinterpret_cast(inputs[2]), // s2_zeros - reinterpret_cast(inputs[3]), // s2_scales - reinterpret_cast(inputs[4]), // s1_scales - reinterpret_cast(inputs[5]), // act_scales - reinterpret_cast(outputs[0]), // C - m, n, k}; - mRunner->gemmPerGroup(params, stream); - } - else - { - ParamsPerChannel params = {reinterpret_cast(inputs[0]), // A - reinterpret_cast(inputs[1]), // B - reinterpret_cast(inputs[2]), // s1_scales - reinterpret_cast(inputs[3]), // s1_szeros - reinterpret_cast(inputs[4]), // act_sums - reinterpret_cast(inputs[5]), // act_scales - reinterpret_cast(outputs[0]), // C - m, n, k}; - mRunner->gemmPerChannel(params, stream); - } - - return 0; -} - -// IPluginV2Ext Methods -nvinfer1::DataType QServeGemmPlugin::getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept -{ - TLLM_CHECK(index == 0); - return mType; -} - -// IPluginV2 Methods - -char const* QServeGemmPlugin::getPluginType() const noexcept -{ - return QSERVE_GEMM_PLUGIN_NAME; -} - -char const* QServeGemmPlugin::getPluginVersion() const noexcept -{ - return QSERVE_GEMM_PLUGIN_VERSION; -} - -int QServeGemmPlugin::getNbOutputs() const noexcept -{ - return 1; -} - -int QServeGemmPlugin::initialize() noexcept -{ - configGemm(); - return 0; -} - -void QServeGemmPlugin::terminate() noexcept {} - -size_t QServeGemmPlugin::getSerializationSize() const noexcept -{ - return sizeof(mQuantMode) + // QuantMode - sizeof(mType) + // dtype - sizeof(mGroupSize) + // GroupSize - sizeof(mDims); // Dimensions -} - -void QServeGemmPlugin::serialize(void* buffer) const noexcept -{ - char *d = static_cast(buffer), *a = d; - write(d, mQuantMode.value()); - write(d, mType); - write(d, mGroupSize); - write(d, mDims); - - TLLM_CHECK(d == a + getSerializationSize()); -} - -void QServeGemmPlugin::destroy() noexcept -{ - // This gets called when the network containing plugin is destroyed - delete this; -} - -void QServeGemmPlugin::configGemm() {} - -/////////////// - -QServeGemmPluginCreator::QServeGemmPluginCreator() -{ - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - mPluginAttributes.push_back(PluginField("type_id", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.push_back(PluginField("group_size", nullptr, PluginFieldType::kINT32)); - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -char const* QServeGemmPluginCreator::getPluginName() const noexcept -{ - return QSERVE_GEMM_PLUGIN_NAME; -} - -char const* QServeGemmPluginCreator::getPluginVersion() const noexcept -{ - return QSERVE_GEMM_PLUGIN_VERSION; -} - -PluginFieldCollection const* QServeGemmPluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -IPluginV2* QServeGemmPluginCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept -{ - // We do not use any fields for now. - - PluginField const* fields = fc->fields; - - // bool perTokenScaling, perChannelScaling; - DataType dtype{}; - int group_size = -1; - // Read configurations from each fields - for (int i = 0; i < fc->nbFields; ++i) - { - char const* attrName = fields[i].name; - if (!strcmp(attrName, "type_id")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - dtype = static_cast(*(static_cast(fields[i].data))); - // Only supports fp16 for now. - assert(dtype == nvinfer1::DataType::kHALF); - } - else if (!strcmp(attrName, "group_size")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - group_size = *static_cast(fields[i].data); - // Currently only support per-channel or g128. - assert(group_size == -1 || group_size == 128); - } - } - try - { - // QServeGemmPluginCreator is unique and shared for an engine generation - // Create plugin profiler with shared tactics map - // auto pluginProfiler = gemmPluginProfileManager.createGemmPluginProfiler(/* inference */ false); - // QuantMode quantMode = QuantMode::fromQuantAlgo("W4A8_QSERVE"); - auto* obj = new QServeGemmPlugin(dtype, group_size); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -IPluginV2* QServeGemmPluginCreator::deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept -{ - // This object will be deleted when the network is destroyed, which will - // call QServeGemmPlugin::destroy() - try - { - // Create plugin profiler with private tactics map which is read from the serialized engine - // auto pluginProfiler = gemmPluginProfileManager.createGemmPluginProfiler(/* inference */ true); - auto* obj = new QServeGemmPlugin(serialData, serialLength); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/qserveGemmPlugin/qserveGemmPlugin.h b/cpp/tensorrt_llm/plugins/qserveGemmPlugin/qserveGemmPlugin.h deleted file mode 100644 index 086460863c4f..000000000000 --- a/cpp/tensorrt_llm/plugins/qserveGemmPlugin/qserveGemmPlugin.h +++ /dev/null @@ -1,112 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. - */ -#pragma once - -#include "tensorrt_llm/common/quantization.h" -#include "tensorrt_llm/plugins/common/gemmPluginProfiler.h" -#include "tensorrt_llm/plugins/common/plugin.h" -#include -#include -#include - -namespace tensorrt_llm::plugins -{ - -using QServeGemmRunnerPtr = std::shared_ptr; - -class QServeGemmPlugin : public BasePlugin -{ -public: - // using PluginProfilerPtr = std::shared_ptr; - - QServeGemmPlugin(void const* data, size_t length); - - QServeGemmPlugin(nvinfer1::DataType dtype, int groupSize); - - ~QServeGemmPlugin() override = default; - - // IPluginV2DynamicExt Methods - nvinfer1::IPluginV2DynamicExt* clone() const noexcept override; - nvinfer1::DimsExprs getOutputDimensions(int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; - bool supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept override; - void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept override; - size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept override; - int enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - - // IPluginV2Ext Methods - nvinfer1::DataType getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept override; - - // IPluginV2 Methods - char const* getPluginType() const noexcept override; - char const* getPluginVersion() const noexcept override; - int getNbOutputs() const noexcept override; - int initialize() noexcept override; - void terminate() noexcept override; - size_t getSerializationSize() const noexcept override; - void serialize(void* buffer) const noexcept override; - void destroy() noexcept override; - -private: - void init(nvinfer1::DataType dtype, int groupSize); - - void configGemm(); - - std::string const mLayerName; - - QServeGemmRunnerPtr mRunner; - - tensorrt_llm::common::QuantMode mQuantMode; // Not used for now - GemmDims mDims{}; - - size_t m_workspaceMaxSize; - - // Only supports fp16 output for now. - nvinfer1::DataType mType; - - int mGroupSize; -}; - -class QServeGemmPluginCreator : public BaseCreator -{ -public: - QServeGemmPluginCreator(); - - QServeGemmPluginCreator(void const* data, size_t length); - - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; - - nvinfer1::IPluginV2* deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept override; - -private: - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; -}; - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/quantizePerTokenPlugin/CMakeLists.txt b/cpp/tensorrt_llm/plugins/quantizePerTokenPlugin/CMakeLists.txt deleted file mode 100755 index 86876224fccd..000000000000 --- a/cpp/tensorrt_llm/plugins/quantizePerTokenPlugin/CMakeLists.txt +++ /dev/null @@ -1,21 +0,0 @@ -# -# SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. -# -file(GLOB SRCS *.cpp) -set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS}) -set(PLUGIN_SOURCES - ${PLUGIN_SOURCES} - PARENT_SCOPE) diff --git a/cpp/tensorrt_llm/plugins/quantizePerTokenPlugin/quantizePerTokenPlugin.cpp b/cpp/tensorrt_llm/plugins/quantizePerTokenPlugin/quantizePerTokenPlugin.cpp deleted file mode 100644 index 23d0b80390e3..000000000000 --- a/cpp/tensorrt_llm/plugins/quantizePerTokenPlugin/quantizePerTokenPlugin.cpp +++ /dev/null @@ -1,353 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. - */ -#include "quantizePerTokenPlugin.h" -#include "tensorrt_llm/kernels/quantization.h" - -using namespace nvinfer1; -using namespace tensorrt_llm::common; -using namespace tensorrt_llm::kernels; -using tensorrt_llm::plugins::QuantizePerTokenPluginCreator; -using tensorrt_llm::plugins::QuantizePerTokenPlugin; - -static char const* QUANTIZE_PER_TOKEN_PLUGIN_VERSION{"1"}; -static char const* QUANTIZE_PER_TOKEN_PLUGIN_NAME{"QuantizePerToken"}; -PluginFieldCollection QuantizePerTokenPluginCreator::mFC{}; -std::vector QuantizePerTokenPluginCreator::mPluginAttributes; - -QuantizePerTokenPlugin::QuantizePerTokenPlugin( - nvinfer1::DataType outputType, QuantMode quantMode, bool clampValEnabled, bool sumPerToken) - : mOutputType{outputType} - , mQuantMode{quantMode} - , mClampValEnabled{clampValEnabled} - , mSumPerToken{sumPerToken} -{ - TLLM_CHECK_WITH_INFO(mOutputType == nvinfer1::DataType::kINT8 || mOutputType == nvinfer1::DataType::kFP8, - "Only int8 or fp8 output type is allowed."); - // Check if the quant mode is valid. - TLLM_CHECK_WITH_INFO(mQuantMode.hasPerTokenScaling(), "The quant mode is not valid."); -} - -// Parameterized constructor -QuantizePerTokenPlugin::QuantizePerTokenPlugin(void const* data, size_t length) -{ - char const *d = reinterpret_cast(data), *a = d; - read(d, mOutputType); - read(d, mQuantMode); - read(d, mClampValEnabled); - read(d, mSumPerToken); - TLLM_CHECK_WITH_INFO(d == a + length, - "Expected length (%d) != real length (%d). This is often " - "caused by using different TensorRT LLM version to build " - "engine and run engine.", - (int) length, (int) (d - a)); -} - -// IPluginV2DynamicExt Methods -nvinfer1::IPluginV2DynamicExt* QuantizePerTokenPlugin::clone() const noexcept -{ - auto* plugin = new QuantizePerTokenPlugin(mOutputType, mQuantMode, mClampValEnabled, mSumPerToken); - plugin->setPluginNamespace(mNamespace.c_str()); - return plugin; -} - -nvinfer1::DimsExprs QuantizePerTokenPlugin::getOutputDimensions( - int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept -{ - try - { - TLLM_CHECK(nbInputs <= 2); - TLLM_CHECK(outputIndex <= 2); - if (outputIndex == 2) - { - // Per token sums. - TLLM_CHECK(mSumPerToken); - } - - if (outputIndex == 0) - { - // Quantized input - return inputs[0]; - } - - DimsExprs ret; - ret.nbDims = inputs[0].nbDims; - for (int ii = 0; ii < ret.nbDims - 1; ++ii) - { - ret.d[ii] = inputs[0].d[ii]; - } - ret.d[ret.nbDims - 1] = exprBuilder.constant(1); - // [M(*), 1] dynamic per token scales or sums - return ret; - } - catch (std::exception const& e) - { - caughtError(e); - } - return DimsExprs{}; -} - -bool QuantizePerTokenPlugin::supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept -{ - if (pos == 0) - { - // activation - return (inOut[pos].type == nvinfer1::DataType::kFLOAT || inOut[pos].type == nvinfer1::DataType::kHALF -#ifdef ENABLE_BF16 - || inOut[pos].type == nvinfer1::DataType::kBF16 -#endif - ) - && inOut[pos].format == TensorFormat::kLINEAR; - } - else if (pos == 1 && mClampValEnabled) - { - // clamp_max_v - return inOut[pos].type == nvinfer1::DataType::kFLOAT && inOut[pos].format == TensorFormat::kLINEAR; - } - else if (pos == 1 + int(mClampValEnabled)) - { - // quantized activation - return inOut[pos].type == mOutputType && inOut[pos].format == TensorFormat::kLINEAR; - } - else if (pos == 2 + int(mClampValEnabled)) - { - // scales - return inOut[pos].type == nvinfer1::DataType::kFLOAT && inOut[pos].format == TensorFormat::kLINEAR; - } - else if (pos == 3 + int(mClampValEnabled)) - { - TLLM_CHECK(mSumPerToken); - // per-token sums - return inOut[pos].type == nvinfer1::DataType::kFLOAT && inOut[pos].format == TensorFormat::kLINEAR; - } - - // Never should be here - assert(false); - return false; -} - -void QuantizePerTokenPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept -{ -} - -size_t QuantizePerTokenPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - return 0; -} - -template -void QuantizePerTokenPlugin::dispatchDataType(void* output, void const* input, void const* clampValPtr, void* scalePtr, - void* sumPtr, int dim0, int dim1, cudaStream_t stream) noexcept -{ - // inputs - // activation [dim0(*), dim1] - // clamp_value [2], contains min val, and max val (optional) - // outputs - // quant [dim0(*), dim1] - // scale_tokens [dim0(*), 1] - - invokePerTokenQuantization(reinterpret_cast(output), reinterpret_cast(input), dim0, dim1, - reinterpret_cast(clampValPtr), reinterpret_cast(scalePtr), - reinterpret_cast(sumPtr), mQuantMode, stream); -} - -int QuantizePerTokenPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) noexcept -{ - // inputs - // activation [M(*), K] - // clamp_value [2], contains min val, and max val (optional) - // outputs - // quant [M(*), K] Quantized activations. - // scale_tokens [M(*), 1] Per-token scales. - // token_sums [M(*), 1] (Optional) Per-token sums of all the channels (before quantization). - - int64_t m = 1; - for (int ii = 0; ii < inputDesc[0].dims.nbDims - 1; ++ii) - { - m *= inputDesc[0].dims.d[ii]; - } - int64_t const k = inputDesc[0].dims.d[inputDesc[0].dims.nbDims - 1]; - - void const* clampValPtr = mClampValEnabled ? inputs[1] : nullptr; - void* sumPtr = mSumPerToken ? outputs[2] : nullptr; - - if (inputDesc[0].type == DataType::kFLOAT && mOutputType == DataType::kINT8) - { - dispatchDataType(outputs[0], inputs[0], clampValPtr, outputs[1], sumPtr, m, k, stream); - } -#ifdef ENABLE_FP8 - else if (inputDesc[0].type == DataType::kFLOAT && mOutputType == DataType::kFP8) - { - dispatchDataType(outputs[0], inputs[0], clampValPtr, outputs[1], sumPtr, m, k, stream); - } -#endif // ENABLE_FP8 - else if (inputDesc[0].type == DataType::kHALF && mOutputType == DataType::kINT8) - { - dispatchDataType(outputs[0], inputs[0], clampValPtr, outputs[1], sumPtr, m, k, stream); - } -#ifdef ENABLE_FP8 - else if (inputDesc[0].type == DataType::kHALF && mOutputType == DataType::kFP8) - { - dispatchDataType(outputs[0], inputs[0], clampValPtr, outputs[1], sumPtr, m, k, stream); - } -#endif // ENABLE_FP8 -#ifdef ENABLE_BF16 - else if (inputDesc[0].type == DataType::kBF16 && mOutputType == DataType::kINT8) - { - dispatchDataType<__nv_bfloat16, int8_t>(outputs[0], inputs[0], clampValPtr, outputs[1], sumPtr, m, k, stream); - } -#ifdef ENABLE_FP8 - else if (inputDesc[0].type == DataType::kBF16 && mOutputType == DataType::kFP8) - { - dispatchDataType<__nv_bfloat16, __nv_fp8_e4m3>( - outputs[0], inputs[0], clampValPtr, outputs[1], sumPtr, m, k, stream); - } -#endif // ENABLE_FP8 -#endif // ENABLE_BF16 - sync_check_cuda_error(stream); - return 0; -} - -// IPluginV2Ext Methods -nvinfer1::DataType QuantizePerTokenPlugin::getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept -{ - TLLM_CHECK(nbInputs >= 1); - TLLM_CHECK(index <= 2); - if (index == 2) - { - // Per token sums. - TLLM_CHECK(mSumPerToken); - } - return index == 0 ? mOutputType : nvinfer1::DataType::kFLOAT; -} - -// IPluginV2 Methods - -char const* QuantizePerTokenPlugin::getPluginType() const noexcept -{ - return QUANTIZE_PER_TOKEN_PLUGIN_NAME; -} - -char const* QuantizePerTokenPlugin::getPluginVersion() const noexcept -{ - return QUANTIZE_PER_TOKEN_PLUGIN_VERSION; -} - -int QuantizePerTokenPlugin::getNbOutputs() const noexcept -{ - return 2 + static_cast(mSumPerToken); -} - -int QuantizePerTokenPlugin::initialize() noexcept -{ - return 0; -} - -void QuantizePerTokenPlugin::terminate() noexcept {} - -size_t QuantizePerTokenPlugin::getSerializationSize() const noexcept -{ - return sizeof(mOutputType) + sizeof(mQuantMode) + sizeof(mClampValEnabled) + sizeof(mSumPerToken); -} - -void QuantizePerTokenPlugin::serialize(void* buffer) const noexcept -{ - char *d = static_cast(buffer), *a = d; - write(d, mOutputType); - write(d, mQuantMode); - write(d, mClampValEnabled); - write(d, mSumPerToken); - TLLM_CHECK(d == a + getSerializationSize()); -} - -void QuantizePerTokenPlugin::destroy() noexcept -{ - // This gets called when the network containing plugin is destroyed - delete this; -} - -/////////////// - -QuantizePerTokenPluginCreator::QuantizePerTokenPluginCreator() -{ - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - mPluginAttributes.emplace_back(PluginField("type_id", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("quant_mode", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("clamp_enabled", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("sum_per_token", nullptr, PluginFieldType::kINT32)); - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -char const* QuantizePerTokenPluginCreator::getPluginName() const noexcept -{ - return QUANTIZE_PER_TOKEN_PLUGIN_NAME; -} - -char const* QuantizePerTokenPluginCreator::getPluginVersion() const noexcept -{ - return QUANTIZE_PER_TOKEN_PLUGIN_VERSION; -} - -PluginFieldCollection const* QuantizePerTokenPluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -IPluginV2* QuantizePerTokenPluginCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept -{ - PluginFieldParser p{fc->nbFields, fc->fields}; - try - { - auto* obj = new QuantizePerTokenPlugin(static_cast(p.getScalar("type_id").value()), - QuantMode(p.getScalar("quant_mode").value()), - static_cast(p.getScalar("clamp_enabled").value()), - static_cast(p.getScalar("sum_per_token").value())); - - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -IPluginV2* QuantizePerTokenPluginCreator::deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept -{ - // This object will be deleted when the network is destroyed, which will - // call QuantizePerTokenPlugin::destroy() - try - { - auto* obj = new QuantizePerTokenPlugin(serialData, serialLength); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} diff --git a/cpp/tensorrt_llm/plugins/quantizePerTokenPlugin/quantizePerTokenPlugin.h b/cpp/tensorrt_llm/plugins/quantizePerTokenPlugin/quantizePerTokenPlugin.h deleted file mode 100644 index 47b218acfd28..000000000000 --- a/cpp/tensorrt_llm/plugins/quantizePerTokenPlugin/quantizePerTokenPlugin.h +++ /dev/null @@ -1,104 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. - */ -#pragma once - -#include "tensorrt_llm/common/quantization.h" -#include "tensorrt_llm/plugins/common/plugin.h" -#include -#include -#include -#include -#include - -namespace tensorrt_llm::plugins -{ - -class QuantizePerTokenPlugin : public BasePlugin -{ -public: - QuantizePerTokenPlugin(nvinfer1::DataType outputType, tensorrt_llm::common::QuantMode quantMode, - bool clampValEnabled, bool sumPerToken); - - QuantizePerTokenPlugin(void const* data, size_t length); - - ~QuantizePerTokenPlugin() override = default; - - // IPluginV2DynamicExt Methods - nvinfer1::IPluginV2DynamicExt* clone() const noexcept override; - nvinfer1::DimsExprs getOutputDimensions(int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; - bool supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept override; - void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept override; - size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept override; - int enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - - template - void dispatchDataType(void* output, void const* input, void const* clampValPtr, void* scalePtr, void* sumPtr, - int dim0, int dim1, cudaStream_t stream) noexcept; - - // IPluginV2Ext Methods - nvinfer1::DataType getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept override; - - // IPluginV2 Methods - char const* getPluginType() const noexcept override; - char const* getPluginVersion() const noexcept override; - int getNbOutputs() const noexcept override; - int initialize() noexcept override; - void terminate() noexcept override; - size_t getSerializationSize() const noexcept override; - void serialize(void* buffer) const noexcept override; - void destroy() noexcept override; - -private: - std::string const mLayerName; - // The quantized output data type. - nvinfer1::DataType mOutputType; - // The quantization mode. - tensorrt_llm::common::QuantMode mQuantMode; - // Do we clamp the input tensor ? - bool mClampValEnabled; - // Do we output the per-token sum? - bool mSumPerToken; -}; - -class QuantizePerTokenPluginCreator : public BaseCreator -{ -public: - QuantizePerTokenPluginCreator(); - - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; - - nvinfer1::IPluginV2* deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept override; - -private: - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; -}; - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/quantizeTensorPlugin/CMakeLists.txt b/cpp/tensorrt_llm/plugins/quantizeTensorPlugin/CMakeLists.txt deleted file mode 100755 index 86876224fccd..000000000000 --- a/cpp/tensorrt_llm/plugins/quantizeTensorPlugin/CMakeLists.txt +++ /dev/null @@ -1,21 +0,0 @@ -# -# SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. -# -file(GLOB SRCS *.cpp) -set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS}) -set(PLUGIN_SOURCES - ${PLUGIN_SOURCES} - PARENT_SCOPE) diff --git a/cpp/tensorrt_llm/plugins/quantizeTensorPlugin/quantizeTensorPlugin.cpp b/cpp/tensorrt_llm/plugins/quantizeTensorPlugin/quantizeTensorPlugin.cpp deleted file mode 100644 index cacb32b809bf..000000000000 --- a/cpp/tensorrt_llm/plugins/quantizeTensorPlugin/quantizeTensorPlugin.cpp +++ /dev/null @@ -1,250 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. - */ -#include "quantizeTensorPlugin.h" -#include "tensorrt_llm/kernels/quantization.h" - -using namespace nvinfer1; -using namespace tensorrt_llm::kernels; -using tensorrt_llm::plugins::QuantizeTensorPluginCreator; -using tensorrt_llm::plugins::QuantizeTensorPlugin; - -static char const* QUANTIZE_TENSOR_PLUGIN_VERSION{"1"}; -static char const* QUANTIZE_TENSOR_PLUGIN_NAME{"QuantizeTensor"}; -PluginFieldCollection QuantizeTensorPluginCreator::mFC{}; -std::vector QuantizeTensorPluginCreator::mPluginAttributes; - -QuantizeTensorPlugin::QuantizeTensorPlugin() {} - -// Parameterized constructor -QuantizeTensorPlugin::QuantizeTensorPlugin(void const* data, size_t length) -{ - char const *d = reinterpret_cast(data), *a = d; - TLLM_CHECK_WITH_INFO(d == a + length, - "Expected length (%d) != real length (%d). This is often " - "caused by using different TensorRT LLM version to build " - "engine and run engine.", - (int) length, (int) (d - a)); -} - -// IPluginV2DynamicExt Methods -nvinfer1::IPluginV2DynamicExt* QuantizeTensorPlugin::clone() const noexcept -{ - return new QuantizeTensorPlugin(*this); -} - -nvinfer1::DimsExprs QuantizeTensorPlugin::getOutputDimensions( - int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept -{ - try - { - TLLM_CHECK(nbInputs == 2); - TLLM_CHECK(outputIndex < 1); - // Quantized input - return inputs[0]; - } - catch (std::exception const& e) - { - caughtError(e); - } - return DimsExprs{}; -} - -bool QuantizeTensorPlugin::supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept -{ - switch (pos) - { - case 0: - // activation - return (inOut[pos].type == nvinfer1::DataType::kFLOAT || inOut[pos].type == nvinfer1::DataType::kHALF -#ifdef ENABLE_BF16 - || inOut[pos].type == nvinfer1::DataType::kBF16 -#endif - ) - && inOut[pos].format == TensorFormat::kLINEAR; - case 1: - // scales - return inOut[pos].type == nvinfer1::DataType::kFLOAT && inOut[pos].format == TensorFormat::kLINEAR; - case 2: - // quantized activation - return inOut[pos].type == nvinfer1::DataType::kINT8 && inOut[pos].format == TensorFormat::kLINEAR; - default: - // Never should be here - TLLM_CHECK(false); - return false; - } -} - -void QuantizeTensorPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept -{ -} - -size_t QuantizeTensorPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - return 0; -} - -int QuantizeTensorPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) noexcept -{ - // inputs - // activation [M(*), K] - // scale [1, 1] - // outputs - // quant [M(*), K] - - int64_t numElts = 1; - for (int ii = 0; ii < inputDesc[0].dims.nbDims; ++ii) - { - numElts *= inputDesc[0].dims.d[ii]; - } - - if (inputDesc[0].type == DataType::kFLOAT) - { - invokeQuantization(reinterpret_cast(outputs[0]), reinterpret_cast(inputs[0]), - numElts, reinterpret_cast(inputs[1]), stream, mProp.maxGridSize[0]); - } - else if (inputDesc[0].type == DataType::kHALF) - { - invokeQuantization(reinterpret_cast(outputs[0]), reinterpret_cast(inputs[0]), - numElts, reinterpret_cast(inputs[1]), stream, mProp.maxGridSize[0]); - } -#ifdef ENABLE_BF16 - else if (inputDesc[0].type == DataType::kBF16) - { - invokeQuantization<__nv_bfloat16>(reinterpret_cast(outputs[0]), - reinterpret_cast<__nv_bfloat16 const*>(inputs[0]), numElts, reinterpret_cast(inputs[1]), - stream, mProp.maxGridSize[0]); - } -#endif - sync_check_cuda_error(stream); - return 0; -} - -// IPluginV2Ext Methods -nvinfer1::DataType QuantizeTensorPlugin::getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept -{ - TLLM_CHECK(nbInputs == 2); - TLLM_CHECK(index == 0); - return nvinfer1::DataType::kINT8; -} - -// IPluginV2 Methods - -char const* QuantizeTensorPlugin::getPluginType() const noexcept -{ - return QUANTIZE_TENSOR_PLUGIN_NAME; -} - -char const* QuantizeTensorPlugin::getPluginVersion() const noexcept -{ - return QUANTIZE_TENSOR_PLUGIN_VERSION; -} - -int QuantizeTensorPlugin::getNbOutputs() const noexcept -{ - return 1; -} - -int QuantizeTensorPlugin::initialize() noexcept -{ - int deviceId = 0; - tensorrt_llm::common::check_cuda_error(cudaGetDevice(&deviceId)); - tensorrt_llm::common::check_cuda_error(cudaGetDeviceProperties(&mProp, deviceId)); - return 0; -} - -void QuantizeTensorPlugin::terminate() noexcept {} - -size_t QuantizeTensorPlugin::getSerializationSize() const noexcept -{ - return 0; -} - -void QuantizeTensorPlugin::serialize(void* buffer) const noexcept -{ - char *d = static_cast(buffer), *a = d; - TLLM_CHECK(d == a + getSerializationSize()); -} - -void QuantizeTensorPlugin::destroy() noexcept -{ - // This gets called when the network containing plugin is destroyed - delete this; -} - -/////////////// - -QuantizeTensorPluginCreator::QuantizeTensorPluginCreator() -{ - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -char const* QuantizeTensorPluginCreator::getPluginName() const noexcept -{ - return QUANTIZE_TENSOR_PLUGIN_NAME; -} - -char const* QuantizeTensorPluginCreator::getPluginVersion() const noexcept -{ - return QUANTIZE_TENSOR_PLUGIN_VERSION; -} - -PluginFieldCollection const* QuantizeTensorPluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -IPluginV2* QuantizeTensorPluginCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept -{ - try - { - auto* obj = new QuantizeTensorPlugin(); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -IPluginV2* QuantizeTensorPluginCreator::deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept -{ - // This object will be deleted when the network is destroyed, which will - // call QuantizeTensorPlugin::destroy() - try - { - auto* obj = new QuantizeTensorPlugin(serialData, serialLength); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} diff --git a/cpp/tensorrt_llm/plugins/quantizeTensorPlugin/quantizeTensorPlugin.h b/cpp/tensorrt_llm/plugins/quantizeTensorPlugin/quantizeTensorPlugin.h deleted file mode 100644 index 6f1ce864ec35..000000000000 --- a/cpp/tensorrt_llm/plugins/quantizeTensorPlugin/quantizeTensorPlugin.h +++ /dev/null @@ -1,92 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. - */ -#pragma once - -#include "tensorrt_llm/common/quantization.h" -#include "tensorrt_llm/plugins/common/plugin.h" -#include -#include -#include -#include -#include - -namespace tensorrt_llm::plugins -{ - -class QuantizeTensorPlugin : public BasePlugin -{ -public: - QuantizeTensorPlugin(); - - QuantizeTensorPlugin(void const* data, size_t length); - - ~QuantizeTensorPlugin() override = default; - - // IPluginV2DynamicExt Methods - nvinfer1::IPluginV2DynamicExt* clone() const noexcept override; - nvinfer1::DimsExprs getOutputDimensions(int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; - bool supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept override; - void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept override; - size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept override; - int enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - - // IPluginV2Ext Methods - nvinfer1::DataType getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept override; - - // IPluginV2 Methods - char const* getPluginType() const noexcept override; - char const* getPluginVersion() const noexcept override; - int getNbOutputs() const noexcept override; - int initialize() noexcept override; - void terminate() noexcept override; - size_t getSerializationSize() const noexcept override; - void serialize(void* buffer) const noexcept override; - void destroy() noexcept override; - -private: - const std::string mLayerName; - cudaDeviceProp mProp; -}; - -class QuantizeTensorPluginCreator : public BaseCreator -{ -public: - QuantizeTensorPluginCreator(); - - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; - - nvinfer1::IPluginV2* deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept override; - -private: - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; -}; - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/quantizeToFP4Plugin/CMakeLists.txt b/cpp/tensorrt_llm/plugins/quantizeToFP4Plugin/CMakeLists.txt deleted file mode 100644 index 86876224fccd..000000000000 --- a/cpp/tensorrt_llm/plugins/quantizeToFP4Plugin/CMakeLists.txt +++ /dev/null @@ -1,21 +0,0 @@ -# -# SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. -# -file(GLOB SRCS *.cpp) -set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS}) -set(PLUGIN_SOURCES - ${PLUGIN_SOURCES} - PARENT_SCOPE) diff --git a/cpp/tensorrt_llm/plugins/quantizeToFP4Plugin/quantizeToFP4Plugin.cpp b/cpp/tensorrt_llm/plugins/quantizeToFP4Plugin/quantizeToFP4Plugin.cpp deleted file mode 100644 index b5eaffeeda2a..000000000000 --- a/cpp/tensorrt_llm/plugins/quantizeToFP4Plugin/quantizeToFP4Plugin.cpp +++ /dev/null @@ -1,301 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. - */ -#include "quantizeToFP4Plugin.h" -#include "pluginUtils.h" -#include "tensorrt_llm/kernels/quantization.h" -#include - -using namespace nvinfer1; -using namespace tensorrt_llm::kernels; -using namespace tensorrt_llm::common; -using tensorrt_llm::plugins::QuantizeToFP4PluginCreator; -using tensorrt_llm::plugins::QuantizeToFP4Plugin; - -constexpr nvinfer1::DataType FP4_DTYPE = nvinfer1::DataType::kFP4; -constexpr nvinfer1::DataType FP8_DTYPE = nvinfer1::DataType::kFP8; - -static char const* QUANT_FP4_PLUGIN_VERSION{"1"}; -static char const* QUANT_FP4_PLUGIN_NAME{"QuantizeToFP4"}; -PluginFieldCollection QuantizeToFP4PluginCreator::mFC{}; -std::vector QuantizeToFP4PluginCreator::mPluginAttributes; - -QuantizeToFP4Plugin::QuantizeToFP4Plugin(){}; - -// Parameterized constructor -QuantizeToFP4Plugin::QuantizeToFP4Plugin(void const* data, size_t length) -{ - char const *d = reinterpret_cast(data), *a = d; - TLLM_CHECK_WITH_INFO(d == a + length, - "Expected length (%d) != real length (%d). This is often " - "caused by using different TensorRT LLM version to build " - "engine and run engine.", - (int) length, (int) (d - a)); -} - -// IPluginV2DynamicExt Methods -nvinfer1::IPluginV2DynamicExt* QuantizeToFP4Plugin::clone() const noexcept -{ - auto* plugin = new QuantizeToFP4Plugin(); - plugin->setPluginNamespace(mNamespace.c_str()); - return plugin; -} - -nvinfer1::DimsExprs QuantizeToFP4Plugin::getOutputDimensions( - int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept -{ - // Quantized output in FP4 datatype. - if (outputIndex == 0) - { - DimsExprs ret; - ret.nbDims = inputs[0].nbDims; - for (int di = 0; di < ret.nbDims; ++di) - { - ret.d[di] = inputs[0].d[di]; - } - // // Div up by 16 as the storage type has 16 FP4 values per element. - // ret.d[ret.nbDims - 1] - // = exprBuilder.operation(DimensionOperation::kCEIL_DIV, *ret.d[ret.nbDims - 1], - // *exprBuilder.constant(16)); - return ret; - } - // Scaling Factors in FP8. - else if (outputIndex == 1) - { - DimsExprs ret; - ret.nbDims = inputs[0].nbDims; - for (int di = 0; di < ret.nbDims; ++di) - { - ret.d[di] = inputs[0].d[di]; - } - // Sequence dimension or token dimension. - // Pad to multiple of 128. - auto dimM - = exprBuilder.operation(DimensionOperation::kCEIL_DIV, *ret.d[ret.nbDims - 2], *exprBuilder.constant(128)); - ret.d[ret.nbDims - 2] = exprBuilder.operation(DimensionOperation::kPROD, *dimM, *exprBuilder.constant(128)); - // Hidden size dimension. - // Div (rounding up) by 16 since 16 elements share one SF and SF padded to k%4==0. - ret.d[ret.nbDims - 1] - = exprBuilder.operation(DimensionOperation::kCEIL_DIV, *ret.d[ret.nbDims - 1], *exprBuilder.constant(16)); - return ret; - } - return DimsExprs{}; -} - -bool QuantizeToFP4Plugin::supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept -{ - // half input + float global_sf + fp4 output (e2m1) + fp8 SF output. - int const totalPoses = 2 + 2; - TLLM_CHECK(0 <= pos && pos < totalPoses); - TLLM_CHECK(nbInputs == 2); - switch (pos) - { - case 0: - return (inOut[pos].type == nvinfer1::DataType::kHALF || inOut[pos].type == nvinfer1::DataType::kBF16 - || inOut[pos].type == nvinfer1::DataType::kFP8) - && (inOut[pos].format == TensorFormat::kLINEAR); - case 1: return (inOut[pos].type == nvinfer1::DataType::kFLOAT) && (inOut[pos].format == TensorFormat::kLINEAR); - case 2: return (inOut[pos].type == FP4_DTYPE) && (inOut[pos].format == TensorFormat::kLINEAR); - case 3: return (inOut[pos].type == FP8_DTYPE) && (inOut[pos].format == TensorFormat::kLINEAR); - default: break; - } - return false; -} - -void QuantizeToFP4Plugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept -{ -} - -size_t QuantizeToFP4Plugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - return 0; -} - -int QuantizeToFP4Plugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) noexcept -{ - // inputs - // input [M(*), N] half data type - // SF scale [1] float data type - // used to scale SF from input range to fp8 range (448.f / (MaxVal of input / 6.f)) - // outputs - // output [M(*), N] fp4 storage (E2M1) - // SF output [M, N / 16] fp8 storage (UE4M3) - - int64_t m64 = 1; - for (int i = 0; i < inputDesc[0].dims.nbDims - 1; ++i) - { - m64 *= inputDesc[0].dims.d[i]; - } - int const m = TLLM_INT32_CAST(m64); - int const n = TLLM_INT32_CAST(inputDesc[0].dims.d[inputDesc[0].dims.nbDims - 1]); - - TLLM_CHECK_WITH_INFO(n % 16 == 0, "the N dimension must be multiple of 16."); - - float const* SFScale = static_cast(inputs[1]); - int64_t* output = reinterpret_cast(outputs[0]); - int32_t* SFoutput = reinterpret_cast(outputs[1]); - - DataType inputDtype = inputDesc[0].type; - - switch (inputDtype) - { - case DataType::kHALF: - { - auto input = reinterpret_cast(inputs[0]); - invokeFP4Quantization(1, m, n, input, SFScale, output, SFoutput, false, QuantizationSFLayout::SWIZZLED, - mMultiProcessorCount, stream); - break; - } - - case DataType::kBF16: - { - auto input = reinterpret_cast<__nv_bfloat16 const*>(inputs[0]); - invokeFP4Quantization(1, m, n, input, SFScale, output, SFoutput, false, QuantizationSFLayout::SWIZZLED, - mMultiProcessorCount, stream); - break; - } - - case DataType::kFP8: - { - auto input = reinterpret_cast<__nv_fp8_e4m3 const*>(inputs[0]); - invokeFP4Quantization(1, m, n, input, SFScale, output, SFoutput, false, QuantizationSFLayout::SWIZZLED, - mMultiProcessorCount, stream); - break; - } - - default: TLLM_LOG_ERROR("only half, bfloat16 and fp8 data type are supported."); break; - } - - // Use UE4M3 scales by default. - return 0; -} - -// IPluginV2Ext Methods -nvinfer1::DataType QuantizeToFP4Plugin::getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept -{ - if (index == 0) - { - // Output 0 quantized output. - return FP4_DTYPE; - } - // Output 1 SF (scaling factors). - return FP8_DTYPE; -} - -// IPluginV2 Methods - -char const* QuantizeToFP4Plugin::getPluginType() const noexcept -{ - return QUANT_FP4_PLUGIN_NAME; -} - -char const* QuantizeToFP4Plugin::getPluginVersion() const noexcept -{ - return QUANT_FP4_PLUGIN_VERSION; -} - -int QuantizeToFP4Plugin::getNbOutputs() const noexcept -{ - return 2; -} - -int QuantizeToFP4Plugin::initialize() noexcept -{ - return 0; -} - -void QuantizeToFP4Plugin::terminate() noexcept {} - -size_t QuantizeToFP4Plugin::getSerializationSize() const noexcept -{ - return 0; -} - -void QuantizeToFP4Plugin::serialize(void* buffer) const noexcept -{ - char *d = static_cast(buffer), *a = d; - TLLM_CHECK(d == a + getSerializationSize()); -} - -void QuantizeToFP4Plugin::destroy() noexcept -{ - // This gets called when the network containing plugin is destroyed - delete this; -} - -/////////////// - -QuantizeToFP4PluginCreator::QuantizeToFP4PluginCreator() -{ - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -char const* QuantizeToFP4PluginCreator::getPluginName() const noexcept -{ - return QUANT_FP4_PLUGIN_NAME; -} - -char const* QuantizeToFP4PluginCreator::getPluginVersion() const noexcept -{ - return QUANT_FP4_PLUGIN_VERSION; -} - -PluginFieldCollection const* QuantizeToFP4PluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -IPluginV2* QuantizeToFP4PluginCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept -{ - try - { - auto* obj = new QuantizeToFP4Plugin(); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -IPluginV2* QuantizeToFP4PluginCreator::deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept -{ - // This object will be deleted when the network is destroyed, which will - // call QuantizeToFP4Plugin::destroy() - try - { - auto* obj = new QuantizeToFP4Plugin(serialData, serialLength); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} diff --git a/cpp/tensorrt_llm/plugins/quantizeToFP4Plugin/quantizeToFP4Plugin.h b/cpp/tensorrt_llm/plugins/quantizeToFP4Plugin/quantizeToFP4Plugin.h deleted file mode 100644 index b584837a447a..000000000000 --- a/cpp/tensorrt_llm/plugins/quantizeToFP4Plugin/quantizeToFP4Plugin.h +++ /dev/null @@ -1,90 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. - */ -#pragma once - -#include "tensorrt_llm/plugins/common/plugin.h" -#include -#include -#include -#include - -namespace tensorrt_llm::plugins -{ - -class QuantizeToFP4Plugin : public BasePlugin -{ -public: - QuantizeToFP4Plugin(); - - QuantizeToFP4Plugin(void const* data, size_t length); - - ~QuantizeToFP4Plugin() override = default; - - // IPluginV2DynamicExt Methods - nvinfer1::IPluginV2DynamicExt* clone() const noexcept override; - nvinfer1::DimsExprs getOutputDimensions(int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; - bool supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept override; - void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept override; - size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept override; - int enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - - // IPluginV2Ext Methods - nvinfer1::DataType getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept override; - - // IPluginV2 Methods - char const* getPluginType() const noexcept override; - char const* getPluginVersion() const noexcept override; - int getNbOutputs() const noexcept override; - int initialize() noexcept override; - void terminate() noexcept override; - size_t getSerializationSize() const noexcept override; - void serialize(void* buffer) const noexcept override; - void destroy() noexcept override; - -private: - const std::string mLayerName; - int const mMultiProcessorCount = tensorrt_llm::common::getMultiProcessorCount(); -}; - -class QuantizeToFP4PluginCreator : public BaseCreator -{ -public: - QuantizeToFP4PluginCreator(); - - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; - - nvinfer1::IPluginV2* deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept override; - -private: - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; -}; - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/rmsnormQuantizationPlugin/CMakeLists.txt b/cpp/tensorrt_llm/plugins/rmsnormQuantizationPlugin/CMakeLists.txt deleted file mode 100755 index 86876224fccd..000000000000 --- a/cpp/tensorrt_llm/plugins/rmsnormQuantizationPlugin/CMakeLists.txt +++ /dev/null @@ -1,21 +0,0 @@ -# -# SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. -# -file(GLOB SRCS *.cpp) -set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS}) -set(PLUGIN_SOURCES - ${PLUGIN_SOURCES} - PARENT_SCOPE) diff --git a/cpp/tensorrt_llm/plugins/rmsnormQuantizationPlugin/rmsnormQuantizationPlugin.cpp b/cpp/tensorrt_llm/plugins/rmsnormQuantizationPlugin/rmsnormQuantizationPlugin.cpp deleted file mode 100644 index 16d0bf2dc356..000000000000 --- a/cpp/tensorrt_llm/plugins/rmsnormQuantizationPlugin/rmsnormQuantizationPlugin.cpp +++ /dev/null @@ -1,452 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. - */ -#include "rmsnormQuantizationPlugin.h" -#include "pluginUtils.h" -#include "tensorrt_llm/kernels/rmsnormKernels.h" - -using namespace nvinfer1; -using namespace tensorrt_llm::kernels; -using namespace tensorrt_llm::common; -using tensorrt_llm::plugins::RmsnormQuantizationPluginCreator; -using tensorrt_llm::plugins::RmsnormQuantizationPlugin; - -static char const* RMSNORM_QUANTIZATION_PLUGIN_VERSION{"1"}; -static char const* RMSNORM_QUANTIZATION_PLUGIN_NAME{"RmsnormQuantization"}; -PluginFieldCollection RmsnormQuantizationPluginCreator::mFC{}; -std::vector RmsnormQuantizationPluginCreator::mPluginAttributes; - -RmsnormQuantizationPlugin::RmsnormQuantizationPlugin(float eps, bool dynamicActivationScaling, bool sumPerToken, - bool clampValEnabled, QuantMode quantMode, nvinfer1::DataType type, nvinfer1::DataType outputType) - : mEps(eps) - , mDynActScaling(dynamicActivationScaling) - , mType(type) - , mOutputType{outputType} - , mClampValEnabled{clampValEnabled} - , mQuantMode{quantMode} - , mSumPerToken(sumPerToken) -{ - TLLM_CHECK_WITH_INFO(mOutputType == nvinfer1::DataType::kINT8 || mOutputType == nvinfer1::DataType::kFP8, - "Only int8 or fp8 output type is allowed."); - // Check if the quant mode is valid. - TLLM_CHECK_WITH_INFO(mQuantMode.hasPerTokenScaling(), "The quant mode is not valid."); -} - -// Parameterized constructor -RmsnormQuantizationPlugin::RmsnormQuantizationPlugin(void const* data, size_t length) -{ - char const *d = reinterpret_cast(data), *a = d; - read(d, mEps); - read(d, mDynActScaling); - read(d, mSumPerToken); - read(d, mClampValEnabled); - read(d, mQuantMode); - read(d, mType); - read(d, mOutputType); - TLLM_CHECK_WITH_INFO(d == a + length, - "Expected length (%d) != real length (%d). This is often " - "caused by using different TensorRT LLM version to build " - "engine and run engine.", - (int) length, (int) (d - a)); -} - -// IPluginV2DynamicExt Methods -nvinfer1::IPluginV2DynamicExt* RmsnormQuantizationPlugin::clone() const noexcept -{ - auto* plugin = new RmsnormQuantizationPlugin( - mEps, mDynActScaling, mSumPerToken, mClampValEnabled, mQuantMode, mType, mOutputType); - plugin->setPluginNamespace(mNamespace.c_str()); - return plugin; -} - -nvinfer1::DimsExprs RmsnormQuantizationPlugin::getOutputDimensions( - int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept -{ - if (outputIndex == 0) - { - // Quantized output - return inputs[outputIndex]; - } - - // Dynamic scaling or per-token sum if enabled. - try - { - if (outputIndex == 1) - { - TLLM_CHECK(mDynActScaling); - } - else if (outputIndex == 2) - { - TLLM_CHECK(mSumPerToken); - } - else - { - TLLM_CHECK(false); - } - - DimsExprs ret; - ret.nbDims = inputs[0].nbDims; - for (int di = 0; di < ret.nbDims - 1; ++di) - { - ret.d[di] = inputs[0].d[di]; - } - ret.d[ret.nbDims - 1] = exprBuilder.constant(1); - return ret; - } - catch (std::exception const& e) - { - caughtError(e); - } - return DimsExprs{}; -} - -bool RmsnormQuantizationPlugin::supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept -{ - int const totalPoses - = 6 + static_cast(mClampValEnabled) + static_cast(mDynActScaling) + static_cast(mSumPerToken); - TLLM_CHECK(0 <= pos && pos < totalPoses); - TLLM_CHECK(nbInputs == 4 + static_cast(mClampValEnabled)); - if (pos < nbInputs) - { - if (pos < 3) - { - // activation, weight, bias - return (inOut[pos].type == mType) && (inOut[pos].format == TensorFormat::kLINEAR); - } - else if (pos == 3) - { - // scale - return (inOut[pos].type == nvinfer1::DataType::kFLOAT) && (inOut[pos].format == TensorFormat::kLINEAR); - } - else if (pos == 4 && mClampValEnabled) - { - // clamp_max_v - return inOut[pos].type == nvinfer1::DataType::kFLOAT && inOut[pos].format == TensorFormat::kLINEAR; - } - } - else if (pos == 4 + int(mClampValEnabled)) - { - // Quantized output - return (inOut[pos].type == mOutputType) && (inOut[pos].format == TensorFormat::kLINEAR); - } - else if (pos == 5 + int(mClampValEnabled)) - { - // Dynamic scaling if enabled - return (inOut[pos].type == nvinfer1::DataType::kFLOAT) && (inOut[pos].format == TensorFormat::kLINEAR); - } - else if (pos == 6 + int(mClampValEnabled)) - { - // Per-token activation sum if enabled - return (inOut[pos].type == nvinfer1::DataType::kFLOAT) && (inOut[pos].format == TensorFormat::kLINEAR); - } - - // Never should be here - TLLM_CHECK_WITH_INFO(false, "The input/output is not supported."); - return false; -} - -void RmsnormQuantizationPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept -{ -} - -size_t RmsnormQuantizationPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - return 0; -} - -template -void RmsnormQuantizationPlugin::dispatchDataType(void* out, void const* input, void const* gamma, void const* beta, - float const eps, int const tokens, int const hidden_dim, cudaStream_t stream, void const* clampValPtr, - void const* scale, void* dynamic_scale, void* sum_per_token, void* normed_output_quant) noexcept -{ - // inputs - // activation [dim0(*), dim1] - // clamp_value [2], contains min val, and max val (optional) - // outputs - // quant [dim0(*), dim1] - // scale_tokens [dim0(*), 1] - - invokeGeneralRmsNorm(reinterpret_cast(out), reinterpret_cast(input), - reinterpret_cast(gamma), reinterpret_cast(beta), eps, tokens, hidden_dim, mQuantMode, - stream, reinterpret_cast(clampValPtr), reinterpret_cast(scale), - reinterpret_cast(dynamic_scale), reinterpret_cast(sum_per_token), - reinterpret_cast(normed_output_quant)); -} - -int RmsnormQuantizationPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) noexcept -{ - // inputs - // input [M(*), N] - // weight [N, ] - // bias [N, ] - // scale_to_int [1] - // clamp_value [2], contains min val, and max val (optional) - // outputs - // output [M(*), N] Normalized activations, potentially with quantization applied. - // dynamic_scaling [M(*), 1] (Optional) Per-token scales if quantization is enabled. - // token_sums [M(*), 1] (Optional) Per-token sums of all the channels (before quantization). - - int64_t m64 = 1; - for (int i = 0; i < inputDesc[0].dims.nbDims - 1; ++i) - { - m64 *= inputDesc[0].dims.d[i]; - } - int const m = TLLM_INT32_CAST(m64); - int const n = TLLM_INT32_CAST(inputDesc[1].dims.d[0]); - - void const* input = inputs[0]; - void const* weight = inputs[1]; - void const* bias = inputs[2]; - void const* scale = inputs[3]; - void const* clampValPtr = mClampValEnabled ? inputs[4] : nullptr; - void* output = outputs[0]; - void* dynamic_scale = mDynActScaling ? outputs[1] : nullptr; - void* sum_per_token = mSumPerToken ? outputs[2] : nullptr; - - if (inputDesc[0].type == DataType::kFLOAT && mOutputType == DataType::kINT8) - { - dispatchDataType( - nullptr, input, weight, bias, mEps, m, n, stream, clampValPtr, scale, dynamic_scale, sum_per_token, output); - } -#ifdef ENABLE_FP8 - else if (inputDesc[0].type == DataType::kFLOAT && mOutputType == DataType::kFP8) - { - dispatchDataType( - nullptr, input, weight, bias, mEps, m, n, stream, clampValPtr, scale, dynamic_scale, sum_per_token, output); - } -#endif // ENABLE_FP8 - else if (inputDesc[0].type == DataType::kHALF && mOutputType == DataType::kINT8) - { - dispatchDataType( - nullptr, input, weight, bias, mEps, m, n, stream, clampValPtr, scale, dynamic_scale, sum_per_token, output); - } -#ifdef ENABLE_FP8 - else if (inputDesc[0].type == DataType::kHALF && mOutputType == DataType::kFP8) - { - dispatchDataType( - nullptr, input, weight, bias, mEps, m, n, stream, clampValPtr, scale, dynamic_scale, sum_per_token, output); - } -#endif // ENABLE_FP8 -#ifdef ENABLE_BF16 - else if (inputDesc[0].type == DataType::kBF16 && mOutputType == DataType::kINT8) - { - dispatchDataType<__nv_bfloat16, int8_t>( - nullptr, input, weight, bias, mEps, m, n, stream, clampValPtr, scale, dynamic_scale, sum_per_token, output); - } -#ifdef ENABLE_FP8 - else if (inputDesc[0].type == DataType::kBF16 && mOutputType == DataType::kFP8) - { - dispatchDataType<__nv_bfloat16, __nv_fp8_e4m3>( - nullptr, input, weight, bias, mEps, m, n, stream, clampValPtr, scale, dynamic_scale, sum_per_token, output); - } -#endif // ENABLE_FP8 -#endif // ENABLE_BF16 - sync_check_cuda_error(stream); - return 0; -} - -// IPluginV2Ext Methods -nvinfer1::DataType RmsnormQuantizationPlugin::getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept -{ - assert(index <= 2); - - if (index == 0) - { - // Output 0 quantized output of layer norm - return mOutputType; - } - if (index == 1) - { - assert(mDynActScaling); - // Output 1 dynamic act scaling - return nvinfer1::DataType::kFLOAT; - } - // index == 2 - { - assert(mDynActScaling && mSumPerToken); - // Output 2 per token sum - return nvinfer1::DataType::kFLOAT; - } -} - -// IPluginV2 Methods - -char const* RmsnormQuantizationPlugin::getPluginType() const noexcept -{ - return RMSNORM_QUANTIZATION_PLUGIN_NAME; -} - -char const* RmsnormQuantizationPlugin::getPluginVersion() const noexcept -{ - return RMSNORM_QUANTIZATION_PLUGIN_VERSION; -} - -int RmsnormQuantizationPlugin::getNbOutputs() const noexcept -{ - return 1 + static_cast(mDynActScaling) + static_cast(mSumPerToken); -} - -int RmsnormQuantizationPlugin::initialize() noexcept -{ - return 0; -} - -void RmsnormQuantizationPlugin::terminate() noexcept {} - -size_t RmsnormQuantizationPlugin::getSerializationSize() const noexcept -{ - return sizeof(mOutputType) + sizeof(mClampValEnabled) + sizeof(mEps) + sizeof(mDynActScaling) + sizeof(mSumPerToken) - + sizeof(mType) + sizeof(mQuantMode); -} - -void RmsnormQuantizationPlugin::serialize(void* buffer) const noexcept -{ - char *d = static_cast(buffer), *a = d; - write(d, mEps); - write(d, mDynActScaling); - write(d, mSumPerToken); - write(d, mClampValEnabled); - write(d, mQuantMode); - write(d, mType); - write(d, mOutputType); - TLLM_CHECK(d == a + getSerializationSize()); -} - -void RmsnormQuantizationPlugin::destroy() noexcept -{ - // This gets called when the network containing plugin is destroyed - delete this; -} - -/////////////// - -RmsnormQuantizationPluginCreator::RmsnormQuantizationPluginCreator() -{ - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - mPluginAttributes.emplace_back(PluginField("eps", nullptr, PluginFieldType::kFLOAT32)); - mPluginAttributes.emplace_back(PluginField("dyn_act_scaling", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("sum_per_token", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("clamp_enabled", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("quant_mode", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("type_id", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("out_type_id", nullptr, PluginFieldType::kINT32)); - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -char const* RmsnormQuantizationPluginCreator::getPluginName() const noexcept -{ - return RMSNORM_QUANTIZATION_PLUGIN_NAME; -} - -char const* RmsnormQuantizationPluginCreator::getPluginVersion() const noexcept -{ - return RMSNORM_QUANTIZATION_PLUGIN_VERSION; -} - -PluginFieldCollection const* RmsnormQuantizationPluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -IPluginV2* RmsnormQuantizationPluginCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept -{ - PluginField const* fields = fc->fields; - nvinfer1::DataType outputType{}; - QuantMode quantMode; - bool clampValEnabled = false; - float eps{}; - nvinfer1::DataType type{}; - bool dynamicActivationScaling{}; - bool sumPerToken{}; - // Read configurations from each fields - for (int i = 0; i < fc->nbFields; ++i) - { - char const* attrName = fields[i].name; - if (!strcmp(attrName, "quant_mode")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - quantMode = QuantMode(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "out_type_id")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - outputType = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "clamp_enabled")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - clampValEnabled = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "eps")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kFLOAT32); - eps = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "type_id")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - type = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "dyn_act_scaling")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - dynamicActivationScaling = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "sum_per_token")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - sumPerToken = static_cast(*(static_cast(fields[i].data))); - } - } - try - { - auto* obj = new RmsnormQuantizationPlugin( - eps, dynamicActivationScaling, sumPerToken, clampValEnabled, quantMode, type, outputType); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -IPluginV2* RmsnormQuantizationPluginCreator::deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept -{ - // This object will be deleted when the network is destroyed, which will - // call RmsnormQuantizationPlugin::destroy() - try - { - auto* obj = new RmsnormQuantizationPlugin(serialData, serialLength); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} diff --git a/cpp/tensorrt_llm/plugins/rmsnormQuantizationPlugin/rmsnormQuantizationPlugin.h b/cpp/tensorrt_llm/plugins/rmsnormQuantizationPlugin/rmsnormQuantizationPlugin.h deleted file mode 100644 index 762a9bb8de1b..000000000000 --- a/cpp/tensorrt_llm/plugins/rmsnormQuantizationPlugin/rmsnormQuantizationPlugin.h +++ /dev/null @@ -1,108 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. - */ -#pragma once - -#include "tensorrt_llm/common/quantization.h" -#include "tensorrt_llm/plugins/common/plugin.h" -#include -#include -#include -#include - -namespace tensorrt_llm::plugins -{ - -class RmsnormQuantizationPlugin : public BasePlugin -{ -public: - RmsnormQuantizationPlugin(float eps, bool dynamicActivationScaling, bool sumPerToken, bool clampValEnabled, - tensorrt_llm::common::QuantMode quantMode, nvinfer1::DataType type, nvinfer1::DataType outputType); - - RmsnormQuantizationPlugin(void const* data, size_t length); - - ~RmsnormQuantizationPlugin() override = default; - - // IPluginV2DynamicExt Methods - nvinfer1::IPluginV2DynamicExt* clone() const noexcept override; - nvinfer1::DimsExprs getOutputDimensions(int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; - bool supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept override; - void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept override; - size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept override; - int enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - - template - void dispatchDataType(void* out, void const* input, void const* gamma, void const* beta, float const eps, - int const tokens, int const hidden_dim, cudaStream_t stream, void const* clampValPtr, void const* scale, - void* dynamic_scale, void* normed_output_quant, void* act_sum) noexcept; - - // IPluginV2Ext Methods - nvinfer1::DataType getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept override; - - // IPluginV2 Methods - char const* getPluginType() const noexcept override; - char const* getPluginVersion() const noexcept override; - int getNbOutputs() const noexcept override; - int initialize() noexcept override; - void terminate() noexcept override; - size_t getSerializationSize() const noexcept override; - void serialize(void* buffer) const noexcept override; - void destroy() noexcept override; - -private: - float mEps; - bool mDynActScaling; - nvinfer1::DataType mType; - - std::string const mLayerName; - // The quantized output data type. - nvinfer1::DataType mOutputType; - // Do we clamp the input tensor ? - bool mClampValEnabled; - // The quantization mode. - tensorrt_llm::common::QuantMode mQuantMode; - // Should we output the sum of channels per-token? (Used by QServe GEMM) - bool mSumPerToken; -}; - -class RmsnormQuantizationPluginCreator : public BaseCreator -{ -public: - RmsnormQuantizationPluginCreator(); - - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; - - nvinfer1::IPluginV2* deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept override; - -private: - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; -}; - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/selectiveScanPlugin/CMakeLists.txt b/cpp/tensorrt_llm/plugins/selectiveScanPlugin/CMakeLists.txt deleted file mode 100644 index 86876224fccd..000000000000 --- a/cpp/tensorrt_llm/plugins/selectiveScanPlugin/CMakeLists.txt +++ /dev/null @@ -1,21 +0,0 @@ -# -# SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. -# -file(GLOB SRCS *.cpp) -set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS}) -set(PLUGIN_SOURCES - ${PLUGIN_SOURCES} - PARENT_SCOPE) diff --git a/cpp/tensorrt_llm/plugins/selectiveScanPlugin/selectiveScanPlugin.cpp b/cpp/tensorrt_llm/plugins/selectiveScanPlugin/selectiveScanPlugin.cpp deleted file mode 100644 index 3e60182f28c2..000000000000 --- a/cpp/tensorrt_llm/plugins/selectiveScanPlugin/selectiveScanPlugin.cpp +++ /dev/null @@ -1,594 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-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. - */ - -#include "selectiveScanPlugin.h" -#include "tensorrt_llm/common/assert.h" - -using namespace nvinfer1; -using namespace tensorrt_llm::kernels; -using namespace tensorrt_llm::common; -using tensorrt_llm::plugins::SelectiveScanPluginCreator; -using tensorrt_llm::plugins::SelectiveScanPlugin; - -static char const* SELECTIVE_SCAN_PLUGIN_VERSION{"1"}; -static char const* SELECTIVE_SCAN_PLUGIN_NAME{"SelectiveScan"}; -PluginFieldCollection SelectiveScanPluginCreator::mFC{}; -std::vector SelectiveScanPluginCreator::mPluginAttributes; - -SelectiveScanPlugin::SelectiveScanPlugin(int dim, int dstate, int dtRank, int nHeads, int nGroups, int chunkSize, - bool deltaSoftplus, nvinfer1::DataType type, bool removePadding, bool pagedState, bool zEnabled, bool isMamba2) - : mDim(dim) - , mDState(dstate) - , mDtRank(dtRank) - , mNHeads(nHeads) - , mNGroups(nGroups) - , mChunkSize(chunkSize) - , mDeltaSoftplus(deltaSoftplus) - , mType(type) - , mRemovePadding(removePadding) - , mPagedState(pagedState) - , mZEnabled(zEnabled) - , mIsMamba2(isMamba2) - , mDriver(tensorrt_llm::common::CUDADriverWrapper::getInstance()) -{ - TLLM_CHECK_WITH_INFO( - (mChunkSize == 256 || mChunkSize == 128) || (!mIsMamba2), "Only support CHUNK_SIZE 256 or 128"); - TLLM_CHECK_WITH_INFO((mType == DataType::kBF16) || (mType == DataType::kFLOAT) || (mType == DataType::kHALF), - "Only support float, half, and bfloat16."); -} - -// Parameterized constructor -SelectiveScanPlugin::SelectiveScanPlugin(void const* data, size_t length) - : mDriver(tensorrt_llm::common::CUDADriverWrapper::getInstance()) -{ - char const *d = reinterpret_cast(data), *a = d; - read(d, mDim); - read(d, mDState); - read(d, mDtRank); - read(d, mNHeads); - read(d, mNGroups); - read(d, mChunkSize); - read(d, mDeltaSoftplus); - read(d, mType); - read(d, mRemovePadding); - read(d, mPagedState); - read(d, mZEnabled); - read(d, mIsMamba2); - TLLM_CHECK(d == a + length); - TLLM_CHECK_WITH_INFO( - (mChunkSize == 256 || mChunkSize == 128) || (!mIsMamba2), "Only support CHUNK_SIZE 256 or 128"); - TLLM_CHECK_WITH_INFO((mType == DataType::kBF16) || (mType == DataType::kFLOAT) || (mType == DataType::kHALF), - "Only support float, half, and bfloat16."); -} - -// IPluginV2DynamicExt Methods -nvinfer1::IPluginV2DynamicExt* SelectiveScanPlugin::clone() const noexcept -{ - auto* plugin = new SelectiveScanPlugin(mDim, mDState, mDtRank, mNHeads, mNGroups, mChunkSize, mDeltaSoftplus, mType, - mRemovePadding, mPagedState, mZEnabled, mIsMamba2); - plugin->setPluginNamespace(mNamespace.c_str()); - return plugin; -} - -// Outputs -// output_tensor: [batch_size, seq_len, dim] or [num_tokens, dim] for remove_input_padding -// state: [batch_size, dstate, dim] -nvinfer1::DimsExprs SelectiveScanPlugin::getOutputDimensions( - int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept -{ - if (outputIndex == 0) - { - if (mIsMamba2) - { - auto ret = inputs[getInputTensorIdx()]; - ret.d[mRemovePadding ? 1 : 2] = exprBuilder.constant(mDim); - return ret; - } - else - { - return inputs[getInputTensorIdx()]; - } - } - return inputs[getStateIdx()]; -} - -bool SelectiveScanPlugin::supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept -{ - if (pos == getHostRequestTypesIdx() || pos == getLastTokenIdsIdx() - || (mRemovePadding && pos == getHostContextLengthIdx()) || (mPagedState && pos == getSlotMappingIdx())) - { - return inOut[pos].type == nvinfer1::DataType::kINT32; - } - else if (pos == getAIdx() || pos == getDeltaBiasIdx() || pos == getDIdx()) - { - return (inOut[pos].type == nvinfer1::DataType::kFLOAT) && (inOut[pos].format == TensorFormat::kLINEAR); - } - else if (mPagedState && pos == getStateIdx()) - { - return inOut[pos].type == nvinfer1::DataType::kINT64; - } - else - { - return (inOut[pos].type == mType) && (inOut[pos].format == TensorFormat::kLINEAR); - } -} - -void SelectiveScanPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept -{ -} - -size_t SelectiveScanPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - if (!mIsMamba2) - return 0; - - int const NUM_BUFFERS = 6; - size_t workspaces[NUM_BUFFERS]; - - if (mRemovePadding) - { - int B = inputs[getLastTokenIdsIdx()].dims.d[0]; - int BxL = inputs[getInputTensorIdx()].dims.d[0]; // num_tokens - int H = mNHeads; - int P = mDim / H; - int G = mNGroups; - int N = mDState; - int Q = mChunkSize; - int BxC = (BxL + Q - 1) / Q + B; - - workspaces[0] = long(BxC) * H * N * P * 2; // g_mxOs_ - workspaces[1] = long(BxC) * H * N * P * 4; // g_mxSt_ in float - workspaces[2] = long(BxC) * H * Q * 4; // g_mxdc_ in float - workspaces[3] = long(BxC) * H * Q * 4; // g_mxdA_ in float - workspaces[4] = long(BxC) * G * Q * Q * 2; // g_mxCB_ - workspaces[5] = 1024; // TMA descs - } - else - { - int B = inputs[getInputTensorIdx()].dims.d[0]; - int L = inputs[getInputTensorIdx()].dims.d[1]; - int H = mNHeads; - int P = mDim / H; - int G = mNGroups; - int N = mDState; - int Q = mChunkSize; - int C = (L + Q - 1) / Q; - - workspaces[0] = long(B * C) * H * N * P * 2; // g_mxOs_ - workspaces[1] = long(B * C) * H * N * P * 4; // g_mxSt_ in float - workspaces[2] = long(B * C) * H * Q * 4; // g_mxdc_ in float - workspaces[3] = long(B * C) * H * Q * 4; // g_mxdA_ in float - workspaces[4] = long(B * C) * G * Q * Q * 2; // g_mxCB_ - workspaces[5] = 1024; // TMA descs - } - - return calculateTotalWorkspaceSize(workspaces, NUM_BUFFERS); -} - -void SelectiveScanPlugin::setSSMParams(SSMParamsBase& params, const size_t batch, const size_t dim, - const size_t maxSeqLen, const size_t numTokens, const size_t dstate, const size_t dtRank, const size_t nHeads, - const size_t nGroups, const size_t chunkSize, void* statePtr, void const* x, void const* delta, - void const* deltaBias, void const* A, void const* BC, void const* D, void const* z, void* osPtr, void* stPtr, - void* dcPtr, void* dAPtr, void* cbPtr, void* descPtr, int const* lastTokenIds, int const* slotMapping, void* out, - bool deltaSoftplus, bool removePadding) -{ - // Reset the parameters - memset(¶ms, 0, sizeof(params)); - - params.batch = batch; - params.dim = dim; - params.max_seqlen = maxSeqLen; - params.num_tokens = numTokens; - params.dstate = dstate; - params.dt_rank = dtRank; - params.nheads = nHeads; - params.ngroups = nGroups; - params.chunk_size = chunkSize; - - params.delta_softplus = deltaSoftplus; - params.remove_padding = removePadding; - params.is_mamba2 = mIsMamba2; - - // Set the pointers and strides. - params.u_ptr = const_cast(x); - params.delta_ptr = const_cast(delta); - params.A_ptr = const_cast(A); - params.BC_ptr = const_cast(BC); - params.D_ptr = const_cast(D); - params.delta_bias_ptr = const_cast(deltaBias); - params.out_ptr = out; - params.x_ptr = statePtr; - params.z_ptr = const_cast(z); - params.Os_ptr = osPtr; - params.St_ptr = stPtr; - params.dc_ptr = dcPtr; - params.dA_ptr = dAPtr; - params.CB_ptr = cbPtr; - params.desc_ptr = descPtr; - params.last_token_ids_ptr = lastTokenIds; - params.slot_mapping_ptr = slotMapping; -} - -template -int SelectiveScanPlugin::enqueueImpl(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) -{ - // inputs - // 0. input_tensor [batch_size, max_seq_len, dim] or [num_tokens, dim] - // 1. state mamba: [batch_size, dstate, dim] or host [1] containing only pointer for paged_state - // mamba2: [batch_size, nheads, dstate, dim] or host [1] containing only pointer for paged_state - // 2. delta, mamba: [batch_size, seq_len, dim] or [num_tokens, dim] for remove_input_padding - // mamba2: [batch_size, seq_len, nheads] or [num_tokens, nheads] for remove_input_padding - // 3. delta_bias, [dim] for mamba, [nheads] for mamba2 - // 4. A, [dstate, dim] for mamba, [nheads] for mamba2 - // 5. BC, mamba: [batch_size, seq_len, dstate * 2] or [num_tokens, dstate * 2] for remove_input_padding - // mamba2: [batch_size, seq_len, ngroups * dstate * 2] or [num_tokens, ngroups * dstate * 2] for - // remove_input_padding - // 6. D, [dim] for mamba, [nheads] for mamba2 - // 7. host_request_types [batch_size] int32. 0: context; 1: generation. - // 8. last_token_ids [batch_size] int32 - // 9. host_context_lengths [batch_size] int32, optional for remove_input_padding - // 10. state_slot_mapping [batch_size] int32, optional for paged state - // 11. z [batch_size, max_seq_len, dim] or [num_tokens, dim] - // outputs - // 0. output_tensor [batch_size, max_seq_len, dim] or [num_tokens, dim] - // 1. state, [batch_size, dstate, dim] for mamba, [batch_size, nheads, dstate, dim] for mamba2 - auto const batch_size = inputDesc[getHostRequestTypesIdx()].dims.d[0]; - int max_seq_len; - if (mRemovePadding) - { - int const* host_context_length = static_cast(inputs[getHostContextLengthIdx()]); - max_seq_len = *std::max_element(host_context_length, host_context_length + batch_size); - } - else - { - max_seq_len = inputDesc[getInputTensorIdx()].dims.d[1]; - } - - // only support context or generation, not for both of them - RequestType const* reqTypes = static_cast(inputs[getHostRequestTypesIdx()]); - - SSMParamsBase ssm_params; - - int const* slotMapping = mPagedState ? static_cast(inputs[getSlotMappingIdx()]) : nullptr; - void const* z = mZEnabled ? inputs[getZIdx()] : nullptr; - - void* statePtr = mPagedState ? *reinterpret_cast(const_cast(inputs[getStateIdx()])) : outputs[1]; - - // Workspace pointer shift - int8_t* workspace_byte_ptr = reinterpret_cast(workspace); - size_t offset = 0; - - T* mxOs = nullptr; - float* mxSt = nullptr; - float* mxdc = nullptr; - float* mxdA = nullptr; - T* mxCB = nullptr; - void* descs = nullptr; - - if (!mIsMamba2 || reqTypes[0] == RequestType::kGENERATION) /* no workspace needed */ - ; - else if (mRemovePadding) - { - int B = inputDesc[getLastTokenIdsIdx()].dims.d[0]; - int BxL = inputDesc[getInputTensorIdx()].dims.d[0]; // num_tokens - int H = mNHeads; - int P = mDim / H; - int G = mNGroups; - int N = mDState; - int Q = mChunkSize; - int BxC = (BxL + Q - 1) / Q + B; - - mxOs = reinterpret_cast(nextWorkspacePtr(workspace_byte_ptr, offset, long(BxC) * H * N * P * 2)); - mxSt = reinterpret_cast(nextWorkspacePtr(workspace_byte_ptr, offset, long(BxC) * H * N * P * 4)); - mxdc = reinterpret_cast(nextWorkspacePtr(workspace_byte_ptr, offset, long(BxC) * H * Q * 4)); - mxdA = reinterpret_cast(nextWorkspacePtr(workspace_byte_ptr, offset, long(BxC) * H * Q * 4)); - mxCB = reinterpret_cast(nextWorkspacePtr(workspace_byte_ptr, offset, long(BxC) * G * Q * Q * 2)); - descs = nextWorkspacePtr(workspace_byte_ptr, offset, 1024); - } - else - { - int B = inputDesc[getInputTensorIdx()].dims.d[0]; - int L = inputDesc[getInputTensorIdx()].dims.d[1]; - int H = mNHeads; - int P = mDim / H; - int G = mNGroups; - int N = mDState; - int Q = mChunkSize; - int C = (L + Q - 1) / Q; - - mxOs = reinterpret_cast(nextWorkspacePtr(workspace_byte_ptr, offset, long(B * C) * H * N * P * 2)); - mxSt = reinterpret_cast(nextWorkspacePtr(workspace_byte_ptr, offset, long(B * C) * H * N * P * 4)); - mxdc = reinterpret_cast(nextWorkspacePtr(workspace_byte_ptr, offset, long(B * C) * H * Q * 4)); - mxdA = reinterpret_cast(nextWorkspacePtr(workspace_byte_ptr, offset, long(B * C) * H * Q * 4)); - mxCB = reinterpret_cast(nextWorkspacePtr(workspace_byte_ptr, offset, long(B * C) * G * Q * Q * 2)); - descs = nextWorkspacePtr(workspace_byte_ptr, offset, 1024); - } - - int numTokens = inputDesc[getInputTensorIdx()].dims.d[0]; - if (!mRemovePadding) - numTokens *= inputDesc[getInputTensorIdx()].dims.d[1]; - - setSSMParams(ssm_params, batch_size, mDim, max_seq_len, numTokens, mDState, mDtRank, mNHeads, mNGroups, mChunkSize, - statePtr, inputs[getInputTensorIdx()], inputs[getDeltaIdx()], inputs[getDeltaBiasIdx()], inputs[getAIdx()], - inputs[getBCIdx()], inputs[getDIdx()], z, mxOs, mxSt, mxdc, mxdA, mxCB, descs, - static_cast(inputs[getLastTokenIdsIdx()]), slotMapping, outputs[0], mDeltaSoftplus, mRemovePadding); - - if (reqTypes[0] == RequestType::kCONTEXT) - { - if (mIsMamba2) - { - invokeChunkScan(ssm_params, stream, mDriver.get()); - } - else - { - invokeSelectiveScan(ssm_params, stream); - } - } - else if (reqTypes[0] == RequestType::kGENERATION) - { - invokeSelectiveScanUpdate(ssm_params, stream); - } - sync_check_cuda_error(stream); - return 0; -} - -int SelectiveScanPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) noexcept -{ - if (isBuilding()) - { - return 0; - } - if (mType == DataType::kHALF) - { - return enqueueImpl(inputDesc, outputDesc, inputs, outputs, workspace, stream); - } - else if (mType == DataType::kFLOAT) - { - return enqueueImpl(inputDesc, outputDesc, inputs, outputs, workspace, stream); - } -#ifdef ENABLE_BF16 - else if (mType == DataType::kBF16) - { - return enqueueImpl<__nv_bfloat16>(inputDesc, outputDesc, inputs, outputs, workspace, stream); - } -#endif - return 0; -} - -// IPluginV2Ext Methods -nvinfer1::DataType SelectiveScanPlugin::getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept -{ - if (index == 0) - { - return inputTypes[getInputTensorIdx()]; - } - else - { - return inputTypes[getStateIdx()]; - } -} - -// IPluginV2 Methods - -char const* SelectiveScanPlugin::getPluginType() const noexcept -{ - return SELECTIVE_SCAN_PLUGIN_NAME; -} - -char const* SelectiveScanPlugin::getPluginVersion() const noexcept -{ - return SELECTIVE_SCAN_PLUGIN_VERSION; -} - -int SelectiveScanPlugin::getNbOutputs() const noexcept -{ - return mPagedState ? 1 : 2; -} - -int SelectiveScanPlugin::initialize() noexcept -{ - return 0; -} - -void SelectiveScanPlugin::terminate() noexcept {} - -size_t SelectiveScanPlugin::getSerializationSize() const noexcept -{ - return sizeof(mDim) + sizeof(mDState) + sizeof(mDtRank) + sizeof(mNHeads) + sizeof(mNGroups) + sizeof(mChunkSize) - + sizeof(mDeltaSoftplus) + sizeof(mType) + sizeof(mRemovePadding) + sizeof(mPagedState) + sizeof(mZEnabled) - + sizeof(mIsMamba2); -} - -void SelectiveScanPlugin::serialize(void* buffer) const noexcept -{ - char *d = static_cast(buffer), *a = d; - write(d, mDim); - write(d, mDState); - write(d, mDtRank); - write(d, mNHeads); - write(d, mNGroups); - write(d, mChunkSize); - write(d, mDeltaSoftplus); - write(d, mType); - write(d, mRemovePadding); - write(d, mPagedState); - write(d, mZEnabled); - write(d, mIsMamba2); - TLLM_CHECK(d == a + getSerializationSize()); -} - -void SelectiveScanPlugin::destroy() noexcept -{ - delete this; -} - -/////////////// - -SelectiveScanPluginCreator::SelectiveScanPluginCreator() -{ - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - mPluginAttributes.emplace_back(PluginField("dim", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("dstate", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("dt_rank", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("nheads", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("ngroups", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("chunk_size", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("delta_softplus", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("type_id", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("remove_input_padding", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("paged_state", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("z_enabled", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("is_mamba2", nullptr, PluginFieldType::kINT8)); - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -char const* SelectiveScanPluginCreator::getPluginName() const noexcept -{ - return SELECTIVE_SCAN_PLUGIN_NAME; -} - -char const* SelectiveScanPluginCreator::getPluginVersion() const noexcept -{ - return SELECTIVE_SCAN_PLUGIN_VERSION; -} - -PluginFieldCollection const* SelectiveScanPluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -IPluginV2* SelectiveScanPluginCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept -{ - PluginField const* fields = fc->fields; - int dim{}; - int dstate{}; - int dtRank{}; - int nHeads{}; - int nGroups{}; - int chunkSize{}; - bool deltaSoftplus{}; - bool removePadding{}; - bool pagedState{}; - bool zEnabled{}; - bool isMamab2{}; - nvinfer1::DataType type{}; - // Read configurations from each fields - for (int i = 0; i < fc->nbFields; ++i) - { - char const* attrName = fields[i].name; - if (!strcmp(attrName, "dim")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - dim = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "dstate")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - dstate = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "dt_rank")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - dtRank = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "nheads")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - nHeads = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "ngroups")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - nGroups = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "chunk_size")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - chunkSize = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "delta_softplus")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT8); - deltaSoftplus = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "type_id")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - type = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "remove_input_padding")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT8); - removePadding = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "paged_state")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT8); - pagedState = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "z_enabled")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT8); - zEnabled = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "is_mamba2")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT8); - isMamab2 = static_cast(*(static_cast(fields[i].data))); - } - } - try - { - auto* obj = new SelectiveScanPlugin(dim, dstate, dtRank, nHeads, nGroups, chunkSize, deltaSoftplus, type, - removePadding, pagedState, zEnabled, isMamab2); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -IPluginV2* SelectiveScanPluginCreator::deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept -{ - // This object will be deleted when the network is destroyed, which will - // call SelectiveScanPlugin::destroy() - try - { - auto* obj = new SelectiveScanPlugin(serialData, serialLength); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} diff --git a/cpp/tensorrt_llm/plugins/selectiveScanPlugin/selectiveScanPlugin.h b/cpp/tensorrt_llm/plugins/selectiveScanPlugin/selectiveScanPlugin.h deleted file mode 100644 index 96cb86fc4cbb..000000000000 --- a/cpp/tensorrt_llm/plugins/selectiveScanPlugin/selectiveScanPlugin.h +++ /dev/null @@ -1,218 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-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. - */ - -#ifndef TRT_SELECTIVE_SCAN_PLUGIN_H -#define TRT_SELECTIVE_SCAN_PLUGIN_H -#include "tensorrt_llm/kernels/selectiveScan/selectiveScan.h" -#include "tensorrt_llm/plugins/common/plugin.h" -#include - -namespace tensorrt_llm::plugins -{ -// batch_size = num_ctx_requests or num_gen_requests -// num_ctx_requests = number of context requests (single sequence per request). -// num_gen_requests = number of generation requests (single sequences per request). -// can not support beam search - -// inputs -// 0. input_tensor [batch_size, seq_len, dim] or [num_tokens, dim] for remove_input_padding -// 1. state, mamba: [batch_size, dstate, dim] or host [1] containing only pointer for paged_state -// mamba2: [batch_size, nheads, dstate, dim] or host [1] containing only pointer for paged_state -// 2. delta, mamba: [batch_size, seq_len, dim] or [num_tokens, dim] for remove_input_padding -// mamba2: [batch_size, seq_len, nheads] or [num_tokens, nheads] for remove_input_padding -// 3. delta_bias, [dim] for mamba, [nheads] for mamba2 -// 4. A, [dstate, dim] for mamba, [nheads] for mamba2 -// 5. BC, mamba: [batch_size, seq_len, dstate * 2] or [num_tokens, dstate * 2] for remove_input_padding -// mamba2: [batch_size, seq_len, ngroups * dstate * 2] or [num_tokens, ngroups * dstate * 2] for -// remove_input_padding -// 6. D, [dim] for mamba, [nheads] for mamba2 -// 7. host_request_types [batch_size] int32. 0: context; 1: generation; 2: none. -// 8. last_token_ids [batch_size] int32 -// 9. host_context_lengths [batch_size] int32, optional for remove_input_padding -// 10. state_slot_mapping [batch_size] int32, optional for paged state -// 11. z [batch_size, seq_len, dim] or [num_tokens, dim] for remove_input_padding -// outputs -// 0. output_tensor [batch_size, seq_len, dim] or [num_tokens, dim] for remove_input_padding -// 1. state, [batch_size, dstate, dim] for mamba, [batch_size, nheads, dstate, dim] for mamba2 - -class SelectiveScanPlugin : public BasePlugin -{ -public: - SelectiveScanPlugin(int dim, int dstate, int dtRank, int nHeads, int nGroups, int chunkSize, bool deltaSoftplus, - nvinfer1::DataType type, bool removePadding, bool pagedState, bool zEnabled, bool isMamba2); - - SelectiveScanPlugin(void const* data, size_t length); - - ~SelectiveScanPlugin() override = default; - - // IPluginV2DynamicExt Methods - nvinfer1::IPluginV2DynamicExt* clone() const noexcept override; - nvinfer1::DimsExprs getOutputDimensions(int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; - bool supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept override; - void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept override; - size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept override; - int enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - template - int enqueueImpl(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream); - - // IPluginV2Ext Methods - nvinfer1::DataType getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept override; - - // IPluginV2 Methods - char const* getPluginType() const noexcept override; - char const* getPluginVersion() const noexcept override; - int getNbOutputs() const noexcept override; - int initialize() noexcept override; - void terminate() noexcept override; - size_t getSerializationSize() const noexcept override; - void serialize(void* buffer) const noexcept override; - void destroy() noexcept override; - - enum class RequestType : int32_t - { - kCONTEXT = 0, - kGENERATION = 1 - }; - -private: - using IndexType = std::int32_t; - - IndexType getInputTensorIdx() const - { - return 0; - }; - - IndexType getStateIdx() const - { - return 1; - }; - - IndexType getDeltaIdx() const - { - return 2; - }; - - IndexType getDeltaBiasIdx() const - { - return 3; - }; - - IndexType getAIdx() const - { - return 4; - }; - - IndexType getBCIdx() const - { - return 5; - }; - - IndexType getDIdx() const - { - return 6; - }; - - IndexType getHostRequestTypesIdx() const - { - return 7; - }; - - IndexType getLastTokenIdsIdx() const - { - return 8; - }; - - IndexType getHostContextLengthIdx() const - { - if (mRemovePadding) - return 9; - else - return 8; - }; - - IndexType getSlotMappingIdx() const - { - if (mPagedState) - return getHostContextLengthIdx() + 1; - else - return getHostContextLengthIdx(); - }; - - IndexType getZIdx() const - { - if (mZEnabled) - return getSlotMappingIdx() + 1; - else - return getSlotMappingIdx(); - }; - - void setSSMParams(tensorrt_llm::kernels::SSMParamsBase& params, - // sizes - const size_t batch, const size_t dim, const size_t maxSeqLen, const size_t numTokens, const size_t dstate, - const size_t dtRank, const size_t nHeads, const size_t nGroups, const size_t chunkSize, - // device pointers - void* statePtr, void const* x, void const* delta, void const* deltaBias, void const* A, void const* BC, - void const* D, void const* z, void* osPtr, void* stPtr, void* dcPtr, void* dAPtr, void* cbPtr, void* descs, - int const* lastTokenIds, int const* slotMapping, void* out, bool deltaSoftplus, bool removePadding); - -private: - int mDim; - int mDState; - int mDtRank; - int mNHeads; - int mNGroups; - int mChunkSize; - bool mDeltaSoftplus; - nvinfer1::DataType mType; - bool mRemovePadding = false; - bool mPagedState = false; - bool mZEnabled = true; - bool mIsMamba2 = false; - std::shared_ptr mDriver; -}; - -class SelectiveScanPluginCreator : public BaseCreator -{ -public: - SelectiveScanPluginCreator(); - - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; - - nvinfer1::IPluginV2* deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept override; - -private: - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; -}; - -} // namespace tensorrt_llm::plugins - -#endif // TRT_SELECTIVE_SCAN_PLUGIN_H diff --git a/cpp/tensorrt_llm/plugins/smoothQuantGemmPlugin/CMakeLists.txt b/cpp/tensorrt_llm/plugins/smoothQuantGemmPlugin/CMakeLists.txt deleted file mode 100755 index 86876224fccd..000000000000 --- a/cpp/tensorrt_llm/plugins/smoothQuantGemmPlugin/CMakeLists.txt +++ /dev/null @@ -1,21 +0,0 @@ -# -# SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. -# -file(GLOB SRCS *.cpp) -set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS}) -set(PLUGIN_SOURCES - ${PLUGIN_SOURCES} - PARENT_SCOPE) diff --git a/cpp/tensorrt_llm/plugins/smoothQuantGemmPlugin/smoothQuantGemmPlugin.cpp b/cpp/tensorrt_llm/plugins/smoothQuantGemmPlugin/smoothQuantGemmPlugin.cpp deleted file mode 100644 index 718d8b7e830d..000000000000 --- a/cpp/tensorrt_llm/plugins/smoothQuantGemmPlugin/smoothQuantGemmPlugin.cpp +++ /dev/null @@ -1,431 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. - */ -#include "smoothQuantGemmPlugin.h" -#include "tensorrt_llm/kernels/weightOnlyBatchedGemv/int8SQ.h" -#include - -using namespace nvinfer1; -using namespace tensorrt_llm::common; -using namespace tensorrt_llm::kernels::cutlass_kernels; -using tensorrt_llm::plugins::SmoothQuantGemmPluginCreator; -using tensorrt_llm::plugins::SmoothQuantGemmPlugin; -using tensorrt_llm::plugins::SmoothQuantGemmPluginProfiler; -using tensorrt_llm::plugins::read; -using tensorrt_llm::plugins::write; - -static char const* SQ_GEMM_PLUGIN_VERSION{"1"}; -static char const* SQ_GEMM_PLUGIN_NAME{"SmoothQuantGemm"}; -PluginFieldCollection SmoothQuantGemmPluginCreator::mFC{}; -std::vector SmoothQuantGemmPluginCreator::mPluginAttributes; - -void SmoothQuantGemmPluginProfiler::runTactic(int m, int n, int k, SmoothQuantGemmPluginProfiler::Config const& tactic, - char* workspace, cudaStream_t const& stream) -{ - int8_t* aTmp = reinterpret_cast(workspace); - int8_t* bTmp = nextWorkspacePtr(aTmp, m * k * sizeof(int8_t)); - void* cTmp = reinterpret_cast(nextWorkspacePtr(bTmp, n * k * sizeof(int8_t))); - float* alphaRowTmp = reinterpret_cast( - nextWorkspacePtr(reinterpret_cast(cTmp), m * n * (mType == nvinfer1::DataType::kFLOAT ? 4 : 2))); - float* alphaColTmp - = reinterpret_cast(nextWorkspacePtr(reinterpret_cast(alphaRowTmp), m * sizeof(float))); - char* workspaceTmp - = reinterpret_cast(nextWorkspacePtr(reinterpret_cast(alphaColTmp), n * sizeof(float))); - - int const wsSize = mRunner->getWorkspaceSize(m, n, k); - - mRunner->gemm( - aTmp, bTmp, mQuantMode, alphaColTmp, alphaRowTmp, cTmp, m, n, k, tactic, workspaceTmp, wsSize, stream); -} - -void SmoothQuantGemmPluginProfiler::computeTmpSize(size_t maxM, size_t n, size_t k) -{ - std::vector workspaces = { - maxM * k * sizeof(int8_t), // A - n * k * sizeof(int8_t), // B - maxM * n * (mType == nvinfer1::DataType::kFLOAT ? 4u : 2u), // C - maxM * sizeof(float), // alphaRow - n * sizeof(float), // alphaCol - mRunner->getWorkspaceSize(maxM, n, k) // workspace - }; - size_t bytes = calculateTotalWorkspaceSize(workspaces.data(), workspaces.size()); - setTmpWorkspaceSizeInBytes(bytes); -} - -std::vector SmoothQuantGemmPluginProfiler::getTactics(int m, int n, int k) const -{ - return mRunner->getConfigs(); -} - -SmoothQuantGemmPlugin::SmoothQuantGemmPlugin( - QuantMode quantMode, nvinfer1::DataType type, SmoothQuantGemmPlugin::PluginProfilerPtr const& pluginProfiler) - : mQuantMode(quantMode) - , mPluginProfiler(pluginProfiler) -{ - init(type); -} - -// Parameterized constructor -SmoothQuantGemmPlugin::SmoothQuantGemmPlugin( - void const* data, size_t length, SmoothQuantGemmPlugin::PluginProfilerPtr const& pluginProfiler) - : mPluginProfiler(pluginProfiler) -{ - char const *d = reinterpret_cast(data), *a = d; - nvinfer1::DataType type; - unsigned int quantMode; - read(d, quantMode); - read(d, type); - read(d, mDims); - - mQuantMode = QuantMode(quantMode); - - init(type); - - mPluginProfiler->deserialize(d, mDims, mGemmId); - - TLLM_CHECK_WITH_INFO(d == a + length, - "Expected length (%d) != real length (%d). This is often " - "caused by using different TensorRT LLM version to build " - "engine and run engine.", - (int) length, (int) (d - a)); -} - -void SmoothQuantGemmPlugin::init(nvinfer1::DataType type) -{ - mType = type; - if (mType == nvinfer1::DataType::kHALF) - { - m_sqGemmRunner = std::make_shared>(); - } - else if (mType == nvinfer1::DataType::kFLOAT) - { - m_sqGemmRunner = std::make_shared>(); - } - else if (mType == nvinfer1::DataType::kINT32) - { - m_sqGemmRunner = std::make_shared>(); - } -#ifdef ENABLE_BF16 - else if (mType == nvinfer1::DataType::kBF16) - { - m_sqGemmRunner = std::make_shared>(); - } -#endif - - mPluginProfiler->setQuantMode(mQuantMode); - - mGemmId = GemmIdCore(mDims.n, mDims.k, mType); -} - -// IPluginV2DynamicExt Methods -nvinfer1::IPluginV2DynamicExt* SmoothQuantGemmPlugin::clone() const noexcept -{ - auto* plugin = new SmoothQuantGemmPlugin(*this); - return plugin; -} - -nvinfer1::DimsExprs SmoothQuantGemmPlugin::getOutputDimensions( - int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept -{ - try - { - TLLM_CHECK(nbInputs == 4); - TLLM_CHECK(outputIndex == 0); - int const nbDimsA = inputs[0].nbDims; - TLLM_CHECK(nbDimsA >= 2); - DimsExprs ret; - ret.nbDims = nbDimsA; - for (int ii = 0; ii < nbDimsA - 1; ++ii) - { - ret.d[ii] = inputs[0].d[ii]; - } - ret.d[nbDimsA - 1] = inputs[1].d[0]; - return ret; - } - catch (std::exception const& e) - { - caughtError(e); - } - return DimsExprs{}; -} - -bool SmoothQuantGemmPlugin::supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept -{ - switch (pos) - { - case 0: - // activation - return inOut[pos].type == nvinfer1::DataType::kINT8 && inOut[pos].format == TensorFormat::kLINEAR; - case 1: - // weights - // Weights stored in checkpoint must have int8 type - return inOut[pos].type == nvinfer1::DataType::kINT8 && inOut[pos].format == TensorFormat::kLINEAR; - case 2: - // scales channels - case 3: - // scales tokens - return inOut[pos].type == nvinfer1::DataType::kFLOAT && inOut[pos].format == TensorFormat::kLINEAR; - case 4: - // out - return inOut[pos].type == mType && inOut[pos].format == TensorFormat::kLINEAR; - default: - // Never should be here - assert(false); - return false; - } -} - -void SmoothQuantGemmPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept -{ - auto const minM = std::accumulate(in[0].min.d, in[0].min.d + in[0].min.nbDims - 1, 1, std::multiplies()); - auto const maxM = std::accumulate(in[0].max.d, in[0].max.d + in[0].max.nbDims - 1, 1, std::multiplies()); - - int const maxK = in[0].max.d[in[0].max.nbDims - 1]; - int const maxN = in[1].max.d[0]; - int const minK = in[0].min.d[in[0].min.nbDims - 1]; - int const minN = in[1].min.d[0]; - - TLLM_CHECK_WITH_INFO(minN == maxN, "Variable out channels is not allowed"); - TLLM_CHECK_WITH_INFO(minK == maxK, "Variable in channels is not allowed"); - - if (!mDims.isInitialized()) - { - mDims = {minM, maxM, maxN, maxK}; - } - mGemmId = {maxN, maxK, mType}; - - m_workspaceMaxSize = m_sqGemmRunner->getWorkspaceSize(maxM, maxN, maxK); -} - -size_t SmoothQuantGemmPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - return m_workspaceMaxSize; -} - -int SmoothQuantGemmPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) noexcept -{ - // inputs - // mat1 [M(*), K] - // mat2 [N, K] - // scale_tokens [M, 1] if has_per_token_scaling else [1, 1] - // scale_channels [1, N] if has_per_channel_scaling else [1, 1] - // outputs - // mat [M(*), N] - int64_t m64 = 1; - for (int ii = 0; ii < inputDesc[0].dims.nbDims - 1; ++ii) - { - m64 *= inputDesc[0].dims.d[ii]; - } - int const m = TLLM_INT32_CAST(m64); - int const n = TLLM_INT32_CAST(inputDesc[1].dims.d[0]); - int const k = TLLM_INT32_CAST(inputDesc[0].dims.d[inputDesc[0].dims.nbDims - 1]); - int const wsSize = m_sqGemmRunner->getWorkspaceSize(m, n, k); - if (m <= 4) - { - tensorrt_llm::kernels::smooth_quant::Params params(reinterpret_cast(inputs[0]), - reinterpret_cast(inputs[1]), reinterpret_cast(inputs[2]), - reinterpret_cast(inputs[3]), reinterpret_cast(outputs[0]), m, n, k, mQuantMode); - if (mType == nvinfer1::DataType::kHALF) - { - tensorrt_llm::kernels::smooth_quant::int8_sq_launcher(params, stream); - } - else if (mType == nvinfer1::DataType::kFLOAT) - { - tensorrt_llm::kernels::smooth_quant::int8_sq_launcher(params, stream); - } -#ifdef ENABLE_BF16 - else if (mType == nvinfer1::DataType::kBF16) - { - tensorrt_llm::kernels::smooth_quant::int8_sq_launcher<__nv_bfloat16>(params, stream); - } -#endif - else if (mType == nvinfer1::DataType::kINT32) - { - tensorrt_llm::kernels::smooth_quant::int8_sq_launcher(params, stream); - } - } - else - { - auto const& bestTactic = mPluginProfiler->getBestConfig(m, mGemmId); - TLLM_CHECK_WITH_INFO(bestTactic, "No valid SQ GEMM tactic"); - m_sqGemmRunner->gemm(reinterpret_cast(inputs[0]), reinterpret_cast(inputs[1]), - mQuantMode, reinterpret_cast(inputs[3]), reinterpret_cast(inputs[2]), - reinterpret_cast(outputs[0]), m, n, k, *bestTactic, reinterpret_cast(workspace), wsSize, - stream); - } - - return 0; -} - -// IPluginV2Ext Methods -nvinfer1::DataType SmoothQuantGemmPlugin::getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept -{ - TLLM_CHECK(index == 0); - return mType; -} - -// IPluginV2 Methods - -char const* SmoothQuantGemmPlugin::getPluginType() const noexcept -{ - return SQ_GEMM_PLUGIN_NAME; -} - -char const* SmoothQuantGemmPlugin::getPluginVersion() const noexcept -{ - return SQ_GEMM_PLUGIN_VERSION; -} - -int SmoothQuantGemmPlugin::getNbOutputs() const noexcept -{ - return 1; -} - -int SmoothQuantGemmPlugin::initialize() noexcept -{ - configGemm(); - return 0; -} - -void SmoothQuantGemmPlugin::terminate() noexcept {} - -size_t SmoothQuantGemmPlugin::getSerializationSize() const noexcept -{ - return sizeof(unsigned int) + // QuantMode - sizeof(nvinfer1::DataType) + // dtype - sizeof(mDims) + // Dimensions - mPluginProfiler->getSerializationSize(mGemmId); // selected tactics container size -} - -void SmoothQuantGemmPlugin::serialize(void* buffer) const noexcept -{ - char *d = static_cast(buffer), *a = d; - write(d, mQuantMode.value()); - write(d, mType); - write(d, mDims); - - mPluginProfiler->serialize(d, mGemmId); - TLLM_CHECK(d == a + getSerializationSize()); -} - -void SmoothQuantGemmPlugin::destroy() noexcept -{ - // This gets called when the network containing plugin is destroyed - delete this; -} - -void SmoothQuantGemmPlugin::configGemm() -{ - mPluginProfiler->profileTactics(m_sqGemmRunner, mType, mDims, mGemmId); -} - -/////////////// - -SmoothQuantGemmPluginCreator::SmoothQuantGemmPluginCreator() -{ - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - mPluginAttributes.emplace_back(PluginField("has_per_channel_scaling", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("has_per_token_scaling", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("type_id", nullptr, PluginFieldType::kINT32)); - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -char const* SmoothQuantGemmPluginCreator::getPluginName() const noexcept -{ - return SQ_GEMM_PLUGIN_NAME; -} - -char const* SmoothQuantGemmPluginCreator::getPluginVersion() const noexcept -{ - return SQ_GEMM_PLUGIN_VERSION; -} - -PluginFieldCollection const* SmoothQuantGemmPluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -IPluginV2* SmoothQuantGemmPluginCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept -{ - PluginField const* fields = fc->fields; - bool perTokenScaling{}; - bool perChannelScaling{}; - nvinfer1::DataType type{}; - // Read configurations from each fields - for (int i = 0; i < fc->nbFields; ++i) - { - char const* attrName = fields[i].name; - if (!strcmp(attrName, "has_per_channel_scaling")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - perChannelScaling = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "has_per_token_scaling")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - perTokenScaling = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "type_id")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - type = static_cast(*(static_cast(fields[i].data))); - } - } - try - { - // SmoothQuantGemmPluginCreator is unique and shared for an engine generation - // Create plugin profiler with shared tactics map - auto pluginProfiler = gemmPluginProfileManager.createGemmPluginProfiler(/* inference */ false); - QuantMode quantMode = QuantMode::fromDescription(true, true, perTokenScaling, perChannelScaling, false, false, - false, false, false, false, false, false, false, false, false, false); - auto* obj = new SmoothQuantGemmPlugin(quantMode, type, pluginProfiler); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -IPluginV2* SmoothQuantGemmPluginCreator::deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept -{ - // This object will be deleted when the network is destroyed, which will - // call SmoothQuantGemmPlugin::destroy() - try - { - // Create plugin profiler with private tactics map which is read from the serialized engine - auto pluginProfiler = gemmPluginProfileManager.createGemmPluginProfiler(/* inference */ true); - auto* obj = new SmoothQuantGemmPlugin(serialData, serialLength, pluginProfiler); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} diff --git a/cpp/tensorrt_llm/plugins/smoothQuantGemmPlugin/smoothQuantGemmPlugin.h b/cpp/tensorrt_llm/plugins/smoothQuantGemmPlugin/smoothQuantGemmPlugin.h deleted file mode 100644 index 3cabf558076b..000000000000 --- a/cpp/tensorrt_llm/plugins/smoothQuantGemmPlugin/smoothQuantGemmPlugin.h +++ /dev/null @@ -1,140 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. - */ -#pragma once - -#include "tensorrt_llm/common/quantization.h" -#include "tensorrt_llm/kernels/cutlass_kernels/int8_gemm/int8_gemm.h" -#include "tensorrt_llm/plugins/common/gemmPluginProfiler.h" -#include "tensorrt_llm/plugins/common/plugin.h" -#include -#include -#include -#include -#include - -namespace tensorrt_llm::plugins -{ - -using perfMapType = std::unordered_map; -using SqGemmRunnerPtr = std::shared_ptr; - -class SmoothQuantGemmPluginProfiler : public GemmPluginProfiler -{ -public: - using Config = tensorrt_llm::cutlass_extensions::CutlassGemmConfig; - - void setQuantMode(tensorrt_llm::common::QuantMode const& quantMode) - { - mQuantMode = quantMode; - } - -protected: - void runTactic(int m, int n, int k, Config const& tactic, char* workspace, cudaStream_t const& stream) override; - - void computeTmpSize(size_t maxM, size_t n, size_t k) override; - - std::vector getTactics(int m, int n, int k) const override; - -private: - tensorrt_llm::common::QuantMode mQuantMode; -}; - -class SmoothQuantGemmPlugin : public BasePlugin -{ -public: - using PluginProfilerPtr = std::shared_ptr; - - SmoothQuantGemmPlugin() = delete; - - SmoothQuantGemmPlugin( - tensorrt_llm::common::QuantMode quantMode, nvinfer1::DataType type, PluginProfilerPtr const& pluginProfiler); - - SmoothQuantGemmPlugin(void const* data, size_t length, PluginProfilerPtr const& pluginProfiler); - - ~SmoothQuantGemmPlugin() override = default; - - // IPluginV2DynamicExt Methods - nvinfer1::IPluginV2DynamicExt* clone() const noexcept override; - nvinfer1::DimsExprs getOutputDimensions(int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; - bool supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept override; - void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept override; - size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept override; - int enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - - // IPluginV2Ext Methods - nvinfer1::DataType getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept override; - - // IPluginV2 Methods - char const* getPluginType() const noexcept override; - char const* getPluginVersion() const noexcept override; - int getNbOutputs() const noexcept override; - int initialize() noexcept override; - void terminate() noexcept override; - size_t getSerializationSize() const noexcept override; - void serialize(void* buffer) const noexcept override; - void destroy() noexcept override; - -private: - void init(nvinfer1::DataType type); - - void configGemm(); - -private: - const std::string mLayerName; - - SqGemmRunnerPtr m_sqGemmRunner; - tensorrt_llm::common::QuantMode mQuantMode; - size_t m_workspaceMaxSize; - - GemmDims mDims{}; - GemmIdCore mGemmId{}; - - PluginProfilerPtr mPluginProfiler; - - nvinfer1::DataType mType; -}; - -class SmoothQuantGemmPluginCreator : public BaseCreator -{ -public: - SmoothQuantGemmPluginCreator(); - - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; - - nvinfer1::IPluginV2* deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept override; - -private: - GemmPluginProfilerManager gemmPluginProfileManager; - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; -}; - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/topkLastDimPlugin/CMakeLists.txt b/cpp/tensorrt_llm/plugins/topkLastDimPlugin/CMakeLists.txt deleted file mode 100644 index 6b4e3d8d9e0f..000000000000 --- a/cpp/tensorrt_llm/plugins/topkLastDimPlugin/CMakeLists.txt +++ /dev/null @@ -1,22 +0,0 @@ -# -# SPDX-FileCopyrightText: Copyright (c) 1993-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. -# - -file(GLOB SRCS *.cpp) -set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS}) -set(PLUGIN_SOURCES - ${PLUGIN_SOURCES} - PARENT_SCOPE) diff --git a/cpp/tensorrt_llm/plugins/topkLastDimPlugin/topkLastDimPlugin.cpp b/cpp/tensorrt_llm/plugins/topkLastDimPlugin/topkLastDimPlugin.cpp deleted file mode 100644 index 072bfc9c8fc4..000000000000 --- a/cpp/tensorrt_llm/plugins/topkLastDimPlugin/topkLastDimPlugin.cpp +++ /dev/null @@ -1,316 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-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. - */ - -#include "topkLastDimPlugin.h" -#include "tensorrt_llm/common/assert.h" - -using namespace nvinfer1; -using namespace tensorrt_llm::kernels; -using namespace tensorrt_llm::common; -using tensorrt_llm::plugins::TopkLastDimPluginCreator; -using tensorrt_llm::plugins::TopkLastDimPlugin; - -static char const* TOPK_LAST_DIM_PLUGIN_VERSION{"1"}; -static char const* TOPK_LAST_DIM_PLUGIN_NAME{"TopkLastDim"}; -PluginFieldCollection TopkLastDimPluginCreator::mFC{}; -std::vector TopkLastDimPluginCreator::mPluginAttributes; - -TopkLastDimPlugin::TopkLastDimPlugin(nvinfer1::DataType type, int32_t k, bool is_largest) - : mType(type) - , mK(k) // To avoid data-dependent shape, enforce K to be non-dynamic - , mIsLargest(is_largest) -{ - TLLM_CHECK_WITH_INFO((mType == DataType::kBF16) || (mType == DataType::kFLOAT) || (mType == DataType::kHALF) - || (mType == DataType::kINT32), - "Only support int, float, half, and bfloat16."); -} - -// Parameterized constructor -TopkLastDimPlugin::TopkLastDimPlugin(void const* data, size_t length) -{ - char const *d = reinterpret_cast(data), *a = d; - read(d, mType); - read(d, mK); - read(d, mIsLargest); - TLLM_CHECK(d == a + length); - TLLM_CHECK_WITH_INFO((mType == DataType::kBF16) || (mType == DataType::kFLOAT) || (mType == DataType::kHALF) - || (mType == DataType::kINT32), - "Only support int, float, half, and bfloat16."); -} - -// IPluginV2DynamicExt Methods -nvinfer1::IPluginV2DynamicExt* TopkLastDimPlugin::clone() const noexcept -{ - auto* plugin = new TopkLastDimPlugin(mType, mK, mIsLargest); - plugin->setPluginNamespace(mNamespace.c_str()); - return plugin; -} - -// Outputs -// out_val or out_idx: [batch_size, K] -nvinfer1::DimsExprs TopkLastDimPlugin::getOutputDimensions( - int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept -{ - TLLM_CHECK_WITH_INFO(outputIndex < 2, "Only 2 outputs."); - nvinfer1::DimsExprs output(inputs[0]); - int numDim = output.nbDims; - output.d[numDim - 1] = exprBuilder.constant(mK); - return output; -} - -bool TopkLastDimPlugin::supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept -{ - bool res = inOut[pos].format == TensorFormat::kLINEAR; - if (pos < 2) // input and out_val tensor must be the same type as the plugin - { - res = res && inOut[pos].type == mType; - } - else if (pos == 2) // out_idx must be int32 - { - res = res && inOut[pos].type == DataType::kINT32; - } - return res; -} - -void TopkLastDimPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept -{ -} - -size_t TopkLastDimPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - // extract shape info and then call helper - auto const batchSize = inputs[getInputTensorIdx()].dims.d[0]; - auto const inputLength = inputs[getInputTensorIdx()].dims.d[1]; - size_t tempStorageBytes{}; - if (mType == DataType::kINT32) - { - tempStorageBytes = invokeComputeTopkLastDimWorkspaceSize(batchSize, inputLength, mK, mIsLargest); - } - else if (mType == DataType::kHALF) - { - tempStorageBytes = invokeComputeTopkLastDimWorkspaceSize(batchSize, inputLength, mK, mIsLargest); - } - else if (mType == DataType::kFLOAT) - { - tempStorageBytes = invokeComputeTopkLastDimWorkspaceSize(batchSize, inputLength, mK, mIsLargest); - } -#ifdef ENABLE_BF16 - else if (mType == DataType::kBF16) - { - tempStorageBytes = invokeComputeTopkLastDimWorkspaceSize<__nv_bfloat16>(batchSize, inputLength, mK, mIsLargest); - } -#endif - return tempStorageBytes; -} - -template -int TopkLastDimPlugin::enqueueImpl(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) -{ - // inputs - // 0. input_tensor [batch_size, inputLength] - // outputs - // 0. output_values [batch_size, k] - // 1. output_indices [batch_size, k] - auto const batchSize = inputDesc[getInputTensorIdx()].dims.d[0]; - auto const inputLength = inputDesc[getInputTensorIdx()].dims.d[1]; - if (batchSize == 0) - { - // nothing to do for empty tensor - return 0; - } - - invokeTopkLastDim( - batchSize, inputLength, mK, mIsLargest, inputs[getInputTensorIdx()], outputs[0], outputs[1], workspace, stream); - - sync_check_cuda_error(stream); - return 0; -} - -int TopkLastDimPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) noexcept -{ - if (mType == DataType::kINT32) - { - return enqueueImpl(inputDesc, outputDesc, inputs, outputs, workspace, stream); - } - else if (mType == DataType::kHALF) - { - return enqueueImpl(inputDesc, outputDesc, inputs, outputs, workspace, stream); - } - else if (mType == DataType::kFLOAT) - { - return enqueueImpl(inputDesc, outputDesc, inputs, outputs, workspace, stream); - } -#ifdef ENABLE_BF16 - else if (mType == DataType::kBF16) - { - return enqueueImpl<__nv_bfloat16>(inputDesc, outputDesc, inputs, outputs, workspace, stream); - } -#endif - return 0; -} - -// IPluginV2Ext Methods -nvinfer1::DataType TopkLastDimPlugin::getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept -{ - TLLM_CHECK_WITH_INFO(index < 2, "Only 2 outputs."); - nvinfer1::DataType data_type; - if (index == 1) - { - data_type = DataType::kINT32; - } - else - { - data_type = inputTypes[getInputTensorIdx()]; - } - return data_type; -} - -// IPluginV2 Methods - -char const* TopkLastDimPlugin::getPluginType() const noexcept -{ - return TOPK_LAST_DIM_PLUGIN_NAME; -} - -char const* TopkLastDimPlugin::getPluginVersion() const noexcept -{ - return TOPK_LAST_DIM_PLUGIN_VERSION; -} - -int TopkLastDimPlugin::getNbOutputs() const noexcept -{ - return 2; -} - -int TopkLastDimPlugin::initialize() noexcept -{ - return 0; -} - -void TopkLastDimPlugin::terminate() noexcept {} - -size_t TopkLastDimPlugin::getSerializationSize() const noexcept -{ - return sizeof(mType) + sizeof(mK) + sizeof(mIsLargest); -} - -void TopkLastDimPlugin::serialize(void* buffer) const noexcept -{ - char *d = static_cast(buffer), *a = d; - write(d, mType); - write(d, mK); - write(d, mIsLargest); - TLLM_CHECK(d == a + getSerializationSize()); -} - -void TopkLastDimPlugin::destroy() noexcept -{ - delete this; -} - -/////////////// - -TopkLastDimPluginCreator::TopkLastDimPluginCreator() -{ - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - mPluginAttributes.emplace_back(PluginField("type_id", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("k", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("is_largest", nullptr, PluginFieldType::kINT32)); - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -char const* TopkLastDimPluginCreator::getPluginName() const noexcept -{ - return TOPK_LAST_DIM_PLUGIN_NAME; -} - -char const* TopkLastDimPluginCreator::getPluginVersion() const noexcept -{ - return TOPK_LAST_DIM_PLUGIN_VERSION; -} - -PluginFieldCollection const* TopkLastDimPluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -IPluginV2* TopkLastDimPluginCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept -{ - PluginField const* fields = fc->fields; - nvinfer1::DataType type{}; - int32_t k{}; - bool is_largest{}; - // Read configurations from each fields - for (int i = 0; i < fc->nbFields; ++i) - { - char const* attrName = fields[i].name; - if (!strcmp(attrName, "type_id")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - type = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "k")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - k = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "is_largest")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - is_largest = static_cast(*(static_cast(fields[i].data))) != 0; - } - } - try - { - auto* obj = new TopkLastDimPlugin(type, k, is_largest); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -IPluginV2* TopkLastDimPluginCreator::deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept -{ - // This object will be deleted when the network is destroyed, which will - // call TopkLastDimPlugin::destroy() - try - { - auto* obj = new TopkLastDimPlugin(serialData, serialLength); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} diff --git a/cpp/tensorrt_llm/plugins/topkLastDimPlugin/topkLastDimPlugin.h b/cpp/tensorrt_llm/plugins/topkLastDimPlugin/topkLastDimPlugin.h deleted file mode 100644 index 0ca38ccfe105..000000000000 --- a/cpp/tensorrt_llm/plugins/topkLastDimPlugin/topkLastDimPlugin.h +++ /dev/null @@ -1,99 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-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. - */ - -#ifndef TRT_TOPK_LAST_DIM_PLUGIN_H -#define TRT_TOPK_LAST_DIM_PLUGIN_H - -#include "tensorrt_llm/kernels/topkLastDim.h" -#include "tensorrt_llm/plugins/common/plugin.h" -#include - -namespace tensorrt_llm::plugins -{ -class TopkLastDimPlugin : public BasePlugin -{ -public: - TopkLastDimPlugin(nvinfer1::DataType type, int32_t k, bool largest); - TopkLastDimPlugin(void const* data, size_t length); - ~TopkLastDimPlugin() override = default; - // IPluginV2DynamicExt Methods - nvinfer1::IPluginV2DynamicExt* clone() const noexcept override; - nvinfer1::DimsExprs getOutputDimensions(int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; - bool supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept override; - void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept override; - size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept override; - int enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - template - int enqueueImpl(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream); - - // IPluginV2Ext Methods - nvinfer1::DataType getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept override; - - // IPluginV2 Methods - char const* getPluginType() const noexcept override; - char const* getPluginVersion() const noexcept override; - int getNbOutputs() const noexcept override; - int initialize() noexcept override; - void terminate() noexcept override; - size_t getSerializationSize() const noexcept override; - void serialize(void* buffer) const noexcept override; - void destroy() noexcept override; - -private: - using IndexType = std::int32_t; - - IndexType getInputTensorIdx() const - { - return 0; - }; - -private: - nvinfer1::DataType mType; - int32_t mK; - bool mIsLargest; -}; - -class TopkLastDimPluginCreator : public BaseCreator -{ -public: - TopkLastDimPluginCreator(); - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; - - nvinfer1::IPluginV2* deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept override; - -private: - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; -}; - -} // namespace tensorrt_llm::plugins - -#endif diff --git a/cpp/tensorrt_llm/plugins/weightOnlyGroupwiseQuantMatmulPlugin/CMakeLists.txt b/cpp/tensorrt_llm/plugins/weightOnlyGroupwiseQuantMatmulPlugin/CMakeLists.txt deleted file mode 100755 index 86876224fccd..000000000000 --- a/cpp/tensorrt_llm/plugins/weightOnlyGroupwiseQuantMatmulPlugin/CMakeLists.txt +++ /dev/null @@ -1,21 +0,0 @@ -# -# SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. -# -file(GLOB SRCS *.cpp) -set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS}) -set(PLUGIN_SOURCES - ${PLUGIN_SOURCES} - PARENT_SCOPE) diff --git a/cpp/tensorrt_llm/plugins/weightOnlyGroupwiseQuantMatmulPlugin/weightOnlyGroupwiseQuantMatmulPlugin.cpp b/cpp/tensorrt_llm/plugins/weightOnlyGroupwiseQuantMatmulPlugin/weightOnlyGroupwiseQuantMatmulPlugin.cpp deleted file mode 100644 index 85f0cf011293..000000000000 --- a/cpp/tensorrt_llm/plugins/weightOnlyGroupwiseQuantMatmulPlugin/weightOnlyGroupwiseQuantMatmulPlugin.cpp +++ /dev/null @@ -1,657 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. - */ -#include "weightOnlyGroupwiseQuantMatmulPlugin.h" - -#include - -using namespace nvinfer1; -using namespace tensorrt_llm::common; -using namespace tensorrt_llm::kernels::cutlass_kernels; -using tensorrt_llm::plugins::WeightOnlyGroupwiseQuantMatmulPluginCreator; -using tensorrt_llm::plugins::WeightOnlyGroupwiseQuantMatmulPlugin; -using tensorrt_llm::plugins::WeightOnlyGroupwiseQuantGemmPluginProfiler; -using tensorrt_llm::plugins::WeightOnlyGemmRunnerPtr; -using tensorrt_llm::plugins::read; -using tensorrt_llm::plugins::write; - -static char const* WOQ_GROUPWISE_MATMUL_PLUGIN_VERSION{"1"}; -static char const* WOQ_GROUPWISE_MATMUL_PLUGIN_NAME{"WeightOnlyGroupwiseQuantMatmul"}; -PluginFieldCollection WeightOnlyGroupwiseQuantMatmulPluginCreator::mFC{}; -std::vector WeightOnlyGroupwiseQuantMatmulPluginCreator::mPluginAttributes; - -void WeightOnlyGroupwiseQuantGemmPluginProfiler::runTactic(int m, int n, int k, - WeightOnlyGroupwiseQuantGemmPluginProfiler::Config const& tactic, char* workspace, cudaStream_t const& stream) -{ - // Quantized weights are packed in FP16 format (INT4*4 -> FP16, INT8*2 -> FP16) - int const originalN = mQuantAlgo & GroupwiseQuantAlgo::INT8_WEIGHT ? n * FP16_INT8_RATIO : n * FP16_INT4_RATIO; - half* actPtr = reinterpret_cast(workspace); - void* weightPtr = nextWorkspacePtr(reinterpret_cast(actPtr), m * k * sizeof(half)); - half* inputScalesPtr - = reinterpret_cast(nextWorkspacePtr(reinterpret_cast(weightPtr), n * k * sizeof(float))); - half* zerosPtr = reinterpret_cast( - nextWorkspacePtr(reinterpret_cast(inputScalesPtr), k * originalN * sizeof(half) / mGroupSize)); - half* biasesPtr = reinterpret_cast( - nextWorkspacePtr(reinterpret_cast(zerosPtr), k * originalN * sizeof(half) / mGroupSize)); - half* outputPtr = reinterpret_cast(nextWorkspacePtr(reinterpret_cast(biasesPtr), n * sizeof(half))); - char* workspacePtr - = reinterpret_cast(nextWorkspacePtr(reinterpret_cast(outputPtr), m * originalN * sizeof(half))); - if ((mQuantAlgo & GroupwiseQuantAlgo::ZERO) == 0) - { - zerosPtr = nullptr; - } - if ((mQuantAlgo & GroupwiseQuantAlgo::BIAS) == 0) - { - biasesPtr = nullptr; - } - - if (tactic.enableCudaKernel) - { - // run CUDA kernel - void const* pre_quant_scale_ptr = nullptr; - bool apply_alpha_in_advance = false; - float alpha = 1.0; - tensorrt_llm::kernels::weight_only::Params params{actPtr, pre_quant_scale_ptr, weightPtr, inputScalesPtr, - zerosPtr, biasesPtr, outputPtr, alpha, m, originalN, k, mGroupSize, mCudaKernelType, - apply_alpha_in_advance}; - tensorrt_llm::kernels::weight_only::kernel_launcher(mArch, params, stream); - } - else - { - // run CUTLASS kernel - int const wsSize = mRunner->getWorkspaceSize(m, originalN, k); - if (mQuantAlgo & GroupwiseQuantAlgo::INT8_WEIGHT) - { - mRunner->gemm(actPtr, reinterpret_cast(weightPtr), inputScalesPtr, zerosPtr, biasesPtr, outputPtr, - m, originalN, k, mGroupSize, tactic, workspacePtr, wsSize, stream); - } - else - { - mRunner->gemm(actPtr, reinterpret_cast(weightPtr), inputScalesPtr, zerosPtr, biasesPtr, - outputPtr, m, originalN, k, mGroupSize, tactic, workspacePtr, wsSize, stream); - } - } -} - -void WeightOnlyGroupwiseQuantGemmPluginProfiler::computeTmpSize(size_t maxM, size_t n, size_t k) -{ - // Quantized weights are packed in FP16 format (INT4*4 -> FP16, INT8*2 -> FP16) - int const originalN = mQuantAlgo & GroupwiseQuantAlgo::INT8_WEIGHT ? n * FP16_INT8_RATIO : n * FP16_INT4_RATIO; - std::vector workspaces = { - maxM * k * sizeof(half), // A - k * n * sizeof(float), // B - k * originalN * sizeof(half) / mGroupSize, // scales - k * originalN * sizeof(half) / mGroupSize, // zeros - originalN * sizeof(half), // biases - maxM * originalN * sizeof(half), // C - mRunner->getWorkspaceSize(maxM, originalN, k) // workspace - }; - size_t bytes = calculateTotalWorkspaceSize(workspaces.data(), workspaces.size()); - setTmpWorkspaceSizeInBytes(bytes); -} - -std::vector WeightOnlyGroupwiseQuantGemmPluginProfiler::getTactics( - int m, int n, int k) const -{ - return mRunner->getConfigs(); -} - -bool WeightOnlyGroupwiseQuantGemmPluginProfiler::checkTactic(int m, int n, int k, Config const& tactic) const -{ - // stop to profile Cuda kernel for m >= 16 - if (tactic.enableCudaKernel) - { - return m < 16; - } - return true; -} - -WeightOnlyGroupwiseQuantMatmulPlugin::WeightOnlyGroupwiseQuantMatmulPlugin(nvinfer1::DataType type, int quant_algo, - int group_size, float alpha, WeightOnlyGroupwiseQuantMatmulPlugin::PluginProfilerPtr const& pluginProfiler) - : mPluginProfiler(pluginProfiler) -{ - init(type, quant_algo, group_size, alpha); -} - -// Parameterized constructor -WeightOnlyGroupwiseQuantMatmulPlugin::WeightOnlyGroupwiseQuantMatmulPlugin( - void const* data, size_t length, WeightOnlyGroupwiseQuantMatmulPlugin::PluginProfilerPtr const& pluginProfiler) - : mPluginProfiler(pluginProfiler) -{ - char const *d = reinterpret_cast(data), *a = d; - nvinfer1::DataType type; - int quant_algo = 0; - int group_size = 0; - float alpha = 1.0f; - read(d, type); - read(d, quant_algo); - read(d, group_size); - read(d, alpha); - read(d, mDims); - - init(type, quant_algo, group_size, alpha); - - mPluginProfiler->deserialize(d, mDims, mGemmId); - - TLLM_CHECK_WITH_INFO(d == a + length, - "Expected length (%d) != real length (%d). This is often " - "caused by using different TensorRT LLM version to build " - "engine and run engine.", - (int) length, (int) (d - a)); -} - -template -using GemmRunner = tensorrt_llm::kernels::cutlass_kernels::CutlassFpAIntBGemmRunner; - -template -WeightOnlyGemmRunnerPtr selectGemmRunnerForZERO(int quant_algo) -{ - if (quant_algo & GroupwiseQuantAlgo::ZERO) - { - return std::make_shared>(); - } - else - { - return std::make_shared>(); - } -} - -template -WeightOnlyGemmRunnerPtr selectGemmRunnerForWeightType(int quant_algo) -{ - if (quant_algo & GroupwiseQuantAlgo::INT8_WEIGHT) - { - return selectGemmRunnerForZERO(quant_algo); - } - else - { - return selectGemmRunnerForZERO(quant_algo); - } -} - -void WeightOnlyGroupwiseQuantMatmulPlugin::init(nvinfer1::DataType type, int quant_algo, int group_size, float alpha) -{ - mArch = tensorrt_llm::common::getSMVersion(); - mType = type; - mQuantAlgo = quant_algo; - mGroupSize = group_size; - - // quant_algo = int8_weight * 16 + fp8_alpha * 8 + pre_quant_scale * 4 + zero * 2 + bias - mPreQuantScaleInputIdx = (quant_algo & GroupwiseQuantAlgo::PRE_QUANT_SCALE) ? 1 : 0; - mWeightInputIdx = mPreQuantScaleInputIdx + 1; - mScalesInputIdx = mWeightInputIdx + 1; - mZerosInputIdx = (quant_algo & GroupwiseQuantAlgo::ZERO) ? mScalesInputIdx + 1 : mScalesInputIdx; - mBiasesInputIdx = (quant_algo & GroupwiseQuantAlgo::BIAS) ? mZerosInputIdx + 1 : mZerosInputIdx; - - if (mType == nvinfer1::DataType::kHALF) - { - // CUTLASS kernel selection - if (quant_algo & GroupwiseQuantAlgo::FP8_ALPHA) - { - mAlpha = alpha; - - // Ada & Hopper style kernels - if (mArch < 89) - { - TLLM_THROW("W4A(fp)8 kernel is unsupported on pre-Ada (sm<89) architectures!"); - } - assert(!(quant_algo & GroupwiseQuantAlgo::INT8_WEIGHT) && "W4A(fp)8 kernel requires INT4 weight!"); - m_weightOnlyGroupwiseGemmRunner - = selectGemmRunnerForZERO<__nv_fp8_e4m3, cutlass::uint4b_t, half>(quant_algo); - } - else - { - m_weightOnlyGroupwiseGemmRunner = selectGemmRunnerForWeightType(quant_algo); - } - // CUDA kernel selection - if (quant_algo & GroupwiseQuantAlgo::INT8_WEIGHT) - { - // INT8 weight - mCudaKernelEnabled = tensorrt_llm::kernels::weight_only::is_supported( - mArch, tensorrt_llm::kernels::weight_only::KernelType::FP16Int8Groupwise); - mCudaKernelType = tensorrt_llm::kernels::weight_only::KernelType::FP16Int8Groupwise; - } - else - { - // INT4 weight - mCudaKernelEnabled = tensorrt_llm::kernels::weight_only::is_supported( - mArch, tensorrt_llm::kernels::weight_only::KernelType::FP16Int4Groupwise); - mCudaKernelType = tensorrt_llm::kernels::weight_only::KernelType::FP16Int4Groupwise; - } - } -#if defined(ENABLE_BF16) - else if (mType == nvinfer1::DataType::kBF16) - { - // CUTLASS kernel selection - if (quant_algo & GroupwiseQuantAlgo::FP8_ALPHA) - { - mAlpha = alpha; - - // FP8 requires at least sm89 devices - if (mArch < 89) - { - TLLM_THROW("W4A(fp)8 kernel is unsupported on pre-Ada (sm<89) architectures!"); - } - assert(!(quant_algo & GroupwiseQuantAlgo::INT8_WEIGHT) && "W4A(fp)8 kernel requires INT4 weight!"); - m_weightOnlyGroupwiseGemmRunner - = selectGemmRunnerForZERO<__nv_fp8_e4m3, cutlass::uint4b_t, __nv_bfloat16, half>(quant_algo); - } - else - { - m_weightOnlyGroupwiseGemmRunner = selectGemmRunnerForWeightType<__nv_bfloat16>(quant_algo); - } - // CUDA kernel selection - if (quant_algo & GroupwiseQuantAlgo::INT8_WEIGHT) - { - // INT8 weight - mCudaKernelEnabled = tensorrt_llm::kernels::weight_only::is_supported( - mArch, tensorrt_llm::kernels::weight_only::KernelType::BF16Int8Groupwise); - mCudaKernelType = tensorrt_llm::kernels::weight_only::KernelType::BF16Int8Groupwise; - } - else - { - // INT4 weight - mCudaKernelEnabled = tensorrt_llm::kernels::weight_only::is_supported( - mArch, tensorrt_llm::kernels::weight_only::KernelType::BF16Int4Groupwise); - mCudaKernelType = tensorrt_llm::kernels::weight_only::KernelType::BF16Int4Groupwise; - } - } -#endif - else - { - TLLM_THROW("Unsupported data type"); - } - mPluginProfiler->setQuantAlgo(mQuantAlgo); - mPluginProfiler->setGroupSize(mGroupSize); - if (mCudaKernelEnabled) - { - mPluginProfiler->setCudaKernelType(mCudaKernelType, mArch); - } - mGemmId = GemmIdCore(mDims.n, mDims.k, mType); -} - -// IPluginV2DynamicExt Methods -nvinfer1::IPluginV2DynamicExt* WeightOnlyGroupwiseQuantMatmulPlugin::clone() const noexcept -{ - auto* plugin = new WeightOnlyGroupwiseQuantMatmulPlugin(*this); - return plugin; -} - -void WeightOnlyGroupwiseQuantMatmulPlugin::configGemm() -{ - mPluginProfiler->profileTactics(m_weightOnlyGroupwiseGemmRunner, mType, mDims, mGemmId, mCudaKernelEnabled); -} - -nvinfer1::DimsExprs WeightOnlyGroupwiseQuantMatmulPlugin::getOutputDimensions( - int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept -{ - - // inputs - // 0 activations [M, K] - // 1 pre-quant scales [K] (optional) - // 2 weights [K, N/2] - // 3 scales [K // group_size, N] - // 4 zeros [K // group_size, N] (optional) - // 5 biases [N] (optional) - - try - { - TLLM_CHECK(nbInputs == mBiasesInputIdx + 1); - TLLM_CHECK(outputIndex == 0); - int const nbDimsA = inputs[0].nbDims; - int const nbDimsB = inputs[mWeightInputIdx].nbDims; - TLLM_CHECK(nbDimsA >= 2); - TLLM_CHECK(nbDimsB == 2); - DimsExprs ret; - ret.nbDims = nbDimsA; - for (int ii = 0; ii < nbDimsA - 1; ++ii) - { - ret.d[ii] = inputs[0].d[ii]; - } - - // int4/int8 weight only quant (INT4*4 -> FP16, INT8*2 -> FP16) - int const weight_multiplier = mQuantAlgo & GroupwiseQuantAlgo::INT8_WEIGHT ? FP16_INT8_RATIO : FP16_INT4_RATIO; - ret.d[nbDimsA - 1] = exprBuilder.constant(inputs[mWeightInputIdx].d[1]->getConstantValue() * weight_multiplier); - - return ret; - } - catch (std::exception const& e) - { - caughtError(e); - } - return DimsExprs{}; -} - -bool WeightOnlyGroupwiseQuantMatmulPlugin::supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept -{ - if (pos < nbInputs + 1) - { - return inOut[pos].type == mType && inOut[pos].format == TensorFormat::kLINEAR; - } - else - { - // Never should be here - assert(false); - return false; - } -} - -void WeightOnlyGroupwiseQuantMatmulPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept -{ - auto const minM = std::accumulate(in[0].min.d, in[0].min.d + in[0].min.nbDims - 1, 1, std::multiplies()); - auto const maxM = std::accumulate(in[0].max.d, in[0].max.d + in[0].max.nbDims - 1, 1, std::multiplies()); - int const maxK = in[0].max.d[in[0].max.nbDims - 1]; - - // Quantized weights are packed in FP16 format (INT4*4 -> FP16, INT8*2 -> FP16) - int const weight_multiplier = mQuantAlgo & GroupwiseQuantAlgo::INT8_WEIGHT ? FP16_INT8_RATIO : FP16_INT4_RATIO; - int const maxN = in[mWeightInputIdx].max.d[1] * weight_multiplier; - - auto const K = maxK; - auto const N = maxN / weight_multiplier; - - if (!mDims.isInitialized()) - { - mDims = {minM, maxM, N, K}; - } - mGemmId = {N, K, mType}; - - size_t smoothedActSize = static_cast(maxM) * static_cast(maxK) - * (in[0].desc.type == nvinfer1::DataType::kFLOAT ? sizeof(float) : sizeof(half)); - m_workspaceMaxSize = smoothedActSize + m_weightOnlyGroupwiseGemmRunner->getWorkspaceSize(maxM, maxN, maxK); -} - -size_t WeightOnlyGroupwiseQuantMatmulPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - return m_workspaceMaxSize; -} - -template -void pre_quant_scale_for_act(int const m, int const k, int const mQuantAlgo, int const mPreQuantScaleInputIdx, - void const* const* inputs, void* workspace, cudaStream_t stream) -{ - // Apply pre-quant per channel scale on activations - if (mQuantAlgo & GroupwiseQuantAlgo::FP8_ALPHA) - { - tensorrt_llm::kernels::apply_per_channel_scale_kernel_launcher( - reinterpret_cast<__nv_fp8_e4m3*>(workspace), reinterpret_cast(inputs[0]), - reinterpret_cast(inputs[mPreQuantScaleInputIdx]), m, k, nullptr, stream); - } - else - { - tensorrt_llm::kernels::apply_per_channel_scale_kernel_launcher( - reinterpret_cast(workspace), reinterpret_cast(inputs[0]), - reinterpret_cast(inputs[mPreQuantScaleInputIdx]), m, k, nullptr, stream); - } -} - -int WeightOnlyGroupwiseQuantMatmulPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) noexcept -{ - // inputs - // 0 activations [M, K] - // 1 pre-quant scales [K] - // 2 weights [K, N/2] - // 3 scales [K // group_size, N] - // 4 zeros [K // group_size, N] - // 5 biases [N] - // outputs - // mat [M, N] - - int64_t m64 = 1; - for (int ii = 0; ii < inputDesc[0].dims.nbDims - 1; ++ii) - { - m64 *= inputDesc[0].dims.d[ii]; - } - int const m = TLLM_INT32_CAST(m64); - int const n = TLLM_INT32_CAST(inputDesc[mWeightInputIdx].dims.d[1]); - int const k = TLLM_INT32_CAST(inputDesc[0].dims.d[inputDesc[0].dims.nbDims - 1]); - - // get best tactic and check if CUDA kernel should be used - bool use_cuda_kernel = false; - auto const& bestTactic = mPluginProfiler->getBestConfig(m, mGemmId); - TLLM_CHECK_WITH_INFO(bestTactic, - "No valid weight only groupwise GEMM tactic(It is usually caused by the failure to execute all " - "candidate configurations of the CUTLASS kernel, please pay attention to the warning information " - "when building the engine.)"); - use_cuda_kernel = bestTactic->enableCudaKernel; - - bool use_pre_quant_scale = mQuantAlgo & GroupwiseQuantAlgo::PRE_QUANT_SCALE; - half const* zeros_ptr - = (mQuantAlgo & GroupwiseQuantAlgo::ZERO) ? reinterpret_cast(inputs[mZerosInputIdx]) : nullptr; - half const* biases_ptr - = (mQuantAlgo & GroupwiseQuantAlgo::BIAS) ? reinterpret_cast(inputs[mBiasesInputIdx]) : nullptr; - half const* act_ptr = reinterpret_cast(inputs[0]); - - if (use_pre_quant_scale && !use_cuda_kernel) - { - // Apply pre-quant per channel scale on activations - act_ptr = reinterpret_cast(workspace); - if (mType == nvinfer1::DataType::kHALF) - { - pre_quant_scale_for_act(m, k, mQuantAlgo, mPreQuantScaleInputIdx, inputs, workspace, stream); - } -#if defined(ENABLE_BF16) - else if (mType == nvinfer1::DataType::kBF16) - { - pre_quant_scale_for_act<__nv_bfloat16>(m, k, mQuantAlgo, mPreQuantScaleInputIdx, inputs, workspace, stream); - } -#endif - } - -#if defined(ENABLE_BF16) - TLLM_CHECK_WITH_INFO(mType == nvinfer1::DataType::kHALF || mType == nvinfer1::DataType::kBF16, - "No valid weightOnlyGropwiseQuantMatmul configuration"); -#else - TLLM_CHECK_WITH_INFO(mType == nvinfer1::DataType::kHALF, "No valid weightOnlyGropwiseQuantMatmul configuration"); -#endif - - // Quantized weights are packed in FP16 format (INT4*4 -> FP16, INT8*2 -> FP16) - int real_n = mQuantAlgo & GroupwiseQuantAlgo::INT8_WEIGHT ? n * FP16_INT8_RATIO : n * FP16_INT4_RATIO; - - if (use_cuda_kernel) - { - // Apply CUDA kernel - void const* pre_quant_scale_ptr = nullptr; - if (use_pre_quant_scale) - pre_quant_scale_ptr = inputs[mPreQuantScaleInputIdx]; - void const* cuda_kernel_act_ptr = inputs[0]; - void const* cuda_kernel_weight_ptr = inputs[mWeightInputIdx]; - void const* cuda_kernel_scales_ptr = inputs[mScalesInputIdx]; - void* cuda_kernel_out_ptr = outputs[0]; - tensorrt_llm::kernels::weight_only::Params params{cuda_kernel_act_ptr, pre_quant_scale_ptr, - cuda_kernel_weight_ptr, cuda_kernel_scales_ptr, zeros_ptr, biases_ptr, cuda_kernel_out_ptr, mAlpha, m, - real_n, k, mGroupSize, mCudaKernelType, static_cast(mQuantAlgo & GroupwiseQuantAlgo::FP8_ALPHA)}; - tensorrt_llm::kernels::weight_only::kernel_launcher(mArch, params, stream); - } - else - { - // Apply CUTLASS kernel - int const ws_bytes = m_weightOnlyGroupwiseGemmRunner->getWorkspaceSize(m, real_n, k); - int32_t* weight_ptr = const_cast(reinterpret_cast(inputs[mWeightInputIdx])); - m_weightOnlyGroupwiseGemmRunner->gemm(act_ptr, weight_ptr, inputs[mScalesInputIdx], zeros_ptr, biases_ptr, - mAlpha, outputs[0], m, real_n, k, mGroupSize, *bestTactic, - reinterpret_cast(workspace) + m * k * sizeof(half), ws_bytes, stream); - } - return 0; -} - -// IPluginV2Ext Methods -nvinfer1::DataType WeightOnlyGroupwiseQuantMatmulPlugin::getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept -{ - TLLM_CHECK(index == 0); - return mType; -} - -// IPluginV2 Methods - -char const* WeightOnlyGroupwiseQuantMatmulPlugin::getPluginType() const noexcept -{ - return WOQ_GROUPWISE_MATMUL_PLUGIN_NAME; -} - -char const* WeightOnlyGroupwiseQuantMatmulPlugin::getPluginVersion() const noexcept -{ - return WOQ_GROUPWISE_MATMUL_PLUGIN_VERSION; -} - -int WeightOnlyGroupwiseQuantMatmulPlugin::getNbOutputs() const noexcept -{ - return 1; -} - -int WeightOnlyGroupwiseQuantMatmulPlugin::initialize() noexcept -{ - configGemm(); - return 0; -} - -void WeightOnlyGroupwiseQuantMatmulPlugin::terminate() noexcept {} - -size_t WeightOnlyGroupwiseQuantMatmulPlugin::getSerializationSize() const noexcept -{ - return sizeof(nvinfer1::DataType) + // mType - sizeof(int) + // mQuantAlgo - sizeof(int) + // mGroupSize - sizeof(float) + // mAlpha - sizeof(mDims) + // Dimensions - mPluginProfiler->getSerializationSize(mGemmId); // selected tactics container size -} - -void WeightOnlyGroupwiseQuantMatmulPlugin::serialize(void* buffer) const noexcept -{ - char *d = static_cast(buffer), *a = d; - write(d, mType); - write(d, mQuantAlgo); - write(d, mGroupSize); - write(d, mAlpha); - write(d, mDims); - - mPluginProfiler->serialize(d, mGemmId); - TLLM_CHECK(d == a + getSerializationSize()); -} - -void WeightOnlyGroupwiseQuantMatmulPlugin::destroy() noexcept -{ - // This gets called when the network containing plugin is destroyed - delete this; -} - -/////////////// - -WeightOnlyGroupwiseQuantMatmulPluginCreator::WeightOnlyGroupwiseQuantMatmulPluginCreator() -{ - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - mPluginAttributes.emplace_back(PluginField("type_id", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("quant_algo", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("group_size", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("alpha", nullptr, PluginFieldType::kFLOAT32)); - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -char const* WeightOnlyGroupwiseQuantMatmulPluginCreator::getPluginName() const noexcept -{ - return WOQ_GROUPWISE_MATMUL_PLUGIN_NAME; -} - -char const* WeightOnlyGroupwiseQuantMatmulPluginCreator::getPluginVersion() const noexcept -{ - return WOQ_GROUPWISE_MATMUL_PLUGIN_VERSION; -} - -PluginFieldCollection const* WeightOnlyGroupwiseQuantMatmulPluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -IPluginV2* WeightOnlyGroupwiseQuantMatmulPluginCreator::createPlugin( - char const* name, PluginFieldCollection const* fc) noexcept -{ - PluginField const* fields = fc->fields; - nvinfer1::DataType type{}; - int QuantAlgo{}; - int GroupSize{}; - float Alpha{}; - // Read configurations from each fields - for (int i = 0; i < fc->nbFields; ++i) - { - char const* attrName = fields[i].name; - if (!strcmp(attrName, "quant_algo")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - QuantAlgo = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "group_size")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - GroupSize = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "type_id")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - type = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "alpha")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kFLOAT32); - Alpha = static_cast(*(static_cast(fields[i].data))); - } - } - try - { - // WeightOnlyGroupwiseQuantMatmulPluginCreator is unique and shared for an engine generation - // Create plugin profiler with shared tactics map - auto pluginProfiler = gemmPluginProfileManager.createGemmPluginProfiler(/* inference */ false); - auto* obj = new WeightOnlyGroupwiseQuantMatmulPlugin(type, QuantAlgo, GroupSize, Alpha, pluginProfiler); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -IPluginV2* WeightOnlyGroupwiseQuantMatmulPluginCreator::deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept -{ - // This object will be deleted when the network is destroyed, which will - // call weightOnlyGroupwiseQuantMatmulPlugin::destroy() - try - { - // Create plugin profiler with private tactics map which is read from the serialized engine - auto pluginProfiler = gemmPluginProfileManager.createGemmPluginProfiler(/* inference */ true); - auto* obj = new WeightOnlyGroupwiseQuantMatmulPlugin(serialData, serialLength, pluginProfiler); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} diff --git a/cpp/tensorrt_llm/plugins/weightOnlyGroupwiseQuantMatmulPlugin/weightOnlyGroupwiseQuantMatmulPlugin.h b/cpp/tensorrt_llm/plugins/weightOnlyGroupwiseQuantMatmulPlugin/weightOnlyGroupwiseQuantMatmulPlugin.h deleted file mode 100644 index 94e98ce0f5c0..000000000000 --- a/cpp/tensorrt_llm/plugins/weightOnlyGroupwiseQuantMatmulPlugin/weightOnlyGroupwiseQuantMatmulPlugin.h +++ /dev/null @@ -1,186 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. - */ -#pragma once - -#include "tensorrt_llm/common/quantization.h" -#include "tensorrt_llm/kernels/cutlass_kernels/fpA_intB_gemm/fpA_intB_gemm.h" -#include "tensorrt_llm/kernels/preQuantScaleKernel.h" -#include "tensorrt_llm/kernels/weightOnlyBatchedGemv//kernelLauncher.h" -#include "tensorrt_llm/plugins/common/gemmPluginProfiler.h" -#include "tensorrt_llm/plugins/common/plugin.h" -#include "tensorrt_llm/plugins/weightOnlyQuantMatmulPlugin/weightOnlyQuantMatmulPlugin.h" - -#include - -#include -#include -#include -#include -#include -#include - -// The blank line here is to avoid clang-format -sort-includes option reordering these two cutlass header files and -// breaking dependencies -#include "cutlass/integer_subbyte.h" - -namespace tensorrt_llm::plugins -{ - -using WeightOnlyGemmRunner = tensorrt_llm::kernels::cutlass_kernels::CutlassFpAIntBGemmRunnerInterface; -using WeightOnlyGemmRunnerPtr = std::shared_ptr; -using KernelType = tensorrt_llm::kernels::weight_only::KernelType; - -class WeightOnlyGroupwiseQuantGemmPluginProfiler - : public GemmPluginProfiler -{ -public: - using Config = tensorrt_llm::cutlass_extensions::CutlassGemmConfig; - - void setQuantAlgo(int quantAlgo) - { - mQuantAlgo = quantAlgo; - } - - void setGroupSize(int groupSize) - { - mGroupSize = groupSize; - } - - void setCudaKernelType(KernelType cudaKernelType, int arch) - { - mCudaKernelType = cudaKernelType; - mArch = arch; - } - -protected: - void runTactic(int m, int n, int k, Config const& tactic, char* workspace, cudaStream_t const& stream) override; - - void computeTmpSize(size_t maxM, size_t n, size_t k) override; - - std::vector getTactics(int m, int n, int k) const override; - - bool checkTactic(int m, int n, int k, Config const& tactic) const override; - -private: - int mQuantAlgo; - int mGroupSize; - KernelType mCudaKernelType; - int mArch; -}; - -class WeightOnlyGroupwiseQuantMatmulPlugin : public BasePlugin -{ -public: - using PluginProfilerPtr = std::shared_ptr; - - WeightOnlyGroupwiseQuantMatmulPlugin() = delete; - - WeightOnlyGroupwiseQuantMatmulPlugin( - nvinfer1::DataType type, int quant_algo, int group_size, float alpha, PluginProfilerPtr const& profiler); - - WeightOnlyGroupwiseQuantMatmulPlugin(void const* data, size_t length, PluginProfilerPtr const& profiler); - - ~WeightOnlyGroupwiseQuantMatmulPlugin() override = default; - - // IPluginV2DynamicExt Methods - nvinfer1::IPluginV2DynamicExt* clone() const noexcept override; - nvinfer1::DimsExprs getOutputDimensions(int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; - bool supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept override; - void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept override; - size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept override; - int enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - - // IPluginV2Ext Methods - nvinfer1::DataType getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept override; - - // IPluginV2 Methods - char const* getPluginType() const noexcept override; - char const* getPluginVersion() const noexcept override; - int getNbOutputs() const noexcept override; - int initialize() noexcept override; - void terminate() noexcept override; - size_t getSerializationSize() const noexcept override; - void serialize(void* buffer) const noexcept override; - void destroy() noexcept override; - -private: - // group_size: 64, 128 - void init(nvinfer1::DataType type, int quant_algo, int group_size, float alpha); - - void configGemm(); - -private: - const std::string mLayerName; - - WeightOnlyGemmRunnerPtr m_weightOnlyGroupwiseGemmRunner; - size_t m_workspaceMaxSize; - nvinfer1::DataType mType; - bool mCudaKernelEnabled; - tensorrt_llm::kernels::weight_only::KernelType mCudaKernelType; - int mArch; - - // When M is smaller than this value, we trigger a fast path - // I.e. a tailored kernel instead of cutlass. - - int mQuantAlgo; - - int mGroupSize; - - float mAlpha = 1.0f; - - int mPreQuantScaleInputIdx; - int mWeightInputIdx; - int mScalesInputIdx; - int mZerosInputIdx; - int mBiasesInputIdx; - - GemmDims mDims{}; - GemmIdCore mGemmId{}; - - PluginProfilerPtr mPluginProfiler; -}; - -class WeightOnlyGroupwiseQuantMatmulPluginCreator : public BaseCreator -{ -public: - WeightOnlyGroupwiseQuantMatmulPluginCreator(); - - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; - - nvinfer1::IPluginV2* deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept override; - -private: - GemmPluginProfilerManager gemmPluginProfileManager; - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; -}; - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/weightOnlyQuantMatmulPlugin/CMakeLists.txt b/cpp/tensorrt_llm/plugins/weightOnlyQuantMatmulPlugin/CMakeLists.txt deleted file mode 100755 index 86876224fccd..000000000000 --- a/cpp/tensorrt_llm/plugins/weightOnlyQuantMatmulPlugin/CMakeLists.txt +++ /dev/null @@ -1,21 +0,0 @@ -# -# SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. -# -file(GLOB SRCS *.cpp) -set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS}) -set(PLUGIN_SOURCES - ${PLUGIN_SOURCES} - PARENT_SCOPE) diff --git a/cpp/tensorrt_llm/plugins/weightOnlyQuantMatmulPlugin/weightOnlyQuantMatmulPlugin.cpp b/cpp/tensorrt_llm/plugins/weightOnlyQuantMatmulPlugin/weightOnlyQuantMatmulPlugin.cpp deleted file mode 100644 index f3ed07fafaff..000000000000 --- a/cpp/tensorrt_llm/plugins/weightOnlyQuantMatmulPlugin/weightOnlyQuantMatmulPlugin.cpp +++ /dev/null @@ -1,507 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. - */ -#include "weightOnlyQuantMatmulPlugin.h" - -#include - -using namespace nvinfer1; -using namespace tensorrt_llm::common; -using namespace tensorrt_llm::kernels::cutlass_kernels; -using tensorrt_llm::plugins::WeightOnlyQuantMatmulPluginCreator; -using tensorrt_llm::plugins::WeightOnlyQuantMatmulPlugin; -using tensorrt_llm::plugins::WeightOnlyQuantGemmPluginProfiler; -using tensorrt_llm::plugins::read; -using tensorrt_llm::plugins::write; - -static char const* WOQ_MATMUL_PLUGIN_VERSION{"1"}; -static char const* WOQ_MATMUL_PLUGIN_NAME{"WeightOnlyQuantMatmul"}; -PluginFieldCollection WeightOnlyQuantMatmulPluginCreator::mFC{}; -std::vector WeightOnlyQuantMatmulPluginCreator::mPluginAttributes; - -void WeightOnlyQuantGemmPluginProfiler::runTactic(int m, int n, int k, - WeightOnlyQuantGemmPluginProfiler::Config const& tactic, char* workspace, cudaStream_t const& stream) -{ - int const originalN = n * getWeightTypeMultiplier(mWeightTypeId); - half* actPtr = reinterpret_cast(workspace); - int8_t* weightPtr - = reinterpret_cast(nextWorkspacePtr(reinterpret_cast(actPtr), m * k * sizeof(half))); - half* scalesPtr - = reinterpret_cast(nextWorkspacePtr(reinterpret_cast(weightPtr), n * k * sizeof(int8_t))); - half* outputPtr - = reinterpret_cast(nextWorkspacePtr(reinterpret_cast(scalesPtr), originalN * sizeof(half))); - char* workspacePtr - = reinterpret_cast(nextWorkspacePtr(reinterpret_cast(outputPtr), m * originalN * sizeof(half))); - - int const wsSize = mRunner->getWorkspaceSize(m, originalN, k); - - if (tactic.enableCudaKernel) - { - // run CUDA kernel - tensorrt_llm::kernels::weight_only::Params params{actPtr, nullptr, weightPtr, scalesPtr, nullptr, nullptr, - outputPtr, 1.f, m, originalN, k, 0, mCudaKernelType}; - tensorrt_llm::kernels::weight_only::kernel_launcher(mArch, params, stream); - } - else - { - // run CUTLASS kernel - if (mWeightTypeId == WeightTypeId::INT8) - { - mRunner->gemm( - actPtr, weightPtr, scalesPtr, outputPtr, m, originalN, k, tactic, workspacePtr, wsSize, stream); - } - else - { - mRunner->gemm(actPtr, reinterpret_cast(weightPtr), scalesPtr, outputPtr, m, originalN, - k, tactic, workspacePtr, wsSize, stream); - } - } -} - -void WeightOnlyQuantGemmPluginProfiler::computeTmpSize(size_t maxM, size_t n, size_t k) -{ - int const originalN = n * getWeightTypeMultiplier(mWeightTypeId); - std::vector workspaces = { - maxM * k * sizeof(half), // A - n * k * sizeof(int8_t), // B - originalN * sizeof(half), // scales - maxM * originalN * sizeof(half), // C - mRunner->getWorkspaceSize(maxM, originalN, k) // workspace - }; - size_t bytes = calculateTotalWorkspaceSize(workspaces.data(), workspaces.size()); - setTmpWorkspaceSizeInBytes(bytes); -} - -std::vector WeightOnlyQuantGemmPluginProfiler::getTactics( - int m, int n, int k) const -{ - return mRunner->getConfigs(); -} - -bool WeightOnlyQuantGemmPluginProfiler::checkTactic(int m, int n, int k, Config const& tactic) const -{ - // stop to profile Cuda kernel for m >= 16 - if (tactic.enableCudaKernel) - { - return m < 16; - } - return true; -} - -WeightOnlyQuantMatmulPlugin::WeightOnlyQuantMatmulPlugin(nvinfer1::DataType type, WeightTypeId weightTypeId, - WeightOnlyQuantMatmulPlugin::PluginProfilerPtr const& pluginProfiler) - : mPluginProfiler(pluginProfiler) -{ - init(type, weightTypeId); -} - -// Parameterized constructor -WeightOnlyQuantMatmulPlugin::WeightOnlyQuantMatmulPlugin( - void const* data, size_t length, WeightOnlyQuantMatmulPlugin::PluginProfilerPtr const& pluginProfiler) - : mPluginProfiler(pluginProfiler) -{ - char const *d = reinterpret_cast(data), *a = d; - nvinfer1::DataType type; - WeightTypeId weightTypeId; - read(d, type); - read(d, weightTypeId); - read(d, mDims); - - init(type, weightTypeId); - - mPluginProfiler->deserialize(d, mDims, mGemmId); - - TLLM_CHECK_WITH_INFO(d == a + length, - "Expected length (%d) != real length (%d). This is often " - "caused by using different TensorRT LLM version to build " - "engine and run engine.", - (int) length, (int) (d - a)); -} - -void WeightOnlyQuantMatmulPlugin::init(nvinfer1::DataType type, WeightTypeId weightTypeId) -{ - mArch = tensorrt_llm::common::getSMVersion(); - mType = type; - mWeightTypeId = weightTypeId; - - if (mWeightTypeId == WeightTypeId::INT8) - { - if (mType == nvinfer1::DataType::kHALF) - { - m_weightOnlyGemmRunner = std::make_shared< - CutlassFpAIntBGemmRunner>(); - mCudaKernelEnabled = tensorrt_llm::kernels::weight_only::is_supported( - mArch, tensorrt_llm::kernels::weight_only::KernelType::FP16Int8PerChannel); - mCudaKernelType = tensorrt_llm::kernels::weight_only::KernelType::FP16Int8PerChannel; - } -#if defined(ENABLE_BF16) - else if (mType == nvinfer1::DataType::kBF16) - { - m_weightOnlyGemmRunner = std::make_shared< - CutlassFpAIntBGemmRunner<__nv_bfloat16, uint8_t, cutlass::WeightOnlyQuantOp::PER_COLUMN_SCALE_ONLY>>(); - mCudaKernelEnabled = tensorrt_llm::kernels::weight_only::is_supported( - mArch, tensorrt_llm::kernels::weight_only::KernelType::BF16Int8PerChannel); - mCudaKernelType = tensorrt_llm::kernels::weight_only::KernelType::BF16Int8PerChannel; - } -#endif - else - { - TLLM_CHECK(false); - } - } - else if (mWeightTypeId == WeightTypeId::INT4) - { - if (mType == nvinfer1::DataType::kHALF) - { - m_weightOnlyGemmRunner = std::make_shared< - CutlassFpAIntBGemmRunner>(); - mCudaKernelEnabled = tensorrt_llm::kernels::weight_only::is_supported( - mArch, tensorrt_llm::kernels::weight_only::KernelType::FP16Int4PerChannel); - mCudaKernelType = tensorrt_llm::kernels::weight_only::KernelType::FP16Int4PerChannel; - } -#if defined(ENABLE_BF16) - else if (mType == nvinfer1::DataType::kBF16) - { - m_weightOnlyGemmRunner = std::make_shared>(); - mCudaKernelEnabled = tensorrt_llm::kernels::weight_only::is_supported( - mArch, tensorrt_llm::kernels::weight_only::KernelType::BF16Int4PerChannel); - mCudaKernelType = tensorrt_llm::kernels::weight_only::KernelType::BF16Int4PerChannel; - } -#endif - else - { - TLLM_CHECK(false); - } - } - else - { - TLLM_CHECK(false); - } - - mPluginProfiler->setWeightTypeId(mWeightTypeId); - if (mCudaKernelEnabled) - { - mPluginProfiler->setCudaKernelType(mCudaKernelType, mArch); - } - mGemmId = GemmIdCore(mDims.n, mDims.k, mType); -} - -// IPluginV2DynamicExt Methods -nvinfer1::IPluginV2DynamicExt* WeightOnlyQuantMatmulPlugin::clone() const noexcept -{ - auto* plugin = new WeightOnlyQuantMatmulPlugin(*this); - return plugin; -} - -void WeightOnlyQuantMatmulPlugin::configGemm() -{ - mPluginProfiler->profileTactics(m_weightOnlyGemmRunner, mType, mDims, mGemmId, mCudaKernelEnabled); -} - -nvinfer1::DimsExprs WeightOnlyQuantMatmulPlugin::getOutputDimensions( - int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept -{ - // input [m1, m2, m3, ... , k] - // weight [k, n] for int8, [k, n/2] for int4 - - try - { - TLLM_CHECK(nbInputs == 3); - TLLM_CHECK(outputIndex == 0); - int const nbDimsA = inputs[0].nbDims; - int const nbDimsB = inputs[1].nbDims; - TLLM_CHECK(nbDimsA >= 2); - TLLM_CHECK(nbDimsB == 2); - DimsExprs ret; - ret.nbDims = nbDimsA; - for (int ii = 0; ii < nbDimsA - 1; ++ii) - { - ret.d[ii] = inputs[0].d[ii]; - } - if (mWeightTypeId == WeightTypeId::INT8) - { - // int8 weight only quant - ret.d[nbDimsA - 1] = exprBuilder.constant(inputs[1].d[1]->getConstantValue()); - } - else - { - // int4 weight only quant - ret.d[nbDimsA - 1] = exprBuilder.constant(inputs[1].d[1]->getConstantValue() * INT8_INT4_RATIO); - } - return ret; - } - catch (std::exception const& e) - { - caughtError(e); - } - return DimsExprs{}; -} - -bool WeightOnlyQuantMatmulPlugin::supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept -{ - switch (pos) - { - case 0: - // activation - return inOut[0].type == mType && inOut[0].format == TensorFormat::kLINEAR; - case 1: - // weights - // Weights are required to be int8, but will be reinterpreted as int4 in enqueue if required - // Weights stored in checkpoint should have int8/int4 type - return inOut[1].type == nvinfer1::DataType::kINT8 && inOut[1].format == TensorFormat::kLINEAR; - case 2: - // scales channels - return inOut[2].type == mType && inOut[2].format == TensorFormat::kLINEAR; - case 3: - // out - return inOut[3].type == mType && inOut[3].format == TensorFormat::kLINEAR; - default: - // Never should be here - assert(false); - return false; - } -} - -void WeightOnlyQuantMatmulPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept -{ - auto const minM = std::accumulate(in[0].min.d, in[0].min.d + in[0].min.nbDims - 1, 1, std::multiplies()); - auto const maxM = std::accumulate(in[0].max.d, in[0].max.d + in[0].max.nbDims - 1, 1, std::multiplies()); - - int const maxK = in[0].max.d[in[0].max.nbDims - 1]; - int const maxN = in[1].max.d[1] * getWeightTypeMultiplier(mWeightTypeId); - - auto const K = maxK; - auto const N = maxN / getWeightTypeMultiplier(mWeightTypeId); - - if (!mDims.isInitialized()) - { - mDims = {minM, maxM, N, K}; - } - - mGemmId = {N, K, mType}; - - m_workspaceMaxSize = m_weightOnlyGemmRunner->getWorkspaceSize(maxM, maxN, maxK); -} - -size_t WeightOnlyQuantMatmulPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - return m_workspaceMaxSize; -} - -int WeightOnlyQuantMatmulPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) noexcept -{ - // inputs - // mat1 [M1, M2,..., K] - // mat2 [K, N] for int8, [K, N/2] for int4 - // scale_channels [N] - // outputs - // mat [M, N] - - int64_t m64 = 1; - for (int ii = 0; ii < inputDesc[0].dims.nbDims - 1; ++ii) - { - m64 *= inputDesc[0].dims.d[ii]; - } - int const m = TLLM_INT32_CAST(m64); - int const n = TLLM_INT32_CAST(inputDesc[1].dims.d[1]); - int const k = TLLM_INT32_CAST(inputDesc[0].dims.d[inputDesc[0].dims.nbDims - 1]); - - if (m == 0) - return 0; - -#if defined(ENABLE_BF16) - TLLM_CHECK_WITH_INFO(mType == nvinfer1::DataType::kHALF || mType == nvinfer1::DataType::kBF16, - "No valid weightOnlyQuantMatmul configuration"); -#else - TLLM_CHECK_WITH_INFO(mType == nvinfer1::DataType::kHALF, "No valid weightOnlyQuantMatmul configuration"); -#endif - int real_n = mWeightTypeId == WeightTypeId::INT4 ? n * INT8_INT4_RATIO : n; - - // get best tactic and check if CUDA kernel should be used - bool use_cuda_kernel = false; - auto const& bestTactic = mPluginProfiler->getBestConfig(m, mGemmId); - TLLM_CHECK_WITH_INFO(bestTactic, - "No valid weight only per-channel GEMM tactic(It is usually caused by the failure to execute all candidate " - "configurations of the CUTLASS kernel, please pay attention to the warning information when building the " - "engine.)"); - use_cuda_kernel = bestTactic->enableCudaKernel; - if (use_cuda_kernel) - { - void const* cuda_kernel_act_ptr = inputs[0]; - void const* cuda_kernel_weight_ptr = inputs[1]; - void const* cuda_kernel_scales_ptr = inputs[2]; - void* cuda_kernel_out_ptr = outputs[0]; - tensorrt_llm::kernels::weight_only::Params params(cuda_kernel_act_ptr, nullptr, cuda_kernel_weight_ptr, - cuda_kernel_scales_ptr, nullptr, nullptr, cuda_kernel_out_ptr, 1.f, m, real_n, k, 0, mCudaKernelType); - tensorrt_llm::kernels::weight_only::kernel_launcher(mArch, params, stream); - } - else - { - int const ws_size = m_weightOnlyGemmRunner->getWorkspaceSize(m, real_n, k); - - m_weightOnlyGemmRunner->gemm(inputs[0], inputs[1], inputs[2], outputs[0], m, real_n, k, *bestTactic, - reinterpret_cast(workspace), ws_size, stream); - } - - return 0; -} - -// IPluginV2Ext Methods -nvinfer1::DataType WeightOnlyQuantMatmulPlugin::getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept -{ - TLLM_CHECK(index == 0); - return mType; -} - -// IPluginV2 Methods - -char const* WeightOnlyQuantMatmulPlugin::getPluginType() const noexcept -{ - return WOQ_MATMUL_PLUGIN_NAME; -} - -char const* WeightOnlyQuantMatmulPlugin::getPluginVersion() const noexcept -{ - return WOQ_MATMUL_PLUGIN_VERSION; -} - -int WeightOnlyQuantMatmulPlugin::getNbOutputs() const noexcept -{ - return 1; -} - -int WeightOnlyQuantMatmulPlugin::initialize() noexcept -{ - configGemm(); - return 0; -} - -void WeightOnlyQuantMatmulPlugin::terminate() noexcept {} - -size_t WeightOnlyQuantMatmulPlugin::getSerializationSize() const noexcept -{ - return sizeof(mWeightTypeId) + // mWeightTypeId - sizeof(nvinfer1::DataType) + // mType - sizeof(mDims) + // Dimensions - mPluginProfiler->getSerializationSize(mGemmId); // selected tactics container size -} - -void WeightOnlyQuantMatmulPlugin::serialize(void* buffer) const noexcept -{ - char *d = static_cast(buffer), *a = d; - write(d, mType); - write(d, mWeightTypeId); - write(d, mDims); - - mPluginProfiler->serialize(d, mGemmId); - TLLM_CHECK(d == a + getSerializationSize()); -} - -void WeightOnlyQuantMatmulPlugin::destroy() noexcept -{ - // This gets called when the network containing plugin is destroyed - delete this; -} - -/////////////// - -WeightOnlyQuantMatmulPluginCreator::WeightOnlyQuantMatmulPluginCreator() -{ - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - mPluginAttributes.emplace_back(PluginField("type_id", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("weight_type_id", nullptr, PluginFieldType::kINT32)); - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -char const* WeightOnlyQuantMatmulPluginCreator::getPluginName() const noexcept -{ - return WOQ_MATMUL_PLUGIN_NAME; -} - -char const* WeightOnlyQuantMatmulPluginCreator::getPluginVersion() const noexcept -{ - return WOQ_MATMUL_PLUGIN_VERSION; -} - -PluginFieldCollection const* WeightOnlyQuantMatmulPluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -IPluginV2* WeightOnlyQuantMatmulPluginCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept -{ - PluginField const* fields = fc->fields; - nvinfer1::DataType type{}; - WeightTypeId weightTypeId{}; - // Read configurations from each fields - for (int i = 0; i < fc->nbFields; ++i) - { - char const* attrName = fields[i].name; - if (!strcmp(attrName, "weight_type_id")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - weightTypeId = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "type_id")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - type = static_cast(*(static_cast(fields[i].data))); - } - } - try - { - // WeightOnlyGroupwiseQuantMatmulPluginCreator is unique and shared for an engine generation - // Create plugin profiler with shared tactics map - auto pluginProfiler = gemmPluginProfileManager.createGemmPluginProfiler(/* inference */ false); - auto* obj = new WeightOnlyQuantMatmulPlugin(type, weightTypeId, pluginProfiler); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -IPluginV2* WeightOnlyQuantMatmulPluginCreator::deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept -{ - // This object will be deleted when the network is destroyed, which will - // call WeightOnlyQuantMatmulPlugin::destroy() - try - { - // Create plugin profiler with private tactics map which is read from the serialized engine - auto pluginProfiler = gemmPluginProfileManager.createGemmPluginProfiler(/* inference */ true); - auto* obj = new WeightOnlyQuantMatmulPlugin(serialData, serialLength, pluginProfiler); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} diff --git a/cpp/tensorrt_llm/plugins/weightOnlyQuantMatmulPlugin/weightOnlyQuantMatmulPlugin.h b/cpp/tensorrt_llm/plugins/weightOnlyQuantMatmulPlugin/weightOnlyQuantMatmulPlugin.h deleted file mode 100644 index 3177d8297d2d..000000000000 --- a/cpp/tensorrt_llm/plugins/weightOnlyQuantMatmulPlugin/weightOnlyQuantMatmulPlugin.h +++ /dev/null @@ -1,175 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. - */ -#pragma once - -#include "tensorrt_llm/common/quantization.h" -#include "tensorrt_llm/kernels/cutlass_kernels/fpA_intB_gemm/fpA_intB_gemm.h" -#include "tensorrt_llm/kernels/weightOnlyBatchedGemv/kernelLauncher.h" -#include "tensorrt_llm/plugins/common/gemmPluginProfiler.h" -#include "tensorrt_llm/plugins/common/plugin.h" - -#include -#include -#include -#include -#include -#include - -// The blank line here is to avoid clang-format -sort-includes option reordering these two cutlass header files and -// breaking dependencies -#include "cutlass/integer_subbyte.h" - -namespace tensorrt_llm::plugins -{ -enum class WeightTypeId -{ - INT8 = 1, - INT4 = 2, -}; - -constexpr int32_t FP16_BITS = 16; -constexpr int32_t INT8_BITS = 8; -constexpr int32_t INT4_BITS = 4; -constexpr int32_t INT8_INT4_RATIO = INT8_BITS / INT4_BITS; -constexpr int32_t FP16_INT4_RATIO = FP16_BITS / INT4_BITS; -constexpr int32_t FP16_INT8_RATIO = FP16_BITS / INT8_BITS; - -inline int32_t getWeightTypeMultiplier(WeightTypeId weightTypeId) -{ - return weightTypeId == WeightTypeId::INT8 ? 1 : INT8_INT4_RATIO; -} - -using WeightOnlyGemmRunner = tensorrt_llm::kernels::cutlass_kernels::CutlassFpAIntBGemmRunnerInterface; -using WeightOnlyGemmRunnerPtr = std::shared_ptr; -using KernelType = tensorrt_llm::kernels::weight_only::KernelType; - -class WeightOnlyQuantGemmPluginProfiler : public GemmPluginProfiler -{ -public: - using Config = tensorrt_llm::cutlass_extensions::CutlassGemmConfig; - - void setWeightTypeId(WeightTypeId weightId) - { - mWeightTypeId = weightId; - } - - void setCudaKernelType(KernelType cudaKernelType, int arch) - { - mCudaKernelType = cudaKernelType; - mArch = arch; - } - -protected: - void runTactic(int m, int n, int k, Config const& tactic, char* workspace, cudaStream_t const& stream) override; - - void computeTmpSize(size_t maxM, size_t n, size_t k) override; - - std::vector getTactics(int m, int n, int k) const override; - - bool checkTactic(int m, int n, int k, Config const& tactic) const override; - -private: - WeightTypeId mWeightTypeId; - KernelType mCudaKernelType; - int mArch; -}; - -class WeightOnlyQuantMatmulPlugin : public BasePlugin -{ -public: - using PluginProfilerPtr = std::shared_ptr; - WeightOnlyQuantMatmulPlugin() = delete; - - WeightOnlyQuantMatmulPlugin(nvinfer1::DataType type, WeightTypeId weightTypeId, PluginProfilerPtr const& profiler); - - WeightOnlyQuantMatmulPlugin(void const* data, size_t length, PluginProfilerPtr const& profiler); - - ~WeightOnlyQuantMatmulPlugin() override = default; - - // IPluginV2DynamicExt Methods - nvinfer1::IPluginV2DynamicExt* clone() const noexcept override; - nvinfer1::DimsExprs getOutputDimensions(int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; - bool supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept override; - void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept override; - size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept override; - int enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - - // IPluginV2Ext Methods - nvinfer1::DataType getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept override; - - // IPluginV2 Methods - char const* getPluginType() const noexcept override; - char const* getPluginVersion() const noexcept override; - int getNbOutputs() const noexcept override; - int initialize() noexcept override; - void terminate() noexcept override; - size_t getSerializationSize() const noexcept override; - void serialize(void* buffer) const noexcept override; - void destroy() noexcept override; - -private: - void init(nvinfer1::DataType type, WeightTypeId weightTypeId); - - void configGemm(); - -private: - const std::string mLayerName; - - WeightOnlyGemmRunnerPtr m_weightOnlyGemmRunner; - size_t m_workspaceMaxSize; - nvinfer1::DataType mType; - WeightTypeId mWeightTypeId; - bool mCudaKernelEnabled; - tensorrt_llm::kernels::weight_only::KernelType mCudaKernelType; - int mArch; - - GemmDims mDims{}; - GemmIdCore mGemmId{}; - - PluginProfilerPtr mPluginProfiler; -}; - -class WeightOnlyQuantMatmulPluginCreator : public BaseCreator -{ -public: - WeightOnlyQuantMatmulPluginCreator(); - - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; - - nvinfer1::IPluginV2* deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept override; - -private: - GemmPluginProfilerManager gemmPluginProfileManager; - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; -}; - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/runtime/CMakeLists.txt b/cpp/tensorrt_llm/runtime/CMakeLists.txt index ca81fbb0f6cd..11a9391c0e69 100644 --- a/cpp/tensorrt_llm/runtime/CMakeLists.txt +++ b/cpp/tensorrt_llm/runtime/CMakeLists.txt @@ -26,7 +26,6 @@ set(SRCS eagleBuffers.cpp explicitDraftTokensBuffers.cpp lookaheadBuffers.cpp - layerProfiler.cpp loraManager.cpp loraUtils.cpp loraModule.cpp @@ -51,9 +50,6 @@ set(SRCS promptTuningParams.cpp runtimeKernels.cu tllmBuffers.cpp - tllmRuntime.cpp - tllmStreamReaders.cpp - tllmLogger.cpp workerPool.cpp worldConfig.cpp virtualMemory.cpp) diff --git a/cpp/tensorrt_llm/runtime/bufferManager.cpp b/cpp/tensorrt_llm/runtime/bufferManager.cpp index 3de42a253158..58257516a387 100644 --- a/cpp/tensorrt_llm/runtime/bufferManager.cpp +++ b/cpp/tensorrt_llm/runtime/bufferManager.cpp @@ -18,6 +18,7 @@ #include "tensorrt_llm/common/assert.h" #include "tensorrt_llm/common/cudaUtils.h" #include "tensorrt_llm/common/memoryUtils.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tllmBuffers.h" #include @@ -37,7 +38,7 @@ BufferManager::BufferManager(CudaStreamPtr stream, bool trimPool) mPool = CudaMemPool::getPrimaryPoolForDevice(mStream->getDevice()); } -BufferManager::IBufferPtr BufferManager::gpu(std::size_t size, nvinfer1::DataType type) const +BufferManager::IBufferPtr BufferManager::gpu(std::size_t size, tensorrt_llm::DataType type) const { if (auto vmAllocator = getVirtualMemoryAllocator()) { @@ -51,7 +52,7 @@ BufferManager::IBufferPtr BufferManager::gpu(std::size_t size, nvinfer1::DataTyp return gpuSync(size, type); } -BufferManager::ITensorPtr BufferManager::gpu(nvinfer1::Dims dims, nvinfer1::DataType type) const +BufferManager::ITensorPtr BufferManager::gpu(tensorrt_llm::Dims dims, tensorrt_llm::DataType type) const { if (auto vmAllocator = getVirtualMemoryAllocator()) { @@ -65,7 +66,7 @@ BufferManager::ITensorPtr BufferManager::gpu(nvinfer1::Dims dims, nvinfer1::Data return gpuSync(dims, type); } -BufferManager::IBufferPtr BufferManager::gpuSync(std::size_t size, nvinfer1::DataType type) +BufferManager::IBufferPtr BufferManager::gpuSync(std::size_t size, tensorrt_llm::DataType type) { if (auto vmAllocator = getVirtualMemoryAllocator()) { @@ -74,7 +75,7 @@ BufferManager::IBufferPtr BufferManager::gpuSync(std::size_t size, nvinfer1::Dat return std::make_unique(size, type, CudaAllocator{}); } -BufferManager::ITensorPtr BufferManager::gpuSync(nvinfer1::Dims dims, nvinfer1::DataType type) +BufferManager::ITensorPtr BufferManager::gpuSync(tensorrt_llm::Dims dims, tensorrt_llm::DataType type) { if (auto vmAllocator = getVirtualMemoryAllocator()) { @@ -83,47 +84,48 @@ BufferManager::ITensorPtr BufferManager::gpuSync(nvinfer1::Dims dims, nvinfer1:: return std::make_unique(dims, type, CudaAllocator{}); } -BufferManager::IBufferPtr BufferManager::cpu(std::size_t size, nvinfer1::DataType type) +BufferManager::IBufferPtr BufferManager::cpu(std::size_t size, tensorrt_llm::DataType type) { return std::make_unique(size, type); } -BufferManager::ITensorPtr BufferManager::cpu(nvinfer1::Dims dims, nvinfer1::DataType type) +BufferManager::ITensorPtr BufferManager::cpu(tensorrt_llm::Dims dims, tensorrt_llm::DataType type) { return std::make_unique(dims, type); } -BufferManager::IBufferPtr BufferManager::pinned(std::size_t size, nvinfer1::DataType type) +BufferManager::IBufferPtr BufferManager::pinned(std::size_t size, tensorrt_llm::DataType type) { return std::make_unique(size, type); } -BufferManager::ITensorPtr BufferManager::pinned(nvinfer1::Dims dims, nvinfer1::DataType type) +BufferManager::ITensorPtr BufferManager::pinned(tensorrt_llm::Dims dims, tensorrt_llm::DataType type) { return std::make_unique(dims, type); } -BufferManager::IBufferPtr BufferManager::pinnedPool(std::size_t size, nvinfer1::DataType type) +BufferManager::IBufferPtr BufferManager::pinnedPool(std::size_t size, tensorrt_llm::DataType type) { return std::make_unique(size, type); } -BufferManager::ITensorPtr BufferManager::pinnedPool(nvinfer1::Dims dims, nvinfer1::DataType type) +BufferManager::ITensorPtr BufferManager::pinnedPool(tensorrt_llm::Dims dims, tensorrt_llm::DataType type) { return std::make_unique(dims, type); } -BufferManager::IBufferPtr BufferManager::managed(std::size_t size, nvinfer1::DataType type) +BufferManager::IBufferPtr BufferManager::managed(std::size_t size, tensorrt_llm::DataType type) { return std::make_unique(size, type); } -BufferManager::ITensorPtr BufferManager::managed(nvinfer1::Dims dims, nvinfer1::DataType type) +BufferManager::ITensorPtr BufferManager::managed(tensorrt_llm::Dims dims, tensorrt_llm::DataType type) { return std::make_unique(dims, type); } -BufferManager::ITensorPtr BufferManager::ipcNvls(std::set ranks, nvinfer1::Dims dims, nvinfer1::DataType type) +BufferManager::ITensorPtr BufferManager::ipcNvls( + std::set ranks, tensorrt_llm::Dims dims, tensorrt_llm::DataType type) { return std::make_unique(dims, type, ranks); } @@ -187,7 +189,7 @@ void BufferManager::copy(IBuffer const& src, IBuffer& dst) const } BufferManager::IBufferPtr BufferManager::allocate( - MemoryType memoryType, std::size_t size, nvinfer1::DataType type) const + MemoryType memoryType, std::size_t size, tensorrt_llm::DataType type) const { switch (memoryType) { @@ -202,7 +204,7 @@ BufferManager::IBufferPtr BufferManager::allocate( } BufferManager::ITensorPtr BufferManager::allocate( - MemoryType memoryType, nvinfer1::Dims dims, nvinfer1::DataType type) const + MemoryType memoryType, tensorrt_llm::Dims dims, tensorrt_llm::DataType type) const { switch (memoryType) { diff --git a/cpp/tensorrt_llm/runtime/bufferView.h b/cpp/tensorrt_llm/runtime/bufferView.h index 236b89d7d455..a001d05f1f8b 100644 --- a/cpp/tensorrt_llm/runtime/bufferView.h +++ b/cpp/tensorrt_llm/runtime/bufferView.h @@ -17,6 +17,7 @@ #pragma once #include "tensorrt_llm/common/assert.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/runtime/iBuffer.h" #include @@ -70,7 +71,7 @@ class BufferView : virtual public IBuffer return mBuffer->getCapacity() - mOffset; } - [[nodiscard]] nvinfer1::DataType getDataType() const override + [[nodiscard]] tensorrt_llm::DataType getDataType() const override { return mBuffer->getDataType(); } diff --git a/cpp/tensorrt_llm/runtime/decoderState.cpp b/cpp/tensorrt_llm/runtime/decoderState.cpp index b5851dc1c2d2..83037b2431cb 100644 --- a/cpp/tensorrt_llm/runtime/decoderState.cpp +++ b/cpp/tensorrt_llm/runtime/decoderState.cpp @@ -16,6 +16,7 @@ #include "tensorrt_llm/runtime/decoderState.h" #include "tensorrt_llm/batch_manager/llmRequest.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/kernels/decodingCommon.h" #include "tensorrt_llm/runtime/runtimeKernels.h" @@ -27,10 +28,10 @@ using TensorPtr = DecoderState::TensorPtr; BeamSearchBuffers::BeamSearchBuffers(BufferManager const& bufferManager) : mOutputBeamHypotheses{} - , mCumLogProbsTmp(bufferManager.emptyTensor(MemoryType::kGPU, nvinfer1::DataType::kFLOAT)) + , mCumLogProbsTmp(bufferManager.emptyTensor(MemoryType::kGPU, tensorrt_llm::DataType::kFLOAT)) { mOutputBeamHypotheses.empty(bufferManager); - mCumLogProbsTmp = bufferManager.emptyTensor(MemoryType::kGPU, nvinfer1::DataType::kFLOAT); + mCumLogProbsTmp = bufferManager.emptyTensor(MemoryType::kGPU, tensorrt_llm::DataType::kFLOAT); int device; cudaGetDevice(&device); @@ -54,8 +55,8 @@ DecoderState::DecoderState() } void DecoderState::setup(SizeType32 maxNumSequences, SizeType32 maxBeamWidth, SizeType32 maxAttentionWindow, - SizeType32 sinkTokenLength, SizeType32 maxSequenceLength, nvinfer1::DataType dtype, ModelConfig const& modelConfig, - WorldConfig const& worldConfig, BufferManager const& bufferManager) + SizeType32 sinkTokenLength, SizeType32 maxSequenceLength, tensorrt_llm::DataType dtype, + ModelConfig const& modelConfig, WorldConfig const& worldConfig, BufferManager const& bufferManager) { TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); setupBuffers(dtype, bufferManager); @@ -64,7 +65,7 @@ void DecoderState::setup(SizeType32 maxNumSequences, SizeType32 maxBeamWidth, Si TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); } -void DecoderState::setupBuffers(nvinfer1::DataType dtype, BufferManager const& bufferManager) +void DecoderState::setupBuffers(tensorrt_llm::DataType dtype, BufferManager const& bufferManager) { TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); auto constexpr nvTokenIdType = TRTDataType::value; @@ -114,7 +115,7 @@ void DecoderState::setupBuffers(nvinfer1::DataType dtype, BufferManager const& b } void DecoderState::setupSpeculativeDecoding(SpeculativeDecodingMode const& speculativeDecodingMode, - SizeType32 maxTokensPerEngineStep, nvinfer1::DataType dtype, ModelConfig const& modelConfig, + SizeType32 maxTokensPerEngineStep, tensorrt_llm::DataType dtype, ModelConfig const& modelConfig, WorldConfig const& worldConfig, BufferManager const& bufferManager) { TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); @@ -124,8 +125,8 @@ void DecoderState::setupSpeculativeDecoding(SpeculativeDecodingMode const& specu TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); } -void DecoderState::setupSpeculativeDecodingBuffers( - SpeculativeDecodingMode const speculativeDecodingMode, nvinfer1::DataType dtype, BufferManager const& bufferManager) +void DecoderState::setupSpeculativeDecodingBuffers(SpeculativeDecodingMode const speculativeDecodingMode, + tensorrt_llm::DataType dtype, BufferManager const& bufferManager) { TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); @@ -151,13 +152,13 @@ void DecoderState::setupSpeculativeDecodingBuffers( if (speculativeDecodingMode.predictsDraftTokens()) { speculativeDecodingOutputs.nextDraftTokens - = bufferManager.emptyTensor(MemoryType::kGPU, nvinfer1::DataType::kINT32); + = bufferManager.emptyTensor(MemoryType::kGPU, tensorrt_llm::DataType::kINT32); if (speculativeDecodingMode.variableDraftLength()) { speculativeDecodingOutputs.nextDraftTokensLen - = bufferManager.emptyTensor(MemoryType::kGPU, nvinfer1::DataType::kINT32); + = bufferManager.emptyTensor(MemoryType::kGPU, tensorrt_llm::DataType::kINT32); speculativeDecodingOutputs.prevDraftTokensLen - = bufferManager.emptyTensor(MemoryType::kGPU, nvinfer1::DataType::kINT32); + = bufferManager.emptyTensor(MemoryType::kGPU, tensorrt_llm::DataType::kINT32); } } if (speculativeDecodingMode.isLookaheadDecoding()) @@ -167,11 +168,11 @@ void DecoderState::setupSpeculativeDecodingBuffers( if (speculativeDecodingMode.needsKVCacheRewind()) { speculativeDecodingOutputs.acceptedTokensLen - = bufferManager.emptyTensor(MemoryType::kGPU, nvinfer1::DataType::kINT32); + = bufferManager.emptyTensor(MemoryType::kGPU, tensorrt_llm::DataType::kINT32); speculativeDecodingOutputs.acceptedLengthsCumSum - = bufferManager.emptyTensor(MemoryType::kGPU, nvinfer1::DataType::kINT32); + = bufferManager.emptyTensor(MemoryType::kGPU, tensorrt_llm::DataType::kINT32); speculativeDecodingOutputs.pathsOffsets - = bufferManager.emptyTensor(MemoryType::kGPU, nvinfer1::DataType::kINT32); + = bufferManager.emptyTensor(MemoryType::kGPU, tensorrt_llm::DataType::kINT32); } dOutput->speculativeDecodingOutputs = speculativeDecodingOutputs; diff --git a/cpp/tensorrt_llm/runtime/decodingLayerWorkspace.cpp b/cpp/tensorrt_llm/runtime/decodingLayerWorkspace.cpp index f8a4fa4e7467..c5098bf777e0 100644 --- a/cpp/tensorrt_llm/runtime/decodingLayerWorkspace.cpp +++ b/cpp/tensorrt_llm/runtime/decodingLayerWorkspace.cpp @@ -15,11 +15,12 @@ */ #include "tensorrt_llm/runtime/decodingLayerWorkspace.h" +#include "tensorrt_llm/common/tllmDataType.h" #include tensorrt_llm::runtime::DecodingLayerWorkspace::DecodingLayerWorkspace(std::shared_ptr bufferManager, - tensorrt_llm::layers::DecoderDomain const& decoderDomain, nvinfer1::DataType logitsType, + tensorrt_llm::layers::DecoderDomain const& decoderDomain, tensorrt_llm::DataType logitsType, size_t workspaceBufferSizeInBytes) : mBufferManager(std::move(bufferManager)) , mBatchSlotsDevice( @@ -82,7 +83,8 @@ void tensorrt_llm::runtime::DecodingLayerWorkspace::resize(size_t minSize) } tensorrt_llm::runtime::DecodingLayerWorkspace::TensorPtr -tensorrt_llm::runtime::DecodingLayerWorkspace::getWorkspaceAsDeviceTensor(ITensor::Shape shape, nvinfer1::DataType type) +tensorrt_llm::runtime::DecodingLayerWorkspace::getWorkspaceAsDeviceTensor( + ITensor::Shape shape, tensorrt_llm::DataType type) { auto const sizeInBytes = ITensor::volume(shape) * BufferDataType(type).getSize(); return std::make_shared>>( diff --git a/cpp/tensorrt_llm/runtime/decodingLayerWorkspace.h b/cpp/tensorrt_llm/runtime/decodingLayerWorkspace.h index c2688b51139f..68d3d54124f5 100644 --- a/cpp/tensorrt_llm/runtime/decodingLayerWorkspace.h +++ b/cpp/tensorrt_llm/runtime/decodingLayerWorkspace.h @@ -19,6 +19,7 @@ #include #include "tensorrt_llm/common/dataType.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/common/workspace.h" #include "tensorrt_llm/layers/decodingParams.h" #include "tensorrt_llm/runtime/bufferManager.h" @@ -39,7 +40,7 @@ class DecodingLayerWorkspace using BufferPtr = IBuffer::SharedPtr; DecodingLayerWorkspace(std::shared_ptr bufferManager, layers::DecoderDomain const& decoderDomain, - nvinfer1::DataType logitsType, size_t workspaceBufferSizeInBytes); + tensorrt_llm::DataType logitsType, size_t workspaceBufferSizeInBytes); DecodingLayerWorkspace() = delete; @@ -71,7 +72,7 @@ class DecodingLayerWorkspace [[nodiscard]] TensorPtr getDeviceRuntimeLogits() const; ///@brief Gets a tensor with the given shape and type at the start of the device workspace. - TensorPtr getWorkspaceAsDeviceTensor(ITensor::Shape shape, nvinfer1::DataType type); + TensorPtr getWorkspaceAsDeviceTensor(ITensor::Shape shape, tensorrt_llm::DataType type); /// @brief A convenience function to copy the content of a standard vector to a device workspace. template @@ -112,7 +113,7 @@ class DecodingLayerWorkspace { size_t lastTensorOffset = 0; auto alignedSizeCalculator - = [&lastTensorOffset](std::pair const& tensorDescriptor) + = [&lastTensorOffset](std::pair const& tensorDescriptor) { auto const& [shape, type] = tensorDescriptor; auto const sizeInBytes = ITensor::volume(shape) * tensorrt_llm::common::getDTypeSize(type); diff --git a/cpp/tensorrt_llm/runtime/eagleBuffers.cpp b/cpp/tensorrt_llm/runtime/eagleBuffers.cpp index 097fd95f49aa..e0f2198c3e58 100644 --- a/cpp/tensorrt_llm/runtime/eagleBuffers.cpp +++ b/cpp/tensorrt_llm/runtime/eagleBuffers.cpp @@ -19,6 +19,7 @@ #include "tensorrt_llm/common/assert.h" #include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/kernels/speculativeDecoding/eagleDecodingKernels.h" #include "tensorrt_llm/kernels/speculativeDecoding/explicitDraftTokensKernels.h" #include "tensorrt_llm/runtime/common.h" @@ -41,50 +42,51 @@ void EagleBuffers::Inputs::create(SizeType32 maxNumSequences, BufferManager cons auto const numEagleLayers = speculativeDecodingModule.getMaxDraftPathLen(); auto constexpr TRTTokenIdType = runtime::TRTDataType::value; - temperatures = manager.gpu(ITensor::makeShape({maxNumSequences}), nvinfer1::DataType::kFLOAT); - randomDataSample = manager.gpu(ITensor::makeShape({maxNumSequences}), nvinfer1::DataType::kFLOAT); + temperatures = manager.gpu(ITensor::makeShape({maxNumSequences}), tensorrt_llm::DataType::kFLOAT); + randomDataSample = manager.gpu(ITensor::makeShape({maxNumSequences}), tensorrt_llm::DataType::kFLOAT); randomDataValidation - = manager.gpu(ITensor::makeShape({maxNumSequences, maxDecodingTokens}), nvinfer1::DataType::kFLOAT); + = manager.gpu(ITensor::makeShape({maxNumSequences, maxDecodingTokens}), tensorrt_llm::DataType::kFLOAT); draftTokens = manager.gpu(ITensor::makeShape({maxNumSequences, maxDecodingDraftTokens}), TRTTokenIdType); - draftLens = manager.gpu(ITensor::makeShape({maxNumSequences}), nvinfer1::DataType::kINT32); + draftLens = manager.gpu(ITensor::makeShape({maxNumSequences}), tensorrt_llm::DataType::kINT32); draftPaths - = manager.gpu(ITensor::makeShape({maxNumSequences, maxNumPaths, maxPathLen}), nvinfer1::DataType::kINT32); + = manager.gpu(ITensor::makeShape({maxNumSequences, maxNumPaths, maxPathLen}), tensorrt_llm::DataType::kINT32); draftPathsHost = BufferManager::pinnedPool( - ITensor::makeShape({maxNumSequences, maxNumPaths, maxPathLen}), nvinfer1::DataType::kINT32); - specDecodingGenerationLengths = manager.gpu(ITensor::makeShape({maxNumSequences}), nvinfer1::DataType::kINT32); + ITensor::makeShape({maxNumSequences, maxNumPaths, maxPathLen}), tensorrt_llm::DataType::kINT32); + specDecodingGenerationLengths = manager.gpu(ITensor::makeShape({maxNumSequences}), tensorrt_llm::DataType::kINT32); specDecodingGenerationLengthsHost - = BufferManager::pinnedPool(ITensor::makeShape({maxNumSequences}), nvinfer1::DataType::kINT32); + = BufferManager::pinnedPool(ITensor::makeShape({maxNumSequences}), tensorrt_llm::DataType::kINT32); specDecodingPackedMasks = manager.gpu(ITensor::makeShape({maxNumSequences, maxDecodingTokens, common::ceilDiv(maxDecodingTokens, 32)}), - nvinfer1::DataType::kINT32); + tensorrt_llm::DataType::kINT32); specDecodingPositionOffsets - = manager.gpu(ITensor::makeShape({maxNumSequences * maxDecodingTokens}), nvinfer1::DataType::kINT32); + = manager.gpu(ITensor::makeShape({maxNumSequences * maxDecodingTokens}), tensorrt_llm::DataType::kINT32); eagleNetCtxRequestTypesHost - = BufferManager::pinnedPool(ITensor::makeShape({maxNumSequences}), nvinfer1::DataType::kINT32); + = BufferManager::pinnedPool(ITensor::makeShape({maxNumSequences}), tensorrt_llm::DataType::kINT32); eagleNetCtxContextLengthsHost - = BufferManager::pinnedPool(ITensor::makeShape({maxNumSequences}), nvinfer1::DataType::kINT32); + = BufferManager::pinnedPool(ITensor::makeShape({maxNumSequences}), tensorrt_llm::DataType::kINT32); eagleNetCtxPastKeyValueLengthsHost - = BufferManager::pinnedPool(ITensor::makeShape({maxNumSequences}), nvinfer1::DataType::kINT32); + = BufferManager::pinnedPool(ITensor::makeShape({maxNumSequences}), tensorrt_llm::DataType::kINT32); eagleNetGenRequestTypesHost - = BufferManager::pinnedPool(ITensor::makeShape({maxNumSequences}), nvinfer1::DataType::kINT32); + = BufferManager::pinnedPool(ITensor::makeShape({maxNumSequences}), tensorrt_llm::DataType::kINT32); eagleNetGenContextLengthsHost - = BufferManager::pinnedPool(ITensor::makeShape({maxNumSequences}), nvinfer1::DataType::kINT32); + = BufferManager::pinnedPool(ITensor::makeShape({maxNumSequences}), tensorrt_llm::DataType::kINT32); eagleNetGenPastKeyValueLengthsHost - = BufferManager::pinnedPool(ITensor::makeShape({maxNumSequences}), nvinfer1::DataType::kINT32); + = BufferManager::pinnedPool(ITensor::makeShape({maxNumSequences}), tensorrt_llm::DataType::kINT32); inputGenTokensHost = BufferManager::pinnedPool( - ITensor::makeShape({maxNumSequences * maxDecodingTokens}), nvinfer1::DataType::kINT32); - chunkedContextNextTokens = manager.gpu(ITensor::makeShape({maxNumSequences}), nvinfer1::DataType::kINT32); - useSpecDecoding = manager.cpu(ITensor::makeShape({1}), nvinfer1::DataType::kINT32); + ITensor::makeShape({maxNumSequences * maxDecodingTokens}), tensorrt_llm::DataType::kINT32); + chunkedContextNextTokens = manager.gpu(ITensor::makeShape({maxNumSequences}), tensorrt_llm::DataType::kINT32); + useSpecDecoding = manager.cpu(ITensor::makeShape({1}), tensorrt_llm::DataType::kINT32); // Eagle-2 - useDynamicTreeHost = BufferManager::pinnedPool(ITensor::makeShape({1}), nvinfer1::DataType::kINT32); - dynamicTreeMaxTopKHost = BufferManager::pinnedPool(ITensor::makeShape({1}), nvinfer1::DataType::kINT32); - prevScores = manager.gpu(ITensor::makeShape({maxNumSequences, maxDecodingDraftTokens}), nvinfer1::DataType::kFLOAT); + useDynamicTreeHost = BufferManager::pinnedPool(ITensor::makeShape({1}), tensorrt_llm::DataType::kINT32); + dynamicTreeMaxTopKHost = BufferManager::pinnedPool(ITensor::makeShape({1}), tensorrt_llm::DataType::kINT32); + prevScores + = manager.gpu(ITensor::makeShape({maxNumSequences, maxDecodingDraftTokens}), tensorrt_llm::DataType::kFLOAT); currentExpandIndices = manager.gpu(ITensor::makeShape({maxNumSequences, maxDecodingDraftTokens}), TRTTokenIdType); allLayersScores = manager.gpu( ITensor::makeShape({maxNumSequences, numEagleLayers, maxDecodingDraftTokens * maxDecodingDraftTokens}), - nvinfer1::DataType::kFLOAT); + tensorrt_llm::DataType::kFLOAT); allLayersDraftTokenIds = manager.gpu( ITensor::makeShape({maxNumSequences, numEagleLayers, maxDecodingDraftTokens * maxDecodingDraftTokens}), TRTTokenIdType); @@ -114,58 +116,63 @@ EagleBuffers::EagleBuffers(SizeType32 maxBatchSize, SizeType32 maxBeamWidth, run auto constexpr TRTTokenIdType = runtime::TRTDataType::value; // input tensors - engineInputs.temperatures = manager.emptyTensor(runtime::MemoryType::kGPU, nvinfer1::DataType::kFLOAT); - engineInputs.posteriorAlpha = manager.emptyTensor(runtime::MemoryType::kGPU, nvinfer1::DataType::kFLOAT); - engineInputs.posteriorThreshold = manager.emptyTensor(runtime::MemoryType::kGPU, nvinfer1::DataType::kFLOAT); - posteriorAlphaHost = manager.emptyTensor(runtime::MemoryType::kPINNEDPOOL, nvinfer1::DataType::kFLOAT); - posteriorThresholdHost = manager.emptyTensor(runtime::MemoryType::kPINNEDPOOL, nvinfer1::DataType::kFLOAT); - greedySamplingHost = BufferManager::pinnedPool(ITensor::makeShape({1}), nvinfer1::DataType::kINT32); + engineInputs.temperatures = manager.emptyTensor(runtime::MemoryType::kGPU, tensorrt_llm::DataType::kFLOAT); + engineInputs.posteriorAlpha = manager.emptyTensor(runtime::MemoryType::kGPU, tensorrt_llm::DataType::kFLOAT); + engineInputs.posteriorThreshold = manager.emptyTensor(runtime::MemoryType::kGPU, tensorrt_llm::DataType::kFLOAT); + posteriorAlphaHost = manager.emptyTensor(runtime::MemoryType::kPINNEDPOOL, tensorrt_llm::DataType::kFLOAT); + posteriorThresholdHost = manager.emptyTensor(runtime::MemoryType::kPINNEDPOOL, tensorrt_llm::DataType::kFLOAT); + greedySamplingHost = BufferManager::pinnedPool(ITensor::makeShape({1}), tensorrt_llm::DataType::kINT32); engineInputs.draftTokens = manager.gpu(ITensor::makeShape({maxNumSequences, maxDecodingDraftTokens}), TRTTokenIdType); - engineInputs.draftLens = manager.gpu(ITensor::makeShape({maxNumSequences}), nvinfer1::DataType::kINT32); + engineInputs.draftLens = manager.gpu(ITensor::makeShape({maxNumSequences}), tensorrt_llm::DataType::kINT32); engineInputs.draftPaths - = manager.gpu(ITensor::makeShape({maxNumSequences, numPaths, pathLen}), nvinfer1::DataType::kINT32); + = manager.gpu(ITensor::makeShape({maxNumSequences, numPaths, pathLen}), tensorrt_llm::DataType::kINT32); engineInputs.specDecodingGenerationLengths - = manager.emptyTensor(runtime::MemoryType::kGPU, nvinfer1::DataType::kINT32); + = manager.emptyTensor(runtime::MemoryType::kGPU, tensorrt_llm::DataType::kINT32); engineInputs.specDecodingPositionOffsets - = manager.emptyTensor(runtime::MemoryType::kGPU, nvinfer1::DataType::kINT32); - engineInputs.specDecodingPackedMasks = manager.emptyTensor(runtime::MemoryType::kGPU, nvinfer1::DataType::kINT32); + = manager.emptyTensor(runtime::MemoryType::kGPU, tensorrt_llm::DataType::kINT32); + engineInputs.specDecodingPackedMasks + = manager.emptyTensor(runtime::MemoryType::kGPU, tensorrt_llm::DataType::kINT32); - engineInputs.randomDataSample = manager.emptyTensor(runtime::MemoryType::kGPU, nvinfer1::DataType::kFLOAT); - engineInputs.randomDataValidation = manager.emptyTensor(runtime::MemoryType::kGPU, nvinfer1::DataType::kFLOAT); + engineInputs.randomDataSample = manager.emptyTensor(runtime::MemoryType::kGPU, tensorrt_llm::DataType::kFLOAT); + engineInputs.randomDataValidation = manager.emptyTensor(runtime::MemoryType::kGPU, tensorrt_llm::DataType::kFLOAT); engineInputs.eagleNetCtxRequestTypesHost - = manager.emptyTensor(runtime::MemoryType::kPINNEDPOOL, nvinfer1::DataType::kINT32); + = manager.emptyTensor(runtime::MemoryType::kPINNEDPOOL, tensorrt_llm::DataType::kINT32); engineInputs.eagleNetCtxContextLengthsHost - = manager.emptyTensor(runtime::MemoryType::kPINNEDPOOL, nvinfer1::DataType::kINT32); + = manager.emptyTensor(runtime::MemoryType::kPINNEDPOOL, tensorrt_llm::DataType::kINT32); engineInputs.eagleNetCtxPastKeyValueLengthsHost - = manager.emptyTensor(runtime::MemoryType::kPINNEDPOOL, nvinfer1::DataType::kINT32); + = manager.emptyTensor(runtime::MemoryType::kPINNEDPOOL, tensorrt_llm::DataType::kINT32); engineInputs.eagleNetGenRequestTypesHost - = manager.emptyTensor(runtime::MemoryType::kPINNEDPOOL, nvinfer1::DataType::kINT32); + = manager.emptyTensor(runtime::MemoryType::kPINNEDPOOL, tensorrt_llm::DataType::kINT32); engineInputs.eagleNetGenContextLengthsHost - = manager.emptyTensor(runtime::MemoryType::kPINNEDPOOL, nvinfer1::DataType::kINT32); + = manager.emptyTensor(runtime::MemoryType::kPINNEDPOOL, tensorrt_llm::DataType::kINT32); engineInputs.eagleNetGenPastKeyValueLengthsHost - = manager.emptyTensor(runtime::MemoryType::kPINNEDPOOL, nvinfer1::DataType::kINT32); - engineInputs.inputGenTokensHost = manager.emptyTensor(runtime::MemoryType::kPINNEDPOOL, nvinfer1::DataType::kINT32); - engineInputs.chunkedContextNextTokens = manager.emptyTensor(runtime::MemoryType::kGPU, nvinfer1::DataType::kINT32); - engineInputs.useSpecDecoding = BufferManager::cpu(ITensor::makeShape({1}), nvinfer1::DataType::kINT32); + = manager.emptyTensor(runtime::MemoryType::kPINNEDPOOL, tensorrt_llm::DataType::kINT32); + engineInputs.inputGenTokensHost + = manager.emptyTensor(runtime::MemoryType::kPINNEDPOOL, tensorrt_llm::DataType::kINT32); + engineInputs.chunkedContextNextTokens + = manager.emptyTensor(runtime::MemoryType::kGPU, tensorrt_llm::DataType::kINT32); + engineInputs.useSpecDecoding = BufferManager::cpu(ITensor::makeShape({1}), tensorrt_llm::DataType::kINT32); bufferCast(*engineInputs.useSpecDecoding)[0] = 1; - chunkedContextNextTokensHost = manager.emptyTensor(runtime::MemoryType::kPINNEDPOOL, nvinfer1::DataType::kINT32); + chunkedContextNextTokensHost + = manager.emptyTensor(runtime::MemoryType::kPINNEDPOOL, tensorrt_llm::DataType::kINT32); // Eagle-2 - engineInputs.useDynamicTreeHost = BufferManager::pinnedPool(ITensor::makeShape({1}), nvinfer1::DataType::kINT32); + engineInputs.useDynamicTreeHost + = BufferManager::pinnedPool(ITensor::makeShape({1}), tensorrt_llm::DataType::kINT32); engineInputs.dynamicTreeMaxTopKHost - = BufferManager::pinnedPool(ITensor::makeShape({1}), nvinfer1::DataType::kINT32); + = BufferManager::pinnedPool(ITensor::makeShape({1}), tensorrt_llm::DataType::kINT32); engineInputs.prevScores - = manager.gpu(ITensor::makeShape({maxNumSequences, maxDecodingDraftTokens}), nvinfer1::DataType::kFLOAT); + = manager.gpu(ITensor::makeShape({maxNumSequences, maxDecodingDraftTokens}), tensorrt_llm::DataType::kFLOAT); engineInputs.currentExpandIndices = manager.gpu(ITensor::makeShape({maxNumSequences, maxDecodingDraftTokens}), TRTTokenIdType); engineInputs.allLayersScores = manager.gpu( ITensor::makeShape({maxNumSequences, numEagleLayers, maxDecodingDraftTokens * maxDecodingDraftTokens}), - nvinfer1::DataType::kFLOAT); + tensorrt_llm::DataType::kFLOAT); engineInputs.allLayersDraftTokenIds = manager.gpu( ITensor::makeShape({maxNumSequences, numEagleLayers, maxDecodingDraftTokens * maxDecodingDraftTokens}), TRTTokenIdType); @@ -176,24 +183,24 @@ EagleBuffers::EagleBuffers(SizeType32 maxBatchSize, SizeType32 maxBeamWidth, run // output tensors engineOutputs.nextDraftTokens = manager.gpu(ITensor::makeShape({maxNumSequences, numPaths, pathLen}), TRTTokenIdType); - engineOutputs.nextDraftLens = manager.gpu(ITensor::makeShape({maxNumSequences}), nvinfer1::DataType::kINT32); + engineOutputs.nextDraftLens = manager.gpu(ITensor::makeShape({maxNumSequences}), tensorrt_llm::DataType::kINT32); engineOutputs.nextDraftPaths - = manager.gpu(ITensor::makeShape({maxNumSequences, numPaths, pathLen}), nvinfer1::DataType::kINT32); + = manager.gpu(ITensor::makeShape({maxNumSequences, numPaths, pathLen}), tensorrt_llm::DataType::kINT32); engineOutputs.acceptedTokens - = manager.gpu(ITensor::makeShape({maxNumSequences, pathLen}), nvinfer1::DataType::kINT32); - engineOutputs.acceptedLens = manager.gpu(ITensor::makeShape({maxNumSequences}), nvinfer1::DataType::kINT32); - engineOutputs.acceptedPaths = manager.gpu(ITensor::makeShape({maxNumSequences}), nvinfer1::DataType::kINT32); + = manager.gpu(ITensor::makeShape({maxNumSequences, pathLen}), tensorrt_llm::DataType::kINT32); + engineOutputs.acceptedLens = manager.gpu(ITensor::makeShape({maxNumSequences}), tensorrt_llm::DataType::kINT32); + engineOutputs.acceptedPaths = manager.gpu(ITensor::makeShape({maxNumSequences}), tensorrt_llm::DataType::kINT32); engineOutputs.chunkedContextNextTokens - = manager.gpu(ITensor::makeShape({maxNumSequences}), nvinfer1::DataType::kINT32); + = manager.gpu(ITensor::makeShape({maxNumSequences}), tensorrt_llm::DataType::kINT32); // helper tensors scanReduceTempStorageBytes = tksd::invokeScanReduceGenerationLengths( maxNumSequences, nullptr, nullptr, 0, nullptr, nullptr, manager.getStream().get()); scanReduceTempStorage = manager.gpu(scanReduceTempStorageBytes); - cumSumGenerationLengths = manager.emptyTensor(runtime::MemoryType::kGPU, nvinfer1::DataType::kINT32); - maxGenerationLength = manager.gpu(ITensor::makeShape({1}), nvinfer1::DataType::kINT32); + cumSumGenerationLengths = manager.emptyTensor(runtime::MemoryType::kGPU, tensorrt_llm::DataType::kINT32); + maxGenerationLength = manager.gpu(ITensor::makeShape({1}), tensorrt_llm::DataType::kINT32); // pre-allocate empty tensors reshape(0, maxNumSequences, modelConfig); @@ -520,15 +527,15 @@ void EagleBuffers::setFromInputs(RequestVector const& contextRequests, RequestVe switch (dtype) { - case nvinfer1::DataType::kFLOAT: + case tensorrt_llm::DataType::kFLOAT: setFromInputs( contextRequests, genRequests, vocabSizePadded, seqSlots, draftBuffers, *eagleModule, manager); break; - case nvinfer1::DataType::kHALF: + case tensorrt_llm::DataType::kHALF: setFromInputs( contextRequests, genRequests, vocabSizePadded, seqSlots, draftBuffers, *eagleModule, manager); break; - case nvinfer1::DataType::kBF16: + case tensorrt_llm::DataType::kBF16: setFromInputs<__nv_bfloat16>( contextRequests, genRequests, vocabSizePadded, seqSlots, draftBuffers, *eagleModule, manager); break; diff --git a/cpp/tensorrt_llm/runtime/explicitDraftTokensBuffers.cpp b/cpp/tensorrt_llm/runtime/explicitDraftTokensBuffers.cpp index ed205ca0e117..89c74e6f9349 100644 --- a/cpp/tensorrt_llm/runtime/explicitDraftTokensBuffers.cpp +++ b/cpp/tensorrt_llm/runtime/explicitDraftTokensBuffers.cpp @@ -18,6 +18,7 @@ #include "tensorrt_llm/common/assert.h" #include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/kernels/speculativeDecoding/explicitDraftTokensKernels.h" #include "tensorrt_llm/runtime/common.h" #include "tensorrt_llm/runtime/iBuffer.h" @@ -40,23 +41,24 @@ void ExplicitDraftTokensBuffers::Inputs::create(SizeType32 maxNumSequences, Buff auto constexpr TRTTokenIdType = runtime::TRTDataType::value; auto const dtype = modelConfig.getDataType(); - maxGenLengthHost = manager.pinned(ITensor::makeShape({1}), nvinfer1::DataType::kINT32); + maxGenLengthHost = manager.pinned(ITensor::makeShape({1}), tensorrt_llm::DataType::kINT32); temperatures = manager.gpu(ITensor::makeShape({maxNumSequences}), dtype); - positionIdsBase = manager.gpu(ITensor::makeShape({maxNumSequences}), nvinfer1::DataType::kINT32); - generationLengths = manager.gpu(ITensor::makeShape({maxNumSequences}), nvinfer1::DataType::kINT32); - generationLengthsHost = manager.pinned(ITensor::makeShape({maxNumSequences}), nvinfer1::DataType::kINT32); + positionIdsBase = manager.gpu(ITensor::makeShape({maxNumSequences}), tensorrt_llm::DataType::kINT32); + generationLengths = manager.gpu(ITensor::makeShape({maxNumSequences}), tensorrt_llm::DataType::kINT32); + generationLengthsHost = manager.pinned(ITensor::makeShape({maxNumSequences}), tensorrt_llm::DataType::kINT32); randomDataSample = manager.gpu(ITensor::makeShape({maxNumSequences}), dtype); randomDataValidation = manager.gpu(ITensor::makeShape({maxNumSequences, maxNumPaths, maxDraftPathLen}), dtype); draftTokens = manager.gpu(ITensor::makeShape({maxNumSequences, maxNumPaths, maxPathLen}), TRTTokenIdType); draftIndices - = manager.gpu(ITensor::makeShape({maxNumSequences, maxNumPaths, maxPathLen}), nvinfer1::DataType::kINT32); + = manager.gpu(ITensor::makeShape({maxNumSequences, maxNumPaths, maxPathLen}), tensorrt_llm::DataType::kINT32); draftProbs = manager.gpu(ITensor::makeShape({maxNumSequences, maxNumPaths, maxDraftPathLen, vocabSizePadded}), dtype); packedMasks = manager.gpu(ITensor::makeShape({maxNumSequences, maxDecodingTokens, common::ceilDiv(maxDecodingTokens, 32)}), - nvinfer1::DataType::kINT32); - positionIds = manager.gpu(ITensor::makeShape({maxNumSequences * maxDecodingTokens}), nvinfer1::DataType::kINT32); - useSpecDecoding = manager.cpu(ITensor::makeShape({1}), nvinfer1::DataType::kINT32); + tensorrt_llm::DataType::kINT32); + positionIds + = manager.gpu(ITensor::makeShape({maxNumSequences * maxDecodingTokens}), tensorrt_llm::DataType::kINT32); + useSpecDecoding = manager.cpu(ITensor::makeShape({1}), tensorrt_llm::DataType::kINT32); } ExplicitDraftTokensBuffers::ExplicitDraftTokensBuffers(SizeType32 maxBatchSize, SizeType32 maxBeamWidth, @@ -81,52 +83,53 @@ ExplicitDraftTokensBuffers::ExplicitDraftTokensBuffers(SizeType32 maxBatchSize, auto const dtype = modelConfig.getDataType(); // input tensors - engineInputs.requestTypesDevice = manager.emptyTensor(runtime::MemoryType::kGPU, nvinfer1::DataType::kINT32); + engineInputs.requestTypesDevice = manager.emptyTensor(runtime::MemoryType::kGPU, tensorrt_llm::DataType::kINT32); engineInputs.temperatures = manager.emptyTensor(runtime::MemoryType::kGPU, dtype); engineInputs.draftTokens = manager.gpu(ITensor::makeShape({maxNumSequences, numBeams, beamLength}), TRTTokenIdType); engineInputs.draftIndices - = manager.gpu(ITensor::makeShape({maxNumSequences, numBeams, beamLength}), nvinfer1::DataType::kINT32); + = manager.gpu(ITensor::makeShape({maxNumSequences, numBeams, beamLength}), tensorrt_llm::DataType::kINT32); engineInputs.draftProbs = manager.gpu(ITensor::makeShape({maxNumSequences, numBeams, beamDraftLength, vocabSizePadded}), dtype); - engineInputs.generationLengths = manager.emptyTensor(runtime::MemoryType::kGPU, nvinfer1::DataType::kINT32); - engineInputs.positionIds = manager.emptyTensor(runtime::MemoryType::kGPU, nvinfer1::DataType::kINT32); - engineInputs.positionOffsets = manager.emptyTensor(runtime::MemoryType::kGPU, nvinfer1::DataType::kINT32); - engineInputs.packedMasks = manager.emptyTensor(runtime::MemoryType::kGPU, nvinfer1::DataType::kINT32); + engineInputs.generationLengths = manager.emptyTensor(runtime::MemoryType::kGPU, tensorrt_llm::DataType::kINT32); + engineInputs.positionIds = manager.emptyTensor(runtime::MemoryType::kGPU, tensorrt_llm::DataType::kINT32); + engineInputs.positionOffsets = manager.emptyTensor(runtime::MemoryType::kGPU, tensorrt_llm::DataType::kINT32); + engineInputs.packedMasks = manager.emptyTensor(runtime::MemoryType::kGPU, tensorrt_llm::DataType::kINT32); engineInputs.randomDataSample = manager.emptyTensor(runtime::MemoryType::kGPU, dtype); engineInputs.randomDataValidation = manager.emptyTensor(runtime::MemoryType::kGPU, dtype); - engineInputs.positionIdsBase = manager.emptyTensor(runtime::MemoryType::kGPU, nvinfer1::DataType::kINT32); - engineInputs.useSpecDecoding = manager.cpu(ITensor::makeShape({1}), nvinfer1::DataType::kINT32); + engineInputs.positionIdsBase = manager.emptyTensor(runtime::MemoryType::kGPU, tensorrt_llm::DataType::kINT32); + engineInputs.useSpecDecoding = manager.cpu(ITensor::makeShape({1}), tensorrt_llm::DataType::kINT32); bufferCast(*engineInputs.useSpecDecoding)[0] = 1; // output tensors engineOutputs.nextDraftTokens = manager.gpu(ITensor::makeShape({maxNumSequences, numBeams, beamLength}), TRTTokenIdType); engineOutputs.nextDraftIndices - = manager.gpu(ITensor::makeShape({maxNumSequences, numBeams, beamLength}), nvinfer1::DataType::kINT32); + = manager.gpu(ITensor::makeShape({maxNumSequences, numBeams, beamLength}), tensorrt_llm::DataType::kINT32); engineOutputs.nextDraftProbs = manager.gpu(ITensor::makeShape({maxNumSequences, numBeams, beamDraftLength, vocabSizePadded}), dtype); - engineOutputs.maxGenToken = manager.gpu(ITensor::makeShape({1}), nvinfer1::DataType::kINT32); - engineOutputs.totalGenToken = manager.gpu(ITensor::makeShape({1}), nvinfer1::DataType::kINT32); + engineOutputs.maxGenToken = manager.gpu(ITensor::makeShape({1}), tensorrt_llm::DataType::kINT32); + engineOutputs.totalGenToken = manager.gpu(ITensor::makeShape({1}), tensorrt_llm::DataType::kINT32); - engineOutputs.nextGenerationLengths = manager.emptyTensor(runtime::MemoryType::kGPU, nvinfer1::DataType::kINT32); - engineOutputs.nextPositionOffsets = manager.emptyTensor(runtime::MemoryType::kGPU, nvinfer1::DataType::kINT32); - engineOutputs.masks = manager.emptyTensor(runtime::MemoryType::kGPU, nvinfer1::DataType::kBOOL); + engineOutputs.nextGenerationLengths + = manager.emptyTensor(runtime::MemoryType::kGPU, tensorrt_llm::DataType::kINT32); + engineOutputs.nextPositionOffsets = manager.emptyTensor(runtime::MemoryType::kGPU, tensorrt_llm::DataType::kINT32); + engineOutputs.masks = manager.emptyTensor(runtime::MemoryType::kGPU, tensorrt_llm::DataType::kBOOL); engineOutputs.nextFlatTokens = manager.emptyTensor(runtime::MemoryType::kGPU, TRTTokenIdType); - engineOutputs.bestPathLengths = manager.emptyTensor(runtime::MemoryType::kGPU, nvinfer1::DataType::kINT32); - engineOutputs.bestPathIndices = manager.emptyTensor(runtime::MemoryType::kGPU, nvinfer1::DataType::kINT32); - engineOutputs.packedPositionIds = manager.emptyTensor(runtime::MemoryType::kGPU, nvinfer1::DataType::kINT32); + engineOutputs.bestPathLengths = manager.emptyTensor(runtime::MemoryType::kGPU, tensorrt_llm::DataType::kINT32); + engineOutputs.bestPathIndices = manager.emptyTensor(runtime::MemoryType::kGPU, tensorrt_llm::DataType::kINT32); + engineOutputs.packedPositionIds = manager.emptyTensor(runtime::MemoryType::kGPU, tensorrt_llm::DataType::kINT32); // helper tensors auto const& stream = manager.getStream(); scanTempStorageBytes = tksd::invokeScanGenerationLengths(nullptr, 0, nullptr, nullptr, maxNumSequences, stream.get()); scanTempStorage = manager.gpu(scanTempStorageBytes); - cumSumGenerationLengths = manager.emptyTensor(runtime::MemoryType::kGPU, nvinfer1::DataType::kINT32); + cumSumGenerationLengths = manager.emptyTensor(runtime::MemoryType::kGPU, tensorrt_llm::DataType::kINT32); // pre-allocate empty tensors reshape(0, maxNumSequences, modelConfig); @@ -295,15 +298,15 @@ void ExplicitDraftTokensBuffers::setFromInputs(SizeType32 numCtxSequences, SizeT switch (dtype) { - case nvinfer1::DataType::kFLOAT: + case tensorrt_llm::DataType::kFLOAT: setFromInputs(numCtxSequences, numGenSequences, vocabSizePadded, seqSlots, draftBuffers, contextPositionIds, *explicitDraftTokensModule, stream); break; - case nvinfer1::DataType::kHALF: + case tensorrt_llm::DataType::kHALF: setFromInputs(numCtxSequences, numGenSequences, vocabSizePadded, seqSlots, draftBuffers, contextPositionIds, *explicitDraftTokensModule, stream); break; - case nvinfer1::DataType::kBF16: + case tensorrt_llm::DataType::kBF16: setFromInputs<__nv_bfloat16>(numCtxSequences, numGenSequences, vocabSizePadded, seqSlots, draftBuffers, contextPositionIds, *explicitDraftTokensModule, stream); break; diff --git a/cpp/tensorrt_llm/runtime/gptDecoder.cpp b/cpp/tensorrt_llm/runtime/gptDecoder.cpp index 930877206462..e1ac1717af45 100644 --- a/cpp/tensorrt_llm/runtime/gptDecoder.cpp +++ b/cpp/tensorrt_llm/runtime/gptDecoder.cpp @@ -21,7 +21,7 @@ #include "tensorrt_llm/layers/dynamicDecodeLayer.h" #include "tensorrt_llm/runtime/decodingLayerWorkspace.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include @@ -121,7 +121,7 @@ void GptDecoder::disableLookahead( template void GptDecoder::setup(SamplingConfig const& samplingConfig, size_t batchSize, TensorConstPtr const& batchSlots, - std::optional const& output, std::optional explicitDraftTokensDType, + std::optional const& output, std::optional explicitDraftTokensDType, std::optional> const& lookaheadPrompt, std::optional> const& lookaheadAlgoConfigs) { diff --git a/cpp/tensorrt_llm/runtime/gptDecoderBatched.cpp b/cpp/tensorrt_llm/runtime/gptDecoderBatched.cpp index c55d02093afc..7b3a12ed7a2c 100644 --- a/cpp/tensorrt_llm/runtime/gptDecoderBatched.cpp +++ b/cpp/tensorrt_llm/runtime/gptDecoderBatched.cpp @@ -22,6 +22,7 @@ #include "tensorrt_llm/batch_manager/decoderBuffers.h" #include "tensorrt_llm/batch_manager/llmRequest.h" #include "tensorrt_llm/common/assert.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/executor/types.h" #include "tensorrt_llm/kernels/decodingKernels.h" #include "tensorrt_llm/runtime/bufferManager.h" @@ -74,7 +75,7 @@ void GptDecoderBatched::disableLookahead(RequestVector const& genRequests, Tenso } void GptDecoderBatched::setup(executor::DecodingMode const& mode, SizeType32 maxNumSequences, SizeType32 maxBeamWidth, - nvinfer1::DataType dtype, ModelConfig const& modelConfig, WorldConfig const& worldConfig) + tensorrt_llm::DataType dtype, ModelConfig const& modelConfig, WorldConfig const& worldConfig) { TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); TLLM_CHECK(maxNumSequences > 0); diff --git a/cpp/tensorrt_llm/runtime/gptJsonConfig.cpp b/cpp/tensorrt_llm/runtime/gptJsonConfig.cpp index 311f63eaf1e7..47310a9a1282 100644 --- a/cpp/tensorrt_llm/runtime/gptJsonConfig.cpp +++ b/cpp/tensorrt_llm/runtime/gptJsonConfig.cpp @@ -20,6 +20,7 @@ #include "modelConfig.h" #include "tensorrt_llm/common/assert.h" #include "tensorrt_llm/common/logger.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/runtime/eagleModule.h" #include "tensorrt_llm/runtime/explicitDraftTokensModule.h" #include "tensorrt_llm/runtime/jsonSerialization.h" @@ -80,14 +81,14 @@ std::optional parseJsonFieldOptional(Json const& json, std::string_vi return value; } -nvinfer1::DataType strToDType(std::string type) +tensorrt_llm::DataType strToDType(std::string type) { - static std::map const typeMap = {{"int64", nvinfer1::DataType::kINT64}, - {"int32", nvinfer1::DataType::kINT32}, {"int", nvinfer1::DataType::kINT32}, - {"float32", nvinfer1::DataType::kFLOAT}, {"bfloat16", nvinfer1::DataType::kBF16}, - {"float16", nvinfer1::DataType::kHALF}, {"bool", nvinfer1::DataType::kBOOL}, - {"uint8", nvinfer1::DataType::kUINT8}, {"int8", nvinfer1::DataType::kINT8}, {"fp8", nvinfer1::DataType::kFP8}, - {"int4", nvinfer1::DataType::kINT4}}; + static std::map const typeMap = {{"int64", tensorrt_llm::DataType::kINT64}, + {"int32", tensorrt_llm::DataType::kINT32}, {"int", tensorrt_llm::DataType::kINT32}, + {"float32", tensorrt_llm::DataType::kFLOAT}, {"bfloat16", tensorrt_llm::DataType::kBF16}, + {"float16", tensorrt_llm::DataType::kHALF}, {"bool", tensorrt_llm::DataType::kBOOL}, + {"uint8", tensorrt_llm::DataType::kUINT8}, {"int8", tensorrt_llm::DataType::kINT8}, + {"fp8", tensorrt_llm::DataType::kFP8}, {"int4", tensorrt_llm::DataType::kINT4}}; TLLM_CHECK_WITH_INFO(typeMap.count(type) > 0, type + " not found in strToDtype."); return typeMap.at(type); @@ -140,14 +141,14 @@ std::vector buildLayerTypes( return result; } -ModelConfig parseMultimodalConfig(Json const& json, nvinfer1::DataType dataType) +ModelConfig parseMultimodalConfig(Json const& json, tensorrt_llm::DataType dataType) { return ModelConfig{128, 10, 10, 0, 1, 128, dataType}; // use dummy values because vision engines of multimodal models does not record this info in config } ModelConfig createModelConfig(Json const& json, bool engineVersionNone, SizeType32 tensorParallelism, - SizeType32 contextParallelism, nvinfer1::DataType dataType) + SizeType32 contextParallelism, tensorrt_llm::DataType dataType) { auto const& config = engineVersionNone ? json.at("builder_config") : json.at("pretrained_config"); auto const multiModalName = parseJsonFieldOptional(config, "model_name"); @@ -248,14 +249,14 @@ ModelConfig createModelConfig(Json const& json, bool engineVersionNone, SizeType modelConfig.setLayerTypes(layerTypes); // Set logits datatype - auto logitsDtype = nvinfer1::DataType::kFLOAT; + auto logitsDtype = tensorrt_llm::DataType::kFLOAT; if (logitsDtypeStr == "float32") { - logitsDtype = nvinfer1::DataType::kFLOAT; + logitsDtype = tensorrt_llm::DataType::kFLOAT; } else if (logitsDtypeStr == "float16") { - logitsDtype = nvinfer1::DataType::kHALF; + logitsDtype = tensorrt_llm::DataType::kHALF; } else { @@ -490,15 +491,15 @@ GptJsonConfig parseJson(InputType&& input) { if (precision == "float32") { - return nvinfer1::DataType::kFLOAT; + return tensorrt_llm::DataType::kFLOAT; } if (precision == "float16") { - return nvinfer1::DataType::kHALF; + return tensorrt_llm::DataType::kHALF; } if (precision == "bfloat16") { - return nvinfer1::DataType::kBF16; + return tensorrt_llm::DataType::kBF16; } TLLM_THROW("Model data type '%s' not supported", precision.c_str()); }(); diff --git a/cpp/tensorrt_llm/runtime/iBuffer.cpp b/cpp/tensorrt_llm/runtime/iBuffer.cpp index 77707a0e4cf8..82574b658b39 100644 --- a/cpp/tensorrt_llm/runtime/iBuffer.cpp +++ b/cpp/tensorrt_llm/runtime/iBuffer.cpp @@ -20,6 +20,7 @@ #include "tensorrt_llm/common/assert.h" #include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/runtime/bufferView.h" #include @@ -48,7 +49,7 @@ IBuffer::UniquePtr IBuffer::slice(IBuffer::SharedPtr buffer, std::size_t offset, return std::make_unique(std::move(buffer), offset, size); } -IBuffer::UniquePtr IBuffer::wrap(void* data, nvinfer1::DataType type, std::size_t size, std::size_t capacity) +IBuffer::UniquePtr IBuffer::wrap(void* data, tensorrt_llm::DataType type, std::size_t size, std::size_t capacity) { TLLM_CHECK_WITH_INFO(size <= capacity, "Requested size is larger than capacity"); auto memoryType = IBuffer::memoryType(data); @@ -91,17 +92,17 @@ char const* IBuffer::getDataTypeName(DataType dataType) { switch (dataType) { - case nvinfer1::DataType::kINT64: return DataTypeTraits::name; - case nvinfer1::DataType::kINT32: return DataTypeTraits::name; - case nvinfer1::DataType::kFLOAT: return DataTypeTraits::name; - case nvinfer1::DataType::kBF16: return DataTypeTraits::name; - case nvinfer1::DataType::kHALF: return DataTypeTraits::name; - case nvinfer1::DataType::kBOOL: return DataTypeTraits::name; - case nvinfer1::DataType::kUINT8: return DataTypeTraits::name; - case nvinfer1::DataType::kINT8: return DataTypeTraits::name; - case nvinfer1::DataType::kFP8: return DataTypeTraits::name; - case nvinfer1::DataType::kINT4: [[fallthrough]] /* do nothing */; - case nvinfer1::DataType::kFP4: [[fallthrough]] /* do nothing */; + case tensorrt_llm::DataType::kINT64: return DataTypeTraits::name; + case tensorrt_llm::DataType::kINT32: return DataTypeTraits::name; + case tensorrt_llm::DataType::kFLOAT: return DataTypeTraits::name; + case tensorrt_llm::DataType::kBF16: return DataTypeTraits::name; + case tensorrt_llm::DataType::kHALF: return DataTypeTraits::name; + case tensorrt_llm::DataType::kBOOL: return DataTypeTraits::name; + case tensorrt_llm::DataType::kUINT8: return DataTypeTraits::name; + case tensorrt_llm::DataType::kINT8: return DataTypeTraits::name; + case tensorrt_llm::DataType::kFP8: return DataTypeTraits::name; + case tensorrt_llm::DataType::kINT4: [[fallthrough]] /* do nothing */; + case tensorrt_llm::DataType::kFP4: [[fallthrough]] /* do nothing */; default: TLLM_THROW("Unknown data type"); } } diff --git a/cpp/tensorrt_llm/runtime/iTensor.cpp b/cpp/tensorrt_llm/runtime/iTensor.cpp index f78b25fdb19a..70b31707130a 100644 --- a/cpp/tensorrt_llm/runtime/iTensor.cpp +++ b/cpp/tensorrt_llm/runtime/iTensor.cpp @@ -18,6 +18,7 @@ #include "tensorrt_llm/common/memoryUtils.h" #include "tensorrt_llm/common/stringUtils.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/runtime/bufferManager.h" #include "tensorrt_llm/runtime/tensorView.h" #include "tensorrt_llm/runtime/tllmBuffers.h" @@ -72,22 +73,22 @@ ITensor::UniquePtr ITensor::slice(SharedPtr tensor, Shape const& offsetDims, ITe return std::make_unique(std::move(tensor), offset, volume(dims), dims); } -ITensor::UniquePtr ITensor::view(IBuffer::SharedPtr buffer, nvinfer1::Dims const& dims) +ITensor::UniquePtr ITensor::view(IBuffer::SharedPtr buffer, tensorrt_llm::Dims const& dims) { auto const size = buffer->getSize(); return std::make_unique(std::move(buffer), 0, size, dims); } -nvinfer1::Dims ITensor::makeShape(std::initializer_list const& dims) +tensorrt_llm::Dims ITensor::makeShape(std::initializer_list const& dims) { - TLLM_CHECK_WITH_INFO(dims.size() <= nvinfer1::Dims::MAX_DIMS, "Number of dimensions is too large"); - nvinfer1::Dims shape{}; + TLLM_CHECK_WITH_INFO(dims.size() <= tensorrt_llm::Dims::MAX_DIMS, "Number of dimensions is too large"); + tensorrt_llm::Dims shape{}; shape.nbDims = static_cast(dims.size()); std::copy(dims.begin(), dims.end(), shape.d); return shape; } -std::string ITensor::toString(nvinfer1::Dims const& dims) +std::string ITensor::toString(tensorrt_llm::Dims const& dims) { if (dims.nbDims < 0) { @@ -103,7 +104,8 @@ std::string ITensor::toString(nvinfer1::Dims const& dims) } } -ITensor::UniquePtr ITensor::wrap(void* data, nvinfer1::DataType type, nvinfer1::Dims const& shape, std::size_t capacity) +ITensor::UniquePtr ITensor::wrap( + void* data, tensorrt_llm::DataType type, tensorrt_llm::Dims const& shape, std::size_t capacity) { auto const size = volumeNonNegative(shape); TLLM_CHECK_WITH_INFO(size <= capacity, "Requested size is larger than capacity"); @@ -230,18 +232,18 @@ std::ostream& tensorrt_llm::runtime::operator<<(std::ostream& out, ITensor const { switch (tensor.getDataType()) { - case nvinfer1::DataType::kFLOAT: printTensor(tensor, out); break; - case nvinfer1::DataType::kHALF: printTensor(tensor, out); break; - case nvinfer1::DataType::kBOOL: printTensor(tensor, out); break; - case nvinfer1::DataType::kINT8: printTensor(tensor, out); break; - case nvinfer1::DataType::kINT32: printTensor(tensor, out); break; - case nvinfer1::DataType::kINT64: printTensor(tensor, out); break; - case nvinfer1::DataType::kUINT8: printTensor(tensor, out); break; + case tensorrt_llm::DataType::kFLOAT: printTensor(tensor, out); break; + case tensorrt_llm::DataType::kHALF: printTensor(tensor, out); break; + case tensorrt_llm::DataType::kBOOL: printTensor(tensor, out); break; + case tensorrt_llm::DataType::kINT8: printTensor(tensor, out); break; + case tensorrt_llm::DataType::kINT32: printTensor(tensor, out); break; + case tensorrt_llm::DataType::kINT64: printTensor(tensor, out); break; + case tensorrt_llm::DataType::kUINT8: printTensor(tensor, out); break; #ifdef ENABLE_BF16 - case nvinfer1::DataType::kBF16: printTensor<__nv_bfloat16, float>(tensor, out); break; + case tensorrt_llm::DataType::kBF16: printTensor<__nv_bfloat16, float>(tensor, out); break; #endif #ifdef ENABLE_FP8 - case nvinfer1::DataType::kFP8: printTensor<__nv_fp8_e4m3, float>(tensor, out); break; + case tensorrt_llm::DataType::kFP8: printTensor<__nv_fp8_e4m3, float>(tensor, out); break; #endif default: TLLM_THROW("Unsupported data type"); } diff --git a/cpp/tensorrt_llm/runtime/ipcUtils.cpp b/cpp/tensorrt_llm/runtime/ipcUtils.cpp index 23a7e28a4f27..48368844f850 100644 --- a/cpp/tensorrt_llm/runtime/ipcUtils.cpp +++ b/cpp/tensorrt_llm/runtime/ipcUtils.cpp @@ -20,7 +20,7 @@ #include "tensorrt_llm/common/workspace.h" #include "tensorrt_llm/runtime/utils/mpiUtils.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include namespace tensorrt_llm::runtime @@ -83,7 +83,7 @@ void IpcMemory::allocateIpcMemory(std::size_t bufferSize, BufferManager const& m // IPC handles. If we want to support stream-ordered allocations here, we need to create another pool with the // correct handle type. auto const ipcAlignedBufferSize = common::alignSize(bufferSize, 1LU << 21); - mBuffer = BufferManager::gpuSync(ipcAlignedBufferSize, nvinfer1::DataType::kUINT8); + mBuffer = BufferManager::gpuSync(ipcAlignedBufferSize, tensorrt_llm::DataType::kUINT8); manager.setZero(*mBuffer); auto* bufferPtr = mBuffer->data(); @@ -149,7 +149,7 @@ AllReduceBuffers::AllReduceBuffers(SizeType32 maxBatchSize, SizeType32 maxBeamWi { auto const tpSize = worldConfig.getTensorParallelism(); mAllReduceCommPtrs = BufferManager::cpu( - ITensor::makeShape({static_cast(7) * tpSize + 3}), nvinfer1::DataType::kINT64); + ITensor::makeShape({static_cast(7) * tpSize + 3}), tensorrt_llm::DataType::kINT64); } else { @@ -178,7 +178,7 @@ AllReduceBuffers::AllReduceBuffers(SizeType32 maxBatchSize, SizeType32 maxBeamWi mAllReduceCommPtrs = BufferManager::cpu(ITensor::makeShape({static_cast(mIpcMemoryHandles.size()) * tpSize + 3}), - nvinfer1::DataType::kINT64); + tensorrt_llm::DataType::kINT64); auto commPtrs = BufferRange(*mAllReduceCommPtrs); // Start from 1 since 0 represents released state for barrier at the beginning of the all_reduce. // The last element is the barrier flag counter. @@ -211,9 +211,9 @@ AllReduceBuffers::AllReduceBuffers(SizeType32 maxBatchSize, SizeType32 maxBeamWi void lamportInitializeAll(void* buffer_0, void* buffer_1, void* buffer_2, size_t size) { #if ENABLE_MULTI_DEVICE - tensorrt_llm::kernels::lamportInitialize(buffer_0, size / sizeof(half), nvinfer1::DataType::kHALF, 0); - tensorrt_llm::kernels::lamportInitialize(buffer_1, size / sizeof(half), nvinfer1::DataType::kHALF, 0); - tensorrt_llm::kernels::lamportInitialize(buffer_2, size / sizeof(half), nvinfer1::DataType::kHALF, 0); + tensorrt_llm::kernels::lamportInitialize(buffer_0, size / sizeof(half), tensorrt_llm::DataType::kHALF, 0); + tensorrt_llm::kernels::lamportInitialize(buffer_1, size / sizeof(half), tensorrt_llm::DataType::kHALF, 0); + tensorrt_llm::kernels::lamportInitialize(buffer_2, size / sizeof(half), tensorrt_llm::DataType::kHALF, 0); cudaDeviceSynchronize(); #endif } diff --git a/cpp/tensorrt_llm/runtime/layerProfiler.cpp b/cpp/tensorrt_llm/runtime/layerProfiler.cpp deleted file mode 100644 index 4c3c9779cedb..000000000000 --- a/cpp/tensorrt_llm/runtime/layerProfiler.cpp +++ /dev/null @@ -1,97 +0,0 @@ -/* - * Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. - * - * 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. - */ - -#include "tensorrt_llm/runtime/layerProfiler.h" -#include -#include -#include -#include - -using namespace tensorrt_llm::runtime; - -void LayerProfiler::reportLayerTime(char const* layerName, float timeMs) noexcept -{ - if (mIterator == mLayers.end()) - { - bool const first = !mLayers.empty() && mLayers.begin()->name == layerName; - mUpdatesCount += mLayers.empty() || first; - if (first) - { - mIterator = mLayers.begin(); - } - else - { - mLayers.emplace_back(); - mLayers.back().name = layerName; - mIterator = mLayers.end() - 1; - } - } - - mIterator->timeMs.push_back(timeMs); - ++mIterator; -} - -float LayerProfiler::getTotalTime() const noexcept -{ - auto const plusLayerTime = [](float accumulator, LayerProfile const& lp) - { return accumulator + std::accumulate(lp.timeMs.begin(), lp.timeMs.end(), 0.F, std::plus()); }; - return std::accumulate(mLayers.begin(), mLayers.end(), 0.0F, plusLayerTime); -} - -std::string LayerProfiler::getLayerProfile() noexcept -{ - std::string const nameHdr(" Layer"); - std::string const timeHdr(" Time(ms)"); - - float const totalTimeMs = getTotalTime(); - - auto const timeLength = timeHdr.size(); - - std::unordered_map layer2times; - std::vector layer_order; - for (auto const& p : mLayers) - { - if (!layer2times.count(p.name)) - { - layer2times[p.name] = 0; - layer_order.push_back(p.name); - } - for (auto const& t : p.timeMs) - { - layer2times[p.name] += t; - } - } - - std::stringstream ss; - ss << "\n=== Per-layer Profile ===\n" << timeHdr << nameHdr << "\n"; - - for (auto const& name : layer_order) - { - if (layer2times[name] == 0.0f) - { - continue; - } - ss << std::setw(timeLength) << std::fixed << std::setprecision(2) << layer2times[name] << " " << name << "\n"; - } - - ss << std::setw(timeLength) << std::fixed << std::setprecision(2) << totalTimeMs << " Total\n"; - ss << "\n"; - - // clear data - mLayers.clear(); - - return ss.str(); -} diff --git a/cpp/tensorrt_llm/runtime/layerProfiler.h b/cpp/tensorrt_llm/runtime/layerProfiler.h deleted file mode 100644 index bcae1546de3d..000000000000 --- a/cpp/tensorrt_llm/runtime/layerProfiler.h +++ /dev/null @@ -1,47 +0,0 @@ -/* - * Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. - * - * 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. - */ - -#pragma once - -#include "tensorrt_llm/runtime/common.h" -#include - -#include - -namespace tensorrt_llm::runtime -{ -struct LayerProfile -{ - std::string name; - std::vector timeMs; -}; - -class LayerProfiler : public nvinfer1::IProfiler -{ - -public: - void reportLayerTime(char const* layerName, float timeMs) noexcept override; - - std::string getLayerProfile() noexcept; - -private: - [[nodiscard]] float getTotalTime() const noexcept; - - std::vector mLayers; - std::vector::iterator mIterator{mLayers.begin()}; - int32_t mUpdatesCount{0}; -}; -} // namespace tensorrt_llm::runtime diff --git a/cpp/tensorrt_llm/runtime/lookaheadBuffers.cpp b/cpp/tensorrt_llm/runtime/lookaheadBuffers.cpp index ef800ef218e4..5e77046c47e0 100644 --- a/cpp/tensorrt_llm/runtime/lookaheadBuffers.cpp +++ b/cpp/tensorrt_llm/runtime/lookaheadBuffers.cpp @@ -16,208 +16,23 @@ */ #include "tensorrt_llm/runtime/lookaheadBuffers.h" -#include "tensorrt_llm/layers/lookaheadDecodingUtils.h" +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/common/tllmDataType.h" namespace tensorrt_llm::runtime { LookaheadDecodingBuffers::LookaheadDecodingBuffers( SizeType32 maxNumSequences, SizeType32 maxTokensPerStep, BufferManager const& bufferManager) - : generationLengths(bufferManager.gpu(ITensor::makeShape({maxNumSequences}), nvinfer1::DataType::kINT32)) + : generationLengths(bufferManager.gpu(ITensor::makeShape({maxNumSequences}), tensorrt_llm::DataType::kINT32)) , positionOffsets( - bufferManager.gpu(ITensor::makeShape({maxNumSequences, maxTokensPerStep}), nvinfer1::DataType::kINT32)) + bufferManager.gpu(ITensor::makeShape({maxNumSequences, maxTokensPerStep}), tensorrt_llm::DataType::kINT32)) , packedMasks(bufferManager.gpu(ITensor::makeShape({maxNumSequences, maxTokensPerStep, static_cast(common::divUp(maxTokensPerStep, 32))}), - nvinfer1::DataType::kINT32)) + tensorrt_llm::DataType::kINT32)) , positionIds( - bufferManager.gpu(ITensor::makeShape({maxNumSequences, maxTokensPerStep}), nvinfer1::DataType::kINT32)) + bufferManager.gpu(ITensor::makeShape({maxNumSequences, maxTokensPerStep}), tensorrt_llm::DataType::kINT32)) { } -LookaheadRuntimeBuffers::LookaheadRuntimeBuffers(SizeType32 maxBatchSize, SizeType32 maxBeamWidth, - BufferManager const& manager, ModelConfig const& modelConfig, WorldConfig const& worldConfig, - executor::DecodingConfig const& /* decodingConfig */, TllmRuntime const& runtime) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - TLLM_CHECK_WITH_INFO(maxBeamWidth == 1, "Lookahead decoding does not support beam search"); - - auto const tokensPerStep = modelConfig.getMaxDecodingTokens(); - auto const numPackedMasks = static_cast(tensorrt_llm::common::divUp(tokensPerStep, 32)); - - cumSumLength = manager.pinned(ITensor::makeShape({1}), nvinfer1::DataType::kINT32); - - packedMasksDevice - = manager.gpu(ITensor::makeShape({maxBatchSize * tokensPerStep, numPackedMasks}), nvinfer1::DataType::kINT32); - positionOffsetsDevice = manager.gpu(ITensor::makeShape({maxBatchSize, tokensPerStep}), nvinfer1::DataType::kINT32); - generationLengthsDevice = manager.gpu(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kINT32); - positionIdsDevice = manager.gpu(ITensor::makeShape({maxBatchSize, tokensPerStep}), nvinfer1::DataType::kINT32); - - packedMaskHost = manager.cpu(packedMasksDevice->getShape(), nvinfer1::DataType::kINT32); - positionOffsetsHost = manager.cpu(positionOffsetsDevice->getShape(), nvinfer1::DataType::kINT32); - generationLengthsHost = manager.cpu(generationLengthsDevice->getShape(), nvinfer1::DataType::kINT32); - positionIdsHost = manager.cpu(positionIdsDevice->getShape(), nvinfer1::DataType::kINT32); - - packedMaskHostCopy = manager.cpu(packedMasksDevice->getShape(), nvinfer1::DataType::kINT32); - positionOffsetsHostCopy = manager.cpu(positionOffsetsDevice->getShape(), nvinfer1::DataType::kINT32); - generationLengthsHostCopy = manager.cpu(generationLengthsDevice->getShape(), nvinfer1::DataType::kINT32); - positionIdsHostCopy = manager.cpu(positionIdsDevice->getShape(), nvinfer1::DataType::kINT32); - - batchSlotsHostCopy = manager.cpu(generationLengthsDevice->getShape(), nvinfer1::DataType::kINT32); - - useSpecDecoding = manager.cpu(ITensor::makeShape({1}), nvinfer1::DataType::kINT32); - bufferCast(*useSpecDecoding)[0] = 1; - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void LookaheadRuntimeBuffers::setFromInputs(SizeType32 numCtxSequences, SizeType32 numGenSequences, - ITensor const& requestTypes, ITensor const& seqSlots, LookaheadDecodingBuffers const& decoderLookaheadBuffers, - TllmRuntime const& runtime, ModelConfig const& modelConfig, WorldConfig const& worldConfig) const -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - auto const& manager = runtime.getBufferManager(); - - auto const tokensPerStep = modelConfig.getMaxDecodingTokens(); - - manager.copy(seqSlots, *batchSlotsHostCopy); - manager.copy(*decoderLookaheadBuffers.generationLengths, *generationLengthsHostCopy); - manager.copy(*decoderLookaheadBuffers.positionOffsets, *positionOffsetsHostCopy); - manager.copy(*decoderLookaheadBuffers.packedMasks, *packedMaskHostCopy); - manager.copy(*decoderLookaheadBuffers.positionIds, *positionIdsHostCopy); - - manager.getStream().synchronize(); - - BufferRange batchSlotsRange(*batchSlotsHostCopy); - BufferRange cumSumLengthRange(*cumSumLength); - - SizeType32 maxGenerationLength = 0; - for (SizeType32 bi = 0; bi < numGenSequences; bi++) - { - SizeType32 gbi = batchSlotsRange[bi + numCtxSequences]; - SizeType32 theLength = BufferRange(*generationLengthsHostCopy)[gbi]; - maxGenerationLength = std::max(maxGenerationLength, theLength); - } - - auto positionOffsetShape = positionOffsetsHost->getShape(); - positionOffsetShape.d[1] = maxGenerationLength; - positionOffsetsHost->reshape(positionOffsetShape); - positionOffsetsDevice->reshape(positionOffsetShape); - - auto positionIdsShape = positionIdsHostCopy->getShape(); - auto positionIdsShape1D = ITensor::makeShape({ITensor::volume(positionIdsShape)}); - positionIdsHostCopy->reshape(positionIdsShape1D); - positionIdsHost->reshape(positionIdsShape1D); - - cumSumLengthRange[0] = 0; - for (SizeType32 bi = 0; bi < numGenSequences; bi++) - { - SizeType32 gbi = batchSlotsRange[bi + numCtxSequences]; - SizeType32 theLength = BufferRange(*generationLengthsHostCopy)[gbi]; - - manager.copy(*ITensor::at(generationLengthsHostCopy, {gbi}), *ITensor::at(generationLengthsHost, {bi})); - - manager.copy(*ITensor::slice(positionOffsetsHostCopy, {gbi, 0}, theLength), - *ITensor::slice(positionOffsetsHost, {bi, 0}, theLength)); - - manager.copy(*ITensor::slice(packedMaskHostCopy, gbi * tokensPerStep, theLength), - *ITensor::slice(packedMaskHost, cumSumLengthRange[0], theLength)); - - manager.copy(*ITensor::slice(positionIdsHostCopy, gbi * tokensPerStep, theLength), - *ITensor::slice(positionIdsHost, cumSumLengthRange[0], theLength)); - - cumSumLengthRange[0] += theLength; - } - - positionIdsHostCopy->reshape(positionIdsShape); - positionIdsHost->reshape(positionIdsShape); - positionIdsDevice->reshape(positionIdsShape); - - manager.copy(*ITensor::slice(generationLengthsHost, 0, numGenSequences), - *ITensor::slice(generationLengthsDevice, 0, numGenSequences)); - manager.copy(*ITensor::slice(positionOffsetsHost, 0, numGenSequences), - *ITensor::slice(positionOffsetsDevice, 0, numGenSequences)); - manager.copy(*ITensor::slice(packedMaskHost, 0, numGenSequences * tokensPerStep), - *ITensor::slice(packedMasksDevice, 0, numGenSequences * tokensPerStep)); - manager.copy( - *ITensor::slice(positionIdsHost, 0, numGenSequences), *ITensor::slice(positionIdsDevice, 0, numGenSequences)); - positionIdsDevice->reshape(ITensor::makeShape({cumSumLengthRange[0]})); - - manager.getStream().synchronize(); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void LookaheadRuntimeBuffers::reshape(SizeType32 numCtxSequences, SizeType32 numGenSequences, SizeType32 tokensPerStep) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - auto const numSequences = numGenSequences; - - auto packedMaskShape = packedMasksDevice->getShape(); - packedMaskShape.d[0] = numSequences * tokensPerStep; - packedMasksDevice->reshape(packedMaskShape); - packedMaskHost->reshape(packedMaskShape); - - auto generationLengthsShape = generationLengthsDevice->getShape(); - generationLengthsShape.d[0] = numSequences; - generationLengthsDevice->reshape(generationLengthsShape); - generationLengthsHost->reshape(generationLengthsShape); - - auto positionOffsetsShape = positionOffsetsDevice->getShape(); - positionOffsetsShape.d[0] = numSequences; - positionOffsetsDevice->reshape(positionOffsetsShape); - positionOffsetsHost->reshape(positionOffsetsShape); - - auto positionIdsShape = positionIdsDevice->getShape(); - positionIdsShape.d[0] = numSequences; - positionIdsDevice->reshape(positionIdsShape); - positionIdsHost->reshape(positionIdsShape); - - auto batchSlotsShape = batchSlotsHostCopy->getShape(); - batchSlotsShape.d[0] = numCtxSequences + numGenSequences; - batchSlotsHostCopy->reshape(batchSlotsShape); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void LookaheadRuntimeBuffers::enableLookaheadDecoding(SizeType32 maxBatchSize, SizeType32 tokensPerStep) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - auto const numPackedMasks = static_cast(tensorrt_llm::common::divUp(tokensPerStep, 32)); - packedMasksDevice->reshape(ITensor::makeShape({maxBatchSize * tokensPerStep, numPackedMasks})); - generationLengthsDevice->reshape(ITensor::makeShape({maxBatchSize})); - positionOffsetsDevice->reshape(ITensor::makeShape({maxBatchSize, tokensPerStep})); - bufferCast(*useSpecDecoding)[0] = 1; - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void LookaheadRuntimeBuffers::disableLookaheadDecoding() -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - packedMasksDevice->reshape(ITensor::makeShape({1, 1})); - generationLengthsDevice->reshape(ITensor::makeShape({1})); - positionOffsetsDevice->reshape(ITensor::makeShape({1, 1})); - bufferCast(*useSpecDecoding)[0] = 0; - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void LookaheadRuntimeBuffers::insertInputTensors( - TensorMap& inputBuffers, TensorMap& /* outputBuffers */, WorldConfig const& /* worldConfig */) const -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - inputBuffers.insert_or_assign("spec_decoding_packed_mask", packedMasksDevice); - inputBuffers.insert_or_assign("spec_decoding_generation_lengths", generationLengthsDevice); - inputBuffers.insert_or_assign("spec_decoding_position_offsets", positionOffsetsDevice); - inputBuffers.insert_or_assign("spec_decoding_use", useSpecDecoding); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - } // namespace tensorrt_llm::runtime diff --git a/cpp/tensorrt_llm/runtime/loraCache.cpp b/cpp/tensorrt_llm/runtime/loraCache.cpp index 3dbb814f058b..36fb0363816f 100644 --- a/cpp/tensorrt_llm/runtime/loraCache.cpp +++ b/cpp/tensorrt_llm/runtime/loraCache.cpp @@ -23,6 +23,7 @@ #include "tensorrt_llm/common/cudaUtils.h" #include "tensorrt_llm/common/logger.h" #include "tensorrt_llm/common/memoryUtils.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/runtime/loraUtils.h" #include #include @@ -537,15 +538,15 @@ void LoraCache::splitTransposeCpu(ITensor& output, ITensor const& input, SizeTyp switch (input.getDataType()) { - case nvinfer1::DataType::kINT32: splitTransposeCpuInner(output, input, tpSize, tpRank); break; - case nvinfer1::DataType::kFLOAT: splitTransposeCpuInner(output, input, tpSize, tpRank); break; - case nvinfer1::DataType::kHALF: splitTransposeCpuInner(output, input, tpSize, tpRank); break; - case nvinfer1::DataType::kINT8: splitTransposeCpuInner(output, input, tpSize, tpRank); break; + case tensorrt_llm::DataType::kINT32: splitTransposeCpuInner(output, input, tpSize, tpRank); break; + case tensorrt_llm::DataType::kFLOAT: splitTransposeCpuInner(output, input, tpSize, tpRank); break; + case tensorrt_llm::DataType::kHALF: splitTransposeCpuInner(output, input, tpSize, tpRank); break; + case tensorrt_llm::DataType::kINT8: splitTransposeCpuInner(output, input, tpSize, tpRank); break; #ifdef ENABLE_FP8 - case nvinfer1::DataType::kFP8: splitTransposeCpuInner<__nv_fp8_e4m3>(output, input, tpSize, tpRank); break; + case tensorrt_llm::DataType::kFP8: splitTransposeCpuInner<__nv_fp8_e4m3>(output, input, tpSize, tpRank); break; #endif // ENABLE_FP8 #ifdef ENABLE_BF16 - case nvinfer1::DataType::kBF16: splitTransposeCpuInner<__nv_bfloat16>(output, input, tpSize, tpRank); break; + case tensorrt_llm::DataType::kBF16: splitTransposeCpuInner<__nv_bfloat16>(output, input, tpSize, tpRank); break; #endif // ENABLE_BF16 default: TLLM_CHECK_WITH_INFO(false, "data type not supported"); } diff --git a/cpp/tensorrt_llm/runtime/loraManager.cpp b/cpp/tensorrt_llm/runtime/loraManager.cpp index 8d7ebe389853..1d25ea20c8e4 100644 --- a/cpp/tensorrt_llm/runtime/loraManager.cpp +++ b/cpp/tensorrt_llm/runtime/loraManager.cpp @@ -26,8 +26,6 @@ #include "tensorrt_llm/runtime/utils/runtimeUtils.h" #include "tensorrt_llm/runtime/worldConfig.h" -#include - namespace tensorrt_llm::runtime { diff --git a/cpp/tensorrt_llm/runtime/loraUtils.cpp b/cpp/tensorrt_llm/runtime/loraUtils.cpp index 3c5e95162474..da7f1475ddec 100644 --- a/cpp/tensorrt_llm/runtime/loraUtils.cpp +++ b/cpp/tensorrt_llm/runtime/loraUtils.cpp @@ -17,6 +17,7 @@ #include "tensorrt_llm/runtime/loraUtils.h" #include "tensorrt_llm/common/assert.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/runtime/common.h" #include "tensorrt_llm/runtime/iTensor.h" #include "tensorrt_llm/runtime/modelConfig.h" @@ -57,7 +58,7 @@ void loraValidateRequestTensorDims(std::optional const& optR keys->getShape().d[0] == expectedBatchSize, "Expected batch dimension to be 1 for each lora request"); TLLM_CHECK_WITH_INFO(weights->getMemoryType() != MemoryType::kGPU, "Expected lora weights to be in CPU memory"); TLLM_CHECK_WITH_INFO(keys->getMemoryType() != MemoryType::kGPU, "Expected lora weights to be in CPU memory"); - TLLM_CHECK_WITH_INFO(keys->getDataType() == nvinfer1::DataType::kINT32, + TLLM_CHECK_WITH_INFO(keys->getDataType() == tensorrt_llm::DataType::kINT32, "Expected lora keys to have TYPE_INT32 but was " + std::string(keys->getDataTypeName())); TLLM_CHECK_WITH_INFO(keys->getShape().d[1] == weights->getShape().d[1], diff --git a/cpp/tensorrt_llm/runtime/ncclCommunicator.cpp b/cpp/tensorrt_llm/runtime/ncclCommunicator.cpp index b76efb75952a..7edfba42f935 100644 --- a/cpp/tensorrt_llm/runtime/ncclCommunicator.cpp +++ b/cpp/tensorrt_llm/runtime/ncclCommunicator.cpp @@ -18,6 +18,7 @@ #include "tensorrt_llm/common/envUtils.h" #include "tensorrt_llm/common/logger.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/runtime/ipcNvlsMemory.h" #include "tensorrt_llm/runtime/utils/multiDeviceUtils.h" @@ -102,19 +103,19 @@ void initNcclCommProbeWithTimeout(ncclUniqueId const& id, int worldSize, int ran } } -ncclDataType_t toNcclType(nvinfer1::DataType dataType) +ncclDataType_t toNcclType(tensorrt_llm::DataType dataType) { switch (dataType) { - case nvinfer1::DataType::kFLOAT: return ncclFloat32; - case nvinfer1::DataType::kHALF: return ncclHalf; - case nvinfer1::DataType::kINT8: return ncclInt8; - case nvinfer1::DataType::kINT32: return ncclInt32; - case nvinfer1::DataType::kUINT8: return ncclUint8; - case nvinfer1::DataType::kINT64: return ncclInt64; - case nvinfer1::DataType::kFP8: return ncclUint8; + case tensorrt_llm::DataType::kFLOAT: return ncclFloat32; + case tensorrt_llm::DataType::kHALF: return ncclHalf; + case tensorrt_llm::DataType::kINT8: return ncclInt8; + case tensorrt_llm::DataType::kINT32: return ncclInt32; + case tensorrt_llm::DataType::kUINT8: return ncclUint8; + case tensorrt_llm::DataType::kINT64: return ncclInt64; + case tensorrt_llm::DataType::kFP8: return ncclUint8; #if ENABLE_BF16 - case nvinfer1::DataType::kBF16: return ncclBfloat16; + case tensorrt_llm::DataType::kBF16: return ncclBfloat16; #endif // ENABLE_BF16 default: TLLM_THROW("Unsupported data type: %d", static_cast(dataType)); } @@ -123,7 +124,7 @@ ncclDataType_t toNcclType(nvinfer1::DataType dataType) } // namespace void NcclCommunicator::send( - void const* sendbuff, size_t count, nvinfer1::DataType dataType, int peer, CudaStream const& stream) const + void const* sendbuff, size_t count, tensorrt_llm::DataType dataType, int peer, CudaStream const& stream) const { #if ENABLE_MULTI_DEVICE TLLM_NCCL_CHECK(ncclSend(sendbuff, count, toNcclType(dataType), peer, mComm, stream.get())); @@ -133,7 +134,7 @@ void NcclCommunicator::send( } void NcclCommunicator::receive( - void* sendbuff, size_t count, nvinfer1::DataType dataType, int peer, CudaStream const& stream) const + void* sendbuff, size_t count, tensorrt_llm::DataType dataType, int peer, CudaStream const& stream) const { #if ENABLE_MULTI_DEVICE TLLM_NCCL_CHECK(ncclRecv(sendbuff, count, toNcclType(dataType), peer, mComm, stream.get())); diff --git a/cpp/tensorrt_llm/runtime/ncclCommunicator.h b/cpp/tensorrt_llm/runtime/ncclCommunicator.h index 76cce4beab8a..21d7f116e95b 100644 --- a/cpp/tensorrt_llm/runtime/ncclCommunicator.h +++ b/cpp/tensorrt_llm/runtime/ncclCommunicator.h @@ -16,6 +16,7 @@ #pragma once +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/runtime/cudaStream.h" #include "tensorrt_llm/runtime/iBuffer.h" #include "tensorrt_llm/runtime/utils/mpiUtils.h" @@ -57,9 +58,10 @@ class NcclCommunicator private: void send( - void const* sendbuff, size_t count, nvinfer1::DataType dataType, int peer, CudaStream const& stream) const; + void const* sendbuff, size_t count, tensorrt_llm::DataType dataType, int peer, CudaStream const& stream) const; - void receive(void* sendbuff, size_t count, nvinfer1::DataType dataType, int peer, CudaStream const& stream) const; + void receive( + void* sendbuff, size_t count, tensorrt_llm::DataType dataType, int peer, CudaStream const& stream) const; static ncclComm_t createComm(int worldSize, int rank, mpi::MpiComm const& mpiComm); diff --git a/cpp/tensorrt_llm/runtime/runtimeKernels.cu b/cpp/tensorrt_llm/runtime/runtimeKernels.cu index 3b3dbcac894a..b22d36052370 100644 --- a/cpp/tensorrt_llm/runtime/runtimeKernels.cu +++ b/cpp/tensorrt_llm/runtime/runtimeKernels.cu @@ -21,7 +21,7 @@ #include "tensorrt_llm/kernels/speculativeDecoding/kvCacheUpdateKernels.h" #include "tensorrt_llm/runtime/runtimeKernels.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include using namespace tensorrt_llm::runtime; @@ -333,13 +333,13 @@ void invokeFillBatch(IBuffer& buffer, IBuffer const& slotIndices, std::size_t sl { switch (buffer.getDataType()) { - case nvinfer1::DataType::kINT32: + case tensorrt_llm::DataType::kINT32: invokeFillBatch(buffer, slotIndices, slotStride, values, stream); break; - case nvinfer1::DataType::kINT8: + case tensorrt_llm::DataType::kINT8: invokeFillBatch(buffer, slotIndices, slotStride, values, stream); break; - case nvinfer1::DataType::kFLOAT: invokeFillBatch(buffer, slotIndices, slotStride, values, stream); break; + case tensorrt_llm::DataType::kFLOAT: invokeFillBatch(buffer, slotIndices, slotStride, values, stream); break; default: TLLM_THROW("data type not supported"); } } @@ -349,13 +349,15 @@ void invokeGatherBatch(IBuffer& buffer, IBuffer const& values, IBuffer const& sl { switch (buffer.getDataType()) { - case nvinfer1::DataType::kINT32: + case tensorrt_llm::DataType::kINT32: invokeGatherBatch(buffer, values, slotIndices, slotStride, stream); break; - case nvinfer1::DataType::kINT8: + case tensorrt_llm::DataType::kINT8: invokeGatherBatch(buffer, values, slotIndices, slotStride, stream); break; - case nvinfer1::DataType::kFLOAT: invokeGatherBatch(buffer, values, slotIndices, slotStride, stream); break; + case tensorrt_llm::DataType::kFLOAT: + invokeGatherBatch(buffer, values, slotIndices, slotStride, stream); + break; default: TLLM_THROW("data type not supported"); } } @@ -408,12 +410,12 @@ void scatterTensor(ITensor& output, ITensor const& input, SizeType32 beamWidth, { switch (input.getDataType()) { - case nvinfer1::DataType::kINT32: invokeScatterTensor(output, input, beamWidth, stream); break; - case nvinfer1::DataType::kFLOAT: invokeScatterTensor(output, input, beamWidth, stream); break; - case nvinfer1::DataType::kHALF: invokeScatterTensor(output, input, beamWidth, stream); break; - case nvinfer1::DataType::kINT8: invokeScatterTensor(output, input, beamWidth, stream); break; + case tensorrt_llm::DataType::kINT32: invokeScatterTensor(output, input, beamWidth, stream); break; + case tensorrt_llm::DataType::kFLOAT: invokeScatterTensor(output, input, beamWidth, stream); break; + case tensorrt_llm::DataType::kHALF: invokeScatterTensor(output, input, beamWidth, stream); break; + case tensorrt_llm::DataType::kINT8: invokeScatterTensor(output, input, beamWidth, stream); break; #ifdef ENABLE_FP8 - case nvinfer1::DataType::kFP8: invokeScatterTensor<__nv_fp8_e4m3>(output, input, beamWidth, stream); break; + case tensorrt_llm::DataType::kFP8: invokeScatterTensor<__nv_fp8_e4m3>(output, input, beamWidth, stream); break; #endif // ENABLE_FP8 default: TLLM_THROW("data type not supported"); } @@ -423,15 +425,15 @@ void tileTensor(ITensor& output, ITensor const& input, SizeType32 beamWidth, Cud { switch (input.getDataType()) { - case nvinfer1::DataType::kINT32: invokeTileTensor(output, input, beamWidth, stream); break; - case nvinfer1::DataType::kFLOAT: invokeTileTensor(output, input, beamWidth, stream); break; - case nvinfer1::DataType::kHALF: invokeTileTensor(output, input, beamWidth, stream); break; + case tensorrt_llm::DataType::kINT32: invokeTileTensor(output, input, beamWidth, stream); break; + case tensorrt_llm::DataType::kFLOAT: invokeTileTensor(output, input, beamWidth, stream); break; + case tensorrt_llm::DataType::kHALF: invokeTileTensor(output, input, beamWidth, stream); break; #ifdef ENABLE_BF16 - case nvinfer1::DataType::kBF16: invokeTileTensor<__nv_bfloat16>(output, input, beamWidth, stream); break; + case tensorrt_llm::DataType::kBF16: invokeTileTensor<__nv_bfloat16>(output, input, beamWidth, stream); break; #endif // ENABLE_BF16 - case nvinfer1::DataType::kINT8: invokeTileTensor(output, input, beamWidth, stream); break; + case tensorrt_llm::DataType::kINT8: invokeTileTensor(output, input, beamWidth, stream); break; #ifdef ENABLE_FP8 - case nvinfer1::DataType::kFP8: invokeTileTensor<__nv_fp8_e4m3>(output, input, beamWidth, stream); break; + case tensorrt_llm::DataType::kFP8: invokeTileTensor<__nv_fp8_e4m3>(output, input, beamWidth, stream); break; #endif // ENABLE_FP8 default: TLLM_THROW("data type not supported"); } @@ -444,22 +446,22 @@ void mergeLogitsFragments(BufferManager const& bufferManager, ITensor& output, { switch (output.getDataType()) { - case nvinfer1::DataType::kFLOAT: + case tensorrt_llm::DataType::kFLOAT: invokeMergeLogitsFragments(bufferManager, output, fragmentsVector, cachePointerDevice, cachePointerHost, firstBatchSlotIdx, microBatchSize, beamWidth, stream, stepOffset); break; - case nvinfer1::DataType::kHALF: + case tensorrt_llm::DataType::kHALF: invokeMergeLogitsFragments(bufferManager, output, fragmentsVector, cachePointerDevice, cachePointerHost, firstBatchSlotIdx, microBatchSize, beamWidth, stream, stepOffset); break; #ifdef ENABLE_BF16 - case nvinfer1::DataType::kBF16: + case tensorrt_llm::DataType::kBF16: invokeMergeLogitsFragments<__nv_bfloat16>(bufferManager, output, fragmentsVector, cachePointerDevice, cachePointerHost, firstBatchSlotIdx, microBatchSize, beamWidth, stream, stepOffset); break; #endif // ENABLE_BF16 #ifdef ENABLE_FP8 - case nvinfer1::DataType::kFP8: + case tensorrt_llm::DataType::kFP8: invokeMergeLogitsFragments<__nv_fp8_e4m3>(bufferManager, output, fragmentsVector, cachePointerDevice, cachePointerHost, firstBatchSlotIdx, microBatchSize, beamWidth, stream, stepOffset); break; diff --git a/cpp/tensorrt_llm/runtime/tensorView.h b/cpp/tensorrt_llm/runtime/tensorView.h index 17e7fb719415..d9e65b0efe23 100644 --- a/cpp/tensorrt_llm/runtime/tensorView.h +++ b/cpp/tensorrt_llm/runtime/tensorView.h @@ -16,6 +16,7 @@ #pragma once +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/runtime/bufferView.h" #include "tensorrt_llm/runtime/iTensor.h" @@ -45,19 +46,19 @@ class TensorView : virtual public ITensor, public BufferView mDims.d[0] = size; } - TensorView(IBuffer::SharedPtr const& buffer, size_t offset, size_t size, nvinfer1::Dims const& dims) + TensorView(IBuffer::SharedPtr const& buffer, size_t offset, size_t size, tensorrt_llm::Dims const& dims) : BufferView{buffer, offset, size} , mDims{dims} { Base::resize(ITensor::volumeNonNegative(dims)); } - [[nodiscard]] nvinfer1::Dims const& getShape() const override + [[nodiscard]] tensorrt_llm::Dims const& getShape() const override { return mDims; } - void reshape(nvinfer1::Dims const& dims) override + void reshape(tensorrt_llm::Dims const& dims) override { Base::resize(ITensor::volumeNonNegative(dims)); mDims = dims; @@ -81,6 +82,6 @@ class TensorView : virtual public ITensor, public BufferView return shape.nbDims > 0 && shape.d[0] > 0 ? ITensor::volume(shape) / shape.d[0] : 0; } - nvinfer1::Dims mDims{}; + tensorrt_llm::Dims mDims{}; }; } // namespace tensorrt_llm::runtime diff --git a/cpp/tensorrt_llm/runtime/tllmBuffers.cpp b/cpp/tensorrt_llm/runtime/tllmBuffers.cpp index ff7ed04001d3..4876d5b87bf6 100644 --- a/cpp/tensorrt_llm/runtime/tllmBuffers.cpp +++ b/cpp/tensorrt_llm/runtime/tllmBuffers.cpp @@ -15,6 +15,7 @@ */ #include "tensorrt_llm/runtime/tllmBuffers.h" +#include "tensorrt_llm/common/tllmDataType.h" namespace tensorrt_llm::runtime { @@ -62,12 +63,12 @@ std::shared_ptr MulticastTensorView::lock() const /////////////////////////////////////// // MulticastTensorView ITensor methods /////////////////////////////////////// -nvinfer1::Dims const& MulticastTensorView::getShape() const +tensorrt_llm::Dims const& MulticastTensorView::getShape() const { return mDims; } -void MulticastTensorView::reshape(nvinfer1::Dims const& dims) +void MulticastTensorView::reshape(tensorrt_llm::Dims const& dims) { auto new_size = nonNegative(volume(dims)); if (new_size > getCapacity()) @@ -102,7 +103,7 @@ std::size_t MulticastTensorView::getCapacity() const return lock()->getCapacity(); } -nvinfer1::DataType MulticastTensorView::getDataType() const +tensorrt_llm::DataType MulticastTensorView::getDataType() const { return lock()->getDataType(); } diff --git a/cpp/tensorrt_llm/runtime/tllmBuffers.h b/cpp/tensorrt_llm/runtime/tllmBuffers.h index faed36537e5c..d023823de5b2 100644 --- a/cpp/tensorrt_llm/runtime/tllmBuffers.h +++ b/cpp/tensorrt_llm/runtime/tllmBuffers.h @@ -27,7 +27,7 @@ #include "tensorrt_llm/runtime/memoryCounters.h" #include "tensorrt_llm/runtime/virtualMemory.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include #include @@ -550,7 +550,7 @@ class GenericBuffer : virtual public IBuffer, TAllocator // Inherit from TAlloca //! //! \brief Construct an empty buffer. //! - explicit GenericBuffer(nvinfer1::DataType type, TAllocator allocator = {}) // NOLINT(*-pro-type-member-init) + explicit GenericBuffer(tensorrt_llm::DataType type, TAllocator allocator = {}) // NOLINT(*-pro-type-member-init) : GenericBuffer{0, type, std::move(allocator)} { } @@ -559,7 +559,7 @@ class GenericBuffer : virtual public IBuffer, TAllocator // Inherit from TAlloca //! \brief Construct a buffer with the specified allocation size in number of elements. //! explicit GenericBuffer( // NOLINT(*-pro-type-member-init) - std::size_t size, nvinfer1::DataType type, TAllocator allocator = {}) + std::size_t size, tensorrt_llm::DataType type, TAllocator allocator = {}) : GenericBuffer{size, size, type, std::move(allocator)} { } @@ -636,7 +636,7 @@ class GenericBuffer : virtual public IBuffer, TAllocator // Inherit from TAlloca //! //! \brief Returns the type of the buffer. //! - [[nodiscard]] nvinfer1::DataType getDataType() const override + [[nodiscard]] tensorrt_llm::DataType getDataType() const override { return mType; } @@ -687,7 +687,8 @@ class GenericBuffer : virtual public IBuffer, TAllocator // Inherit from TAlloca } protected: - explicit GenericBuffer(std::size_t size, std::size_t capacity, nvinfer1::DataType type, TAllocator allocator = {}) + explicit GenericBuffer( + std::size_t size, std::size_t capacity, tensorrt_llm::DataType type, TAllocator allocator = {}) : TAllocator{std::move(allocator)} , mSize{size} , mCapacity{capacity} @@ -700,14 +701,14 @@ class GenericBuffer : virtual public IBuffer, TAllocator // Inherit from TAlloca private: std::size_t mSize{0}, mCapacity{0}; - nvinfer1::DataType mType; + tensorrt_llm::DataType mType; void* mBuffer; }; class MulticastBuffer : virtual public IBuffer { public: - explicit MulticastBuffer(nvinfer1::DataType type, std::set const& ranks) + explicit MulticastBuffer(tensorrt_llm::DataType type, std::set const& ranks) : mSize(0) , mCapacity(0) , mType(type) @@ -716,7 +717,7 @@ class MulticastBuffer : virtual public IBuffer TLLM_CHECK(ranks.size() > 1); } - explicit MulticastBuffer(size_t size, nvinfer1::DataType type, std::set const& ranks) + explicit MulticastBuffer(size_t size, tensorrt_llm::DataType type, std::set const& ranks) : mSize(0) , mCapacity(0) , mType(type) @@ -817,7 +818,7 @@ class MulticastBuffer : virtual public IBuffer return mCapacity; } - [[nodiscard]] nvinfer1::DataType getDataType() const override + [[nodiscard]] tensorrt_llm::DataType getDataType() const override { return mType; } @@ -853,7 +854,7 @@ class MulticastBuffer : virtual public IBuffer private: std::size_t mSize = 0; std::size_t mCapacity = 0; - nvinfer1::DataType mType; + tensorrt_llm::DataType mType; std::set mRanks; IpcNvlsHandle* mHandle; }; @@ -882,7 +883,7 @@ class GenericTensor : virtual public ITensor, public GenericBuffer //! //! \brief Construct an empty tensor. //! - explicit GenericTensor(nvinfer1::DataType type, TAllocator allocator = {}) + explicit GenericTensor(tensorrt_llm::DataType type, TAllocator allocator = {}) : Base{type, std::move(allocator)} { mDims.nbDims = 0; @@ -891,14 +892,14 @@ class GenericTensor : virtual public ITensor, public GenericBuffer //! //! \brief Construct a tensor with the specified allocation dimensions. //! - explicit GenericTensor(nvinfer1::Dims const& dims, nvinfer1::DataType type, TAllocator allocator = {}) + explicit GenericTensor(tensorrt_llm::Dims const& dims, tensorrt_llm::DataType type, TAllocator allocator = {}) : Base{nonNegative(volume(dims)), type, std::move(allocator)} , mDims{dims} { } explicit GenericTensor( - nvinfer1::Dims const& dims, std::size_t capacity, nvinfer1::DataType type, TAllocator allocator = {}) + tensorrt_llm::Dims const& dims, std::size_t capacity, tensorrt_llm::DataType type, TAllocator allocator = {}) : Base{nonNegative(volume(dims)), capacity, type, std::move(allocator)} , mDims{dims} { @@ -923,12 +924,12 @@ class GenericTensor : virtual public ITensor, public GenericBuffer return *this; } - [[nodiscard]] nvinfer1::Dims const& getShape() const override + [[nodiscard]] tensorrt_llm::Dims const& getShape() const override { return mDims; } - void reshape(nvinfer1::Dims const& dims) override + void reshape(tensorrt_llm::Dims const& dims) override { Base::resize(nonNegative(volume(dims))); mDims = dims; @@ -946,7 +947,7 @@ class GenericTensor : virtual public ITensor, public GenericBuffer } private: - nvinfer1::Dims mDims{}; + tensorrt_llm::Dims mDims{}; }; // Forward declaration @@ -971,9 +972,9 @@ class MulticastTensorView : virtual public ITensor ///////////////////// // ITensor methods ///////////////////// - [[nodiscard]] nvinfer1::Dims const& getShape() const override; + [[nodiscard]] tensorrt_llm::Dims const& getShape() const override; - void reshape(nvinfer1::Dims const& dims) override; + void reshape(tensorrt_llm::Dims const& dims) override; ///////////////////// // IBuffer methods @@ -983,7 +984,7 @@ class MulticastTensorView : virtual public ITensor [[nodiscard]] std::size_t getCapacity() const override; - [[nodiscard]] nvinfer1::DataType getDataType() const override; + [[nodiscard]] tensorrt_llm::DataType getDataType() const override; [[nodiscard]] MemoryType getMemoryType() const override; @@ -1016,7 +1017,7 @@ class MulticastTensorView : virtual public ITensor std::weak_ptr mTensor; ViewType mViewType; - nvinfer1::Dims mDims{}; + tensorrt_llm::Dims mDims{}; }; class MulticastTensor : virtual public ITensor, @@ -1026,13 +1027,13 @@ class MulticastTensor : virtual public ITensor, public: using Base = MulticastBuffer; - explicit MulticastTensor(nvinfer1::DataType type, std::set const& ranks) + explicit MulticastTensor(tensorrt_llm::DataType type, std::set const& ranks) : Base(type, ranks) { mDims.nbDims = 0; } - explicit MulticastTensor(nvinfer1::Dims const& dims, nvinfer1::DataType type, std::set const& ranks) + explicit MulticastTensor(tensorrt_llm::Dims const& dims, tensorrt_llm::DataType type, std::set const& ranks) : Base(nonNegative(volume(dims)), type, ranks) , mDims(dims) { @@ -1068,12 +1069,12 @@ class MulticastTensor : virtual public ITensor, ///////////////////// // ITensor methods ///////////////////// - [[nodiscard]] nvinfer1::Dims const& getShape() const override + [[nodiscard]] tensorrt_llm::Dims const& getShape() const override { return mDims; } - void reshape(nvinfer1::Dims const& dims) override + void reshape(tensorrt_llm::Dims const& dims) override { Base::resize(nonNegative(volume(dims))); mDims = dims; @@ -1091,7 +1092,7 @@ class MulticastTensor : virtual public ITensor, } private: - nvinfer1::Dims mDims{}; + tensorrt_llm::Dims mDims{}; }; using DeviceTensor = GenericTensor; diff --git a/cpp/tensorrt_llm/runtime/tllmLogger.cpp b/cpp/tensorrt_llm/runtime/tllmLogger.cpp deleted file mode 100644 index 586ab2f4ae95..000000000000 --- a/cpp/tensorrt_llm/runtime/tllmLogger.cpp +++ /dev/null @@ -1,62 +0,0 @@ -/* - * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. - * - * 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. - */ -#include "tllmLogger.h" -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/logger.h" - -using namespace tensorrt_llm::runtime; -namespace tc = tensorrt_llm::common; - -void TllmLogger::log(nvinfer1::ILogger::Severity severity, nvinfer1::AsciiChar const* msg) noexcept -{ - switch (severity) - { - case nvinfer1::ILogger::Severity::kINTERNAL_ERROR: - case nvinfer1::ILogger::Severity::kERROR: TLLM_LOG_ERROR(msg); break; - case nvinfer1::ILogger::Severity::kWARNING: TLLM_LOG_WARNING(msg); break; - case nvinfer1::ILogger::Severity::kINFO: TLLM_LOG_INFO(msg); break; - case nvinfer1::ILogger::Severity::kVERBOSE: TLLM_LOG_DEBUG(msg); break; - default: TLLM_LOG_TRACE(msg); break; - } -} - -nvinfer1::ILogger::Severity TllmLogger::getLevel() -{ - auto* const logger = tc::Logger::getLogger(); - switch (logger->getLevel()) - { - case tc::Logger::Level::ERROR: return nvinfer1::ILogger::Severity::kERROR; - case tc::Logger::Level::WARNING: return nvinfer1::ILogger::Severity::kWARNING; - case tc::Logger::Level::INFO: return nvinfer1::ILogger::Severity::kINFO; - case tc::Logger::Level::DEBUG: - case tc::Logger::Level::TRACE: return nvinfer1::ILogger::Severity::kVERBOSE; - default: return nvinfer1::ILogger::Severity::kINTERNAL_ERROR; - } -} - -void TllmLogger::setLevel(nvinfer1::ILogger::Severity level) -{ - auto* const logger = tc::Logger::getLogger(); - switch (level) - { - case nvinfer1::ILogger::Severity::kINTERNAL_ERROR: - case nvinfer1::ILogger::Severity::kERROR: logger->setLevel(tc::Logger::Level::ERROR); break; - case nvinfer1::ILogger::Severity::kWARNING: logger->setLevel(tc::Logger::Level::WARNING); break; - case nvinfer1::ILogger::Severity::kINFO: logger->setLevel(tc::Logger::Level::INFO); break; - case nvinfer1::ILogger::Severity::kVERBOSE: logger->setLevel(tc::Logger::Level::TRACE); break; - default: TLLM_THROW("Unsupported severity"); - } -} diff --git a/cpp/tensorrt_llm/runtime/tllmRuntime.cpp b/cpp/tensorrt_llm/runtime/tllmRuntime.cpp deleted file mode 100644 index 7c2ca4747213..000000000000 --- a/cpp/tensorrt_llm/runtime/tllmRuntime.cpp +++ /dev/null @@ -1,831 +0,0 @@ -/* - * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. - * - * 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. - */ -#include "tllmRuntime.h" -#include "common.h" -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/logger.h" -#include "tensorrt_llm/common/nvtxUtils.h" -#include "tensorrt_llm/common/safetensors.h" -#include "tensorrt_llm/executor/tensor.h" -#include "tensorrt_llm/kernels/userbuffers/ub_interface.h" -#include "tensorrt_llm/runtime/utils/mpiUtils.h" -#include "tllmLogger.h" -#include "tllmStreamReaders.h" - -#include "nlohmann/json.hpp" -#include - -#include -#include -#include -#include -#include -#include -#include -#include -#include - -using namespace tensorrt_llm::runtime; -using TensorMap = StringPtrMap; - -namespace -{ -static_assert(std::is_signed::value, "SizeType32 must be signed"); - -nvinfer1::Dims shapeToDims(std::vector const& shape) -{ - TLLM_CHECK(shape.size() <= nvinfer1::Dims::MAX_DIMS); - nvinfer1::Dims dims; - auto constexpr dim_max = std::numeric_limits::max(); - dims.nbDims = static_cast(shape.size()); - for (std::size_t i = 0; i < shape.size(); ++i) - { - // shape[i] >= 0 because it has unsigned type. Check upper bound: - TLLM_CHECK(shape[i] <= static_cast(dim_max)); - dims.d[i] = static_cast(shape[i]); - } - return dims; -} - -std::vector dimsToShape(nvinfer1::Dims const& dims) -{ - TLLM_CHECK(dims.nbDims >= 0); - std::vector shape(dims.nbDims); - for (std::int32_t i = 0; i < dims.nbDims; ++i) - { - TLLM_CHECK(dims.d[i] >= 0); - shape[i] = static_cast(dims.d[i]); - } - return shape; -} - -tensorrt_llm::runtime::TllmLogger defaultLogger{}; - -void setWeightStreaming(nvinfer1::ICudaEngine& engine, float const gpuWeightsPercent) -{ - if (gpuWeightsPercent < 1) - { - int64_t streamableSize = engine.getStreamableWeightsSize(); - int64_t budget = gpuWeightsPercent * streamableSize; - TLLM_LOG_INFO("Set gpu weights percent to %f, which is %lld bytes. Valid range: %lld bytes - %lld bytes.", - gpuWeightsPercent, budget, 0, streamableSize); - engine.setWeightStreamingBudgetV2(budget); - } -} - -class LayerInfo -{ -public: - LayerInfo(std::optional name, std::string type) - : name(std::move(name)) - , type(std::move(type)){}; - std::optional name; - std::string type; -}; - -void assessLikelihoodOfRuntimeAllocation( - nvinfer1::ICudaEngine const& engine, nvinfer1::IEngineInspector const& engineInspector) - -{ - TLLM_LOG_INFO("Inspecting the engine to identify potential runtime issues..."); - auto const profilingVerbosity = engine.getProfilingVerbosity(); - if (profilingVerbosity != nvinfer1::ProfilingVerbosity::kDETAILED) - { - TLLM_LOG_INFO( - "The profiling verbosity of the engine does not allow this analysis to proceed. Re-build the engine with " - "'detailed' profiling verbosity to get more diagnostics."); - return; - } - auto const* const layerTypeKey = "LayerType"; - auto const* const nameKey = "Name"; - auto const numLayers = engine.getNbLayers(); - TLLM_LOG_INFO("Model has %i layers.", numLayers); - std::vector indexes(numLayers); - std::iota(indexes.begin(), indexes.end(), 0); - std::vector> layerInfos(numLayers); - std::transform(indexes.cbegin(), indexes.cend(), layerInfos.begin(), - [&](SizeType32 const idx) - { - auto const* const layerInfo - = engineInspector.getLayerInformation(idx, nvinfer1::LayerInformationFormat::kJSON); - - // Needs to be copied explicitly, see documentation of `getLayerInformation`. - auto const layerInfoCopy = std::string(layerInfo); - auto const jsonLayerInfo = nlohmann::json::parse(layerInfoCopy); - auto const layerJsonType = jsonLayerInfo.type(); - if (layerJsonType != nlohmann::detail::value_t::object) - { - return std::optional{}; - } - if (!jsonLayerInfo.contains(layerTypeKey)) - { - return std::optional{}; - } - auto const& typeJson = jsonLayerInfo.at(layerTypeKey); - if (typeJson.type() != nlohmann::detail::value_t::string) - { - return std::optional{}; - } - std::optional name{}; - if (jsonLayerInfo.contains(nameKey)) - { - auto const& nameJson = jsonLayerInfo.at(nameKey); - auto const nameJsonType = nameJson.type(); - if (nameJsonType == nlohmann::detail::value_t::string) - { - name = nameJson.get(); - } - } - return std::make_optional(LayerInfo{name, typeJson.get()}); - }); - auto const layersWithInfoEnd = std::partition( - layerInfos.begin(), layerInfos.end(), [](std::optional const& info) { return info.has_value(); }); - if (layersWithInfoEnd == layerInfos.begin()) - { - TLLM_LOG_INFO("Engine layer infos could not be parsed into useful information."); - return; - } - auto const allocateLayersEnd = std::partition(layerInfos.begin(), layersWithInfoEnd, - [](std::optional const& info) { return info.value().type == "allocate"; }); - auto numWarnings = 0; - for (auto layerInfo = layerInfos.begin(); layerInfo != allocateLayersEnd; layerInfo++) - { - auto constexpr maxNumWarnings = 25; - if (numWarnings < maxNumWarnings) - { - auto const layerName = layerInfo->value().name.value_or(""); - TLLM_LOG_WARNING( - "Layer '%s' has type '%s', which could lead to large runtime memory allocations. Performance " - "might be degraded and / or you might run out of memory.", - layerName.c_str(), layerInfo->value().type.c_str()); - } - numWarnings++; - } - if (numWarnings > 0) - { - TLLM_LOG_WARNING( - "There were a total of %i layers with type 'allocate'. Some warnings might have been silenced to keep the " - "output concise.", - numWarnings); - } -} - -} // namespace - -TllmRuntime::TllmRuntime(RawEngine const& rawEngine, nvinfer1::ILogger* logger, bool useGpuDirectStorage, - float gpuWeightsPercent, bool useShapeInference) - : mStream(std::make_shared()) - , mBufferManager{mStream, true} // Ensure to trim the memory pool on destruction. - , mRuntime{nvinfer1::createInferRuntime(static_cast(logger) ? *logger : defaultLogger)} - , mUseShapeInference{useShapeInference} - , mUserBufferEnabled{false} -{ - auto const startTime = std::chrono::high_resolution_clock::now(); - - switch (rawEngine.getType()) - { - case RawEngine::Type::FilePath: - { - if (useGpuDirectStorage) - { - TLLM_LOG_INFO("GDS is used to load the engine!"); - auto reader = GDSStreamReader(rawEngine.getPath()); - mEngine.reset(mRuntime->deserializeCudaEngine(reader)); - } - else - { - auto reader = StreamReader(rawEngine.getPath()); - mEngine.reset(mRuntime->deserializeCudaEngine(reader)); - } - break; - } - case RawEngine::Type::AddressWithSize: - mEngine.reset(mRuntime->deserializeCudaEngine(rawEngine.getAddress(), rawEngine.getSize())); - break; - case RawEngine::Type::HostMemory: - mEngine.reset( - mRuntime->deserializeCudaEngine(rawEngine.getHostMemory()->data(), rawEngine.getHostMemory()->size())); - break; - default: TLLM_THROW("Unsupported raw engine type."); - } - - auto const elapsedMs - = std::chrono::duration_cast(std::chrono::high_resolution_clock::now() - startTime); - - TLLM_LOG_INFO("Engine load time %lld ms", elapsedMs); - - TLLM_CHECK_WITH_INFO(mEngine != nullptr, "Failed to deserialize cuda engine."); - mEngineInspector.reset(mEngine->createEngineInspector()); - assessLikelihoodOfRuntimeAllocation(*mEngine, *mEngineInspector); - setWeightStreaming(getEngine(), gpuWeightsPercent); - auto const devMemorySize = mEngine->getDeviceMemorySizeV2(); - mEngineBuffer = mBufferManager.gpu(devMemorySize); - // Print context memory size for CI/CD to track. - TLLM_LOG_INFO("[MemUsageChange] Allocated %.2f MiB for execution context memory.", - static_cast(devMemorySize) / 1048576.0); - - cacheTensorNames(); -} - -void TllmRuntime::cacheTensorNames() -{ - for (std::int32_t i = 0; i < mEngine->getNbIOTensors(); ++i) - { - auto const* const name = mEngine->getIOTensorName(i); - if (mEngine->getTensorIOMode(name) == nvinfer1::TensorIOMode::kINPUT) - { - mInputTensorNames.emplace_back(name); - } - else if (mEngine->getTensorIOMode(name) == nvinfer1::TensorIOMode::kOUTPUT) - { - mOutputTensorNames.emplace_back(name); - } - } -} - -nvinfer1::IExecutionContext& TllmRuntime::addContext(std::int32_t profileIndex) -{ - TLLM_CHECK(0 <= profileIndex && profileIndex < mEngine->getNbOptimizationProfiles()); - mContexts.emplace_back(mEngine->createExecutionContextWithoutDeviceMemory()); - if (!mContexts.back()) - { - if (mEngine->getStreamableWeightsSize() > 0) - { - TLLM_THROW("Failed to allocate memory for weights. Please try reducing --gpu_weights_percent."); - } - else - { - TLLM_THROW("Internal Error: Failed to create an execution context."); - } - } - auto& context = *mContexts.back(); - context.setDeviceMemoryV2(mEngineBuffer->data(), static_cast(mEngineBuffer->getCapacity())); - - if (tensorrt_llm::common::Logger::getLogger()->isEnabled(tensorrt_llm::common::Logger::TRACE) - && mContexts.size() == 1) - { - // Print engine information only once - printEngineInfo(); - } - - context.setOptimizationProfileAsync(profileIndex, mStream->get()); - // If nvtx verbosity is DETAILED, print an info about potential perf overhead. - if (context.getNvtxVerbosity() == nvinfer1::ProfilingVerbosity::kDETAILED) - { - TLLM_LOG_INFO( - "The engine was built with kDETAILED profiling verbosity, which may result in small overheads at runtime."); - } - return context; -} - -void TllmRuntime::printEngineInfo() -{ - auto& context = *(mContexts[0]); - int const nIO = mEngine->getNbIOTensors(); // Count of input / output tensor - int const nOP = mEngine->getNbOptimizationProfiles(); // Count of Optimization Profile - std::size_t maxNameWidth = 0; - std::size_t maxShapeWidth = 0; - - // Get information of engine input / output - std::vector tensorNameList{}; - tensorNameList.reserve(nIO); - for (int i = 0; i < nIO; ++i) - { - tensorNameList.emplace_back(mEngine->getIOTensorName(i)); - } - std::vector> tensorInfo(nIO); // Tensor Information Vector - std::vector>> profileInfo(nIO); // Tensor Optimization Profile Vector - for (int i = 0; i < nIO; ++i) - { - auto const& name = tensorNameList[i]; - char const* nameC{name.c_str()}; // name of C-style - maxNameWidth = std::max(maxNameWidth, name.size()); - tensorInfo[i]["mode"] = mEngine->getTensorIOMode(nameC) == nvinfer1::TensorIOMode::kINPUT ? "I" : "O"; - tensorInfo[i]["location"] - = mEngine->getTensorLocation(nameC) == nvinfer1::TensorLocation::kDEVICE ? "GPU" : "CPU"; - tensorInfo[i]["data_type"] = dataTypeToString(mEngine->getTensorDataType(nameC)); - tensorInfo[i]["build_shape"] = shapeToString(mEngine->getTensorShape(nameC)); - maxShapeWidth = std::max(maxShapeWidth, tensorInfo[i]["build_shape"].size()); - if (tensorInfo[i]["mode"] == "I") - { - std::vector> topPerTensor(nOP); - for (int k = 0; k < nOP; ++k) - { - if (tensorInfo[i]["location"] == std::string("GPU")) - { - std::vector top(3); - top[0] = mEngine->getProfileShape(nameC, k, nvinfer1::OptProfileSelector::kMIN); - top[1] = mEngine->getProfileShape(nameC, k, nvinfer1::OptProfileSelector::kOPT); - top[2] = mEngine->getProfileShape(nameC, k, nvinfer1::OptProfileSelector::kMAX); - topPerTensor[k] = top; - maxShapeWidth = std::max(maxShapeWidth, shapeToString(top[2]).size()); - } - else - { - // Shape input tensor, not used in TRT-LLM support yet - std::vector top(3); - int const nDim = mEngine->getTensorShape(nameC).nbDims; - nvinfer1::Dims64 tensorShape{nDim, {-1}}; - int const* pos = nullptr; - pos = mEngine->getProfileTensorValues(nameC, k, nvinfer1::OptProfileSelector::kMIN); - std::copy(pos, pos + nDim, tensorShape.d); - top[0] = tensorShape; - pos = mEngine->getProfileTensorValues(nameC, k, nvinfer1::OptProfileSelector::kOPT); - std::copy(pos, pos + nDim, tensorShape.d); - top[1] = tensorShape; - pos = mEngine->getProfileTensorValues(nameC, k, nvinfer1::OptProfileSelector::kMAX); - std::copy(pos, pos + nDim, tensorShape.d); - top[2] = tensorShape; - topPerTensor[k] = top; - } - } - profileInfo[i] = topPerTensor; - } - else - { - profileInfo[i] = std::vector>(nOP); - } - } - // Set input shape to get output shape - for (int k = 0; k < nOP; ++k) - { - for (int j = 0; j < 3; ++j) // Min, Opt, Max - { - for (int i = 0; i < nIO; ++i) - { - auto const& name = tensorNameList[i]; - char const* nameC = name.c_str(); - if (tensorInfo[i]["mode"] == "I") - { - if (tensorInfo[i]["location"] == std::string("GPU")) - { - context.setInputShape(nameC, profileInfo[i][k][j]); - } - else - { - // Shape input tensor, not used in TRT-LLM support yet - context.setInputTensorAddress(nameC, profileInfo[i][k][j].d); - } - } - else - { - TLLM_CHECK_WITH_INFO(context.allInputDimensionsSpecified(), "Input dimensions not specified"); - TLLM_CHECK_WITH_INFO(context.allInputShapesSpecified(), "Input shapes not specified"); - if (tensorInfo[i]["location"] == std::string("GPU")) - { - profileInfo[i][k].push_back(context.getTensorShape(nameC)); - } - else - { - // Shape input tensor, not used in TRT-LLM support yet - int const nDim = mEngine->getTensorShape(nameC).nbDims; - nvinfer1::Dims64 tensorShape{nDim, {}}; - int const* pos = reinterpret_cast(context.getTensorAddress(nameC)); - std::copy(pos, pos + nDim, tensorShape.d); - profileInfo[i][k].push_back(tensorShape); - } - } - } - } - } - - // Print information of engine input / output - std::string info; - TLLM_LOG_TRACE("Information of engine input / output."); - TLLM_LOG_TRACE(std::string(maxNameWidth + maxShapeWidth + 24, '=')); - info = alignText("Name", maxNameWidth) + "|I/O|Location|DataType|" + alignText("Shape", maxShapeWidth) + "|"; - TLLM_LOG_TRACE(info.c_str()); - TLLM_LOG_TRACE(std::string(maxNameWidth + maxShapeWidth + 24, '-')); - for (int i = 0; i < nIO; ++i) - { - info = alignText(tensorNameList[i], maxNameWidth, false) + "|"; - info += alignText(tensorInfo[i]["mode"], 3) + "|"; - info += alignText(tensorInfo[i]["location"], 8) + "|"; - info += alignText(tensorInfo[i]["data_type"], 8) + "|"; - info += alignText(tensorInfo[i]["build_shape"], maxShapeWidth) + "|"; - TLLM_LOG_TRACE(info.c_str()); - } - TLLM_LOG_TRACE(std::string(maxNameWidth + maxShapeWidth + 24, '=')); - // Print information of optimization profile - TLLM_LOG_TRACE("Information of optimization profile."); - for (int k = 0; k < nOP; ++k) - { - TLLM_LOG_TRACE("Optimization Profile %d:", k); - TLLM_LOG_TRACE(std::string(maxNameWidth + maxShapeWidth * 3 + 4, '=')); - info = alignText("Name", maxNameWidth) + "|"; - info += alignText("Min", maxShapeWidth) + "|"; - info += alignText("Opt", maxShapeWidth) + "|"; - info += alignText("Max", maxShapeWidth) + "|"; - TLLM_LOG_TRACE(info.c_str()); - TLLM_LOG_TRACE(std::string(maxNameWidth + maxShapeWidth * 3 + 4, '-')); - for (int i = 0; i < nIO; ++i) - { - auto const& top = profileInfo[i][k]; - info = alignText(tensorNameList[i], maxNameWidth, false) + "|"; - info += alignText(shapeToString(top[0]), maxShapeWidth) + "|"; - info += alignText(shapeToString(top[1]), maxShapeWidth) + "|"; - info += alignText(shapeToString(top[2]), maxShapeWidth) + "|"; - TLLM_LOG_TRACE(info.c_str()); - } - TLLM_LOG_TRACE(std::string(maxNameWidth + maxShapeWidth * 3 + 4, '=')); - } -} - -void TllmRuntime::printContextInfo(SizeType32 contextIndex) -{ - auto const& context = *(mContexts[contextIndex]); - int const nIO = mEngine->getNbIOTensors(); // Count of input / output tensor - std::size_t maxNameWidth = 0; - std::size_t maxShapeWidth = 0; - std::vector> tensorInfo(nIO); - for (int i = 0; i < nIO; ++i) - { - auto const name = std::string(mEngine->getIOTensorName(i)); - bool const isInput = mEngine->getTensorIOMode(name.c_str()) == nvinfer1::TensorIOMode::kINPUT; - auto const shape = shapeToString(context.getTensorShape(name.c_str())); - tensorInfo[i] = std::make_tuple(name, isInput, shape); - maxNameWidth = std::max(maxNameWidth, name.size()); - maxShapeWidth = std::max(maxShapeWidth, shape.size()); - // Shape input tensor is not considered in TRT-LLM yet - } - - TLLM_LOG_TRACE("Information of context input / output."); - TLLM_LOG_TRACE("Using Optimization Profile: %d", contextIndex); - TLLM_LOG_TRACE(std::string(maxNameWidth + maxShapeWidth + 6, '=')); - std::string info = alignText("Name", maxNameWidth) + "|I/O|" + alignText("Shape", maxShapeWidth) + "|"; - TLLM_LOG_TRACE(info.c_str()); - TLLM_LOG_TRACE(std::string(maxNameWidth + maxShapeWidth + 6, '-')); - for (int i = 0; i < nIO; ++i) - { - auto const& [name, isInput, shape] = tensorInfo[i]; - info = alignText(name, maxNameWidth, false) + "|"; - info += alignText(isInput ? "I" : "O", 3) + "|"; - info += alignText(shape, maxShapeWidth) + "|"; - TLLM_LOG_TRACE(info.c_str()); - } - TLLM_LOG_TRACE(std::string(maxNameWidth + maxShapeWidth + 6, '=')); -} - -void TllmRuntime::clearContexts() -{ - for (auto& context : mContexts) - { - context.reset(); - } - mContexts.clear(); -} - -bool TllmRuntime::executeContext(SizeType32 contextIndex) const -{ - NVTX3_FUNC_RANGE(); - auto& context = getContext(contextIndex); - auto res = context.enqueueV3(mStream->get()); - sync_check_cuda_error(mStream->get()); - return res; -} - -void TllmRuntime::setInputTensorsImpl(SizeType32 contextIndex, TensorMap const& tensorMap, bool throwOnMiss) -{ - NVTX3_FUNC_RANGE(); - auto& context = getContext(contextIndex); - for (auto const& name : mInputTensorNames) - { - auto const pos = tensorMap.find(name); - if (pos == tensorMap.end()) - { - if (throwOnMiss) - { - auto expectedShape = mEngine->getTensorShape(name.c_str()); - TLLM_THROW("Input tensor '%s' not found; expected shape: %s", name.c_str(), - ITensor::toString(expectedShape).c_str()); - } - else - { - continue; - } - } - - auto const& tensor = pos->second; - auto const tensorDtype = tensor->getDataType(); - auto const engineDtype = mEngine->getTensorDataType(name.c_str()); - // WAR: TRT does not support mixed FP8 and FP16 input, so engine expects FP16 tensors. - TLLM_CHECK_WITH_INFO(tensorDtype == engineDtype - || (tensorDtype == nvinfer1::DataType::kFP8 && engineDtype == nvinfer1::DataType::kHALF), - "%s: expected type %d, provided type %d", name.c_str(), static_cast(engineDtype), - static_cast(tensorDtype)); - - auto tensorShape = tensor->getShape(); - - // Change shape of `cache_indirection` for Variable-Beam-Width-Search - // TODO: remove this hack if beamWidth of each request are passed into GptAttentionPlugin by input tensor - if (name == "cache_indirection" && mCurrentBeamWidths.size() > 0) - { - SizeType32 const beamWidth = getCurrentBeamWidth(); - if (tensorShape.d[1] != beamWidth) - { - tensorShape.d[1] = beamWidth; - TLLM_LOG_TRACE("Change shape of cache_indirection to %s", ITensor::toString(tensorShape).c_str()); - } - } - - auto const setInputShapeSuccess = context.setInputShape(name.c_str(), tensorShape); - if (!setInputShapeSuccess) - { - auto const minShape - = mEngine->getProfileShape(name.c_str(), contextIndex, nvinfer1::OptProfileSelector::kMIN); - auto const maxShape - = mEngine->getProfileShape(name.c_str(), contextIndex, nvinfer1::OptProfileSelector::kMAX); - - TLLM_THROW("Tensor '%s' has invalid shape %s, expected in range min %s, max %s", name.c_str(), - ITensor::toString(tensorShape).c_str(), ITensor::toString(minShape).c_str(), - ITensor::toString(maxShape).c_str()); - } - auto* const data = tensor->data(); - if (static_cast(data)) - { - context.setInputTensorAddress(name.c_str(), data); - } - else - { - TLLM_CHECK_WITH_INFO(tensor->getSize() == 0, std::string("Invalid data for tensor: ") + name); - // TensorRT runtime does not support nullptr. - if (!mDummyTensor) - { - mDummyTensor = mBufferManager.gpu(ITensor::makeShape({1})); - } - context.setInputTensorAddress(name.c_str(), mDummyTensor->data()); - } - } -} - -void TllmRuntime::setStaticInputTensors(TensorMap const& tensorMap) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_FUNC_RANGE(); - - TLLM_CHECK_WITH_INFO(getNbContexts() > 0, "Contexts should be created before calling setStaticInputTensors"); - for (auto contextIndex = 0; contextIndex < getNbContexts(); ++contextIndex) - { - setInputTensorsImpl(contextIndex, tensorMap, false); - } - - // move static input tensor names to separate vector - auto const begin = mInputTensorNames.begin(); - auto end = mInputTensorNames.end(); - for (auto const& [name, tensor] : tensorMap) - { - end = std::remove(begin, end, name); - } - mInputTensorNames.erase(end, mInputTensorNames.end()); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TllmRuntime::setInputTensors(SizeType32 contextIndex, TensorMap const& tensorMap) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_FUNC_RANGE(); - setInputTensorsImpl(contextIndex, tensorMap, true); - - auto& context = getContext(contextIndex); - if (mUseShapeInference) - { - NVTX3_SCOPED_RANGE(infer_shapes); - char const* missing = nullptr; - auto const nbMissing = context.inferShapes(1, &missing); - if (nbMissing > 0) - { - TLLM_THROW("Input shape not specified: %s", missing); - } - else if (nbMissing < 0) - { - TLLM_THROW("Invalid input shape"); - } - } - - { - NVTX3_SCOPED_RANGE(final_checks); - TLLM_CHECK_WITH_INFO(context.allInputDimensionsSpecified(), "Input dimensions not specified"); - TLLM_CHECK_WITH_INFO(context.allInputShapesSpecified(), "Input shapes not specified"); - } - - // Print shape of input / output tensors for the TRT engine - if (tensorrt_llm::common::Logger::getLogger()->isEnabled(tensorrt_llm::common::Logger::TRACE)) - { - printContextInfo(contextIndex); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TllmRuntime::setOutputTensors(SizeType32 contextIndex, TensorMap& tensorMap) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_FUNC_RANGE(); - if (isUserBufferEnabled()) - { - // This function will identify the output tensors in the network that need to be bound as UB buffers - // and bind the corresponding buffers to them based on their names. - setUserBufferTensors(contextIndex, tensorMap); - } - - auto& context = getContext(contextIndex); - for (auto const& name : mOutputTensorNames) - { - auto const engineDtype = mEngine->getTensorDataType(name.c_str()); - auto const pos = tensorMap.find(name); - if (pos != tensorMap.end()) - { - auto const& tensor = pos->second; - auto const tensorDtype = tensor->getDataType(); - // WAR: TRT does not support mixed FP8 and FP16 input, so engine expects FP16 tensors. - TLLM_CHECK_WITH_INFO(tensorDtype == engineDtype - || (tensorDtype == nvinfer1::DataType::kFP8 && engineDtype == nvinfer1::DataType::kHALF), - "%s: expected type %d, provided type %d", name.c_str(), static_cast(engineDtype), - static_cast(tensorDtype)); - - if (mUseShapeInference) - { - auto const dims = context.getTensorShape(name.c_str()); - tensor->reshape(dims); - } - context.setTensorAddress(name.c_str(), tensor->data()); - } - else if (mUseShapeInference) - { - auto const dims = context.getTensorShape(name.c_str()); - auto tensor = ITensor::SharedPtr(mBufferManager.gpu(dims, engineDtype)); - tensorMap.insert(pos, std::make_pair(name, tensor)); - context.setTensorAddress(name.c_str(), tensor->data()); - } - else - { - TLLM_THROW("Tensor %s is not found in tensorMap and shape inference is not allowed", name.c_str()); - } - } - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TllmRuntime::setUserBufferTensors(SizeType32 contextIndex, TensorMap& tensorMap) -{ - auto startsWith = [](std::string const& str, std::string const& prefix) -> bool - { return str.size() > prefix.size() && str.compare(0, prefix.size(), prefix) == 0; }; - std::string const prefix(tensorrt_llm::runtime::ub::tensor_prefix); - auto& context = getContext(contextIndex); - for (auto const& name : mOutputTensorNames) - { - auto const pos = tensorMap.find(name); - if (pos != tensorMap.end() || !startsWith(name, prefix)) - { - continue; - } - auto const engineDtype = mEngine->getTensorDataType(name.c_str()); - auto const dims = context.getTensorShape(name.c_str()); - void* ubBuffer = nullptr; - if (name[prefix.size()] == '0') - { - ubBuffer = tensorrt_llm::runtime::ub::ub_get(0).addr; - } - else if (name[prefix.size()] == '1') - { - ubBuffer = tensorrt_llm::runtime::ub::ub_get(1).addr; - } - else if (name[prefix.size()] == '2') - { - ubBuffer = tensorrt_llm::runtime::ub::ub_get(2).addr; - } - else - { - TLLM_CHECK(false); - } - auto tensor = ITensor::SharedPtr(ITensor::wrap(ubBuffer, engineDtype, dims)); - tensorMap.insert(pos, std::make_pair(name, tensor)); - context.setTensorAddress(name.c_str(), ubBuffer); - } -} - -void TllmRuntime::initializeUserBuffer(tensorrt_llm::runtime::WorldConfig const& world_config, SizeType32 maxBatchSize, - SizeType32 maxBeamWidth, SizeType32 maxSequenceLength, SizeType32 hiddenSize, - std::optional maxNumTokens) -{ - auto startsWith = [](std::string const& str, std::string const& prefix) -> bool - { return str.size() > prefix.size() && str.compare(0, prefix.size(), prefix) == 0; }; - std::string const prefix(tensorrt_llm::runtime::ub::tensor_prefix); - bool useNVFP4Model = false; - for (auto const& name : mOutputTensorNames) - { - if (startsWith(name, prefix)) - { - mUserBufferEnabled = true; - if (name[prefix.size()] == '2') - { - useNVFP4Model = true; - break; - } - } - } - if (!mUserBufferEnabled) - { - return; - } - // The hidden size returned by ModelConfig is the real hidden size divided by the TP size. - auto const tpSize = world_config.getTensorParallelism(); - size_t const realHiddenSize = hiddenSize * tpSize; - size_t const tokensNum = maxNumTokens.value_or(maxBatchSize * maxBeamWidth * maxSequenceLength); - TLLM_CHECK(tokensNum > 0); - size_t const elemNum = tokensNum * realHiddenSize; - TLLM_LOG_INFO("[UserBuffer] MaxBatchSize %d, maxBeamWidth %d, maxSequenceLength %d, maxNumTokens %d, select %lu", - maxBatchSize, maxBeamWidth, maxSequenceLength, maxNumTokens.has_value() ? maxNumTokens.value() : 0, tokensNum); - tensorrt_llm::runtime::ub::ub_initialize(world_config); - tensorrt_llm::runtime::ub::ub_allocate(elemNum * sizeof(half)); - tensorrt_llm::runtime::ub::ub_allocate(elemNum * sizeof(half)); - if (useNVFP4Model) - { - tensorrt_llm::runtime::ub::ub_allocate(elemNum * sizeof(uint8_t) / 16); - } -} - -CudaStream const& TllmRuntime::getStream() const -{ - return *mStream; -} - -bool TllmRuntime::hasLayerProfiler(SizeType32 contextId) const -{ - return mContexts[contextId]->getProfiler() != nullptr; -} - -void TllmRuntime::setLayerProfiler() -{ - mLayerProfiler = std::make_unique(); - for (auto& context : mContexts) - { - context->setProfiler(mLayerProfiler.get()); - context->setEnqueueEmitsProfile(false); - } -} - -std::string TllmRuntime::getLayerProfileInfo() const -{ - TLLM_CHECK(mLayerProfiler); - return mLayerProfiler->getLayerProfile(); -} - -void TllmRuntime::reportToProfiler(SizeType32 contextId) -{ - mContexts[contextId]->reportToProfiler(); -} - -void TllmRuntime::loadManagedWeights(RawEngine const& rawEngine, int localRank) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_FUNC_RANGE(); - auto& engine = getEngine(); - auto& manager = getBufferManager(); - if (rawEngine.getManagedWeightsMapOpt().has_value()) - { - TLLM_LOG_DEBUG("Loading managed weights from raw engine"); - auto executorMap = rawEngine.getManagedWeightsMapOpt().value(); - for (auto const& [name, weight] : executorMap) - { - TLLM_LOG_DEBUG("Loading managed weight: %s", name.c_str()); - auto iTensor = tensorrt_llm::executor::detail::toITensor(weight); - auto weightsDevice = std::shared_ptr{manager.copyFrom(*iTensor, MemoryType::kGPU)}; - mManagedWeightsMap.insert(std::make_pair(name, weightsDevice)); - } - } - else - { - TLLM_LOG_DEBUG("Loading managed weights from file"); - auto const enginePath = rawEngine.getPathOpt(); - TLLM_CHECK_WITH_INFO(enginePath.has_value(), "Engine path is not set."); - auto weightPath - = enginePath->parent_path() / ("rank" + std::to_string(localRank) + "_managed_weights.safetensors"); - auto managed_weights = common::safetensors::ISafeTensor::open(weightPath.string().c_str()); - for (auto const& name : managed_weights->keys()) - { - TLLM_LOG_DEBUG("Loading managed weight: %s", name.c_str()); - auto const weight = managed_weights->getTensor(name.c_str()); - TLLM_CHECK(weight->dtype() == engine.getTensorDataType(name.c_str())); - auto weightsDevice - = std::shared_ptr{manager.allocate(MemoryType::kGPU, weight->trtDims(), weight->dtype())}; - manager.copy(weight->data(), *weightsDevice, MemoryType::kCPU); - mManagedWeightsMap.insert(std::make_pair(name, weightsDevice)); - } - } - setStaticInputTensors(mManagedWeightsMap); - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} diff --git a/cpp/tensorrt_llm/runtime/tllmRuntime.h b/cpp/tensorrt_llm/runtime/tllmRuntime.h deleted file mode 100644 index dfef06d8b45f..000000000000 --- a/cpp/tensorrt_llm/runtime/tllmRuntime.h +++ /dev/null @@ -1,243 +0,0 @@ -/* - * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. - * - * 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. - */ -#pragma once - -#include "tensorrt_llm/batch_manager/llmRequest.h" -#include "tensorrt_llm/runtime/bufferManager.h" -#include "tensorrt_llm/runtime/common.h" -#include "tensorrt_llm/runtime/iTensor.h" -#include "tensorrt_llm/runtime/layerProfiler.h" -#include "tensorrt_llm/runtime/rawEngine.h" -#include "tensorrt_llm/runtime/worldConfig.h" -#include - -#include -#include -#include -#include -#include - -namespace tensorrt_llm::runtime -{ -class TllmRuntime -{ -public: - using TensorMap = StringPtrMap; - - explicit TllmRuntime(RawEngine const& rawEngine, nvinfer1::ILogger* logger, bool useGpuDirectStorage = false, - float gpuWeightsPercent = 1.0f, bool useShapeInference = true); - - SizeType32 getNbContexts() const - { - return static_cast(mContexts.size()); - } - - nvinfer1::IExecutionContext& getContext(SizeType32 contextIndex) const - { - return *mContexts.at(contextIndex); - } - - SizeType32 getNbProfiles() const - { - return static_cast(mEngine->getNbOptimizationProfiles()); - } - - /// @brief If multiple TensorRT optimization profiles are built in the engine, this function selects the - /// corresponding profile that is going to be used based on the runtime shape, for now, TensorRT LLM only split - /// multiple profiles on the num_tokens dimension, hence the profile index is selected based on which profile - /// handles the actual num_tokens - /// @return The index of the selected TensorRT optimization profile - [[nodiscard]] SizeType32 getOptProfileId(int numTokens, std::vector const& splitPoints) const - { - if (getNbProfiles() == 1) - { - return 0; - } - auto const it = std::lower_bound(splitPoints.begin(), splitPoints.end(), numTokens); - auto const optProfileId = std::distance(splitPoints.begin(), it); - return optProfileId; - } - - nvinfer1::IExecutionContext& addContext(std::int32_t profileIndex); - - void clearContexts(); - - /// @brief Set input tensors from tensorMap for all contexts. - /// @details The function can be used to set static input tensors for all iterations. If a tensor was set this way, - /// it doesn't need to included in calls to setInputTensors anymore. - void setStaticInputTensors(TensorMap const& tensorMap); - - /// @brief Set input tensors from tensorMap for context at contextIndex. - /// @details The function expects that all input tensors (excluding the ones set by setStaticInputTensors) are - /// contained in the tensorMap. If a tensor is missing, has a bad shape or type, it will throw. - void setInputTensors(SizeType32 contextIndex, TensorMap const& tensorMap); - - /// @brief Set output tensors from tensorMap for context at contextIndex. - /// @details The function expects that all output tensors are contained in the tensorMap. If a tensor is missing and - /// shape inference is enabled, it will allocate the tensor on GPU and insert it into the tensorMap. Otherwise it - /// will throw. - void setOutputTensors(SizeType32 contextIndex, TensorMap& tensorMap); - - bool executeContext(SizeType32 contextIndex) const; - - CudaStream const& getStream() const; - - BufferManager::CudaStreamPtr getStreamPtr() - { - return mStream; - } - - nvinfer1::ICudaEngine& getEngine() - { - return *mEngine; - } - - nvinfer1::ICudaEngine const& getEngine() const - { - return *mEngine; - } - - nvinfer1::IEngineInspector& getEngineInspector() - { - return *mEngineInspector; - } - - nvinfer1::IEngineInspector const& getEngineInspector() const - { - return *mEngineInspector; - } - - BufferManager& getBufferManager() - { - return mBufferManager; - } - - BufferManager const& getBufferManager() const - { - return mBufferManager; - } - - void setLayerProfiler(); - bool hasLayerProfiler(SizeType32 contextId) const; - std::string getLayerProfileInfo() const; - void reportToProfiler(SizeType32 contextId); - void loadManagedWeights(RawEngine const& rawEngine, int localRank); - void initializeUserBuffer(tensorrt_llm::runtime::WorldConfig const& world_config, SizeType32 maxBatchSize, - SizeType32 maxBeamWidth, SizeType32 maxSequenceLength, SizeType32 hiddenSize, - std::optional maxNumTokens); - - bool isUserBufferEnabled() const - { - return mUserBufferEnabled; - } - - void setCurrentBeamWidths(std::vector const& beamWidth) noexcept - { - mCurrentBeamWidths = beamWidth; - } - - [[nodiscard]] SizeType32 const& getCurrentBeamWidth() const noexcept - { - // At present, all requests of a batch must have the same beam width in one generation step (or they will not - // be batched together). So, the beam widths in `mCurrentBeamWidths` are the same. - // Corresponding changes must be done if Diverse-Beam-Width-Search (DBWS, requests with diverse beam width in - // a batch in one generation step) is supported in the future. - TLLM_CHECK_WITH_INFO(mCurrentBeamWidths.size() > 0, "`mCurrentBeamWidths` is empty."); - bool const isEqual = std::all_of(mCurrentBeamWidths.begin(), mCurrentBeamWidths.end(), - [&](int elem) { return elem == mCurrentBeamWidths.front(); }); - TLLM_CHECK_WITH_INFO(isEqual, "beam widths in `mCurrentBeamWidths` are not all equal."); - return mCurrentBeamWidths.front(); - } - -private: - void cacheTensorNames(); - - void setInputTensorsImpl(SizeType32 contextIndex, TensorMap const& tensorMap, bool throwOnMiss); - - void setUserBufferTensors(SizeType32 contextIndex, TensorMap& tensorMap); - - void printEngineInfo(); - - void printContextInfo(SizeType32 contextIndex); - - // Tool functions for `printEngineInfo()`. - static std::string shapeToString(nvinfer1::Dims64 const& dim) - { - std::string output("("); - if (dim.nbDims == 0) - { - return output + ")"; - } - for (int i = 0; i < dim.nbDims - 1; ++i) - { - output += std::to_string(dim.d[i]) + ", "; - } - output += std::to_string(dim.d[dim.nbDims - 1]) + ")"; - return output; - } - - static std::string dataTypeToString(nvinfer1::DataType type) - { - switch (type) - { - case nvinfer1::DataType::kINT64: return "INT64"; - case nvinfer1::DataType::kINT32: return "INT32"; - case nvinfer1::DataType::kFLOAT: return "FP32"; - case nvinfer1::DataType::kBF16: return "BF16"; - case nvinfer1::DataType::kHALF: return "FP16"; - case nvinfer1::DataType::kBOOL: return "BOOL"; - case nvinfer1::DataType::kUINT8: return "UINT8"; - case nvinfer1::DataType::kINT8: return "INT8"; - case nvinfer1::DataType::kFP8: return "FP8"; - case nvinfer1::DataType::kINT4: return "INT4"; - case nvinfer1::DataType::kFP4: return "FP4"; - default: return "UNKNOWN"; - } - return ""; - } - - static std::string alignText( - std::string const& text, int const width, bool const bCenter = true, char const blank = ' ') - { - int textLen = text.size(); - int padLeft = 0; - int padRight = 0; - padLeft = bCenter ? (width - textLen) / 2 : 0; - padRight = width - padLeft - textLen; - return std::string(padLeft, blank) + text + std::string(padRight, blank); - } - - BufferManager::CudaStreamPtr mStream; - BufferManager mBufferManager; - std::unique_ptr mRuntime; - std::unique_ptr mEngine; - BufferManager::IBufferPtr mEngineBuffer; - std::vector> mContexts; - std::unique_ptr mDummyTensor; - std::unique_ptr mEngineInspector; - std::unique_ptr mLayerProfiler; - bool mUseShapeInference; - TensorMap mManagedWeightsMap; - // List of input tensor names. - // Names of static tensors are removed from this list when setStaticInputTensors is called. - std::vector mInputTensorNames; - std::vector mOutputTensorNames; - - bool mUserBufferEnabled; - // For Variable-Beam-Width-Search - std::vector mCurrentBeamWidths; -}; -} // namespace tensorrt_llm::runtime diff --git a/cpp/tensorrt_llm/runtime/tllmStreamReaders.cpp b/cpp/tensorrt_llm/runtime/tllmStreamReaders.cpp deleted file mode 100644 index 55440bbe714f..000000000000 --- a/cpp/tensorrt_llm/runtime/tllmStreamReaders.cpp +++ /dev/null @@ -1,217 +0,0 @@ -/* - * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. - * - * 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. - */ - -#include "tllmStreamReaders.h" -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/logger.h" - -#include -#include -#include -#include -#include -#include -#include - -// Non-GDS StreamReader - -StreamReader::StreamReader(std::filesystem::path fp) -{ - mFile.open(fp.string(), std::ios::binary | std::ios::in); - TLLM_CHECK_WITH_INFO(mFile.good(), std::string("Error opening engine file: " + fp.string())); -} - -StreamReader::~StreamReader() -{ - if (mFile.is_open()) - { - mFile.close(); - } -} - -int64_t StreamReader::read(void* destination, int64_t nbBytes) -{ - if (!mFile.good()) - { - return -1; - } - - mFile.read(static_cast(destination), nbBytes); - - return mFile.gcount(); -} - -// StreamReader using GDS - -GDSStreamReader::GDSStreamReader(std::filesystem::path const& filePath) -{ - auto const start_time = std::chrono::high_resolution_clock::now(); - initializeDriver(); - auto const elapsed_ms - = std::chrono::duration_cast(std::chrono::high_resolution_clock::now() - start_time); - - TLLM_LOG_INFO("GDS driver initialization time %lld ms", elapsed_ms); - - open(filePath); -} - -bool GDSStreamReader::open(std::string const& filepath) -{ - if (!initializeDriver()) - { - TLLM_LOG_INFO("Failed to initialize cuFile driver"); - return false; - } - - int32_t const ret = ::open(filepath.c_str(), O_CREAT | O_RDWR | O_DIRECT, 0664); - - if (ret < 0) - { - TLLM_LOG_INFO("Failed to open engine file"); - return false; - } - - mFd = ret; - mFileSize = lseek(mFd, 0, SEEK_END); - lseek(mFd, 0, SEEK_SET); - - CUfileDescr_t fileDescr; - memset((void*) &fileDescr, 0, sizeof(fileDescr)); - fileDescr.handle.fd = mFd; - fileDescr.type = CU_FILE_HANDLE_TYPE_OPAQUE_FD; - - CUfileError_t gdsStatus = cuFileHandleRegister(&mFileHandle, &fileDescr); - - if (gdsStatus.err != CU_FILE_SUCCESS) - { - TLLM_LOG_INFO("Failed to cuFileHandleRegister"); - ::close(mFd); - return false; - } - return true; -} - -void GDSStreamReader::close() -{ - if (mFd >= 0) - { - ::close(mFd); - mFd = -1; - } -} - -GDSStreamReader::~GDSStreamReader() -{ - if (mFileHandle) - { - cuFileHandleDeregister(mFileHandle); - mFileHandle = nullptr; - } - - if (mDriverInitialized) - { - cuFileDriverClose(); - } -} - -bool GDSStreamReader::seek(int64_t offset, nvinfer1::SeekPosition where) noexcept -{ - switch (where) - { - case nvinfer1::SeekPosition::kSET: mCursor = offset; return true; - case nvinfer1::SeekPosition::kCUR: mCursor += offset; return true; - case nvinfer1::SeekPosition::kEND: mCursor = -offset; return true; - default: return false; - } - return true; -} - -int64_t GDSStreamReader::read(void* dest, int64_t bytes, cudaStream_t stream) noexcept -{ - cudaPointerAttributes attributes{}; - if (cudaPointerGetAttributes(&attributes, dest) != cudaSuccess) - { - TLLM_LOG_INFO("cudaPointerGetAttributes failed"); - } - - off_t destOffset = 0; - void* destBase = dest; - - if (attributes.type == cudaMemoryTypeDevice) - { - CUdeviceptr cuDest = reinterpret_cast(dest); - CUdeviceptr cuBufBase = 0; - size_t cuBufSize = 0; - - cuMemGetAddressRange(&cuBufBase, &cuBufSize, cuDest); - destOffset += cuDest - cuBufBase; - destBase = reinterpret_cast(cuBufBase); - } - cuFileRead(this->mFileHandle, destBase, bytes, mCursor, destOffset); - - mCursor += bytes; - return bytes; -} - -void GDSStreamReader::reset() -{ - lseek(mFd, 0, SEEK_SET); - mCursor = 0; -} - -[[nodiscard]] bool GDSStreamReader::isOpen() const -{ - bool open = mFd >= 0; - return open; -} - -bool GDSStreamReader::initializeDriver() -{ - if (mDriverInitialized) - { - return true; - } - - mCuFileLibHandle = dlopen("libcufile.so", RTLD_LAZY | RTLD_GLOBAL); - if (!mCuFileLibHandle) - { - TLLM_LOG_INFO("Failed to dlopen libcufile.so"); - return false; - } - - // Load the required functions - *reinterpret_cast(&cuFileDriverOpen) = dlsym(mCuFileLibHandle, "cuFileDriverOpen"); - *reinterpret_cast(&cuFileHandleRegister) = dlsym(mCuFileLibHandle, "cuFileHandleRegister"); - *reinterpret_cast(&cuFileHandleDeregister) = dlsym(mCuFileLibHandle, "cuFileHandleDeregister"); - *reinterpret_cast(&cuFileDriverClose) = dlsym(mCuFileLibHandle, "cuFileDriverClose"); - *reinterpret_cast(&cuFileRead) = dlsym(mCuFileLibHandle, "cuFileRead"); - - if (!cuFileDriverOpen || !cuFileHandleRegister || !cuFileHandleDeregister || !cuFileDriverClose || !cuFileRead) - { - TLLM_LOG_INFO("Failed to dlsym libcufile.so"); - return false; - } - - CUfileError_t gdsStatus = cuFileDriverOpen(); - if (gdsStatus.err != CU_FILE_SUCCESS) - { - TLLM_LOG_INFO("cuFileDriverOpen failed"); - return false; - } - - mDriverInitialized = true; - return true; -} diff --git a/cpp/tensorrt_llm/runtime/tllmStreamReaders.h b/cpp/tensorrt_llm/runtime/tllmStreamReaders.h deleted file mode 100644 index 943f0bb3e32e..000000000000 --- a/cpp/tensorrt_llm/runtime/tllmStreamReaders.h +++ /dev/null @@ -1,71 +0,0 @@ -/* - * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. - * - * 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. - */ -#pragma once - -#include - -#include -#include -#include - -class StreamReader final : public nvinfer1::IStreamReader -{ -public: - StreamReader(std::filesystem::path fp); - - virtual ~StreamReader(); - - int64_t read(void* destination, int64_t nbBytes) final; - -private: - std::ifstream mFile; -}; - -class GDSStreamReader final : public nvinfer1::IStreamReaderV2 -{ -public: - explicit GDSStreamReader(std::filesystem::path const& filePath); - - virtual ~GDSStreamReader(); - - void close(); - - [[nodiscard]] bool isOpen() const; - - bool open(std::string const& filepath); - - int64_t read(void* dest, int64_t bytes, cudaStream_t stream) noexcept final; - - void reset(); - - bool seek(int64_t offset, nvinfer1::SeekPosition where) noexcept final; - -private: - bool initializeDriver(); - - void* mCuFileLibHandle{}; - CUfileHandle_t mFileHandle{nullptr}; - bool mDriverInitialized{false}; - int32_t mFd{-1}; - int64_t mCursor{0}; - int64_t mFileSize{0}; - - CUfileError_t (*cuFileDriverOpen)(){}; - CUfileError_t (*cuFileHandleRegister)(CUfileHandle_t*, CUfileDescr_t*){}; - CUfileError_t (*cuFileHandleDeregister)(CUfileHandle_t){}; - CUfileError_t (*cuFileDriverClose)(){}; - ssize_t (*cuFileRead)(CUfileHandle_t, void*, size_t, int64_t, int64_t){}; -}; diff --git a/cpp/tensorrt_llm/runtime/utils/debugUtils.cu b/cpp/tensorrt_llm/runtime/utils/debugUtils.cu index 661dacd9a7ac..d4aaa8244c11 100644 --- a/cpp/tensorrt_llm/runtime/utils/debugUtils.cu +++ b/cpp/tensorrt_llm/runtime/utils/debugUtils.cu @@ -26,6 +26,7 @@ #include "tensorrt_llm/common/cudaUtils.h" #include "tensorrt_llm/common/memoryUtils.h" +#include "tensorrt_llm/common/tllmDataType.h" #include #include @@ -167,7 +168,7 @@ template bool tensorHasInvalid(ITensor const& tensor, BufferManager const& manager, std::string const& infoStr) { printLogitsKeyInfo(tensor, infoStr); - auto foundInvalid = BufferManager::pinnedPool(ITensor::makeShape({1}), nvinfer1::DataType::kINT32); + auto foundInvalid = BufferManager::pinnedPool(ITensor::makeShape({1}), tensorrt_llm::DataType::kINT32); auto foundInvalidPtr = bufferCast(*foundInvalid); foundInvalidPtr[0] = 0; auto const size = tensor.getSize(); @@ -184,24 +185,24 @@ template bool tensorHasInvalid<__nv_fp8_e4m3>( ITensor const& tensor, BufferManager const& manager, std::string const& infoStr); bool tensorHasInvalid( - size_t M, size_t K, nvinfer1::DataType type, void const* data, cudaStream_t stream, std::string const& infoStr) + size_t M, size_t K, tensorrt_llm::DataType type, void const* data, cudaStream_t stream, std::string const& infoStr) { auto tensorView = ITensor::wrap( const_cast(data), type, ITensor::makeShape({static_cast(M), static_cast(K)})); auto manager = BufferManager(std::make_shared(stream)); - if (type == nvinfer1::DataType::kFLOAT) + if (type == tensorrt_llm::DataType::kFLOAT) { return tensorHasInvalid(*tensorView, manager, infoStr); } - else if (type == nvinfer1::DataType::kHALF) + else if (type == tensorrt_llm::DataType::kHALF) { return tensorHasInvalid(*tensorView, manager, infoStr); } - else if (type == nvinfer1::DataType::kBF16) + else if (type == tensorrt_llm::DataType::kBF16) { return tensorHasInvalid<__nv_bfloat16>(*tensorView, manager, infoStr); } - else if (type == nvinfer1::DataType::kFP8) + else if (type == tensorrt_llm::DataType::kFP8) { return tensorHasInvalid<__nv_fp8_e4m3>(*tensorView, manager, infoStr); } diff --git a/cpp/tensorrt_llm/runtime/utils/numpyUtils.cpp b/cpp/tensorrt_llm/runtime/utils/numpyUtils.cpp index 6f95704455d3..931fbf0203da 100644 --- a/cpp/tensorrt_llm/runtime/utils/numpyUtils.cpp +++ b/cpp/tensorrt_llm/runtime/utils/numpyUtils.cpp @@ -20,9 +20,9 @@ #include "tensorrt_llm/common/logger.h" #include "tensorrt_llm/common/memoryUtils.h" #include "tensorrt_llm/common/stringUtils.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/runtime/bufferManager.h" #include "tensorrt_llm/runtime/iTensor.h" -#include #include #include @@ -34,9 +34,9 @@ namespace tc = tensorrt_llm::common; namespace tensorrt_llm::runtime::utils { -std::string getNumpyTypeDesc(nvinfer1::DataType type) +std::string getNumpyTypeDesc(tensorrt_llm::DataType type) { - using dt = nvinfer1::DataType; + using dt = tensorrt_llm::DataType; static std::unordered_map const type_map{{dt::kBOOL, "?"}, {dt::kUINT8, "u1"}, {dt::kINT8, "i1"}, {dt::kINT32, "i4"}, {dt::kINT64, "i8"}, {dt::kHALF, "f2"}, {dt::kFLOAT, "f4"}}; @@ -51,11 +51,11 @@ std::string getNumpyTypeDesc(nvinfer1::DataType type) return type_map.count(type) > 0 ? type_map.at(type) : "x"; } -nvinfer1::DataType typeFromNumpyDesc(std::string const& type) +tensorrt_llm::DataType typeFromNumpyDesc(std::string const& type) { TLLM_LOG_DEBUG("numpy type: %s", type.c_str()); - using dt = nvinfer1::DataType; + using dt = tensorrt_llm::DataType; static std::unordered_map const type_map{{"?", dt::kBOOL}, {"u1", dt::kUINT8}, {"i1", dt::kINT8}, {"i4", dt::kINT32}, {"i8", dt::kINT64}, {"f2", dt::kHALF}, {"f4", dt::kFLOAT}}; TLLM_CHECK_WITH_INFO(type_map.count(type) > 0, "numpy data type '" + type + "' not supported"); @@ -102,7 +102,7 @@ void parseNpyIntro(FILE*& f_ptr, uint32_t& header_len, uint32_t& start_data) start_data = 8 + 2 * npy_major + header_len; } -int parseNpyHeader(FILE*& f_ptr, uint32_t header_len, nvinfer1::DataType& type, std::vector& shapeVec) +int parseNpyHeader(FILE*& f_ptr, uint32_t header_len, tensorrt_llm::DataType& type, std::vector& shapeVec) { char* header_c = (char*) malloc(header_len * sizeof(char)); TLLM_CHECK_WITH_INFO(header_c != nullptr, "Failed to allocate memory for npy header"); @@ -168,11 +168,11 @@ int parseNpyHeader(FILE*& f_ptr, uint32_t header_len, nvinfer1::DataType& type, uint32_t header_len, start_data; utils::parseNpyIntro(f_ptr, header_len, start_data); - nvinfer1::DataType type; + tensorrt_llm::DataType type; std::vector shape; utils::parseNpyHeader(f_ptr, header_len, type, shape); - nvinfer1::Dims dims; + tensorrt_llm::Dims dims; dims.nbDims = shape.size(); std::copy(shape.begin(), shape.end(), dims.d); @@ -203,10 +203,10 @@ void saveNpy(BufferManager const& manager, ITensor const& tensor, std::string co auto const dtype = tensor.getDataType(); #ifdef ENABLE_BF16 - if (dtype == nvinfer1::DataType::kBF16) + if (dtype == tensorrt_llm::DataType::kBF16) { TLLM_CHECK(where == MemoryType::kGPU); - auto tensorFp32 = manager.gpu(shape, nvinfer1::DataType::kFLOAT); + auto tensorFp32 = manager.gpu(shape, tensorrt_llm::DataType::kFLOAT); auto dataFp32 = bufferCast(*tensorFp32); auto dataBf16 = bufferCast<__nv_bfloat16 const>(tensor); tc::invokeCudaD2DcpyConvert(dataFp32, dataBf16, tensorSize); diff --git a/cpp/tensorrt_llm/runtime/utils/runtimeUtils.h b/cpp/tensorrt_llm/runtime/utils/runtimeUtils.h index f131ab3419bf..da94419eff64 100644 --- a/cpp/tensorrt_llm/runtime/utils/runtimeUtils.h +++ b/cpp/tensorrt_llm/runtime/utils/runtimeUtils.h @@ -25,7 +25,6 @@ namespace tensorrt_llm::runtime { -class TllmRuntime; namespace utils { diff --git a/cpp/tensorrt_llm/runtime/virtualMemory.cpp b/cpp/tensorrt_llm/runtime/virtualMemory.cpp index 0d08012a29d8..c2ca710db9ac 100644 --- a/cpp/tensorrt_llm/runtime/virtualMemory.cpp +++ b/cpp/tensorrt_llm/runtime/virtualMemory.cpp @@ -16,6 +16,7 @@ #include "tensorrt_llm/runtime/virtualMemory.h" #include "bufferManager.h" +#include "tensorrt_llm/common/tllmDataType.h" #include #include @@ -141,8 +142,8 @@ void OffloadConfigurator::teardown(CUmemGenericAllocationHandle, bool destructin { switch (mBackType) { - case MemoryType::kCPU: mBackedStorage = BufferManager::cpu(mSize, nvinfer1::DataType::kINT8); break; - case MemoryType::kPINNED: mBackedStorage = BufferManager::pinned(mSize, nvinfer1::DataType::kINT8); break; + case MemoryType::kCPU: mBackedStorage = BufferManager::cpu(mSize, tensorrt_llm::DataType::kINT8); break; + case MemoryType::kPINNED: mBackedStorage = BufferManager::pinned(mSize, tensorrt_llm::DataType::kINT8); break; default: TLLM_THROW("Unknown memory type: %d", static_cast(mBackType)); } } diff --git a/cpp/tensorrt_llm/testing/CMakeLists.txt b/cpp/tensorrt_llm/testing/CMakeLists.txt deleted file mode 100644 index 646302929817..000000000000 --- a/cpp/tensorrt_llm/testing/CMakeLists.txt +++ /dev/null @@ -1,38 +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. -include(FetchContent) - -set(SRCS modelSpec.cpp) - -include_directories(${API_INCLUDE_DIR}/tensorrt_llm/runtime) - -if(NOT WIN32) - # additional warnings - # - # Ignore overloaded-virtual warning. We intentionally change parameters of - # some methods in derived class. - set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wall") - if(WARNING_IS_ERROR) - message(STATUS "Treating warnings as errors in GCC compilation") - set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Werror") - endif() -else() # Windows - # warning level 4 - set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} /W4") -endif() - -add_library(testing_src OBJECT ${SRCS}) -set_property(TARGET testing_src PROPERTY POSITION_INDEPENDENT_CODE ON) -set_property(TARGET testing_src PROPERTY CUDA_RESOLVE_DEVICE_SYMBOLS ON) diff --git a/cpp/tensorrt_llm/testing/modelSpec.cpp b/cpp/tensorrt_llm/testing/modelSpec.cpp deleted file mode 100644 index bc868a157a1c..000000000000 --- a/cpp/tensorrt_llm/testing/modelSpec.cpp +++ /dev/null @@ -1,303 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2023-2025 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. - */ - -#include "modelSpec.h" -#include "tensorrt_llm/common/dataType.h" - -#include - -namespace tensorrt_llm::testing -{ - -std::string ModelSpec::getQuantMethodString() const -{ - switch (mQuantMethod) - { - case QuantMethod::kNONE: - // Bypass here. - break; - case QuantMethod::kSMOOTH_QUANT: return "sq"; break; - default: throw std::runtime_error("Unsupported quant method"); break; - } - - return ""; -} - -std::string ModelSpec::getKVCacheTypeString() const -{ - switch (mKVCacheType) - { - case KVCacheType::kDISABLED: return "no-cache"; break; - case KVCacheType::kPAGED: return "paged"; break; - case KVCacheType::kCONTINUOUS: return "continuous"; break; - default: throw std::runtime_error("Unsupported KV cache type"); break; - } - - return ""; -} - -std::string ModelSpec::getSpeculativeDecodingModeString() const -{ - if (mSpecDecodingMode.isLookaheadDecoding()) - { - return "la-decoding"; - } - else if (mSpecDecodingMode.isDraftTokensExternal()) - { - return "draft-tokens"; - } - else if (mSpecDecodingMode.isNone()) - { - // Bypass here. - } - else if (mSpecDecodingMode.isExplicitDraftTokens()) - { - return "explicit-draft-tokens"; - } - else if (mSpecDecodingMode.isMedusa()) - { - return "medusa"; - } - else if (mSpecDecodingMode.isEagle()) - { - return "eagle"; - } - else - { - throw std::runtime_error("Unsupported decoding mode"); - } - - return ""; -} - -std::string ModelSpec::getCapacitySchedulerString() const -{ - if (mCapacitySchedulerPolicy) - { - if (mCapacitySchedulerPolicy.value() == tensorrt_llm::executor::CapacitySchedulerPolicy::kMAX_UTILIZATION) - { - return "MaxUtilization"; - } - else if (mCapacitySchedulerPolicy.value() - == tensorrt_llm::executor::CapacitySchedulerPolicy::kGUARANTEED_NO_EVICT) - { - return "GuaranteedNoEvict"; - } - else if (mCapacitySchedulerPolicy.value() == tensorrt_llm::executor::CapacitySchedulerPolicy::kSTATIC_BATCH) - { - return "StaticBatch"; - } - else - { - throw std::runtime_error("Unsupported capacity scheduler"); - } - } - return ""; -} - -std::string ModelSpec::getInputFile() const -{ - return mInputFile; -} - -std::string ModelSpec::getModelPath() const -{ - std::vector ret; - - ret.emplace_back(getDtypeString()); - - if (mUseGptAttentionPlugin || mUseMambaPlugin) - { - if (mUseGptAttentionPlugin && mUseMambaPlugin) - { - throw std::runtime_error("Cannot use both GPT attention plugin and MAMBA plugin"); - } - - ret.emplace_back("plugin"); - } - else - { - ret.emplace_back("default"); - } - - if (mUsePackedInput) - { - ret.emplace_back("packed"); - } - - ret.emplace_back(getKVCacheTypeString()); - - if (mMaxInputLength) - { - ret.emplace_back("in" + std::to_string(mMaxInputLength)); - } - - ret.emplace_back(getSpeculativeDecodingModeString()); - - if (mUseLoraPlugin) - { - ret.emplace_back("lora"); - } - - ret.emplace_back(getQuantMethodString()); - - if (mUseMultipleProfiles) - { - ret.emplace_back("nprofiles"); - } - - if (mGatherLogits) - { - ret.emplace_back("gather"); - } - - auto finalRet = std::accumulate(ret.begin(), ret.end(), std::string(), - [](std::string& a, std::string& b) - { - if (a.empty()) - { - return b; - } - else - { - return b.empty() ? a : a + "_" + b; - } - }); - - return finalRet; -} - -std::string ModelSpec::getResultsFileInternal(OutputContentType outputContentType) const -{ - std::vector ret; - - if (mInputFile == "input_tokens_long.npy") - { - ret.emplace_back("output_tokens_long"); - } - else - { - ret.emplace_back("output_tokens"); - } - - if (mMaxOutputLength) - { - ret.emplace_back("out" + std::to_string(mMaxOutputLength)); - } - - ret.emplace_back(getDtypeString()); - - if (mUseGptAttentionPlugin || mUseMambaPlugin) - { - if (mUseGptAttentionPlugin && mUseMambaPlugin) - { - throw std::runtime_error("Cannot use both GPT attention plugin and MAMBA plugin"); - } - ret.emplace_back("plugin"); - } - - if (mUsePackedInput) - { - ret.emplace_back("packed"); - } - - ret.emplace_back(getKVCacheTypeString()); - - ret.emplace_back(getQuantMethodString()); - - if (mGatherLogits) - { - ret.emplace_back("gather"); - } - - ret.emplace_back("tp" + std::to_string(mTPSize)); - - ret.emplace_back("pp" + std::to_string(mPPSize)); - - ret.emplace_back("cp" + std::to_string(mCPSize)); - - if (mEnableContextFMHAFp32Acc) - { - ret.emplace_back("fmhafp32acc"); - } - - switch (outputContentType) - { - case OutputContentType::kNONE: - // Bypass here. - break; - case OutputContentType::kCONTEXT_LOGITS: ret.emplace_back("logits_context"); break; - case OutputContentType::kGENERATION_LOGITS: ret.emplace_back("logits_generation"); break; - case OutputContentType::kLOG_PROBS: ret.emplace_back("log_probs"); break; - case OutputContentType::kCUM_LOG_PROBS: ret.emplace_back("cum_log_probs"); break; - default: throw std::runtime_error("Unsupported output content type"); break; - } - - auto finalRet = std::accumulate(ret.begin(), ret.end(), std::string(), - [](std::string& a, std::string& b) - { - if (a.empty()) - { - return b; - } - else - { - return b.empty() ? a : a + "_" + b; - } - }); - return finalRet + ".npy"; -} - -std::string ModelSpec::getResultsFile() const -{ - return mOtherModelSpecToCompare ? mOtherModelSpecToCompare->getResultsFileInternal(OutputContentType::kNONE) - : getResultsFileInternal(OutputContentType::kNONE); -} - -std::string ModelSpec::getGenerationLogitsFile() const -{ - return mOtherModelSpecToCompare - ? mOtherModelSpecToCompare->getResultsFileInternal(OutputContentType::kGENERATION_LOGITS) - : getResultsFileInternal(OutputContentType::kGENERATION_LOGITS); -} - -std::string ModelSpec::getContextLogitsFile() const -{ - return mOtherModelSpecToCompare - ? mOtherModelSpecToCompare->getResultsFileInternal(OutputContentType::kCONTEXT_LOGITS) - : getResultsFileInternal(OutputContentType::kCONTEXT_LOGITS); -} - -std::string ModelSpec::getCumLogProbsFile() const -{ - return mOtherModelSpecToCompare - ? mOtherModelSpecToCompare->getResultsFileInternal(OutputContentType::kCUM_LOG_PROBS) - : getResultsFileInternal(OutputContentType::kCUM_LOG_PROBS); -} - -std::string ModelSpec::getLogProbsFile() const -{ - return mOtherModelSpecToCompare ? mOtherModelSpecToCompare->getResultsFileInternal(OutputContentType::kLOG_PROBS) - : getResultsFileInternal(OutputContentType::kLOG_PROBS); -} - -std::string ModelSpec::getDtypeString() const -{ - return tensorrt_llm::common::getDtypeString(mDataType); -} - -} // namespace tensorrt_llm::testing diff --git a/cpp/tensorrt_llm/testing/modelSpec.h b/cpp/tensorrt_llm/testing/modelSpec.h deleted file mode 100644 index 5b6f88dcd135..000000000000 --- a/cpp/tensorrt_llm/testing/modelSpec.h +++ /dev/null @@ -1,342 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2023-2025 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. - */ - -#pragma once - -#include "NvInfer.h" -#include "tensorrt_llm/runtime/common.h" -#include "tensorrt_llm/runtime/modelConfig.h" -#include "tensorrt_llm/runtime/speculativeDecodingMode.h" - -#include -#include - -namespace tensorrt_llm::testing -{ - -using tensorrt_llm::runtime::SizeType32; -using tensorrt_llm::runtime::SpeculativeDecodingMode; -using KVCacheType = tensorrt_llm::runtime::ModelConfig::KVCacheType; - -enum class QuantMethod -{ - kNONE, - kSMOOTH_QUANT, -}; - -enum class OutputContentType -{ - kNONE, - kCONTEXT_LOGITS, - kGENERATION_LOGITS, - kLOG_PROBS, - kCUM_LOG_PROBS -}; - -class ModelSpec -{ -public: - ModelSpec(std::string const& inputFile, nvinfer1::DataType dtype, - std::shared_ptr otherModelSpecToCompare = nullptr) - : mInputFile{std::move(inputFile)} - , mDataType{dtype} - , mOtherModelSpecToCompare(otherModelSpecToCompare) - { - } - - ModelSpec& setInputFile(std::string const& inputFile) - { - mInputFile = inputFile; - return *this; - } - - ModelSpec& useGptAttentionPlugin() - { - mUseGptAttentionPlugin = true; - return *this; - } - - ModelSpec& usePackedInput() - { - mUsePackedInput = true; - return *this; - } - - ModelSpec& setKVCacheType(KVCacheType kvCacheType) - { - mKVCacheType = kvCacheType; - return *this; - } - - ModelSpec& setKVCacheReuse(bool kvCacheReuse) - { - mKVCacheReuse = kvCacheReuse; - return *this; - } - - ModelSpec& useDecoderPerRequest() - { - mDecoderPerRequest = true; - return *this; - } - - ModelSpec& useTensorParallelism(int tensorParallelism) - { - mTPSize = tensorParallelism; - return *this; - } - - ModelSpec& usePipelineParallelism(int pipelineParallelism) - { - mPPSize = pipelineParallelism; - return *this; - } - - ModelSpec& useContextParallelism(int contextParallelism) - { - mCPSize = contextParallelism; - return *this; - } - - ModelSpec& setDraftTokens(SizeType32 maxDraftTokens) - { - mMaxDraftTokens = maxDraftTokens; - return *this; - } - - ModelSpec& useAcceptByLogits() - { - mAcceptDraftByLogits = true; - return *this; - } - - ModelSpec& useMambaPlugin() - { - mUseMambaPlugin = true; - return *this; - } - - ModelSpec& gatherLogits() - { - mGatherLogits = true; - return *this; - } - - ModelSpec& replaceLogits() - { - mReplaceLogits = true; - return *this; - } - - ModelSpec& returnLogProbs() - { - mReturnLogProbs = true; - return *this; - } - - ModelSpec& smokeTest() - { - mSmokeTest = true; - return *this; - } - - ModelSpec& useMedusa() - { - mSpecDecodingMode = SpeculativeDecodingMode::Medusa(); - return *this; - } - - ModelSpec& useEagle() - { - mSpecDecodingMode = SpeculativeDecodingMode::Eagle(); - return *this; - } - - ModelSpec& useLookaheadDecoding() - { - mSpecDecodingMode = SpeculativeDecodingMode::LookaheadDecoding(); - return *this; - } - - ModelSpec& useExplicitDraftTokensDecoding() - { - mSpecDecodingMode = SpeculativeDecodingMode::ExplicitDraftTokens(); - return *this; - } - - ModelSpec& useDraftTokensExternalDecoding() - { - mSpecDecodingMode = SpeculativeDecodingMode::DraftTokensExternal(); - return *this; - } - - [[nodiscard]] bool useLogits() const - { - return mGatherLogits || mReplaceLogits; - } - - ModelSpec& useMultipleProfiles() - { - mUseMultipleProfiles = true; - return *this; - } - - ModelSpec& enableContextFMHAFp32Acc() - { - mEnableContextFMHAFp32Acc = true; - return *this; - } - - [[nodiscard]] bool getEnableContextFMHAFp32Acc() const - { - return mEnableContextFMHAFp32Acc; - } - - ModelSpec& setMaxInputLength(SizeType32 maxInputLength) - { - mMaxInputLength = maxInputLength; - return *this; - } - - ModelSpec& setMaxOutputLength(SizeType32 maxOutputLength) - { - mMaxOutputLength = maxOutputLength; - return *this; - } - - ModelSpec& setQuantMethod(QuantMethod quantMethod) - { - mQuantMethod = quantMethod; - return *this; - } - - ModelSpec& useLoraPlugin() - { - mUseLoraPlugin = true; - return *this; - } - - ModelSpec& collectGenerationLogitsFile() - { - mCollectGenerationLogits = true; - return *this; - } - - ModelSpec& collectContextLogitsFile() - { - mCollectContextLogits = true; - return *this; - } - - ModelSpec& collectCumLogProbsFile() - { - mCollectCumLogProbs = true; - return *this; - } - - ModelSpec& collectLogProbsFile() - { - mCollectLogProbs = true; - return *this; - } - - ModelSpec& capacitySchedulerPolicy(tensorrt_llm::executor::CapacitySchedulerPolicy policy) - { - mCapacitySchedulerPolicy = policy; - return *this; - } - - friend std::ostream& operator<<(std::ostream& os, ModelSpec const& modelSpec) - { - return os << modelSpec.getModelPath(); - } - - // Computed properties - [[nodiscard]] std::string getInputFile() const; - - [[nodiscard]] std::string getModelPath() const; - - [[nodiscard]] std::string getResultsFileInternal( - OutputContentType outputContentType = OutputContentType::kNONE) const; - - [[nodiscard]] std::string getResultsFile() const; - [[nodiscard]] std::string getGenerationLogitsFile() const; - - [[nodiscard]] std::string getContextLogitsFile() const; - - [[nodiscard]] std::string getCumLogProbsFile() const; - - [[nodiscard]] std::string getLogProbsFile() const; - - [[nodiscard]] std::string getDtypeString() const; - - [[nodiscard]] std::string getQuantMethodString() const; - - [[nodiscard]] std::string getKVCacheTypeString() const; - - [[nodiscard]] std::string getSpeculativeDecodingModeString() const; - - [[nodiscard]] std::string getCapacitySchedulerString() const; - - static ModelSpec getDefaultModelSpec() - { - static ModelSpec modelSpec{"input_tokens.npy", nvinfer1::DataType::kHALF}; - modelSpec.useGptAttentionPlugin().setKVCacheType(KVCacheType::kPAGED).usePackedInput(); - - return modelSpec; - } - - std::string mInputFile; - nvinfer1::DataType mDataType; - - bool mUseGptAttentionPlugin{false}; - bool mUsePackedInput{false}; - KVCacheType mKVCacheType{KVCacheType::kCONTINUOUS}; - bool mKVCacheReuse{false}; - bool mDecoderPerRequest{false}; - int mPPSize{1}; - int mTPSize{1}; - int mCPSize{1}; - int mMaxDraftTokens{0}; - bool mAcceptDraftByLogits{false}; - bool mUseMambaPlugin{false}; - bool mGatherLogits{false}; - bool mReplaceLogits{false}; - bool mReturnLogProbs{false}; - bool mSmokeTest{false}; - bool mUseMultipleProfiles{false}; - int mMaxInputLength{0}; - int mMaxOutputLength{0}; - bool mUseLoraPlugin{false}; - bool mEnableContextFMHAFp32Acc{false}; - - // Flags to store whether model spec wants collect these outputs, you could call getXXXFile() if you need the name. - bool mCollectGenerationLogits{false}; - bool mCollectContextLogits{false}; - bool mCollectCumLogProbs{false}; - bool mCollectLogProbs{false}; - QuantMethod mQuantMethod{QuantMethod::kNONE}; - - SpeculativeDecodingMode mSpecDecodingMode{SpeculativeDecodingMode::None()}; - - std::optional mCapacitySchedulerPolicy{std::nullopt}; - - // Sometimes, we need to compare with another model spec for golden results. - std::shared_ptr mOtherModelSpecToCompare{nullptr}; -}; - -}; // namespace tensorrt_llm::testing diff --git a/cpp/tensorrt_llm/thop/CMakeLists.txt b/cpp/tensorrt_llm/thop/CMakeLists.txt index c628a9a152c0..b95426f580fe 100644 --- a/cpp/tensorrt_llm/thop/CMakeLists.txt +++ b/cpp/tensorrt_llm/thop/CMakeLists.txt @@ -31,8 +31,11 @@ endif() add_library(th_utils STATIC thUtils.cpp) set_property(TARGET th_utils PROPERTY POSITION_INDEPENDENT_CODE ON) set_property(TARGET th_utils PROPERTY CUDA_RESOLVE_DEVICE_SYMBOLS ON) -target_link_libraries(th_utils PUBLIC ${TORCH_LIBRARIES} ${CUBLAS_LIB} - ${CURAND_LIB}) +# Declare the dependency on the main shared library explicitly so consumers +# (e.g. thUtilsTest) place it after th_utils on the link line; this was +# previously satisfied transitively via the removed TensorRT plugin target. +target_link_libraries(th_utils PUBLIC ${SHARED_TARGET} ${TORCH_LIBRARIES} + ${CUBLAS_LIB} ${CURAND_LIB}) # TODO This does not compile with internal cutlass MOE gemm add_library( diff --git a/cpp/tensorrt_llm/thop/IndexerTopKOp.cpp b/cpp/tensorrt_llm/thop/IndexerTopKOp.cpp index 1981c417dbae..fc29041bbf93 100644 --- a/cpp/tensorrt_llm/thop/IndexerTopKOp.cpp +++ b/cpp/tensorrt_llm/thop/IndexerTopKOp.cpp @@ -19,7 +19,6 @@ #include "tensorrt_llm/kernels/IndexerTopK.h" -// #include // #include // #include // #include diff --git a/cpp/tensorrt_llm/thop/allgatherOp.cpp b/cpp/tensorrt_llm/thop/allgatherOp.cpp index 0d92aa966901..5f7d4571d258 100644 --- a/cpp/tensorrt_llm/thop/allgatherOp.cpp +++ b/cpp/tensorrt_llm/thop/allgatherOp.cpp @@ -20,7 +20,6 @@ #include "tensorrt_llm/runtime/utils/mpiUtils.h" #include "tensorrt_llm/runtime/utils/pgUtils.h" -#include #include #include #include diff --git a/cpp/tensorrt_llm/thop/allreduceOp.cpp b/cpp/tensorrt_llm/thop/allreduceOp.cpp index 4096116fdb42..38cc79379f6d 100644 --- a/cpp/tensorrt_llm/thop/allreduceOp.cpp +++ b/cpp/tensorrt_llm/thop/allreduceOp.cpp @@ -23,6 +23,7 @@ #include "tensorrt_llm/common/ncclUtils.h" #include "tensorrt_llm/common/nvmlWrapper.h" #include "tensorrt_llm/common/opUtils.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/kernels/communicationKernels/MiniMaxReduceRMSKernel.h" #include "tensorrt_llm/kernels/communicationKernels/allReduceFusionKernels.h" #include "tensorrt_llm/kernels/communicationKernels/customLowPrecisionAllReduceKernels.h" @@ -61,7 +62,6 @@ #include #include -// using namespace nvinfer1; using tensorrt_llm::kernels::AllReduceFusionOp; using tensorrt_llm::kernels::AllReduceStrategyType; using tensorrt_llm::mpi::MpiTag; @@ -234,8 +234,8 @@ std::set getLocalGroupTorch(std::set const& group) class AllreduceOp { public: - AllreduceOp( - std::set group, nvinfer1::DataType type, AllReduceStrategyType strategy, AllReduceFusionOp op, float eps) + AllreduceOp(std::set group, tensorrt_llm::DataType type, AllReduceStrategyType strategy, AllReduceFusionOp op, + float eps) : mGroup(std::move(group)) , mIsNVLINKSupported(false) , mIsP2PSupported(false) @@ -248,7 +248,7 @@ class AllreduceOp } AllreduceOp(std::set group, c10::intrusive_ptr const& process_group_, - nvinfer1::DataType type, AllReduceStrategyType strategy, AllReduceFusionOp op, float eps) + tensorrt_llm::DataType type, AllReduceStrategyType strategy, AllReduceFusionOp op, float eps) : mGroup(std::move(group)) , mIsNVLINKSupported(false) , mIsP2PSupported(false) @@ -348,7 +348,7 @@ class AllreduceOp { TORCH_CHECK(norm_weight, "norm_weight is required for residual rms norm allreduce"); TORCH_CHECK(!bias, "bias is not supported for residual rms norm allreduce"); - TORCH_CHECK(mType == nvinfer1::DataType::kHALF || mType == nvinfer1::DataType::kBF16); + TORCH_CHECK(mType == tensorrt_llm::DataType::kHALF || mType == tensorrt_llm::DataType::kBF16); auto [norm_out, ub_buffer1] = torch_ext::create_userbuffers_tensor(input.sizes(), input.scalar_type()); tensorrt_llm::kernels::ub::allreduce2_userbuff_rmsnorm_launcher(ub_buffer0.handle, 0, ub_buffer1.handle, 0, size, hidden_size, nullptr, norm_weight.value().data_ptr(), mEps, residual.value().data_ptr(), @@ -1461,7 +1461,7 @@ class AllreduceOp bool mIsNVLINKSupported; bool mIsP2PSupported; bool mIsMNNVLSupported; - nvinfer1::DataType mType; + tensorrt_llm::DataType mType; AllReduceStrategyType mStrategy; AllReduceFusionOp mOp; float mEps; diff --git a/cpp/tensorrt_llm/thop/attentionOp.cpp b/cpp/tensorrt_llm/thop/attentionOp.cpp index 848554c512d6..b675a99aa775 100644 --- a/cpp/tensorrt_llm/thop/attentionOp.cpp +++ b/cpp/tensorrt_llm/thop/attentionOp.cpp @@ -18,6 +18,7 @@ #include "tensorrt_llm/common/attentionOp.h" #include "tensorrt_llm/common/attentionWorkspace.h" #include "tensorrt_llm/common/dataType.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/kernels/flashMLA/flash_mla.h" #include "tensorrt_llm/kernels/gptKernels.h" #include "tensorrt_llm/kernels/mlaKernels.h" @@ -1123,7 +1124,7 @@ void attention(torch::Tensor q, std::optional k, std::optional k, std::optional>(); } } - else if (dtype == nvinfer1::DataType::kFLOAT) + else if (dtype == tensorrt_llm::DataType::kFLOAT) { TLLM_CHECK(out_dtype == torch::kFloat32); runner = std::make_shared>(); } #ifdef ENABLE_BF16 - else if (dtype == nvinfer1::DataType::kBF16) + else if (dtype == tensorrt_llm::DataType::kBF16) { if (is_fp8_out) { @@ -1168,7 +1169,7 @@ void attention(torch::Tensor q, std::optional k, std::optional(); op->mType = dtype; - op->mFMHAForceFP32Acc = dtype == nvinfer1::DataType::kBF16; + op->mFMHAForceFP32Acc = dtype == tensorrt_llm::DataType::kBF16; op->mLayerIdx = local_layer_idx; op->mNumHeads = num_heads; op->mNumKVHeads = num_kv_heads; @@ -1431,7 +1432,7 @@ bool attention_supports_nvfp4_output(int64_t const num_heads, int64_t const num_ } auto op = std::make_shared(); - op->mType = nvinfer1::DataType::kHALF; + op->mType = tensorrt_llm::DataType::kHALF; op->mNumHeads = num_heads; op->mNumKVHeads = num_kv_heads; op->mHeadSize = head_size; diff --git a/cpp/tensorrt_llm/thop/cublasFp4ScaledMM.cpp b/cpp/tensorrt_llm/thop/cublasFp4ScaledMM.cpp index a9ad46ad8f04..8e4da99dbadc 100644 --- a/cpp/tensorrt_llm/thop/cublasFp4ScaledMM.cpp +++ b/cpp/tensorrt_llm/thop/cublasFp4ScaledMM.cpp @@ -16,7 +16,6 @@ #include "tensorrt_llm/common/cublasMMWrapper.h" #include "tensorrt_llm/common/cudaUtils.h" -#include "tensorrt_llm/plugins/common/plugin.h" #include "tensorrt_llm/thop/outputTensor.h" #include "tensorrt_llm/thop/thUtils.h" #include "userbuffersTensor.h" diff --git a/cpp/tensorrt_llm/thop/cublasScaledMM.cpp b/cpp/tensorrt_llm/thop/cublasScaledMM.cpp index dea6f51363e5..62f51f7b06f9 100644 --- a/cpp/tensorrt_llm/thop/cublasScaledMM.cpp +++ b/cpp/tensorrt_llm/thop/cublasScaledMM.cpp @@ -18,8 +18,6 @@ #include "tensorrt_llm/common/cublasMMWrapper.h" #include "tensorrt_llm/common/cudaUtils.h" #include "tensorrt_llm/kernels/userbuffers/ub_interface.h" -#include "tensorrt_llm/plugins/common/plugin.h" -#include "tensorrt_llm/plugins/gemmPlugin/gemmPlugin.h" #include "tensorrt_llm/runtime/torchUtils.h" #include "tensorrt_llm/thop/outputTensor.h" #include "tensorrt_llm/thop/thUtils.h" diff --git a/cpp/tensorrt_llm/thop/dynamicDecodeOp.cpp b/cpp/tensorrt_llm/thop/dynamicDecodeOp.cpp index 8e9e817bbb51..228b2c614ab7 100644 --- a/cpp/tensorrt_llm/thop/dynamicDecodeOp.cpp +++ b/cpp/tensorrt_llm/thop/dynamicDecodeOp.cpp @@ -16,6 +16,7 @@ #include "tensorrt_llm/thop/dynamicDecodeOp.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/executor/types.h" #include "tensorrt_llm/kernels/decodingCommon.h" #include "tensorrt_llm/runtime/bufferManager.h" @@ -54,7 +55,7 @@ FtDynamicDecode::FtDynamicDecode(size_t const maxBatchSize, size_t const maxB auto bufferManager = std::make_shared(cudaStreamPtr); mFinishedSum = bufferManager->pinnedPool( - tr::ITensor::makeShape({static_cast(maxBatchSize)}), nvinfer1::DataType::kINT32); + tr::ITensor::makeShape({static_cast(maxBatchSize)}), tensorrt_llm::DataType::kINT32); mDynamicDecodeLayer = std::make_shared>(tle::DecodingMode::Auto(), decodingDomain, bufferManager); mBatchSlots = tr::getDefaultBatchSlots(maxBatchSize); diff --git a/cpp/tensorrt_llm/thop/groupRmsNormOp.cpp b/cpp/tensorrt_llm/thop/groupRmsNormOp.cpp index c408a8c286fb..c1f3f41c5a1c 100644 --- a/cpp/tensorrt_llm/thop/groupRmsNormOp.cpp +++ b/cpp/tensorrt_llm/thop/groupRmsNormOp.cpp @@ -17,6 +17,7 @@ #include "tensorrt_llm/common/cudaUtils.h" #include "tensorrt_llm/common/dataType.h" #include "tensorrt_llm/common/opUtils.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/kernels/groupRmsNormKernels/groupRmsNormKernels.h" #include "tensorrt_llm/runtime/torchUtils.h" #include "tensorrt_llm/thop/thUtils.h" @@ -100,9 +101,9 @@ void groupRMSNormBase(torch::TensorList const& inputs, torch::TensorList const& /* Handle dtype conversion */ \ switch (dtype) \ { \ - case torch::ScalarType::Half: params.dtype = nvinfer1::DataType::kHALF; break; \ - case torch::ScalarType::BFloat16: params.dtype = nvinfer1::DataType::kBF16; break; \ - case torch::ScalarType::Float: params.dtype = nvinfer1::DataType::kFLOAT; break; \ + case torch::ScalarType::Half: params.dtype = tensorrt_llm::DataType::kHALF; break; \ + case torch::ScalarType::BFloat16: params.dtype = tensorrt_llm::DataType::kBF16; break; \ + case torch::ScalarType::Float: params.dtype = tensorrt_llm::DataType::kFLOAT; break; \ default: TORCH_CHECK(false, "Unsupported data type"); \ } \ tensorrt_llm::kernels::group_rms_norm::GroupRMSNormBaseKernelLauncher(params); \ @@ -181,9 +182,9 @@ void groupRMSNormLargeBatch(torch::TensorList const& inputs, torch::TensorList c // Handle dtype conversion switch (dtype) { - case torch::ScalarType::Half: params.dtype = nvinfer1::DataType::kHALF; break; - case torch::ScalarType::BFloat16: params.dtype = nvinfer1::DataType::kBF16; break; - case torch::ScalarType::Float: params.dtype = nvinfer1::DataType::kFLOAT; break; + case torch::ScalarType::Half: params.dtype = tensorrt_llm::DataType::kHALF; break; + case torch::ScalarType::BFloat16: params.dtype = tensorrt_llm::DataType::kBF16; break; + case torch::ScalarType::Float: params.dtype = tensorrt_llm::DataType::kFLOAT; break; default: TORCH_CHECK(false, "Unsupported data type"); } @@ -260,9 +261,9 @@ void groupRMSNormHeuristic(torch::TensorList const& inputs, torch::TensorList co /* Handle dtype conversion */ \ switch (dtype) \ { \ - case torch::ScalarType::Half: params.dtype = nvinfer1::DataType::kHALF; break; \ - case torch::ScalarType::BFloat16: params.dtype = nvinfer1::DataType::kBF16; break; \ - case torch::ScalarType::Float: params.dtype = nvinfer1::DataType::kFLOAT; break; \ + case torch::ScalarType::Half: params.dtype = tensorrt_llm::DataType::kHALF; break; \ + case torch::ScalarType::BFloat16: params.dtype = tensorrt_llm::DataType::kBF16; break; \ + case torch::ScalarType::Float: params.dtype = tensorrt_llm::DataType::kFLOAT; break; \ default: TORCH_CHECK(false, "Unsupported data type"); \ } \ \ diff --git a/cpp/tensorrt_llm/thop/loraOp.cpp b/cpp/tensorrt_llm/thop/loraOp.cpp index b35ca2608625..6957987be6b8 100644 --- a/cpp/tensorrt_llm/thop/loraOp.cpp +++ b/cpp/tensorrt_llm/thop/loraOp.cpp @@ -18,6 +18,7 @@ #include "tensorrt_llm/common/cublasMMWrapper.h" #include "tensorrt_llm/common/cudaUtils.h" #include "tensorrt_llm/common/opUtils.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/kernels/cuda_graph_grouped_gemm.h" #include "tensorrt_llm/kernels/lora/lora.h" #include "tensorrt_llm/kernels/lora/loraGroupGEMMParamFillRowReorderFusion.h" @@ -151,11 +152,11 @@ std::vector lora_grouped_gemm(th::Tensor const& input, th::Tensor co { outHiddenSizes[i] = output_hidden_sizes[i]; } - nvinfer1::DataType loraRuntimeDataType; + tensorrt_llm::DataType loraRuntimeDataType; switch (input.scalar_type()) { - case torch::kFloat16: loraRuntimeDataType = nvinfer1::DataType::kHALF; break; - case torch::kBFloat16: loraRuntimeDataType = nvinfer1::DataType::kBF16; break; + case torch::kFloat16: loraRuntimeDataType = tensorrt_llm::DataType::kHALF; break; + case torch::kBFloat16: loraRuntimeDataType = tensorrt_llm::DataType::kBF16; break; default: throw std::invalid_argument("Invalid dtype, only supports float16, bfloat16"); } @@ -221,11 +222,11 @@ void lora_grouped_gemm_cuda_graph(th::Tensor const& lora_in_sizes, // [layer_mod auto* splitk_offsets_gpu = reinterpret_cast(const_cast(splitk_offsets.data_ptr())); // Get data type - nvinfer1::DataType loraRuntimeDataType; + tensorrt_llm::DataType loraRuntimeDataType; switch (dtype) { - case torch::kFloat16: loraRuntimeDataType = nvinfer1::DataType::kHALF; break; - case torch::kBFloat16: loraRuntimeDataType = nvinfer1::DataType::kBF16; break; + case torch::kFloat16: loraRuntimeDataType = tensorrt_llm::DataType::kHALF; break; + case torch::kBFloat16: loraRuntimeDataType = tensorrt_llm::DataType::kBF16; break; default: TORCH_CHECK(false, "Invalid dtype, only supports float16, bfloat16, got %s", c10::toString(dtype)); } @@ -301,11 +302,11 @@ void lora_group_gemm_param_fill_row_reorder_fusion(th::Tensor const& in_sizes, / int32_t const module_count = static_cast(in_sizes.size(0)); // Get data type info - nvinfer1::DataType loraRuntimeDataType; + tensorrt_llm::DataType loraRuntimeDataType; switch (dtype) { - case torch::kFloat16: loraRuntimeDataType = nvinfer1::DataType::kHALF; break; - case torch::kBFloat16: loraRuntimeDataType = nvinfer1::DataType::kBF16; break; + case torch::kFloat16: loraRuntimeDataType = tensorrt_llm::DataType::kHALF; break; + case torch::kBFloat16: loraRuntimeDataType = tensorrt_llm::DataType::kBF16; break; default: TORCH_CHECK(false, "Invalid dtype, only supports float16, bfloat16, got %s", c10::toString(dtype)); } diff --git a/cpp/tensorrt_llm/thop/moeAlltoAllOp.cpp b/cpp/tensorrt_llm/thop/moeAlltoAllOp.cpp index 373f936c4c5d..e985eb943ee1 100644 --- a/cpp/tensorrt_llm/thop/moeAlltoAllOp.cpp +++ b/cpp/tensorrt_llm/thop/moeAlltoAllOp.cpp @@ -15,6 +15,7 @@ */ #include "tensorrt_llm/common/envUtils.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h" #include "tensorrt_llm/runtime/utils/mpiUtils.h" #include "tensorrt_llm/thop/moeAlltoAllMeta.h" @@ -484,20 +485,20 @@ torch::Tensor moeA2ACombineOp(torch::Tensor const& payload, int64_t localNumToke TORCH_CHECK(epRank >= 0 && epRank < epSize, "epRank must be in the range [0, epSize)"); TORCH_CHECK(topK > 0 && topK <= kMaxTopK, "topK must be in the range (0, kMaxTopK]"); - // Map torch dtype to nvinfer1::DataType - nvinfer1::DataType nvDtype = nvinfer1::DataType::kFLOAT; + // Map torch dtype to tensorrt_llm::DataType + tensorrt_llm::DataType nvDtype = tensorrt_llm::DataType::kFLOAT; auto scalarType = payload.scalar_type(); if (scalarType == at::kHalf) { - nvDtype = nvinfer1::DataType::kHALF; + nvDtype = tensorrt_llm::DataType::kHALF; } else if (scalarType == at::kBFloat16) { - nvDtype = nvinfer1::DataType::kBF16; + nvDtype = tensorrt_llm::DataType::kBF16; } else if (scalarType == at::kFloat) { - nvDtype = nvinfer1::DataType::kFLOAT; + nvDtype = tensorrt_llm::DataType::kFLOAT; } else { diff --git a/cpp/tensorrt_llm/thop/moeOp.cpp b/cpp/tensorrt_llm/thop/moeOp.cpp index 4a938455488b..81c975241a2e 100644 --- a/cpp/tensorrt_llm/thop/moeOp.cpp +++ b/cpp/tensorrt_llm/thop/moeOp.cpp @@ -30,6 +30,7 @@ #include "tensorrt_llm/common/config.h" #include "tensorrt_llm/common/dataType.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/common/workspace.h" #include "tensorrt_llm/kernels/cuda_graph_grouped_gemm.h" #include "tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_gemm.h" @@ -87,7 +88,8 @@ enum class MoeLoraRequestType : int32_t // --------------------------------------------------------------------------- inline void moeLoraGroupedGemmRunImpl(::tensorrt_llm::kernels::cutlass_kernels::MoeLoraGroupedGemmModule const& mod, int64_t num_permuted_tokens, int64_t in_hidden_size, int64_t max_lora_rank, int64_t dtype_bytes, - int64_t splitk_slices, void const* input_base, void* output_base, nvinfer1::DataType data_type, cudaStream_t stream) + int64_t splitk_slices, void const* input_base, void* output_base, tensorrt_llm::DataType data_type, + cudaStream_t stream) { TLLM_CHECK_WITH_INFO(mod.permuted_ranks_dev != nullptr, "Grouped-GEMM LoRA module is missing permuted ranks buffer (forgot to populate grouped_gemm?)."); @@ -1206,17 +1208,17 @@ class FusedMoeRunner : public torch::CustomClassHolder // ===== LoRA helpers ===== - // Map a torch dtype to the TRT-LLM nvinfer1::DataType used to size the + // Map a torch dtype to the TRT-LLM tensorrt_llm::DataType used to size the // grouped-GEMM low-rank scratch. Kept as a const member (not static) so the // FP8 case can read mOutputDtype to pick the fp16/bf16 LoRA compute dtype. - nvinfer1::DataType loraTypeFromActDtype(c10::ScalarType dtype) const + tensorrt_llm::DataType loraTypeFromActDtype(c10::ScalarType dtype) const { switch (dtype) { - case c10::ScalarType::Half: return nvinfer1::DataType::kHALF; - case c10::ScalarType::Float: return nvinfer1::DataType::kFLOAT; + case c10::ScalarType::Half: return tensorrt_llm::DataType::kHALF; + case c10::ScalarType::Float: return tensorrt_llm::DataType::kFLOAT; #ifdef ENABLE_BF16 - case c10::ScalarType::BFloat16: return nvinfer1::DataType::kBF16; + case c10::ScalarType::BFloat16: return tensorrt_llm::DataType::kBF16; #endif #ifdef ENABLE_FP8 case c10::ScalarType::Float8_e4m3fn: diff --git a/cpp/tensorrt_llm/thop/noAuxTcOp.cpp b/cpp/tensorrt_llm/thop/noAuxTcOp.cpp index e445206e1d78..4dfb20072734 100644 --- a/cpp/tensorrt_llm/thop/noAuxTcOp.cpp +++ b/cpp/tensorrt_llm/thop/noAuxTcOp.cpp @@ -20,7 +20,6 @@ #include "tensorrt_llm/kernels/noAuxTcKernels.h" -// #include // #include // #include // #include diff --git a/cpp/tensorrt_llm/thop/reducescatterOp.cpp b/cpp/tensorrt_llm/thop/reducescatterOp.cpp index 40f89e40ff75..a50ca1862f76 100644 --- a/cpp/tensorrt_llm/thop/reducescatterOp.cpp +++ b/cpp/tensorrt_llm/thop/reducescatterOp.cpp @@ -20,7 +20,6 @@ #include "tensorrt_llm/runtime/utils/mpiUtils.h" #include "tensorrt_llm/runtime/utils/pgUtils.h" -#include #include #include #if ENABLE_MULTI_DEVICE diff --git a/cpp/tensorrt_llm/thop/thUtils.cpp b/cpp/tensorrt_llm/thop/thUtils.cpp index 97fe6acaab7b..c151414127fa 100644 --- a/cpp/tensorrt_llm/thop/thUtils.cpp +++ b/cpp/tensorrt_llm/thop/thUtils.cpp @@ -15,7 +15,7 @@ */ #include "tensorrt_llm/thop/thUtils.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include TRTLLM_NAMESPACE_BEGIN @@ -25,12 +25,12 @@ namespace torch_ext tensorrt_llm::runtime::ITensor::Shape convert_shape(torch::Tensor tensor) { - constexpr auto trtMaxDims = nvinfer1::Dims::MAX_DIMS; + constexpr auto trtMaxDims = tensorrt_llm::Dims::MAX_DIMS; auto const torchTensorNumDims = tensor.dim(); TLLM_CHECK_WITH_INFO(torchTensorNumDims <= trtMaxDims, "TensorRT supports at most %i tensor dimensions. Found a Torch tensor with %li dimensions.", trtMaxDims, torchTensorNumDims); - auto result = nvinfer1::Dims{}; + auto result = tensorrt_llm::Dims{}; result.nbDims = static_cast(torchTensorNumDims); for (int i = 0; i < torchTensorNumDims; i++) { diff --git a/cpp/tensorrt_llm/thop/trtllmGenQKVProcessOp.cpp b/cpp/tensorrt_llm/thop/trtllmGenQKVProcessOp.cpp index 5720951e2720..bc80d8446445 100644 --- a/cpp/tensorrt_llm/thop/trtllmGenQKVProcessOp.cpp +++ b/cpp/tensorrt_llm/thop/trtllmGenQKVProcessOp.cpp @@ -18,6 +18,7 @@ #include "tensorrt_llm/common/attentionOp.h" #include "tensorrt_llm/common/cudaUtils.h" #include "tensorrt_llm/common/quantization.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/kernels/gptKernels.h" #include "tensorrt_llm/kernels/kvCacheUtils.h" #include "tensorrt_llm/kernels/unfusedAttentionKernels.h" @@ -401,16 +402,16 @@ trtllmGenContextPreprocess(torch::Tensor qkv_input, torch::Tensor workspace, tor switch (qkvDtype) { - case nvinfer1::DataType::kFLOAT: + case tensorrt_llm::DataType::kFLOAT: tensorrt_llm::kernels::invokeQKVPreprocessing( reinterpret_cast&>(qkvParams), stream); break; - case nvinfer1::DataType::kHALF: + case tensorrt_llm::DataType::kHALF: tensorrt_llm::kernels::invokeQKVPreprocessing( reinterpret_cast&>(qkvParams), stream); break; #ifdef ENABLE_BF16 - case nvinfer1::DataType::kBF16: + case tensorrt_llm::DataType::kBF16: tensorrt_llm::kernels::invokeQKVPreprocessing( reinterpret_cast&>(qkvParams), stream); break; @@ -549,16 +550,16 @@ void trtllmGenContextPostprocess(torch::Tensor qkv_input, torch::Tensor workspac switch (qkvDtype) { - case nvinfer1::DataType::kFLOAT: + case tensorrt_llm::DataType::kFLOAT: tensorrt_llm::kernels::invokeKvCachePostprocessing( reinterpret_cast&>(qkvParams), stream); break; - case nvinfer1::DataType::kHALF: + case tensorrt_llm::DataType::kHALF: tensorrt_llm::kernels::invokeKvCachePostprocessing( reinterpret_cast&>(qkvParams), stream); break; #ifdef ENABLE_BF16 - case nvinfer1::DataType::kBF16: + case tensorrt_llm::DataType::kBF16: tensorrt_llm::kernels::invokeKvCachePostprocessing( reinterpret_cast&>(qkvParams), stream); break; @@ -733,16 +734,16 @@ trtllmGenGenerationPreprocess(torch::Tensor qkv_input, torch::Tensor workspace, switch (qkvDtype) { - case nvinfer1::DataType::kFLOAT: + case tensorrt_llm::DataType::kFLOAT: tensorrt_llm::kernels::invokeQKVPreprocessing( reinterpret_cast&>(qkvParams), stream); break; - case nvinfer1::DataType::kHALF: + case tensorrt_llm::DataType::kHALF: tensorrt_llm::kernels::invokeQKVPreprocessing( reinterpret_cast&>(qkvParams), stream); break; #ifdef ENABLE_BF16 - case nvinfer1::DataType::kBF16: + case tensorrt_llm::DataType::kBF16: tensorrt_llm::kernels::invokeQKVPreprocessing( reinterpret_cast&>(qkvParams), stream); break; diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index 0a06d40ee85e..99e16ea8b147 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -25,8 +25,13 @@ include_directories( ${PROJECT_SOURCE_DIR}/include ${cutlass_source_dir}/include ${cutlass_source_dir}/tools/util/include - ${PROJECT_SOURCE_DIR}/tests/batch_manager - ${PROJECT_SOURCE_DIR}/tests/utils) + ${PROJECT_SOURCE_DIR}/tests/batch_manager) + +# Tests previously inherited the MPI include dirs transitively through the +# removed TensorRT plugin target's PUBLIC includes. +if(ENABLE_MULTI_DEVICE) + include_directories(${MPI_C_INCLUDE_DIRS}) +endif() set(TOP_LEVEL_DIR "${PROJECT_SOURCE_DIR}/..") @@ -38,13 +43,12 @@ function(add_gtest test_name test_src) ${ARGN}) add_executable(${test_name} ${test_src}) - target_link_libraries(${test_name} PUBLIC gmock_main TensorRT::OnnxParser) + target_link_libraries(${test_name} PUBLIC gmock_main) if(NOT ARGS_NO_GTEST_MAIN) target_link_libraries(${test_name} PUBLIC gtest_main) endif() if(NOT ARGS_NO_TLLM_LINKAGE) - target_link_libraries(${test_name} PUBLIC ${SHARED_TARGET} - nvinfer_plugin_tensorrt_llm) + target_link_libraries(${test_name} PUBLIC ${SHARED_TARGET}) if(WIN32) target_link_libraries(${test_name} PRIVATE context_attention_src) endif() @@ -66,6 +70,4 @@ function(add_gtest test_name test_src) add_dependencies(google-tests ${test_name}) endfunction() -add_subdirectory(utils) add_subdirectory(unit_tests) -add_subdirectory(e2e_tests) diff --git a/cpp/tests/e2e_tests/CMakeLists.txt b/cpp/tests/e2e_tests/CMakeLists.txt deleted file mode 100644 index f5deb048a180..000000000000 --- a/cpp/tests/e2e_tests/CMakeLists.txt +++ /dev/null @@ -1,17 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2023-2025 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. - -add_subdirectory(batch_manager) -add_subdirectory(executor) diff --git a/cpp/tests/e2e_tests/batch_manager/CMakeLists.txt b/cpp/tests/e2e_tests/batch_manager/CMakeLists.txt deleted file mode 100644 index 875e12eb975f..000000000000 --- a/cpp/tests/e2e_tests/batch_manager/CMakeLists.txt +++ /dev/null @@ -1,22 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2023-2025 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. - -# guidedDecoderTest requires model tokenizer info, so it's easier to run it with -# e2e tests instead of unit tests. -add_gtest(guidedDecoderTest guidedDecoderTest.cpp) -add_gtest(trtEncoderModelTest trtEncoderModelTest.cpp) -add_gtest(trtGptModelTest trtGptModelTest.cpp) -add_gtest(trtGptModelRealDecoderTest trtGptModelRealDecoderTest.cpp) -target_link_libraries(trtGptModelRealDecoderTest PRIVATE testingUtils) diff --git a/cpp/tests/e2e_tests/batch_manager/guidedDecoderTest.cpp b/cpp/tests/e2e_tests/batch_manager/guidedDecoderTest.cpp deleted file mode 100644 index 7b262cacb27d..000000000000 --- a/cpp/tests/e2e_tests/batch_manager/guidedDecoderTest.cpp +++ /dev/null @@ -1,227 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2024-2025 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. - */ - -#ifndef TOP_LEVEL_DIR -#error "Define TOP_LEVEL_DIR" -#endif - -#include -#include -#include - -#include "tensorrt_llm/batch_manager/common.h" -#include "tensorrt_llm/batch_manager/decoderBuffers.h" -#include "tensorrt_llm/batch_manager/guidedDecoder.h" -#include "tensorrt_llm/batch_manager/llmRequest.h" -#include "tensorrt_llm/executor/executor.h" - -using namespace tensorrt_llm::runtime; -using namespace tensorrt_llm::batch_manager; -namespace texec = tensorrt_llm::executor; - -namespace -{ -auto const TEST_RESOURCE_PATH = std::filesystem::path{TOP_LEVEL_DIR} / "cpp/tests/resources"; -auto const DATA_PATH = TEST_RESOURCE_PATH / "data"; -auto const GPT_XGRAMMAR_TOKENIZER_INFO_PATH = DATA_PATH / "gpt2" / "xgrammar_tokenizer_info.json"; -auto const LLAMA_XGRAMMAR_TOKENIZER_INFO_PATH = DATA_PATH / "Llama-3.2-1B" / "xgrammar_tokenizer_info.json"; -} // namespace - -class GuidedDecoderTest : public ::testing::Test -{ -public: - using TensorPtr = ITensor::SharedPtr; - using VecTokens = std::vector; - using RequestIdType = std::uint64_t; - using RequestVector = std::vector>; - - void SetUp() override - { - mStream = std::make_shared(); - mRuntimeBufferManager = std::make_shared(mStream); - } - - void TearDown() override {} - - void initData(std::filesystem::path tokenizerInfoPath, SizeType32 vocabSizePadded, VecTokens outputIds, - std::vector expectedNumRejected) - { - mLogitsDtype = nvinfer1::DataType::kFLOAT; - mMaxNumRequests = 16; - - mVocabSizePadded = vocabSizePadded; - auto const tokenizerInfo = nlohmann::json::parse(std::ifstream{tokenizerInfoPath}); - auto const encodedVocab = tokenizerInfo["encoded_vocab"].template get>(); - auto const tokenizerStr = tokenizerInfo["tokenizer_str"].template get(); - auto const stopTokenIds = tokenizerInfo["stop_token_ids"].template get>(); - texec::GuidedDecodingConfig guidedDecodingConfig( - texec::GuidedDecodingConfig::GuidedDecodingBackend::kXGRAMMAR, encodedVocab, tokenizerStr, stopTokenIds); - mGuidedDecoder = std::make_shared( - guidedDecodingConfig, mMaxNumRequests, mVocabSizePadded, mLogitsDtype, *mRuntimeBufferManager); - - mLogits.resize(mMaxNumRequests); - mLogitsHost.resize(mMaxNumRequests); - for (int i = 0; i < mMaxNumRequests; i++) - { - mLogits[i] = mRuntimeBufferManager->gpu(ITensor::makeShape({mVocabSizePadded}), mLogitsDtype); - mLogitsHost[i] = BufferManager::pinned(ITensor::makeShape({mVocabSizePadded}), mLogitsDtype); - } - - mOutputIds = outputIds; - mExpectedNumRejected = expectedNumRejected; - } - - void resetLogits() - { - for (int i = 0; i < mMaxNumRequests; i++) - { - auto logitsHostData = bufferCast(*mLogitsHost[i]); - for (int j = 0; j < mVocabSizePadded; j++) - { - logitsHostData[j] = 0.0f; - } - mRuntimeBufferManager->copy(*(mLogitsHost[i]), *(mLogits[i])); - } - } - - void syncLogitsToHost() - { - for (int i = 0; i < mMaxNumRequests; i++) - { - mRuntimeBufferManager->copy(*(mLogits[i]), *(mLogitsHost[i])); - } - } - - int32_t countRejected(int i) - { - int32_t numRejected = 0; - for (int j = 0; j < mVocabSizePadded; j++) - { - auto logitsHostData = bufferCast(*mLogitsHost[i]); - if (logitsHostData[j] < -1e6) - { - numRejected++; - } - } - return numRejected; - } - - void runTest() - { - auto llmReq1 = std::make_shared(1, 100, std::make_shared(10), SamplingConfig(), false); - texec::GuidedDecodingParams guidedDecodingParams(texec::GuidedDecodingParams::GuideType::kJSON); - llmReq1->setGuidedDecodingParams(guidedDecodingParams); - llmReq1->mSeqSlot = 1; - - auto llmReq2 = std::make_shared(1, 100, std::make_shared(10), SamplingConfig(), false); - llmReq2->mSeqSlot = 2; - - RequestVector contextRequests{llmReq1, llmReq2}; - RequestVector generationRequests{}; - ScheduledRequests scheduledRequests{contextRequests, generationRequests}; - DecoderInputBuffers decoderInputBuffers(mMaxNumRequests, 1, *mRuntimeBufferManager); - - for (auto const& requests : {scheduledRequests.contextRequests, scheduledRequests.generationRequests}) - { - for (auto const& llmReq : requests) - { - decoderInputBuffers.decoderRequests.push_back(llmReq); - } - } - decoderInputBuffers.decoderLogits = mLogits; - - // Context phase - resetLogits(); - mGuidedDecoder->build(scheduledRequests); - mGuidedDecoder->execute(decoderInputBuffers, *mRuntimeBufferManager); - syncLogitsToHost(); - mRuntimeBufferManager->getStream().synchronize(); - - // Move request to generation phase - contextRequests.pop_back(); - contextRequests.pop_back(); - llmReq1->setState(LlmRequestState::kGENERATION_IN_PROGRESS); - generationRequests.push_back(llmReq1); - llmReq2->setState(LlmRequestState::kGENERATION_IN_PROGRESS); - generationRequests.push_back(llmReq2); - - decoderInputBuffers.decoderRequests.clear(); - for (auto const& requests : {scheduledRequests.contextRequests, scheduledRequests.generationRequests}) - { - for (auto const& llmReq : requests) - { - decoderInputBuffers.decoderRequests.push_back(llmReq); - } - } - - EXPECT_EQ(countRejected(0), mExpectedNumRejected[0]); - EXPECT_EQ(countRejected(1), 0); - - // Generation phase - for (int i = 0; i < mOutputIds.size(); i++) - { - llmReq1->addNewToken(mOutputIds[i], 0); - llmReq2->addNewToken(mOutputIds[i], 0); - - resetLogits(); - mGuidedDecoder->build(scheduledRequests); - mGuidedDecoder->execute(decoderInputBuffers, *mRuntimeBufferManager); - syncLogitsToHost(); - mRuntimeBufferManager->getStream().synchronize(); - - EXPECT_EQ(countRejected(0), mExpectedNumRejected[i + 1]); - EXPECT_EQ(countRejected(1), 0); - } - } - -private: - SizeType32 mMaxNumRequests; - SizeType32 mVocabSizePadded; - nvinfer1::DataType mLogitsDtype; - - std::vector mLogits; // [mBatchSize, mVocabSizePadded] - std::vector mLogitsHost; // [mBatchSize, mVocabSizePadded] - - std::shared_ptr mRuntimeBufferManager; - std::shared_ptr mStream; - std::shared_ptr mGuidedDecoder; - - VecTokens mOutputIds; - std::vector mExpectedNumRejected; -}; - -TEST_F(GuidedDecoderTest, GptTokenizer) -{ - VecTokens outputIds{4895, 824, 312, 1298, 366, 27743, 7934, 49793, 1600, 366, 12961, 19703, 4668, 1298, 366, 54, - 4537, 17, 12, 17469, 7919, 1600, 366, 3903, 10394, 1298, 366, 1485, 405, 41022, 20662}; - std::vector expectedNumRejected{50251, 219, 219, 219, 48558, 219, 219, 219, 219, 50191, 219, 219, 219, 219, - 48558, 219, 219, 219, 219, 219, 219, 219, 50191, 219, 219, 219, 48558, 219, 219, 219, 219, 50256}; - initData(GPT_XGRAMMAR_TOKENIZER_INFO_PATH, 50257, outputIds, expectedNumRejected); - runTest(); -} - -TEST_F(GuidedDecoderTest, LlamaTokenizer) -{ - VecTokens outputIds{6377, 893, 333, 1115, 376, 27247, 6779, 7898, 545, 613, 376, 8926, 17830, 1115, 376, 29956, - 7228, 29906, 29899, 10399, 7734, 613, 376, 4980, 2103, 1115, 376, 29896, 29941, 29900, 29900, 341, 29890, 567, - 9092}; - std::vector expectedNumRejected{128235, 128235, 128235, 128235, 128235, 128235, 128235, 128235, 128235, - 128235, 128235, 128235, 128235, 128235, 128235, 128235, 128235, 128235, 128235, 128235, 128235, 128235, 128235, - 128235, 128235, 128235, 128235, 128235, 128235, 128235, 128235, 128235, 128235, 128235, 128235, 128235}; - initData(LLAMA_XGRAMMAR_TOKENIZER_INFO_PATH, 128256, outputIds, expectedNumRejected); - runTest(); -} diff --git a/cpp/tests/e2e_tests/batch_manager/trtEncoderModelTest.cpp b/cpp/tests/e2e_tests/batch_manager/trtEncoderModelTest.cpp deleted file mode 100644 index dce09539bd45..000000000000 --- a/cpp/tests/e2e_tests/batch_manager/trtEncoderModelTest.cpp +++ /dev/null @@ -1,219 +0,0 @@ - -/* - * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. - * - * 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. - */ - -#ifndef TOP_LEVEL_DIR -#error "Define TOP_LEVEL_DIR" -#endif - -#include "tensorrt_llm/batch_manager/trtEncoderModel.h" -#include "tensorrt_llm/executor/executor.h" -#include "tensorrt_llm/plugins/api/tllmPlugin.h" -#include "tensorrt_llm/runtime/gptJsonConfig.h" -#include "tensorrt_llm/runtime/rawEngine.h" -#include "tensorrt_llm/runtime/tllmLogger.h" -#include "tensorrt_llm/runtime/utils/numpyUtils.h" -#include "tensorrt_llm/runtime/utils/runtimeUtils.h" - -#include -#include - -#include -#include - -using namespace tensorrt_llm::runtime; -namespace fs = std::filesystem; - -using TensorPtr = ITensor::SharedPtr; - -namespace -{ -auto const TEST_RESOURCE_PATH = fs::path{TOP_LEVEL_DIR} / "cpp/tests/resources"; -auto const ENC_DEC_BASE = TEST_RESOURCE_PATH / "models/enc_dec/trt_engines"; -auto const ENC_DEC_ENGINE_BASE = TEST_RESOURCE_PATH / "models/enc_dec/trt_engines"; -auto const BART_TP1_PP1_ENCODER_RMPAD_DIR = "bart-large-cnn/1-gpu/float16/tp1/encoder"; -auto const BART_TP2_PP1_ENCODER_RMPAD_DIR = "bart-large-cnn/2-gpu/float16/tp2/encoder"; -auto const BART_TP2_PP2_ENCODER_RMPAD_DIR = "bart-large-cnn/4-gpu/float16/tp2/encoder"; -auto const T5_TP1_PP1_ENCODER_RMPAD_DIR = "t5-small/1-gpu/float16/tp1/encoder"; -auto const ENC_DEC_DATA_BASE = TEST_RESOURCE_PATH / "data/enc_dec"; -} // namespace - -namespace tensorrt_llm::batch_manager -{ - -class EncoderModelTestSingleGPU : public ::testing::Test // NOLINT(cppcoreguidelines-pro-type-member-init) -{ -protected: - EncoderModelTestSingleGPU(std::filesystem::path const& modelPath) - : mModelConfig(1, 2, 1, 1, 1, 1, nvinfer1::DataType::kFLOAT) - , mModelPath(modelPath) - { - } - - EncoderModelTestSingleGPU() - : EncoderModelTestSingleGPU(ENC_DEC_ENGINE_BASE / T5_TP1_PP1_ENCODER_RMPAD_DIR) - { - } - - void SetUp() override - { - std::filesystem::path trtEnginePath = mModelPath; - - mBeamWidth = 1; - - mLogger = std::make_shared(); - - initTrtLlmPlugins(mLogger.get()); - - auto const json = GptJsonConfig::parse(trtEnginePath / "config.json"); - mModelConfig = json.getModelConfig(); - mWorldConfig = WorldConfig::mpi(json.getGpusPerNode(), json.getTensorParallelism(), - json.getPipelineParallelism(), json.getContextParallelism()); - mVocabSizePadded = mModelConfig.getVocabSizePadded(mWorldConfig.getSize()); - - auto const enginePath = trtEnginePath / json.engineFilename(mWorldConfig); - auto const dtype = mModelConfig.getDataType(); - - ASSERT_TRUE(fs::exists(enginePath)); - mEngineBuffer = utils::loadEngine(enginePath.string()); - - mStream = std::make_unique(); - mManager = std::make_unique(mStream); - } - - void TearDown() override {} - - int32_t mMaxNumRequests; - int32_t mMaxSeqLen; - int32_t mBeamWidth; - int32_t mVocabSizePadded; - // SamplingConfig mSamplingConfig; - std::string mDataPath; - std::shared_ptr mLogger; - ModelConfig mModelConfig; - WorldConfig mWorldConfig; - std::vector mEngineBuffer; - std::unique_ptr mManager; - BufferManager::CudaStreamPtr mStream; - std::filesystem::path mModelPath; -}; - -// test for TP2PP2 -class TrtEncoderModelTestMultiGPU : public EncoderModelTestSingleGPU -{ -protected: - TrtEncoderModelTestMultiGPU() - : EncoderModelTestSingleGPU(ENC_DEC_ENGINE_BASE / BART_TP2_PP2_ENCODER_RMPAD_DIR) - { - } -}; - -namespace -{ - -void runEncoderTest(std::unique_ptr& bufferManager, ModelConfig const& modelConfig, - WorldConfig const& worldConfig, std::vector const& engineBuffer, - std::shared_ptr& logger) -{ - using VecTokens = LlmRequest::VecTokens; - using TokenIdType = LlmRequest::TokenIdType; - - auto inputsIdsHost - = utils::loadNpy(*bufferManager, (ENC_DEC_DATA_BASE / "input_ids.npy").string(), MemoryType::kCPU); - auto inputsIdsPtr = bufferCast(*inputsIdsHost); - auto inputLengthsHost - = utils::loadNpy(*bufferManager, (ENC_DEC_DATA_BASE / "input_lengths.npy").string(), MemoryType::kCPU); - auto inputLengthsPtr = bufferCast(*inputLengthsHost); - auto encoderOutput - = utils::loadNpy(*bufferManager, (ENC_DEC_DATA_BASE / "encoder_output.npy").string(), MemoryType::kCPU); - auto encoderOutputPrt = bufferCast(*encoderOutput); - - SizeType32 const nbRequests = inputLengthsHost->getShape().d[0]; - SizeType32 const stride = inputsIdsHost->getShape().d[1]; - SizeType32 const hiddenSize = encoderOutput->getShape().d[1]; - ASSERT_EQ(nbRequests, inputsIdsHost->getShape().d[0]); - - // std::vector> inputIds(nbRequests); - RequestVector requestList; - for (SizeType32 i = 0; i < nbRequests; i++) - { - SizeType32 length = inputLengthsPtr[i]; - auto currentInputId = std::make_shared(0); - currentInputId->insert(currentInputId->end(), inputsIdsPtr, inputsIdsPtr + length); - executor::Request req(*currentInputId, 1); - req.setEncoderInputTokenIds(*currentInputId); - auto request = std::make_shared(i, req); - inputsIdsPtr += stride; - requestList.push_back(request); - } - - tensorrt_llm::executor::ExecutorConfig executorConfig{}; - auto trtEncoderModel = std::make_shared( - modelConfig, worldConfig, runtime::RawEngine(engineBuffer.data(), engineBuffer.size()), logger, executorConfig); - - trtEncoderModel->forward(requestList); - - if (worldConfig.isLastPipelineParallelRank() && worldConfig.getTensorParallelRank() == 0) - { - auto arrayEqual = [](auto it0, auto it1, SizeType32 length) - { - SizeType32 nbNotEqual = 0; - for (SizeType32 i = 0; i < length; i++) - { - auto v0 = static_cast(*it0); - auto v1 = static_cast(*it1); - if (std::abs(v0 - v1) > 1e-3) - { - nbNotEqual++; - } - it0++; - it1++; - } - return static_cast(nbNotEqual) / length; - }; - ASSERT_EQ(requestList.size(), inputLengthsHost->getShape().d[0]); - { - auto curLengthPtr = inputLengthsPtr; - auto curOutPtr = encoderOutputPrt; - for (auto const& req : requestList) - { - ASSERT_TRUE(req->getEncoderOutputHost()) << "Encoder output is empty!"; - EXPECT_EQ(req->getState(), LlmRequestState::kCONTEXT_INIT); - auto actualOut = bufferCast(*(req->getEncoderOutputHost())); - auto unequalFraction = arrayEqual(curOutPtr, actualOut, *curLengthPtr); - EXPECT_TRUE(unequalFraction == 0) - << "Req " << req->mRequestId << ": " << unequalFraction << " of outputs are different"; - curOutPtr += *curLengthPtr * hiddenSize; - curLengthPtr++; - } - } - } -} - -} // Anonymous namespace - -TEST_F(EncoderModelTestSingleGPU, Forward) -{ - runEncoderTest(mManager, mModelConfig, mWorldConfig, mEngineBuffer, mLogger); -} - -TEST_F(TrtEncoderModelTestMultiGPU, Forward) -{ - - runEncoderTest(mManager, mModelConfig, mWorldConfig, mEngineBuffer, mLogger); -} - -} // namespace tensorrt_llm::batch_manager diff --git a/cpp/tests/e2e_tests/batch_manager/trtGptModelRealDecoderTest.cpp b/cpp/tests/e2e_tests/batch_manager/trtGptModelRealDecoderTest.cpp deleted file mode 100644 index f401a31305d3..000000000000 --- a/cpp/tests/e2e_tests/batch_manager/trtGptModelRealDecoderTest.cpp +++ /dev/null @@ -1,1741 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2023-2025 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. - */ - -#include "tensorrt_llm/batch_manager/trtGptModel.h" -#include "tensorrt_llm/batch_manager/trtGptModelFactory.h" -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/memoryUtils.h" -#include "tensorrt_llm/executor/executor.h" -#include "tensorrt_llm/plugins/api/tllmPlugin.h" -#include "tensorrt_llm/runtime/common.h" -#include "tensorrt_llm/runtime/tllmLogger.h" -#include "tensorrt_llm/runtime/utils/mpiUtils.h" -#include "tensorrt_llm/runtime/utils/numpyUtils.h" -#include "tensorrt_llm/testing/modelSpec.h" -#include "tests/utils/common.h" - -#include -#include - -#include -#include -#include -#include -#include - -using namespace tensorrt_llm::testing; -using namespace tensorrt_llm::runtime; -using namespace tensorrt_llm::runtime::utils; -using namespace tensorrt_llm::batch_manager; -namespace fs = std::filesystem; -namespace tc = tensorrt_llm::common; -namespace texec = tensorrt_llm::executor; -using tensorrt_llm::testing::ModelSpec; -using tensorrt_llm::testing::KVCacheType; -using tensorrt_llm::testing::QuantMethod; - -namespace -{ -using TensorPtr = tensorrt_llm::runtime::ITensor::SharedPtr; - -auto constexpr GPT_MODEL_DIR = "gpt2"; -auto constexpr GPTJ_MODEL_DIR = "gpt-j-6b"; -auto constexpr LLAMA_MODEL_DIR = "Llama-3.2-1B"; -auto constexpr MEDUSA_MODEL_DIR = "vicuna-7b-medusa"; -auto constexpr EAGLE_MODEL_DIR = "vicuna-7b-eagle"; -auto constexpr MAMBA_MODEL_DIR = "mamba-2.8b-hf"; -auto constexpr RECURRENTGEMMA_MODEL_DIR = "recurrentgemma-2b"; -auto constexpr EXPLICIT_DRAFT_MODEL_DIR = "vicuna-7b-redrafter"; -auto constexpr CHATGLM_MODEL_DIR = "chatglm-6b"; -auto constexpr GLM_MODEL_DIR = "glm-10b"; - -auto constexpr FP8_GPT_ATTENTION_PLUGIN_IFB_PACKED_PATH = "fp8-plugin"; - -auto constexpr INPUT_FILE = "input_tokens.npy"; -auto constexpr INPUT_LLAMA_FILE = "input_tokens_llama.npy"; -auto constexpr INPUT_VICUNA_FILE = "input_vicuna.npy"; -auto constexpr LONG_INPUT_FILE = "input_tokens_long.npy"; -auto constexpr CHATGLM_INPUT_FILE = "input_tokens_chatglm-6b.npy"; -auto constexpr GLM_INPUT_FILE = "input_tokens_glm-10b.npy"; - -auto constexpr LLAMA_END_ID = 128001; -auto constexpr LLAMA_PAD_ID = 128001; - -struct ModelParams -{ - char const* baseDir; - ModelIds ids; - - friend std::ostream& operator<<(std::ostream& os, ModelParams const& modelParams) - { - return os << "baseDir: " << modelParams.baseDir << ", ids: (" << modelParams.ids.padId << "," - << modelParams.ids.endId << ")"; - } -}; - -} // namespace - -class TrtModelRealDecoderTest : public ::testing::Test // NOLINT(cppcoreguidelines-pro-type-member-init) -{ -protected: - TrtModelRealDecoderTest() {} - - void SetUp() override - { - mDeviceCount = tc::getDeviceCount(); - if (mDeviceCount == 0) - { - GTEST_SKIP() << "No GPUs found"; - } - - mLogger = std::make_shared(); - - initTrtLlmPlugins(mLogger.get()); - } - - void TearDown() override {} - - int mDeviceCount{}; - std::shared_ptr mLogger{}; -}; - -enum class TrtGptModelIfbTestType -{ - BULK, - WAVEFRONT, - RANDOM -}; - -namespace -{ - -void verifyOutput(RequestList const& finishedRequestList, - std::unordered_map const& beamWidthTestData, std::vector const& givenInputLengths, - SizeType32 nbGivenInputs, ModelSpec const& modelSpec) -{ - auto const checkRawLogits = modelSpec.mOtherModelSpecToCompare ? false : modelSpec.mGatherLogits; - auto const smokeTest = modelSpec.mSmokeTest; - auto const returnLogProbs = modelSpec.mReturnLogProbs; - auto const checkAcceptedTokenLogits = modelSpec.mAcceptDraftByLogits; - - if (smokeTest) - { - return; - } - - for (auto const& llmReqPtr : finishedRequestList) - { - auto const& llmReq = *llmReqPtr; - auto const requestId = llmReq.mRequestId; - auto const [givenInputIdx, givenInputLength] - = getRequestGivenInputIdxLength(requestId, nbGivenInputs, givenInputLengths); - auto const reqBeamWidth = llmReq.mSamplingConfig.beamWidth; - auto const& testData = beamWidthTestData.at(reqBeamWidth); - auto const* const expectedOutputData = bufferCast(*testData.expectedOutputIds); - auto const expectedOutputLengths = testData.expectedOutputLengths; - auto const acceptedDraftTokensLengths = testData.acceptedDraftTokensLengths; - auto const endId = testData.endIds[givenInputIdx]; - auto const maxSeqLen = testData.maxSeqLen; - auto const draftLogits = testData.draftLogits; - auto const expectedGenerationLogits = testData.expectedGenerationLogits; - auto const expectedContextLogits = testData.expectedContextLogits; - auto const expectedCumLogProbs = testData.expectedCumLogProbs; - auto const expectedLogProbs = testData.expectedLogProbs; - auto const draftTokens = llmReq.getDraftTokens(); - auto const isDraftTokensExternal = modelSpec.mSpecDecodingMode.isDraftTokensExternal(); - auto const inputLength = givenInputLength + static_cast(isDraftTokensExternal); - - for (auto beam = 0; beam < reqBeamWidth; ++beam) - { - auto const expectedOutputLength = expectedOutputLengths[givenInputIdx * reqBeamWidth + beam]; - auto const predictedTokens = llmReq.getTokens(beam); - - auto numPredTokens = static_cast(predictedTokens.size() - inputLength); - if (isDraftTokensExternal && !draftTokens->empty()) - { - numPredTokens - = std::min(numPredTokens, acceptedDraftTokensLengths[givenInputIdx * reqBeamWidth + beam] + 1); - } - if (modelSpec.mSpecDecodingMode.isMedusa() || modelSpec.mSpecDecodingMode.isLookaheadDecoding() - || modelSpec.mSpecDecodingMode.isExplicitDraftTokens() || modelSpec.mSpecDecodingMode.isEagle()) - { - // WAR to ensure bulk execution of spec decoding. - // We hope that no request in batch can finish 2x faster than any other request. - // For the cases when BS < 8, some predicted tokens are mismatched to reference data. - numPredTokens /= 2; - } - - if (modelSpec.mKVCacheType == KVCacheType::kDISABLED) - { - EXPECT_EQ(numPredTokens, 1) << "b: " << requestId << " beam: " << beam; - } - else - { - EXPECT_EQ(predictedTokens.size(), expectedOutputLength) << "b: " << requestId << " beam: " << beam; - } - - bool anyMismatch = false; - for (auto i = 0; i < numPredTokens; ++i) - { - // Use the expected data for that beamWidth - auto const expectIndex = tc::flat_index3(givenInputIdx, beam, inputLength + i, reqBeamWidth, maxSeqLen); - - auto const expectedToken = expectedOutputData[expectIndex]; - if (expectedToken == endId) - { - break; - } - auto const predictIndex = inputLength + i; - auto const predictedToken = predictedTokens.at(predictIndex); - EXPECT_EQ(predictedToken, expectedToken) << "b: " << requestId << " beam: " << beam << " i: " << i; - anyMismatch |= (predictedToken != expectedToken); - } - EXPECT_FALSE(anyMismatch) << "b: " << requestId << " beam: " << beam; - - if (returnLogProbs) - { - auto cumLogProbs = llmReq.getCumLogProbs(); - auto* const reqExpectedCumLogProbs = bufferCast(*expectedCumLogProbs[requestId]); - EXPECT_TRUE(almostEqual(reqExpectedCumLogProbs[beam], cumLogProbs[beam])); - - auto logProbs = llmReq.getLogProbs(beam); - auto expectedLogProbsBeam = std::shared_ptr(ITensor::slice(expectedLogProbs[requestId], beam, 1)); - expectedLogProbsBeam->squeeze(0); - auto* const reqExpectedLogProbs = bufferCast(*expectedLogProbsBeam); - - for (auto i = 0; i < numPredTokens; ++i) - { - EXPECT_TRUE(almostEqual(reqExpectedLogProbs[inputLength + i], logProbs[i], 5e-2, 5e-2)) - << "expectedLogProbs : " << reqExpectedLogProbs[inputLength + i] - << " logProbs : " << logProbs[i]; - } - } - - if (checkAcceptedTokenLogits && llmReq.hasDraftTokens()) - { - TLLM_CHECK_WITH_INFO(reqBeamWidth == 1, "speculative decoding only works for beam width == 1"); - - TensorPtr const& acceptedTokensLogits = llmReq.getGenerationLogitsHost(); - auto const acceptedTokensLogitsShape = acceptedTokensLogits->getShape(); - - EXPECT_EQ(acceptedTokensLogitsShape.nbDims, 3); - EXPECT_EQ(1, acceptedTokensLogitsShape.d[0]); - EXPECT_EQ(numPredTokens, acceptedTokensLogitsShape.d[1]); - - TensorPtr const& expectedLogits = ITensor::slice(expectedGenerationLogits[requestId], 1, numPredTokens); - - // For hyperparameters - // Greater tolerance for the accepted logits of the target model. - float atol = 0.f; - float rtol = 0.01f; - EXPECT_TRUE(compareLogits(*expectedLogits, *acceptedTokensLogits, atol, rtol)); - } - - if (checkRawLogits) - { - // Check generation logits - TensorPtr const& expectedGenerationLogitsSliced - = ITensor::slice(expectedGenerationLogits[requestId], 0, numPredTokens); - - TensorPtr const& llmReqGeneration = llmReq.getGenerationLogitsHost(); - auto llmReqGenerationShape = llmReqGeneration->getShape(); - - TensorPtr generationLogitsBeam = nullptr; - if (llmReq.isStreaming()) - { - // Expect generation logits shape: [outputLength, beamWidth, vocabSizePad] - EXPECT_EQ(reqBeamWidth, llmReqGenerationShape.d[1]); - EXPECT_EQ(reqBeamWidth, 1); // Streaming mode does not support beam > 1 - llmReqGeneration->squeeze(1); // [outputLength, vocabSizePad] - generationLogitsBeam = llmReqGeneration; - } - else - { - // Expect generation logits shape: [beamWidth, outputLength, vocabSizePad] - EXPECT_EQ(reqBeamWidth, llmReqGenerationShape.d[0]); - generationLogitsBeam - = std::shared_ptr(ITensor::slice(llmReqGeneration, beam, 1)); // [1, outputLength, vocabSizePad] - generationLogitsBeam->squeeze(0); // [outputLength, vocabSizePad] - } - TensorPtr const& generationLogitsSliced = ITensor::slice(generationLogitsBeam, 0, numPredTokens); - EXPECT_TRUE(compareLogits(*expectedGenerationLogitsSliced, *generationLogitsSliced)); - } - } - - if (checkRawLogits) - { - // Check context logits - TensorPtr const& llmReqContext = llmReq.getContextLogitsHost(); - auto llmReqContextShape = llmReqContext->getShape(); - EXPECT_EQ(llmReqContextShape.nbDims, 2); - EXPECT_EQ(llmReq.mPromptLen, llmReqContextShape.d[0]); - EXPECT_TRUE(compareLogits(*expectedContextLogits[requestId], *llmReqContext)); - } - } -} - -// Pick a different endId at random from one of the expected tokens -std::vector pickRandomEndIds(TestData const& testData, std::vector const& givenInputLengths, - SizeType32 const maxNewTokens, bool replaceLogits) -{ - auto const nbGivenInputs = testData.nbGivenInputs; - auto const beamWidth = testData.beamWidth; - auto* const expectedOutputData = bufferCast(*testData.expectedOutputIds); - - std::vector endIds; - - // For IFB, pick one of the output tokens as endId - for (SizeType32 bi = 0; bi < nbGivenInputs; ++bi) - { - TokenIdType skippedEndId0 = 0; - TokenIdType skippedEndId1 = 0; - SizeType32 endIdIndex = 0; - TokenIdType endId = 0; - auto const endIdRow = bi; - auto const inputLength = givenInputLengths.at(endIdRow); - do - { - auto const endIdBeam = std::rand() % beamWidth; - auto const firstOutputIndex - = tc::flat_index3(endIdRow, endIdBeam, inputLength, beamWidth, testData.maxSeqLen); - // We do not use the 1st token for EndId because of Speculative Decoding test design - // We skip 1st token because minLength is 1 - auto const endIdCol = 2 + (std::rand() % std::max(maxNewTokens - 2, 1)); - endIdIndex = firstOutputIndex + endIdCol; - skippedEndId0 = expectedOutputData[firstOutputIndex]; - skippedEndId1 = expectedOutputData[firstOutputIndex + 1]; - endId = expectedOutputData[endIdIndex]; - } while (endId == skippedEndId0 || endId == skippedEndId1); - // Workaround: The first example has endIdIndex 14, where the generation logits are almost same at - // token ids 257 and 373, which causes unstable generation results. Hence, we use the one previous - // token as endId. - if (bi == 0 && !replaceLogits) - { - endId = expectedOutputData[endIdIndex - 1]; - } - endIds.push_back(endId); - } - - return endIds; -} - -TestData loadTestData(ModelSpec const& modelSpec, ModelIds const modelIds, BeamResult const& beamResult, - ITensor const& givenInput, SizeType32 const maxBeamWidth, bool const useRandomEndId, bool const replaceLogits, - BufferManager& manager) -{ - auto const [givenInputLengths, nbGivenInputs, maxInputLength] = getGivenInputLengths(givenInput, modelIds.padId); - auto const& [beamWidth, resultsFile, contextLogitsFile, genLogitsFile, cumLogProbsFile, logProbsFile] = beamResult; - - TestData testData{nbGivenInputs, beamWidth}; - testData.expectedOutputIds = loadNpy(manager, resultsFile.string(), MemoryType::kCPU); - - auto* const expectedOutputData = bufferCast(*testData.expectedOutputIds); - - auto const& outputShape = testData.expectedOutputIds->getShape(); - EXPECT_EQ(outputShape.nbDims, 2); - EXPECT_EQ(nbGivenInputs * beamWidth, outputShape.d[0]); - testData.maxSeqLen = static_cast(outputShape.d[1]); - EXPECT_LE(maxInputLength, testData.maxSeqLen); - EXPECT_LE(beamWidth, maxBeamWidth); - - auto const maxNewTokens = testData.maxSeqLen - maxInputLength; - - std::srand(42); - - if (useRandomEndId) - { - testData.endIds = pickRandomEndIds(testData, givenInputLengths, maxNewTokens, replaceLogits); - } - else - { - testData.endIds.insert(testData.endIds.end(), nbGivenInputs, modelIds.endId); - } - - if (modelSpec.useLogits()) - { - testData.loadContextLogits(contextLogitsFile, givenInputLengths, manager); - } - if (modelSpec.useLogits() || modelSpec.mAcceptDraftByLogits) - { - testData.loadGenerationLogits(genLogitsFile, manager); - } - if (modelSpec.mReturnLogProbs) - { - testData.loadLogProbs(cumLogProbsFile, logProbsFile, manager); - } - - for (SizeType32 bi = 0; bi < nbGivenInputs; ++bi) - { - auto const endId = testData.endIds[bi]; - for (SizeType32 beam = 0; beam < beamWidth; ++beam) - { - SizeType32 expectedLen = givenInputLengths[bi] + maxNewTokens; - for (SizeType32 si = givenInputLengths[bi]; si < testData.maxSeqLen; ++si) - { - auto const expectIndex = tc::flat_index2((bi * beamWidth + beam), si, testData.maxSeqLen); - if (expectedOutputData[expectIndex] == endId) - { - expectedLen = si; - break; - } - } - // Fill new EOS token to the expected data - for (SizeType32 si = expectedLen; si < testData.maxSeqLen; ++si) - { - auto const expectIndex = tc::flat_index2((bi * beamWidth + beam), si, testData.maxSeqLen); - expectedOutputData[expectIndex] = endId; - } - - testData.expectedOutputLengths[bi * beamWidth + beam] = expectedLen; - } - } - - if (modelSpec.mMaxDraftTokens > 0) - { - testData.makeDraft( - modelSpec.mMaxDraftTokens, modelSpec.mAcceptDraftByLogits, genLogitsFile, givenInputLengths, manager); - } - - return testData; -} - -std::tuple, std::unordered_map> loadTestData(ModelSpec const& modelSpec, - ModelIds const modelIds, BeamResults const& resultsFilesBeamWidths, ITensor const& givenInput, - SizeType32 const maxBeamWidth, bool const useRandomEndId, bool const replaceLogits, BufferManager& manager) -{ - // Map between beam width, and expected results for that beam width - std::unordered_map beamWidthTestData; - std::vector beamWidths; - - for (auto const& beamResult : resultsFilesBeamWidths) - { - auto const beamWidth = beamResult.beamWidth; - - EXPECT_EQ(std::find(beamWidths.begin(), beamWidths.end(), beamWidth), beamWidths.end()); - beamWidths.push_back(beamWidth); - - auto testData = loadTestData( - modelSpec, modelIds, beamResult, givenInput, maxBeamWidth, useRandomEndId, replaceLogits, manager); - beamWidthTestData.emplace(beamWidth, std::move(testData)); - } - - return {std::move(beamWidths), std::move(beamWidthTestData)}; -} - -RequestList runGptModelInference(std::shared_ptr& trtGptModel, std::vector const& beamWidths, - std::unordered_map const& beamWidthTestData, SizeType32 batchSize, SizeType32 nbGivenInputs, - SizeType32 maxInputLength, SizeType32 padId, std::vector const& givenInputLengths, - TokenIdType const* givenInputData, ModelSpec const& modelSpec, TrtGptModelIfbTestType testType, int maxReqPerStep, - bool prepopulateKVCache, bool enableStreamingMode, bool enableBlockReuse) -{ - // Fill the requests using givenInput - // requestList will have batchSize requests - RequestList requestList; - - SizeType32 requestId = 0; - RequestList finishedRequestList; - std::vector reqVec; - // Advance the requests until they are all finished - if (COMM_SESSION.getRank() == 0) - { - SizeType32 numReq = 0; - while (numReq < batchSize) - { - // Add appropriate number of requests in each iteration. For WAVEFRONT, this is always 1. - // For RANDOM, it could be any integer <= maxReqPerStep including 0. - SizeType32 reqThisStep{0}; - switch (testType) - { - case TrtGptModelIfbTestType::WAVEFRONT: reqThisStep = 1; break; - case TrtGptModelIfbTestType::RANDOM: reqThisStep = rand() % (maxReqPerStep + 1); break; - case TrtGptModelIfbTestType::BULK: [[fallthrough]]; - default: reqThisStep = batchSize; break; - } - reqThisStep = std::min(reqThisStep, (batchSize - numReq)); - reqVec.push_back(reqThisStep); - numReq += reqThisStep; - } - } - COMM_SESSION.bcast(reqVec, 0); - - SizeType32 reqVecIdx = 0; - while (requestId < batchSize || !requestList.empty()) - { - SizeType32 reqThisStep = reqVecIdx < reqVec.size() ? reqVec[reqVecIdx++] : 0; - for (SizeType32 req = 0; req < reqThisStep; req++) - { - // Alternate between beamWidths - SizeType32 beamWidth = beamWidths.at(requestId % beamWidths.size()); - auto const& testData = beamWidthTestData.at(beamWidth); - auto const* const expectedOutputData = bufferCast(*testData.expectedOutputIds); - auto const maxSeqLen = testData.maxSeqLen; - - SamplingConfig samplingConfig{beamWidth}; - samplingConfig.temperature = std::vector{1.0f}; - samplingConfig.minLength = std::vector{1}; - samplingConfig.randomSeed = std::vector{static_cast(42ull)}; - samplingConfig.topK = std::vector{1}; - samplingConfig.topP = std::vector{0.0f}; - samplingConfig.draftAcceptanceThreshold = std::vector{0.3f}; - samplingConfig.noRepeatNgramSize = std::vector{1 << 30}; - - auto const [givenInputIdx, inputLength] - = getRequestGivenInputIdxLength(requestId, nbGivenInputs, givenInputLengths); - SizeType32 endId = testData.endIds[givenInputIdx]; - - auto maxNewTokens = maxSeqLen - maxInputLength; - // Run model only to produce a single token and prepopulate KV cache - if (prepopulateKVCache || modelSpec.mKVCacheType == KVCacheType::kDISABLED) - { - maxNewTokens = 1; - } - auto const* const seqBegin = givenInputData + givenInputIdx * maxInputLength; - auto tokens = std::make_shared>(seqBegin, seqBegin + inputLength); - if (!prepopulateKVCache && modelSpec.mMaxDraftTokens > 0) - { - // Append the 1st predicted token to the prompt to get the match with prepopulated KV cache - auto const expectIndex = tc::flat_index3(givenInputIdx, 0, inputLength, 1, maxSeqLen); - auto expectedToken = expectedOutputData[expectIndex]; - tokens->push_back(expectedToken); - // subtract this token from maxNewTokens - maxNewTokens -= 1; - } - auto r = std::make_shared(requestId, maxNewTokens, tokens, samplingConfig, false, endId, padId); - - auto const& draftTokens = testData.draftTokens[givenInputIdx]; - auto draftLogits = modelSpec.mAcceptDraftByLogits - ? std::make_optional(testData.draftLogits[givenInputIdx]) - : std::nullopt; - if (!prepopulateKVCache && !draftTokens.empty()) - { - r->setDraftTokens(std::make_shared>(draftTokens)); - r->setDraftLogits(draftLogits); - } - - SizeType32 maxDraftTokens{0}; - if (trtGptModel->getModelConfig().hasSpeculativeDecodingModule()) - { - maxDraftTokens - = trtGptModel->getModelConfig().getSpeculativeDecodingModulePtr()->getMaxDecodingDraftTokens(); - } - r->validate(trtGptModel->getMaxInputLen(), trtGptModel->getMaxSequenceLen(), maxDraftTokens, - trtGptModel->getVocabSizePadded(), std::nullopt, enableBlockReuse); - - if (enableStreamingMode) - { - r->setReturnAllGeneratedTokens(true); // Test allGeneratedTokens in this test - r->setStreaming(true); - } - - auto const vocabSizePadded - = trtGptModel->getModelConfig().getVocabSizePadded(trtGptModel->getWorldConfig().getSize()); - auto const logitDatatype = trtGptModel->getLogitDataType(); - if (modelSpec.mGatherLogits) - { - r->setReturnContextLogits(true); - r->setReturnGenerationLogits(true); - r->allocContextLogitsHost(vocabSizePadded, logitDatatype); - r->allocGenerationLogitsHost(vocabSizePadded, logitDatatype); - } - - if (!prepopulateKVCache && modelSpec.mAcceptDraftByLogits && !draftTokens.empty()) - { - r->allocTargetModelAcceptedTokenLogitsHost(vocabSizePadded, logitDatatype); - r->setReturnGenerationLogits(true); - } - - if (modelSpec.mReplaceLogits) - { - LlmRequest::LogitsPostProcessor logitsCb - = [&testData](uint64_t rId, tensorrt_llm::runtime::ITensor::SharedPtr& logits, - LlmRequest::BeamTokens const& tokens, - tensorrt_llm::runtime::BufferManager::CudaStreamPtr streamPtr, std::optional cId) - { - auto const expectedGenerationLogits = testData.expectedGenerationLogits[rId]; - auto const expectedContextLogits = testData.expectedContextLogits[rId]; - auto const acceptedDraftTokensLengths = testData.acceptedDraftTokensLengths[rId]; - - auto const beamWidth = tokens.size(); - TLLM_CHECK_WITH_INFO(beamWidth == 1, "Logits substitution is not supported for beam search"); - - auto const genLogitsOffset = tokens[0].size() - expectedContextLogits->getShape().d[0]; - // TODO: Avoid static cast in TRT 10.0 - auto const numLogits = static_cast(logits->getShape().d[0]); - auto const numVerifyLogits = std::min(numLogits, acceptedDraftTokensLengths + 1); - - TensorPtr logitsSlice = ITensor::slice(logits, 0, numVerifyLogits); - - auto manager = BufferManager(streamPtr); - TensorPtr logitsHost = manager.copyFrom(*logitsSlice, MemoryType::kCPU); - manager.getStream().synchronize(); - - TensorPtr refLogitsHost - = ITensor::slice(expectedGenerationLogits, genLogitsOffset, numVerifyLogits); - - EXPECT_TRUE(compareLogits(*refLogitsHost, *logitsHost, 0.f, 1e-2)) << "reqId: " << rId; - - manager.copy(*refLogitsHost, *logitsSlice); - }; - - r->mLogitsPostProcessor = logitsCb; - } - - if (modelSpec.mReturnLogProbs) - { - r->setReturnLogProbs(true); - } - requestList.push_back(r); - ++requestId; - } - - // Advance all active requests by one step - trtGptModel->forwardAsync(requestList); - trtGptModel->forwardSync(); - - // Check which requests are done, move them out - for (auto it = requestList.cbegin(); it != requestList.cend();) - { - if ((*it)->isGenerationCompleteState()) - { - finishedRequestList.push_back(*it); - requestList.erase(it++); - } - else - { - ++it; - } - } - } - return finishedRequestList; -} - -void runIfbTest(fs::path const& modelPath, ModelSpec const& modelSpec, ModelIds const modelIds, - TrtGptModelType modelType, std::vector const& batchSizes, BeamResults const& resultsFilesBeamWidths, - TrtGptModelIfbTestType testType, int maxReqPerStep, texec::ExecutorConfig const& executorConfig, - bool enableStreamingMode, bool useRandomEndId) -{ - auto manager = BufferManager(std::make_shared()); - auto const padId = modelIds.padId; - - // Load input data - ASSERT_TRUE(fs::exists(DATA_PATH)); - auto const inputPath = DATA_PATH / modelSpec.mInputFile; - auto const& givenInput = loadNpy(manager, inputPath.string(), MemoryType::kCPU); - auto [givenInputLengths, nbGivenInputs, maxInputLength] = getGivenInputLengths(*givenInput, padId); - auto const* const givenInputData = bufferCast(*givenInput); - - auto const& inputShape = givenInput->getShape(); - ASSERT_EQ(inputShape.nbDims, 2); - ASSERT_GT(inputShape.d[0], 0); - - auto const maxBeamWidth = executorConfig.getMaxBeamWidth(); - // Load expected outputs for each beam width value - auto [beamWidths, beamWidthTestData] = loadTestData(modelSpec, modelIds, resultsFilesBeamWidths, *givenInput, - maxBeamWidth, useRandomEndId, modelSpec.mReplaceLogits, manager); - - int const worldSize = modelSpec.mTPSize * modelSpec.mPPSize * modelSpec.mCPSize; - auto const worldConfig = WorldConfig::mpi(worldSize, modelSpec.mTPSize, modelSpec.mPPSize, modelSpec.mCPSize); - - ASSERT_TRUE(fs::exists(modelPath)); - - for (auto batchSize : batchSizes) - { - std::cout << "=== batchSize:" << batchSize << " ===\n"; - - auto trtGptModel = TrtGptModelFactory::create(modelPath, modelType, executorConfig, false); - - if (modelSpec.mKVCacheType == KVCacheType::kDISABLED) - { - ASSERT_FALSE(trtGptModel->hasKVCacheManager()); - } - - // Prepopulate KV cache for speculative decoding test - bool const prepopulateKVCache = modelSpec.mMaxDraftTokens > 0; - auto finishedRequestList = runGptModelInference(trtGptModel, beamWidths, beamWidthTestData, batchSize, - nbGivenInputs, maxInputLength, padId, givenInputLengths, givenInputData, modelSpec, testType, maxReqPerStep, - prepopulateKVCache, enableStreamingMode, modelSpec.mKVCacheReuse); - - if (prepopulateKVCache) - { - // Call the 2nd time with prefilled KV cache - finishedRequestList = runGptModelInference(trtGptModel, beamWidths, beamWidthTestData, batchSize, - nbGivenInputs, maxInputLength, padId, givenInputLengths, givenInputData, modelSpec, testType, - maxReqPerStep, false, enableStreamingMode, modelSpec.mKVCacheReuse); - } - - // WAR: disabled verification because of switched beams for different batch composition - if (worldConfig.isFirstPipelineParallelRank() - && (testType == TrtGptModelIfbTestType::BULK || maxBeamWidth == 1)) - { - bool shouldVerify = true; - - if (testType == TrtGptModelIfbTestType::BULK) - { - if (modelSpec.mKVCacheType == KVCacheType::kDISABLED && maxBeamWidth != 1) - { - // For disabled KV cache, only verify when maxBeamWidth is 1, the reason is we only compare with - // results with KV cache enabled case and usually, beams search results locate in last token while - // disabled KV cache only get exactly one new token. - shouldVerify = false; - } - } - - if (shouldVerify) - { - verifyOutput(finishedRequestList, beamWidthTestData, givenInputLengths, nbGivenInputs, modelSpec); - } - } - } -} - -struct BeamConfig -{ - SizeType32 maxBeamWidth; - std::vector beamWidths; -}; - -} // namespace - -using ParamType = std::tuple, // 5. maxTokensInPagedKvCache - std::optional, // 6. freeGpuMemoryFraction - bool, // 7. enableTrtOverlap - bool, // 8. enableChunkedContext - bool, // 9. enableStreamingMode - bool, // 10. enableCudaGraphMode - std::optional, // 11. hostCacheSize - bool, // 12. useRandomEndId - std::vector, // 13. batchSizes - std::optional // 14. maxNumTokens - >; - -std::string generateTestName(testing::TestParamInfo const& info) -{ - auto const modelSpec = std::get<1>(info.param); - std::string name; - switch (modelSpec.mDataType) - { - case nvinfer1::DataType::kFLOAT: name.append("Float"); break; - case nvinfer1::DataType::kHALF: name.append("Half"); break; - case nvinfer1::DataType::kINT8: name.append("Int8"); break; - case nvinfer1::DataType::kINT32: name.append("Int32"); - case nvinfer1::DataType::kBOOL: name.append("Bool"); break; - case nvinfer1::DataType::kUINT8: name.append("UInt8"); break; - case nvinfer1::DataType::kFP8: name.append("Float8"); break; - case nvinfer1::DataType::kBF16: name.append("BFloat16"); break; - case nvinfer1::DataType::kINT4: name.append("Int4"); break; - case nvinfer1::DataType::kFP4: name.append("Fp4"); break; - default: throw std::runtime_error("Unsupported DataType"); break; - } - - auto const modelType = std::get<2>(info.param); - switch (modelType) - { - case TrtGptModelType::InflightBatching: name.append("IbModel"); break; - case TrtGptModelType::InflightFusedBatching: name.append("FusedIbModel"); break; - default: name.append("DefaultModel"); break; - } - - switch (modelSpec.mKVCacheType) - { - case KVCacheType::kCONTINUOUS: name.append("ContinuousKVCache"); break; - case KVCacheType::kPAGED: name.append("PagedKVCache"); break; - case KVCacheType::kDISABLED: name.append("NoKVCache"); break; - default: throw std::runtime_error("Unknown KVCacheType"); break; - } - - auto const testType = std::get<3>(info.param); - switch (testType) - { - case TrtGptModelIfbTestType::BULK: name.append("Bulk"); break; - case TrtGptModelIfbTestType::WAVEFRONT: name.append("Wavefront"); break; - case TrtGptModelIfbTestType::RANDOM: name.append("Random"); break; - default: name.append("DefaultTest"); break; - } - BeamConfig const beamConfig = std::get<4>(info.param); - name.append("MaxBeamWidth" + std::to_string(beamConfig.maxBeamWidth)); - for (auto const beamWdith : beamConfig.beamWidths) - { - name.append("Bw" + std::to_string(beamWdith)); - } - - auto const maxTokensInPagedKvCache = std::get<5>(info.param); - if (maxTokensInPagedKvCache.has_value()) - { - name.append("KvCacheSize" + std::to_string(maxTokensInPagedKvCache.value())); - } - - auto const freeGpuMemoryFraction = std::get<6>(info.param); - if (freeGpuMemoryFraction.has_value()) - { - name.append("GpuFrac"); - } - - auto const enableTrtOverlap = std::get<7>(info.param); - if (enableTrtOverlap) - { - name.append("TrtOverlap"); - } - - auto const enableChunkedContext = std::get<8>(info.param); - if (enableChunkedContext) - { - name.append("Chunked"); - } - - if (modelSpec.mTPSize > 1) - { - name.append("TP" + std::to_string(modelSpec.mTPSize)); - } - - if (modelSpec.mPPSize > 1) - { - name.append("PP" + std::to_string(modelSpec.mPPSize)); - } - - if (modelSpec.mCPSize > 1) - { - name.append("CP" + std::to_string(modelSpec.mCPSize)); - } - - auto const useRandomEndId = std::get<12>(info.param); - if (useRandomEndId) - { - name.append("EndId"); - } - - if (modelSpec.mMaxDraftTokens > 0) - { - name.append("DraftTokens" + std::to_string(modelSpec.mMaxDraftTokens)); - } - - if (modelSpec.mAcceptDraftByLogits) - { - name.append("AcceptByLogits"); - } - - if (modelSpec.mCapacitySchedulerPolicy) - { - name.append(modelSpec.getCapacitySchedulerString()); - } - - auto const enableStreamingMode = std::get<9>(info.param); - if (enableStreamingMode) - { - name.append("Streaming"); - } - - auto const enableCudaGraphMode = std::get<10>(info.param); - if (enableCudaGraphMode) - { - name.append("CudaGraph"); - } - - auto const enableHostCache = std::get<11>(info.param); - if (enableHostCache) - { - name.append("SecondaryOffloading"); - } - - return name; -} - -class ParamTest : public TrtModelRealDecoderTest, public ::testing::WithParamInterface -{ -}; - -TEST_P(ParamTest, Test) -{ - - auto const& beamConfig = std::get<4>(GetParam()); - auto const& beamWidths = beamConfig.beamWidths; - - auto const modelParams = std::get<0>(GetParam()); - auto const modelIds = modelParams.ids; - auto const* const modelDir = modelParams.baseDir; - auto const modelSpec = std::get<1>(GetParam()); - - auto const useRandomEndId = std::get<12>(GetParam()); - - auto const batchSizes = std::get<13>(GetParam()); - - std::ostringstream gpuSizePath; - gpuSizePath << "tp" << modelSpec.mTPSize << "-pp" << modelSpec.mPPSize << "-cp" << modelSpec.mCPSize; - gpuSizePath << "-gpu"; - - auto const modelPath{ENGINE_PATH / modelDir / modelSpec.getModelPath() / gpuSizePath.str()}; - - auto const inputPath = DATA_PATH / modelSpec.mInputFile; - - BeamResults beamResults; - beamResults.reserve(beamWidths.size()); - for (auto beamWidth : beamWidths) - { - fs::path resultsPath - = DATA_PATH / modelDir / ((beamWidth == 1) ? "sampling" : "beam_search_" + std::to_string(beamWidth)); - fs::path generationLogitsPath - = modelSpec.mCollectGenerationLogits ? (resultsPath / modelSpec.getGenerationLogitsFile()).string() : ""; - fs::path contextLogitsPath - = modelSpec.mCollectContextLogits ? (resultsPath / modelSpec.getContextLogitsFile()).string() : ""; - fs::path cumLogProbsPath - = modelSpec.mCollectCumLogProbs ? (resultsPath / modelSpec.getCumLogProbsFile()).string() : ""; - fs::path logProbsPath = modelSpec.mCollectLogProbs ? (resultsPath / modelSpec.getLogProbsFile()).string() : ""; - - beamResults.emplace_back(beamWidth, (resultsPath / modelSpec.getResultsFile()).string(), contextLogitsPath, - generationLogitsPath, cumLogProbsPath, logProbsPath); - } - - auto const modelType = std::get<2>(GetParam()); - auto const testType = std::get<3>(GetParam()); - auto const enableStreamingMode = std::get<9>(GetParam()); - auto const cudaGraphMode = std::get<10>(GetParam()); - - if (!(modelSpec.mUsePackedInput - && (modelSpec.mKVCacheType == KVCacheType::kPAGED || modelSpec.mKVCacheType == KVCacheType::kDISABLED))) - { - GTEST_SKIP() << "Inflight batching requires packed input and (paged KV cache or disabled KV cache)."; - } - - if (!modelSpec.mUsePackedInput && useRandomEndId) - { - GTEST_SKIP() << "Test does not support endId test with padded inputs"; - } - - for (auto beamWidth : beamWidths) - { - if (useRandomEndId && beamWidth > 1) - { - GTEST_SKIP() << "Test does not support endId test with beam search"; - } - - if (modelSpec.mMaxDraftTokens > 0 && beamWidth > 1) - { - GTEST_SKIP() << "Target model in speculative decoding does not support beam search"; - } - } - - auto executorConfig = texec::ExecutorConfig{}; - - auto const maxTokens = std::get<5>(GetParam()); - auto const enableBlockReuse = modelSpec.mMaxDraftTokens > 0 || modelSpec.mKVCacheReuse; - auto const freeGpuMemoryFraction = std::get<6>(GetParam()); - auto const hostCacheSize = std::get<11>(GetParam()); - auto const kvCacheConfig = texec::KvCacheConfig{ - enableBlockReuse, maxTokens, std::nullopt, std::nullopt, freeGpuMemoryFraction, hostCacheSize}; - executorConfig.setKvCacheConfig(kvCacheConfig); - - executorConfig.setEnableTrtOverlap(std::get<7>(GetParam())); - executorConfig.setEnableChunkedContext(std::get<8>(GetParam())); - auto const maxNumTokens = std::get<14>(GetParam()); - if (maxNumTokens.has_value()) - { - executorConfig.setMaxNumTokens(maxNumTokens.value()); - } - executorConfig.setNormalizeLogProbs(false); - executorConfig.setMaxBeamWidth(beamConfig.maxBeamWidth); - executorConfig.setGatherGenerationLogits(modelSpec.mCollectGenerationLogits); - auto extendedRuntimePerfKnobConfig = texec::ExtendedRuntimePerfKnobConfig{}; - extendedRuntimePerfKnobConfig.setCudaGraphMode(cudaGraphMode); - executorConfig.setExtendedRuntimePerfKnobConfig(extendedRuntimePerfKnobConfig); - - auto const capacitySchedulerPolicy - = modelSpec.mCapacitySchedulerPolicy.value_or(texec::CapacitySchedulerPolicy::kMAX_UTILIZATION); - executorConfig.setSchedulerConfig(texec::SchedulerConfig{capacitySchedulerPolicy}); - - if (modelSpec.mSpecDecodingMode == SpeculativeDecodingMode::LookaheadDecoding()) - { - auto decodingConfig = texec::DecodingConfig{}; - decodingConfig.setLookaheadDecodingConfig(texec::LookaheadDecodingConfig(5, 5, 5)); - executorConfig.setDecodingConfig(decodingConfig); - } - - for (auto beamWidth : beamWidths) - { - if (executorConfig.getEnableTrtOverlap() && beamWidth > 1) - { - GTEST_SKIP() << "TrtOverlap is not supported with beam search"; - } - } - - if (executorConfig.getEnableTrtOverlap() && modelSpec.mMaxDraftTokens > 0) - { - GTEST_SKIP() << "TrtOverlap is not supported with speculative decoding"; - } - - // Warning: This should be the last check before running the test. - // It will initialize MPI which can take significant time. - if (modelSpec.mTPSize * modelSpec.mPPSize * modelSpec.mCPSize != COMM_SESSION.getSize()) - { - GTEST_SKIP() << "Model's world size " << modelSpec.mPPSize * modelSpec.mTPSize * modelSpec.mCPSize - << " is not equal to the system world size"; - } - - runIfbTest(modelPath, modelSpec, modelIds, modelType, batchSizes, beamResults, testType, 2, executorConfig, - enableStreamingMode, useRandomEndId); -} - -auto constexpr gptModelParams = ModelParams{GPT_MODEL_DIR, ModelIds{50256, 50256}}; - -std::shared_ptr getGptDraftTestsCompareModelSpec() -{ - auto pModelSpec = std::make_shared(INPUT_FILE, nvinfer1::DataType::kHALF); - pModelSpec->useGptAttentionPlugin(); - pModelSpec->gatherLogits(); - pModelSpec->usePackedInput(); - pModelSpec->setKVCacheType(KVCacheType::kPAGED); - - return pModelSpec; -} - -std::shared_ptr getMedusaTestsCompareModelSpec() -{ - auto pModelSpec = std::make_shared(LONG_INPUT_FILE, nvinfer1::DataType::kHALF); - pModelSpec->useGptAttentionPlugin(); - pModelSpec->usePackedInput(); - pModelSpec->setKVCacheType(KVCacheType::kPAGED); - pModelSpec->setMaxOutputLength(128); - - return pModelSpec; -} - -std::shared_ptr getEagleTestsCompareModelSpec() -{ - auto pModelSpec = std::make_shared(LONG_INPUT_FILE, nvinfer1::DataType::kHALF); - pModelSpec->useGptAttentionPlugin(); - pModelSpec->usePackedInput(); - pModelSpec->setKVCacheType(KVCacheType::kPAGED); - pModelSpec->setMaxOutputLength(128); - - return pModelSpec; -} - -std::shared_ptr getGptChunkedContextTestsCompareModelSpec() -{ - auto pModelSpec = std::make_shared(LONG_INPUT_FILE, nvinfer1::DataType::kHALF); - pModelSpec->useGptAttentionPlugin(); - pModelSpec->usePackedInput(); - pModelSpec->setKVCacheType(KVCacheType::kPAGED); - pModelSpec->setMaxInputLength(128); - - return pModelSpec; -} - -INSTANTIATE_TEST_SUITE_P(GptTests, ParamTest, - testing::Combine(testing::Values(gptModelParams), - testing::Values( - // - ModelSpec{INPUT_FILE, nvinfer1::DataType::kHALF} - .useGptAttentionPlugin() - .setKVCacheType(KVCacheType::kPAGED) - .usePackedInput(), - ModelSpec{INPUT_FILE, nvinfer1::DataType::kHALF, - []() -> std::shared_ptr - { - auto pModelSpec = std::make_shared(INPUT_FILE, nvinfer1::DataType::kHALF); - pModelSpec->useGptAttentionPlugin().setKVCacheType(KVCacheType::kPAGED).usePackedInput(); - return pModelSpec; - }()} - .useGptAttentionPlugin() - .setKVCacheType(KVCacheType::kDISABLED) - .usePackedInput()), - testing::Values(TrtGptModelType::InflightFusedBatching), - testing::Values( - TrtGptModelIfbTestType::BULK, TrtGptModelIfbTestType::WAVEFRONT, TrtGptModelIfbTestType::RANDOM), - testing::Values( - // TODO: enable more tests when mixed beam width is supported - BeamConfig{1, {1}}, BeamConfig{2, {2}} // , BeamConfig{2, {1, 2}} - ), - testing::Values(std::nullopt, 1280), // maxTokensInPagedKvCache - testing::Values(std::nullopt, 0.4), // freeGpuMemoryFraction - testing::Values(true), // enableTrtOverlap - testing::Values(false), // enableChunkedContext - testing::Values(false), // enableStreamingMode - testing::Values(false), // enableCudaGraphMode - testing::Values(std::nullopt), // hostCacheSize - testing::Values(false), // useRandomEndId - testing::Values(std::vector{1, 2, 8}), // batchSizes - testing::Values(std::nullopt) // maxNumTokens - ), - generateTestName); - -INSTANTIATE_TEST_SUITE_P(GptRandomEndIdTests, ParamTest, - testing::Combine(testing::Values(gptModelParams), - testing::Values( - // - ModelSpec{INPUT_FILE, nvinfer1::DataType::kHALF} - .useGptAttentionPlugin() - .setKVCacheType(KVCacheType::kPAGED) - .usePackedInput()), - testing::Values(TrtGptModelType::InflightFusedBatching), - testing::Values( - TrtGptModelIfbTestType::BULK, TrtGptModelIfbTestType::WAVEFRONT, TrtGptModelIfbTestType::RANDOM), - testing::Values( - // TODO: enable more tests when mixed beam width is supported - BeamConfig{1, {1}}, BeamConfig{2, {2}} // , BeamConfig{2, {1, 2}} - ), - testing::Values(std::nullopt, 1280), // maxTokensInPagedKvCache - testing::Values(std::nullopt, 0.4), // freeGpuMemoryFraction - testing::Values(true), // enableTrtOverlap - testing::Values(false), // enableChunkedContext - testing::Values(false), // enableStreamingMode - testing::Values(false), // enableCudaGraphMode - testing::Values(std::nullopt), // hostCacheSize - testing::Values(true), // useRandomEndId - testing::Values(std::vector{1, 2, 8}), // batchSizes - testing::Values(std::nullopt) // maxNumTokens - ), - generateTestName); - -INSTANTIATE_TEST_SUITE_P(GptKVOffloadingTest, ParamTest, - testing::Combine(testing::Values(gptModelParams), - testing::Values( - // - ModelSpec{LONG_INPUT_FILE, nvinfer1::DataType::kHALF} - .useGptAttentionPlugin() - .setKVCacheType(KVCacheType::kPAGED) - .usePackedInput() - .setKVCacheReuse(true)), - testing::Values(TrtGptModelType::InflightFusedBatching), - testing::Values( - TrtGptModelIfbTestType::BULK, TrtGptModelIfbTestType::WAVEFRONT, TrtGptModelIfbTestType::RANDOM), - testing::Values(BeamConfig{1, {1}}), - testing::Values(256), // maxTokensInPagedKvCache - testing::Values(std::nullopt, 0.4), // freeGpuMemoryFraction - testing::Values(true), // enableTrtOverlap - testing::Values(false), // enableChunkedContext - testing::Values(false), // enableStreamingMode - testing::Values(false), // enableCudaGraphMode - testing::Values(100000000), // hostCacheSize - testing::Values(false, true), // useRandomEndId - testing::Values(std::vector{1, 2, 8}), // batchSizes - testing::Values(std::nullopt) // maxNumTokens - ), - generateTestName); - -INSTANTIATE_TEST_SUITE_P(GptCudaGraphTests, ParamTest, - testing::Combine(testing::Values(gptModelParams), - testing::Values( - // - ModelSpec{INPUT_FILE, nvinfer1::DataType::kHALF} - .useGptAttentionPlugin() - .setKVCacheType(KVCacheType::kPAGED) - .usePackedInput() - .capacitySchedulerPolicy(texec::CapacitySchedulerPolicy::kSTATIC_BATCH), - ModelSpec{INPUT_FILE, nvinfer1::DataType::kHALF} - .useGptAttentionPlugin() - .setKVCacheType(KVCacheType::kPAGED) - .usePackedInput() - .capacitySchedulerPolicy(texec::CapacitySchedulerPolicy::kMAX_UTILIZATION)), - testing::Values(TrtGptModelType::InflightBatching, TrtGptModelType::InflightFusedBatching), - testing::Values( - TrtGptModelIfbTestType::BULK, TrtGptModelIfbTestType::WAVEFRONT, TrtGptModelIfbTestType::RANDOM), - testing::Values( - // TODO: enable more tests when mixed beam width is supported - BeamConfig{1, {1}}, BeamConfig{2, {2}} // , BeamConfig{2, {1, 2}} - ), - testing::Values(std::nullopt), // maxTokensInPagedKvCache - testing::Values(0.4), // freeGpuMemoryFraction - testing::Values(false), // enableTrtOverlap - testing::Values(true), // enableChunkedContext - testing::Values(false), // enableStreamingMode - testing::Values(true), // enableCudaGraphMode - testing::Values(std::nullopt), // hostCacheSize - testing::Values(false), // useRandomEndId - testing::Values(std::vector{1, 2, 8}), // batchSizes - testing::Values(std::nullopt) // maxNumTokens - ), - generateTestName); - -INSTANTIATE_TEST_SUITE_P(GptSwitchBwTests, ParamTest, - testing::Combine(testing::Values(gptModelParams), - testing::Values( - // - ModelSpec{INPUT_FILE, nvinfer1::DataType::kHALF} - .useGptAttentionPlugin() - .setKVCacheType(KVCacheType::kPAGED) - .usePackedInput()), - testing::Values(TrtGptModelType::InflightFusedBatching), - testing::Values( - TrtGptModelIfbTestType::BULK, TrtGptModelIfbTestType::WAVEFRONT, TrtGptModelIfbTestType::RANDOM), - testing::Values( - // TODO: enable more tests when mixed beam width is supported - BeamConfig{2, {1}} // , BeamConfig{2, {1, 2}} - ), - testing::Values(std::nullopt), // maxTokensInPagedKvCache - testing::Values(0.4), // freeGpuMemoryFraction - testing::Values(false), // enableTrtOverlap - testing::Values(true), // enableChunkedContext - testing::Values(false), // enableStreamingMode - testing::Values(false), // enableCudaGraphMode - testing::Values(std::nullopt), // hostCacheSize - testing::Values(false), // useRandomEndId - testing::Values(std::vector{4}), // batchSizes - testing::Values(std::nullopt) // maxNumTokens - ), - generateTestName); - -INSTANTIATE_TEST_SUITE_P(GptNProfilesTests, ParamTest, - testing::Combine(testing::Values(gptModelParams), - testing::Values(ModelSpec{INPUT_FILE, nvinfer1::DataType::kHALF} - .useGptAttentionPlugin() - .usePackedInput() - .setKVCacheType(KVCacheType::kPAGED) - .useMultipleProfiles()), - testing::Values(TrtGptModelType::InflightFusedBatching), testing::Values(TrtGptModelIfbTestType::BULK), - testing::Values( - // TODO: enable more tests when mixed beam width is supported - BeamConfig{1, {1}}, BeamConfig{2, {2}} // , BeamConfig{2, {1, 2}} - ), - testing::Values(std::nullopt, 1280), // maxTokensInPagedKvCache - testing::Values(std::nullopt, 0.4), // freeGpuMemoryFraction - testing::Values(true), // enableTrtOverlap - testing::Values(true), // enableChunkedContext - testing::Values(false), // enableStreamingMode - testing::Values(false), // enableCudaGraphMode - testing::Values(std::nullopt), // hostCacheSize - testing::Values(true), // useRandomEndId - testing::Values(std::vector{1, 2, 8}), // batchSizes - testing::Values(std::nullopt) // maxNumTokens - ), - generateTestName); - -INSTANTIATE_TEST_SUITE_P(GptSqTests, ParamTest, - testing::Combine(testing::Values(gptModelParams), - testing::Values(ModelSpec{INPUT_FILE, nvinfer1::DataType::kHALF} - .useGptAttentionPlugin() - .usePackedInput() - .setKVCacheType(KVCacheType::kPAGED) - .setQuantMethod(QuantMethod::kSMOOTH_QUANT), - ModelSpec{INPUT_FILE, nvinfer1::DataType::kHALF, - []() -> std::shared_ptr - { - auto pModelSpec = std::make_shared(INPUT_FILE, nvinfer1::DataType::kHALF); - pModelSpec->useGptAttentionPlugin() - .usePackedInput() - .setKVCacheType(KVCacheType::kPAGED) - .setQuantMethod(QuantMethod::kSMOOTH_QUANT); - return pModelSpec; - }()} - .useGptAttentionPlugin() - .usePackedInput() - .setKVCacheType(KVCacheType::kDISABLED) - .setQuantMethod(QuantMethod::kSMOOTH_QUANT)), - testing::Values(TrtGptModelType::InflightFusedBatching), - testing::Values( - TrtGptModelIfbTestType::BULK, TrtGptModelIfbTestType::WAVEFRONT, TrtGptModelIfbTestType::RANDOM), - testing::Values( - // TODO: enable more tests when mixed beam width is supported - // FIXME: disabled flaky beam search tests (https://nvbugspro.nvidia.com/bug/4646234) - BeamConfig{1, {1}} //, BeamConfig{2, {2}} - ), - testing::Values(std::nullopt), // maxTokensInPagedKvCache - testing::Values(0.4), // freeGpuMemoryFraction - testing::Values(false), // enableTrtOverlap - testing::Values(false), // enableChunkedContext - testing::Values(false), // enableStreamingMode - testing::Values(false), // enableCudaGraphMode - testing::Values(std::nullopt), // hostCacheSize - testing::Values(false), // useRandomEndId - testing::Values(std::vector{1, 2, 8}), // batchSizes - testing::Values(std::nullopt) // maxNumTokens - ), - generateTestName); - -// disabled because paused requests generate different tokens after resuming -INSTANTIATE_TEST_SUITE_P(DISABLED_GptChunkedContextTests, ParamTest, - testing::Combine(testing::Values(gptModelParams), - testing::Values( - // - ModelSpec{INPUT_FILE, nvinfer1::DataType::kHALF, getGptChunkedContextTestsCompareModelSpec()} - .useGptAttentionPlugin() - .usePackedInput() - .setKVCacheType(KVCacheType::kPAGED) - .setMaxInputLength(128)), - testing::Values(TrtGptModelType::InflightFusedBatching), - testing::Values(TrtGptModelIfbTestType::BULK), // TrtGptModelIfbTestType - testing::Values(BeamConfig{1, {1}}), // beam config - testing::Values(257), // maxTokensInPagedKvCache - testing::Values(0.4), // freeGpuMemoryFraction - testing::Values(false), // enableTrtOverlap - testing::Values(true), // enableChunkedContext - testing::Values(false), // enableStreamingMode - testing::Values(false), // enableCudaGraphMode - testing::Values(std::nullopt), // hostCacheSize - testing::Values(false), // useRandomEndId - testing::Values(std::vector{1, 2, 8}), // batchSizes - testing::Values(std::nullopt) // maxNumTokens - ), - generateTestName); - -INSTANTIATE_TEST_SUITE_P(GptChunkedLongContextTests, ParamTest, - testing::Combine(testing::Values(gptModelParams), - testing::Values( - // - ModelSpec{LONG_INPUT_FILE, nvinfer1::DataType::kHALF} - .useGptAttentionPlugin() - .usePackedInput() - .setKVCacheType(KVCacheType::kPAGED) - .setMaxInputLength(128), - ModelSpec{LONG_INPUT_FILE, nvinfer1::DataType::kHALF} - .useGptAttentionPlugin() - .usePackedInput() - .setKVCacheType(KVCacheType::kPAGED) - .useDraftTokensExternalDecoding() - .setDraftTokens(5)), - testing::Values(TrtGptModelType::InflightFusedBatching), - testing::Values(TrtGptModelIfbTestType::BULK, TrtGptModelIfbTestType::WAVEFRONT, - TrtGptModelIfbTestType::RANDOM), // TrtGptModelIfbTestType - testing::Values(BeamConfig{1, {1}}), // beam config - testing::Values(std::nullopt), // maxTokensInPagedKvCache - testing::Values(0.4), // freeGpuMemoryFraction - testing::Values(true), // enableTrtOverlap - testing::Values(true), // enableChunkedContext - testing::Values(false), // enableStreamingMode - testing::Values(false), // enableCudaGraphMode - testing::Values(std::nullopt), // hostCacheSize - testing::Values(false), // useRandomEndId - testing::Values(std::vector{1, 2, 8}), // batchSizes - testing::Values(64) // maxNumTokens - ), - generateTestName); - -INSTANTIATE_TEST_SUITE_P(GptDraftTests, ParamTest, - testing::Combine(testing::Values(gptModelParams), - testing::Values( - // - ModelSpec{INPUT_FILE, nvinfer1::DataType::kHALF, getGptDraftTestsCompareModelSpec()} - .useGptAttentionPlugin() - .usePackedInput() - .setKVCacheType(KVCacheType::kPAGED) - .useDraftTokensExternalDecoding() - .setDraftTokens(5) - .replaceLogits() - .collectGenerationLogitsFile() - .collectContextLogitsFile(), - ModelSpec{INPUT_FILE, nvinfer1::DataType::kHALF, getGptDraftTestsCompareModelSpec()} - .useGptAttentionPlugin() - .usePackedInput() - .setKVCacheType(KVCacheType::kPAGED) - .useDraftTokensExternalDecoding() - .setDraftTokens(5) - .useAcceptByLogits() - .replaceLogits() - .collectGenerationLogitsFile() - .collectContextLogitsFile()), - testing::Values(TrtGptModelType::InflightFusedBatching), - testing::Values( - TrtGptModelIfbTestType::BULK, TrtGptModelIfbTestType::WAVEFRONT, TrtGptModelIfbTestType::RANDOM), - testing::Values(BeamConfig{1, {1}}), // beamConfig - testing::Values(std::nullopt), // maxTokensInPagedKvCache - testing::Values(0.4), // freeGpuMemoryFraction - testing::Values(false), // enableTrtOverlap - testing::Values(true), // enableChunkedContext - testing::Values(false), // enableStreamingMode - testing::Values(false), // enableCudaGraphMode - testing::Values(std::nullopt), // hostCacheSize - testing::Values(false, true), // useRandomEndId - testing::Values(std::vector{1, 2, 8}), // batchSizes - testing::Values(std::nullopt) // maxNumTokens - ), - generateTestName); - -INSTANTIATE_TEST_SUITE_P(GptLogitsTests, ParamTest, - testing::Combine(testing::Values(gptModelParams), - testing::Values( - // modelSpec - ModelSpec{INPUT_FILE, nvinfer1::DataType::kHALF} - .useGptAttentionPlugin() - .usePackedInput() - .setKVCacheType(KVCacheType::kPAGED) - .gatherLogits() - .collectGenerationLogitsFile() - .collectContextLogitsFile()), - testing::Values(TrtGptModelType::InflightBatching, TrtGptModelType::InflightFusedBatching), // modelType - testing::Values(TrtGptModelIfbTestType::BULK, TrtGptModelIfbTestType::WAVEFRONT, - TrtGptModelIfbTestType::RANDOM), // testType - testing::Values(BeamConfig{1, {1}}), // beamConfig - testing::Values(std::nullopt), // maxTokensInPagedKvCache - testing::Values(0.4), // freeGpuMemoryFraction - testing::Values(false), // enableTrtOverlap - testing::Values(true), // enableChunkedContext - testing::Values(false, true), // enableStreamingMode - testing::Values(false), // enableCudaGraphMode - testing::Values(std::nullopt), // hostCacheSize - testing::Values(true), // useRandomEndId - testing::Values(std::vector{1, 2, 8}), // batchSizes - testing::Values(std::nullopt) // maxNumTokens - ), - generateTestName); - -INSTANTIATE_TEST_SUITE_P(GptLogProbsTests, ParamTest, - testing::Combine(testing::Values(gptModelParams), - testing::Values( - // modelSpec - ModelSpec{INPUT_FILE, nvinfer1::DataType::kHALF} - .useGptAttentionPlugin() - .usePackedInput() - .setKVCacheType(KVCacheType::kPAGED) - .returnLogProbs() - .collectCumLogProbsFile() - .collectLogProbsFile()), - testing::Values(TrtGptModelType::InflightFusedBatching), // modelType - testing::Values(TrtGptModelIfbTestType::BULK, TrtGptModelIfbTestType::WAVEFRONT, - TrtGptModelIfbTestType::RANDOM), // testType - testing::Values(BeamConfig{1, {1}}), // beamConfig - testing::Values(std::nullopt), // maxTokensInPagedKvCache - testing::Values(0.4), // freeGpuMemoryFraction - testing::Values(false), // enableTrtOverlap - testing::Values(true), // enableChunkedContext - testing::Values(false), // enableStreamingMode - testing::Values(false), // enableCudaGraphMode - testing::Values(std::nullopt), // hostCacheSize - testing::Values(false), // useRandomEndId - testing::Values(std::vector{1, 2, 8}), // batchSizes - testing::Values(std::nullopt) // maxNumTokens - ), - generateTestName); - -INSTANTIATE_TEST_SUITE_P(GptjTests, ParamTest, - testing::Combine(testing::Values(ModelParams{GPTJ_MODEL_DIR, {50256, 50256}}), - testing::Values( - // - ModelSpec{INPUT_FILE, nvinfer1::DataType::kHALF} - .useGptAttentionPlugin() - .setKVCacheType(KVCacheType::kCONTINUOUS) - .usePackedInput(), - ModelSpec{INPUT_FILE, nvinfer1::DataType::kHALF} - .useGptAttentionPlugin() - .setKVCacheType(KVCacheType::kPAGED) - .usePackedInput() - - ), - testing::Values(TrtGptModelType::InflightFusedBatching), - // WAR: disable wavefront and random tests on because of switched beams - testing::Values(TrtGptModelIfbTestType::BULK - /* , TrtGptModelIfbTestType::WAVEFRONT, TrtGptModelIfbTestType::RANDOM */), - testing::Values( - // TODO: enable more tests when mixed beam width is supported - BeamConfig{1, {1}}, BeamConfig{2, {2}} // , BeamConfig{2, {1, 2}} - ), - testing::Values(std::nullopt), // maxTokensInPagedKvCache - testing::Values(0.4), // freeGpuMemoryFraction - testing::Values(false), // enableTrtOverlap - testing::Values(false), // enableChunkedContext - testing::Values(false), // enableStreamingMode - testing::Values(false), // enableCudaGraphMode - testing::Values(std::nullopt), // hostCacheSize - testing::Values(false), // useRandomEndId - testing::Values(std::vector{1, 2, 8}), // batchSizes - testing::Values(std::nullopt) // maxNumTokens - ), - generateTestName); - -INSTANTIATE_TEST_SUITE_P(MambaTests, ParamTest, - testing::Combine(testing::Values(ModelParams{MAMBA_MODEL_DIR, {0, 1}}), - testing::Values( - // - ModelSpec{INPUT_FILE, nvinfer1::DataType::kHALF} - .useGptAttentionPlugin() - .setKVCacheType(KVCacheType::kCONTINUOUS) - .usePackedInput(), - ModelSpec{INPUT_FILE, nvinfer1::DataType::kHALF} - .useGptAttentionPlugin() - .setKVCacheType(KVCacheType::kPAGED) - .usePackedInput() - - ), - testing::Values(TrtGptModelType::InflightBatching), - testing::Values( - TrtGptModelIfbTestType::BULK, TrtGptModelIfbTestType::WAVEFRONT, TrtGptModelIfbTestType::RANDOM), - testing::Values(BeamConfig{1, {1}}), - testing::Values(std::nullopt), // maxTokensInPagedKvCache - testing::Values(0.4), // freeGpuMemoryFraction - testing::Values(false), // enableTrtOverlap - testing::Values(false), // enableChunkedContext - testing::Values(false), // enableStreamingMode - testing::Values(false), // enableCudaGraphMode - testing::Values(std::nullopt), // hostCacheSize - testing::Values(false), // useRandomEndId - testing::Values(std::vector{1, 2, 8}), // batchSizes - testing::Values(std::nullopt) // maxNumTokens - ), - generateTestName); - -INSTANTIATE_TEST_SUITE_P(RecurrentGemmaTests, ParamTest, - testing::Combine(testing::Values(ModelParams{RECURRENTGEMMA_MODEL_DIR, {0, 1}}), - testing::Values(ModelSpec{INPUT_FILE, nvinfer1::DataType::kHALF} - .useGptAttentionPlugin() - .setKVCacheType(KVCacheType::kPAGED) - .usePackedInput() - - ), - testing::Values(TrtGptModelType::InflightBatching), - testing::Values( - TrtGptModelIfbTestType::BULK, TrtGptModelIfbTestType::WAVEFRONT, TrtGptModelIfbTestType::RANDOM), - testing::Values(BeamConfig{1, {1}}), - testing::Values(std::nullopt), // maxTokensInPagedKvCache - testing::Values(0.4), // freeGpuMemoryFraction - testing::Values(false), // enableTrtOverlap - testing::Values(false), // enableChunkedContext - testing::Values(false), // enableStreamingMode - testing::Values(false), // enableCudaGraphMode - testing::Values(std::nullopt), // hostCacheSize - testing::Values(false), // useRandomEndId - testing::Values(std::vector{1, 2, 8}), // batchSizes - testing::Values(std::nullopt) // maxNumTokens - ), - generateTestName); - -INSTANTIATE_TEST_SUITE_P(LlamaTests, ParamTest, - testing::Combine(testing::Values(ModelParams{LLAMA_MODEL_DIR, {LLAMA_END_ID, LLAMA_PAD_ID}}), - testing::Values( - // - ModelSpec{INPUT_LLAMA_FILE, nvinfer1::DataType::kHALF} - .useGptAttentionPlugin() - .setKVCacheType(KVCacheType::kPAGED) - .usePackedInput(), - ModelSpec{INPUT_LLAMA_FILE, nvinfer1::DataType::kHALF} - .useGptAttentionPlugin() - .usePackedInput() - .setKVCacheType(KVCacheType::kPAGED) - .usePipelineParallelism(4), - ModelSpec{INPUT_LLAMA_FILE, nvinfer1::DataType::kHALF} - .useGptAttentionPlugin() - .usePackedInput() - .setKVCacheType(KVCacheType::kPAGED) - .useTensorParallelism(4), - ModelSpec{INPUT_LLAMA_FILE, nvinfer1::DataType::kHALF} - .useGptAttentionPlugin() - .usePackedInput() - .setKVCacheType(KVCacheType::kPAGED) - .usePipelineParallelism(2) - .useTensorParallelism(2) - - ), - testing::Values(TrtGptModelType::InflightFusedBatching), - testing::Values( - TrtGptModelIfbTestType::BULK, TrtGptModelIfbTestType::WAVEFRONT, TrtGptModelIfbTestType::RANDOM), - testing::Values( - // TODO: enable more tests when mixed beam width is supported - BeamConfig{1, {1}}, BeamConfig{2, {2}} // , BeamConfig{2, {1, 2}} - ), - testing::Values(std::nullopt), // maxTokensInPagedKvCache - testing::Values(0.4), // freeGpuMemoryFraction - testing::Values(false), // enableTrtOverlap - testing::Values(false), // enableChunkedContext - testing::Values(false), // enableStreamingMode - testing::Values(false), // enableCudaGraphMode - testing::Values(std::nullopt), // hostCacheSize - testing::Values(false), // useRandomEndId - testing::Values(std::vector{1, 2, 8}), // batchSizes - testing::Values(std::nullopt) // maxNumTokens - ), - generateTestName); - -INSTANTIATE_TEST_SUITE_P(ChatGlmTests, ParamTest, - testing::Combine(testing::Values(ModelParams{CHATGLM_MODEL_DIR, {130005, 3}}), - testing::Values( - // - ModelSpec{CHATGLM_INPUT_FILE, nvinfer1::DataType::kHALF} - .useGptAttentionPlugin() - .usePackedInput() - .setKVCacheType(KVCacheType::kPAGED)), - testing::Values(TrtGptModelType::InflightFusedBatching), - testing::Values( - TrtGptModelIfbTestType::BULK, TrtGptModelIfbTestType::WAVEFRONT, TrtGptModelIfbTestType::RANDOM), - testing::Values(BeamConfig{1, {1}}), - testing::Values(std::nullopt), // maxTokensInPagedKvCache - testing::Values(0.4), // freeGpuMemoryFraction - testing::Values(false), // enableTrtOverlap - testing::Values(false, true), // enableChunkedContext - testing::Values(false), // enableStreamingMode - testing::Values(false), // enableCudaGraphMode - testing::Values(std::nullopt), // hostCacheSize - testing::Values(false), // useRandomEndId - testing::Values(std::vector{1, 2, 8}), // batchSizes - testing::Values(std::nullopt) // maxNumTokens - ), - generateTestName); - -// ChatGlm0Tests is for glm-10b. -INSTANTIATE_TEST_SUITE_P(ChatGlm0Tests, ParamTest, - testing::Combine(testing::Values(ModelParams{GLM_MODEL_DIR, {50258, 50256}}), - testing::Values( - // - ModelSpec{GLM_INPUT_FILE, nvinfer1::DataType::kHALF} - .useGptAttentionPlugin() - .usePackedInput() - .setKVCacheType(KVCacheType::kPAGED)), - testing::Values(TrtGptModelType::InflightFusedBatching), - testing::Values( - TrtGptModelIfbTestType::BULK, TrtGptModelIfbTestType::WAVEFRONT, TrtGptModelIfbTestType::RANDOM), - testing::Values(BeamConfig{1, {1}}), - testing::Values(std::nullopt), // maxTokensInPagedKvCache - testing::Values(0.4), // freeGpuMemoryFraction - testing::Values(false), // enableTrtOverlap - testing::Values(false), // enableChunkedContext - testing::Values(false), // enableStreamingMode - testing::Values(false), // enableCudaGraphMode - testing::Values(std::nullopt), // hostCacheSize - testing::Values(false), // useRandomEndId - testing::Values(std::vector{1, 2, 8}), // batchSizes - testing::Values(std::nullopt) // maxNumTokens - ), - generateTestName); - -// https://nvbugspro.nvidia.com/bug/4640177 -// WAVEFRONT and RANDOM are disabled because of the accuracy mismatch -INSTANTIATE_TEST_SUITE_P(MedusaTests, ParamTest, - testing::Combine(testing::Values(ModelParams{MEDUSA_MODEL_DIR, {2, 2}}), - testing::Values( - // - ModelSpec{INPUT_VICUNA_FILE, nvinfer1::DataType::kHALF, getMedusaTestsCompareModelSpec()} - .useGptAttentionPlugin() - .usePackedInput() - .setKVCacheType(KVCacheType::kPAGED) - .useMedusa()), - testing::Values(TrtGptModelType::InflightFusedBatching), - testing::Values( - TrtGptModelIfbTestType::BULK, TrtGptModelIfbTestType::WAVEFRONT, TrtGptModelIfbTestType::RANDOM), - testing::Values(BeamConfig{1, {1}}), - testing::Values(std::nullopt), // maxTokensInPagedKvCache - testing::Values(0.4), // freeGpuMemoryFraction - testing::Values(false), // enableTrtOverlap - testing::Values(true), // enableChunkedContext - testing::Values(false), // enableStreamingMode - testing::Values(true, false), // enableCudaGraphMode - testing::Values(std::nullopt), // hostCacheSize - testing::Values(false), // useRandomEndId - testing::Values(std::vector{8}), // batchSizes - testing::Values(std::nullopt) // maxNumTokens - ), - generateTestName); - -INSTANTIATE_TEST_SUITE_P(EagleTests, ParamTest, - testing::Combine(testing::Values(ModelParams{EAGLE_MODEL_DIR, {2, 2}}), - testing::Values( - // - ModelSpec{INPUT_VICUNA_FILE, nvinfer1::DataType::kHALF, getEagleTestsCompareModelSpec()} - .useGptAttentionPlugin() - .usePackedInput() - .setKVCacheType(KVCacheType::kPAGED) - .useEagle()), - testing::Values(TrtGptModelType::InflightFusedBatching), - testing::Values( - TrtGptModelIfbTestType::BULK, TrtGptModelIfbTestType::WAVEFRONT, TrtGptModelIfbTestType::RANDOM), - testing::Values(BeamConfig{1, {1}}), - testing::Values(std::nullopt), // maxTokensInPagedKvCache - testing::Values(0.4), // freeGpuMemoryFraction - testing::Values(false), // enableTrtOverlap - testing::Values(true), // enableChunkedContext - testing::Values(false), // enableStreamingMode - testing::Values(true, false), // enableCudaGraphMode - testing::Values(std::nullopt), // hostCacheSize - testing::Values(false), // useRandomEndId - testing::Values(std::vector{8}), // batchSizes - testing::Values(std::nullopt) // maxNumTokens - ), - generateTestName); - -INSTANTIATE_TEST_SUITE_P(LlamaLookaheadDecodingTests, ParamTest, - testing::Combine(testing::Values(ModelParams{LLAMA_MODEL_DIR, {LLAMA_END_ID, LLAMA_PAD_ID}}), - testing::Values( - // - ModelSpec{INPUT_LLAMA_FILE, nvinfer1::DataType::kHALF} - .useGptAttentionPlugin() - .usePackedInput() - .setKVCacheType(KVCacheType::kPAGED) - .useLookaheadDecoding()), - testing::Values(TrtGptModelType::InflightFusedBatching), - testing::Values( - TrtGptModelIfbTestType::BULK, TrtGptModelIfbTestType::WAVEFRONT, TrtGptModelIfbTestType::RANDOM), - testing::Values(BeamConfig{1, {1}}), // beamConfig - testing::Values(std::nullopt), // maxTokensInPagedKvCache - testing::Values(0.4), // freeGpuMemoryFraction - testing::Values(false), // enableTrtOverlap - testing::Values(false), // enableChunkedContext - testing::Values(false), // enableStreamingMode - testing::Values(false), // enableCudaGraphMode - testing::Values(std::nullopt), // hostCacheSize - testing::Values(true), // useRandomEndId - testing::Values(std::vector{1, 16}), // batchSizes - testing::Values(std::nullopt) // maxNumTokens - ), - - generateTestName); - -INSTANTIATE_TEST_SUITE_P(ExplicitDraftTokensDecodingTests, ParamTest, - testing::Combine(testing::Values(ModelParams{EXPLICIT_DRAFT_MODEL_DIR, {2, 2}}), - testing::Values( - // - ModelSpec{INPUT_VICUNA_FILE, nvinfer1::DataType::kHALF} - .useGptAttentionPlugin() - .usePackedInput() - .setKVCacheType(KVCacheType::kPAGED) - .useExplicitDraftTokensDecoding() - .setMaxOutputLength(128)), - testing::Values(TrtGptModelType::InflightFusedBatching), - testing::Values( - TrtGptModelIfbTestType::BULK, TrtGptModelIfbTestType::WAVEFRONT, TrtGptModelIfbTestType::RANDOM), - testing::Values(BeamConfig{1, {1}}), // beamConfig - testing::Values(std::nullopt), // maxTokensInPagedKvCache - testing::Values(0.4), // freeGpuMemoryFraction - testing::Values(false), // enableTrtOverlap - testing::Values(true), // enableChunkedContext - testing::Values(false), // enableStreamingMode - testing::Values(false), // enableCudaGraphMode - testing::Values(std::nullopt), // hostCacheSize - testing::Values(false), // useRandomEndId - testing::Values(std::vector{8}), // batchSizes - testing::Values(std::nullopt) // maxNumTokens - ), - - generateTestName); - -#ifdef ENABLE_FP8 -// Using IFB-enabled engine -INSTANTIATE_TEST_SUITE_P(GptjFP8Tests, ParamTest, - testing::Combine(testing::Values(ModelParams{GPTJ_MODEL_DIR, {50256, 50256}}), - testing::Values( - // - ModelSpec{INPUT_FILE, nvinfer1::DataType::kFP8} - .useGptAttentionPlugin() - .setKVCacheType(KVCacheType::kPAGED) - .usePackedInput() - - ), - testing::Values(TrtGptModelType::InflightFusedBatching), - testing::Values( - TrtGptModelIfbTestType::BULK, TrtGptModelIfbTestType::WAVEFRONT, TrtGptModelIfbTestType::RANDOM), - testing::Values( - // TODO: enable more tests when supported - BeamConfig{1, {1}} // , BeamConfig{2, {2}}, BeamConfig{2, {1, 2}} - ), - testing::Values(std::nullopt), // maxTokensInPagedKvCache - testing::Values(0.4), // freeGpuMemoryFraction - testing::Values(false), // enableTrtOverlap - testing::Values(true), // enableChunkedContext - testing::Values(false), // enableStreamingMode - testing::Values(false), // enableCudaGraphMode - testing::Values(std::nullopt), // hostCacheSize - testing::Values(false), // useRandomEndId - testing::Values(std::vector{1, 2, 8}), // batchSizes - testing::Values(std::nullopt) // maxNumTokens - ), - generateTestName); - -#endif diff --git a/cpp/tests/e2e_tests/batch_manager/trtGptModelTest.cpp b/cpp/tests/e2e_tests/batch_manager/trtGptModelTest.cpp deleted file mode 100644 index 268e9bf9a238..000000000000 --- a/cpp/tests/e2e_tests/batch_manager/trtGptModelTest.cpp +++ /dev/null @@ -1,1328 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2023-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. - */ - -#ifndef TOP_LEVEL_DIR -#error "Define TOP_LEVEL_DIR" -#endif - -#include "tensorrt_llm/batch_manager/trtGptModel.h" -#include "tensorrt_llm/batch_manager/kvCacheManager.h" -#include "tensorrt_llm/batch_manager/trtGptModelInflightBatching.h" -#include "tensorrt_llm/plugins/api/tllmPlugin.h" -#include "tensorrt_llm/runtime/gptJsonConfig.h" -#include "tensorrt_llm/runtime/rawEngine.h" -#include "tensorrt_llm/runtime/tllmLogger.h" -#include "tensorrt_llm/testing/modelSpec.h" - -#include -#include - -#include -#include -#include - -using ::testing::ElementsAre; -using namespace tensorrt_llm::runtime; -namespace fs = std::filesystem; -using tensorrt_llm::testing::ModelSpec; -using tensorrt_llm::testing::KVCacheType; - -using TensorPtr = ITensor::SharedPtr; - -namespace -{ -auto const TEST_RESOURCE_PATH = fs::path{TOP_LEVEL_DIR} / "cpp/tests/resources"; -auto const ENGINE_PATH = TEST_RESOURCE_PATH / "models/rt_engine"; -auto const GPT_MODEL_PATH = ENGINE_PATH / "gpt2"; -auto const LLAMA_MODEL_PATH = ENGINE_PATH / "Llama-3.2-1B"; -} // namespace - -namespace tensorrt_llm::batch_manager -{ - -class TrtGptModelTest : public ::testing::Test // NOLINT(cppcoreguidelines-pro-type-member-init) -{ -protected: - TrtGptModelTest(std::filesystem::path const& modelPath) - : mModelConfig(1, 1, 1, 0, 1, 1, nvinfer1::DataType::kFLOAT) - , mModelPath(modelPath) - { - } - - TrtGptModelTest() - : TrtGptModelTest(GPT_MODEL_PATH / GetModelSpec().getModelPath() / "tp1-pp1-cp1-gpu") - { - } - - static ModelSpec& GetModelSpec() - { - static ModelSpec modelSpec{"input_tokens.npy", nvinfer1::DataType::kHALF}; - modelSpec.useGptAttentionPlugin().usePackedInput().setKVCacheType(KVCacheType::kPAGED); - return modelSpec; - } - - void SetUp() override - { - std::filesystem::path trtEnginePath = mModelPath; - - mBeamWidth = 1; - - mLogger = std::make_shared(); - - initTrtLlmPlugins(mLogger.get()); - - auto const json = GptJsonConfig::parse(trtEnginePath / "config.json"); - mModelConfig = json.getModelConfig(); - mMaxNumRequests = mModelConfig.getMaxBatchSize(); - mMaxSeqLen = mModelConfig.getMaxSequenceLen(); - mWorldConfig = WorldConfig::mpi(); - mVocabSizePadded = mModelConfig.getVocabSizePadded(mWorldConfig.getSize()); - - auto const enginePath = trtEnginePath / json.engineFilename(mWorldConfig); - auto const dtype = mModelConfig.getDataType(); - - mRawEngine.reset(new RawEngine(enginePath)); - - mSamplingConfig.temperature = std::vector{1.0f}; - mSamplingConfig.minLength = std::vector{1}; - mSamplingConfig.randomSeed = std::vector{static_cast(42ul)}; - mSamplingConfig.topK = std::vector{0}; - mSamplingConfig.topP = std::vector{0.0f}; - mSamplingConfig.noRepeatNgramSize = std::vector{1 << 30}; - - mStream = std::make_unique(); - mManager = std::make_unique(mStream); - } - - void TearDown() override {} - - // Thin wrapper around the private TrtGptModelInflightBatching::changeBeamWidth(). - static void changeBeamWidth(std::shared_ptr const& model, SizeType32 beamWidth) - { - model->changeBeamWidth(beamWidth); - } - - void forwardRequestsToCompletion( - std::shared_ptr const& trtGptModel, RequestList& requestList, SizeType32 maxNumIterations) - { - SizeType32 numFinished = 0; - SizeType32 numIterations = 0; - while (numFinished < requestList.size() && numIterations < maxNumIterations) - { - if (numIterations > maxNumIterations) - { - FAIL() << "Iterations never finished"; - } - trtGptModel->forwardAsync(requestList); - trtGptModel->forwardSync(); - numFinished = 0; - for (auto& request : requestList) - { - if (request->isGenerationCompleteState()) - { - ++numFinished; - } - } - ++numIterations; - } - } - - int32_t mMaxNumRequests; - int32_t mMaxSeqLen; - int32_t mBeamWidth; - int32_t mVocabSizePadded; - SamplingConfig mSamplingConfig; - std::string mDataPath; - std::shared_ptr mLogger; - ModelConfig mModelConfig; - WorldConfig mWorldConfig; - std::unique_ptr mRawEngine; - std::unique_ptr mManager; - BufferManager::CudaStreamPtr mStream; - std::filesystem::path mModelPath; -}; - -class TrtGptModelLoraTest : public TrtGptModelTest -{ -protected: - TrtGptModelLoraTest() - : TrtGptModelTest(GPT_MODEL_PATH / GetModelSpec().getModelPath() / "tp1-pp1-cp1-gpu") - { - } - - static ModelSpec& GetModelSpec() - { - static ModelSpec modelSpec{"input_tokens.npy", nvinfer1::DataType::kHALF}; - modelSpec.useGptAttentionPlugin().usePackedInput().setKVCacheType(KVCacheType::kPAGED).useLoraPlugin(); - return modelSpec; - } -}; - -TEST_F(TrtGptModelTest, Forward) -{ - SamplingConfig inSamplingConfig; - inSamplingConfig.temperature = std::vector{2.0f}; - int correlationId = 0; - auto maxNewTokens = 4; - auto tokens = std::make_shared>(std::initializer_list{1, 2, 3, 4}); - auto llmRequest = std::make_shared(correlationId, maxNewTokens, tokens, inSamplingConfig, false); - - RequestList requestList{llmRequest}; - - auto& manager = *mManager; - std::vector newTokensHost(mMaxNumRequests, 5); - TensorPtr const fakeNewTokens - = manager.copyFrom(newTokensHost, ITensor::makeShape({mMaxNumRequests, 1}), MemoryType::kGPU); - - std::vector finished(mMaxNumRequests, false); - - executor::ExecutorConfig executorConfig; - executorConfig.setEnableTrtOverlap(false); - executorConfig.setMaxBeamWidth(mBeamWidth); - executorConfig.setSchedulerConfig(executor::SchedulerConfig{executor::CapacitySchedulerPolicy::kMAX_UTILIZATION}); - - auto trtGptModel = std::make_shared( - mLogger, mModelConfig, mWorldConfig, *mRawEngine, true, executorConfig, false); - - // Generate one token for the requests in request_table - // We need to sync with decoder - trtGptModel->forwardAsync(requestList); - trtGptModel->forwardSync(); - - EXPECT_EQ(requestList.size(), 1); - EXPECT_EQ(requestList.front()->getState(), LlmRequestState::kGENERATION_IN_PROGRESS); - EXPECT_EQ(requestList.front()->getNumTokens(0), 5); - EXPECT_EQ(requestList.front()->getMaxNumGeneratedTokens(), 1); - EXPECT_THAT(requestList.front()->getTokens(0), ElementsAre(1, 2, 3, 4, 2)); -} - -TEST_F(TrtGptModelTest, ChangeBeamWidthClearsCudaGraphCache) -{ - if (mModelConfig.getMaxBeamWidth() < 2) - { - GTEST_SKIP() << "Engine was built with max_beam_width < 2; cannot exercise changeBeamWidth()."; - } - - executor::ExecutorConfig executorConfig; - executorConfig.setEnableTrtOverlap(false); - // Configure the executor for max beam width = 2 so we can transition between - // operating beam widths 1 and 2. - executorConfig.setMaxBeamWidth(2); - executorConfig.setSchedulerConfig(executor::SchedulerConfig{executor::CapacitySchedulerPolicy::kMAX_UTILIZATION}); - - auto extendedRuntimePerfKnobConfig = executor::ExtendedRuntimePerfKnobConfig{}; - extendedRuntimePerfKnobConfig.setCudaGraphMode(true); - extendedRuntimePerfKnobConfig.setCudaGraphCacheSize(8); - executorConfig.setExtendedRuntimePerfKnobConfig(extendedRuntimePerfKnobConfig); - - auto trtGptModel = std::make_shared( - mLogger, mModelConfig, mWorldConfig, *mRawEngine, true, executorConfig, false); - - // Run a single beam=1 request to completion. After at least one generation step - // the model captures and caches a CUDA graph for the subsequent batch state. - SamplingConfig samplingConfig; - samplingConfig.beamWidth = 1; - samplingConfig.temperature = std::vector{1.0f}; - auto tokens = std::make_shared>(std::initializer_list{1, 2, 3, 4}); - auto llmRequest = std::make_shared( - /*requestId=*/0, /*maxNewTokens=*/4, tokens, samplingConfig, /*isStreaming=*/false); - RequestList requestList{llmRequest}; - - forwardRequestsToCompletion(trtGptModel, requestList, /*maxNumIterations=*/8); - - // Cache must have been populated by the captured generation graph(s). - EXPECT_GT(trtGptModel->numCachedCudaGraphs(), 0) - << "Expected the CUDA graph executor cache to be populated after running a " - "beam=1 request to completion."; - - // Drop the completed request before changing beam width (changeBeamWidth requires - // no in-flight requests). - requestList.clear(); - - // Switch operating beam width via the fixture's friend-access helper. - changeBeamWidth(trtGptModel, 2); - - EXPECT_EQ(trtGptModel->numCachedCudaGraphs(), 0) - << "changeBeamWidth() must invalidate the CUDA graph executor cache. Stale " - "cudaGraphExec_t instances captured against the previous decoder state " - "would otherwise be replayed against freshly allocated memory."; -} - -TEST_F(TrtGptModelLoraTest, Forward) -{ - SamplingConfig inSamplingConfig; - inSamplingConfig.temperature = std::vector{2.0f}; - int correlationId = 0; - auto maxNewTokens = 4; - auto tokens = std::make_shared>(std::initializer_list{1, 2, 3, 4}); - auto llmRequest = std::make_shared(correlationId, maxNewTokens, tokens, inSamplingConfig, false); - - RequestList requestList{llmRequest}; - - auto& manager = *mManager; - std::vector newTokensHost(mMaxNumRequests, 5); - TensorPtr const fakeNewTokens - = manager.copyFrom(newTokensHost, ITensor::makeShape({mMaxNumRequests, 1}), MemoryType::kGPU); - - std::vector finished(mMaxNumRequests, false); - - executor::ExecutorConfig executorConfig; - executorConfig.setEnableTrtOverlap(false); - executorConfig.setMaxBeamWidth(mBeamWidth); - executorConfig.setSchedulerConfig(executor::SchedulerConfig{executor::CapacitySchedulerPolicy::kMAX_UTILIZATION}); - - auto trtGptModel = std::make_shared( - mLogger, mModelConfig, mWorldConfig, *mRawEngine, true, executorConfig, false); - - // Generate one token for the requests in request_table - trtGptModel->forwardAsync(requestList); - trtGptModel->forwardSync(); - - EXPECT_EQ(requestList.size(), 1); - EXPECT_EQ(requestList.front()->getState(), LlmRequestState::kGENERATION_IN_PROGRESS); - EXPECT_EQ(requestList.front()->getNumTokens(0), 5); - EXPECT_EQ(requestList.front()->getMaxNumGeneratedTokens(), 1); - EXPECT_THAT(requestList.front()->getTokens(0), ElementsAre(1, 2, 3, 4, 2)); -} - -TEST_F(TrtGptModelTest, ForwardMaxNewTokens) -{ - executor::ExecutorConfig executorConfig; - executorConfig.setEnableTrtOverlap(false); - executorConfig.setMaxBeamWidth(mBeamWidth); - executorConfig.setSchedulerConfig( - executor::SchedulerConfig{executor::CapacitySchedulerPolicy::kGUARANTEED_NO_EVICT}); - - executor::KvCacheConfig kvCacheConfig; - kvCacheConfig.setMaxTokens(10000); - executorConfig.setKvCacheConfig(kvCacheConfig); - - auto trtGptModel = std::make_shared( - mLogger, mModelConfig, mWorldConfig, *mRawEngine, true, executorConfig, false); - - SamplingConfig inSamplingConfig; - inSamplingConfig.temperature = std::vector{2.0f}; - int correlationId = 0; - auto maxNewTokens = 4; - auto tokens = std::make_shared>(256); - std::iota(std::begin(*tokens), std::end(*tokens), 1); - auto llmRequest = std::make_shared(correlationId, maxNewTokens, tokens, inSamplingConfig, false); - - int correlationId2 = 2; - auto maxNewTokens2 = 8; - auto llmRequest2 = std::make_shared(correlationId2, maxNewTokens2, tokens, inSamplingConfig, false); - - RequestList requestList{llmRequest, llmRequest2}; - - auto& manager = *mManager; - std::vector finished(mMaxNumRequests, false); - - // Generate one token for the requests in request_table - // We call forward twice because the first call doesn't sync with decoder - SizeType32 maxNumIterations = 13; - forwardRequestsToCompletion(trtGptModel, requestList, maxNumIterations); - - for (auto& request : requestList) - { - auto outputTokens = request->getTokens(0); - if (request->mRequestId == correlationId) - { - EXPECT_EQ(outputTokens.size(), tokens->size() + maxNewTokens); - } - if (request->mRequestId == correlationId2) - { - EXPECT_EQ(outputTokens.size(), tokens->size() + maxNewTokens2); - } - } -} - -TEST_F(TrtGptModelTest, MaxNumTokensInChunked) -{ - executor::ExecutorConfig executorConfig; - executorConfig.setEnableTrtOverlap(false); - executorConfig.setEnableChunkedContext(true); - executorConfig.setMaxBeamWidth(mBeamWidth); - executorConfig.setSchedulerConfig( - executor::SchedulerConfig{executor::CapacitySchedulerPolicy::kGUARANTEED_NO_EVICT}); - - auto modelConfig = mModelConfig; - mModelConfig.setMaxNumTokens(200); - - auto trtGptModelIfb = std::make_shared( - mLogger, mModelConfig, mWorldConfig, *mRawEngine, true, executorConfig, false); - std::vector> trtGptModels{trtGptModelIfb}; - - for (auto trtGptModel : trtGptModels) - { - SamplingConfig inSamplingConfig; - inSamplingConfig.temperature = std::vector{2.0f}; - int correlationId = 0; - auto maxNewTokens = 4; - auto tokens = std::make_shared>(256); - std::iota(std::begin(*tokens), std::end(*tokens), 1); - auto llmRequest = std::make_shared(correlationId, maxNewTokens, tokens, inSamplingConfig, false); - - int correlationId2 = 2; - auto maxNewTokens2 = 8; - auto llmRequest2 = std::make_shared(correlationId2, maxNewTokens2, tokens, inSamplingConfig, false); - - RequestList requestList{llmRequest, llmRequest2}; - - auto& manager = *mManager; - std::vector finished(mMaxNumRequests, false); - - // Generate one token for the requests in request_table - // We call forward twice because the first call doesn't sync with decoder - SizeType32 maxNumIterations = 13; - forwardRequestsToCompletion(trtGptModel, requestList, maxNumIterations); - - for (auto& request : requestList) - { - auto outputTokens = request->getTokens(0); - if (request->mRequestId == correlationId) - { - EXPECT_EQ(outputTokens.size(), tokens->size() + maxNewTokens); - } - if (request->mRequestId == correlationId2) - { - EXPECT_EQ(outputTokens.size(), tokens->size() + maxNewTokens2); - } - } - } -} - -TEST_F(TrtGptModelTest, ForwardEndId) -{ - executor::ExecutorConfig executorConfig; - executorConfig.setEnableTrtOverlap(false); - executorConfig.setMaxBeamWidth(mBeamWidth); - executorConfig.setSchedulerConfig( - executor::SchedulerConfig{executor::CapacitySchedulerPolicy::kGUARANTEED_NO_EVICT}); - - executor::KvCacheConfig kvCacheConfig; - kvCacheConfig.setMaxTokens(10000); - executorConfig.setKvCacheConfig(kvCacheConfig); - - auto trtGptModel = std::make_shared( - mLogger, mModelConfig, mWorldConfig, *mRawEngine, true, executorConfig, false); - - SamplingConfig inSamplingConfig; - inSamplingConfig.temperature = std::vector{2.0f}; - int correlationId = 0; - auto maxNewTokens = 4; - auto endId = 107; - auto tokens = std::make_shared>(256); - std::iota(std::begin(*tokens), std::end(*tokens), 1); - auto llmRequest = std::make_shared(correlationId, maxNewTokens, tokens, inSamplingConfig, false, endId); - - int correlationId2 = 2; - auto maxNewTokens2 = 8; - auto llmRequest2 - = std::make_shared(correlationId2, maxNewTokens2, tokens, inSamplingConfig, false, endId); - - RequestList requestList{llmRequest, llmRequest2}; - - auto& manager = *mManager; - std::vector finished(mMaxNumRequests, false); - - // Generate one token for the requests in request_table - // We call forward twice because the first call doesn't sync with decoder - SizeType32 maxNumIterations = 13; - forwardRequestsToCompletion(trtGptModel, requestList, maxNumIterations); - - for (auto& request : requestList) - { - auto outputTokens = request->getTokens(0); - // endId token is generated at 2nd iteration, so expect 1 output token - if (request->mRequestId == correlationId) - { - EXPECT_EQ(outputTokens.size(), tokens->size() + 1); - } - if (request->mRequestId == correlationId2) - { - EXPECT_EQ(outputTokens.size(), tokens->size() + 1); - } - } -} - -TEST_F(TrtGptModelTest, ForwardNoEoS) -{ - executor::ExecutorConfig executorConfig; - executorConfig.setEnableTrtOverlap(false); - executorConfig.setMaxBeamWidth(mBeamWidth); - executorConfig.setSchedulerConfig(executor::SchedulerConfig{executor::CapacitySchedulerPolicy::kSTATIC_BATCH}); - - executor::KvCacheConfig kvCacheConfig; - kvCacheConfig.setMaxTokens(10000); - executorConfig.setKvCacheConfig(kvCacheConfig); - - auto trtGptModel = std::make_shared( - mLogger, mModelConfig, mWorldConfig, *mRawEngine, true, executorConfig, false); - - SamplingConfig inSamplingConfig; - inSamplingConfig.topP = {0.9}; - inSamplingConfig.temperature = {0.6}; - inSamplingConfig.minLength = {5}; - - auto tokens = std::make_shared>(256); - std::iota(std::begin(*tokens), std::end(*tokens), 1); - - RequestList requestList; - for (auto requestIdx = 0; requestIdx < mMaxNumRequests; requestIdx++) - { - auto llmRequest = std::make_shared(requestIdx, 8, tokens, inSamplingConfig, false, -1); - requestList.push_back(llmRequest); - } - - auto& manager = *mManager; - std::vector finished(mMaxNumRequests, false); - - // Generate one token for the requests in request_table - // We call forward twice because the first call doesn't sync with decoder - SizeType32 maxNumIterations = 13; - forwardRequestsToCompletion(trtGptModel, requestList, maxNumIterations); -} - -TEST_F(TrtGptModelTest, ForwardFinished) -{ - SamplingConfig inSamplingConfig; - inSamplingConfig.temperature = std::vector{2.0f}; - int correlationId = 0; - auto maxNewTokens = 2; - auto tokens = std::make_shared>(std::initializer_list{10, 9, 8, 7, 6}); - auto llmRequest = std::make_shared(correlationId, maxNewTokens, tokens, inSamplingConfig, false); - - RequestList requestList{llmRequest}; - - int mForwardCount = 0; - - auto& manager = *mManager; - std::vector newTokensHost(mMaxNumRequests, 5); - TensorPtr const fakeNewTokens - = manager.copyFrom(newTokensHost, ITensor::makeShape({mMaxNumRequests, 1}), MemoryType::kGPU); - - std::vector newTokensHost2(mMaxNumRequests, 4); - TensorPtr const fakeNewTokens2 - = manager.copyFrom(newTokensHost2, ITensor::makeShape({mMaxNumRequests, 1}), MemoryType::kGPU); - - // Below are only used if beam > 1 - // So we are just returning tensors with the correct shape, content is not important - std::vector outputIdsHost(mMaxNumRequests * (5 + 2), 5); - TensorPtr const fakeOutputIds - = manager.copyFrom(outputIdsHost, ITensor::makeShape({mMaxNumRequests, 1, 5 + 2}), MemoryType::kGPU); - - std::vector finishedFalse(mMaxNumRequests, false); - std::vector finishedTrue(mMaxNumRequests, true); - - executor::ExecutorConfig executorConfig; - executorConfig.setEnableTrtOverlap(false); - executorConfig.setMaxBeamWidth(mBeamWidth); - executorConfig.setSchedulerConfig(executor::SchedulerConfig{executor::CapacitySchedulerPolicy::kMAX_UTILIZATION}); - - auto trtGptModel = std::make_shared( - mLogger, mModelConfig, mWorldConfig, *mRawEngine, true, executorConfig, false); - - // Generate one token for the requests in request_table - trtGptModel->forwardAsync(requestList); - trtGptModel->forwardSync(); - - EXPECT_EQ(requestList.size(), 1); - EXPECT_EQ(requestList.front()->getState(), LlmRequestState::kGENERATION_IN_PROGRESS); - EXPECT_EQ(requestList.front()->getNumTokens(0), 6); - EXPECT_EQ(requestList.front()->getMaxNumGeneratedTokens(), 1); - EXPECT_THAT(requestList.front()->getTokens(0), ElementsAre(10, 9, 8, 7, 6, 10)); - - // Generate one more token - trtGptModel->forwardAsync(requestList); - trtGptModel->forwardSync(); - - EXPECT_EQ(requestList.size(), 1); - EXPECT_EQ(requestList.front()->getState(), LlmRequestState::kGENERATION_COMPLETE); - EXPECT_EQ(requestList.front()->getNumTokens(0), 7); - EXPECT_EQ(requestList.front()->getMaxNumGeneratedTokens(), 2); - EXPECT_THAT(requestList.front()->getTokens(0), ElementsAre(10, 9, 8, 7, 6, 10, 6)); -} - -TEST_F(TrtGptModelTest, ForwardStopWords) -{ - executor::ExecutorConfig executorConfig; - executorConfig.setEnableTrtOverlap(false); - executorConfig.setMaxBeamWidth(mBeamWidth); - executorConfig.setSchedulerConfig( - executor::SchedulerConfig{executor::CapacitySchedulerPolicy::kGUARANTEED_NO_EVICT}); - - executor::KvCacheConfig kvCacheConfig; - kvCacheConfig.setMaxTokens(10000); - executorConfig.setKvCacheConfig(kvCacheConfig); - - auto trtGptModel = std::make_shared( - mLogger, mModelConfig, mWorldConfig, *mRawEngine, true, executorConfig, false); - - SamplingConfig inSamplingConfig; - inSamplingConfig.temperature = std::vector{2.0f}; - int correlationId = 0; - auto maxNewTokens = 4; - auto tokens = std::make_shared>(std::initializer_list{10, 9, 8, 7, 6}); - std::optional endId(std::nullopt); - std::optional padId(std::nullopt); - std::optional embeddingBias(std::nullopt); - std::optional badWordsList(std::nullopt); - - auto& manager = *mManager; - // No stop words - { - auto llmRequest = std::make_shared(correlationId, maxNewTokens, tokens, inSamplingConfig, false); - RequestList requestList{llmRequest}; - trtGptModel->forwardAsync(requestList); - trtGptModel->forwardSync(); - trtGptModel->forwardAsync(requestList); - trtGptModel->forwardSync(); - trtGptModel->forwardAsync(requestList); - trtGptModel->forwardSync(); - trtGptModel->forwardAsync(requestList); - trtGptModel->forwardSync(); - EXPECT_EQ(requestList.front()->getState(), LlmRequestState::kGENERATION_COMPLETE); - EXPECT_THAT(requestList.front()->getTokens(0), ElementsAre(10, 9, 8, 7, 6, 10, 6, 10, 6)); - } - // With stop words - { - TensorPtr stopWordsList = manager.cpu(ITensor::makeShape({1, 2, 3}), nvinfer1::DataType::kINT32); - auto stopWordsPtr = bufferCast(*stopWordsList); - // make 10, 6 10 the tokens for the stop word: - stopWordsPtr[0] = 10; - stopWordsPtr[1] = 6; - stopWordsPtr[2] = 10; - stopWordsPtr[3] = 3; - stopWordsPtr[4] = -1; - stopWordsPtr[5] = -1; - - auto llmRequest = std::make_shared(correlationId, maxNewTokens, tokens, inSamplingConfig, false, - endId, padId, embeddingBias, badWordsList, stopWordsList); - RequestList requestList{llmRequest}; - trtGptModel->forwardAsync(requestList); - trtGptModel->forwardSync(); - trtGptModel->forwardAsync(requestList); - trtGptModel->forwardSync(); - trtGptModel->forwardAsync(requestList); - trtGptModel->forwardSync(); - EXPECT_EQ(requestList.front()->getState(), LlmRequestState::kGENERATION_COMPLETE); - EXPECT_THAT(requestList.front()->getTokens(0), ElementsAre(10, 9, 8, 7, 6, 10, 6, 10)); - } - - // With stop words - { - TensorPtr stopWordsList = manager.cpu(ITensor::makeShape({1, 2, 1}), nvinfer1::DataType::kINT32); - auto stopWordsPtr = bufferCast(*stopWordsList); - // make 10 is the token for the stop word: - stopWordsPtr[0] = 10; - stopWordsPtr[1] = 1; - - auto llmRequest = std::make_shared(correlationId, maxNewTokens, tokens, inSamplingConfig, false, - endId, padId, embeddingBias, badWordsList, stopWordsList); - RequestList requestList{llmRequest}; - trtGptModel->forwardAsync(requestList); - trtGptModel->forwardSync(); - EXPECT_EQ(requestList.front()->getState(), LlmRequestState::kGENERATION_COMPLETE); - EXPECT_THAT(requestList.front()->getTokens(0), ElementsAre(10, 9, 8, 7, 6, 10)); - } - - // Multiple requests, each with different stop words - { - // Request w/o stop words - auto llmRequest = std::make_shared(1, maxNewTokens, tokens, inSamplingConfig, false); - - TensorPtr stopWordsList2 = manager.cpu(ITensor::makeShape({1, 2, 1}), nvinfer1::DataType::kINT32); - { - auto stopWordsPtr = bufferCast(*stopWordsList2); - stopWordsPtr[0] = 10; - stopWordsPtr[1] = 1; - } - auto llmRequest2 = std::make_shared(2, maxNewTokens, tokens, inSamplingConfig, false, endId, padId, - embeddingBias, badWordsList, stopWordsList2); - - TensorPtr stopWordsList3 = manager.cpu(ITensor::makeShape({1, 2, 3}), nvinfer1::DataType::kINT32); - { - auto stopWordsPtr = bufferCast(*stopWordsList3); - stopWordsPtr[0] = 10; - stopWordsPtr[1] = 6; - stopWordsPtr[2] = 10; - stopWordsPtr[3] = 3; - stopWordsPtr[4] = -1; - stopWordsPtr[5] = -1; - } - auto llmRequest3 = std::make_shared(3, maxNewTokens, tokens, inSamplingConfig, false, endId, padId, - embeddingBias, badWordsList, stopWordsList3); - - RequestList requestList{llmRequest, llmRequest2, llmRequest3}; - - SizeType32 maxNumIterations(5); - forwardRequestsToCompletion(trtGptModel, requestList, maxNumIterations); - - for (auto& request : requestList) - { - auto outputTokens = request->getTokens(0); - if (request->mRequestId == 1) - { - EXPECT_EQ(outputTokens.size(), tokens->size() + maxNewTokens); - EXPECT_THAT(request->getTokens(0), ElementsAre(10, 9, 8, 7, 6, 10, 6, 10, 6)); - } - if (request->mRequestId == 2) - { - EXPECT_EQ(outputTokens.size(), tokens->size() + 1); - EXPECT_THAT(request->getTokens(0), ElementsAre(10, 9, 8, 7, 6, 10)); - } - if (request->mRequestId == 3) - { - EXPECT_EQ(outputTokens.size(), tokens->size() + 3); - EXPECT_THAT(request->getTokens(0), ElementsAre(10, 9, 8, 7, 6, 10, 6, 10)); - } - } - } -} - -TEST_F(TrtGptModelTest, ForwardBadWords) -{ - executor::ExecutorConfig executorConfig; - executorConfig.setEnableTrtOverlap(false); - executorConfig.setMaxBeamWidth(mBeamWidth); - executorConfig.setSchedulerConfig( - executor::SchedulerConfig{executor::CapacitySchedulerPolicy::kGUARANTEED_NO_EVICT}); - - executor::KvCacheConfig kvCacheConfig; - kvCacheConfig.setMaxTokens(10000); - executorConfig.setKvCacheConfig(kvCacheConfig); - - auto trtGptModel = std::make_shared( - mLogger, mModelConfig, mWorldConfig, *mRawEngine, true, executorConfig, false); - - SamplingConfig inSamplingConfig; - inSamplingConfig.temperature = std::vector{2.0f}; - int correlationId = 0; - auto maxNewTokens = 4; - auto tokens = std::make_shared>(std::initializer_list{10, 9, 8, 7, 6}); - std::optional endId(std::nullopt); - std::optional padId(std::nullopt); - std::optional embeddingBias(std::nullopt); - std::optional stopWordsList(std::nullopt); - - auto& manager = *mManager; - // No bad words - { - auto llmRequest = std::make_shared(correlationId, maxNewTokens, tokens, inSamplingConfig, false); - RequestList requestList{llmRequest}; - - SizeType32 maxNumIterations = 5; - forwardRequestsToCompletion(trtGptModel, requestList, maxNumIterations); - EXPECT_EQ(requestList.front()->getState(), LlmRequestState::kGENERATION_COMPLETE); - EXPECT_THAT(requestList.front()->getTokens(0), ElementsAre(10, 9, 8, 7, 6, 10, 6, 10, 6)); - } - // With bad words, multiple tokens - { - TensorPtr badWordsList = manager.cpu(ITensor::makeShape({1, 2, 3}), nvinfer1::DataType::kINT32); - auto badWordsPtr = bufferCast(*badWordsList); - // make 10, 6 10 the tokens for the bad word: - badWordsPtr[0] = 10; - badWordsPtr[1] = 6; - badWordsPtr[2] = 10; - badWordsPtr[3] = 3; - badWordsPtr[4] = -1; - badWordsPtr[5] = -1; - - auto llmRequest = std::make_shared(correlationId, maxNewTokens, tokens, inSamplingConfig, false, - endId, padId, embeddingBias, badWordsList, stopWordsList); - RequestList requestList{llmRequest}; - SizeType32 maxNumIterations = 5; - forwardRequestsToCompletion(trtGptModel, requestList, maxNumIterations); - EXPECT_EQ(requestList.front()->getState(), LlmRequestState::kGENERATION_COMPLETE); - // Token at position 7 should be different than 10 - EXPECT_NE(requestList.front()->getTokens(0).at(7), 10); - } - - // With bad words single token - { - TensorPtr badWordsList = manager.cpu(ITensor::makeShape({1, 2, 1}), nvinfer1::DataType::kINT32); - auto badWordsPtr = bufferCast(*badWordsList); - // make 10 is the token for the bad word: - badWordsPtr[0] = 10; - badWordsPtr[1] = 1; - - auto llmRequest = std::make_shared(correlationId, maxNewTokens, tokens, inSamplingConfig, false, - endId, padId, embeddingBias, badWordsList, stopWordsList); - RequestList requestList{llmRequest}; - SizeType32 maxNumIterations = 5; - forwardRequestsToCompletion(trtGptModel, requestList, maxNumIterations); - EXPECT_EQ(requestList.front()->getState(), LlmRequestState::kGENERATION_COMPLETE); - EXPECT_NE(requestList.front()->getTokens(0).at(5), 10); - } - - // Multiple requests, each with different bad words - { - // Request w/o bad words - auto llmRequest = std::make_shared(1, maxNewTokens, tokens, inSamplingConfig, false); - - TensorPtr badWordsList2 = manager.cpu(ITensor::makeShape({1, 2, 1}), nvinfer1::DataType::kINT32); - { - auto badWordsPtr = bufferCast(*badWordsList2); - badWordsPtr[0] = 10; - badWordsPtr[1] = 1; - } - auto llmRequest2 = std::make_shared(2, maxNewTokens, tokens, inSamplingConfig, false, endId, padId, - embeddingBias, badWordsList2, stopWordsList); - - TensorPtr badWordsList3 = manager.cpu(ITensor::makeShape({1, 2, 3}), nvinfer1::DataType::kINT32); - { - auto badWordsPtr = bufferCast(*badWordsList3); - badWordsPtr[0] = 10; - badWordsPtr[1] = 6; - badWordsPtr[2] = 10; - badWordsPtr[3] = 3; - badWordsPtr[4] = -1; - badWordsPtr[5] = -1; - } - auto llmRequest3 = std::make_shared(3, maxNewTokens, tokens, inSamplingConfig, false, endId, padId, - embeddingBias, badWordsList3, stopWordsList); - - RequestList requestList{llmRequest, llmRequest2, llmRequest3}; - - SizeType32 maxNumIterations(6); - forwardRequestsToCompletion(trtGptModel, requestList, maxNumIterations); - - for (auto& request : requestList) - { - auto outputTokens = request->getTokens(0); - if (request->mRequestId == 1) - { - EXPECT_EQ(outputTokens.size(), tokens->size() + maxNewTokens); - EXPECT_THAT(request->getTokens(0), ElementsAre(10, 9, 8, 7, 6, 10, 6, 10, 6)); - } - if (request->mRequestId == 2) - { - EXPECT_EQ(outputTokens.size(), tokens->size() + maxNewTokens); - EXPECT_NE(request->getTokens(0).at(5), 10); - } - if (request->mRequestId == 3) - { - EXPECT_EQ(outputTokens.size(), tokens->size() + maxNewTokens); - EXPECT_NE(request->getTokens(0).at(7), 10); - } - } - } -} - -TEST_F(TrtGptModelTest, ForwardEmbeddingBias) -{ - executor::ExecutorConfig executorConfig; - executorConfig.setEnableTrtOverlap(false); - executorConfig.setMaxBeamWidth(mBeamWidth); - executorConfig.setSchedulerConfig( - executor::SchedulerConfig{executor::CapacitySchedulerPolicy::kGUARANTEED_NO_EVICT}); - - executor::KvCacheConfig kvCacheConfig; - kvCacheConfig.setMaxTokens(10000); - executorConfig.setKvCacheConfig(kvCacheConfig); - - auto trtGptModelIfb = std::make_shared( - mLogger, mModelConfig, mWorldConfig, *mRawEngine, true, executorConfig, false); - - std::vector> trtGptModels{trtGptModelIfb}; - - for (auto& trtGptModel : trtGptModels) - { - SamplingConfig inSamplingConfig; - inSamplingConfig.temperature = std::vector{2.0f}; - int correlationId = 0; - auto maxNewTokens = 4; - auto tokens = std::make_shared>(std::initializer_list{10, 9, 8, 7, 6}); - std::optional endId(std::nullopt); - std::optional padId(std::nullopt); - std::optional badWordsList(std::nullopt); - std::optional stopWordsList(std::nullopt); - - auto& manager = *mManager; - // No bad words - { - auto llmRequest - = std::make_shared(correlationId, maxNewTokens, tokens, inSamplingConfig, false); - RequestList requestList{llmRequest}; - - SizeType32 maxNumIterations = 5; - forwardRequestsToCompletion(trtGptModel, requestList, maxNumIterations); - EXPECT_EQ(requestList.front()->getState(), LlmRequestState::kGENERATION_COMPLETE); - EXPECT_THAT(requestList.front()->getTokens(0), ElementsAre(10, 9, 8, 7, 6, 10, 6, 10, 6)); - } - // With embedding bias - { - TensorPtr embeddingBias - = manager.cpu(ITensor::makeShape({1, mVocabSizePadded}), nvinfer1::DataType::kFLOAT); - auto embeddingBiasPtr = bufferCast(*embeddingBias); - for (SizeType32 vi = 0; vi < mVocabSizePadded; ++vi) - { - embeddingBiasPtr[vi] = 0.f; - } - // bias all words to the 10th token - embeddingBiasPtr[10] = std::numeric_limits::max(); - - auto llmRequest = std::make_shared(correlationId, maxNewTokens, tokens, inSamplingConfig, false, - endId, padId, embeddingBias, badWordsList, stopWordsList); - RequestList requestList{llmRequest}; - SizeType32 maxNumIterations = 5; - forwardRequestsToCompletion(trtGptModel, requestList, maxNumIterations); - EXPECT_EQ(requestList.front()->getState(), LlmRequestState::kGENERATION_COMPLETE); - // All tokens should become 10 after applying bias - EXPECT_EQ(requestList.front()->getTokens(0).at(5), 10); - EXPECT_EQ(requestList.front()->getTokens(0).at(6), 10); - EXPECT_EQ(requestList.front()->getTokens(0).at(7), 10); - EXPECT_EQ(requestList.front()->getTokens(0).at(8), 10); - } - - // Multiple requests, each with different bias - { - // Request w/o bias - auto llmRequest = std::make_shared(1, maxNewTokens, tokens, inSamplingConfig, false); - - TensorPtr embeddingBias1 - = manager.cpu(ITensor::makeShape({1, mVocabSizePadded}), nvinfer1::DataType::kFLOAT); - auto embeddingBias1Ptr = bufferCast(*embeddingBias1); - for (SizeType32 vi = 0; vi < mVocabSizePadded; ++vi) - { - embeddingBias1Ptr[vi] = 0.f; - } - // bias all words to the 10th token - embeddingBias1Ptr[10] = std::numeric_limits::max(); - - auto llmRequest2 = std::make_shared(2, maxNewTokens, tokens, inSamplingConfig, false, endId, - padId, embeddingBias1, badWordsList, stopWordsList); - - TensorPtr embeddingBias2 - = manager.cpu(ITensor::makeShape({1, mVocabSizePadded}), nvinfer1::DataType::kFLOAT); - auto embeddingBias2Ptr = bufferCast(*embeddingBias2); - for (SizeType32 vi = 0; vi < mVocabSizePadded; ++vi) - { - embeddingBias2Ptr[vi] = 0.f; - } - // bias all words to the 100th token - embeddingBias2Ptr[100] = std::numeric_limits::max(); - - auto llmRequest3 = std::make_shared(3, maxNewTokens, tokens, inSamplingConfig, false, endId, - padId, embeddingBias2, badWordsList, stopWordsList); - - RequestList requestList{llmRequest, llmRequest2, llmRequest3}; - - SizeType32 maxNumIterations(6); - forwardRequestsToCompletion(trtGptModel, requestList, maxNumIterations); - - for (auto& request : requestList) - { - auto outputTokens = request->getTokens(0); - if (request->mRequestId == 1) - { - EXPECT_EQ(outputTokens.size(), tokens->size() + maxNewTokens); - EXPECT_THAT(request->getTokens(0), ElementsAre(10, 9, 8, 7, 6, 10, 6, 10, 6)); - } - if (request->mRequestId == 2) - { - EXPECT_EQ(outputTokens.size(), tokens->size() + maxNewTokens); - EXPECT_THAT(request->getTokens(0), ElementsAre(10, 9, 8, 7, 6, 10, 10, 10, 10)); - } - if (request->mRequestId == 3) - { - EXPECT_EQ(outputTokens.size(), tokens->size() + maxNewTokens); - EXPECT_THAT(request->getTokens(0), ElementsAre(10, 9, 8, 7, 6, 100, 100, 100, 100)); - } - } - } - } -} - -class TrtGptModelIfbHelper : public TrtGptModelInflightBatching -{ -public: - using TrtGptModelInflightBatching::TrtGptModelInflightBatching; - - [[nodiscard]] std::shared_ptr getKVCacheManager() const - { - return TrtGptModelInflightBatching::getKVCacheManager(); - } - - [[nodiscard]] SizeType32 getMaxAttentionWindow() const - { - return TrtGptModelInflightBatching::getMaxAttentionWindow(); - } -}; - -TEST_F(TrtGptModelTest, KVCacheReuseChunked) -{ - executor::ExecutorConfig executorConfig; - executorConfig.setEnableTrtOverlap(false); - executorConfig.setEnableChunkedContext(true); - executorConfig.setMaxBeamWidth(mBeamWidth); - executorConfig.setSchedulerConfig( - executor::SchedulerConfig{executor::CapacitySchedulerPolicy::kGUARANTEED_NO_EVICT}); - - executor::KvCacheConfig kvCacheConfig; - kvCacheConfig.setEnableBlockReuse(true); - executorConfig.setKvCacheConfig(kvCacheConfig); - - mModelConfig.setMaxNumTokens(384); - - for (int const numBlocksExpectedReused : {1, 2}) - { - auto trtGptModelIfb = std::make_shared( - mLogger, mModelConfig, mWorldConfig, *mRawEngine, true, executorConfig, false); - auto const cacheManager = trtGptModelIfb->getKVCacheManager(); - auto const tokensPerBlock = cacheManager->getTokensPerBlock(); - constexpr int numPrefillBlocks = 2; - - SamplingConfig inSamplingConfig; - inSamplingConfig.temperature = std::vector{2.0f}; - constexpr int correlationId = 0; - constexpr int maxNewTokens = 4; - - auto tokens = std::make_shared>(tokensPerBlock * numPrefillBlocks); - std::iota(std::begin(*tokens), std::end(*tokens), 1); - auto subTokens = std::make_shared>( - tokens->begin(), tokens->begin() + numBlocksExpectedReused * tokensPerBlock); - // Add new token to "start" a new block. - subTokens->push_back(0); - { - auto llmRequest - = std::make_shared(correlationId, maxNewTokens, tokens, inSamplingConfig, false); - RequestList requests{llmRequest}; - forwardRequestsToCompletion(trtGptModelIfb, requests, 6); - EXPECT_EQ(llmRequest->isGenerationCompleteState(), true); - } - for (size_t i = 1; i <= 2; ++i) - { - auto llmRequest - = std::make_shared(correlationId, maxNewTokens, subTokens, inSamplingConfig, false); - RequestList req{llmRequest}; - forwardRequestsToCompletion(trtGptModelIfb, req, 5); - EXPECT_EQ(cacheManager->getBlockManager().getNumReusedBlocks(), i * numBlocksExpectedReused); - } - } -} - -TEST_F(TrtGptModelTest, PauseRequestStats) -{ - SamplingConfig inSamplingConfig; - inSamplingConfig.temperature = std::vector{2.0f}; - int correlationId = 0; - auto maxNewTokens = 3; - auto tokens = std::make_shared>(std::initializer_list{1, 2, 3, 4}); - auto llmRequest = std::make_shared(correlationId, maxNewTokens, tokens, inSamplingConfig, false, - std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, - std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, - std::nullopt, std::nullopt, std::nullopt, std::nullopt, false, false, false, std::nullopt, std::nullopt, false, - std::nullopt, false, std::nullopt, false, std::nullopt, executor::Request::kDefaultPriority, std::nullopt, - std::nullopt, std::nullopt, LlmRequestType::LLMREQUEST_TYPE_CONTEXT_AND_GENERATION, std::nullopt, 1, - std::nullopt, std::nullopt, true /* returnPerfMetrics */); - - RequestList requestList{llmRequest}; - - executor::ExecutorConfig executorConfig; - executorConfig.setEnableTrtOverlap(false); - executorConfig.setMaxBeamWidth(mBeamWidth); - executorConfig.setSchedulerConfig(executor::SchedulerConfig{executor::CapacitySchedulerPolicy::kMAX_UTILIZATION}); - - auto trtGptModel = std::make_shared( - mLogger, mModelConfig, mWorldConfig, *mRawEngine, true, executorConfig, false); - - // Generate one token for the requests in request_table - // We need to sync with decoder - trtGptModel->forwardAsync(requestList); - trtGptModel->forwardSync(); - - EXPECT_EQ(requestList.size(), 1); - EXPECT_EQ(requestList.front()->getState(), LlmRequestState::kGENERATION_IN_PROGRESS); - EXPECT_EQ(requestList.front()->getNumTokens(0), 5); - EXPECT_EQ(requestList.front()->getMaxNumGeneratedTokens(), 1); - EXPECT_THAT(requestList.front()->getTokens(0), ElementsAre(1, 2, 3, 4, 2)); - - auto perfMetrics = requestList.front()->getPerfMetrics(); - auto zero = executor::RequestPerfMetrics::TimePoint{}; - - EXPECT_NE(perfMetrics.timingMetrics.arrivalTime, zero); - EXPECT_NE(perfMetrics.timingMetrics.firstScheduledTime, zero); - EXPECT_NE(perfMetrics.timingMetrics.firstTokenTime, zero); - EXPECT_EQ(perfMetrics.timingMetrics.lastTokenTime, zero); - EXPECT_EQ(perfMetrics.firstIter, 0); - EXPECT_EQ(perfMetrics.iter, 0); - EXPECT_EQ(perfMetrics.lastIter, std::nullopt); - - // Pause the request - trtGptModel->terminateRequest(llmRequest, true); - - // Resume work - trtGptModel->forwardAsync(requestList); - trtGptModel->forwardSync(); - - // Generate one more token - EXPECT_EQ(requestList.size(), 1); - EXPECT_EQ(requestList.front()->getState(), LlmRequestState::kGENERATION_IN_PROGRESS); - EXPECT_EQ(requestList.front()->getNumTokens(0), 6); - EXPECT_EQ(requestList.front()->getMaxNumGeneratedTokens(), 1); - EXPECT_THAT(requestList.front()->getTokens(0), ElementsAre(1, 2, 3, 4, 2, 4)); - - auto newPerfMetrics = requestList.front()->getPerfMetrics(); - EXPECT_EQ(newPerfMetrics.firstIter, 0); - EXPECT_EQ(newPerfMetrics.iter, 1); - EXPECT_EQ(newPerfMetrics.lastIter, std::nullopt); - - // Check that firstScheduledTime and firstTokenTime are the same - EXPECT_EQ(perfMetrics.timingMetrics.firstScheduledTime, newPerfMetrics.timingMetrics.firstScheduledTime); - EXPECT_EQ(perfMetrics.timingMetrics.firstTokenTime, newPerfMetrics.timingMetrics.firstTokenTime); - - // Pause the request - trtGptModel->terminateRequest(llmRequest, true); - - // Resume work - trtGptModel->forwardAsync(requestList); - trtGptModel->forwardSync(); - - // Generate last token - EXPECT_EQ(requestList.size(), 1); - EXPECT_EQ(requestList.front()->getState(), LlmRequestState::kGENERATION_COMPLETE); - EXPECT_EQ(requestList.front()->getNumTokens(0), 7); - EXPECT_EQ(requestList.front()->getMaxNumGeneratedTokens(), 1); - EXPECT_THAT(requestList.front()->getTokens(0), ElementsAre(1, 2, 3, 4, 2, 4, 2)); - - auto endPerfMetrics = requestList.front()->getPerfMetrics(); - EXPECT_EQ(endPerfMetrics.firstIter, 0); - EXPECT_EQ(endPerfMetrics.iter, 2); - EXPECT_EQ(endPerfMetrics.lastIter, 2); - - // Check that firstScheduledTime and firstTokenTime are the same - EXPECT_EQ(perfMetrics.timingMetrics.firstScheduledTime, endPerfMetrics.timingMetrics.firstScheduledTime); - EXPECT_EQ(perfMetrics.timingMetrics.firstTokenTime, endPerfMetrics.timingMetrics.firstTokenTime); -} - -class TrtGptModelLogitsTest : public TrtGptModelTest -{ -protected: - TrtGptModelLogitsTest() - : TrtGptModelTest(GPT_MODEL_PATH / GetModelSpec().getModelPath() / "tp1-pp1-cp1-gpu") - { - } - - static ModelSpec& GetModelSpec() - { - static ModelSpec modelSpec{"input_tokens.npy", nvinfer1::DataType::kHALF}; - modelSpec.useGptAttentionPlugin().usePackedInput().setKVCacheType(KVCacheType::kPAGED).gatherLogits(); - return modelSpec; - } -}; - -TEST_F(TrtGptModelLogitsTest, ReturnContextLogitsWithChunkedContext) -{ - // General config - int correlationId = 0; - auto maxNewTokens = 4; - int const worldSize = 1; - auto const vocabSizePadded = mModelConfig.getVocabSizePadded(worldSize); - - SamplingConfig inSamplingConfig; - - // Different prompt length - for (int const promptLength : {10, 128, 200, 250, 256}) - { - RequestList finishList; - for (bool enableChunkedContext : {false, true}) - { - auto modelConfig = mModelConfig; - if (enableChunkedContext) - { - modelConfig.setMaxNumTokens(128); - } - - executor::ExecutorConfig executorConfig; - executorConfig.setEnableTrtOverlap(false); - executorConfig.setMaxBeamWidth(mBeamWidth); - executorConfig.setEnableChunkedContext(enableChunkedContext); - executorConfig.setSchedulerConfig( - executor::SchedulerConfig{executor::CapacitySchedulerPolicy::kGUARANTEED_NO_EVICT}); - - executor::KvCacheConfig kvCacheConfig; - kvCacheConfig.setEnableBlockReuse(true); - executorConfig.setKvCacheConfig(kvCacheConfig); - - auto trtGptModelIfb = std::make_shared( - mLogger, modelConfig, mWorldConfig, *mRawEngine, true, executorConfig, false); - - // Prepare input tokens - std::vector input_ids; - for (int i = 1; i <= promptLength; i++) - { - input_ids.push_back(i); - } - auto tokens = std::make_shared>(input_ids); - - auto llmRequest - = std::make_shared(correlationId, maxNewTokens, tokens, inSamplingConfig, false); - TensorPtr contextLogitsHost = BufferManager::cpu( - ITensor::makeShape({llmRequest->mPromptLen, vocabSizePadded}), nvinfer1::DataType::kFLOAT); - - llmRequest->setContextLogitsHost(contextLogitsHost); - llmRequest->setReturnContextLogits(true); - - RequestList requestList{llmRequest}; - forwardRequestsToCompletion(trtGptModelIfb, requestList, 6); - - finishList.push_back(llmRequest); - } - EXPECT_EQ(finishList.size(), 2); - - float const* const disableChunkedContextLogits - = bufferCast(*(finishList.front()->getContextLogitsHost())); - float const* const enableChunkedContextLogits = bufferCast(*(finishList.back()->getContextLogitsHost())); - - for (int tokenIdx = 0; tokenIdx < promptLength; tokenIdx++) - { - for (int vocabIdx = 0; vocabIdx < vocabSizePadded; vocabIdx++) - { - size_t idx = tokenIdx * vocabSizePadded + vocabIdx; - EXPECT_NEAR(disableChunkedContextLogits[idx], enableChunkedContextLogits[idx], 1e-0) - << "tokenIdx=" << tokenIdx << " vocabIdx=" << vocabIdx; - } - } - finishList.clear(); - } -} - -class LlamaModelLADTest : public TrtGptModelTest -{ -protected: - LlamaModelLADTest() - : TrtGptModelTest(LLAMA_MODEL_PATH / GetModelSpec().getModelPath() / "tp1-pp1-cp1-gpu") - { - } - - static ModelSpec& GetModelSpec() - { - static ModelSpec modelSpec = ModelSpec{"input_tokens.npy", nvinfer1::DataType::kHALF} - .useGptAttentionPlugin() - .usePackedInput() - .setKVCacheType(KVCacheType::kPAGED) - .useLookaheadDecoding(); - return modelSpec; - } -}; - -TEST_F(LlamaModelLADTest, SeamlessLookaheadDecoding) -{ - GTEST_SKIP() << "Will enable this test when we have a force LAD support."; - SizeType32 requestId = 0; - for (bool const initLADConfig : {true, false}) - { - RequestList requestList{}; - for (SizeType32 i = 0; i < 8; ++i) - { - SamplingConfig inSamplingConfig; - int correlationId = requestId; - auto maxNewTokens = 8; - auto tokens = std::make_shared>(std::initializer_list{1, 2, 3, 4}); - auto llmRequest - = std::make_shared(correlationId, maxNewTokens, tokens, inSamplingConfig, false); - requestList.emplace_back(std::move(llmRequest)); - requestId += 1; - } - - executor::ExecutorConfig executorConfig; - executorConfig.setEnableChunkedContext(false); - executorConfig.setEnableTrtOverlap(false); - executorConfig.setMaxBeamWidth(1); - executorConfig.setSchedulerConfig( - executor::SchedulerConfig{executor::CapacitySchedulerPolicy::kMAX_UTILIZATION}); - if (initLADConfig) - { - executor::DecodingConfig decodingConfig; - decodingConfig.setLookaheadDecodingConfig(executor::LookaheadDecodingConfig(5, 5, 5)); - executorConfig.setDecodingConfig(decodingConfig); - } - - auto trtGptModel = std::make_shared( - mLogger, mModelConfig, mWorldConfig, *mRawEngine, true, executorConfig, false); - - // Generate tokens for the requests in request_table - // We need to sync with decoder - trtGptModel->forwardAsync(requestList); - trtGptModel->forwardSync(); - EXPECT_EQ(trtGptModel->getSpeculativeDecodingMode().isLookaheadDecoding(), true); - - // Add new requests - for (SizeType32 i = 0; i < 4; ++i) - { - SamplingConfig inSamplingConfig; - int correlationId = requestId; - auto maxNewTokens = 8; - auto tokens = std::make_shared>(std::initializer_list{1, 2, 3, 4}); - auto llmRequest - = std::make_shared(correlationId, maxNewTokens, tokens, inSamplingConfig, false); - requestList.emplace_back(std::move(llmRequest)); - requestId += 1; - } - trtGptModel->forwardAsync(requestList); - trtGptModel->forwardSync(); - EXPECT_EQ(trtGptModel->getSpeculativeDecodingMode().isLookaheadDecoding(), false); - - // Complete all of the requests - SizeType32 maxNumIterations = 8; - forwardRequestsToCompletion(trtGptModel, requestList, maxNumIterations); - - // Run new requests with lookahead - requestList.clear(); - for (SizeType32 i = 0; i < 4; ++i) - { - SamplingConfig inSamplingConfig; - int correlationId = requestId; - auto maxNewTokens = 8; - auto tokens = std::make_shared>(std::initializer_list{1, 2, 3, 4}); - auto llmRequest - = std::make_shared(correlationId, maxNewTokens, tokens, inSamplingConfig, false); - requestList.emplace_back(std::move(llmRequest)); - requestId += 1; - } - trtGptModel->forwardAsync(requestList); - trtGptModel->forwardSync(); - EXPECT_EQ(trtGptModel->getSpeculativeDecodingMode().isLookaheadDecoding(), true); - forwardRequestsToCompletion(trtGptModel, requestList, maxNumIterations); - requestList.clear(); - } -} - -TEST_F(TrtGptModelTest, ClampSeqLenToAttentionWindow) -{ - auto constexpr maxAttentionWindow = 65536; - auto constexpr maxSequenceLen = maxAttentionWindow + 1; - - executor::KvCacheConfig kvCacheConfig; - kvCacheConfig.setMaxAttentionWindowVec(std::vector{maxAttentionWindow}); - kvCacheConfig.setFreeGpuMemoryFraction(0.0001); // minuscule amount of memory to force a clamp - - executor::ExecutorConfig executorConfig; - executorConfig.setKvCacheConfig(kvCacheConfig); - executorConfig.setMaxBeamWidth(mBeamWidth); - - auto modelConfig = mModelConfig; - modelConfig.setMaxSequenceLen(maxSequenceLen); - - auto trtGptModel = std::make_shared( - mLogger, modelConfig, mWorldConfig, *mRawEngine, true, executorConfig, false); - EXPECT_LT(trtGptModel->getMaxAttentionWindow(), maxAttentionWindow); - EXPECT_EQ(trtGptModel->getMaxSequenceLen(), trtGptModel->getMaxAttentionWindow()); -} - -} // namespace tensorrt_llm::batch_manager diff --git a/cpp/tests/e2e_tests/executor/CMakeLists.txt b/cpp/tests/e2e_tests/executor/CMakeLists.txt deleted file mode 100644 index 4813c92584fc..000000000000 --- a/cpp/tests/e2e_tests/executor/CMakeLists.txt +++ /dev/null @@ -1,22 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2023-2025 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. - -add_gtest(executorMockTest executorMockTest.cpp) -add_gtest(executorTest executorTest.cpp) -target_link_libraries(executorTest PRIVATE testingUtils) -add_gtest(encDecTest encDecTest.cpp) -target_link_libraries(encDecTest PRIVATE testingUtils) -add_gtest(disaggExecutorTest disaggExecutorTest.cpp) -target_link_libraries(disaggExecutorTest PRIVATE testingUtils) diff --git a/cpp/tests/e2e_tests/executor/disaggExecutor.h b/cpp/tests/e2e_tests/executor/disaggExecutor.h deleted file mode 100644 index 6b3a529ca16e..000000000000 --- a/cpp/tests/e2e_tests/executor/disaggExecutor.h +++ /dev/null @@ -1,840 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2023-2025 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. - */ - -#pragma once - -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/logger.h" -#include "tensorrt_llm/common/stringUtils.h" -#include "tensorrt_llm/common/utils.h" -#include "tensorrt_llm/executor/dataTransceiverState.h" -#include "tensorrt_llm/executor/disaggServerUtil.h" -#include "tensorrt_llm/executor/executor.h" -#include "tensorrt_llm/executor/requestWithId.h" -#include "tensorrt_llm/executor/serializeUtils.h" -#include "tensorrt_llm/executor/types.h" -#include "tensorrt_llm/runtime/utils/mpiUtils.h" - -#include -#include -#include - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -using namespace tensorrt_llm::executor; -using namespace tensorrt_llm::executor::disagg_executor; -namespace su = tensorrt_llm::executor::serialize_utils; - -namespace tensorrt_llm::testing::disaggexecutor -{ - -constexpr int32_t kM_INSTANCE_ID_TAG{12024}; -constexpr int32_t kM_CONTROLLER_ID_TAG{22024}; -constexpr int32_t kM_INSTANCE_DATA_TAG{32024}; -constexpr int32_t kM_CONTROLLER_DATA_TAG{42024}; - -enum class MessageID : uint64_t -{ - PENDING_CONTEXT_REQUEST = 1, - PENDING_GENERATION_REQUEST = 2, - PENDING_FULL_REQUEST = 3, - CONTEXT_RESPONSE = 4, - GENERATION_RESPONSE = 5, - - TERMINATION = 6, -}; - -enum DisaggRole : uint32_t -{ - DISAGG_CONTEXT = 1, - DISAGG_GENERATION = 2, - DISAGG_MIXED = DISAGG_CONTEXT | DISAGG_GENERATION, - DISAGG_LEADER = 4, - DISAGG_CONTROLLER = 8, -}; - -struct RequestsData -{ - std::vector requests; -}; - -static std::vector serializeResponseWithIds(std::vector const& responseWithIds) -{ - size_t totalSize = 0; - totalSize += sizeof(size_t); - for (auto const& responseWithId : responseWithIds) - { - totalSize += su::serializedSize(responseWithId.gid); - totalSize += su::serializedSize(responseWithId.response); - } - - std::vector buffer(totalSize); - std::stringbuf strbuf{std::ios_base::out | std::ios_base::in}; - strbuf.pubsetbuf(buffer.data(), static_cast(buffer.size())); - std::ostream ostream{&strbuf}; - - su::serialize(responseWithIds.size(), ostream); - for (auto const& responseWithId : responseWithIds) - { - su::serialize(responseWithId.gid, ostream); - su::serialize(responseWithId.response, ostream); - } - return buffer; -} - -static std::vector deserializeResponseWithIds(std::vector& buffer) -{ - std::vector responseWithIds; - su::VectorWrapBuf strbuf{buffer}; - std::istream istream{&strbuf}; - auto numReq = su::deserialize(istream); - for (int64_t req = 0; req < numReq; ++req) - { - auto const id = su::deserialize(istream); - responseWithIds.emplace_back(ResponseWithId{Serialization::deserializeResponse(istream), id}); - } - return responseWithIds; -} - -struct ResponsesData -{ - std::vector response; -}; - -using MessageData = std::variant; - -struct Message -{ - MessageID id; - MessageData data; -}; - -class MessageQueue -{ -public: - void push(Message&& message) - { - std::lock_guard lock(mMutex); - mQueue.push(std::move(message)); - mCv.notify_one(); - } - - Message pop() - { - std::unique_lock lock(mMutex); - mCv.wait(lock, [this] { return !mQueue.empty(); }); - Message message = std::move(mQueue.front()); - mQueue.pop(); - return message; - } - -private: - std::queue mQueue; - std::mutex mMutex; - std::condition_variable mCv; -}; - -class DisaggExecutorLeader -{ -public: - DisaggExecutorLeader(std::filesystem::path const& modelPath, ModelType modelType, - ExecutorConfig const& executorConfig, bool isController, bool isContext, bool isGeneration, int numRequests, - std::vector& participatIds, std::vector const& participantDeviceIdsThisInstance, int worldRank) - : mNumRequests(numRequests) - , mWorldRanksInstances(participatIds) - , mDeviceIdsThisInstance(participantDeviceIdsThisInstance) - , mWorldRank(worldRank) - , mShutdown(false) - , mWorldComm(tensorrt_llm::mpi::MpiComm::world()) - - { - -#if ENABLE_MULTI_DEVICE - - auto world_size = mWorldComm.getSize(); - mRolesPerRank.resize(world_size); - - if (isContext) - { - mRole |= DisaggRole::DISAGG_CONTEXT; - } - if (isGeneration) - { - mRole |= DisaggRole::DISAGG_GENERATION; - } - - if (!mWorldRanksInstances.empty() && mWorldRank == mWorldRanksInstances.front()) - { - mRole |= DisaggRole::DISAGG_LEADER; - } - - if (isController) - { - mRole |= DisaggRole::DISAGG_CONTROLLER; - } - - bool needExecutor = (std::find(mWorldRanksInstances.begin(), mWorldRanksInstances.end(), worldRank) - != mWorldRanksInstances.end()); - if (needExecutor) - { - ExecutorConfig executorConfigC = executorConfig; - - auto parallelConfig = executorConfigC.getParallelConfig().value_or(ParallelConfig{}); - std::vector participantIds = mWorldRanksInstances; - - parallelConfig.setParticipantIds(participantIds); - TLLM_CHECK(parallelConfig.getCommunicationMode() == tensorrt_llm::executor::CommunicationMode::kLEADER); - parallelConfig.setCommunicationType(tensorrt_llm::executor::CommunicationType::kMPI); - parallelConfig.setDeviceIds(mDeviceIdsThisInstance); - executorConfigC.setParallelConfig(parallelConfig); - - mExecutor = std::make_unique(modelPath, modelType, executorConfigC); - } - - TLLM_CHECK(mWorldRanksInstances.size() == mDeviceIdsThisInstance.size()); - - mWorldComm.allgather(&mRole, mRolesPerRank.data(), 1, tensorrt_llm::mpi::MpiType::kUINT32); - - generateRoles(); - - if (isController) - { - mControllerSendThread = std::thread(&DisaggExecutorLeader::ControllerSendThread, this); - mControllerRecvThread = std::thread(&DisaggExecutorLeader::ControllerRecvThread, this); - } - if (isLeaderInstance()) - { - mInstanceRecvThread = std::thread(&DisaggExecutorLeader::InstanceLeaderRecvThread, this); - mInstanceSendThread = std::thread(&DisaggExecutorLeader::InstanceLeaderSendThread, this); - mInstanceLoopThread = std::thread(&DisaggExecutorLeader::InstanceLeaderLoopThread, this); - } -#else - TLLM_THROW("DisaggExecutor only support being compiled with ENABLE_MULTI_DEVICE"); - -#endif - } - - bool isControllerRank() const - { - return mRole & DISAGG_CONTROLLER; - } - - bool isContextRank() const - { - return mRole & DISAGG_CONTEXT; - } - - bool isGenerationRank() const - { - return mRole & DISAGG_GENERATION; - } - - bool isLeaderInstance() const - { - return mRole & DISAGG_LEADER; - } - - std::vector enqueueRequests(std::vector const& llmRequests) - - { - if (!isControllerRank()) - { - return {}; - } - - std::vector requestWithIds; - std::vector requestWithIdsFull; // full request, not disaggregated - std::vector reqIds; - for (auto const& req : llmRequests) - { - IdType id = generatedControlId(); - reqIds.push_back(id); - - RequestWithId reqWithId{req, id}; - if (req.getRequestType() == RequestType::REQUEST_TYPE_CONTEXT_ONLY) - { - requestWithIds.push_back(std::move(reqWithId)); - } - else - { - TLLM_CHECK(req.getRequestType() == RequestType::REQUEST_TYPE_CONTEXT_AND_GENERATION); - requestWithIdsFull.push_back(std::move(reqWithId)); - } - - mRequestMap.insert(std::make_pair(id, req)); - } - - if (!requestWithIds.empty()) - { - Message message{MessageID::PENDING_CONTEXT_REQUEST, MessageData{RequestsData{requestWithIds}}}; - mControllerSendQueue.push(std::move(message)); - } - if (!requestWithIdsFull.empty()) - { - Message message{MessageID::PENDING_FULL_REQUEST, MessageData{RequestsData{requestWithIdsFull}}}; - mControllerSendQueue.push(std::move(message)); - } - - return reqIds; - } - - std::vector awaitResponses(std::optional const& timeout) - { - // wait for responseQueue , modify reqid- - std::vector responses; - std::unique_lock lck(mResponsesMtx); - auto pred = [&mShutdown = mShutdown, &resp = this->mResponses]() -> bool { return !resp.empty() || mShutdown; }; - auto storeResponses = [this, &resp = this->mResponses, &responses]() - { - for (auto it = resp.cbegin(); it != resp.cend();) - { - responses.insert(responses.end(), it->second.begin(), it->second.end()); - resp.erase(it++); - } - }; - - if (timeout) - { - if (mResponsesCv.wait_for(lck, timeout.value(), pred)) - { - storeResponses(); - } - } - else - { - mResponsesCv.wait(lck, pred); - storeResponses(); - } - return responses; - } - - std::deque getLatestRequestStats() - { - if (mExecutor && mExecutor->canEnqueueRequests()) - { - return mExecutor->getLatestRequestStats(); - } - return {}; - } - - void shutDown() - { - if (mShutdown) - { - return; - } - - if (isControllerRank()) - { - std::call_once(mHasSendTerminFlag, - [&]() - { - MessageID terminationMessage = MessageID::TERMINATION; - std::vector isSend(mWorldComm.getSize(), false); - for (auto&& leaderRanks : {mContextLeaderRanks, mGenerationLeaderRanks}) - { - for (auto&& leaderRank : leaderRanks) - { - if (isSend[leaderRank]) - { - continue; - } - mWorldComm.sendRawTag(&terminationMessage, 1, tensorrt_llm::mpi::MpiType::kUINT64, - leaderRank, kM_CONTROLLER_ID_TAG); - isSend[leaderRank] = true; - } - } - - mWorldComm.sendRawTag(&terminationMessage, 1, tensorrt_llm::mpi::MpiType::kUINT64, mControllerRank, - kM_INSTANCE_ID_TAG); - }); - // end recv thread; - } - mShutdown = true; - - // end send thread - if (isControllerRank()) - { - mControllerSendQueue.push({MessageID::TERMINATION, {}}); - } - mInstanceSendQueue.push({MessageID::TERMINATION, {}}); - } - - ~DisaggExecutorLeader() - { - - if (isControllerRank()) - { - shutDown(); - } - - if (isLeaderInstance()) - { - if (mInstanceSendThread.joinable()) - { - mInstanceSendThread.join(); - } - if (mInstanceRecvThread.joinable()) - { - mInstanceRecvThread.join(); - } - if (mInstanceLoopThread.joinable()) - { - mInstanceLoopThread.join(); - } - } - - if (isControllerRank()) - { - if (mControllerSendThread.joinable()) - { - mControllerSendThread.join(); - } - if (mControllerRecvThread.joinable()) - { - mControllerRecvThread.join(); - } - } - - if (!isControllerRank()) - { - mExecutor->shutdown(); - } - if (isControllerRank() && isLeaderInstance()) - { - mExecutor->shutdown(); - } - - shutDown(); - } - -private: - tensorrt_llm::mpi::MpiComm const& mWorldComm; - std::unique_ptr mExecutor; - std::thread mInstanceSendThread; - std::thread mInstanceRecvThread; - std::thread mInstanceLoopThread; - std::thread mControllerSendThread; - std::thread mControllerRecvThread; - int mNumRequests; - std::map mRequestMap; - std::map mGenIdToContextPhase; - std::unordered_map mInstanceIdToGlobalId; - std::mutex mIdToGlbalMutex; - - std::vector mWorldRanksInstances; - - int mWorldRank; - int mControllerRank = 0; - uint32_t mRole = 0; - std::vector mRolesPerRank; - std::vector mContextLeaderRanks; - std::vector mGenerationLeaderRanks; - - IdType mLastId = 1; - MessageQueue mControllerSendQueue; - MessageQueue mInstanceSendQueue; - - std::atomic mShutdown; - - // Ready responses - std::unordered_map> mResponses; - mutable std::mutex mResponsesMtx; - std::condition_variable mResponsesCv; - - std::vector mDeviceIdsThisInstance; - std::once_flag mHasSendTerminFlag; - - void appendNewResponses(std::vector& newResponses) - { - { - std::scoped_lock lck(mResponsesMtx); - for (auto& responseWithId : newResponses) - { - // global id to Result - responseWithId.response = Response(responseWithId.gid, responseWithId.response.getResult()); - - mResponses[responseWithId.gid].emplace_back(responseWithId.response); - } - } - mResponsesCv.notify_all(); - } - - void generateRoles() - { - int contextNum = 0; - int genrationNum = 0; - int controllerNum = 0; - for (int rank = 0; rank < mRolesPerRank.size(); rank++) - { - uint32_t role = mRolesPerRank[rank]; - if (role & DISAGG_LEADER) - { - if (role & DISAGG_CONTEXT) - { - contextNum++; - mContextLeaderRanks.push_back(rank); - } - if (role & DISAGG_GENERATION) - { - genrationNum++; - mGenerationLeaderRanks.push_back(rank); - } - } - if (role & DISAGG_CONTROLLER) - { - controllerNum++; - mControllerRank = rank; - } - } - TLLM_CHECK_WITH_INFO(controllerNum == 1, "only one rank is controller but get %d controllerNum", controllerNum); - TLLM_LOG_INFO("leader ctx: %s, gen: %s", common::vec2str(mContextLeaderRanks).c_str(), - common::vec2str(mGenerationLeaderRanks).c_str()); - } - - IdType generatedControlId() - { - return (mLastId++ % UINT64_MAX); - } - - int selectContextLeaderRank() - { - static int leaderRank = 0; - leaderRank = (leaderRank + 1) % mContextLeaderRanks.size(); - return mContextLeaderRanks[leaderRank]; - } - - int selectGenerationLeaderRank() - { - - // TODO: for same reqId , need select specific generationLeader - static int leaderRank = 0; - leaderRank = (leaderRank + 1) % mGenerationLeaderRanks.size(); - return mGenerationLeaderRanks[leaderRank]; - } - - void ControllerSendThread() - { - // send request to context reqid - // and send context pahse to generation - - TLLM_CUDA_CHECK( - cudaSetDevice(mDeviceIdsThisInstance.at(COMM_SESSION.getRank() % (mDeviceIdsThisInstance.size())))); - tensorrt_llm::common::setThreadName("ControllerSendThread"); - - while (!mShutdown) - { - auto message = mControllerSendQueue.pop(); - if (message.id == MessageID::TERMINATION) - { - - TLLM_LOG_DEBUG("controller get termination message in sendQueue"); - break; - } - if (message.id == MessageID::PENDING_CONTEXT_REQUEST) - { - - auto& reqWithIds = std::get(message.data); - auto packed = RequestWithId::serializeReqWithIds(reqWithIds.requests); - int contextRank = selectContextLeaderRank(); - - mWorldComm.sendRawTag( - &message.id, 1, tensorrt_llm::mpi::MpiType::kUINT64, contextRank, kM_CONTROLLER_ID_TAG); - - mWorldComm.sendRawTag(packed.data(), packed.size(), tensorrt_llm::mpi::MpiType::kCHAR, contextRank, - kM_CONTROLLER_DATA_TAG); - } - else if (message.id == MessageID::PENDING_GENERATION_REQUEST - || message.id == MessageID::PENDING_FULL_REQUEST) - { - - auto& reqWithIds = std::get(message.data); - auto packed = RequestWithId::serializeReqWithIds(reqWithIds.requests); - int generationRank = selectGenerationLeaderRank(); - - mWorldComm.sendRawTag( - &message.id, 1, tensorrt_llm::mpi::MpiType::kUINT64, generationRank, kM_CONTROLLER_ID_TAG); - - mWorldComm.sendRawTag(packed.data(), packed.size(), tensorrt_llm::mpi::MpiType::kCHAR, generationRank, - kM_CONTROLLER_DATA_TAG); - } - else - { - TLLM_THROW("rank:%d, size:%d controller send Invalid message id:%ld", mWorldComm.getRank(), - mWorldComm.getSize(), static_cast(message.id)); - } - } - } - - void ControllerRecvThread() - { -#if ENABLE_MULTI_DEVICE - tensorrt_llm::common::setThreadName("ControllerRecvThread"); - - // recv response from context and push to sendQueue - // recv response from generation and push to responseQueue and notify awaitResponse - TLLM_CUDA_CHECK( - cudaSetDevice(mDeviceIdsThisInstance.at(COMM_SESSION.getRank() % (mDeviceIdsThisInstance.size())))); - - while (!mShutdown) - { - - MPI_Message msg = nullptr; - MPI_Status status; - - mWorldComm.mprobeRawTag(MPI_ANY_SOURCE, kM_INSTANCE_ID_TAG, &msg, &status); - - auto sourceRank{status.MPI_SOURCE}; - int32_t count = 0; - MPICHECK(MPI_Get_count(&status, MPI_UINT64_T, &count)); - TLLM_CHECK(count == 1); - - MessageID messageId; - MPICHECK(MPI_Mrecv(&messageId, count, MPI_UINT64_T, &msg, &status)); - - if (messageId == MessageID::TERMINATION) - { - TLLM_LOG_DEBUG("controller received termination message***************\n"); - break; - } - if (messageId == MessageID::CONTEXT_RESPONSE) - { - mWorldComm.mprobeRawTag(sourceRank, kM_INSTANCE_DATA_TAG, &msg, &status); - MPICHECK(MPI_Get_count(&status, MPI_CHAR, &count)); - std::vector buffer(count); - MPICHECK(MPI_Mrecv(buffer.data(), count, MPI_CHAR, &msg, &status)); - auto responseWithIds = deserializeResponseWithIds(buffer); - // enqueueTo sendQueue like enqueuRequest. . modify requestType and set ContextPhaseParams - // and push to sendQueue. - std::vector requestWithIds; - for (auto&& responseWithId : responseWithIds) - { - auto reqId = responseWithId.gid; - auto& request = mRequestMap.at(reqId); - - request.setRequestType(RequestType::REQUEST_TYPE_GENERATION_ONLY); - request.setContextPhaseParams(responseWithId.response.getResult().contextPhaseParams.value()); - requestWithIds.push_back(RequestWithId{request, reqId}); - } - mControllerSendQueue.push({MessageID::PENDING_GENERATION_REQUEST, RequestsData{requestWithIds}}); - } - - else if (messageId == MessageID::GENERATION_RESPONSE) - { - - mWorldComm.mprobeRawTag(sourceRank, kM_INSTANCE_DATA_TAG, &msg, &status); - MPICHECK(MPI_Get_count(&status, MPI_CHAR, &count)); - std::vector buffer(count); - MPICHECK(MPI_Mrecv(buffer.data(), count, MPI_CHAR, &msg, &status)); - - auto responseWithIds = deserializeResponseWithIds(buffer); - appendNewResponses(responseWithIds); - } - else - { - TLLM_THROW("rank:%d, size:%d controller recv Invalid message id:%ld", mWorldComm.getRank(), - mWorldComm.getSize(), static_cast(messageId)); - } - } -#endif - } - - void InstanceLeaderSendThread() - { - tensorrt_llm::common::setThreadName("InstanceLeaderSendThread"); - - TLLM_CUDA_CHECK( - cudaSetDevice(mDeviceIdsThisInstance.at(COMM_SESSION.getRank() % (mDeviceIdsThisInstance.size())))); - - // pop senQueue and send response to controller - - while (!mShutdown) - { - auto message = mInstanceSendQueue.pop(); - if (message.id == MessageID::CONTEXT_RESPONSE || message.id == MessageID::GENERATION_RESPONSE) - { - auto& responseWithIds = std::get(message.data); - auto packed = serializeResponseWithIds(responseWithIds.response); - - mWorldComm.sendRawTag( - &message.id, 1, tensorrt_llm::mpi::MpiType::kUINT64, mControllerRank, kM_INSTANCE_ID_TAG); - mWorldComm.sendRawTag(packed.data(), packed.size(), tensorrt_llm::mpi::MpiType::kCHAR, mControllerRank, - kM_INSTANCE_DATA_TAG); - } - else if (message.id == MessageID::TERMINATION) - { - // break; no send - TLLM_LOG_DEBUG( - "ranK:%d ,size:%d ,isContext:%d... Context or Generation leader get termination message in " - "sendQueue***************\n", - mWorldComm.getRank(), mWorldComm.getSize(), int(isContextRank())); - break; - } - else - { - TLLM_THROW("rank:%d, size:%d InstanceLeaderSendThread send Invalid message id:%ld", - mWorldComm.getRank(), mWorldComm.getSize(), static_cast(message.id)); - } - } - } - - void InstanceLeaderRecvThread() - { - -#if ENABLE_MULTI_DEVICE - tensorrt_llm::common::setThreadName("InstanceLeaderRecvThread"); - - TLLM_CUDA_CHECK( - cudaSetDevice(mDeviceIdsThisInstance.at(COMM_SESSION.getRank() % (mDeviceIdsThisInstance.size())))); - - // recv request from controller and enqueRequest to executor - while (!mShutdown) - { - MPI_Message msg; - MPI_Status status; - auto sourceRank{mControllerRank}; - mWorldComm.mprobeRawTag(sourceRank, kM_CONTROLLER_ID_TAG, &msg, &status); - - int32_t count; - MPICHECK(MPI_Get_count(&status, MPI_UINT64_T, &count)); - TLLM_CHECK(count == 1); - - MessageID messageId; - MPICHECK(MPI_Mrecv(&messageId, count, MPI_UINT64_T, &msg, &status)); - - if (messageId == MessageID::TERMINATION) - { - TLLM_LOG_DEBUG( - "ranK:%d ,size:%d ,isContext:%d ... Context or Generation leader recv termination message in " - "InstanceLeaderRecvThread***************\n", - mWorldComm.getRank(), mWorldComm.getSize(), int(isContextRank())); - shutDown(); - break; - } - if (messageId == MessageID::PENDING_CONTEXT_REQUEST || messageId == MessageID::PENDING_GENERATION_REQUEST - || messageId == MessageID::PENDING_FULL_REQUEST) - { - mWorldComm.mprobeRawTag(sourceRank, kM_CONTROLLER_DATA_TAG, &msg, &status); - MPICHECK(MPI_Get_count(&status, MPI_CHAR, &count)); - std::vector buffer(count); - MPICHECK(MPI_Mrecv(buffer.data(), count, MPI_CHAR, &msg, &status)); - auto requestWithIds = RequestWithId::deserializeReqWithIds(buffer); - for (auto&& requestWithId : requestWithIds) - { - - auto globalReqId = requestWithId.id; - if (isContextRank() && messageId == MessageID::PENDING_CONTEXT_REQUEST) - { - TLLM_CHECK(requestWithId.req.getRequestType() == RequestType::REQUEST_TYPE_CONTEXT_ONLY); - } - else if (isGenerationRank() - && (messageId == MessageID::PENDING_GENERATION_REQUEST - || messageId == MessageID::PENDING_FULL_REQUEST)) - { - if (messageId == MessageID::PENDING_GENERATION_REQUEST) - { - TLLM_CHECK(requestWithId.req.getRequestType() == RequestType::REQUEST_TYPE_GENERATION_ONLY); - } - else // PENDING_FULL_REQUEST - { - TLLM_CHECK( - requestWithId.req.getRequestType() == RequestType::REQUEST_TYPE_CONTEXT_AND_GENERATION); - } - } - else - { - TLLM_THROW("rank:%d, size:%d InstanceLeaderRecvThread recv Invalid message id:%ld", - mWorldComm.getRank(), mWorldComm.getSize(), static_cast(messageId)); - } - auto reqId = mExecutor->enqueueRequest(requestWithId.req); - { - std::scoped_lock lock{mIdToGlbalMutex}; - mInstanceIdToGlobalId[reqId] = globalReqId; - } - } - } - else - { - TLLM_THROW("rank:%d, size:%d InstanceLeaderRecvThread send Invalid message id:%ld", - mWorldComm.getRank(), mWorldComm.getSize(), static_cast(messageId)); - } - } -#endif - } - - void InstanceLeaderLoopThread() - { - - tensorrt_llm::common::setThreadName("InstanceLeaderLoopThread"); - - TLLM_CUDA_CHECK( - cudaSetDevice(mDeviceIdsThisInstance.at(COMM_SESSION.getRank() % (mDeviceIdsThisInstance.size())))); - - // loop awaitResponse and enqueue into sendQueue - while (!mShutdown) - { - std::chrono::milliseconds waitTime(1); - - auto responses = mExecutor->awaitResponses(waitTime); - if (responses.empty()) - { - continue; - } - std::vector responseWithIdsContext; - std::vector responseWithIdsGeneration; - for (auto&& response : responses) - { - auto reqId = response.getRequestId(); - IdType globalId{0}; - { - std::scoped_lock lock{mIdToGlbalMutex}; - globalId = mInstanceIdToGlobalId[reqId]; - } - TLLM_CHECK(globalId != 0); - auto const& result = response.getResult(); - if (result.contextPhaseParams.has_value()) - { - responseWithIdsContext.emplace_back(response, globalId); - } - else - { - responseWithIdsGeneration.emplace_back(response, globalId); - } - } - - if (isContextRank()) - { - mInstanceSendQueue.push({MessageID::CONTEXT_RESPONSE, ResponsesData{responseWithIdsContext}}); - } - if (isGenerationRank()) - { - mInstanceSendQueue.push({MessageID::GENERATION_RESPONSE, ResponsesData{responseWithIdsGeneration}}); - } - } - } -}; -} // namespace tensorrt_llm::testing::disaggexecutor diff --git a/cpp/tests/e2e_tests/executor/disaggExecutorTest.cpp b/cpp/tests/e2e_tests/executor/disaggExecutorTest.cpp deleted file mode 100644 index 0eb05d2cc807..000000000000 --- a/cpp/tests/e2e_tests/executor/disaggExecutorTest.cpp +++ /dev/null @@ -1,1437 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2023-2025 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. - */ - -#include "disaggExecutor.h" -#include "executorTest.h" -#include "tensorrt_llm/common/envUtils.h" -#include "tensorrt_llm/runtime/utils/numpyUtils.h" -#include "tests/utils/common.h" - -#include -#include - -namespace tr = tensorrt_llm::runtime; - -using namespace tensorrt_llm::testing; - -namespace -{ -auto constexpr LLAMA_INPUT_FILE = "input_tokens_llama.npy"; -auto constexpr LLAMA_VOCAB_SIZE_PADDED = 128256; -auto constexpr LLAMA_END_ID = 128001; -auto constexpr LLAMA_PAD_ID = 128001; - -using CondDisaggParamsType = std::tuple; // modelName - -enum class InstanceRole : int -{ - kCONTEXT = 1, - kGENERATION = 0, - kMIXED = 2 -}; - -using DisaggParamsType = std::tuple< // - int, // processNum - std::vector, // modelNames - std::vector>, // participantIdsEachInstance - std::vector>, // participantDeviceIdsEachInstance - std::vector, // instanceRoles - int // controllerRank - >; - -std::string convertToString(std::vector> const& vec) -{ - std::ostringstream oss; - oss << "XX"; - - for (size_t i = 0; i < vec.size(); ++i) - { - for (size_t j = 0; j < vec[i].size(); ++j) - { - oss << vec[i][j]; - if (j < vec[i].size() - 1) - { - oss << "_"; - } - } - if (i < vec.size() - 1) - { - oss << "X_X"; - } - } - - oss << "XX"; - return oss.str(); -}; - -std::string convertToString(std::vector const& vec) -{ - std::ostringstream oss; - oss << "XX"; - - for (size_t j = 0; j < vec.size(); ++j) - { - oss << static_cast(vec[j]); - if (j < vec.size() - 1) - { - oss << "_"; - } - } - - oss << "XX"; - return oss.str(); -}; - -std::string generateTestNameDisaggParams(testing::TestParamInfo const& info) -{ - auto const processNum = std::get<0>(info.param); - auto const modelNames = std::get<1>(info.param); - auto const participantIdsEachInstance = std::get<2>(info.param); // std::vector> - auto const participantDeviceIdsEachInstance = std::get<3>(info.param); // std::vector>; - auto const instanceRoles = std::get<4>(info.param); // std::vector ; //1 is context , 0 is generation - auto const controllerRank = std::get<5>(info.param); - - std::string name = "DisaggExecutorTest_"; - - name.append("ProcessNum_" + std::to_string(processNum)); - // name.append("_contextModel_" + contextModel + "_genModel_" + genModel); - name.append("_modelNames_"); - for (auto&& modelName : modelNames) - { - name.append(modelName).append("_"); - } - - name.append("_controllerRank_" + std::to_string(controllerRank)); - - name.append("_ranks_").append(convertToString(participantIdsEachInstance)); - name.append("_devices_").append(convertToString(participantDeviceIdsEachInstance)); - name.append("_roles_").append(convertToString(instanceRoles)); - name.append("_controllerRank_" + std::to_string(controllerRank)); - - return name; -} - -std::string generateTestNameCondDisaggParams(testing::TestParamInfo const& info) -{ - auto const modelName = std::get<0>(info.param); - return "Model_" + modelName; -} - -class DisaggParamsTest : public GptExecutorTest, public ::testing::WithParamInterface -{ -}; - -class DisaggOrchestratorParamsTest : public GptExecutorTest, public ::testing::WithParamInterface -{ -}; - -class ConditionalDisaggParamsTest : public GptExecutorTest, public ::testing::WithParamInterface -{ -}; - -void verifyGenerateDistStats(std::deque const& iterationStats) -{ - for (auto const& iteration : iterationStats) - { - for (auto const& requestStats : iteration.requestStats) - { - // exclude context only requests for mixed server - if (requestStats.stage == RequestStage::kGENERATION_COMPLETE && requestStats.numGeneratedTokens > 1) - { - EXPECT_TRUE(requestStats.disServingStats.has_value()); - EXPECT_GT(requestStats.disServingStats.value().kvCacheTransferMS, 0.0); - } - if (requestStats.stage != RequestStage::kQUEUED) - { - EXPECT_TRUE(requestStats.disServingStats.has_value()); - } - else - { - EXPECT_FALSE(requestStats.disServingStats.has_value()); - } - } - } -} -} // namespace - -void runDisaggTest(tensorrt_llm::testing::disaggexecutor::DisaggExecutorLeader& executor, - tensorrt_llm::runtime::BufferManager& manager, ITensor const& givenInput, ModelIds const& modelIds, - FlakyTestInfo const& flakyTestInfo, bool streaming, SizeType32 const vocabSizePadded, BeamResult const& beamResult, - OutputConfig const& outConfig, bool isSpeculativeDecoding, int maxWaitMs, BatchingType batchingType, - bool returnAllGeneratedTokens) -{ - - auto& comm = tensorrt_llm::mpi::MpiComm::world(); - auto const worldRank = comm.getRank(); - auto const worldSize = comm.getSize(); - auto const beamWidth = beamResult.beamWidth; - - std::unordered_map reqIdToBatchId; - std::unordered_map> tokens; - auto [givenInputLengths, nbGivenInputs, maxInputLength] = getGivenInputLengths(givenInput, modelIds.padId); - auto const* const givenInputData = tr::bufferCast(givenInput); - - auto const& inputShape = givenInput.getShape(); - ASSERT_EQ(inputShape.nbDims, 2); - ASSERT_GT(inputShape.d[0], 0); - - // Load expected outputs for each beam width value - auto testData = TestData::loadTestData(beamResult, givenInput, beamWidth, manager, outConfig, modelIds); - auto const maxSeqLen = testData.maxSeqLen; - - // Load expected outputs and inputs - SizeType32 numRequests = static_cast(givenInputLengths.size()); - SizeType32 maxRequests = numRequests; - std::vector requests; - std::vector reqMaxNewTokens; - SizeType32 const numReturnSequences = 1; - - for (SizeType32 req = 0; req < maxRequests; ++req) - { - SizeType32 inputLen = givenInputLengths.at(req); - auto maxNewTokens = maxSeqLen - maxInputLength; - reqMaxNewTokens.push_back(maxNewTokens); - SizeType32 endId = -1; - auto const* const seqBegin = givenInputData + req * maxInputLength; - VecTokens tokens(seqBegin, seqBegin + inputLen); - auto samplingConfig = tensorrt_llm::executor::SamplingConfig(beamWidth); - samplingConfig.setNumReturnSequences(numReturnSequences); - auto request = Request( - VecTokens(seqBegin, seqBegin + inputLen), maxNewTokens, streaming, samplingConfig, outConfig, endId); - request.setReturnAllGeneratedTokens(returnAllGeneratedTokens); - request.setRequestType(RequestType::REQUEST_TYPE_CONTEXT_ONLY); - requests.emplace_back(std::move(request)); - } - - if (executor.isControllerRank()) - { - std::vector reqIds; - - for (int i = 0; i < requests.size(); ++i) - { - std::vector resultTokens; - resultTokens.reserve(numReturnSequences); - for (SizeType32 seqIdx = 0; seqIdx < numReturnSequences; ++seqIdx) - { - resultTokens.emplace_back(beamWidth); - } - auto retReqId = executor.enqueueRequests({requests[i]}); - reqIds.push_back(retReqId.front()); - tokens[i] = std::move(resultTokens); - reqIdToBatchId[retReqId.front()] = i; - } - - // Get the new tokens for each requests - int32_t numFinished = 0; - int iter = 0; - SizeType32 numResponses = 0; - while (numFinished < maxRequests && iter < maxWaitMs) - { - std::chrono::milliseconds waitTime(1); - auto responses = executor.awaitResponses(waitTime); - for (auto& response : responses) - { - numResponses++; - if (!response.hasError()) - { - auto result = response.getResult(); - numFinished += result.isFinal; - auto batchId = reqIdToBatchId.at(response.getRequestId()); - auto seqIdx = result.sequenceIndex; - - auto& contextLogits = result.contextLogits; - auto& genLogits = result.generationLogits; - auto& outputTokenIds = result.outputTokenIds; - - EXPECT_EQ(result.finishReasons.size(), beamWidth); - for (SizeType32 beam = 0; beam < beamWidth; ++beam) - { - auto& newTokens = outputTokenIds.at(beam); - auto& reqTokens = tokens.at(batchId).at(seqIdx).at(beam); - - reqTokens.insert(reqTokens.end(), newTokens.begin(), newTokens.end()); - // FinishReason is only supported for bw=1 and inflight batching. - if (beamWidth == 1 && batchingType == BatchingType::kINFLIGHT) - { - EXPECT_EQ(result.finishReasons.at(beam), - result.isFinal ? FinishReason::kLENGTH : FinishReason::kNOT_FINISHED); - } - } - - auto& cumLogProbs = result.cumLogProbs; - auto& logProbs = result.logProbs; - auto& beamTokens = tokens.at(batchId).at(seqIdx); - testData.verifyLogProbs(outConfig.returnLogProbs, streaming, outConfig.excludeInputFromOutput, - givenInputLengths.at(batchId), beamWidth, beamTokens, cumLogProbs, logProbs, batchId, - flakyTestInfo); - - testData.validateContextLogits(outConfig.returnContextLogits, givenInputLengths.at(batchId), - beamWidth, contextLogits, vocabSizePadded, batchId); - testData.validateGenerationLogits(outConfig.returnGenerationLogits, result.isFinal, streaming, - outConfig.excludeInputFromOutput, givenInputLengths.at(batchId), reqMaxNewTokens.at(batchId), - beamWidth, beamTokens, genLogits, vocabSizePadded, batchId, returnAllGeneratedTokens); - } - else - { - // Allow response with error only if awaitResponse processed a terminated request id - std::string err = "ReqId " + std::to_string(response.getRequestId()) - + " has already been processed and was terminated."; - EXPECT_EQ(response.getErrorMsg(), err); - } - } - ++iter; - } - EXPECT_LT(iter, maxWaitMs); - testData.verifyOutput(tokens, givenInputLengths, streaming, outConfig.excludeInputFromOutput, flakyTestInfo, - isSpeculativeDecoding, beamWidth, numReturnSequences, false); - } - comm.barrier(); - if (executor.isGenerationRank()) - { - verifyGenerateDistStats(executor.getLatestRequestStats()); - } -} - -void runDisaggTest(DisaggExecutorOrchestrator& executor, tensorrt_llm::runtime::BufferManager& manager, - ITensor const& givenInput, ModelIds const& modelIds, FlakyTestInfo const& flakyTestInfo, bool streaming, - SizeType32 const vocabSizePadded, BeamResult const& beamResult, OutputConfig const& outConfig, - bool isSpeculativeDecoding, int maxWaitMs, BatchingType batchingType, bool returnAllGeneratedTokens) -{ - - auto& comm = tensorrt_llm::mpi::MpiComm::world(); - auto const worldRank = comm.getRank(); - auto const worldSize = comm.getSize(); - auto const beamWidth = beamResult.beamWidth; - - std::unordered_map reqIdToBatchId; - std::unordered_map> tokens; - // std::unordered_map gGenIdIdTogContextId; - auto [givenInputLengths, nbGivenInputs, maxInputLength] = getGivenInputLengths(givenInput, modelIds.padId); - auto const* const givenInputData = tr::bufferCast(givenInput); - - auto const& inputShape = givenInput.getShape(); - ASSERT_EQ(inputShape.nbDims, 2); - ASSERT_GT(inputShape.d[0], 0); - - // Load expected outputs for each beam width value - auto testData = TestData::loadTestData(beamResult, givenInput, beamWidth, manager, outConfig, modelIds); - auto const maxSeqLen = testData.maxSeqLen; - - // Load expected outputs and inputs - SizeType32 numRequests = static_cast(givenInputLengths.size()); - SizeType32 maxRequests = numRequests; - std::vector requests; - std::vector reqMaxNewTokens; - SizeType32 const numReturnSequences = 1; - - for (SizeType32 req = 0; req < maxRequests; ++req) - { - SizeType32 inputLen = givenInputLengths.at(req); - auto maxNewTokens = maxSeqLen - maxInputLength; - reqMaxNewTokens.push_back(maxNewTokens); - SizeType32 endId = -1; - auto const* const seqBegin = givenInputData + req * maxInputLength; - VecTokens tokens(seqBegin, seqBegin + inputLen); - auto samplingConfig = tensorrt_llm::executor::SamplingConfig(beamWidth); - samplingConfig.setNumReturnSequences(numReturnSequences); - auto request = Request( - VecTokens(seqBegin, seqBegin + inputLen), maxNewTokens, streaming, samplingConfig, outConfig, endId); - request.setReturnAllGeneratedTokens(returnAllGeneratedTokens); - request.setRequestType(RequestType::REQUEST_TYPE_CONTEXT_ONLY); - requests.emplace_back(std::move(request)); - } - - if (worldRank == 0) - { - std::vector reqIds; - - for (int i = 0; i < requests.size(); ++i) - { - std::vector resultTokens; - resultTokens.reserve(numReturnSequences); - for (SizeType32 seqIdx = 0; seqIdx < numReturnSequences; ++seqIdx) - { - resultTokens.emplace_back(beamWidth); - } - auto retReqId = executor.enqueueContext({requests[i]}, std::nullopt); - reqIds.push_back(retReqId.front()); - tokens[i] = std::move(resultTokens); - reqIdToBatchId[retReqId.front()] = i; - } - - int32_t numContextFinished = 0; - int contextIter = 0; - while (numContextFinished < maxRequests && contextIter < maxWaitMs) - { - std::chrono::milliseconds waitTime(1); - - auto contextResponses = executor.awaitContextResponses(waitTime); - contextIter++; - numContextFinished += contextResponses.size(); - - for (auto&& responseWithId : contextResponses) - { - auto contextGid = responseWithId.gid; - int batchId = reqIdToBatchId[contextGid]; - auto&& request = requests[batchId]; - request.setRequestType(RequestType::REQUEST_TYPE_GENERATION_ONLY); - request.setContextPhaseParams(responseWithId.response.getResult().contextPhaseParams.value()); - executor.enqueueGeneration({request}, {responseWithId.gid}, std::nullopt); - } - } - // Get the new tokens for each requests - int32_t numFinished = 0; - int iter = 0; - SizeType32 numResponses = 0; - while (numFinished < maxRequests && iter < maxWaitMs) - { - std::chrono::milliseconds waitTime(1); - auto responses = executor.awaitGenerationResponses(waitTime); - for (auto& responseWithId : responses) - { - numResponses++; - if (!responseWithId.response.hasError()) - { - auto result = responseWithId.response.getResult(); - numFinished += result.isFinal; - auto batchId = reqIdToBatchId.at(responseWithId.gid); - auto seqIdx = result.sequenceIndex; - - auto& contextLogits = result.contextLogits; - auto& genLogits = result.generationLogits; - auto& outputTokenIds = result.outputTokenIds; - - EXPECT_EQ(result.finishReasons.size(), beamWidth); - for (SizeType32 beam = 0; beam < beamWidth; ++beam) - { - auto& newTokens = outputTokenIds.at(beam); - auto& reqTokens = tokens.at(batchId).at(seqIdx).at(beam); - - reqTokens.insert(reqTokens.end(), newTokens.begin(), newTokens.end()); - // FinishReason is only supported for bw=1 and inflight batching. - if (beamWidth == 1 && batchingType == BatchingType::kINFLIGHT) - { - EXPECT_EQ(result.finishReasons.at(beam), - result.isFinal ? FinishReason::kLENGTH : FinishReason::kNOT_FINISHED); - } - } - - auto& cumLogProbs = result.cumLogProbs; - auto& logProbs = result.logProbs; - auto& beamTokens = tokens.at(batchId).at(seqIdx); - testData.verifyLogProbs(outConfig.returnLogProbs, streaming, outConfig.excludeInputFromOutput, - givenInputLengths.at(batchId), beamWidth, beamTokens, cumLogProbs, logProbs, batchId, - flakyTestInfo); - - testData.validateContextLogits(outConfig.returnContextLogits, givenInputLengths.at(batchId), - beamWidth, contextLogits, vocabSizePadded, batchId); - testData.validateGenerationLogits(outConfig.returnGenerationLogits, result.isFinal, streaming, - outConfig.excludeInputFromOutput, givenInputLengths.at(batchId), reqMaxNewTokens.at(batchId), - beamWidth, beamTokens, genLogits, vocabSizePadded, batchId, returnAllGeneratedTokens); - } - else - { - // Allow response with error only if awaitResponse processed a terminated request id - std::string err = "ReqId " + std::to_string(responseWithId.gid) - + " has already been processed and was terminated."; - EXPECT_EQ(responseWithId.response.getErrorMsg(), err); - } - } - ++iter; - } - EXPECT_LT(iter, maxWaitMs); - testData.verifyOutput(tokens, givenInputLengths, streaming, outConfig.excludeInputFromOutput, flakyTestInfo, - isSpeculativeDecoding, beamWidth, numReturnSequences, false); - } - comm.barrier(); -} - -TEST_P(DisaggParamsTest, DisaggTokenComparison) -{ - -#if ENABLE_MULTI_DEVICE - - if (!(tensorrt_llm::common::getEnvUseUCXKvCache())) - { - setenv("UCX_TLS", "^cuda_ipc", 1); // disable cuda_ipc for testing for mpi - } - else - { - setenv("UCX_TCP_CM_REUSEADDR", "y", - 1); // tests creates and destroies ucxCacheCommunicatoers frequently, so listener ports must be reused - } - auto const processNum = std::get<0>(GetParam()); - auto const modelNames = std::get<1>(GetParam()); - auto const participantIdsEachInstance = std::get<2>(GetParam()); // std::vector> - auto const participantDeviceIdsEachInstance = std::get<3>(GetParam()); // std::vector>; - auto const instanceRoles - = std::get<4>(GetParam()); // std::vector ; //1 is context , 0 is generation, 2 is mixed - auto const controllerRank = std::get<5>(GetParam()); - - // params_check - auto const& world_comm = tensorrt_llm::mpi::MpiComm::world(); - int const commRank = world_comm.getRank(); - int const commSize = world_comm.getSize(); - if (commSize != processNum) - { - GTEST_SKIP() << " need " << processNum << " processes but got " << commSize << " mpi processes, skip test."; - } - ASSERT_EQ(participantIdsEachInstance.size(), participantDeviceIdsEachInstance.size()); - SizeType32 instanceNum = participantIdsEachInstance.size(); - ASSERT_EQ(instanceNum, instanceRoles.size()); - ASSERT_EQ(instanceNum, modelNames.size()); - - std::unordered_set deviceIdsSet; - for (auto const& ids : participantDeviceIdsEachInstance) - { - for (auto const& id : ids) - { - deviceIdsSet.insert(id); - } - } - if (mDeviceCount < deviceIdsSet.size()) - { - GTEST_SKIP() << " need " << deviceIdsSet.size() << " devices but got " << mDeviceCount - << " devices, skip test."; - } - - ASSERT_GE(controllerRank, 0); - ASSERT_LT(controllerRank, commSize); - int ranksNum = 0; - std::unordered_map rankCounter; - std::unordered_map deviceCounter; - SizeType32 deviceRuseNum = 1; - bool isContext = false; - bool isGeneration = false; - std::vector participatntIds; - std::vector deviceIds; - std::string modelName; - bool isController = (commRank == controllerRank); - for (SizeType32 i = 0; i < instanceNum; i++) - { - auto const& ranksThisInstance = participantIdsEachInstance[i]; - auto const& devicesThisInstance = participantDeviceIdsEachInstance[i]; - - ASSERT_EQ(ranksThisInstance.size(), devicesThisInstance.size()); - SizeType32 rankNumThisInstance = ranksThisInstance.size(); - ASSERT_GT(rankNumThisInstance, 0); - ranksNum += rankNumThisInstance; - for (SizeType32 j = 0; j < rankNumThisInstance; j++) - { - rankCounter[ranksThisInstance[j]]++; - deviceCounter[devicesThisInstance[j]]++; - ASSERT_GE(rankCounter[ranksThisInstance[j]], 1); - deviceRuseNum = std::max(deviceCounter[devicesThisInstance[j]], deviceRuseNum); - ASSERT_GE(ranksThisInstance[j], 0); - ASSERT_LT(ranksThisInstance[j], commSize); - - if (commRank == ranksThisInstance[j]) - { - participatntIds = ranksThisInstance; - deviceIds = devicesThisInstance; - isContext = instanceRoles[i] == InstanceRole::kCONTEXT || instanceRoles[i] == InstanceRole::kMIXED; - isGeneration - = instanceRoles[i] == InstanceRole::kGENERATION || instanceRoles[i] == InstanceRole::kMIXED; - // modelName = isContext ? contextModel : genModel; - modelName = modelNames[i]; - } - } - } - ASSERT_GE(ranksNum, commSize); - - OutputConfig outConfig; - int const beamWidth = 1; - BeamResult beamResult{beamWidth}; - - bool streaming = false; - int const maxBeamWidth = 1; - ASSERT_TRUE(fs::exists(DATA_PATH)); - - fs::path modelPath; - // set defaults and adjust if needed by different models - fs::path inputPath = DATA_PATH / "input_tokens.npy"; - ModelIds modelIds{50256, 50256}; - SizeType32 vocabSizePadded{50257}; // gpt vocabSizePadded - bool isSpeculativeDecoding{false}; - - // NOTE: This can be used to disable checks for certain prompt batch entries - FlakyTestInfo flakyTestInfo; - - if (modelName == "gpt") - { - auto const resultsPath - = GPT_DATA_PATH / ((beamWidth == 1) ? "sampling" : "beam_search_" + std::to_string(beamWidth)); - if (outConfig.returnContextLogits || outConfig.returnGenerationLogits) - { - modelPath = GPT_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_GATHER_DIR() / "tp1-pp1-cp1-gpu"; - beamResult.resultsFile = resultsPath / PathUtil::FP16_PLUGIN_PACKED_PAGED_GATHER_RESULT_FILE(); - beamResult.contextLogitsFile = resultsPath / PathUtil::FP16_PLUGIN_PACKED_PAGED_CONTEXT_LOGITS_FILE(); - beamResult.genLogitsFile = resultsPath / PathUtil::FP16_PLUGIN_PACKED_PAGED_GENERATION_LOGITS_FILE(); - if (outConfig.returnLogProbs) - { - beamResult.cumLogProbsFile - = resultsPath / PathUtil::FP16_PLUGIN_PACKED_PAGED_GATHER_CUM_LOG_PROBS_FILE(); - beamResult.logProbsFile = resultsPath / PathUtil::FP16_PLUGIN_PACKED_PAGED_GATHER_LOG_PROBS_FILE(); - } - } - else - { - modelPath = GPT_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp1-pp1-cp1-gpu"; - beamResult.resultsFile = resultsPath / PathUtil::FP16_PLUGIN_PACKED_PAGED_RESULT_FILE(); - if (outConfig.returnLogProbs) - { - beamResult.cumLogProbsFile = resultsPath / PathUtil::FP16_PLUGIN_PACKED_PAGED_CUM_LOG_PROBS_FILE(); - beamResult.logProbsFile = resultsPath / PathUtil::FP16_PLUGIN_PACKED_PAGED_LOG_PROBS_FILE(); - } - } - } - else if (modelName == "llama_tp4_pp1_cp1" || modelName == "llama_tp1_pp4_cp1" || modelName == "llama_tp2_pp2_cp1" - || modelName == "llama_tp1_pp2_cp1" || modelName == "llama_tp2_pp1_cp1" || modelName == "llama_tp1_pp1_cp1") - { - inputPath = DATA_PATH / LLAMA_INPUT_FILE; - vocabSizePadded = LLAMA_VOCAB_SIZE_PADDED; - - auto const resultsPath - = LLAMA_DATA_PATH / ((beamWidth == 1) ? "sampling" : "beam_search_" + std::to_string(beamWidth)); - modelIds.padId = LLAMA_PAD_ID; - modelIds.endId = LLAMA_END_ID; - if (modelName == "llama_tp4_pp1_cp1") - { - beamResult.resultsFile = resultsPath / PathUtil::FP16_PLUGIN_PACKED_PAGED_RESULT_TP4_PP1_FILE(); - modelPath = LLAMA_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp4-pp1-cp1-gpu"; - } - else if (modelName == "llama_tp1_pp4_cp1") - { - beamResult.resultsFile = resultsPath / PathUtil::FP16_PLUGIN_PACKED_PAGED_RESULT_TP1_PP4_FILE(); - modelPath = LLAMA_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp1-pp4-cp1-gpu"; - } - else if (modelName == "llama_tp1_pp2_cp1") - { - beamResult.resultsFile = resultsPath / PathUtil::FP16_PLUGIN_PACKED_PAGED_RESULT_TP1_PP2_FILE(); - modelPath = LLAMA_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp1-pp2-cp1-gpu"; - } - else if (modelName == "llama_tp2_pp1_cp1") - { - beamResult.resultsFile = resultsPath / PathUtil::FP16_PLUGIN_PACKED_PAGED_RESULT_TP2_PP1_FILE(); - modelPath = LLAMA_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp2-pp1-cp1-gpu"; - } - else if (modelName == "llama_tp2_pp2_cp1") - { - beamResult.resultsFile = resultsPath / PathUtil::FP16_PLUGIN_PACKED_PAGED_RESULT_TP2_PP2_FILE(); - modelPath = LLAMA_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp2-pp2-cp1-gpu"; - } - else if (modelName == "llama_tp1_pp1_cp1") - { - beamResult.resultsFile = resultsPath / PathUtil::FP16_PLUGIN_PACKED_PAGED_RESULT_TP2_PP2_FILE(); - modelPath = LLAMA_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp1-pp1-cp1-gpu"; - } - } - else - { - TLLM_THROW("Unrecognized modelName"); - } - - // Warning: This should be the last check before running the test. - // It will initialize MPI which can take significant time. - if (modelName == "llama_tp4_pp1_cp1" || modelName == "llama_tp1_pp4_cp1" || modelName == "llama_tp2_pp2_cp1" - || modelName == "llama_tp1_pp2_cp1" || modelName == "llama_tp2_pp1_cp1") - { - if (outConfig.returnLogProbs || outConfig.returnContextLogits || outConfig.returnGenerationLogits) - { - GTEST_SKIP() << "Skipping logits and log probs tests for mpi runs"; - } - } - - // Returning logits will bring higher latency - if (streaming && (outConfig.returnContextLogits || outConfig.returnGenerationLogits)) - { - mMaxWaitMs = 20000; - } - - auto executorConfig = ExecutorConfig(maxBeamWidth); - FloatType freeGpuMemoryFraction = 0.9f / (deviceRuseNum); // context and gen instance run on same device - KvCacheConfig kvCacheConfig{true, std::nullopt, std::nullopt, std::nullopt, freeGpuMemoryFraction}; - executorConfig.setKvCacheConfig(kvCacheConfig); - executorConfig.setRequestStatsMaxIterations(1000); - executorConfig.setCacheTransceiverConfig( - texec::CacheTransceiverConfig(texec::CacheTransceiverConfig::BackendType::DEFAULT)); - auto manager = tr::BufferManager(std::make_shared()); - auto const& givenInput = tr::utils::loadNpy(manager, inputPath.string(), tr::MemoryType::kCPU); - auto [givenInputLengths, nbGivenInputs, maxInputLength] = getGivenInputLengths(*givenInput, modelIds.padId); - world_comm.barrier(); - auto disaggExecutor = tensorrt_llm::testing::disaggexecutor::DisaggExecutorLeader(modelPath, - ModelType::kDECODER_ONLY, executorConfig, isController, isContext, isGeneration, givenInputLengths.size(), - participatntIds, deviceIds, commRank); - - runDisaggTest(disaggExecutor, manager, *givenInput, modelIds, flakyTestInfo, streaming, vocabSizePadded, beamResult, - outConfig, isSpeculativeDecoding, mMaxWaitMs, executorConfig.getBatchingType(), false); - -#else - - GTEST_SKIP() << "Skipping DisaggExecutor Test"; - -#endif -} - -TEST_P(DisaggOrchestratorParamsTest, DisaggTokenComparison) -{ - -#if ENABLE_MULTI_DEVICE - - if (!(tensorrt_llm::common::getEnvUseUCXKvCache())) - { - setenv("UCX_TLS", "^cuda_ipc", 1); // disable cuda_ipc for testing for mpi - } - else - { - setenv("UCX_TCP_CM_REUSEADDR", "y", - 1); // tests creates and destroies ucxCacheCommunicatoers frequently, so listener ports must be reused - } - auto const processNum = std::get<0>(GetParam()); - auto const modelNames = std::get<1>(GetParam()); - auto const participantIdsEachInstance = std::get<2>(GetParam()); // std::vector> - auto const participantDeviceIdsEachInstance = std::get<3>(GetParam()); // std::vector>; - auto const instanceRoles = std::get<4>(GetParam()); // std::vector ; //1 is context , 0 is generation - auto const controllerRank = std::get<5>(GetParam()); - - // params_check - auto const& world_comm = tensorrt_llm::mpi::MpiComm::world(); - int const commRank = world_comm.getRank(); - int const commSize = world_comm.getSize(); - if (commSize != processNum) - { - GTEST_SKIP() << " need " << processNum << " processes but got " << commSize << " mpi processes, skip test."; - } - - bool spawnProcess = false; - if (commSize == 1) - { - spawnProcess = true; - if (mDeviceCount < 4) - { - GTEST_SKIP() << "DisaggExecutorTest requires at least 4 GPUs"; - } - ASSERT_TRUE(tensorrt_llm::common::getEnvUseUCXKvCache() || tensorrt_llm::common::getEnvUseNixlKvCache()); - } - - ASSERT_EQ(participantIdsEachInstance.size(), participantDeviceIdsEachInstance.size()); - SizeType32 instanceNum = participantIdsEachInstance.size(); - ASSERT_EQ(instanceNum, instanceRoles.size()); - ASSERT_EQ(instanceNum, modelNames.size()); - - std::unordered_set deviceIdsSet; - for (auto const& ids : participantDeviceIdsEachInstance) - { - for (auto const& id : ids) - { - deviceIdsSet.insert(id); - } - } - if (mDeviceCount < deviceIdsSet.size()) - { - GTEST_SKIP() << " need " << deviceIdsSet.size() << " devices but got " << mDeviceCount - << " devices, skip test."; - } - - ASSERT_GE(controllerRank, 0); - ASSERT_LT(controllerRank, commSize); - std::string modelName = modelNames[0]; - bool isController = (commRank == controllerRank); - std::vector contextModels; - std::vector genModels; - - auto getModelPath = [=](std::string modelNN) - { - fs::path retPath; - if (modelNN == "llama_tp4_pp1") - { - retPath = LLAMA_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp4-pp1-cp1-gpu"; - } - else if (modelNN == "llama_tp1_pp4") - { - retPath = LLAMA_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp1-pp4-cp1-gpu"; - } - else if (modelNN == "llama_tp1_pp2") - { - retPath = LLAMA_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp1-pp2-cp1-gpu"; - } - else if (modelNN == "llama_tp2_pp1") - { - retPath = LLAMA_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp2-pp1-cp1-gpu"; - } - else if (modelNN == "llama_tp2_pp2") - { - retPath = LLAMA_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp2-pp2-cp1-gpu"; - } - else if (modelNN == "llama_tp1_pp1") - { - retPath = LLAMA_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp1-pp1-cp1-gpu"; - } - return retPath; - }; - for (SizeType32 i = 0; i < instanceNum; i++) - { - if (instanceRoles[i] == InstanceRole::kCONTEXT) - { - contextModels.push_back(getModelPath(modelNames[i])); - } - else - { - genModels.push_back(getModelPath(modelNames[i])); - } - } - - OutputConfig outConfig; - int const beamWidth = 1; - BeamResult beamResult{beamWidth}; - - bool streaming = false; - int const maxBeamWidth = 1; - ASSERT_TRUE(fs::exists(DATA_PATH)); - - fs::path modelPath; - // set defaults and adjust if needed by different models - fs::path inputPath = DATA_PATH / "input_tokens.npy"; - ModelIds modelIds{50256, 50256}; - SizeType32 vocabSizePadded{50257}; // gpt vocabSizePadded - bool isSpeculativeDecoding{false}; - - // NOTE: This can be used to disable checks for certain prompt batch entries - FlakyTestInfo flakyTestInfo; - if (modelName == "llama_tp4_pp1" || modelName == "llama_tp1_pp4" || modelName == "llama_tp2_pp2" - || modelName == "llama_tp1_pp2" || modelName == "llama_tp2_pp1" || modelName == "llama_tp1_pp1") - { - inputPath = DATA_PATH / LLAMA_INPUT_FILE; - vocabSizePadded = LLAMA_VOCAB_SIZE_PADDED; - - auto const resultsPath - = LLAMA_DATA_PATH / ((beamWidth == 1) ? "sampling" : "beam_search_" + std::to_string(beamWidth)); - modelIds.padId = LLAMA_PAD_ID; - modelIds.endId = LLAMA_END_ID; - if (modelName == "llama_tp4_pp1") - { - beamResult.resultsFile = resultsPath / PathUtil::FP16_PLUGIN_PACKED_PAGED_RESULT_TP4_PP1_FILE(); - modelPath = LLAMA_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp4-pp1-cp1-gpu"; - } - else if (modelName == "llama_tp1_pp4") - { - beamResult.resultsFile = resultsPath / PathUtil::FP16_PLUGIN_PACKED_PAGED_RESULT_TP1_PP4_FILE(); - modelPath = LLAMA_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp1-pp4-cp1-gpu"; - } - else if (modelName == "llama_tp1_pp2") - { - beamResult.resultsFile = resultsPath / PathUtil::FP16_PLUGIN_PACKED_PAGED_RESULT_TP1_PP2_FILE(); - modelPath = LLAMA_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp1-pp2-cp1-gpu"; - } - else if (modelName == "llama_tp2_pp1") - { - beamResult.resultsFile = resultsPath / PathUtil::FP16_PLUGIN_PACKED_PAGED_RESULT_TP2_PP1_FILE(); - modelPath = LLAMA_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp2-pp1-cp1-gpu"; - } - else if (modelName == "llama_tp2_pp2") - { - beamResult.resultsFile = resultsPath / PathUtil::FP16_PLUGIN_PACKED_PAGED_RESULT_TP2_PP2_FILE(); - modelPath = LLAMA_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp2-pp2-cp1-gpu"; - } - else if (modelName == "llama_tp1_pp1") - { - beamResult.resultsFile = resultsPath / PathUtil::FP16_PLUGIN_PACKED_PAGED_RESULT_TP2_PP2_FILE(); - modelPath = LLAMA_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp1-pp1-cp1-gpu"; - } - } - - else - { - TLLM_THROW("Unrecognized modelName"); - } - - // Warning: This should be the last check before running the test. - // It will initialize MPI which can take significant time. - if (modelName == "llama_tp4_pp1" || modelName == "llama_tp1_pp4" || modelName == "llama_tp2_pp2" - || modelName == "llama_tp1_pp2" || modelName == "llama_tp2_pp1") - { - if (outConfig.returnLogProbs || outConfig.returnContextLogits || outConfig.returnGenerationLogits) - { - GTEST_SKIP() << "Skipping logits and log probs tests for mpi runs"; - } - } - - // Returning logits will bring higher latency - if (streaming && (outConfig.returnContextLogits || outConfig.returnGenerationLogits)) - { - mMaxWaitMs = 20000; - } - - auto manager = tr::BufferManager(std::make_shared()); - auto const& givenInput = tr::utils::loadNpy(manager, inputPath.string(), tr::MemoryType::kCPU); - auto [givenInputLengths, nbGivenInputs, maxInputLength] = getGivenInputLengths(*givenInput, modelIds.padId); - world_comm.barrier(); - auto contextNum = contextModels.size(); - auto genNum = genModels.size(); - // int deviceCount = -1; - // TLLM_CUDA_CHECK(cudaGetDeviceCount(&deviceCount)); - bool isOrchestrator = commRank == 0; - std::vector ctxExecutorConfigs; - std::vector genExecutorConfigs; - for (int in = 0; in < instanceNum; in++) - { - tensorrt_llm::executor::SchedulerConfig schedulerConfig(CapacitySchedulerPolicy::kMAX_UTILIZATION); - KvCacheConfig kvCacheConfig{true, std::nullopt, std::nullopt, std::nullopt, 0.2}; - - tensorrt_llm::executor::ExecutorConfig executorConfig(maxBeamWidth, schedulerConfig, kvCacheConfig); - tensorrt_llm::executor::OrchestratorConfig orchestratorConfig{ - isOrchestrator, PathUtil::EXECUTOR_WORKER_PATH(), nullptr, spawnProcess}; - - tensorrt_llm::executor::ParallelConfig parallelConfig{tensorrt_llm::executor::CommunicationType::kMPI, - tensorrt_llm::executor::CommunicationMode::kORCHESTRATOR, participantDeviceIdsEachInstance.at(in), - spawnProcess ? std::nullopt : std::optional>(participantIdsEachInstance.at(in)), - orchestratorConfig}; - executorConfig.setParallelConfig(parallelConfig); - executorConfig.setCacheTransceiverConfig( - texec::CacheTransceiverConfig(texec::CacheTransceiverConfig::BackendType::DEFAULT)); - if (in < contextNum) - { - ctxExecutorConfigs.push_back(executorConfig); - } - else - { - genExecutorConfigs.push_back(executorConfig); - } - } - auto disaggExecutor - = DisaggExecutorOrchestrator(contextModels, genModels, ctxExecutorConfigs, genExecutorConfigs, true, true); - - runDisaggTest(disaggExecutor, manager, *givenInput, modelIds, flakyTestInfo, streaming, vocabSizePadded, beamResult, - outConfig, isSpeculativeDecoding, mMaxWaitMs, BatchingType::kINFLIGHT, false); - -#else - - GTEST_SKIP() << "Skipping DisaggExecutor Test"; - -#endif -} - -TEST_P(ConditionalDisaggParamsTest, DisaggTokenComparison) -{ -#if ENABLE_MULTI_DEVICE - if (!tensorrt_llm::common::getEnvUseUCXKvCache()) - { - setenv("UCX_TLS", "^cuda_ipc", 1); // disable cuda_ipc for testing for mpi - } - auto constexpr processNum = 2; - auto constexpr deviceNum = 2; - auto const& modelName = std::get<0>(GetParam()); - auto constexpr controllerRank = 0; - - // params_check - auto const& world_comm = tensorrt_llm::mpi::MpiComm::world(); - int const commRank = world_comm.getRank(); - int const commSize = world_comm.getSize(); - if (commSize != processNum) - { - GTEST_SKIP() << " need " << processNum << " processes but got " << commSize << " mpi processes, skip test."; - } - if (mDeviceCount < deviceNum) - { - GTEST_SKIP() << " need " << deviceNum << " devices but got " << mDeviceCount << " devices, skip test."; - } - - bool isContext = commRank == 0; - bool isGeneration = commRank == 1; - std::vector participatntIds = {commRank}; - std::vector deviceIds = {commRank}; - bool isController = (commRank == controllerRank); - - OutputConfig outConfig(false, false, false, false, false, false); - int const beamWidth = 1; - BeamResult beamResult{beamWidth}; - - bool streaming = false; - int const maxBeamWidth = 1; - ASSERT_TRUE(fs::exists(DATA_PATH)); - - fs::path modelPath; - // set defaults and adjust if needed by different models - fs::path inputPath = DATA_PATH / "input_tokens.npy"; - ModelIds modelIds{50256, 50256}; - SizeType32 vocabSizePadded{50257}; // gpt vocabSizePadded - bool isSpeculativeDecoding{false}; - - // NOTE: This can be used to disable checks for certain prompt batch entries - FlakyTestInfo flakyTestInfo; - - if (modelName == "gpt") - { - auto const resultsPath - = GPT_DATA_PATH / ((beamWidth == 1) ? "sampling" : "beam_search_" + std::to_string(beamWidth)); - modelPath = GPT_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp1-pp1-cp1-gpu"; - beamResult.resultsFile = resultsPath / PathUtil::FP16_PLUGIN_PACKED_PAGED_RESULT_FILE(); - } - else if (modelName == "llama_tp1_pp1_cp1") - { - inputPath = DATA_PATH / LLAMA_INPUT_FILE; - vocabSizePadded = LLAMA_VOCAB_SIZE_PADDED; - - auto const resultsPath - = LLAMA_DATA_PATH / ((beamWidth == 1) ? "sampling" : "beam_search_" + std::to_string(beamWidth)); - modelIds.padId = LLAMA_PAD_ID; - modelIds.endId = LLAMA_END_ID; - beamResult.resultsFile = resultsPath / PathUtil::FP16_PLUGIN_PACKED_PAGED_RESULT_TP1_PP1_FILE(); - modelPath = LLAMA_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp1-pp1-cp1-gpu"; - } - else - { - TLLM_THROW("Unrecognized modelName"); - } - - auto executorConfig = ExecutorConfig(maxBeamWidth); - FloatType freeGpuMemoryFraction = 0.9f; - KvCacheConfig kvCacheConfig{true, std::nullopt, std::nullopt, std::nullopt, freeGpuMemoryFraction}; - executorConfig.setKvCacheConfig(kvCacheConfig); - executorConfig.setRequestStatsMaxIterations(1000); - executorConfig.setCacheTransceiverConfig( - texec::CacheTransceiverConfig(CacheTransceiverConfig::BackendType::DEFAULT)); - auto manager = tr::BufferManager(std::make_shared()); - auto const& givenInput = tr::utils::loadNpy(manager, inputPath.string(), tr::MemoryType::kCPU); - auto [givenInputLengths, nbGivenInputs, maxInputLength] = getGivenInputLengths(*givenInput, modelIds.padId); - world_comm.barrier(); - auto executor = tensorrt_llm::testing::disaggexecutor::DisaggExecutorLeader(modelPath, ModelType::kDECODER_ONLY, - executorConfig, isController, isContext, isGeneration, givenInputLengths.size(), participatntIds, deviceIds, - commRank); - - std::unordered_map reqIdToBatchId; - std::unordered_map> tokens; - auto const* const givenInputData = tr::bufferCast(*givenInput); - - auto const& inputShape = givenInput->getShape(); - ASSERT_EQ(inputShape.nbDims, 2); - ASSERT_GT(inputShape.d[0], 0); - - // Load expected outputs for each beam width value - auto testData = TestData::loadTestData(beamResult, *givenInput, beamWidth, manager, outConfig, modelIds); - auto const maxSeqLen = testData.maxSeqLen; - - // Load expected outputs and inputs - SizeType32 numRequests = static_cast(givenInputLengths.size()); - SizeType32 maxRequests = numRequests; - std::vector requests; - std::vector reqMaxNewTokens; - SizeType32 const numReturnSequences = 1; - - for (SizeType32 req = 0; req < maxRequests; ++req) - { - SizeType32 inputLen = givenInputLengths.at(req); - auto maxNewTokens = maxSeqLen - maxInputLength; - reqMaxNewTokens.push_back(maxNewTokens); - SizeType32 endId = -1; - auto const* const seqBegin = givenInputData + req * maxInputLength; - VecTokens tokens(seqBegin, seqBegin + inputLen); - auto samplingConfig = tensorrt_llm::executor::SamplingConfig(beamWidth); - samplingConfig.setNumReturnSequences(numReturnSequences); - auto request = Request( - VecTokens(seqBegin, seqBegin + inputLen), maxNewTokens, streaming, samplingConfig, outConfig, endId); - request.setReturnAllGeneratedTokens(false); - // setting request type to context/full by condition - if (req % 2 == 0) - { - request.setRequestType(RequestType::REQUEST_TYPE_CONTEXT_ONLY); - } - else - { - request.setRequestType(RequestType::REQUEST_TYPE_CONTEXT_AND_GENERATION); - } - requests.emplace_back(std::move(request)); - } - - if (isController) - { - std::vector reqIds; - - for (int i = 0; i < requests.size(); ++i) - { - std::vector resultTokens; - resultTokens.reserve(numReturnSequences); - for (SizeType32 seqIdx = 0; seqIdx < numReturnSequences; ++seqIdx) - { - resultTokens.emplace_back(beamWidth); - } - auto retReqId = executor.enqueueRequests({requests[i]}); - reqIds.push_back(retReqId.front()); - tokens[i] = std::move(resultTokens); - reqIdToBatchId[retReqId.front()] = i; - } - - // Get the new tokens for each requests - int32_t numFinished = 0; - int iter = 0; - SizeType32 numResponses = 0; - while (numFinished < maxRequests && iter < mMaxWaitMs) - { - std::chrono::milliseconds waitTime(1); - auto responses = executor.awaitResponses(waitTime); - for (auto& response : responses) - { - numResponses++; - if (!response.hasError()) - { - auto result = response.getResult(); - numFinished += result.isFinal; - auto batchId = reqIdToBatchId.at(response.getRequestId()); - auto seqIdx = result.sequenceIndex; - - auto& outputTokenIds = result.outputTokenIds; - - EXPECT_EQ(result.finishReasons.size(), beamWidth); - for (SizeType32 beam = 0; beam < beamWidth; ++beam) - { - auto& newTokens = outputTokenIds.at(beam); - auto& reqTokens = tokens.at(batchId).at(seqIdx).at(beam); - - reqTokens.insert(reqTokens.end(), newTokens.begin(), newTokens.end()); - // FinishReason is only supported for bw=1 and inflight batching. - if (beamWidth == 1 && executorConfig.getBatchingType() == BatchingType::kINFLIGHT) - { - EXPECT_EQ(result.finishReasons.at(beam), - result.isFinal ? FinishReason::kLENGTH : FinishReason::kNOT_FINISHED); - } - } - } - else - { - // Allow response with error only if awaitResponse processed a terminated request id - std::string err = "ReqId " + std::to_string(response.getRequestId()) - + " has already been processed and was terminated."; - EXPECT_EQ(response.getErrorMsg(), err); - } - } - ++iter; - } - EXPECT_LT(iter, mMaxWaitMs); - testData.verifyOutput(tokens, givenInputLengths, streaming, outConfig.excludeInputFromOutput, flakyTestInfo, - isSpeculativeDecoding, beamWidth, numReturnSequences, false); - } - world_comm.barrier(); -#else - GTEST_SKIP() << "Skipping DisaggExecutor Test"; -#endif -} - -INSTANTIATE_TEST_SUITE_P(GptDisaggSymmetricExecutorTest, DisaggParamsTest, - testing::Combine( // - testing::Values(2), // processNum - testing::Values(std::vector{"gpt", "gpt"}), // modelNames - testing::Values(std::vector>{{0}, {1}}), // participantIdsEachInstance - testing::Values(std::vector>{{0}, {1}}), // participantDeviceIdsEachInstance - testing::Values(std::vector{InstanceRole::kCONTEXT, InstanceRole::kGENERATION}), // instanceRoles - testing::Values(0, 1) // controllerRank - ), - generateTestNameDisaggParams); - -INSTANTIATE_TEST_SUITE_P(GptDisaggSymmetricExecutorMixedTest, DisaggParamsTest, - testing::Combine( // - testing::Values(2), // processNum - testing::Values(std::vector{"gpt", "gpt"}), // modelNames - testing::Values(std::vector>{{0}, {1}}), // participantIdsEachInstance - testing::Values(std::vector>{{0}, {1}}), // participantDeviceIdsEachInstance - testing::Values(std::vector{InstanceRole::kMIXED, InstanceRole::kMIXED}), // instanceRoles - testing::Values(1) // controllerRank - ), - generateTestNameDisaggParams); - -INSTANTIATE_TEST_SUITE_P(GptSingleDeviceDisaggSymmetricExecutorTest, DisaggParamsTest, - testing::Combine( // - testing::Values(2), // processNum - testing::Values(std::vector{"gpt", "gpt"}), // modelNames - testing::Values(std::vector>{{0}, {1}}), // participantIdsEachInstance - testing::Values(std::vector>{{0}, {0}}), // participantDeviceIdsEachInstance - testing::Values(std::vector{InstanceRole::kCONTEXT, InstanceRole::kGENERATION}), // instanceRoles - testing::Values(0) // controllerRank - ), - generateTestNameDisaggParams); - -INSTANTIATE_TEST_SUITE_P(GptSingleDeviceDisaggSymmetricExecutorMixedTest, DisaggParamsTest, - testing::Combine( // - testing::Values(2), // processNum - testing::Values(std::vector{"gpt", "gpt"}), // modelNames - testing::Values(std::vector>{{0}, {1}}), // participantIdsEachInstance - testing::Values(std::vector>{{0}, {0}}), // participantDeviceIdsEachInstance - testing::Values(std::vector{InstanceRole::kMIXED, InstanceRole::kMIXED}), // instanceRoles - testing::Values(1) // controllerRank - ), - generateTestNameDisaggParams); - -INSTANTIATE_TEST_SUITE_P(GptConditionalDisaggSymmetricExecutorTest, ConditionalDisaggParamsTest, - testing::Combine(testing::Values("gpt")), generateTestNameCondDisaggParams); - -INSTANTIATE_TEST_SUITE_P(LlamaConditionalDisaggSymmetricExecutorTest, ConditionalDisaggParamsTest, - testing::Combine(testing::Values("llama_tp1_pp1_cp1")), generateTestNameCondDisaggParams); - -INSTANTIATE_TEST_SUITE_P(LlamaTP2DisaggSymmetricExecutorTest, DisaggParamsTest, - testing::Combine( // - testing::Values(4), // processNum - testing::Values(std::vector{"llama_tp2_pp1_cp1", "llama_tp2_pp1_cp1"}), // modelNames - testing::Values(std::vector>{{0, 1}, {2, 3}}), // participantIdsEachInstance - testing::Values(std::vector>{{0, 1}, {2, 3}}), // participantDeviceIdsEachInstance - testing::Values(std::vector{InstanceRole::kCONTEXT, InstanceRole::kGENERATION}), // instanceRoles - testing::Values(0) // controllerRank - ), - generateTestNameDisaggParams); - -INSTANTIATE_TEST_SUITE_P(LlamaPP2DisaggSymmetricExecutorTest, DisaggParamsTest, - testing::Combine( // - testing::Values(4), // processNum - testing::Values(std::vector{"llama_tp1_pp2_cp1", "llama_tp1_pp2_cp1"}), // modelNames - testing::Values(std::vector>{{0, 1}, {2, 3}}), // participantIdsEachInstance - testing::Values(std::vector>{{1, 0}, {3, 2}}), // participantDeviceIdsEachInstance - testing::Values(std::vector{InstanceRole::kCONTEXT, InstanceRole::kGENERATION}), // instanceRoles - testing::Values(0) // controllerRank - ), - generateTestNameDisaggParams); - -INSTANTIATE_TEST_SUITE_P(LlamaTP2DisaggSymmetricExecutorMixedTest, DisaggParamsTest, - testing::Combine( // - testing::Values(2), // processNum - testing::Values(std::vector{"llama_tp2_pp1_cp1"}), // modelNames - testing::Values(std::vector>{{0, 1}}), // participantIdsEachInstance - testing::Values(std::vector>{{0, 1}}), // participantDeviceIdsEachInstance - testing::Values(std::vector{InstanceRole::kMIXED}), // instanceRoles - testing::Values(0) // controllerRank - ), - generateTestNameDisaggParams); - -INSTANTIATE_TEST_SUITE_P(LlamaPP2DisaggSymmetricExecutorMixedTest, DisaggParamsTest, - testing::Combine( // - testing::Values(2), // processNum - testing::Values(std::vector{"llama_tp1_pp2_cp1"}), // modelNames - testing::Values(std::vector>{{0, 1}}), // participantIdsEachInstance - testing::Values(std::vector>{{0, 1}}), // participantDeviceIdsEachInstance - testing::Values(std::vector{InstanceRole::kMIXED}), // instanceRoles - testing::Values(0) // controllerRank - ), - generateTestNameDisaggParams); - -INSTANTIATE_TEST_SUITE_P(LlamaTP2PP2DisaggSymmetricExecutorTest, DisaggParamsTest, - testing::Combine( // - testing::Values(8), // processNum - testing::Values(std::vector{"llama_tp2_pp2_cp1", "llama_tp2_pp2_cp1"}), // modelNames - testing::Values(std::vector>{{0, 1, 2, 3}, {4, 5, 6, 7}}), // participantIdsEachInstance - testing::Values(std::vector>{{2, 3, 0, 1}, {2, 3, 0, 1}}), // participantDeviceIdsEachInstance - testing::Values(std::vector{InstanceRole::kCONTEXT, InstanceRole::kGENERATION}), // instanceRoles - testing::Values(0) // controllerRank - ), - generateTestNameDisaggParams); - -INSTANTIATE_TEST_SUITE_P(LlamaConPP2GenTP2DisaggAsymmetricExecutorTest, DisaggParamsTest, - testing::Combine( // - testing::Values(4), // processNum - testing::Values(std::vector{"llama_tp1_pp2_cp1", "llama_tp2_pp1_cp1"}), // modelNames - testing::Values(std::vector>{{0, 1}, {2, 3}}), // (1,0) (2,3) // participantIdsEachInstance - testing::Values(std::vector>{{1, 0}, {2, 3}}), // participantDeviceIdsEachInstance - testing::Values(std::vector{InstanceRole::kCONTEXT, InstanceRole::kGENERATION}), // instanceRoles - testing::Values(0) // controllerRank - ), - generateTestNameDisaggParams); - -INSTANTIATE_TEST_SUITE_P(LlamaConTP2GenPP2DisaggAsymmetricExecutorTest, DisaggParamsTest, - testing::Combine( // - testing::Values(4), // processNum - testing::Values(std::vector{"llama_tp2_pp1_cp1", "llama_tp1_pp2_cp1"}), // modelNames - testing::Values(std::vector>{{0, 1}, {2, 3}}), // (0,1), (3,2)// participantIdsEachInstance - testing::Values(std::vector>{{0, 1}, {3, 2}}), // participantDeviceIdsEachInstance - testing::Values(std::vector{InstanceRole::kCONTEXT, InstanceRole::kGENERATION}), // instanceRoles - testing::Values(0) // controllerRank - ), - generateTestNameDisaggParams); - -INSTANTIATE_TEST_SUITE_P(LlamaConTP2PP2GenPP2DisaggAsymmetricExecutorTest, DisaggParamsTest, - testing::Combine( // - testing::Values(6), // processNum - testing::Values(std::vector{"llama_tp2_pp2_cp1", "llama_tp1_pp2_cp1"}), // modelNames - testing::Values( - std::vector>{{0, 1, 2, 3}, {4, 5}}), // (2,3,0,1) , (5,4)// participantIdsEachInstance - testing::Values(std::vector>{{2, 3, 0, 1}, {1, 0}}), // participantDeviceIdsEachInstance - testing::Values(std::vector{InstanceRole::kCONTEXT, InstanceRole::kGENERATION}), // instanceRoles - testing::Values(0) // controllerRank - ), - generateTestNameDisaggParams); - -INSTANTIATE_TEST_SUITE_P(LlamaConTP2PP2GenTP2DisaggAsymmetricExecutorTest, DisaggParamsTest, - testing::Combine( // - testing::Values(6), // processNum - testing::Values(std::vector{"llama_tp2_pp2_cp1", "llama_tp2_pp1_cp1"}), // modelNames - testing::Values( - std::vector>{{0, 1, 2, 3}, {4, 5}}), // (2,3,0,1), (4,5)// participantIdsEachInstance - testing::Values(std::vector>{{2, 3, 0, 1}, {0, 1}}), // participantDeviceIdsEachInstance - testing::Values(std::vector{InstanceRole::kCONTEXT, InstanceRole::kGENERATION}), // instanceRoles - testing::Values(0) // controllerRank - ), - generateTestNameDisaggParams); -INSTANTIATE_TEST_SUITE_P(LlamaConTP2PP1GenTP2PP2DisaggAsymmetricExecutorTest, DisaggParamsTest, - testing::Combine( // - testing::Values(6), // processNum - testing::Values(std::vector{"llama_tp2_pp1_cp1", "llama_tp2_pp2_cp1"}), // modelNames - testing::Values( - std::vector>{{0, 1}, {2, 3, 4, 5}}), // (0,1) , (4,5,2,3)%4// participantIdsEachInstance - testing::Values(std::vector>{{0, 1}, {0, 1, 2, 3}}), // participantDeviceIdsEachInstance - testing::Values(std::vector{InstanceRole::kCONTEXT, InstanceRole::kGENERATION}), // instanceRoles - testing::Values(0) // controllerRank - ), - generateTestNameDisaggParams); - -INSTANTIATE_TEST_SUITE_P(LlamaConTP2GenPP4DisaggAsymmetricExecutorTest, DisaggParamsTest, - testing::Combine( // - testing::Values(6), // processNum - testing::Values(std::vector{"llama_tp2_pp1_cp1", "llama_tp1_pp4_cp1"}), // modelNames - testing::Values( - std::vector>{{4, 5}, {0, 1, 2, 3}}), // (4,5) ,(3,2,1,0)// participantIdsEachInstance - testing::Values(std::vector>{{0, 1}, {3, 2, 1, 0}}), // participantDeviceIdsEachInstance - testing::Values(std::vector{InstanceRole::kCONTEXT, InstanceRole::kGENERATION}), // instanceRoles - testing::Values(0) // controllerRank - ), - generateTestNameDisaggParams); - -INSTANTIATE_TEST_SUITE_P(LlamaCon4TP1Gen1TP4DisaggAsymmetricExecutorTest, DisaggParamsTest, - testing::Combine( // - testing::Values(8), // processNum - testing::Values(std::vector{"llama_tp1_pp1_cp1", "llama_tp1_pp1_cp1", "llama_tp1_pp1_cp1", - "llama_tp1_pp1_cp1", "llama_tp4_pp1_cp1"}), // modelNames - testing::Values(std::vector>{{0}, {1}, {2}, {3}, {4, 5, 6, 7}}), // participantIdsEachInstance - testing::Values( - std::vector>{{0}, {1}, {2}, {3}, {0, 1, 2, 3}}), // participantDeviceIdsEachInstance - testing::Values(std::vector{InstanceRole::kCONTEXT, InstanceRole::kCONTEXT, - InstanceRole::kCONTEXT, InstanceRole::kCONTEXT, InstanceRole::kGENERATION}), // instanceRoles - testing::Values(4) // controllerRank - ), - generateTestNameDisaggParams); - -INSTANTIATE_TEST_SUITE_P(LlamaCon2TP1Gen2TP2AndPP2DisaggAsymmetricExecutorTest, DisaggParamsTest, - testing::Combine( // - testing::Values(6), // processNum - testing::Values(std::vector{ - "llama_tp1_pp1_cp1", "llama_tp1_pp1_cp1", "llama_tp2_pp1_cp1", "llama_tp1_pp2_cp1"}), // modelNames - testing::Values(std::vector>{{0}, {1}, {2, 3}, {4, 5}}), // participantIdsEachInstance - testing::Values(std::vector>{{0}, {1}, {2, 3}, {1, 0}}), // participantDeviceIdsEachInstance - testing::Values(std::vector{InstanceRole::kCONTEXT, InstanceRole::kCONTEXT, - InstanceRole::kGENERATION, InstanceRole::kGENERATION}), // instanceRoles - testing::Values(0) // controllerRank - ), - generateTestNameDisaggParams); - -INSTANTIATE_TEST_SUITE_P(LlamaCon2TP1Gen2PP2DisaggAsymmetricExecutorTest, DisaggParamsTest, - testing::Combine( // - testing::Values(6), // processNum - testing::Values(std::vector{ - "llama_tp1_pp1_cp1", "llama_tp1_pp1_cp1", "llama_tp1_pp2_cp1", "llama_tp1_pp2_cp1"}), // modelNames - testing::Values(std::vector>{{0}, {1}, {2, 3}, {4, 5}}), // participantIdsEachInstance - testing::Values(std::vector>{{0}, {1}, {3, 2}, {1, 0}}), // participantDeviceIdsEachInstance - testing::Values(std::vector{InstanceRole::kCONTEXT, InstanceRole::kCONTEXT, - InstanceRole::kGENERATION, InstanceRole::kGENERATION}), // instanceRoles - testing::Values(0) // controllerRank - ), - generateTestNameDisaggParams); - -INSTANTIATE_TEST_SUITE_P(LlamaCon4TP1Gen1TP2PP2DisaggAsymmetricExecutorTest, DisaggParamsTest, - testing::Combine( // - testing::Values(8), // processNum - testing::Values(std::vector{"llama_tp1_pp1_cp1", "llama_tp1_pp1_cp1", "llama_tp1_pp1_cp1", - "llama_tp1_pp1_cp1", "llama_tp2_pp2_cp1"}), // modelNames - testing::Values(std::vector>{{0}, {1}, {2}, {3}, {4, 5, 6, 7}}), // participantIdsEachInstance - testing::Values( - std::vector>{{0}, {1}, {2}, {3}, {2, 3, 0, 1}}), // participantDeviceIdsEachInstance - testing::Values(std::vector{InstanceRole::kCONTEXT, InstanceRole::kCONTEXT, - InstanceRole::kCONTEXT, InstanceRole::kCONTEXT, InstanceRole::kGENERATION}), // instanceRoles - testing::Values(4) // controllerRank - ), - generateTestNameDisaggParams); - -INSTANTIATE_TEST_SUITE_P(LlamaCon2TP1Gen2TP2DisaaggOrchestrator, DisaggOrchestratorParamsTest, - testing::Combine( // - testing::Values(7), // processNum - testing::Values( - std::vector{"llama_tp1_pp1", "llama_tp1_pp1", "llama_tp2_pp1", "llama_tp2_pp1"}), // modelNames - testing::Values(std::vector>{{1}, {2}, {3, 4}, {5, 6}}), // participantIdsEachInstance - testing::Values(std::vector>{{0}, {1}, {2, 3}, {0, 1}}), // participantDeviceIdsEachInstance - testing::Values(std::vector{InstanceRole::kCONTEXT, InstanceRole::kCONTEXT, - InstanceRole::kGENERATION, InstanceRole::kGENERATION}), // instanceRoles - testing::Values(0) // controllerRank - ), - generateTestNameDisaggParams); -// for disaggOrchestrator 1->0, 2->1, 3->2, 4->3, 5->0, 6->1 - -INSTANTIATE_TEST_SUITE_P(LlamaCon2TP2Gen2TP1DisaaggOrchestrator, DisaggOrchestratorParamsTest, - testing::Combine( // - testing::Values(7), // processNum - testing::Values( - std::vector{"llama_tp2_pp1", "llama_tp2_pp1", "llama_tp1_pp1", "llama_tp1_pp1"}), // modelNames - testing::Values(std::vector>{{1, 2}, {3, 4}, {5}, {6}}), // participantIdsEachInstance - testing::Values(std::vector>{{0, 1}, {2, 3}, {0}, {1}}), // participantDeviceIdsEachInstance - testing::Values(std::vector{InstanceRole::kCONTEXT, InstanceRole::kCONTEXT, - InstanceRole::kGENERATION, InstanceRole::kGENERATION}), // instanceRoles - testing::Values(0) // controllerRank - ), - generateTestNameDisaggParams); - -INSTANTIATE_TEST_SUITE_P(LlamaCon2TP1Gen2PP2DisaaggOrchestrator, DisaggOrchestratorParamsTest, - testing::Combine( // - testing::Values(7), // processNum - testing::Values( - std::vector{"llama_tp1_pp1", "llama_tp1_pp1", "llama_tp1_pp2", "llama_tp1_pp2"}), // modelNames - testing::Values(std::vector>{{1}, {2}, {3, 4}, {5, 6}}), // participantIdsEachInstance - testing::Values(std::vector>{{0}, {1}, {3, 2}, {1, 0}}), // participantDeviceIdsEachInstance - testing::Values(std::vector{InstanceRole::kCONTEXT, InstanceRole::kCONTEXT, - InstanceRole::kGENERATION, InstanceRole::kGENERATION}), // instanceRoles - testing::Values(0) // controllerRank - ), - generateTestNameDisaggParams); - -INSTANTIATE_TEST_SUITE_P(LlamaCon2TP1Gen1TP2PP2DisaaggOrchestrator, DisaggOrchestratorParamsTest, - testing::Combine( // - testing::Values(7), // processNum - testing::Values(std::vector{"llama_tp1_pp1", "llama_tp1_pp1", "llama_tp2_pp2"}), // modelNames - testing::Values(std::vector>{{1}, {2}, {3, 4, 5, 6}}), // participantIdsEachInstance - testing::Values(std::vector>{{0}, {1}, {0, 1, 2, 3}}), // participantDeviceIdsEachInstance - testing::Values(std::vector{ - InstanceRole::kCONTEXT, InstanceRole::kCONTEXT, InstanceRole::kGENERATION}), // instanceRoles - testing::Values(0) // controllerRank - ), - generateTestNameDisaggParams); - -INSTANTIATE_TEST_SUITE_P(LlamaCon2TP2Gen2TP1DisaggSpawnOrchestrator, DisaggOrchestratorParamsTest, - testing::Combine( // - testing::Values(1), // processNum - testing::Values( - std::vector{"llama_tp2_pp1", "llama_tp2_pp1", "llama_tp1_pp1", "llama_tp1_pp1"}), // modelNames - testing::Values(std::vector>{{1, 2}, {3, 4}, {5}, {6}}), // participantIdsEachInstance - testing::Values(std::vector>{{0, 1}, {2, 3}, {0}, {1}}), // participantDeviceIdsEachInstance - testing::Values(std::vector{InstanceRole::kCONTEXT, InstanceRole::kCONTEXT, - InstanceRole::kGENERATION, InstanceRole::kGENERATION}), // instanceRoles - testing::Values(0) // controllerRank - ), - generateTestNameDisaggParams); - -INSTANTIATE_TEST_SUITE_P(LlamaCon2TP1Gen2PP2DisaggSpawnOrchestrator, DisaggOrchestratorParamsTest, - testing::Combine( // - testing::Values(1), // processNum - testing::Values( - std::vector{"llama_tp1_pp1", "llama_tp1_pp1", "llama_tp1_pp2", "llama_tp1_pp2"}), // modelNames - testing::Values(std::vector>{{1}, {2}, {3, 4}, {5, 6}}), // participantIdsEachInstance - testing::Values(std::vector>{{0}, {1}, {3, 2}, {1, 0}}), // participantDeviceIdsEachInstance - testing::Values(std::vector{InstanceRole::kCONTEXT, InstanceRole::kCONTEXT, - InstanceRole::kGENERATION, InstanceRole::kGENERATION}), // instanceRoles - testing::Values(0) // controllerRank - ), - generateTestNameDisaggParams); diff --git a/cpp/tests/e2e_tests/executor/encDecTest.cpp b/cpp/tests/e2e_tests/executor/encDecTest.cpp deleted file mode 100644 index 0095ae30ad40..000000000000 --- a/cpp/tests/e2e_tests/executor/encDecTest.cpp +++ /dev/null @@ -1,387 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2023-2025 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. - */ - -#include "executorTest.h" - -#include "tensorrt_llm/common/logger.h" -#include "tensorrt_llm/executor/executor.h" -#include "tensorrt_llm/executor/types.h" -#include "tensorrt_llm/runtime/iBuffer.h" -#include "tensorrt_llm/runtime/iTensor.h" -#include "tensorrt_llm/runtime/utils/mpiUtils.h" -#include "tensorrt_llm/runtime/utils/numpyUtils.h" -#include "tensorrt_llm/testing/modelSpec.h" -#include "tests/utils/common.h" - -#include -#include -#include - -#include -#include -#include -#include -#include - -namespace tr = tensorrt_llm::runtime; - -using namespace tensorrt_llm::testing; -using namespace tensorrt_llm::executor; -using namespace std::chrono_literals; -using tensorrt_llm::testing::KVCacheType; - -namespace -{ - -std::string getEncDecEnginePath(std::string const& modelName, SizeType32 tp, SizeType32 pp, SizeType32 cp) -{ - return modelName + '/' + std::to_string(tp * pp * cp) + "-gpu/float16"; -} - -TokenIdType getDecTokenFromJsonConfig(std::filesystem::path decEnginePath, std::string const& token_name) -{ - TokenIdType tokenId = 0; - try - { - std::ifstream decoderJsonConfigPath(decEnginePath / "config.json"); - auto const decoderPretrainedConfig - = nlohmann::json::parse(decoderJsonConfigPath, nullptr, true, true).at("pretrained_config"); - tokenId = decoderPretrainedConfig.at(token_name).template get(); - } - catch (nlohmann::json::out_of_range& e) - { - TLLM_LOG_ERROR( - "Parameter %s cannot be found from decoder config.json in pretrained_config. Using default id 0.", - token_name.c_str()); - } - catch (nlohmann::json::type_error const& e) - { - TLLM_LOG_ERROR( - "Parameter %s has a different type from decoder config.json in pretrained_config. Using default id 0.", - token_name.c_str()); - } - return tokenId; -} - -} // namespace - -using EncDecParamsType = std::tuple>; - -std::string generateTestNameEncDec(testing::TestParamInfo const& info) -{ - auto modelName = std::get<0>(info.param); - auto const beamWidth = std::get<1>(info.param); - auto const maxNewTokens = std::get<2>(info.param); - auto const tp = std::get<3>(info.param); - auto const pp = std::get<4>(info.param); - - // GTEST does not allow '-' in its test name - for (auto& c : modelName) - { - if (c == '-') - { - c = '_'; - } - } - - std::string name = "EncDecTest"; - name.append("_" + modelName); - name.append("_BeamWidth" + std::to_string(beamWidth)); - name.append("_MaxNewTokens" + std::to_string(maxNewTokens)); - name.append("_TP" + std::to_string(tp)); - name.append("_PP" + std::to_string(pp)); - return name; -} - -bool isLanguageAdapterName(std::string const& modelName) -{ - return modelName == LANGUAGE_ADAPTER_NAME; -} - -class EncDecParamsTest : public GptExecutorTest, public ::testing::WithParamInterface -{ -}; - -TEST_P(EncDecParamsTest, validEncDecCtor) -{ - auto const modelName = std::get<0>(GetParam()); - SizeType32 const beamWidth = std::get<1>(GetParam()); - SizeType32 const maxNewTokens = std::get<2>(GetParam()); - SizeType32 const tp = std::get<3>(GetParam()); - SizeType32 const pp = std::get<4>(GetParam()); - SizeType32 const cp = std::get<5>(GetParam()); - - auto const enginePathName = getEncDecEnginePath(modelName, tp, pp, cp); - std::filesystem::path encEnginePath = ENC_DEC_ENGINE_BASE / enginePathName / "encoder"; - std::filesystem::path decEnginePath = ENC_DEC_ENGINE_BASE / enginePathName / "decoder"; - ExecutorConfig executorConfig{}; - FloatType freeGpuMemoryFraction = 0.4f; - FloatType crossKvCacheFraction = 0.4f; - KvCacheConfig kvCacheConfig{false, std::nullopt, std::nullopt, std::nullopt, freeGpuMemoryFraction}; - kvCacheConfig.setCrossKvCacheFraction(crossKvCacheFraction); - executorConfig.setKvCacheConfig(kvCacheConfig); - auto executor = Executor(encEnginePath, decEnginePath, ModelType::kENCODER_DECODER, executorConfig); -} - -TEST_P(EncDecParamsTest, Forward) -{ - bool constexpr VERBOSE = false; - auto const modelName = std::get<0>(GetParam()); - SizeType32 const beamWidth = std::get<1>(GetParam()); - SizeType32 const maxNewTokens = std::get<2>(GetParam()); - SizeType32 const tp = std::get<3>(GetParam()); - SizeType32 const pp = std::get<4>(GetParam()); - SizeType32 const cp = std::get<5>(GetParam()); - - // Parameters for language adapter test - SizeType32 const numLanguages = std::get<6>(GetParam()); - std::vector languageAdapterUids = std::get<7>(GetParam()); - - bool const streaming = false; - - auto const enginePathName = getEncDecEnginePath(modelName, tp, pp, cp); - std::filesystem::path encEnginePath = ENC_DEC_ENGINE_BASE / enginePathName / "encoder"; - std::filesystem::path decEnginePath = ENC_DEC_ENGINE_BASE / enginePathName / "decoder"; - - // load ground truth input & output data - auto manager = tr::BufferManager(std::make_shared()); - auto inputsIdsHost - = tr::utils::loadNpy(manager, (ENC_DEC_DATA_BASE / "input_ids.npy").string(), tr::MemoryType::kCPU); - auto inputsIdsPtr = tr::bufferCast(*inputsIdsHost); - auto inputLengthsHost - = tr::utils::loadNpy(manager, (ENC_DEC_DATA_BASE / "input_lengths.npy").string(), tr::MemoryType::kCPU); - auto inputLengthsPtr = tr::bufferCast(*inputLengthsHost); - auto encoderOutputHost - = tr::utils::loadNpy(manager, (ENC_DEC_DATA_BASE / "encoder_output.npy").string(), tr::MemoryType::kCPU); - auto encoderOutputPtr = tr::bufferCast(*encoderOutputHost); - auto decoderOutputHost = tr::utils::loadNpy(manager, - (ENC_DEC_DATA_BASE / "output_ids_beam").string() + std::to_string(beamWidth) + ".npy", tr::MemoryType::kCPU); - auto decoderOutputPtr = tr::bufferCast(*decoderOutputHost); - - // Rank and size info - auto& comm = tensorrt_llm::mpi::MpiComm::world(); - auto const worldRank = comm.getRank(); - auto const worldSize = comm.getSize(); - - // create executor - BatchingType const batchingType = BatchingType::kINFLIGHT; - FloatType freeGpuMemoryFraction = 0.5f; - FloatType crossKvCacheFraction = 0.5f; - KvCacheConfig kvCacheConfig{false, std::nullopt, std::nullopt, std::nullopt, freeGpuMemoryFraction}; - kvCacheConfig.setCrossKvCacheFraction(crossKvCacheFraction); - - ExecutorConfig executorConfig{beamWidth}; - executorConfig.setBatchingType(batchingType); - executorConfig.setKvCacheConfig(kvCacheConfig); - executorConfig.setNormalizeLogProbs(false); - - // TODO: OrchestratorMode test does not pass - bool const useOrchestratorMode = (tp * pp) > worldSize; - std::optional orchestratorConfig = std::nullopt; - if (useOrchestratorMode) - { - orchestratorConfig = OrchestratorConfig(true, PathUtil::EXECUTOR_WORKER_PATH()); - } - auto parallelConfig = ParallelConfig(CommunicationType::kMPI, - useOrchestratorMode ? CommunicationMode::kORCHESTRATOR : CommunicationMode::kLEADER, std::nullopt, std::nullopt, - orchestratorConfig); - executorConfig.setParallelConfig(parallelConfig); - - auto executor = Executor(encEnginePath, decEnginePath, ModelType::kENCODER_DECODER, executorConfig); - - OutputConfig outConfig; - outConfig.excludeInputFromOutput = false; - outConfig.returnLogProbs = false; - outConfig.returnGenerationLogits = false; - outConfig.returnContextLogits = false; - outConfig.returnEncoderOutput = false; - - TokenIdType bosId = getDecTokenFromJsonConfig(decEnginePath, "bos_token_id"); - TokenIdType padId = getDecTokenFromJsonConfig(decEnginePath, "pad_token_id"); - TokenIdType eosId = getDecTokenFromJsonConfig(decEnginePath, "eos_token_id"); - TokenIdType decoderStartTokenId = getDecTokenFromJsonConfig(decEnginePath, "decoder_start_token_id"); - - bool const isLanguageAdapterTest = isLanguageAdapterName(modelName); - // create requests - SizeType32 const nbRequests = inputLengthsHost->getShape().d[0]; - std::vector requests; - for (int i = 0, cumInputLen = 0; i < nbRequests; i++) - { - auto encoderInput = VecTokens(&inputsIdsPtr[cumInputLen], - &inputsIdsPtr[cumInputLen] + inputLengthsPtr[i]); // assume inputIds is flattened / no-padding - cumInputLen += inputLengthsPtr[i]; - auto decoderInput = VecTokens{decoderStartTokenId}; - Request req(decoderInput, maxNewTokens, streaming, tensorrt_llm::executor::SamplingConfig(beamWidth), outConfig, - eosId, padId); - req.setEncoderInputTokenIds(encoderInput); - if (isLanguageAdapterTest) - { - req.setLanguageAdapterUid(languageAdapterUids[i]); - } - requests.emplace_back(req); - } - - using namespace std::chrono; - - // enqueue requests - if (worldRank == 0) - { - auto tik = high_resolution_clock::now(); - std::vector reqIds = executor.enqueueRequests(std::move(requests)); - - // get responses - milliseconds waitTime(5000); - auto responsesAll = executor.awaitResponses(reqIds, waitTime); - auto tok = high_resolution_clock::now(); - TLLM_LOG_DEBUG("TRT-LLM C++ E2E time %d ms", duration_cast(tok - tik).count()); - TLLM_LOG_DEBUG("Number of responses: %d", responsesAll.size()); - - int32_t numFinished = 0; - int iter = 0; - SizeType32 numResponses = 0; - std::unordered_map> outputTokens; - for_each(reqIds.begin(), reqIds.end(), - [&outputTokens, &beamWidth](auto const& id) - { - TLLM_LOG_DEBUG("Request IDs: %d", id); - outputTokens[id] = {}; - for (int i = 0; i < beamWidth; i++) - { - outputTokens[id].emplace_back(VecTokens{}); - } - }); - for (int i = 0; i < reqIds.size(); i++) - { - auto& responses = responsesAll[i]; - for (auto& response : responses) - { - numResponses++; - if (!response.hasError()) - { - auto result = response.getResult(); - numFinished += result.isFinal; - for (int beam = 0; beam < beamWidth; beam++) - { - auto& resTokens = result.outputTokenIds.at(beam); - auto& outTokens = outputTokens.at(response.getRequestId()).at(beam); - outTokens.insert(outTokens.end(), std::make_move_iterator(resTokens.begin()), - std::make_move_iterator(resTokens.end())); - } - } - else - { - // Allow response with error only if awaitResponse processed a terminated request id - std::string err = "ReqId " + std::to_string(response.getRequestId()) - + " has already been processed and was terminated."; - EXPECT_EQ(response.getErrorMsg(), err); - } - } - } - - // print output & check correctness with ground truth - for (auto const& [reqId, tokens] : outputTokens) - { - SizeType32 gtMaxLength = decoderOutputHost->getShape().d[1]; - auto gtOutput = decoderOutputPtr + (reqId - 1) * gtMaxLength; - - if constexpr (VERBOSE) - { - std::cout << ">>> Request ID: " << reqId << std::endl; - for (int beam = 0; beam < beamWidth; beam++) - { - std::cout << "output tokens, beam " << beam << ", output length " << tokens[beam].size() << ": " - << std::endl; - for_each(tokens[beam].begin(), tokens[beam].end(), - [](auto const& token) { std::cout << token << ", "; }); - std::cout << std::endl; - } - std::cout << "ground truth tokens: " << std::endl; - - SizeType32 gtLength = 0; - for (int i = 0; i < gtMaxLength; i++) - { - if (gtOutput[i] != eosId) - { - std::cout << gtOutput[i] << ", "; - gtLength++; - } - } - std::cout << std::endl; - std::cout << "ground truth length: " << gtLength << std::endl; - } - - // check token-by-token match between beam 0 & ground truth - ASSERT_TRUE(tokens.size() <= gtMaxLength) - << "Request ID " << reqId << "'s generated length is longer than ground truth length " << gtMaxLength; - for (int i = 0; i < gtMaxLength; i++) - { - if (outConfig.excludeInputFromOutput) - { - // if results exclude decoder start token, skip it in ground truth too - continue; - } - if (i < tokens[0].size()) - { - ASSERT_EQ(tokens[0][i], gtOutput[i]) - << "Generated token id: " << tokens[0][i] << " v.s. ground truth: " << gtOutput[i]; - } - else - { - ASSERT_EQ(gtOutput[i], eosId) << "Request ID " << reqId << "'s generated length " << tokens.size() - << " is shorter than ground truth length " << gtMaxLength; - } - } - } - } -} - -INSTANTIATE_TEST_SUITE_P(T5BasicTest, EncDecParamsTest, - testing::Combine(testing::Values(T5_NAME), testing::Values(1), testing::Values(64), testing::Values(1), - testing::Values(1), testing::Values(1), testing::Values(0), testing::Values(std::vector{})), - generateTestNameEncDec); - -INSTANTIATE_TEST_SUITE_P(T5Beam2Test, EncDecParamsTest, - testing::Combine(testing::Values(T5_NAME), testing::Values(2), testing::Values(64), testing::Values(1), - testing::Values(1), testing::Values(1), testing::Values(0), testing::Values(std::vector{})), - generateTestNameEncDec); - -INSTANTIATE_TEST_SUITE_P(T5MultiGPUTest, EncDecParamsTest, - testing::Combine(testing::Values(T5_NAME), testing::Values(1), testing::Values(64), testing::Values(4), - testing::Values(1), testing::Values(1), testing::Values(0), testing::Values(std::vector{})), - generateTestNameEncDec); - -INSTANTIATE_TEST_SUITE_P(BartBasicTest, EncDecParamsTest, - testing::Combine(testing::Values(BART_NAME), testing::Values(1), testing::Values(64), testing::Values(1), - testing::Values(1), testing::Values(1), testing::Values(0), testing::Values(std::vector{})), - generateTestNameEncDec); - -INSTANTIATE_TEST_SUITE_P(BartBeam2Test, EncDecParamsTest, - testing::Combine(testing::Values(BART_NAME), testing::Values(2), testing::Values(64), testing::Values(1), - testing::Values(1), testing::Values(1), testing::Values(0), testing::Values(std::vector{})), - generateTestNameEncDec); - -INSTANTIATE_TEST_SUITE_P(BartMultiGPUTest, EncDecParamsTest, - testing::Combine(testing::Values(BART_NAME), testing::Values(1), testing::Values(64), testing::Values(4), - testing::Values(1), testing::Values(1), testing::Values(0), testing::Values(std::vector{})), - generateTestNameEncDec); - -INSTANTIATE_TEST_SUITE_P(LanguageAdapterBasicTest, EncDecParamsTest, - testing::Combine(testing::Values(LANGUAGE_ADAPTER_NAME), testing::Values(1), testing::Values(64), - testing::Values(1), testing::Values(1), testing::Values(1), testing::Values(4), - testing::Values(std::vector{2, 3})), - generateTestNameEncDec); diff --git a/cpp/tests/e2e_tests/executor/executorMockTest.cpp b/cpp/tests/e2e_tests/executor/executorMockTest.cpp deleted file mode 100644 index 0f0176f1ed6b..000000000000 --- a/cpp/tests/e2e_tests/executor/executorMockTest.cpp +++ /dev/null @@ -1,1028 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2023-2025 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. - */ - -#include "executorTest.h" - -#include "tensorrt_llm/batch_manager/trtGptModel.h" -#include "tensorrt_llm/executor/executor.h" -#include "tensorrt_llm/executor/types.h" -#include "tensorrt_llm/runtime/utils/mpiUtils.h" -#include "tensorrt_llm/testing/modelSpec.h" -#include "tests/utils/common.h" - -#include -#include -#include - -#include -#include -#include -#include -#include -#include -#include - -using ::testing::_; -using ::testing::Invoke; - -namespace tr = tensorrt_llm::runtime; -namespace tb = tensorrt_llm::batch_manager; - -using namespace tensorrt_llm::testing; -using namespace tensorrt_llm::executor; -using namespace std::chrono_literals; -using tensorrt_llm::testing::KVCacheType; - -class MockedModel : public Model -{ - using LlmRequestPtr = std::shared_ptr; - using RequestList = std::list; - -public: - MOCK_METHOD(void, forwardSync, (), ()); - MOCK_METHOD(void, forwardAsync, (RequestList const&), ()); - MOCK_METHOD(void, terminateRequest, (std::shared_ptr const& llmRequest, bool pause), ()); - MOCK_METHOD( - void, terminateRequestSync, (std::shared_ptr const& llmRequest, FinishReason finishReason), ()); - MOCK_METHOD(SizeType32, getMaxNumSequences, (), (const)); - MOCK_METHOD(SizeType32, getMaxInputLen, (), (const)); - MOCK_METHOD(SizeType32, getHiddenSize, (), (const)); - MOCK_METHOD(SizeType32, getMaxSequenceLen, (), (const)); - MOCK_METHOD(SizeType32, getVocabSizePadded, (), (const)); - MOCK_METHOD(SizeType32, getMaxDraftLen, (), (const)); - MOCK_METHOD(SizeType32, getNumMicroBatches, (), (const)); - MOCK_METHOD(SizeType32, getOperatingBeamWidth, (), (const)); - MOCK_METHOD(nvinfer1::DataType, getLogitDataType, (), (const)); - MOCK_METHOD(nvinfer1::DataType, getTensorDataType, (std::string const&), (const)); - MOCK_METHOD(nvinfer1::Dims, getTensorShape, (std::string const&), (const)); - MOCK_METHOD(void, getCurrentIterationStats, (IterationStats&), (const)); - MOCK_METHOD(void, getCurrentRequestStats, (RequestStatsPerIteration&), (const)); - MOCK_METHOD(DebugTensorsPerIteration, getCurrentDebugTensors, (), (const)); - MOCK_METHOD(tr::WorldConfig const&, getWorldConfig, (), (const)); - MOCK_METHOD(tr::ModelConfig const&, getModelConfig, (), (const)); - MOCK_METHOD(tr::BufferManager const&, getBufferManager, (), (const)); - MOCK_METHOD(tr::BufferManager::CudaStreamPtr, getRuntimeStreamPtr, (), (const)); - MOCK_METHOD(IterationType, getIterCounter, (), (const, noexcept)); - MOCK_METHOD(bool, hasSpeculativeDecodingFastLogits, (), (const, noexcept)); - MOCK_METHOD(bool, getGatherGenerationLogits, (), (const)); - MOCK_METHOD(void, updatePeftCache, (LlmRequestPtr const& llmReqeust), ()); - MOCK_METHOD(void, setLogitsPostProcessorBatched, (std::optional), ()); - MOCK_METHOD(void, setReplicateLogitsPostProcessor, (bool), ()); - MOCK_METHOD(bool, getReplicateLogitsPostProcessor, (), (const)); - MOCK_METHOD(bool, hasGuidedDecoder, (), (const, noexcept)); - MOCK_METHOD(void, resetIterationStats, (), ()); - MOCK_METHOD( - std::shared_ptr, getKVCacheManager, (), ()); - MOCK_METHOD(std::shared_ptr, - getKVCacheManager, (), (const)); - MOCK_METHOD(SizeType32, getMaxCapacityBatchSize, (SizeType32, SizeType32), (const)); -}; - -using ParamType = std::tuple; - -std::string generateTestName(testing::TestParamInfo const& info) -{ - auto const streaming = std::get<0>(info.param); - auto const excludeInputFromOutput = std::get<1>(info.param); - auto const beamWidth = std::get<2>(info.param); - std::string name = "ExecutorTest"; - if (streaming) - { - name += "Streaming"; - } - if (excludeInputFromOutput) - { - name += "ExclInput"; - } - name.append("BW" + std::to_string(beamWidth)); - return name; -} - -class ParamTest : public GptExecutorTest, public ::testing::WithParamInterface -{ -}; - -TEST_P(ParamTest, MockedModel) -{ - using LlmRequestPtr = std::shared_ptr; - using RequestList = std::list; - - bool const streaming = std::get<0>(GetParam()); - bool const excludeInputFromOutput = std::get<1>(GetParam()); - auto const beamWidth = std::get<2>(GetParam()); - OutputConfig outConfig; - outConfig.excludeInputFromOutput = excludeInputFromOutput; - auto model = std::make_shared(); - - EXPECT_CALL(*model, terminateRequest(_, _)).Times(0); - EXPECT_CALL(*model, getVocabSizePadded()).Times(0); - EXPECT_CALL(*model, getLogitDataType()).Times(0); - tr::ModelConfig dummyModelConfig(0, 0, 0, 0, 1, 0, nvinfer1::DataType::kHALF); - EXPECT_CALL(*model, getModelConfig()) - .WillRepeatedly(Invoke([&]() -> tr::ModelConfig const& { return dummyModelConfig; })); - SizeType32 callCount = 0; - EXPECT_CALL(*model, forwardAsync(_)) - .WillRepeatedly(Invoke( - [&](RequestList const& requestList) - { - for (auto const& llmReq : requestList) - { - // Don't add any tokens to simulate no output tokens - llmReq->addNewTokens(VecTokens(beamWidth, 1)); - llmReq->setState(tb::LlmRequestState::kGENERATION_IN_PROGRESS); - if (llmReq->getMaxNumGeneratedTokens() >= llmReq->mMaxNewTokens) - { - llmReq->setState(tb::LlmRequestState::kGENERATION_COMPLETE); - } - } - callCount++; - })); - - EXPECT_CALL(*model, getMaxNumSequences()).WillRepeatedly(Invoke([&]() { return 10; })); - EXPECT_CALL(*model, getMaxInputLen()).WillRepeatedly(Invoke([&]() { return 10; })); - EXPECT_CALL(*model, getMaxSequenceLen()).WillRepeatedly(Invoke([&]() { return 10; })); - EXPECT_CALL(*model, getMaxDraftLen()).WillRepeatedly(Invoke([&]() { return 0; })); - EXPECT_CALL(*model, getVocabSizePadded()).WillRepeatedly(Invoke([&]() { return 80000; })); - tr::WorldConfig const dummyWorldConfig; - EXPECT_CALL(*model, getWorldConfig()) - .WillRepeatedly(Invoke([&]() -> tr::WorldConfig const& { return dummyWorldConfig; })); - EXPECT_CALL(*model, getCurrentIterationStats(_)).WillRepeatedly(Invoke([&](IterationStats& /*stats*/) { return; })); - EXPECT_CALL(*model, getCurrentRequestStats(_)) - .WillRepeatedly(Invoke([&](RequestStatsPerIteration& /*stats*/) { return; })); - - ExecutorConfig const executorConfig(beamWidth); - auto executor = Executor(model, executorConfig); - - // Create the request - constexpr SizeType32 maxNewTokens = 5; - VecTokens const inputTokens{1, 2, 3, 4}; - auto request - = Request(inputTokens, maxNewTokens, streaming, tensorrt_llm::executor::SamplingConfig(beamWidth), outConfig); - - // Enqueue the request - auto requestId = executor.enqueueRequest(request); - - bool done = false; - int iter = 0; - while (!done && iter < mMaxWaitMs) - { - std::chrono::milliseconds waitTime(1); - auto responses = executor.awaitResponses(requestId, waitTime); - for (auto& response : responses) - { - auto const& result = response.getResult(); - done = result.isFinal; - } - ++iter; - } - - EXPECT_LT(iter, mMaxWaitMs); - EXPECT_EQ(callCount, maxNewTokens); -} - -TEST_F(GptExecutorTest, MockedModelMaxQueueSize) -{ - using LlmRequestPtr = std::shared_ptr; - using RequestList = std::list; - - auto model = std::make_shared(); - - EXPECT_CALL(*model, terminateRequest(_, _)).Times(0); - EXPECT_CALL(*model, terminateRequestSync(_, _)).Times(0); - EXPECT_CALL(*model, getVocabSizePadded()).Times(0); - EXPECT_CALL(*model, getLogitDataType()).Times(0); - - SizeType32 callCount = 0; - EXPECT_CALL(*model, forwardAsync(_)) - .WillRepeatedly(Invoke( - [&](RequestList const& requestList) - { - for (auto const& llmReq : requestList) - { - // Sleep to allow queue to fill up - std::this_thread::sleep_for(std::chrono::milliseconds(200)); - // Don't add any tokens to simulate no output tokens - llmReq->addNewTokens({1}); - llmReq->setState(tb::LlmRequestState::kGENERATION_IN_PROGRESS); - if (llmReq->getMaxNumGeneratedTokens() >= llmReq->mMaxNewTokens) - { - llmReq->setState(tb::LlmRequestState::kGENERATION_COMPLETE); - } - } - callCount++; - })); - - EXPECT_CALL(*model, getMaxNumSequences()).WillRepeatedly(Invoke([&]() { return 10; })); - EXPECT_CALL(*model, getMaxInputLen()).WillRepeatedly(Invoke([&]() { return 10; })); - EXPECT_CALL(*model, getMaxSequenceLen()).WillRepeatedly(Invoke([&]() { return 10; })); - EXPECT_CALL(*model, getMaxDraftLen()).WillRepeatedly(Invoke([&]() { return 0; })); - EXPECT_CALL(*model, getVocabSizePadded()).WillRepeatedly(Invoke([&]() { return 80000; })); - tr::WorldConfig const dummyWorldConfig; - EXPECT_CALL(*model, getWorldConfig()) - .WillRepeatedly(Invoke([&]() -> tr::WorldConfig const& { return dummyWorldConfig; })); - EXPECT_CALL(*model, getCurrentIterationStats(_)).WillRepeatedly(Invoke([&](IterationStats& /*stats*/) { return; })); - EXPECT_CALL(*model, getCurrentRequestStats(_)) - .WillRepeatedly(Invoke([&](RequestStatsPerIteration& /*stats*/) { return; })); - tr::ModelConfig dummyModelConfig(0, 0, 0, 0, 1, 0, nvinfer1::DataType::kHALF); - EXPECT_CALL(*model, getModelConfig()) - .WillRepeatedly(Invoke([&]() -> tr::ModelConfig const& { return dummyModelConfig; })); - SizeType32 maxQueueSize = 6; - ExecutorConfig executorConfig; - executorConfig.setMaxQueueSize(maxQueueSize); - - auto executor = Executor(model, executorConfig); - - // Create the request - SizeType32 const maxNewTokens = 5; - VecTokens const inputTokens{1, 2, 3, 4}; - auto request = Request(inputTokens, maxNewTokens); - - // Enqueue as many requests as the queue can manage - for (int i = 0; i < maxQueueSize; i++) - { - auto requestId = executor.enqueueRequest(request); - } - try - { - auto requestId = executor.enqueueRequest(request); - - FAIL() << "Expected TllmException"; - } - catch (std::exception const& e) - { - EXPECT_THAT(e.what(), testing::HasSubstr("Maximum queue size of 6 has been reached, please try again later")); - } - - // Wait for requests to get scheduled to free up space in queue - std::this_thread::sleep_for(std::chrono::milliseconds(maxQueueSize * 200)); - auto requestId = executor.enqueueRequest(request); - - try - { - auto samplingConfig = SamplingConfig(1); - samplingConfig.setNumReturnSequences(maxQueueSize); - auto request = Request(inputTokens, maxNewTokens, false, samplingConfig); - auto requestId = executor.enqueueRequest(request); - FAIL() << "Expected TllmException"; - } - catch (std::exception const& e) - { - EXPECT_THAT(e.what(), testing::HasSubstr("Maximum queue size of 6 has been reached, please try again later")); - } -} - -TEST_F(GptExecutorTest, MockedModelReqStatsBug) -{ - using LlmRequestPtr = std::shared_ptr; - using RequestList = std::list; - - bool streaming = false; - bool excludeInputFromOutput = false; - OutputConfig outConfig; - outConfig.excludeInputFromOutput = excludeInputFromOutput; - auto model = std::make_shared(); - - EXPECT_CALL(*model, terminateRequest(_, _)).Times(0); - EXPECT_CALL(*model, getVocabSizePadded()).Times(0); - EXPECT_CALL(*model, getLogitDataType()).Times(0); - EXPECT_CALL(*model, updatePeftCache(_)).WillRepeatedly(Invoke([&]() { return; })); - - SizeType32 callCount = 0; - RequestList currentReq; - EXPECT_CALL(*model, forwardAsync(_)) - .WillRepeatedly(Invoke( - [&](RequestList const& requestList) - { - currentReq = requestList; - for (auto const& llmReq : requestList) - { - // Don't add any tokens to simulate no output tokens - llmReq->addNewTokens({1}); - llmReq->setState(tb::LlmRequestState::kGENERATION_IN_PROGRESS); - } - callCount++; - })); - - EXPECT_CALL(*model, forwardSync()) - .WillRepeatedly(Invoke( - [&]() - { - for (auto const& llmReq : currentReq) - { - if (llmReq->getMaxNumGeneratedTokens() >= llmReq->mMaxNewTokens) - { - llmReq->setState(tb::LlmRequestState::kGENERATION_COMPLETE); - } - } - return; - })); - - EXPECT_CALL(*model, getMaxNumSequences()).WillRepeatedly(Invoke([&]() { return 10; })); - EXPECT_CALL(*model, getMaxInputLen()).WillRepeatedly(Invoke([&]() { return 10; })); - EXPECT_CALL(*model, getMaxSequenceLen()).WillRepeatedly(Invoke([&]() { return 10; })); - EXPECT_CALL(*model, getMaxDraftLen()).WillRepeatedly(Invoke([&]() { return 0; })); - EXPECT_CALL(*model, getVocabSizePadded()).WillRepeatedly(Invoke([&]() { return 80000; })); - tr::WorldConfig const dummyWorldConfig; - EXPECT_CALL(*model, getWorldConfig()) - .WillRepeatedly(Invoke([&]() -> tr::WorldConfig const& { return dummyWorldConfig; })); - EXPECT_CALL(*model, getCurrentIterationStats(_)).WillRepeatedly(Invoke([&](IterationStats& stats) { return; })); - EXPECT_CALL(*model, getCurrentRequestStats(_)) - .WillRepeatedly(Invoke([&](RequestStatsPerIteration& stats) { return; })); - - SizeType32 beamWidth = 1; - ExecutorConfig executorConfig(beamWidth); - executorConfig.setRequestStatsMaxIterations(1000); - auto executor = Executor(model, executorConfig); - - // Create the request - SizeType32 maxNewTokens = 5; - VecTokens inputTokens{1, 2, 3, 4}; - int numRequests = 10000; - auto request - = Request(inputTokens, maxNewTokens, streaming, tensorrt_llm::executor::SamplingConfig(beamWidth), outConfig); - - auto done = std::atomic{false}; - auto statsThreadDone = false; - // Spawn a thread that continuously get stats - auto statsThread = std::thread( - [&executor, &done, &statsThreadDone]() - { - while (!done) - { - auto reqStats = executor.getLatestRequestStats(); - std::this_thread::sleep_for(std::chrono::microseconds(10)); - } - statsThreadDone = true; - }); - - // Spawn a thread that enqueues the requests - std::vector requestIds; - auto enqueueThread = std::thread( - [&executor, &requestIds, &request, &done, numRequests]() - { - for (int i = 0; i < numRequests; ++i) - { - requestIds.push_back(executor.enqueueRequest(request)); - } - done = true; - }); - enqueueThread.join(); - ASSERT_EQ(requestIds.size(), numRequests); - - // Wait for stats thread to be done, fail otherwise - int iter = 0; - while (!statsThreadDone && iter < mMaxWaitMs) - { - std::chrono::milliseconds waitTime(1); - std::this_thread::sleep_for(std::chrono::milliseconds(waitTime)); - iter++; - } - ASSERT_TRUE(statsThreadDone); - statsThread.join(); -} - -TEST_F(GptExecutorTest, MockedModelEvictRestartValidityTest) -{ - using LlmRequestPtr = std::shared_ptr; - using RequestList = std::list; - - constexpr bool excludeInputFromOutput = false; - OutputConfig outConfig; - outConfig.excludeInputFromOutput = excludeInputFromOutput; - auto model = std::make_shared(); - - EXPECT_CALL(*model, terminateRequest(_, _)).Times(0); - EXPECT_CALL(*model, getVocabSizePadded()).Times(0); - EXPECT_CALL(*model, getLogitDataType()).Times(0); - EXPECT_CALL(*model, updatePeftCache(_)).WillRepeatedly(Invoke([&]() { return; })); - tr::ModelConfig dummyModelConfig(0, 0, 0, 0, 1, 0, nvinfer1::DataType::kHALF); - EXPECT_CALL(*model, getModelConfig()) - .WillRepeatedly(Invoke([&]() -> tr::ModelConfig const& { return dummyModelConfig; })); - SizeType32 callCount = 0; - RequestList currentReq; - EXPECT_CALL(*model, forwardAsync(_)) - .WillRepeatedly(Invoke( - [&](RequestList const& requestList) - { - currentReq = requestList; - for (auto const& llmReq : requestList) - { - // Don't add any tokens to simulate no output tokens - llmReq->addNewTokens({1}); - llmReq->setState(tb::LlmRequestState::kGENERATION_IN_PROGRESS); - } - callCount++; - })); - - EXPECT_CALL(*model, forwardSync()) - .WillRepeatedly(Invoke( - [&]() - { - for (auto const& llmReq : currentReq) - { - if (llmReq->getMaxNumGeneratedTokens() >= llmReq->mMaxNewTokens) - { - llmReq->setState(tb::LlmRequestState::kGENERATION_COMPLETE); - } - } - return; - })); - - EXPECT_CALL(*model, getMaxNumSequences()).WillRepeatedly(Invoke([&]() { return 10; })); - EXPECT_CALL(*model, getMaxInputLen()).WillRepeatedly(Invoke([&]() { return 6; })); - EXPECT_CALL(*model, getMaxSequenceLen()).WillRepeatedly(Invoke([&]() { return 10; })); - EXPECT_CALL(*model, getMaxDraftLen()).WillRepeatedly(Invoke([&]() { return 0; })); - EXPECT_CALL(*model, getVocabSizePadded()).WillRepeatedly(Invoke([&]() { return 80000; })); - tr::WorldConfig const dummyWorldConfig; - EXPECT_CALL(*model, getWorldConfig()) - .WillRepeatedly(Invoke([&]() -> tr::WorldConfig const& { return dummyWorldConfig; })); - EXPECT_CALL(*model, getCurrentIterationStats(_)).WillRepeatedly(Invoke([&](IterationStats& /*stats*/) { return; })); - EXPECT_CALL(*model, getCurrentRequestStats(_)) - .WillRepeatedly(Invoke([&](RequestStatsPerIteration& /*stats*/) { return; })); - - SizeType32 const beamWidth = 1; - ExecutorConfig executorConfig(beamWidth, - SchedulerConfig(CapacitySchedulerPolicy::kMAX_UTILIZATION)); // Condition 1 : MAX_UTILIZATION scheduling policy - executorConfig.setEnableChunkedContext(false); // Condition 2 : Chunked context disabled - executorConfig.setRequestStatsMaxIterations(1000); - auto executor = Executor(model, executorConfig); - - // Create the request - constexpr bool streaming = true; // Condition 3 : Streaming enabled - SizeType32 const maxNewTokens = 5; - VecTokens const tooLongInputTokens{1, 2, 3, 4, 5}; // Condition 4 : prompt input len + maxNewTokens > MaxInputLen - auto tooLongRequest = Request( - tooLongInputTokens, maxNewTokens, streaming, tensorrt_llm::executor::SamplingConfig(beamWidth), outConfig); - - // Enqueue the request - auto longRequestId = executor.enqueueRequest(tooLongRequest); - bool done = false; - int iter = 0; - while (!done && iter < mMaxWaitMs) - { - std::chrono::milliseconds waitTime(1); - auto responses = executor.awaitResponses(longRequestId, waitTime); - for (auto& response : responses) - { - EXPECT_EQ(response.hasError(), true); - EXPECT_THAT(response.getErrorMsg(), - testing::HasSubstr("sequence length is potentially greater than max input length")); - done = true; - } - ++iter; - } -} - -#if ENABLE_MULTI_DEVICE -// This test can be run manually to test multiGPU execution -// mpirun --allow-run-as-root -n 5 ./executorTest --gtest_filter="*MockedModelMultiGpu/ExecutorTest" -// Number of MPI ranks can be greater than tp - -TEST_P(ParamTest, MockedModelMultiGpu) -{ - auto const& world = tensorrt_llm::mpi::MpiComm::world(); - auto const worldRank = world.getRank(); - auto const worldSize = world.getSize(); - - // In this test, allow worldSize to be greater than tp = 4 - // If so, set participant ids to be the last 4 ranks - SizeType32 const tp = std::min(4, worldSize); - - using LlmRequestPtr = std::shared_ptr; - using RequestList = std::list; - - bool const streaming = std::get<0>(GetParam()); - bool const excludeInputFromOutput = std::get<1>(GetParam()); - auto const beamWidth = std::get<2>(GetParam()); - OutputConfig outConfig; - outConfig.excludeInputFromOutput = excludeInputFromOutput; - auto model = std::make_shared(); - - // Create the request - constexpr SizeType32 maxNewTokens = 5; - VecTokens const inputTokens{1, 2, 3, 4}; - auto request - = Request(inputTokens, maxNewTokens, streaming, tensorrt_llm::executor::SamplingConfig(beamWidth), outConfig); - - EXPECT_CALL(*model, terminateRequest(_, _)).Times(0); - EXPECT_CALL(*model, getVocabSizePadded()).Times(0); - EXPECT_CALL(*model, getLogitDataType()).Times(0); - tr::ModelConfig dummyModelConfig(0, 0, 0, 0, 1, 0, nvinfer1::DataType::kHALF); - EXPECT_CALL(*model, getModelConfig()) - .WillRepeatedly(Invoke([&]() -> tr::ModelConfig const& { return dummyModelConfig; })); - SizeType32 callCount = 0; - SizeType32 reqCallCount = 0; - EXPECT_CALL(*model, forwardAsync(_)) - .WillRepeatedly(Invoke( - [&](RequestList const& requestList) - { - for (auto const& llmReq : requestList) - { - EXPECT_EQ(llmReq->getTokens().size(), beamWidth); - // Verify that all MPI ranks get the expected request, even though only rank 0 actually gets the - // request - if (reqCallCount == 0) - { - EXPECT_EQ(llmReq->getOrigPromptLen(), request.getInputTokenIds().size()); - for (int i = 0; i < llmReq->getOrigPromptLen(); ++i) - { - EXPECT_EQ(llmReq->getTokens(beamWidth - 1).at(i), request.getInputTokenIds().at(i)); - } - } - EXPECT_EQ(llmReq->isStreaming(), request.getStreaming()); - EXPECT_EQ(llmReq->mMaxNewTokens, request.getMaxTokens()); - EXPECT_EQ( - llmReq->getTokens(beamWidth - 1).size(), request.getInputTokenIds().size() + reqCallCount); - - SizeType32 tokenId = 1; - COMM_SESSION.bcastValue(tokenId, 0); - // Don't add any tokens to simulate no output tokens - // Simulate leader rank communicating with comm session - VecTokens const newTokens(beamWidth, tokenId); - llmReq->addNewTokens(newTokens); - llmReq->setState(tb::LlmRequestState::kGENERATION_IN_PROGRESS); - if (llmReq->getMaxNumGeneratedTokens() >= llmReq->mMaxNewTokens) - { - llmReq->setState(tb::LlmRequestState::kGENERATION_COMPLETE); - } - reqCallCount++; - } - callCount++; - })); - - EXPECT_CALL(*model, getMaxNumSequences()).WillRepeatedly(Invoke([&]() { return 10; })); - EXPECT_CALL(*model, getMaxInputLen()).WillRepeatedly(Invoke([&]() { return 10; })); - EXPECT_CALL(*model, getMaxSequenceLen()).WillRepeatedly(Invoke([&]() { return 10; })); - EXPECT_CALL(*model, getVocabSizePadded()).WillRepeatedly(Invoke([&]() { return 80000; })); - EXPECT_CALL(*model, getCurrentIterationStats(_)).WillRepeatedly(Invoke([&](IterationStats& /*stats*/) { return; })); - EXPECT_CALL(*model, getCurrentRequestStats(_)) - .WillRepeatedly(Invoke([&](RequestStatsPerIteration& /*stats*/) { return; })); - - tr::WorldConfig dummyWorldConfig = tr::WorldConfig(tp, 1, 1, worldRank, tp); - EXPECT_CALL(*model, getWorldConfig()) - .WillRepeatedly(Invoke([&]() -> tr::WorldConfig const& { return dummyWorldConfig; })); - - ParallelConfig parallelConfig; - - // Set participant ids to be of size tp, starting at worldSize - 1 - std::vector participantIds; - participantIds.reserve(tp); - for (int i = 0; i < tp; ++i) - { - participantIds.push_back(worldSize - tp + i); - } - bool const isLeader = (worldRank == participantIds.front()); - parallelConfig.setParticipantIds(participantIds); - - bool const isWorker = (std::find(participantIds.begin(), participantIds.end(), worldRank) != participantIds.end()); - - // Set device ids - std::vector deviceIds(tp); - std::iota(deviceIds.begin(), deviceIds.end(), 0); - parallelConfig.setDeviceIds(deviceIds); - - ExecutorConfig executorConfig(beamWidth); - executorConfig.setParallelConfig(parallelConfig); - auto executor = Executor(model, executorConfig); - - EXPECT_EQ(isWorker, executor.isParticipant()); - - // Enqueue the request - IdType requestId = 0; - if (isLeader) - { - requestId = executor.enqueueRequest(request); - - SizeType32 numResponses{0}; - bool done = false; - int iter = 0; - while (!done && iter < mMaxWaitMs) - { - std::chrono::milliseconds waitTime(1); - auto responses = executor.awaitResponses(waitTime); - for (auto& response : responses) - { - ++numResponses; - auto const& result = response.getResult(); - EXPECT_EQ(result.outputTokenIds.size(), beamWidth); - auto expectedSize = streaming ? (beamWidth > 1 ? numResponses : 1) - : (maxNewTokens + (excludeInputFromOutput ? 0 : inputTokens.size())); - EXPECT_EQ(result.outputTokenIds.at(beamWidth - 1).size(), expectedSize); - done = result.isFinal; - } - ++iter; - } - - EXPECT_LT(iter, mMaxWaitMs); - EXPECT_EQ(numResponses, streaming ? maxNewTokens : 1); - EXPECT_EQ(callCount, maxNewTokens); - } -} -#endif // ENABLE_MULTI_DEVICE - -TEST_F(GptExecutorTest, MockedModelWithError) -{ - using LlmRequestPtr = std::shared_ptr; - using RequestList = std::list; - - struct MockedModelParams - { - SizeType32 maxInputLen; - SizeType32 maxSeqLen; - SizeType32 expectedTerminateCnt; - SizeType32 expectedForwardCnt; - bool computeGenLogits; - bool computeContextLogits; - std::string expectedError; - }; - - std::vector mockedModelParams; - // Mocked error in forward call - mockedModelParams.emplace_back(MockedModelParams{10, 20, 1, 1, true, true, "mocked error"}); - // prompt longer than maxInputLen - mockedModelParams.emplace_back(MockedModelParams{1, 20, 0, 0, true, true, "exceeds maximum input length"}); - // Model doesn't support context logits output - mockedModelParams.emplace_back( - MockedModelParams{10, 20, 0, 0, false, true, "gather_generation_logits must be enabled"}); - // Model doesn't support gen logits output - mockedModelParams.emplace_back( - MockedModelParams{10, 20, 0, 0, true, false, "need to build engine with gather_context"}); - - for (auto const& mockedModelParam : mockedModelParams) - { - auto model = std::make_shared(); - SizeType32 beamWidth = 1; - - // One request should be terminated - EXPECT_CALL(*model, terminateRequest(_, _)).Times(mockedModelParam.expectedTerminateCnt); - EXPECT_CALL(*model, getVocabSizePadded()).WillRepeatedly(Invoke([&]() { return 1024; })); - EXPECT_CALL(*model, getLogitDataType()).WillRepeatedly(Invoke([&]() { return nvinfer1::DataType::kFLOAT; })); - EXPECT_CALL(*model, getCurrentIterationStats(_)).WillRepeatedly(Invoke([&](IterationStats& stats) { return; })); - EXPECT_CALL(*model, getCurrentRequestStats(_)) - .WillRepeatedly(Invoke([&](RequestStatsPerIteration& stats) { return; })); - - SizeType32 callCount = 0; - EXPECT_CALL(*model, forwardAsync(_)) - .WillRepeatedly(Invoke( - [&](RequestList const&) - { - callCount++; - // There was a bug where we were missing a notify call when errors were encountered - // and this test was not catching it, probably because the error was reported - // before the first call to awaitResponses. So we add a sleep here to make sure - // the awaitResponses is called before the error is thrown - std::this_thread::sleep_for(std::chrono::milliseconds(200)); - throw std::runtime_error("mocked error"); - })); - - EXPECT_CALL(*model, getMaxNumSequences()).WillRepeatedly(Invoke([&]() { return 10; })); - EXPECT_CALL(*model, getMaxInputLen()).WillRepeatedly(Invoke([&]() { return mockedModelParam.maxInputLen; })); - EXPECT_CALL(*model, getMaxSequenceLen()).WillRepeatedly(Invoke([&]() { return mockedModelParam.maxSeqLen; })); - EXPECT_CALL(*model, getMaxDraftLen()).WillRepeatedly(Invoke([&]() { return 0; })); - tr::WorldConfig const dummyWorldConfig; - EXPECT_CALL(*model, getWorldConfig()) - .WillRepeatedly(Invoke([&]() -> tr::WorldConfig const& { return dummyWorldConfig; })); - tr::ModelConfig dummyModelConfig(0, 0, 0, 0, 1, 0, nvinfer1::DataType::kHALF); - dummyModelConfig.computeContextLogits(mockedModelParam.computeContextLogits); - dummyModelConfig.computeGenerationLogits(mockedModelParam.computeGenLogits); - EXPECT_CALL(*model, getModelConfig()) - .WillRepeatedly(Invoke([&]() -> tr::ModelConfig const& { return dummyModelConfig; })); - EXPECT_CALL(*model, getGatherGenerationLogits()) - .WillRepeatedly(Invoke([&]() -> bool { return mockedModelParam.computeGenLogits; })); - EXPECT_CALL(*model, getCurrentIterationStats(_)).WillRepeatedly(Invoke([&](IterationStats& stats) { return; })); - EXPECT_CALL(*model, getCurrentRequestStats(_)) - .WillRepeatedly(Invoke([&](RequestStatsPerIteration& stats) { return; })); - EXPECT_CALL(*model, getIterCounter()).WillRepeatedly(Invoke([&]() -> IterationType { return 0; })); - - ExecutorConfig executorConfig(beamWidth); - auto executor = Executor(model, executorConfig); - - // Create the request - SizeType32 maxNewTokens = 5; - VecTokens inputTokens{1, 2, 3, 4}; - - OutputConfig outConfig; - outConfig.returnContextLogits = true; - outConfig.returnGenerationLogits = true; - - auto streaming = false; - auto request = Request( - inputTokens, maxNewTokens, streaming, tensorrt_llm::executor::SamplingConfig(beamWidth), outConfig); - - // Enqueue the request - auto requestId = executor.enqueueRequest(std::move(request)); - - bool done = false; - auto responses = executor.awaitResponses(requestId); - for (auto& response : responses) - { - if (!response.hasError()) - { - FAIL() << "Expecting an error to be received"; - } - else - { - auto err = response.getErrorMsg(); - EXPECT_THAT(err, testing::HasSubstr(mockedModelParam.expectedError)); - done = true; - } - } - - EXPECT_TRUE(done); - EXPECT_EQ(callCount, mockedModelParam.expectedForwardCnt); - } -} - -TEST_F(GptExecutorTest, MockedModelCancelRequest) -{ - using LlmRequestPtr = std::shared_ptr; - using RequestList = std::list; - - constexpr bool streaming = true; - auto model = std::make_shared(); - - std::unordered_map reqIdsToTerminate; - // Two requests with one child request (3 in total) should be terminated - EXPECT_CALL(*model, terminateRequestSync(_, _)) - .Times(3) - .WillRepeatedly(Invoke([&](LlmRequestPtr const& llmRequest, FinishReason finishReason) - { reqIdsToTerminate.try_emplace(llmRequest->mRequestId, finishReason); })); - EXPECT_CALL(*model, terminateRequest(_, _)).Times(3); - EXPECT_CALL(*model, getVocabSizePadded()).Times(0); - EXPECT_CALL(*model, getLogitDataType()).Times(0); - tr::WorldConfig const dummyWorldConfig; - EXPECT_CALL(*model, getWorldConfig()) - .WillRepeatedly(Invoke([&]() -> tr::WorldConfig const& { return dummyWorldConfig; })); - EXPECT_CALL(*model, getCurrentIterationStats(_)).WillRepeatedly(Invoke([&](IterationStats& /*stats*/) { return; })); - EXPECT_CALL(*model, getCurrentRequestStats(_)) - .WillRepeatedly(Invoke([&](RequestStatsPerIteration& /*stats*/) { return; })); - tr::ModelConfig dummyModelConfig(0, 0, 0, 0, 1, 0, nvinfer1::DataType::kHALF); - EXPECT_CALL(*model, getModelConfig()) - .WillRepeatedly(Invoke([&]() -> tr::ModelConfig const& { return dummyModelConfig; })); - - SizeType32 callCount = 0; - std::unordered_map callCountPerSeq; - EXPECT_CALL(*model, forwardAsync(_)) - .WillRepeatedly(Invoke( - [&](RequestList const& requestList) - { - std::this_thread::sleep_for(std::chrono::milliseconds(10)); - - for (auto const& llmReq : requestList) - { - if (llmReq->isGenerationCompleteState()) - { - continue; - } - // Don't add any tokens to simulate no output tokens - llmReq->addNewTokens({1}); - llmReq->setState(tb::LlmRequestState::kGENERATION_IN_PROGRESS); - if (llmReq->getMaxNumGeneratedTokens() >= llmReq->mMaxNewTokens) - { - llmReq->setState(tb::LlmRequestState::kGENERATION_COMPLETE); - } - if (callCountPerSeq.find(llmReq->mRequestId) != callCountPerSeq.end()) - { - callCountPerSeq[llmReq->mRequestId]++; - } - else - { - callCountPerSeq[llmReq->mRequestId] = 1; - } - - if (reqIdsToTerminate.count(llmReq->mRequestId) != 0U) - { - if (!llmReq->isGenerationToCompleteState()) - { - model->terminateRequest(llmReq, false); - llmReq->finishByReason(reqIdsToTerminate[llmReq->mRequestId]); - llmReq->clearGeneratedTokens(); - } - reqIdsToTerminate.erase(llmReq->mRequestId); - } - } - callCount++; - })); - - EXPECT_CALL(*model, getMaxNumSequences()).WillRepeatedly(Invoke([&]() { return 10; })); - EXPECT_CALL(*model, getMaxInputLen()).WillRepeatedly(Invoke([&]() { return 100; })); - EXPECT_CALL(*model, getMaxSequenceLen()).WillRepeatedly(Invoke([&]() { return 200; })); - EXPECT_CALL(*model, getVocabSizePadded()).WillRepeatedly(Invoke([&]() { return 80000; })); - - SizeType32 const beamWidth = 1; - ExecutorConfig const executorConfig(beamWidth); - auto executor = Executor(model, executorConfig); - - // Create the request - SizeType32 const maxNewTokens = 150; - VecTokens const inputTokens{1, 2, 3, 4}; - auto request = Request(inputTokens, maxNewTokens, streaming); - - // Enqueue the request - auto requestId = executor.enqueueRequest(std::move(request)); - - // Cancel the request - std::this_thread::sleep_for(std::chrono::milliseconds(30)); - executor.cancelRequest(requestId); - - bool done = false; - int iter = 0; - while (!done && iter < mMaxWaitMs) - { - std::chrono::milliseconds waitTime(1); - auto responses = executor.awaitResponses(requestId, waitTime); - for (auto& response : responses) - { - - if (response.hasError()) - { - FAIL() << "Not expecting an error to be received"; - } - - auto const& result = response.getResult(); - done = result.isFinal; - if (done) - { - for (SizeType32 beamIdx = 0; beamIdx < beamWidth; ++beamIdx) - { - EXPECT_EQ(result.finishReasons[beamIdx], FinishReason::kCANCELLED); - } - } - } - ++iter; - } - - EXPECT_LT(iter, mMaxWaitMs); - // Expecting to receiving fewer tokens than maxNewTokens - EXPECT_LT(callCount, maxNewTokens); - - // Create the request having child requests. - auto samplingConfig2 = SamplingConfig(1); - samplingConfig2.setNumReturnSequences(2); - auto request2 = Request(inputTokens, maxNewTokens, streaming, samplingConfig2); - - // Reset call count. - callCount = 0; - callCountPerSeq.clear(); - - // Enqueue the request - auto requestId2 = executor.enqueueRequest(request2); - - // Cancel the request - std::this_thread::sleep_for(std::chrono::milliseconds(30)); - executor.cancelRequest(requestId2); - - done = false; - iter = 0; - while (!done && iter < mMaxWaitMs) - { - std::chrono::milliseconds waitTime(1); - auto responses = executor.awaitResponses(requestId2, waitTime); - for (auto& response : responses) - { - - if (response.hasError()) - { - FAIL() << "Not expecting an error to be received"; - } - - auto const& result = response.getResult(); - done = result.isFinal; - if (done) - { - EXPECT_EQ(result.finishReasons[0], FinishReason::kCANCELLED); - } - } - ++iter; - } - - EXPECT_LT(iter, mMaxWaitMs); - for (auto& [reqId, count] : callCountPerSeq) - { - // Expecting to receiving fewer tokens than maxNewTokens - EXPECT_LT(count, maxNewTokens) << "Failed at request id: " << reqId; - } -} - -TEST_F(GptExecutorTest, MockedModelNumReturns) -{ - using LlmRequestPtr = std::shared_ptr; - using RequestList = std::list; - - SizeType32 const maxBeamWidth = 4; - OutputConfig const outConfig; - auto model = std::make_shared(); - - EXPECT_CALL(*model, terminateRequest(_, _)).Times(0); - EXPECT_CALL(*model, getVocabSizePadded()).Times(0); - EXPECT_CALL(*model, getLogitDataType()).Times(0); - tr::WorldConfig const dummyWorldConfig; - EXPECT_CALL(*model, getWorldConfig()) - .WillRepeatedly(Invoke([&]() -> tr::WorldConfig const& { return dummyWorldConfig; })); - EXPECT_CALL(*model, getCurrentIterationStats(_)).WillRepeatedly(Invoke([&](IterationStats& /*stats*/) { return; })); - EXPECT_CALL(*model, getCurrentRequestStats(_)) - .WillRepeatedly(Invoke([&](RequestStatsPerIteration& /*stats*/) { return; })); - tr::ModelConfig dummyModelConfig(0, 0, 0, 0, 1, 0, nvinfer1::DataType::kHALF); - EXPECT_CALL(*model, getModelConfig()) - .WillRepeatedly(Invoke([&]() -> tr::ModelConfig const& { return dummyModelConfig; })); - SizeType32 callCount = 0; - EXPECT_CALL(*model, forwardAsync(_)) - .WillRepeatedly(Invoke( - [&](RequestList const& requestList) - { - for (auto const& llmReq : requestList) - { - // Don't add any tokens to simulate no output tokens - auto numBeams = llmReq->mSamplingConfig.getNumReturnBeams(); - llmReq->addNewTokens(VecTokens(numBeams, 1)); - llmReq->setState(tb::LlmRequestState::kGENERATION_IN_PROGRESS); - if (llmReq->getMaxNumGeneratedTokens() >= llmReq->mMaxNewTokens) - { - llmReq->setState(tb::LlmRequestState::kGENERATION_COMPLETE); - } - } - callCount++; - })); - - EXPECT_CALL(*model, getMaxNumSequences()).WillRepeatedly(Invoke([&]() { return 10; })); - EXPECT_CALL(*model, getMaxInputLen()).WillRepeatedly(Invoke([&]() { return 10; })); - EXPECT_CALL(*model, getMaxSequenceLen()).WillRepeatedly(Invoke([&]() { return 20; })); - EXPECT_CALL(*model, getVocabSizePadded()).WillRepeatedly(Invoke([&]() { return 80000; })); - - ExecutorConfig const executorConfig(maxBeamWidth); - auto executor = Executor(model, executorConfig); - - // Create the request - SizeType32 const maxNewTokens = 5; - VecTokens const inputTokens{1, 2, 3, 4}; - constexpr bool streaming = false; - - auto samplingConfig1 = SamplingConfig(1); - samplingConfig1.setNumReturnSequences(3); - auto request1 = Request(inputTokens, maxNewTokens, streaming, samplingConfig1, outConfig); - auto samplingConfig2 = SamplingConfig(4); - auto request2 = Request(inputTokens, maxNewTokens, streaming, samplingConfig2, outConfig); - auto samplingConfig3 = SamplingConfig(4); - samplingConfig3.setNumReturnSequences(2); - auto request3 = Request(inputTokens, maxNewTokens, streaming, samplingConfig3, outConfig); - - // Enqueue the request - auto requestId1 = executor.enqueueRequest(request1); - auto requestId2 = executor.enqueueRequest(request2); - auto requestId3 = executor.enqueueRequest(request3); - - // Expecting one response in beam search. Instead, numReturnSequences limits the number of beams to return. - std::unordered_map expectedNumResponses{{requestId1, 3}, {requestId2, 1}, {requestId3, 1}}; - std::unordered_map const expectedNumBeams{{requestId1, 1}, {requestId2, 4}, {requestId3, 2}}; - - std::unordered_map numResponses{{requestId1, 0}, {requestId2, 0}, {requestId3, 0}}; - std::unordered_map numBeams{{requestId1, 0}, {requestId2, 0}, {requestId3, 0}}; - int numFinished = 0; - int iter = 0; - while (numFinished < 3 && iter < mMaxWaitMs) - { - std::chrono::milliseconds waitTime(1); - auto responses = executor.awaitResponses(waitTime); - for (auto& response : responses) - { - auto const& result = response.getResult(); - auto reqId = response.getRequestId(); - numFinished += result.isFinal; - numResponses[reqId]++; - numBeams[reqId] = result.outputTokenIds.size(); - } - ++iter; - } - - EXPECT_LT(iter, mMaxWaitMs); - EXPECT_EQ(numFinished, 3); - for (auto& [reqId, numResp] : numResponses) - { - EXPECT_EQ(numResp, expectedNumResponses[reqId]); - } - for (auto& [reqId, numResp] : numResponses) - { - EXPECT_EQ(numResp, expectedNumResponses[reqId]); - } -} - -INSTANTIATE_TEST_SUITE_P(GptExecutorTest, ParamTest, - testing::Combine(testing::Values(false, true), // streaming - testing::Values(false, true), // excludeInputFromOutput - testing::Values(1, 2) // beamWidth - ), - generateTestName); diff --git a/cpp/tests/e2e_tests/executor/executorTest.cpp b/cpp/tests/e2e_tests/executor/executorTest.cpp deleted file mode 100644 index e1227970cb71..000000000000 --- a/cpp/tests/e2e_tests/executor/executorTest.cpp +++ /dev/null @@ -1,4671 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2023-2025 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. - */ - -#ifndef TOP_LEVEL_DIR -#error "Define TOP_LEVEL_DIR" -#endif - -#include "executorTest.h" - -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/logger.h" -#include "tensorrt_llm/common/memoryUtils.h" -#include "tensorrt_llm/executor/dataTransceiverState.h" -#include "tensorrt_llm/executor/requestWithId.h" -#include "tensorrt_llm/executor/types.h" -#include "tensorrt_llm/executor/version.h" -#include "tensorrt_llm/runtime/gptJsonConfig.h" -#include "tensorrt_llm/runtime/iBuffer.h" -#include "tensorrt_llm/runtime/iTensor.h" -#include "tensorrt_llm/runtime/tllmLogger.h" -#include "tensorrt_llm/runtime/utils/mpiUtils.h" -#include "tensorrt_llm/runtime/utils/numpyUtils.h" -#include "tensorrt_llm/testing/modelSpec.h" -#include "tests/utils/common.h" - -#include -#include -#include - -#include -#include -#include -#include -#include -#include -#include -#include - -namespace tr = tensorrt_llm::runtime; -namespace tc = tensorrt_llm::common; - -using namespace tensorrt_llm::testing; -using namespace tensorrt_llm::executor; -using namespace std::chrono_literals; -namespace fs = std::filesystem; -using tensorrt_llm::testing::KVCacheType; -using tensorrt_llm::testing::ModelSpec; - -namespace -{ - -auto const LORA_DATA_PATH = DATA_PATH / "lora-test-weights-gpt2-tp1"; -auto const LORA_WEIGHTS_FILE = LORA_DATA_PATH / "source.npy"; -auto const LORA_CONFIG_FILE = LORA_DATA_PATH / "config.npy"; - -auto constexpr LLAMA_INPUT_FILE = "input_tokens_llama.npy"; -auto constexpr LLAMA_VOCAB_SIZE_PADDED = 128256; -auto constexpr LLAMA_PAD_ID = 128001; -auto constexpr LLAMA_END_ID = 128001; - -} // namespace - -void testInvalidCtor(std::filesystem::path const& enginePath, ModelType modelType, ExecutorConfig executorConfig, - std::string expectedErrMsg = "") -{ - try - { - auto executor = Executor(enginePath, modelType, executorConfig); - - FAIL() << "Expected TllmException"; - } - catch (std::exception const& e) - { - EXPECT_THAT(e.what(), testing::HasSubstr(expectedErrMsg)); - } -} - -TEST_F(GptExecutorTest, version) -{ - EXPECT_STRNE(kTensorRtLlmVersion, "@TRTLLM_VERSION@"); - EXPECT_STREQ(kTensorRtLlmVersion, version()); -} - -TEST_F(GptExecutorTest, validCtor) -{ - SizeType32 beamWidth = 1; - auto executorConfig = ExecutorConfig(beamWidth); - auto trtEnginePath = (GPT_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp1-pp1-cp1-gpu"); - auto executor = Executor(trtEnginePath, ModelType::kDECODER_ONLY, executorConfig); -} - -TEST_F(GptExecutorTest, invalidCtor) -{ - SizeType32 beamWidth = 1; - auto executorConfig = ExecutorConfig(beamWidth); - auto trtEnginePath = GPT_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp1-pp1-cp1-gpu"; - std::filesystem::path invalidPath{"Bla"}; - - // Invalid path - { - testInvalidCtor(invalidPath, ModelType::kDECODER_ONLY, executorConfig, "File does not exist"); - } -} - -TEST_F(GptExecutorTest, enqueueAfterShutdown) -{ - SizeType32 beamWidth = 1; - auto executorConfig = ExecutorConfig(beamWidth); - auto trtEnginePath = (GPT_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp1-pp1-cp1-gpu"); - auto executor = Executor(trtEnginePath, ModelType::kDECODER_ONLY, executorConfig); - - SizeType32 maxNewTokens = 5; - VecTokens inputTokens{1, 2, 3, 4}; - auto request = Request(inputTokens, maxNewTokens, false, tensorrt_llm::executor::SamplingConfig(beamWidth)); - auto requestId = executor.enqueueRequest(request); - - bool done = false; - int iter = 0; - while (!done && iter < mMaxWaitMs) - { - std::chrono::milliseconds waitTime(1); - auto responses = executor.awaitResponses(requestId, waitTime); - for (auto& response : responses) - { - if (response.hasError()) - { - FAIL(); - } - else - { - done = response.getResult().isFinal; - } - } - ++iter; - } - EXPECT_LT(iter, mMaxWaitMs); - - executor.shutdown(); - - EXPECT_FALSE(executor.canEnqueueRequests()); - - std::string expErrMsg{"Shutdown called"}; - EXPECT_THAT([&]() { auto reqId = executor.enqueueRequest(request); }, - testing::Throws( - testing::Property(&tensorrt_llm::common::TllmException::what, testing::HasSubstr(expErrMsg)))); - EXPECT_THAT([&]() { auto resp = executor.awaitResponses(); }, - testing::Throws( - testing::Property(&tensorrt_llm::common::TllmException::what, testing::HasSubstr(expErrMsg)))); - EXPECT_THAT([&]() { auto stats = executor.getLatestIterationStats(); }, - testing::Throws( - testing::Property(&tensorrt_llm::common::TllmException::what, testing::HasSubstr(expErrMsg)))); - EXPECT_THAT([&]() { auto stats = executor.getLatestRequestStats(); }, - testing::Throws( - testing::Property(&tensorrt_llm::common::TllmException::what, testing::HasSubstr(expErrMsg)))); - EXPECT_THAT([&]() { executor.cancelRequest(requestId); }, - testing::Throws( - testing::Property(&tensorrt_llm::common::TllmException::what, testing::HasSubstr(expErrMsg)))); -} - -TEST_F(GptExecutorTest, missingPeftTask) -{ - SizeType32 beamWidth = 1; - auto executorConfig = ExecutorConfig(beamWidth); - auto trtEnginePath = GPT_MODEL_PATH / PathUtil::FP16_GPT_LORA_DIR() / "tp1-pp1-cp1-gpu"; - auto executor = Executor(trtEnginePath, ModelType::kDECODER_ONLY, executorConfig); - - // Create the request - SizeType32 maxNewTokens = 5; - VecTokens inputTokens{1, 2, 3, 4}; - auto request = Request(inputTokens, maxNewTokens, false, tensorrt_llm::executor::SamplingConfig(beamWidth)); - auto loraConfig = LoraConfig{10}; - request.setLoraConfig(loraConfig); - - auto requestId = executor.enqueueRequest(request); - - bool done = false; - std::chrono::milliseconds waitTime(mMaxWaitMs); - auto responses = executor.awaitResponses(requestId, waitTime); - for (auto& response : responses) - { - if (response.hasError()) - { - auto err = response.getErrorMsg(); - EXPECT_EQ(err, std::string("LoRA task 10 not found in cache. Please send LoRA weights with request")); - done = true; - } - else - { - FAIL() << "Expects error due to missing Lora weights"; - } - } - EXPECT_TRUE(done); -} - -TEST_F(GptExecutorTest, ReturnAcceptedTokenLogits) -{ - SizeType32 constexpr beamWidth{1}; - SizeType32 constexpr vocabSizePadded{50257}; // gpt vocabSizePadded - - // Create executor config - auto executorConfig = ExecutorConfig(beamWidth); - executorConfig.setGatherGenerationLogits(true); - - // Enable kv cache reuse of executorConfig - bool enableBlockReuse = true; - FloatType freeGpuMemoryFraction = 0.4; - auto kvCacheConfig - = KvCacheConfig(enableBlockReuse, std::nullopt, std::nullopt, std::nullopt, freeGpuMemoryFraction); - executorConfig.setKvCacheConfig(kvCacheConfig); - - // Create executor - auto trtEnginePath - = (GPT_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DRAFT_TOKENS_DIR() / "tp1-pp1-cp1-gpu"); - auto executor = Executor(trtEnginePath, ModelType::kDECODER_ONLY, executorConfig); - - // Create request - SizeType32 maxNewTokens = 5; - VecTokens inputTokens{1, 2, 3, 4, 5, 6, 7, 8}; - - std::vector streamingOptions{false, true}; - - for (auto streaming : streamingOptions) - { - auto request = Request(inputTokens, maxNewTokens, streaming, tensorrt_llm::executor::SamplingConfig(beamWidth)); - - // Set draft tokens - auto draftTokens = VecTokens{9, 10, 11, 12, 13}; // draft tokens - auto draftLength = draftTokens.size(); - FloatType const acceptanceThreshold = 0.00001f; // Ensure the draft token can be accepted - auto externalDraftTokensConfig = ExternalDraftTokensConfig(draftTokens, std::nullopt, acceptanceThreshold); - request.setExternalDraftTokensConfig(externalDraftTokensConfig); - - // Set return accepted token logits for this request - OutputConfig outConfig; - outConfig.returnGenerationLogits = true; - request.setOutputConfig(outConfig); - - // Enqueue this request - auto requestId = executor.enqueueRequest(request); - - bool done = false; - int iter = 0; - while (!done && iter < 5000) - { - std::chrono::milliseconds waitTime(mMaxWaitMs); - auto responses = executor.awaitResponses(requestId, waitTime); - for (auto& response : responses) - { - if (response.hasError()) - { - FAIL(); - } - else - { - auto result = response.getResult(); - done = result.isFinal; - auto& genLogits = result.generationLogits; - EXPECT_TRUE(genLogits.has_value()); - - // Expected shape: (1, numAcceptedDraftToken, vocabSizePadded) - auto const& acceptedTokenLogitsShape = genLogits->getShape(); - EXPECT_EQ(acceptedTokenLogitsShape.size(), 3); - EXPECT_EQ(acceptedTokenLogitsShape[0], 1); - EXPECT_LE(acceptedTokenLogitsShape[1], draftLength); // number of accepted tokens - EXPECT_EQ(acceptedTokenLogitsShape[2], vocabSizePadded); // vocabSizePadded - } - } - ++iter; - } - } -} - -TEST_F(GptExecutorTest, GenerationLogitsEarlyStop) -{ - SizeType32 constexpr beamWidth{1}; - SizeType32 constexpr vocabSizePadded{50257}; // gpt vocabSizePadded - auto constexpr streaming = false; - - ExtendedRuntimePerfKnobConfig perfKnobConfig = ExtendedRuntimePerfKnobConfig(); - - // Create executor config - auto executorConfig = ExecutorConfig(beamWidth); - executorConfig.setExtendedRuntimePerfKnobConfig(perfKnobConfig); - executorConfig.setGatherGenerationLogits(true); - - // Create executor - auto trtEnginePath = (GPT_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_GATHER_DIR() / "tp1-pp1-cp1-gpu"); - auto executor = Executor(trtEnginePath, ModelType::kDECODER_ONLY, executorConfig); - - auto const inputPath = DATA_PATH / "input_tokens.npy"; - ModelIds modelIds{50256, 50256}; - - auto manager = tr::BufferManager(std::make_shared()); - auto const& givenInput = tr::utils::loadNpy(manager, inputPath.string(), tr::MemoryType::kCPU); - auto [givenInputLengths, nbGivenInputs, maxInputLength] = getGivenInputLengths(*givenInput, modelIds.padId); - auto const* const givenInputData = tr::bufferCast(*givenInput); - - auto const& inputShape = givenInput->getShape(); - ASSERT_EQ(inputShape.nbDims, 2); - ASSERT_GT(inputShape.d[0], 0); - - BeamResult beamResult{beamWidth}; - auto const resultsPath - = GPT_DATA_PATH / ((beamWidth == 1) ? "sampling" : "beam_search_" + std::to_string(beamWidth)); - beamResult.resultsFile = resultsPath / PathUtil::FP16_PLUGIN_PACKED_PAGED_GATHER_RESULT_FILE(); - beamResult.contextLogitsFile = resultsPath / PathUtil::FP16_PLUGIN_PACKED_PAGED_CONTEXT_LOGITS_FILE(); - beamResult.genLogitsFile = resultsPath / PathUtil::FP16_PLUGIN_PACKED_PAGED_GENERATION_LOGITS_FILE(); - - // Set return generation logits for this request - OutputConfig outConfig; - outConfig.returnGenerationLogits = true; - outConfig.excludeInputFromOutput = true; - - // Load expected outputs for each beam width value - auto testData = TestData::loadTestData(beamResult, *givenInput, beamWidth, manager, outConfig, modelIds); - auto const maxSeqLen = testData.maxSeqLen; - - // Load expected outputs and inputs - std::vector requests; - std::vector reqMaxNewTokens; - - auto constexpr reqIdx = 0; - SizeType32 inputLen = givenInputLengths.at(reqIdx); - auto maxNewTokens = maxSeqLen - maxInputLength; - reqMaxNewTokens.push_back(maxNewTokens); - auto const* const seqBegin = givenInputData + reqIdx * maxInputLength; - - auto request = Request(VecTokens(seqBegin, seqBegin + inputLen), maxNewTokens, streaming, - tensorrt_llm::executor::SamplingConfig(beamWidth), outConfig, modelIds.endId); - // copy request - auto request2 = request; - - auto const expectedOutputData = tr::BufferRange(*testData.expectedOutputIds); - auto const expectedOutputLengths = testData.expectedOutputLengths; - auto const endPos = expectedOutputLengths[reqIdx] - 3; - auto const endIndex = tc::flat_index3(reqIdx, beamWidth - 1, endPos, beamWidth, maxSeqLen); - auto const endToken = expectedOutputData[endIndex]; - - // Set end id to stop early - request.setEndId(endToken); - requests.emplace_back(std::move(request)); - - // Set stop words to stop early - request2.setStopWords({{endToken}}); - requests.emplace_back(std::move(request2)); - - // Enqueue requests - auto requestIds = executor.enqueueRequests(requests); - - std::map expectedNewTokens; - expectedNewTokens[requestIds.at(0)] = endPos - inputLen; - expectedNewTokens[requestIds.at(1)] = endPos - inputLen + 1; - - std::map expectedFinishReason; - expectedFinishReason[requestIds.at(0)] = FinishReason::kEND_ID; - expectedFinishReason[requestIds.at(1)] = FinishReason::kSTOP_WORDS; - - std::map done; - std::for_each(requestIds.begin(), requestIds.end(), [&done](auto id) { done[id] = false; }); - int iter = 0; - while (!(std::all_of(done.begin(), done.end(), [](auto x) { return x.second; })) && iter < 5000) - { - std::chrono::milliseconds waitTime(mMaxWaitMs); - auto responses = executor.awaitResponses(waitTime); - for (auto& response : responses) - { - if (response.hasError()) - { - FAIL(); - } - else - { - auto const reqId = response.getRequestId(); - auto const& result = response.getResult(); - EXPECT_TRUE(result.isFinal); - done.at(reqId) = result.isFinal; - - // only 1 beam - auto const& outputIds = result.outputTokenIds.at(0); - EXPECT_EQ(outputIds.size(), expectedNewTokens.at(reqId)) << "req " << reqId; - - auto const& finishReason = result.finishReasons.at(0); - EXPECT_EQ(finishReason, expectedFinishReason.at(reqId)) << "req " << reqId; - - auto const& genLogits = result.generationLogits; - EXPECT_TRUE(genLogits.has_value()); - - // Expected shape: (1, numAcceptedDraftToken, vocabSizePadded) - auto const& generationLogitsShape = genLogits->getShape(); - EXPECT_EQ(generationLogitsShape.size(), 3); - EXPECT_EQ(generationLogitsShape[0], 1); - EXPECT_LE(generationLogitsShape[1], maxNewTokens); - EXPECT_EQ(generationLogitsShape[2], vocabSizePadded); - - auto const genLogitsTensor = detail::toITensor(*genLogits); - genLogitsTensor->squeeze(0); // only 1 beam - - for (size_t outputIdx = 0; outputIdx < expectedNewTokens.at(reqId); ++outputIdx) - { - // logits argmax should be equal to tokenId - auto const genLogitsSlice = tr::ITensor::slice(genLogitsTensor, outputIdx, 1); - auto const genLogitsRange = tr::BufferRange(*genLogitsSlice); - auto const* maxPos = std::max_element(genLogitsRange.begin(), genLogitsRange.end()); - auto const maxIdx = std::distance(genLogitsRange.begin(), maxPos); - - auto const tokenId = outputIds.at(outputIdx); - // Observed token mismatch at index 2 after building GPT engine with TRT builder optimization - // level 3. The testcase is sensitive to slight variation in kernel computation, so we skip checking - // for token id at index 2. - if (outputIdx != 2) - { - EXPECT_EQ(tokenId, maxIdx) << "req " << reqId << " outputIdx " << outputIdx; - } - } - } - } - ++iter; - } -} - -TEST_F(GptExecutorTest, GenerationChangeEndId) -{ - SizeType32 constexpr beamWidth{2}; - SizeType32 constexpr vocabSizePadded{50257}; // gpt vocabSizePadded - auto constexpr streaming = false; - - ExtendedRuntimePerfKnobConfig perfKnobConfig = ExtendedRuntimePerfKnobConfig(); - perfKnobConfig.setEnableContextFMHAFP32Acc(true); // use fmha fp32 acc for better accuracy - - // Create executor config - auto executorConfig = ExecutorConfig(beamWidth); - executorConfig.setExtendedRuntimePerfKnobConfig(perfKnobConfig); - - // Create executor - auto trtEnginePath = (GPT_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_GATHER_DIR() / "tp1-pp1-cp1-gpu"); - auto executor = Executor(trtEnginePath, ModelType::kDECODER_ONLY, executorConfig); - - auto const inputPath = DATA_PATH / "input_tokens.npy"; - ModelIds modelIds{50256, 50256}; - - auto manager = tr::BufferManager(std::make_shared()); - auto const& givenInput = tr::utils::loadNpy(manager, inputPath.string(), tr::MemoryType::kCPU); - auto [givenInputLengths, nbGivenInputs, maxInputLength] = getGivenInputLengths(*givenInput, modelIds.padId); - auto const* const givenInputData = tr::bufferCast(*givenInput); - - auto const& inputShape = givenInput->getShape(); - ASSERT_EQ(inputShape.nbDims, 2); - ASSERT_GT(inputShape.d[0], 0); - - BeamResult beamResult{beamWidth}; - auto const resultsPath - = GPT_DATA_PATH / ((beamWidth == 1) ? "sampling" : "beam_search_" + std::to_string(beamWidth)); - beamResult.resultsFile = resultsPath / PathUtil::FP16_PLUGIN_PACKED_PAGED_GATHER_CONTEXTFMHAFP32ACC_RESULT_FILE(); - - // Just return tokens for check - OutputConfig outConfig; - outConfig.excludeInputFromOutput = true; - - // Load expected outputs for each beam width value - auto testData = TestData::loadTestData(beamResult, *givenInput, beamWidth, manager, outConfig, modelIds); - auto const maxSeqLen = testData.maxSeqLen; - - // Load expected outputs and inputs - std::vector requests; - std::vector reqMaxNewTokens; - - // Only use the first request to test - auto constexpr reqIdx = 0; - SizeType32 inputLen = givenInputLengths.at(reqIdx); - auto maxNewTokens = maxSeqLen - maxInputLength; - reqMaxNewTokens.push_back(maxNewTokens); - auto const* const seqBegin = givenInputData + reqIdx * maxInputLength; - - // Use customized `EndId` to enqueue once - auto request = Request(VecTokens(seqBegin, seqBegin + inputLen), maxNewTokens, streaming, - tensorrt_llm::executor::SamplingConfig(beamWidth), outConfig, modelIds.endId); - - TokenIdType customizedEndId = *(seqBegin + 1); // Use a token appeared in ground-truth - request.setEndId(customizedEndId); - requests.emplace_back(std::move(request)); - - auto requestIds = executor.enqueueRequests(requests); - std::chrono::milliseconds waitTime(mMaxWaitMs); - auto responses = executor.awaitResponses(waitTime); - if (responses.at(0).hasError()) - { - FAIL(); - } - requests.clear(); - - // Change back to default `EndId` to enqueue again, and check the output - request = Request(VecTokens(seqBegin, seqBegin + inputLen), maxNewTokens, streaming, - tensorrt_llm::executor::SamplingConfig(beamWidth), outConfig, modelIds.endId); - - auto const expectedOutputData = tr::BufferRange(*testData.expectedOutputIds); - auto const expectedOutputLengths = testData.expectedOutputLengths; - auto const endPos = expectedOutputLengths[reqIdx]; - auto const endIndex = tc::flat_index3(reqIdx, beamWidth, endPos, beamWidth, maxSeqLen); - auto const endToken = expectedOutputData[endIndex]; - - request.setEndId(endToken); - requests.emplace_back(std::move(request)); - requestIds = executor.enqueueRequests(requests); - auto const requestId = requestIds.at(0); - - std::map expectedNewTokens; - expectedNewTokens[requestId] = endPos - inputLen; - - std::map expectedFinishReason; - expectedFinishReason[requestId] = FinishReason::kLENGTH; - - std::map done; - std::for_each(requestIds.begin(), requestIds.end(), [&done](auto id) { done[id] = false; }); - int iter = 0; - while (!(std::all_of(done.begin(), done.end(), [](auto x) { return x.second; })) && iter < 5000) - { - std::chrono::milliseconds waitTime(mMaxWaitMs); - auto responses = executor.awaitResponses(waitTime); - auto& response = responses.at(0); - if (response.hasError()) - { - FAIL(); - } - else - { - auto const reqId = response.getRequestId(); - auto const& result = response.getResult(); - EXPECT_TRUE(result.isFinal); - done.at(reqId) = result.isFinal; - - bool anyMismatch = false; - for (int i = 0; i < result.outputTokenIds.size(); ++i) - { - auto const& outputIds = result.outputTokenIds.at(i); - EXPECT_EQ(outputIds.size(), expectedNewTokens.at(reqId)) << "req " << reqId; - anyMismatch |= outputIds.size() != expectedNewTokens.at(reqId); - - auto const& finishReason = result.finishReasons.at(i); - EXPECT_EQ(finishReason, expectedFinishReason.at(reqId)) << "req " << reqId; - anyMismatch |= finishReason != expectedFinishReason.at(reqId); - - if (anyMismatch) - { - break; - } - - for (int j = 0; j < outputIds.size(); ++j) - { - auto const resultToken = outputIds[j]; - auto const groundTruthToken = expectedOutputData[maxSeqLen * i + inputLen + j]; - EXPECT_EQ(resultToken, groundTruthToken); - anyMismatch |= resultToken != groundTruthToken; - } - } - EXPECT_FALSE(anyMismatch); - } - ++iter; - } -} - -// stream, excludeInputFromOutput, beamWidth -using ParamType = std::tuple; -// useOrchestratorMode, beamWidth, modelName -using ParamCancelReqType = std::tuple; -// modelName -using LeaderApiUsageType = std::tuple; -// iterStatsMaxIterations, useOrchestratorMode -using ParamStatsType = std::tuple; -// streaming, beamWidth, computeLogProbs, excludeInputInOutput, returnContextLogits, returnGenerationLogits, modelName, -// useOrchestratorMode, returnAllGeneratedTokens, numReturnSequences -using AllParamsType = std::tuple; -// modelName, batched, replicated -using LogitsProcParamsType = std::tuple; -// modelName -using GuidedDecodingParamsType = std::tuple; -// modelName, useOrchestratorMode, beamWidth -using TimeoutTestParamsType = std::tuple; - -std::string generateTestName(testing::TestParamInfo const& info) -{ - auto const streaming = std::get<0>(info.param); - auto const excludeInputFromOutput = std::get<1>(info.param); - auto const beamWidth = std::get<2>(info.param); - std::string name = "ExecutorTest"; - if (streaming) - { - name += "Streaming"; - } - if (excludeInputFromOutput) - { - name += "ExclInput"; - } - name.append("BW" + std::to_string(beamWidth)); - return name; -} - -std::string generateTestNameCancelReq(testing::TestParamInfo const& info) -{ - auto const& useOrchestratorMode = std::get<0>(info.param); - auto const beamWidth = std::get<1>(info.param); - auto const modelName = std::get<2>(info.param); - std::string name = "ExecutorTest"; - name.append("BW" + std::to_string(beamWidth)); - name.append("_" + modelName + "_"); - - if (useOrchestratorMode) - { - name.append("OrchMode"); - } - else - { - name.append("LeaderMode"); - } - return name; -} - -std::string generateTestNameLeaderApiUsage(testing::TestParamInfo const& info) -{ - auto const modelName = std::get<0>(info.param); - std::string name = "ExecutorTest"; - name.append("_" + modelName); - return name; -} - -std::string generateTestNameLogitsProc(testing::TestParamInfo const& info) -{ - auto const modelName = std::get<0>(info.param); - bool const batched = std::get<1>(info.param); - bool const replicated = std::get<2>(info.param); - std::string name = "ExecutorTest"; - name.append("_" + modelName); - if (batched) - { - name.append("_Batched"); - } - if (replicated) - { - name.append("_Replicated"); - } - return name; -} - -std::string generateTestNameGuidedDecoding(testing::TestParamInfo const& info) -{ - auto const modelName = std::get<0>(info.param); - std::string name = "ExecutorTest"; - name.append("_" + modelName); - return name; -} - -std::string generateTestNameTimeoutTest(testing::TestParamInfo const& info) -{ - auto const modelName = std::get<0>(info.param); - auto const& useOrchestratorMode = std::get<1>(info.param); - auto const beamWidth = std::get<2>(info.param); - - std::string name = "ExecutorTest"; - name.append("_" + modelName); - - if (useOrchestratorMode) - { - name.append("_OrchMode"); - } - else - { - name.append("_LeaderMode"); - } - name.append("_BW" + std::to_string(beamWidth)); - return name; -} - -std::string generateTestNameStats(testing::TestParamInfo const& info) -{ - int iterStatsMaxIterations = std::get<0>(info.param); - auto const& useOrchestratorMode = std::get<1>(info.param); - std::string name = "ExecutorTest_"; - name.append(std::to_string(iterStatsMaxIterations) + "_"); - if (useOrchestratorMode) - { - name.append("OrchMode"); - } - else - { - name.append("LeaderMode"); - } - return name; -} - -std::string generateTestNameAllParams(testing::TestParamInfo const& info) -{ - auto const streaming = std::get<0>(info.param); - auto const& beamWidth = std::get<1>(info.param); - auto const& computeLogProbs = std::get<2>(info.param); - auto const& excludeInputInOutput = std::get<3>(info.param); - auto const& returnContextLogits = std::get<4>(info.param); - auto const& returnGenerationLogits = std::get<5>(info.param); - auto const modelName = std::get<6>(info.param); - auto const& useOrchestratorMode = std::get<7>(info.param); - auto const& returnAllGeneratedTokens = std::get<8>(info.param); - auto const& numReturnSequences = std::get<9>(info.param); - - std::string name = "ExecutorTest_"; - - if (streaming) - { - name += "Streaming"; - } - - name.append("_BW" + std::to_string(beamWidth)); - name.append("Nseq" + std::to_string(numReturnSequences)); - - if (computeLogProbs) - { - name.append("LogProbs"); - } - if (excludeInputInOutput) - { - name.append("ExcludeInput"); - } - if (returnContextLogits) - { - name.append("ContextLogits"); - } - if (returnGenerationLogits) - { - name.append("GenerationLogits"); - } - name.append("_" + modelName + "_"); - if (useOrchestratorMode) - { - name.append("OrchMode"); - } - else - { - name.append("LeaderMode"); - } - - if (returnAllGeneratedTokens) - { - name.append("returnAllGeneratedTokens"); - } - return name; -} - -class ParamTest : public GptExecutorTest, public ::testing::WithParamInterface -{ -}; - -class ParamStatsTest : public GptExecutorTest, public ::testing::WithParamInterface -{ -}; - -class AllParamsTest : public GptExecutorTest, public ::testing::WithParamInterface -{ -}; - -class ParamCancelReqTest : public GptExecutorTest, public ::testing::WithParamInterface -{ -}; - -class LeaderApiUsageTest : public GptExecutorTest, public ::testing::WithParamInterface -{ -}; - -class LogitsProcParamsTest : public GptExecutorTest, public ::testing::WithParamInterface -{ -}; - -class GuidedDecodingParamsTest : public GptExecutorTest, public ::testing::WithParamInterface -{ -}; - -class TimeoutTest : public GptExecutorTest, public ::testing::WithParamInterface -{ -}; - -TEST_F(GptExecutorTest, GetLatestStats) -{ - bool streaming = false; - bool excludeInputFromOutput = false; - OutputConfig outConfig; - outConfig.excludeInputFromOutput = excludeInputFromOutput; - - SizeType32 beamWidth = 1; - auto executorConfig = ExecutorConfig(beamWidth); - auto trtEnginePath = GPT_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp1-pp1-cp1-gpu"; - auto executor = Executor(trtEnginePath, ModelType::kDECODER_ONLY, executorConfig); - - // Create the request - SizeType32 maxNewTokens = 5; - VecTokens inputTokens{1, 2, 3, 4}; - auto request - = Request(inputTokens, maxNewTokens, streaming, tensorrt_llm::executor::SamplingConfig(beamWidth), outConfig); - auto requestId = executor.enqueueRequest(std::move(request)); - - bool done = false; - int iter = 0; - while (!done && iter < mMaxWaitMs) - { - std::chrono::milliseconds waitTime(1); - auto responses = executor.awaitResponses(requestId, waitTime); - for (auto& response : responses) - { - if (response.hasError()) - { - FAIL(); - } - else - { - done = response.getResult().isFinal; - } - } - ++iter; - } - EXPECT_LT(iter, mMaxWaitMs); - - // Expect 6 non-empty iterations - auto stats = executor.getLatestIterationStats(); - EXPECT_EQ(stats.size(), 6); - uint64_t currentIter = 0; - for (auto const& stat : stats) - { - EXPECT_EQ(stat.timestamp.size(), 26); - EXPECT_EQ(stat.iter, currentIter); - if (currentIter != 5) - { - EXPECT_EQ(stat.numActiveRequests, 1); - } - else - { - // For the last iteration the number of active requests - // should be zero. - EXPECT_EQ(stat.numActiveRequests, 0); - } - EXPECT_EQ(stat.maxNumActiveRequests, 64); - // Very loose check to make sure the memory stats are valid - EXPECT_GT(stat.gpuMemUsage, 16); - EXPECT_GT(stat.cpuMemUsage, 16); - EXPECT_GT(stat.pinnedMemUsage, 16); - - // Stats for KV cache - EXPECT_TRUE(stat.kvCacheStats.has_value()); - KvCacheStats const& kvStats = stat.kvCacheStats.value(); - EXPECT_GT(kvStats.maxNumBlocks, 0); - EXPECT_GT(kvStats.freeNumBlocks, 0); - EXPECT_EQ(kvStats.usedNumBlocks, currentIter == maxNewTokens ? 0 : 1); - EXPECT_GT(kvStats.tokensPerBlock, 0); - EXPECT_GT(kvStats.allocTotalBlocks, 0); - EXPECT_GT(kvStats.allocNewBlocks, 0); - EXPECT_GE(kvStats.reusedBlocks, 0); - EXPECT_GE(kvStats.missedBlocks, 0); - EXPECT_GE(kvStats.cacheHitRate, 0); - - // Stats for inflight batching - EXPECT_TRUE(stat.inflightBatchingStats.has_value() && !stat.staticBatchingStats.has_value()); - InflightBatchingStats const& modelStats = stat.inflightBatchingStats.value(); - EXPECT_EQ(modelStats.numScheduledRequests, currentIter == maxNewTokens ? 0 : 1); - EXPECT_EQ(modelStats.numContextRequests, currentIter == 0 ? 1 : 0); - EXPECT_EQ(modelStats.numGenRequests, currentIter == 0 || currentIter == maxNewTokens ? 0 : 1); - EXPECT_EQ(modelStats.numPausedRequests, 0); - EXPECT_EQ(modelStats.numCtxTokens, currentIter == 0 ? inputTokens.size() : 0); - EXPECT_EQ(modelStats.microBatchId, 0); - EXPECT_NEAR( - modelStats.avgNumDecodedTokensPerIter, currentIter == 0 || currentIter == maxNewTokens ? 0.f : 1.f, 1e-9f); - - auto jsonStr = JsonSerialization::toJsonStr(stat); - EXPECT_THAT(jsonStr, testing::HasSubstr("\"iter\":" + std::to_string(currentIter))); - EXPECT_THAT(jsonStr, testing::HasSubstr("\"staticBatchingStats\":null")); - EXPECT_THAT(jsonStr, testing::HasSubstr("\"numCtxTokens\":" + std::to_string(modelStats.numCtxTokens))); - EXPECT_THAT(jsonStr, testing::HasSubstr("\"numGenRequests\":" + std::to_string(modelStats.numGenRequests))); - - ++currentIter; - } -} - -TEST_F(GptExecutorTest, GetLatestStatsWithMultipleRequests) -{ - bool streaming = false; - bool excludeInputFromOutput = false; - OutputConfig outConfig; - outConfig.excludeInputFromOutput = excludeInputFromOutput; - - SizeType32 beamWidth = 1; - auto executorConfig = ExecutorConfig(beamWidth); - auto trtEnginePath = GPT_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp1-pp1-cp1-gpu"; - auto executor = Executor(trtEnginePath, ModelType::kDECODER_ONLY, executorConfig); - - // Create the requests - SizeType32 const numRequests = 2; - std::vector maxNewTokens{3, 5}; - std::vector inputTokens{{1, 2, 3, 4}, {5, 6, 7}}; - std::vector reqIds; - for (SizeType32 ireq = 0; ireq < numRequests; ++ireq) - { - auto request = Request(inputTokens[ireq], maxNewTokens[ireq], streaming, - tensorrt_llm::executor::SamplingConfig(beamWidth), outConfig); - auto requestId = executor.enqueueRequest(std::move(request)); - reqIds.emplace_back(requestId); - // sleep for 10 ms before sending the next request - std::this_thread::sleep_for(std::chrono::milliseconds(20)); - } - - for (SizeType32 ireq = 0; ireq < numRequests; ++ireq) - { - auto requestId = reqIds[ireq]; - bool done = false; - int iter = 0; - while (!done && iter < mMaxWaitMs) - { - std::chrono::milliseconds waitTime(1); - auto responses = executor.awaitResponses(requestId, waitTime); - for (auto& response : responses) - { - if (response.hasError()) - { - FAIL(); - } - else - { - done = response.getResult().isFinal; - } - } - ++iter; - } - EXPECT_LT(iter, mMaxWaitMs); - } - - // NOTES: - // Expect at least max(maxNewTokens) i.e. 5 non-empty iterations - // 4th iteration should have numCompletedRequests to be 1. - // Depending on the timing, first iteration will either have: - // 2 active requests - // or - // 1 active requests and 1 queued requests - auto stats = executor.getLatestIterationStats(); - EXPECT_GT(stats.size(), 0); // make sure we have at least 1 stat before the accessing 0-th element - if (stats[0].numActiveRequests == 2) - { - // we cannot reliably check queue latency since both started in the same iteration - // there should be exactly 5 non-empty iterations - EXPECT_EQ(stats.size(), 5); - // only check numCompletedRequests in 4th iteration - EXPECT_EQ(stats[3].numCompletedRequests, 1); - // 1st iteration shall record all 2 requests queueing time; - EXPECT_EQ(stats[0].numNewActiveRequests, 2); - // all rest iterations shall not return any queueing time; - for (int i = 1; i < stats.size(); ++i) - { - EXPECT_EQ(stats[i].numNewActiveRequests, 0); - } - } - else - { - // there should be more than 5 non-empty iterations since 2nd request started after 1st iteration - EXPECT_GT(stats.size(), 5); - // 1st request's completion is at 4th iteration - EXPECT_EQ(stats[3].numCompletedRequests, 1); - // 1st iteration record 1 request's queueing time; - EXPECT_EQ(stats[0].numNewActiveRequests, 1); - // the iteration where 2nd request became active, queue latency must be > 0 - uint64_t currentIter = 0; - for (auto const& stat : stats) - { - // To check when 2nd request becomes active, we need to think about 2 cases: - // - it overlaps with first request - // => only check queue time in this case - // - it doesn't overlap with the first request (e.g. 1st request ended too fast) - // => little to no queue time, cannot check reliably - // so we only check for queue time when numActiveRequests > 1 i.e. overlap happened after first iteration - if (stat.numActiveRequests > 1) - { - EXPECT_GT(currentIter, 0); // it must be after 1st iteration - EXPECT_GT(stat.newActiveRequestsQueueLatencyMS, 0); - // 2nd request record queueing time in this iteration - EXPECT_EQ(stat.numNewActiveRequests, 1); - break; - } - ++currentIter; - } - } -} - -TEST_F(GptExecutorTest, GetLatestRequestStats) -{ - bool streaming = false; - bool excludeInputFromOutput = false; - OutputConfig outConfig; - outConfig.excludeInputFromOutput = excludeInputFromOutput; - - SizeType32 beamWidth = 1; - auto executorConfig = ExecutorConfig(beamWidth); - executorConfig.setRequestStatsMaxIterations(1000); - executorConfig.setEnableChunkedContext(true); - auto trtEnginePath = GPT_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp1-pp1-cp1-gpu"; - auto executor = Executor(trtEnginePath, ModelType::kDECODER_ONLY, executorConfig); - - // Create the requests - std::vector> requestParams = { - // {maxNewTokens, inputTokens} - {5, {1, 2, 3, 4}}, {4, {1, 1, 2, 3, 5}}, {1, {1}}, - {8, VecTokens(383, 1)} // Long enough to be chunked into multiple iterations - }; - std::vector requests; - for (auto requestParam : requestParams) - { - requests.emplace_back(requestParam.second, requestParam.first, streaming, - tensorrt_llm::executor::SamplingConfig(beamWidth), outConfig); - } - auto requestIdsVec = executor.enqueueRequests(std::move(requests)); - std::map requestIdToIndex; - std::set activeRequests; - for (SizeType32 i = 0; i < requestIdsVec.size(); ++i) - { - auto requestId = requestIdsVec[i]; - activeRequests.insert(requestId); - requestIdToIndex[requestId] = i; - } - - int iter = 0; - while (!activeRequests.empty() && iter < mMaxWaitMs) - { - for (auto i = activeRequests.begin(); i != activeRequests.end();) - { - auto requestId = *i; - bool thisDone = false; - std::chrono::milliseconds waitTime(1); - auto responses = executor.awaitResponses(requestId, waitTime); - for (auto& response : responses) - { - if (response.hasError()) - { - // Allow response with error only if awaitResponse processed a terminated request id - std::string err = "ReqId " + std::to_string(response.getRequestId()) - + " has already been processed and was terminated."; - EXPECT_EQ(response.getErrorMsg(), err); - } - else - { - thisDone = response.getResult().isFinal; - } - } - if (thisDone) - { - // Erase completed request and move to the next one - i = activeRequests.erase(i); - } - else - { - ++i; - } - } - ++iter; - } - EXPECT_LT(iter, mMaxWaitMs); - - // Expect 5 non-empty iterations - // Note: The 6th iteration with the last finished request will be reported - // but might be unavailable when getLatestRequestStats is called since - // it could be updated after the final response has been sent. - auto stats = executor.getLatestRequestStats(); - EXPECT_GE(stats.size(), 5); - SizeType32 currentIter = 0; - auto invalidStart = std::numeric_limits::max(); - std::vector genStart(requestParams.size(), invalidStart); // The iteration index when generation started - std::set completedRequests; - for (auto stat = stats.begin(); stat != stats.begin() + 5; ++stat) - { - auto jsonStrIter = JsonSerialization::toJsonStr(*stat); - EXPECT_EQ(stat->iter, currentIter); - EXPECT_THAT(jsonStrIter, testing::HasSubstr("\"iter\":" + std::to_string(currentIter))); - EXPECT_EQ(stat->requestStats.size() + completedRequests.size(), requestParams.size()); - for (auto rStat : stat->requestStats) - { - auto jsonStr = JsonSerialization::toJsonStr(rStat); - // Only a few requests here so all of them should be scheduled. A separate test - // GetLatestRequestStatsScheduling will target the scheduling stats. - if (rStat.stage != RequestStage::kGENERATION_COMPLETE) - { - EXPECT_TRUE(rStat.scheduled); - EXPECT_THAT(jsonStr, testing::HasSubstr("\"scheduled\":true")); - } - EXPECT_TRUE(!rStat.paused); - EXPECT_THAT(jsonStr, testing::HasSubstr("\"paused\":false")); - EXPECT_TRUE(requestIdToIndex.count(rStat.id)); - EXPECT_THAT(jsonStr, testing::HasSubstr("\"id\":" + std::to_string(rStat.id))); - auto requestIndex = requestIdToIndex[rStat.id]; - auto contextSize = requestParams[requestIndex].second.size(); - if (rStat.contextPrefillPosition == contextSize) // Check generation phase - { - bool firstIteration{false}; - // Context phase is done - EXPECT_TRUE(rStat.stage == RequestStage::kGENERATION_IN_PROGRESS - || rStat.stage == RequestStage::kGENERATION_COMPLETE); - EXPECT_THAT(jsonStr, testing::HasSubstr("\"stage\":\"GENERATION")); - if (genStart[requestIndex] == invalidStart) - { - // Just started generation - genStart[requestIndex] = currentIter; - firstIteration = true; - } - - // One token per iteration - EXPECT_TRUE(currentIter - genStart[requestIndex] == rStat.numGeneratedTokens); - EXPECT_NEAR(rStat.avgNumDecodedTokensPerIter, firstIteration ? 0.f : 1.0f, 1e-9); - if (rStat.stage == RequestStage::kGENERATION_COMPLETE) - { - EXPECT_TRUE(requestParams[requestIndex].first >= rStat.numGeneratedTokens); - completedRequests.insert(requestIndex); - } - else - { - EXPECT_FALSE(completedRequests.count(requestIndex)); - } - } - else if (rStat.contextPrefillPosition < contextSize) // Check context phase - { - // Must be chunked - SizeType32 const maxChunkSize = 128; - EXPECT_TRUE(rStat.contextPrefillPosition % maxChunkSize == 0); - // Context phase is on-going - EXPECT_TRUE(rStat.stage == RequestStage::kCONTEXT_IN_PROGRESS); - // No tokens are generated - EXPECT_TRUE(0 == rStat.numGeneratedTokens); - } - else - { - FAIL() << "Out-of-boundary contextPrefillPosition in stats: " << rStat.contextPrefillPosition - << " out of " << contextSize; - } - // Sanity check that disaggregated serving stats is not set in typical use case - EXPECT_FALSE(rStat.disServingStats.has_value()); - } - ++currentIter; - } - // We should have visited all requests. - // Take into consideration the last request has not been reported - EXPECT_EQ(completedRequests.size() + 1, requestParams.size()); -} - -TEST_F(GptExecutorTest, GetLatestRequestStatsScheduling) -{ - // Specifically test the case where there are too many requests to be scheduled for a iteration - bool streaming = false; - bool excludeInputFromOutput = false; - OutputConfig outConfig; - outConfig.excludeInputFromOutput = excludeInputFromOutput; - - SizeType32 beamWidth = 1; - auto executorConfig = ExecutorConfig(beamWidth); - executorConfig.setRequestStatsMaxIterations(1000); - executorConfig.setEnableChunkedContext(true); - auto trtEnginePath = GPT_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp1-pp1-cp1-gpu"; - auto executor = Executor(trtEnginePath, ModelType::kDECODER_ONLY, executorConfig); - - // Create 100 requests. Note the max batch size for this model is 64 so some requests won't be scheduled right away. - std::vector> requestParams(100, {5, {1, 2, 3, 4}}); - std::vector requests; - requests.reserve(requestParams.size()); - for (auto requestParam : requestParams) - { - requests.emplace_back(requestParam.second, requestParam.first, streaming, - tensorrt_llm::executor::SamplingConfig(beamWidth), outConfig); - } - auto requestIdsVec = executor.enqueueRequests(std::move(requests)); - std::map requestIdToIndex; - std::set activeRequests; - for (SizeType32 i = 0; i < requestIdsVec.size(); ++i) - { - auto requestId = requestIdsVec[i]; - activeRequests.insert(requestId); - requestIdToIndex[requestId] = i; - } - - int iter = 0; - while (!activeRequests.empty() && iter < mMaxWaitMs) - { - for (auto i = activeRequests.begin(); i != activeRequests.end();) - { - auto requestId = *i; - bool thisDone = false; - std::chrono::milliseconds waitTime(1); - auto responses = executor.awaitResponses(requestId, waitTime); - for (auto& response : responses) - { - if (response.hasError()) - { - // Allow response with error only if awaitResponse processed a terminated request id - std::string err = "ReqId " + std::to_string(response.getRequestId()) - + " has already been processed and was terminated."; - EXPECT_EQ(response.getErrorMsg(), err); - } - else - { - thisDone = response.getResult().isFinal; - } - } - if (thisDone) - { - // Erase completed request and move to the next one - i = activeRequests.erase(i); - } - else - { - ++i; - } - } - ++iter; - } - EXPECT_LT(iter, mMaxWaitMs); - - auto stats = executor.getLatestRequestStats(); - SizeType32 numFinished = 0; - SizeType32 const maxActiveSize = 64; // Decided by the model - - // The 6th iteration request stat may or may not be available when getLatestRequestStats - // is called. When there are no other active or inTransmission requests, there will be - // another request stats to properly reset all the statistics to zero. - for (auto stat = stats.begin(); stat != stats.begin() + 5; ++stat) - { - SizeType32 numReqs = 0; - SizeType32 numReqsActive = 0; - SizeType32 numReqsQueued = 0; - SizeType32 numReqsJustDone = 0; - for (auto rStat : stat->requestStats) - { - ++numReqs; - numReqsActive += rStat.scheduled ? 1 : 0; - numReqsQueued += rStat.stage == RequestStage::kQUEUED ? 1 : 0; - numReqsJustDone += rStat.stage == RequestStage::kGENERATION_COMPLETE ? 1 : 0; - } - EXPECT_EQ(numReqs, numReqsActive + numReqsQueued + numReqsJustDone); - EXPECT_EQ(numReqs + numFinished, requestParams.size()); // Should report all unfinished requests - EXPECT_TRUE(numReqsActive <= maxActiveSize); // Not all requests are active due to max active size limit. - numFinished += numReqsJustDone; - } -} - -TEST_F(GptExecutorTest, GetRequestStatsMultipleRequests) -{ - bool streaming = false; - bool excludeInputFromOutput = false; - OutputConfig outConfig; - outConfig.excludeInputFromOutput = excludeInputFromOutput; - - SizeType32 beamWidth = 1; - auto executorConfig = ExecutorConfig(beamWidth); - executorConfig.setRequestStatsMaxIterations(1000); - auto trtEnginePath = GPT_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp1-pp1-cp1-gpu"; - auto executor = Executor(trtEnginePath, ModelType::kDECODER_ONLY, executorConfig); - - auto sendRequestWaitForResponseFn = [&]() - { - Request request({1, 2, 3}, 5); - auto requestId = executor.enqueueRequest(request); - bool isFinalResponse = false; - while (!isFinalResponse) - { - std::chrono::milliseconds waitTime(1); - auto responses = executor.awaitResponses(requestId, waitTime); - for (auto response : responses) - { - if (response.getResult().isFinal) - { - isFinalResponse = true; - break; - } - } - } - return requestId; - }; - - std::unordered_map requestIdToGenerationComplete; - auto updateStats = [&]() - { - auto stats = executor.getLatestRequestStats(); - for (auto& stat : stats) - { - for (auto const& request : stat.requestStats) - { - // only check and aggregate results when request is completed - if (request.stage == RequestStage::kGENERATION_COMPLETE) - { - requestIdToGenerationComplete[request.id] += 1; - } - } - } - }; - - auto requestId = sendRequestWaitForResponseFn(); - requestIdToGenerationComplete[requestId] = 0; - updateStats(); - - requestId = sendRequestWaitForResponseFn(); - requestIdToGenerationComplete[requestId] = 0; - updateStats(); - - for (auto [key, value] : requestIdToGenerationComplete) - { - EXPECT_EQ(value, 1); - } -} - -TEST_F(GptExecutorTest, BatchSizeTuning) -{ - bool streaming = false; - bool excludeInputFromOutput = false; - OutputConfig outConfig; - outConfig.excludeInputFromOutput = excludeInputFromOutput; - - SizeType32 beamWidth = 1; - auto executorConfig = ExecutorConfig(beamWidth); - executorConfig.setRequestStatsMaxIterations(1000); - executorConfig.setEnableChunkedContext(true); - - DynamicBatchConfig dynamicBatchConfig(true, false, 1); // Set window size to 1 - SchedulerConfig schedulerConfig(CapacitySchedulerPolicy::kGUARANTEED_NO_EVICT, std::nullopt, dynamicBatchConfig); - executorConfig.setSchedulerConfig(schedulerConfig); - - auto trtEnginePath = GPT_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp1-pp1-cp1-gpu"; - auto executor = Executor(trtEnginePath, ModelType::kDECODER_ONLY, executorConfig); - - std::vector tunerRecommendedBatchSizes; - - for (size_t i = 0; i <= 8; ++i) - { - auto inputLength = 1 << i; // Note that for this model max input len is 383 - Request request( - VecTokens(inputLength, 2), 5, streaming, tensorrt_llm::executor::SamplingConfig(beamWidth), outConfig); - auto requestId = executor.enqueueRequest(std::move(request)); - // Wait for current request to finish - while (true) - { - std::chrono::milliseconds waitTime(1); - auto responses = executor.awaitResponses(requestId, waitTime); - bool done = false; - if (responses.size() != 0) - { - EXPECT_TRUE(responses.size() == 1); - auto response = responses[0]; - EXPECT_FALSE(response.hasError()); - if (response.getResult().isFinal) - { - break; - } - } - } - auto reqStats = executor.getLatestIterationStats(); - EXPECT_TRUE(reqStats.size() > 0); - auto lastStat = reqStats.back(); - tunerRecommendedBatchSizes.push_back(lastStat.maxBatchSizeTunerRecommended); - } - - EXPECT_TRUE(tunerRecommendedBatchSizes.size() > 0); - // It's supposed to be decreasing when input length increases - EXPECT_TRUE(*tunerRecommendedBatchSizes.begin() > *tunerRecommendedBatchSizes.rbegin()); -} - -TEST_F(GptExecutorTest, GetLatestDebugTensors) -{ - bool streaming = false; - bool excludeInputFromOutput = false; - OutputConfig outConfig; - outConfig.excludeInputFromOutput = excludeInputFromOutput; - - SizeType32 maxNewTokens = 5; - - tensorrt_llm::executor::DebugConfig debugConfig; - debugConfig.setDebugTensorNames({{"sequence_length"}}); - debugConfig.setDebugTensorsMaxIterations(maxNewTokens); - - SizeType32 beamWidth = 1; - auto executorConfig = ExecutorConfig(beamWidth); - executorConfig.setDebugConfig(debugConfig); - - auto trtEnginePath = GPT_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp1-pp1-cp1-gpu"; - auto executor = Executor(trtEnginePath, ModelType::kDECODER_ONLY, executorConfig); - - // Create the request - VecTokens inputTokens{1, 2, 3, 4}; - auto request - = Request(inputTokens, maxNewTokens, streaming, tensorrt_llm::executor::SamplingConfig(beamWidth), outConfig); - auto requestId = executor.enqueueRequest(request); - - bool done = false; - int iter = 0; - while (!done && iter < mMaxWaitMs) - { - std::chrono::milliseconds waitTime(1); - auto responses = executor.awaitResponses(requestId, waitTime); - for (auto& response : responses) - { - if (response.hasError()) - { - FAIL(); - } - else - { - done = response.getResult().isFinal; - } - } - ++iter; - } - EXPECT_LT(iter, mMaxWaitMs); - - auto stream = std::make_shared(); - - // Expect 5 non-empty iterations - auto debugTensors = executor.getLatestDebugTensors(); - EXPECT_EQ(debugTensors.size(), 5); - uint64_t currentIter = 0; - for (auto const& debugIteration : debugTensors) - { - EXPECT_EQ(debugIteration.iter, currentIter); - EXPECT_EQ(debugIteration.debugTensors.size(), 2); - - { - auto it = debugIteration.debugTensors.find("request_ids"); - EXPECT_NE(it, debugIteration.debugTensors.end()); - auto const& tensor = it->second; - auto const& shape = tensor.getShape(); - EXPECT_EQ(shape.size(), 1); - EXPECT_EQ(shape[0], 1); - EXPECT_EQ(tensor.getSize(), 1); - auto const* dataPtr = static_cast(tensor.getData()); - EXPECT_EQ(dataPtr[0], 1) << "currentIter " << currentIter; - } - { - auto it = debugIteration.debugTensors.find("sequence_length"); - EXPECT_NE(it, debugIteration.debugTensors.end()); - auto const& tensor = it->second; - auto const& shape = tensor.getShape(); - EXPECT_EQ(shape.size(), 1); - EXPECT_EQ(tensor.getSize(), 1); - auto tensorHost = tensor.copyToCpu(stream); - auto const* dataPtr = static_cast(tensorHost.getData()); - EXPECT_EQ(dataPtr[0], inputTokens.size() + currentIter); - } - - ++currentIter; - } -} - -TEST_P(ParamTest, SingleRequestDemo) -{ - bool const streaming = std::get<0>(GetParam()); - bool const excludeInputFromOutput = std::get<1>(GetParam()); - auto const beamWidth = std::get<2>(GetParam()); - OutputConfig outConfig; - outConfig.excludeInputFromOutput = excludeInputFromOutput; - - auto executorConfig = ExecutorConfig(beamWidth); - auto trtEnginePath = GPT_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp1-pp1-cp1-gpu"; - auto executor = Executor(trtEnginePath, ModelType::kDECODER_ONLY, executorConfig); - - // Create the request - SizeType32 maxNewTokens = 5; - VecTokens inputTokens{1, 2, 3, 4}; - auto request - = Request(inputTokens, maxNewTokens, streaming, tensorrt_llm::executor::SamplingConfig(beamWidth), outConfig); - - // Enqueue the request - auto requestId = executor.enqueueRequest(request); - - // Get the new tokens - VecTokens tokens; - SizeType32 numResponses{0}; - bool done = false; - int iter = 0; - std::chrono::milliseconds waitTime(1); - while (!done && iter < mMaxWaitMs) - { - auto responses = executor.awaitResponses(requestId, waitTime); - for (auto& response : responses) - { - ++numResponses; - if (response.hasError()) - { - // This request failed for some reason, get error msg - std::string errStr - = "Request id " + std::to_string(requestId) + " failed with err " + response.getErrorMsg(); - FAIL(); - } - - auto result = response.getResult(); - done = result.isFinal; - auto& newTokens = result.outputTokenIds.at(beamWidth - 1); - auto const expectedSize = streaming ? (beamWidth > 1 ? numResponses : 1) - : (maxNewTokens + (excludeInputFromOutput ? 0 : inputTokens.size())); - EXPECT_EQ(newTokens.size(), expectedSize); - - if (streaming && beamWidth > 1) - { - // replace tokens - tokens = newTokens; - } - else - { - // Append tokens - tokens.insert(tokens.end(), newTokens.begin(), newTokens.end()); - } - } - ++iter; - } - EXPECT_LT(iter, mMaxWaitMs); - EXPECT_EQ(numResponses, streaming ? maxNewTokens : 1); - EXPECT_EQ( - tokens.size(), streaming ? maxNewTokens : (excludeInputFromOutput ? 0 : inputTokens.size()) + maxNewTokens); - - // Expect awaitResponse to return error message because the request is already terminated (isFinal = True) - auto response = executor.awaitResponses(requestId, waitTime).at(0); - EXPECT_TRUE(response.hasError()); - std::string err - = "ReqId " + std::to_string(response.getRequestId()) + " has already been processed and was terminated."; - EXPECT_EQ(response.getErrorMsg(), err); -} - -TEST_P(ParamTest, MultipleRequestDemo) -{ - bool const streaming = std::get<0>(GetParam()); - bool const excludeInputFromOutput = std::get<1>(GetParam()); - auto const beamWidth = std::get<2>(GetParam()); - OutputConfig outConfig; - outConfig.excludeInputFromOutput = excludeInputFromOutput; - SizeType32 numRequests = 20; - - auto executorConfig = ExecutorConfig(beamWidth); - auto trtEnginePath = GPT_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp1-pp1-cp1-gpu"; - auto executor = Executor(trtEnginePath, ModelType::kDECODER_ONLY, executorConfig); - - SizeType32 maxPromptLen = 20; - SizeType32 maxMaxNewTokens = 20; - - SizeType32 endId = -1; - // Enqueue the requests - std::unordered_map tokens; - std::unordered_map expectedNumTokens; - std::unordered_map expectedNumResponses; - for (SizeType32 req = 0; req < numRequests; ++req) - { - SizeType32 promptLen = rand() % maxPromptLen + 1; - SizeType32 maxNewTokens = rand() % maxMaxNewTokens + 1; - - auto request = Request(VecTokens(promptLen, 1), maxNewTokens, streaming, - tensorrt_llm::executor::SamplingConfig(beamWidth), outConfig, endId); - auto reqId = executor.enqueueRequest(std::move(request)); - tokens[reqId] = {}; - expectedNumTokens[reqId] = ((streaming || excludeInputFromOutput) ? 0 : promptLen) + maxNewTokens; - expectedNumResponses[reqId] = streaming ? maxNewTokens : 1; - } - - // Get the new tokens for each requests - int32_t numFinished = 0; - int iter = 0; - std::unordered_map numResponses; - while (numFinished < numRequests && iter < mMaxWaitMs) - { - std::chrono::milliseconds waitTime(1); - auto responses = executor.awaitResponses(waitTime); - for (auto& response : responses) - { - auto reqId = response.getRequestId(); - ++numResponses[reqId]; - if (!response.hasError()) - { - auto result = response.getResult(); - numFinished += result.isFinal; - auto& newTokens = result.outputTokenIds.at(beamWidth - 1); - auto const expectedSize - = streaming ? (beamWidth > 1 ? numResponses[reqId] : 1) : expectedNumTokens[reqId]; - EXPECT_EQ(newTokens.size(), expectedSize); - - auto& reqTokens = tokens.at(response.getRequestId()); - if (streaming && beamWidth > 1) - { - reqTokens = newTokens; - } - else - { - reqTokens.insert(reqTokens.end(), newTokens.begin(), newTokens.end()); - } - - for (SizeType32 b = 0; b < beamWidth; ++b) - { - EXPECT_EQ(result.finishReasons.at(b), - result.isFinal ? FinishReason::kLENGTH : FinishReason::kNOT_FINISHED); - } - } - else - { - // Allow response with error only if awaitResponse processed a terminated request id - std::string err = "ReqId " + std::to_string(response.getRequestId()) - + " has already been processed and was terminated."; - EXPECT_EQ(response.getErrorMsg(), err); - } - } - ++iter; - } - EXPECT_LT(iter, mMaxWaitMs); - - // Check that number of tokens matches expectations - for (auto const& [reqId, numTokens] : expectedNumTokens) - { - EXPECT_EQ(expectedNumResponses[reqId], numResponses[reqId]) << "reqId " << reqId; - EXPECT_EQ(expectedNumTokens[reqId], tokens[reqId].size()) << "reqId " << reqId; - } -} - -TEST_P(ParamStatsTest, MultipleRequestStats) -{ - bool streaming = false; - bool excludeInputFromOutput = false; - OutputConfig outConfig; - outConfig.excludeInputFromOutput = excludeInputFromOutput; - SizeType32 numRequests = 100; - auto iterStatsMaxIterations = std::get<0>(GetParam()); - bool useOrchestratorMode = std::get<1>(GetParam()); - - SizeType32 beamWidth = 1; - auto executorConfig = ExecutorConfig(beamWidth); - executorConfig.setIterStatsMaxIterations(iterStatsMaxIterations); - auto trtEnginePath = GPT_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp1-pp1-cp1-gpu"; - - std::optional orchestratorConfig = std::nullopt; - if (useOrchestratorMode) - { - orchestratorConfig = OrchestratorConfig(true, PathUtil::EXECUTOR_WORKER_PATH()); - } - auto parallelConfig = ParallelConfig(CommunicationType::kMPI, - useOrchestratorMode ? CommunicationMode::kORCHESTRATOR : CommunicationMode::kLEADER, std::nullopt, std::nullopt, - orchestratorConfig); - executorConfig.setParallelConfig(parallelConfig); - - auto executor = Executor(trtEnginePath, ModelType::kDECODER_ONLY, executorConfig); - - SizeType32 maxPromptLen = 20; - SizeType32 maxMaxNewTokens = 20; - - SizeType32 endId = -1; - // Enqueue the requests - std::unordered_map tokens; - std::unordered_map expectedNumTokens; - for (SizeType32 req = 0; req < numRequests; ++req) - { - SizeType32 promptLen = rand() % maxPromptLen + 1; - SizeType32 maxNewTokens = rand() % maxMaxNewTokens + 1; - - auto request = Request(VecTokens(promptLen, 1), maxNewTokens, streaming, - tensorrt_llm::executor::SamplingConfig(beamWidth), outConfig, endId); - auto reqId = executor.enqueueRequest(std::move(request)); - tokens[reqId] = {}; - expectedNumTokens[reqId] = (streaming ? 0 : (excludeInputFromOutput ? 0 : promptLen)) + maxNewTokens; - } - - std::atomic statsThreadDone = false; - std::atomic numFinished = 0; - std::deque iterStatsReceived; - // Spawn a thread that continuously get stats - auto statsThread = std::thread( - [&executor, &numFinished, numRequests, &iterStatsReceived, &statsThreadDone]() - { - while (numFinished < numRequests) - { - auto reqStats = executor.getLatestIterationStats(); - iterStatsReceived.insert(iterStatsReceived.end(), std::make_move_iterator(reqStats.begin()), - std::make_move_iterator(reqStats.end())); - std::this_thread::sleep_for(std::chrono::milliseconds(10)); - } - statsThreadDone = true; - }); - - // Get the new tokens for each requests - int iter = 0; - SizeType32 numResponses = 0; - while (numFinished < numRequests && iter < mMaxWaitMs) - { - std::chrono::milliseconds waitTime(1); - auto responses = executor.awaitResponses(waitTime); - for (auto& response : responses) - { - numResponses++; - if (!response.hasError()) - { - auto result = response.getResult(); - numFinished += result.isFinal; - auto& newTokens = result.outputTokenIds.at(beamWidth - 1); - auto& reqTokens = tokens.at(response.getRequestId()); - reqTokens.insert(reqTokens.end(), std::make_move_iterator(newTokens.begin()), - std::make_move_iterator(newTokens.end())); - } - else - { - // Allow response with error only if awaitResponse processed a terminated request id - std::string err = "ReqId " + std::to_string(response.getRequestId()) - + " has already been processed and was terminated."; - EXPECT_EQ(response.getErrorMsg(), err); - } - } - ++iter; - } - EXPECT_LT(iter, mMaxWaitMs); - - // Check that number of tokens matches expectations - for (auto const& [reqId, numTokens] : expectedNumTokens) - { - EXPECT_EQ(expectedNumTokens[reqId], tokens[reqId].size()) << "reqId " << reqId; - } - - // Wait for stats thread to be done, fail otherwise - iter = 0; - while (!statsThreadDone && iter < mMaxWaitMs) - { - std::chrono::milliseconds waitTime(1); - std::this_thread::sleep_for(std::chrono::milliseconds(waitTime)); - iter++; - } - ASSERT_TRUE(statsThreadDone); - if (iterStatsMaxIterations > 0) - { - ASSERT_GT(iterStatsReceived.size(), 1); - - for (auto stats : iterStatsReceived) - { - EXPECT_GT(stats.numActiveRequests, 0); - TLLM_LOG_INFO("%d %d", stats.iter, stats.numActiveRequests); - - EXPECT_TRUE(stats.inflightBatchingStats.has_value()); - if (stats.inflightBatchingStats.has_value()) - { - EXPECT_GT(stats.inflightBatchingStats.value().numScheduledRequests, 0); - } - } - } - - statsThread.join(); -} - -TEST_P(ParamTest, MultipleRequestBatchResponses) -{ - bool const streaming = std::get<0>(GetParam()); - bool const excludeInputFromOutput = std::get<1>(GetParam()); - auto const beamWidth = std::get<2>(GetParam()); - OutputConfig outConfig; - outConfig.excludeInputFromOutput = excludeInputFromOutput; - SizeType32 constexpr numRequests{20}; - - auto executorConfig = ExecutorConfig(beamWidth); - auto trtEnginePath = GPT_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp1-pp1-cp1-gpu"; - auto executor = Executor(trtEnginePath, ModelType::kDECODER_ONLY, executorConfig); - - SizeType32 constexpr maxPromptLen{20}; - SizeType32 constexpr maxMaxNewTokens{20}; - - SizeType32 endId = -1; - // Enqueue the requests - std::unordered_map tokens; - std::unordered_map expectedNumTokens; - std::vector requestIds; - for (SizeType32 req = 0; req < numRequests; ++req) - { - SizeType32 promptLen = rand() % maxPromptLen + 1; - SizeType32 maxNewTokens = rand() % maxMaxNewTokens + 1; - - auto request = Request(VecTokens(promptLen, 1), maxNewTokens, streaming, - tensorrt_llm::executor::SamplingConfig(beamWidth), outConfig, endId); - auto reqId = executor.enqueueRequest(std::move(request)); - requestIds.push_back(reqId); - tokens[reqId] = {}; - expectedNumTokens[reqId] = (streaming ? 0 : (excludeInputFromOutput ? 0 : promptLen)) + maxNewTokens; - } - - // Get the new tokens for each requests - int32_t numFinished = 0; - int iter = 0; - SizeType32 numResponses = 0; - std::chrono::milliseconds waitTime(1); - while (numFinished < numRequests && iter < mMaxWaitMs) - { - auto idResponses = executor.awaitResponses(requestIds, waitTime); - for (unsigned i = 0; i < requestIds.size(); ++i) - { - auto& responses = idResponses[i]; - for (auto& response : responses) - { - numResponses++; - if (!response.hasError()) - { - auto result = response.getResult(); - numFinished += result.isFinal; - auto& newTokens = result.outputTokenIds.at(beamWidth - 1); - auto& reqTokens = tokens.at(response.getRequestId()); - if (streaming && beamWidth > 1) - { - reqTokens = newTokens; - } - else - { - reqTokens.insert(reqTokens.end(), newTokens.begin(), newTokens.end()); - } - } - else - { - // Allow response with error only if awaitResponse processed a terminated request id - std::string err = "ReqId " + std::to_string(response.getRequestId()) - + " has already been processed and was terminated."; - EXPECT_EQ(response.getErrorMsg(), err); - } - } - } - ++iter; - } - EXPECT_LT(iter, mMaxWaitMs); - - // Rerun awaitResponses again and we expect to only see terminated request id error. - auto idResponses = executor.awaitResponses(requestIds, waitTime); - for (auto const& responses : idResponses) - { - for (auto& response : responses) - { - EXPECT_TRUE(response.hasError()); - std::string err = "ReqId " + std::to_string(response.getRequestId()) - + " has already been processed and was terminated."; - EXPECT_EQ(response.getErrorMsg(), err); - } - } - - // Check that number of tokens matches expectations - for (auto const& [reqId, numTokens] : expectedNumTokens) - { - EXPECT_EQ(expectedNumTokens[reqId], tokens[reqId].size()) << "reqId " << reqId; - } -} - -TEST_P(ParamTest, GetNumResponsesReadyTest) -{ - bool const streaming = std::get<0>(GetParam()); - bool const excludeInputFromOutput = std::get<1>(GetParam()); - auto const beamWidth = std::get<2>(GetParam()); - OutputConfig outConfig; - outConfig.excludeInputFromOutput = excludeInputFromOutput; - - auto executorConfig = ExecutorConfig(beamWidth); - auto trtEnginePath = GPT_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp1-pp1-cp1-gpu"; - auto executor = Executor(trtEnginePath, ModelType::kDECODER_ONLY, executorConfig); - - SizeType32 maxNumRequests = 50; - SizeType32 maxPromptLen = 20; - SizeType32 maxMaxNewTokens = 20; - - SizeType32 numRequests = rand() % maxNumRequests + 1; - SizeType32 numExpectedResponses = 0; - std::map reqNumExpectedResponses; - std::vector ids; - for (SizeType32 req = 0; req < numRequests; ++req) - { - SizeType32 promptLen = rand() % maxPromptLen + 1; - SizeType32 maxNewTokens = rand() % maxMaxNewTokens + 1; - - auto request = Request(VecTokens(promptLen, 1), maxNewTokens, streaming, - tensorrt_llm::executor::SamplingConfig(beamWidth), outConfig); - auto id = executor.enqueueRequest(std::move(request)); - ids.emplace_back(id); - reqNumExpectedResponses[id] = streaming ? maxNewTokens : 1; - numExpectedResponses += reqNumExpectedResponses.at(id); - } - - SizeType32 iter = 0; - SizeType32 numReady = 0; - while (numReady < numExpectedResponses && iter < mMaxWaitMs) - { - numReady = 0; - for (auto id : ids) - { - numReady += executor.getNumResponsesReady(id); - } - std::this_thread::sleep_for(std::chrono::milliseconds(1)); - ++iter; - } - EXPECT_LT(iter, mMaxWaitMs); - // Expect one response per request - for (auto id : ids) - { - SizeType32 numReady = executor.getNumResponsesReady(id); - EXPECT_EQ(numReady, reqNumExpectedResponses.at(id)); - } - auto numResponsesReady = executor.getNumResponsesReady(); - EXPECT_EQ(numResponsesReady, numExpectedResponses); -} - -namespace -{ - -void runTest(Executor& executor, fs::path const& inputPath, ModelIds const& modelIds, - FlakyTestInfo const& flakyTestInfo, bool streaming, SizeType32 const vocabSizePadded, BeamResult const& beamResult, - OutputConfig const& outConfig, bool isSpeculativeDecoding, int maxWaitMs, bool returnAllGeneratedTokens, - SizeType32 const numReturnSequences, bool isNonGreedySampling, SizeType32 const modelParallelism) -{ - auto const beamWidth = beamResult.beamWidth; - - auto manager = tr::BufferManager(std::make_shared()); - auto const& givenInput = tr::utils::loadNpy(manager, inputPath.string(), tr::MemoryType::kCPU); - auto [givenInputLengths, nbGivenInputs, maxInputLength] = getGivenInputLengths(*givenInput, modelIds.padId); - auto const* const givenInputData = tr::bufferCast(*givenInput); - - auto const& inputShape = givenInput->getShape(); - ASSERT_EQ(inputShape.nbDims, 2); - ASSERT_GT(inputShape.d[0], 0); - - // Load expected outputs for each beam width value - auto testData = TestData::loadTestData(beamResult, *givenInput, beamWidth, manager, outConfig, modelIds); - auto const maxSeqLen = testData.maxSeqLen; - - // Load expected outputs and inputs - SizeType32 numRequests = static_cast(givenInputLengths.size()); - SizeType32 maxRequests = numRequests; - std::vector requests; - std::vector reqMaxNewTokens; - - auto samplingConfig = tensorrt_llm::executor::SamplingConfig(beamWidth); - // top-k will be set by a large number to test non-identical N sequences. - if (isNonGreedySampling) - { - samplingConfig.setTopK(32); - } - samplingConfig.setNumReturnSequences(numReturnSequences); - - for (SizeType32 req = 0; req < maxRequests; ++req) - { - SizeType32 inputLen = givenInputLengths.at(req); - auto maxNewTokens = maxSeqLen - maxInputLength; - reqMaxNewTokens.push_back(maxNewTokens); - SizeType32 endId = -1; - auto const* const seqBegin = givenInputData + req * maxInputLength; - VecTokens tokens(seqBegin, seqBegin + inputLen); - auto request = Request( - VecTokens(seqBegin, seqBegin + inputLen), maxNewTokens, streaming, samplingConfig, outConfig, endId); - request.setReturnAllGeneratedTokens(returnAllGeneratedTokens); - requests.emplace_back(std::move(request)); - } - - auto& comm = tensorrt_llm::mpi::MpiComm::world(); - auto const worldRank = comm.getRank(); - - // Expected return sizes. - auto const numSequences = beamWidth > 1 ? 1 : numReturnSequences; - auto const numReturnBeams = std::min(beamWidth, numReturnSequences); - - if (worldRank == 0) - { - auto const reqIds = executor.enqueueRequests(requests); - - std::unordered_map> tokens; - std::unordered_map reqIdToBatchId; - - for (SizeType32 req = 0; req < reqIds.size(); ++req) - { - std::vector resultTokens(numSequences, BeamTokens(numReturnBeams)); - tokens[req] = std::move(resultTokens); - reqIdToBatchId[reqIds.at(req)] = req; - } - - // Get the new tokens for each requests - int32_t numFinished = 0; - int iter = 0; - std::unordered_map numResponses; - while (numFinished < maxRequests && iter < maxWaitMs) - { - std::chrono::milliseconds waitTime(1); - auto responses = executor.awaitResponses(waitTime); - for (auto& response : responses) - { - auto batchId = reqIdToBatchId.at(response.getRequestId()); - numResponses[batchId]++; - if (!response.hasError()) - { - auto result = response.getResult(); - numFinished += result.isFinal; - auto seqIdx = result.sequenceIndex; - - auto const& contextLogits = result.contextLogits; - auto const& genLogits = result.generationLogits; - auto const& outputTokenIds = result.outputTokenIds; - - EXPECT_EQ(result.finishReasons.size(), numReturnBeams); - for (SizeType32 beam = 0; beam < numReturnBeams; ++beam) - { - auto const& newTokens = outputTokenIds.at(beam); - auto& reqTokens = tokens.at(batchId).at(seqIdx).at(beam); - - if (!returnAllGeneratedTokens) - { - reqTokens.insert(reqTokens.end(), newTokens.begin(), newTokens.end()); - } - else - { - EXPECT_EQ(newTokens.size(), - (numResponses.at(batchId) + numReturnSequences - 1) / numReturnSequences); - reqTokens = newTokens; - } - // FinishReason is only supported for bw=1 and inflight batching. - if (beamWidth == 1) - { - EXPECT_EQ(result.finishReasons.at(beam), - result.isSequenceFinal ? FinishReason::kLENGTH : FinishReason::kNOT_FINISHED); - } - } - - auto const& cumLogProbs = result.cumLogProbs; - auto const& logProbs = result.logProbs; - auto const& beamTokens = tokens.at(batchId).at(seqIdx); - EXPECT_EQ(beamTokens.size(), numReturnBeams); - - if (!isNonGreedySampling) - { - float const logitsAtol = modelParallelism > 1 ? 1e-1 : 1e-2; - float const logitsRtol = modelParallelism > 1 ? 1e-2 : 1e-3; - - testData.verifyLogProbs(outConfig.returnLogProbs, streaming, outConfig.excludeInputFromOutput, - givenInputLengths.at(batchId), beamWidth, beamTokens, cumLogProbs, logProbs, batchId, - flakyTestInfo); - testData.validateContextLogits(outConfig.returnContextLogits, givenInputLengths.at(batchId), - beamWidth, contextLogits, vocabSizePadded, batchId, logitsAtol, logitsRtol); - testData.validateGenerationLogits(outConfig.returnGenerationLogits, result.isSequenceFinal, - streaming, outConfig.excludeInputFromOutput, givenInputLengths.at(batchId), - reqMaxNewTokens.at(batchId), beamWidth, beamTokens, genLogits, vocabSizePadded, batchId, - returnAllGeneratedTokens, logitsAtol, logitsRtol); - } - - // Ignore first iteration as it doesn't use draft tokens - if (outConfig.returnPerfMetrics && isSpeculativeDecoding - && result.requestPerfMetrics.value().iter > 0) - { - auto& specDecMetrics = result.requestPerfMetrics.value().speculativeDecoding; - // 4 draft tokens are used per step - EXPECT_EQ(specDecMetrics.totalDraftTokens, result.requestPerfMetrics.value().iter.value() * 4); - EXPECT_EQ(specDecMetrics.acceptanceRate, - static_cast(specDecMetrics.totalAcceptedDraftTokens) - / specDecMetrics.totalDraftTokens); - } - } - else - { - // Allow response with error only if awaitResponse processed a terminated request id - std::string err = "ReqId " + std::to_string(response.getRequestId()) - + " has already been processed and was terminated."; - EXPECT_EQ(response.getErrorMsg(), err); - } - } - ++iter; - } - EXPECT_LT(iter, maxWaitMs); - testData.verifyOutput(tokens, givenInputLengths, streaming, outConfig.excludeInputFromOutput, flakyTestInfo, - isSpeculativeDecoding, beamWidth, numSequences, isNonGreedySampling); - } -} - -void runTest(fs::path const& modelPath, ExecutorConfig const& executorConfig, fs::path const& inputPath, - ModelIds const& modelIds, FlakyTestInfo const& flakyTestInfo, bool streaming, SizeType32 const vocabSizePadded, - BeamResult const& beamResult, OutputConfig const& outConfig, bool isSpeculativeDecoding, int maxWaitMs, - bool returnAllGeneratedTokens, SizeType32 const numReturnSequences, bool isNonGreedySampling, - SizeType32 const modelParallelism) -{ - auto executor = Executor{modelPath, ModelType::kDECODER_ONLY, executorConfig}; - - runTest(executor, inputPath, modelIds, flakyTestInfo, streaming, vocabSizePadded, beamResult, outConfig, - isSpeculativeDecoding, maxWaitMs, returnAllGeneratedTokens, numReturnSequences, isNonGreedySampling, - modelParallelism); -} - -ExecutorConfig createExecutorConfig(SizeType32 maxBeamWidth, bool useOrchestratorMode, bool gatherGenerationLogits, - std::optional> deviceIds = std::nullopt, - std::optional> participantIds = std::nullopt) -{ - // Note: we reduce memory fraction for cases that return context/generation logits which require more free - // memory - FloatType constexpr freeGpuMemoryFraction{0.5F}; - KvCacheConfig kvCacheConfig(false, std::nullopt, std::nullopt, std::nullopt, freeGpuMemoryFraction); - auto executorConfig = ExecutorConfig(maxBeamWidth); - executorConfig.setKvCacheConfig(kvCacheConfig); - executorConfig.setNormalizeLogProbs(false); - executorConfig.setGatherGenerationLogits(gatherGenerationLogits); - - std::optional orchestratorConfig = std::nullopt; - if (useOrchestratorMode) - { - orchestratorConfig = OrchestratorConfig(true, PathUtil::EXECUTOR_WORKER_PATH()); - } - auto parallelConfig = ParallelConfig(CommunicationType::kMPI, - useOrchestratorMode ? CommunicationMode::kORCHESTRATOR : CommunicationMode::kLEADER, std::move(deviceIds), - std::move(participantIds), orchestratorConfig); - executorConfig.setParallelConfig(parallelConfig); - - return executorConfig; -} - -} // namespace - -TEST_P(AllParamsTest, TokenComparison) -{ - auto const streaming = std::get<0>(GetParam()); - auto const& beamWidth = std::get<1>(GetParam()); - OutputConfig outConfig; - outConfig.returnLogProbs = std::get<2>(GetParam()); - outConfig.excludeInputFromOutput = std::get<3>(GetParam()); - outConfig.returnContextLogits = std::get<4>(GetParam()); - outConfig.returnGenerationLogits = std::get<5>(GetParam()); - auto const modelName = std::get<6>(GetParam()); - auto const useOrchestratorMode = std::get<7>(GetParam()); - auto const returnAllGeneratedTokens = std::get<8>(GetParam()); - auto const numReturnSequences = std::get<9>(GetParam()); - if (returnAllGeneratedTokens && !streaming) - { - GTEST_SKIP() << "Test does not support returnAllGeneratedTokens without streaming"; - } - - std::optional> participantIds = std::nullopt; - - BeamResult beamResult{beamWidth}; - - ASSERT_TRUE(fs::exists(DATA_PATH)); - - fs::path modelPath; - // set defaults and adjust if needed by different models - fs::path inputPath = DATA_PATH / "input_tokens.npy"; - ModelIds modelIds{50256, 50256}; - bool isSpeculativeDecoding{false}; - - SizeType32 vocabSizePadded = 50257; - - // NOTE: This can be used to disable checks for certain prompt batch entries - FlakyTestInfo flakyTestInfo; - - if (modelName == "gpt") - { - auto const resultsPath - = GPT_DATA_PATH / ((beamWidth == 1) ? "sampling" : "beam_search_" + std::to_string(beamWidth)); - if (outConfig.returnContextLogits || outConfig.returnGenerationLogits) - { - modelPath = GPT_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_GATHER_DIR() / "tp1-pp1-cp1-gpu"; - beamResult.resultsFile = resultsPath / PathUtil::FP16_PLUGIN_PACKED_PAGED_GATHER_RESULT_FILE(); - beamResult.contextLogitsFile = resultsPath / PathUtil::FP16_PLUGIN_PACKED_PAGED_CONTEXT_LOGITS_FILE(); - beamResult.genLogitsFile = resultsPath / PathUtil::FP16_PLUGIN_PACKED_PAGED_GENERATION_LOGITS_FILE(); - if (outConfig.returnLogProbs) - { - beamResult.cumLogProbsFile - = resultsPath / PathUtil::FP16_PLUGIN_PACKED_PAGED_GATHER_CUM_LOG_PROBS_FILE(); - beamResult.logProbsFile = resultsPath / PathUtil::FP16_PLUGIN_PACKED_PAGED_GATHER_LOG_PROBS_FILE(); - } - } - else - { - modelPath = GPT_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp1-pp1-cp1-gpu"; - beamResult.resultsFile = resultsPath / PathUtil::FP16_PLUGIN_PACKED_PAGED_RESULT_FILE(); - if (outConfig.returnLogProbs) - { - beamResult.cumLogProbsFile = resultsPath / PathUtil::FP16_PLUGIN_PACKED_PAGED_CUM_LOG_PROBS_FILE(); - beamResult.logProbsFile = resultsPath / PathUtil::FP16_PLUGIN_PACKED_PAGED_LOG_PROBS_FILE(); - } - } - } - else if (modelName == "llama_tp4_pp1_cp1" || modelName == "llama_tp1_pp4_cp1" || modelName == "llama_tp2_pp2_cp1" - || modelName == "llama_tp1_pp2_cp1") - { - inputPath = DATA_PATH / LLAMA_INPUT_FILE; - modelIds.padId = LLAMA_PAD_ID; - modelIds.endId = LLAMA_END_ID; - - vocabSizePadded = LLAMA_VOCAB_SIZE_PADDED; - - auto const resultsPath - = LLAMA_DATA_PATH / ((beamWidth == 1) ? "sampling" : "beam_search_" + std::to_string(beamWidth)); - if (modelName == "llama_tp4_pp1_cp1") - { - beamResult.resultsFile = resultsPath / PathUtil::FP16_PLUGIN_PACKED_PAGED_RESULT_TP4_PP1_FILE(); - modelPath = LLAMA_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp4-pp1-cp1-gpu"; - } - else if (modelName == "llama_tp1_pp4_cp1") - { - beamResult.resultsFile = resultsPath / PathUtil::FP16_PLUGIN_PACKED_PAGED_RESULT_TP1_PP4_FILE(); - modelPath = LLAMA_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp1-pp4-cp1-gpu"; - } - else if (modelName == "llama_tp1_pp2_cp1") - { - beamResult.resultsFile = resultsPath / PathUtil::FP16_PLUGIN_PACKED_PAGED_RESULT_TP1_PP2_FILE(); - modelPath = LLAMA_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp1-pp2-cp1-gpu"; - } - else if (modelName == "llama_tp2_pp2_cp1") - { - beamResult.resultsFile = resultsPath / PathUtil::FP16_PLUGIN_PACKED_PAGED_RESULT_TP2_PP2_FILE(); - modelPath = LLAMA_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp2-pp2-cp1-gpu"; - } - beamResult.genLogitsFile = resultsPath / PathUtil::FP16_PLUGIN_PACKED_PAGED_GENERATION_LOGITS_TP4_PP1_FILE(); - if (outConfig.returnLogProbs) - { - beamResult.cumLogProbsFile = resultsPath / PathUtil::FP16_PLUGIN_PACKED_PAGED_CUM_LOG_PROBS_TP4_PP1_FILE(); - beamResult.logProbsFile = resultsPath / PathUtil::FP16_PLUGIN_PACKED_PAGED_LOG_PROBS_TP4_PP1_FILE(); - } - } - else if (modelName == "medusa") - { - TLLM_CHECK_WITH_INFO(beamWidth == 1, "Medusa does not support beam search."); - auto const resultsPath = MEDUSA_DATA_PATH / "sampling"; - auto modelSpec = ModelSpec::getDefaultModelSpec() - .useMedusa() - .setInputFile("input_tokens_long.npy") - .setMaxOutputLength(128); - beamResult.resultsFile = resultsPath / modelSpec.getResultsFile(); - modelPath = MEDUSA_MODEL_PATH / modelSpec.getModelPath() / "tp1-pp1-cp1-gpu"; - - inputPath = DATA_PATH / "input_vicuna.npy"; - modelIds.padId = 2; - modelIds.endId = 2; - isSpeculativeDecoding = true; - outConfig.returnPerfMetrics = true; - } - else if (modelName == "chatglm" || modelName == "chatglm2" || modelName == "chatglm3" || modelName == "glm") - { - fs::path resultsPath; - if (modelName == "chatglm") - { - resultsPath = CHATGLM_DATA_PATH; - modelPath = CHATGLM_MODEL_PATH; - } - else if (modelName == "chatglm2") - { - resultsPath = CHATGLM2_DATA_PATH; - modelPath = CHATGLM2_MODEL_PATH; - } - else if (modelName == "chatglm3") - { - resultsPath = CHATGLM3_DATA_PATH; - modelPath = CHATGLM3_MODEL_PATH; - } - else if (modelName == "glm") - { - resultsPath = GLM_DATA_PATH; - modelPath = GLM_MODEL_PATH; - } - resultsPath /= (beamWidth == 1) ? "sampling" : "beam_search_" + std::to_string(beamWidth); - beamResult.resultsFile = resultsPath / PathUtil::FP16_PLUGIN_PACKED_PAGED_RESULT_FILE(); - modelPath = modelPath / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp1-pp1-cp1-gpu"; - - char versionChatglm{0}; - if (size_t index = modelPath.string().find("chatglm"); index != std::string::npos) - { - versionChatglm = modelPath.string()[index + 7]; - std::string const vChatglmString - = (versionChatglm == '-') ? std::string("") : std::string(1, versionChatglm); - inputPath = DATA_PATH / ("input_tokens_chatglm" + vChatglmString + "-6b.npy"); - modelIds.padId = (versionChatglm == '-') ? 3 : 0; - modelIds.endId = (versionChatglm == '-') ? 130005 : 2; - } - else if (size_t index = modelPath.string().find("glm-10b"); index != std::string::npos) - { - inputPath = DATA_PATH / "input_tokens_glm-10b.npy"; - modelIds.padId = 50256; - modelIds.endId = 50258; - } - - if (versionChatglm != 0) - { - flakyTestInfo.batchIdBeams.insert(std::make_pair(1, 0)); - } - } - else - { - TLLM_THROW("Unrecognized modelName"); - } - - if (streaming && beamWidth > 1) - { - GTEST_SKIP() << "Test does not support streaming with beam search"; - } - - // Warning: This should be the last check before running the test. - // It will initialize MPI which can take significant time. - if (modelName == "llama_tp4_pp1_cp1" || modelName == "llama_tp1_pp4_cp1" || modelName == "llama_tp2_pp2_cp1" - || modelName == "llama_tp1_pp2_cp1") - { - // For llama model, only run for multiple GPUs - // This is detected by setting an env variable when running the test - char const* val = getenv("RUN_LLAMA_MULTI_GPU"); - if (val == nullptr) - { - GTEST_SKIP() << "Skipping Llama test"; - } - - if (outConfig.returnContextLogits) - { - GTEST_SKIP() << "Skipping context logits tests for mpi runs"; - } - - // Check that it was launched with right number of MPI ranks - if (!useOrchestratorMode && COMM_SESSION.getSize() != 4) - { - // No orchestrator, need worldSize to match TP*PP - FAIL() << "Leader mode and world size is not equal to 4"; - } - if (useOrchestratorMode && COMM_SESSION.getSize() != 1) - { - // No orchestrator, need worldSize to match TP*PP - FAIL() << "Orchestrator mode and World size is not equal to 1"; - } - } - auto decoderJsonConfig = tensorrt_llm::runtime::GptJsonConfig::parse(modelPath / "config.json"); - - auto const modelTP = decoderJsonConfig.getTensorParallelism(); - auto const modelPP = decoderJsonConfig.getPipelineParallelism(); - auto const modelParallelism = modelTP * modelPP; - int deviceCount = -1; - TLLM_CUDA_CHECK(cudaGetDeviceCount(&deviceCount)); - std::optional> deviceIds = std::vector(modelParallelism); - for (auto i = 0; i < deviceIds->size(); i++) - { - deviceIds->at(i) = i % deviceCount; - } - if (modelName == "llama_tp1_pp2_cp1") - { - auto const& session = tensorrt_llm::mpi::MpiComm::world(); - if (session.getSize() != 4) - { - FAIL() << "Llama-tp1-pp2 is intended solely for testing coexisting engines within the same MPI world," - " which requires a session size of 4. However, the current session size is " - << session.getSize() << " ."; - } - if (session.getRank() / 2 == 0) - { - participantIds = std::vector{0, 1}; - deviceIds = std::vector{0, 1}; - } - else - { - participantIds = std::vector{2, 3}; - deviceIds = std::vector{2, 3}; - } - } - - if (modelPP > 1) - { - std::reverse(deviceIds->begin(), deviceIds->end()); - if (modelTP > 1) - { - for (SizeType32 ppRank = 0; ppRank < modelPP; ppRank++) - { - std::reverse(deviceIds->begin() + ppRank * modelTP, deviceIds->begin() + (ppRank + 1) * modelPP); - } - } - } - - // Returning logits will bring higher latency - if (streaming && (outConfig.returnContextLogits || outConfig.returnGenerationLogits)) - { - mMaxWaitMs = 20000; - } - - auto executorConfig = createExecutorConfig(beamWidth, useOrchestratorMode, outConfig.returnGenerationLogits, - std::move(deviceIds), std::move(participantIds)); - - runTest(modelPath, executorConfig, inputPath, modelIds, flakyTestInfo, streaming, vocabSizePadded, beamResult, - outConfig, isSpeculativeDecoding, mMaxWaitMs, returnAllGeneratedTokens, numReturnSequences, false, - modelParallelism); -} - -TEST_F(GptExecutorTest, ChangeBeamWidth) -{ - SizeType32 constexpr maxBeamWidth{2}; - auto executorConfig = ExecutorConfig(maxBeamWidth); - - auto trtEnginePath = (GPT_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp1-pp1-cp1-gpu"); - auto executor = Executor(trtEnginePath, ModelType::kDECODER_ONLY, executorConfig); - - SizeType32 constexpr beamWidth1{1}; - SizeType32 constexpr beamWidth2{2}; - SizeType32 constexpr maxNewTokens{2}; - VecTokens inputTokens{1, 2, 3, 4}; - - // Create requests with different beam widths - std::vector requests; - requests.emplace_back(inputTokens, maxNewTokens, false, tensorrt_llm::executor::SamplingConfig(beamWidth1)); - requests.emplace_back(inputTokens, maxNewTokens, false, tensorrt_llm::executor::SamplingConfig(beamWidth1)); - requests.emplace_back(inputTokens, maxNewTokens, false, tensorrt_llm::executor::SamplingConfig(beamWidth2)); - requests.emplace_back(inputTokens, maxNewTokens, false, tensorrt_llm::executor::SamplingConfig(beamWidth1)); - - auto requestIds = executor.enqueueRequests(requests); - - int numFinished = 0; - int iter = 0; - while (numFinished < 4 && iter < mMaxWaitMs) - { - std::chrono::milliseconds waitTime(1); - auto responses = executor.awaitResponses(waitTime); - for (auto& response : responses) - { - if (response.hasError()) - { - auto err = response.getErrorMsg(); - std::cout << "err:" << err << std::endl; - FAIL() << "Should not get a response with error"; - } - else - { - auto result = response.getResult(); - numFinished += static_cast(result.isFinal); - } - } - ++iter; - } - EXPECT_LT(iter, mMaxWaitMs); - - auto stats = executor.getLatestIterationStats(); - uint64_t currentIter = 0; - for (auto const& stat : stats) - { - // TODO: enable this check when stats are cleaned - // EXPECT_EQ(stat.iter, currentIter); - if (stat.iter < 2) - { - // req 1 and 2 run with same beam width - EXPECT_EQ(stat.numActiveRequests, 2); - } - else if (stat.numActiveRequests != 0) // TODO: remove this check when stats are cleaned - { - // req 3 or 4 run width different beam width - EXPECT_EQ(stat.numActiveRequests, 1); - } - - ++currentIter; - } -} - -void doTokenComparisonChangeBeamWidth(bool enableReuse, SizeType32 maxWaitMs) -{ - SizeType32 constexpr maxBeamWidth{2}; - SizeType32 constexpr vocabSizePadded{50257}; // gpt vocabSizePadded - auto constexpr streaming = false; - - // Create executor config - auto kvCacheConfig = KvCacheConfig(enableReuse); - auto executorConfig = ExecutorConfig(maxBeamWidth, SchedulerConfig(), kvCacheConfig); - - // Create executor - auto trtEnginePath = (GPT_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_GATHER_DIR() / "tp1-pp1-cp1-gpu"); - auto executor = Executor(trtEnginePath, ModelType::kDECODER_ONLY, executorConfig); - - auto const inputPath = DATA_PATH / "input_tokens.npy"; - ModelIds modelIds{50256, 50256}; - - OutputConfig outConfig; - FlakyTestInfo flakyTestInfo; - bool constexpr isSpeculativeDecoding{false}; - - for (SizeType32 beamWidth : {1, 2}) - { - TLLM_LOG_INFO("Running beam width: %d", beamWidth); - BeamResult beamResult{beamWidth}; - auto const resultsPath - = GPT_DATA_PATH / ((beamWidth == 1) ? "sampling" : "beam_search_" + std::to_string(beamWidth)); - beamResult.resultsFile = resultsPath / PathUtil::FP16_PLUGIN_PACKED_PAGED_GATHER_RESULT_FILE(); - beamResult.contextLogitsFile = resultsPath / PathUtil::FP16_PLUGIN_PACKED_PAGED_CONTEXT_LOGITS_FILE(); - beamResult.genLogitsFile = resultsPath / PathUtil::FP16_PLUGIN_PACKED_PAGED_GENERATION_LOGITS_FILE(); - - auto const numReturnSequences = beamWidth; - - runTest(executor, inputPath, modelIds, flakyTestInfo, streaming, vocabSizePadded, beamResult, outConfig, - isSpeculativeDecoding, maxWaitMs, false, numReturnSequences, false, 1); - } -} - -TEST_F(GptExecutorTest, TokenComparisonChangeBeamWidth) -{ - doTokenComparisonChangeBeamWidth(false, mMaxWaitMs); -} - -TEST_F(GptExecutorTest, TokenComparisonChangeBeamWidthBlockReuse) -{ - doTokenComparisonChangeBeamWidth(true, mMaxWaitMs); -} - -TEST_F(GptExecutorTest, NReturnRandomness) -{ - SizeType32 constexpr maxBeamWidth{1}; - SizeType32 constexpr numReturnSequences{2}; - SizeType32 constexpr vocabSizePadded{50257}; // gpt vocabSizePadded - auto constexpr streaming = false; - - // Create executor config - auto executorConfig = ExecutorConfig(maxBeamWidth); - - // Create executor - auto trtEnginePath = (GPT_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_GATHER_DIR() / "tp1-pp1-cp1-gpu"); - auto executor = Executor(trtEnginePath, ModelType::kDECODER_ONLY, executorConfig); - - auto const inputPath = DATA_PATH / "input_tokens.npy"; - ModelIds modelIds{50256, 50256}; - - OutputConfig outConfig; - FlakyTestInfo flakyTestInfo; - bool constexpr isSpeculativeDecoding{false}; - - BeamResult beamResult{maxBeamWidth}; - auto const resultsPath = GPT_DATA_PATH / "sampling"; - beamResult.resultsFile = resultsPath / PathUtil::FP16_PLUGIN_PACKED_PAGED_GATHER_RESULT_FILE(); - beamResult.contextLogitsFile = resultsPath / PathUtil::FP16_PLUGIN_PACKED_PAGED_CONTEXT_LOGITS_FILE(); - beamResult.genLogitsFile = resultsPath / PathUtil::FP16_PLUGIN_PACKED_PAGED_GENERATION_LOGITS_FILE(); - - runTest(executor, inputPath, modelIds, flakyTestInfo, streaming, vocabSizePadded, beamResult, outConfig, - isSpeculativeDecoding, mMaxWaitMs, false, 1, true, 1); -} - -TEST_F(GptExecutorTest, TimedOut) -{ - SizeType32 beamWidth = 1; - auto executorConfig = ExecutorConfig(beamWidth); - auto trtEnginePath = GPT_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp1-pp1-cp1-gpu"; - auto executor = Executor(trtEnginePath, ModelType::kDECODER_ONLY, executorConfig); - - // No requests enqueued, expect no responses - auto numResponsesReady = executor.getNumResponsesReady(); - EXPECT_EQ(numResponsesReady, 0); - - std::chrono::milliseconds waitTime(10); - auto responses = executor.awaitResponses(waitTime); - EXPECT_EQ(responses.size(), 0); -} - -TEST_F(GptExecutorTest, MaxSeqIdleMicrosecondsError) -{ - auto executorConfig = ExecutorConfig(1); - // Request will time out - executorConfig.setMaxSeqIdleMicroseconds(1); - auto trtEnginePath = (GPT_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp1-pp1-cp1-gpu"); - auto executor = Executor(trtEnginePath, ModelType::kDECODER_ONLY, executorConfig); - - SizeType32 constexpr maxNewTokens{5}; - VecTokens inputTokens{1, 2, 3, 4}; - - std::vector requests; - requests.emplace_back(inputTokens, maxNewTokens, false); - - auto requestIds = executor.enqueueRequests(requests); - - bool done = false; - int iter = 0; - while (!done && iter < mMaxWaitMs) - { - std::chrono::milliseconds waitTime(1); - auto responses = executor.awaitResponses(waitTime); - for (auto& response : responses) - { - if (response.hasError()) - { - auto err = response.getErrorMsg(); - std::cout << "err:" << err << std::endl; - EXPECT_THAT(err, testing::HasSubstr("Unable to get batch slot for request ID")); - done = true; - } - else - { - FAIL() << "Should get a response with error"; - } - } - ++iter; - } - EXPECT_LT(iter, mMaxWaitMs); -} - -void logitsProcessorMixedReqsTest(std::string const& modelDir, SizeType32 worldRank, SizeType32 maxWaitMs, - bool replicated, std::optional> deviceIds); - -TEST_P(LogitsProcParamsTest, All) -{ - auto const modelName = std::get<0>(GetParam()); - auto const batched = std::get<1>(GetParam()); - auto const replicated = std::get<2>(GetParam()); - - std::string modelDir; - int tp_size = 1, pp_size = 1, cp_size = 1; - std::optional> deviceIds = std::nullopt; - - if (modelName == "llama_tp1_pp1_cp1") - { - modelDir = "tp1-pp1-cp1-gpu"; - } - else if (modelName == "llama_tp4_pp1_cp1") - { - modelDir = "tp4-pp1-cp1-gpu"; - tp_size = 4; - } - else if (modelName == "llama_tp1_pp4_cp1") - { - modelDir = "tp1-pp4-cp1-gpu"; - pp_size = 4; - deviceIds = std::vector{3, 2, 1, 0}; - } - else if (modelName == "llama_tp2_pp2_cp1") - { - modelDir = "tp2-pp2-cp1-gpu"; - tp_size = pp_size = 2; - deviceIds = std::vector{2, 3, 0, 1}; - } - else - { - TLLM_THROW("Unrecognized modelName"); - } - std::filesystem::path modelPath = LLAMA_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / modelDir; - - auto& comm = tensorrt_llm::mpi::MpiComm::world(); - auto const worldRank = comm.getRank(); - auto const worldSize = comm.getSize(); - - if (tp_size * pp_size * cp_size != 1) - { - // Run multi GPU test only when env variable is set - char const* val = getenv("RUN_LLAMA_MULTI_GPU"); - if (val == NULL) - { - GTEST_SKIP() << "Skipping multi-gpu logits post processor test"; - } - - if (worldSize != 4) - { - FAIL() << "Leader mode and world size is not equal to 4"; - } - } - else - { - // This has no effect for single-GPU tests - if (replicated) - { - GTEST_SKIP() << "Skipping single-gpu replicated logits post processor test"; - } - } - - // Configuration options - bool const streaming = false; - bool excludeInputFromOutput = false; - OutputConfig outConfig; - outConfig.excludeInputFromOutput = excludeInputFromOutput; - SizeType32 numRequests = 20; - IdType const kClientId = 1234; - - SizeType32 beamWidth = 1; - SizeType32 maxPromptLen = 20; - SizeType32 maxMaxNewTokens = 20; - - SizeType32 constexpr endId{2}; - SizeType32 constexpr vocabSizePadded{32000}; // llama-7b vocabSizePadded - // We just use tokenIdCalculator to generate a token_id based on request index, output position and max new tokens. - // Then LogitsPostProcessor set all other logits except the generated token_id to large negative value. - // So the output token should be the generated token by tokenIdCalculator. - auto tokenIdCalculator = [endId, vocabSizePadded](IdType req, SizeType32 pos) - { - SizeType32 tokenId = (req * 1000 + pos) % vocabSizePadded; - if (tokenId == endId) - { - tokenId = 0; - } - return tokenId; - }; - - std::unordered_map tokens; - std::unordered_map expectedNumTokens; - std::unordered_map expectedOutputTokens; - - // Enqueue the requests - auto enqueueRequests = [&](Executor& executor, std::optional logitsProcessorName, - std::optional logitsProcessor = std::nullopt) - { - tokens.clear(); - expectedNumTokens.clear(); - expectedOutputTokens.clear(); - - for (SizeType32 req = 0; req < numRequests; ++req) - { - SizeType32 promptLen = rand() % maxPromptLen + 1; - SizeType32 maxNewTokens = rand() % maxMaxNewTokens + 1; - - auto request = Request(VecTokens(promptLen, 1), maxNewTokens, streaming, - tensorrt_llm::executor::SamplingConfig(beamWidth), outConfig, endId); - request.setClientId(kClientId); - if (logitsProcessorName) - { - request.setLogitsPostProcessorName(logitsProcessorName.value()); - } - else if (logitsProcessor) - { - request.setLogitsPostProcessor(logitsProcessor.value()); - } - auto reqId = executor.enqueueRequest(std::move(request)); - tokens[reqId] = {}; - expectedNumTokens[reqId] = (streaming ? 0 : (excludeInputFromOutput ? 0 : promptLen)) + maxNewTokens; - expectedOutputTokens[reqId] = {}; - if (!streaming && !excludeInputFromOutput) - { - expectedOutputTokens[reqId].resize(promptLen, 1); - } - for (SizeType32 outputPos = 0; outputPos < maxNewTokens; ++outputPos) - { - SizeType32 outputTokenId = tokenIdCalculator(reqId, outputPos + promptLen); - expectedOutputTokens[reqId].push_back(outputTokenId); - } - } - }; - - // Get the new tokens for each requests - auto collectResponses = [&](Executor& executor) - { - int32_t numFinished = 0; - int iter = 0; - SizeType32 numResponses = 0; - while (numFinished < numRequests && iter < mMaxWaitMs) - { - std::chrono::milliseconds waitTime(1); - auto responses = executor.awaitResponses(waitTime); - for (auto& response : responses) - { - numResponses++; - if (!response.hasError()) - { - EXPECT_EQ(response.getClientId().value(), kClientId); - auto result = response.getResult(); - numFinished += result.isFinal; - auto& newTokens = result.outputTokenIds.at(beamWidth - 1); - auto& reqTokens = tokens.at(response.getRequestId()); - reqTokens.insert(reqTokens.end(), std::make_move_iterator(newTokens.begin()), - std::make_move_iterator(newTokens.end())); - } - else - { - // Allow response with error only if awaitResponse processed a terminated request id - std::string err = "ReqId " + std::to_string(response.getRequestId()) - + " has already been processed and was terminated."; - EXPECT_EQ(response.getErrorMsg(), err); - } - } - ++iter; - } - EXPECT_LT(iter, mMaxWaitMs); - }; - - // Check that tokens matches expectations - auto checkOutput = [&]() - { - for (auto const& [reqId, numTokens] : expectedNumTokens) - { - EXPECT_EQ(expectedNumTokens[reqId], tokens[reqId].size()) << "reqId " << reqId; - for (SizeType32 tokenPos = 0; - tokenPos < std::min(expectedNumTokens[reqId], tokens[reqId].size()); ++tokenPos) - { - EXPECT_EQ(expectedOutputTokens[reqId][tokenPos], tokens[reqId][tokenPos]) - << "reqId=" << reqId << ", tokenPos=" << tokenPos; - } - } - }; - - // Test non-batched logits processor - std::string const logitsProcessorName = "SelectToken"; - - auto logitsPostProcessorFn = [&](IdType reqId, Tensor& logits, BeamTokens const& tokens, StreamPtr const& streamPtr, - std::optional clientId) - { - if (replicated) - { - EXPECT_TRUE(worldRank <= tp_size - 1); - } - else - { - EXPECT_TRUE(worldRank == 0); - } - EXPECT_TRUE(clientId.value() == kClientId); - SizeType32 numTokens = tokens.at(0).size(); - SizeType32 pos = numTokens; - SizeType32 outputTokenId = tokenIdCalculator(reqId, pos); - auto logitsDataType = logits.getDataType(); - EXPECT_TRUE(logitsDataType == DataType::kFP16 || logitsDataType == DataType::kBF16 - || logitsDataType == DataType::kFP32); - // logits has shape [draftLength + 1, reqBeamWidth, vocabSize] - auto logitsCpu = tensorrt_llm::executor::Tensor::cpu(logitsDataType, logits.getShape()); - auto* dataPtr = logitsCpu.getData(); - auto eltSize = logitsCpu.getSizeInBytes() / logitsCpu.getSize(); - EXPECT_TRUE(eltSize == 2 || eltSize == 4); - if (eltSize == 2) - { - auto* dataPtrU16 = static_cast(dataPtr); - uint16_t hugeNegValue = logitsDataType == DataType::kFP16 ? 0xFBFF : 0xFF7F; // a huge negative value - for (size_t i = 0; i < logitsCpu.getSize(); ++i) - { - dataPtrU16[i] = hugeNegValue; - } - dataPtrU16[outputTokenId] = 0; - } - else - { - auto* dataPtrFloat = static_cast(dataPtr); - for (size_t i = 0; i < logitsCpu.getSize(); ++i) - { - dataPtrFloat[i] = -HUGE_VALF; - } - dataPtrFloat[outputTokenId] = 0.0f; - } - - logits.setFrom(logitsCpu, streamPtr); - }; - - if (!batched) - { - auto executorConfig = ExecutorConfig(beamWidth); - LogitsPostProcessorConfig logitsProcConfig{ - std::unordered_map{ - {logitsProcessorName, logitsPostProcessorFn}}, - std::nullopt, replicated}; - executorConfig.setLogitsPostProcessorConfig(logitsProcConfig); - if (deviceIds.has_value()) - { - auto parallelConfig = executorConfig.getParallelConfig().value_or(ParallelConfig()); - parallelConfig.setDeviceIds(deviceIds.value()); - executorConfig.setParallelConfig(parallelConfig); - } - auto executor = Executor(modelPath, ModelType::kDECODER_ONLY, executorConfig); - - if (worldRank == 0) - { - enqueueRequests(executor, logitsProcessorName); - collectResponses(executor); - checkOutput(); - - if (!replicated || tp_size == 1) - { - // Dynamic logits postprocessor must be used with replicate=false or no tensor parallelism. - enqueueRequests(executor, std::nullopt, logitsPostProcessorFn); - collectResponses(executor); - checkOutput(); - } - } - } - - // Test batched logits processor - auto logitsPostProcessorBatchedFn - = [logitsPostProcessorFn](std::vector const& reqIdBatch, std::vector& logitsBatch, - std::vector> const& tokensBatch, StreamPtr const& streamPtr, - std::vector> const& clientIdBatch) - { - for (int sample = 0; sample < reqIdBatch.size(); sample++) - { - logitsPostProcessorFn( - reqIdBatch[sample], logitsBatch[sample], tokensBatch[sample], streamPtr, clientIdBatch[sample]); - } - }; - - if (batched) - { - auto batchedExecutorConfig = ExecutorConfig(beamWidth); - if (deviceIds.has_value()) - { - auto parallelConfig = batchedExecutorConfig.getParallelConfig().value_or(ParallelConfig()); - - parallelConfig.setDeviceIds(deviceIds.value()); - batchedExecutorConfig.setParallelConfig(parallelConfig); - } - LogitsPostProcessorConfig logitsProcConfig{std::nullopt, logitsPostProcessorBatchedFn, replicated}; - batchedExecutorConfig.setLogitsPostProcessorConfig(logitsProcConfig); - - auto batchedExecutor = Executor(modelPath, ModelType::kDECODER_ONLY, batchedExecutorConfig); - - if (worldRank == 0) - { - enqueueRequests(batchedExecutor, Request::kBatchedPostProcessorName); - collectResponses(batchedExecutor); - checkOutput(); - } - } - - if (!batched) - { - logitsProcessorMixedReqsTest(modelDir, worldRank, mMaxWaitMs, replicated, std::move(deviceIds)); - } -} - -// Test for mixing requests with and without logits processor. -void logitsProcessorMixedReqsTest(std::string const& modelDir, SizeType32 worldRank, SizeType32 maxWaitMs, - bool replicated, std::optional> deviceIds) -{ - std::string const logitsProcessorName = "dummy"; - auto logitsPostProcessorFn = [&](IdType reqId, Tensor& logits, BeamTokens const& tokens, StreamPtr const& streamPtr, - std::optional clientId) - { - // Dummy callback that does not modify logits - assert(!clientId.has_value()); - }; - - LogitsPostProcessorConfig logitsProcConfig{ - std::unordered_map{ - {logitsProcessorName, logitsPostProcessorFn}}, - std::nullopt, replicated}; - - // Create executor - SizeType32 beamWidth = 1; - auto executorConfig = ExecutorConfig(beamWidth); - executorConfig.setLogitsPostProcessorConfig(logitsProcConfig); - if (deviceIds.has_value()) - { - auto parallelConfig = executorConfig.getParallelConfig().value_or(ParallelConfig()); - - parallelConfig.setDeviceIds(deviceIds.value()); - executorConfig.setParallelConfig(parallelConfig); - } - std::filesystem::path modelPath = LLAMA_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / modelDir; - auto executor = Executor(modelPath, ModelType::kDECODER_ONLY, executorConfig); - - if (worldRank == 0) - { - SizeType32 numRequests = 2; - SizeType32 promptLen = 5; - - // First request with no LP and many output tokens - auto request1 = Request(VecTokens(promptLen, 1), 25); - // Second request with LP and few output tokens - auto request2 = Request(VecTokens(promptLen, 1), 5); - request2.setLogitsPostProcessorName(logitsProcessorName); - - // Enqueue requests - auto reqId1 = executor.enqueueRequest(request1); - auto reqId2 = executor.enqueueRequest(request2); - - // Wait for responses - int32_t numFinished = 0; - int iter = 0; - SizeType32 numResponses = 0; - while (numFinished < numRequests && iter < maxWaitMs) - { - std::chrono::milliseconds waitTime(1); - auto responses = executor.awaitResponses(waitTime); - for (auto& response : responses) - { - numResponses++; - if (!response.hasError()) - { - auto result = response.getResult(); - numFinished += result.isFinal; - } - else - { - // Allow response with error only if awaitResponse processed a terminated request id - std::string err = "ReqId " + std::to_string(response.getRequestId()) - + " has already been processed and was terminated."; - EXPECT_EQ(response.getErrorMsg(), err); - } - } - ++iter; - } - EXPECT_LT(iter, maxWaitMs); - } -} - -TEST_F(GptExecutorTest, LogitsPostProcessorThrow) -{ - SizeType32 beamWidth = 1; - auto executorConfig = ExecutorConfig(beamWidth); - auto trtEnginePath = GPT_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp1-pp1-cp1-gpu"; - auto executor = Executor(trtEnginePath, ModelType::kDECODER_ONLY, executorConfig); - - std::string const logitsProcessorName = "UnExistProcessor"; - - auto request - = Request(VecTokens(10, 1), 10, false, tensorrt_llm::executor::SamplingConfig(beamWidth), OutputConfig()); - request.setLogitsPostProcessorName(logitsProcessorName); - EXPECT_THROW({ auto reqId = executor.enqueueRequest(std::move(request)); }, tensorrt_llm::common::TllmException); -} - -static Response executeDraftRequest(Executor& executor) -{ - OutputConfig outputConfig; - outputConfig.returnGenerationLogits = true; - - // Create the request - SizeType32 maxNewTokens = 4; - VecTokens inputTokens{1, 2, 3, 4}; - - Request request{std::move(inputTokens), maxNewTokens}; - request.setOutputConfig(outputConfig); - - // Enqueue the request - auto requestId = executor.enqueueRequest(std::move(request)); - - // Wait for the response - auto responses = executor.awaitResponses(requestId); - - return responses.at(0); -} - -static Response executeTargetRequest(Executor& executor, Result const& draftResult) -{ - // Create the request - SizeType32 maxNewTokens = 5; - VecTokens inputTokens{1, 2, 3, 4}; - - Request request{std::move(inputTokens), maxNewTokens}; - - VecTokens const& outputTokenIds = draftResult.outputTokenIds.at(0); - VecTokens draftTokens(outputTokenIds.end() - 4, outputTokenIds.end()); - - auto const& logitsInfo = draftResult.specDecFastLogitsInfo.value(); - auto logitsTensor = logitsInfo.toTensor(); - - ExternalDraftTokensConfig draftTokensConfig( - std::move(draftTokens), logitsTensor, std::nullopt /* acceptance threshold */, true /* fastLogits */); - request.setExternalDraftTokensConfig(draftTokensConfig); - - // Enqueue the request - auto requestId = executor.enqueueRequest(std::move(request)); - - // Wait for the response - auto responses = executor.awaitResponses(requestId); - - return responses.at(0); -} - -class SpeculativeDecodingTest : public GptExecutorTest -{ -}; - -TEST_F(SpeculativeDecodingTest, SpecDecFastLogits) -{ - SizeType32 beamWidth = 1; - auto executorConfig = ExecutorConfig(beamWidth); - auto trtDraftEnginePath - = GPT_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_GATHER_DIR() / "tp1-pp1-cp1-gpu"; - auto trtEnginePath - = GPT_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DRAFT_TOKENS_DIR() / "tp1-pp1-cp1-gpu"; - - FloatType freeGpuMemoryFraction = 0.3; - auto kvCacheConfig - = KvCacheConfig(true /* enableBlockReuse */, std::nullopt, std::nullopt, std::nullopt, freeGpuMemoryFraction); - executorConfig.setKvCacheConfig(kvCacheConfig); - - tensorrt_llm::mpi::initialize(tensorrt_llm::mpi::MpiThreadSupport::THREAD_MULTIPLE); - int const worldSize = tensorrt_llm::mpi::MpiComm::world().getSize(); - ASSERT_EQ(worldSize, 3); - int const myRank = tensorrt_llm::mpi::MpiComm::world().getRank(); - bool const isOrchestrator = (myRank == 0); - - auto orchestratorConfig - = OrchestratorConfig(isOrchestrator, "" /* workerExecutablePath */, nullptr, false /* spawnPrcesses */); - auto parallelConfig = ParallelConfig( - CommunicationType::kMPI, CommunicationMode::kORCHESTRATOR, std::nullopt, std::nullopt, orchestratorConfig); - executorConfig.setParallelConfig(parallelConfig); - - auto specDecConfig = SpeculativeDecodingConfig(true /* fastLogits */); - executorConfig.setSpecDecConfig(specDecConfig); - - std::unique_ptr draftExecutor; - std::unique_ptr targetExecutor; - - if (isOrchestrator) - { - auto executorConfigDraft = executorConfig; - parallelConfig.setParticipantIds({1}); - executorConfigDraft.setParallelConfig(parallelConfig); - - draftExecutor = std::make_unique(trtDraftEnginePath, ModelType::kDECODER_ONLY, executorConfigDraft); - - parallelConfig.setParticipantIds({2}); - executorConfig.setParallelConfig(parallelConfig); - - targetExecutor = std::make_unique(trtEnginePath, ModelType::kDECODER_ONLY, executorConfig); - } - else if (myRank == 1) // draft model process - { - parallelConfig.setParticipantIds({1}); - parallelConfig.setDeviceIds({0}); - executorConfig.setParallelConfig(parallelConfig); - executorConfig.setGatherGenerationLogits(true); - draftExecutor = std::make_unique(trtDraftEnginePath, ModelType::kDECODER_ONLY, executorConfig); - } - else if (myRank == 2) // target model process - { - parallelConfig.setParticipantIds({2}); - parallelConfig.setDeviceIds({0}); - executorConfig.setParallelConfig(parallelConfig); - draftExecutor = std::make_unique(trtEnginePath, ModelType::kDECODER_ONLY, executorConfig); - } - - if (isOrchestrator) - { - auto response = executeDraftRequest(*draftExecutor); - ASSERT_FALSE(response.hasError()); - response = executeTargetRequest(*targetExecutor, response.getResult()); - ASSERT_FALSE(response.hasError()); - } -} - -TEST_F(GptExecutorTest, OrchestratorMaxQueueSize) -{ - auto trtEnginePath = GPT_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp1-pp1-cp1-gpu"; - SizeType32 maxQueueSize = 6; - ExecutorConfig executorConfig; - executorConfig.setMaxQueueSize(maxQueueSize); - auto orchestratorConfig = OrchestratorConfig(true, PathUtil::EXECUTOR_WORKER_PATH()); - auto parallelConfig = ParallelConfig( - CommunicationType::kMPI, CommunicationMode::kORCHESTRATOR, std::nullopt, std::nullopt, orchestratorConfig); - executorConfig.setParallelConfig(parallelConfig); - - auto executor = Executor(trtEnginePath, ModelType::kDECODER_ONLY, executorConfig); - - // Create the request - SizeType32 maxNewTokens = 100; - VecTokens inputTokens{1, 2, 3, 4}; - auto request = Request(inputTokens, maxNewTokens); - std::vector requestIds; - auto numberOfRequests = maxQueueSize * 5; - requestIds.reserve(numberOfRequests); - - // Enqueue more requests than the queue can manage - for (int i = 0; i < numberOfRequests; i++) - { - auto requestId = executor.enqueueRequest(request); - requestIds.emplace_back(requestId); - } - - auto responseVectors = executor.awaitResponses(std::move(requestIds)); - bool failedWithFullQueue = false; - for (auto& responseVector : responseVectors) - { - for (auto& response : responseVector) - { - if (response.hasError()) - { - EXPECT_THAT(response.getErrorMsg(), - testing::HasSubstr("Maximum queue size of 6 has been reached, please try again later")); - failedWithFullQueue = true; - } - } - } - EXPECT_TRUE(failedWithFullQueue) << "Expected requests to fail due to maximum queue size reached"; - - // Wait for requests to get scheduled to free up space in queue - std::this_thread::sleep_for(std::chrono::milliseconds(maxQueueSize * 200)); - auto requestId = executor.enqueueRequest(std::move(request)); - auto responses = executor.awaitResponses(requestId); - for (auto& response : responses) - { - EXPECT_FALSE(response.hasError()); - } -} - -TEST_F(GptExecutorTest, SingleRequestInvalidInputs) -{ - bool streaming = true; - - SizeType32 beamWidth = 1; - auto executorConfig = ExecutorConfig(beamWidth); - auto trtEnginePath = GPT_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp1-pp1-cp1-gpu"; - auto executor = Executor(trtEnginePath, ModelType::kDECODER_ONLY, executorConfig); - - // Create the request - SizeType32 maxNewTokens = 5; - VecTokens inputTokens{1, 2, 3, 4}; - - std::vector expectedErrMsgs; - std::vector requests; - - // Invalid embedding bias shape - { - requests.emplace_back(inputTokens, maxNewTokens, streaming); - auto embeddingBias = Tensor::cpu(DataType::kFP32, {1}); - requests.back().setEmbeddingBias(embeddingBias); - expectedErrMsgs.emplace_back("embedding bias shape is not as expected"); - } - - for (auto req = 0; req < requests.size(); ++req) - { - auto& request = requests.at(req); - auto const& expectedErrMsg = expectedErrMsgs.at(req); - - auto requestId = executor.enqueueRequest(std::move(request)); - - // Try to get the new tokens - bool done = false; - int iter = 0; - while (!done && iter < mMaxWaitMs) - { - std::chrono::milliseconds waitTime(1); - auto responses = executor.awaitResponses(requestId, waitTime); - for (auto& response : responses) - { - if (response.hasError()) - { - - auto err = response.getErrorMsg(); - EXPECT_THAT(err, testing::HasSubstr(expectedErrMsg)); - done = true; - } - else - { - FAIL() << "Expected an err: " << expectedErrMsg; - } - } - ++iter; - } - EXPECT_EQ(done, true); - } -} - -TEST_F(GptExecutorTest, ExecutorKVCacheManager) -{ - - bool streaming = true; - int numRequests = 3; - - SizeType32 beamWidth = 1; - SizeType32 maxNewTokens = 5; - auto executorConfig = ExecutorConfig(beamWidth); - auto kvCacheConfig = KvCacheConfig(true, 128); - kvCacheConfig.setEventBufferMaxSize(1024); - executorConfig.setKvCacheConfig(kvCacheConfig); - - auto trtEnginePath = GPT_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp1-pp1-cp1-gpu"; - auto executor = Executor(trtEnginePath, ModelType::kDECODER_ONLY, executorConfig); - - auto kvCacheManager = *executor.getKVCacheEventManager(); - - // Created event should be available before any requests. - auto events = kvCacheManager->getLatestEvents(std::chrono::seconds(1)); - EXPECT_EQ(events.size(), 1); - EXPECT_TRUE(std::holds_alternative(events.front().data)); - - // Create requests - std::vector requests; - for (int request = 0; request < 3; request++) - { - VecTokens inputTokens; - for (int i = 0; i < 63; i++) - { - inputTokens.emplace_back(i + request); - } - requests.emplace_back(inputTokens, maxNewTokens, streaming); - } - - for (auto req = 0; req < requests.size(); ++req) - { - auto& request = requests.at(req); - - auto requestId = executor.enqueueRequest(std::move(request)); - - // Get the new tokens - bool done = false; - int iter = 0; - while (!done && iter < mMaxWaitMs) - { - std::chrono::milliseconds waitTime(1); - auto responses = executor.awaitResponses(requestId, waitTime); - for (auto& response : responses) - { - if (response.hasError()) - { - // This request failed for some reason, get error msg - std::string errStr - = "Request id " + std::to_string(requestId) + " failed with err " + response.getErrorMsg(); - FAIL(); - } - else - { - auto result = response.getResult(); - done = result.isFinal; - if (done) - { - std::this_thread::sleep_for(std::chrono::milliseconds(100)); - auto events = kvCacheManager->getLatestEvents(std::chrono::milliseconds(100)); - if (req == 0) - { - EXPECT_EQ(events.size(), 3); - - // Store the first context block - EXPECT_EQ(std::get(events.front().data).parentHash, std::nullopt); - EXPECT_EQ(std::get(events.front().data).blocks.size(), 1); - events.pop_front(); - // Store the second (now completed) context block and the partial decode block. - EXPECT_EQ(std::get(events.front().data).blocks.size(), 1); - EXPECT_EQ(std::get(events.back().data).blocks.size(), 1); - EXPECT_EQ(std::get(events.front().data).blocks[0].blockHash, - std::get(events.back().data).parentHash); - } - else - { - EXPECT_EQ(events.size(), 5); - - // Remove a block to make room for the second context block. On the second request, we need - // to remove 2 blocks. - EXPECT_EQ(std::get(events.front().data).blockHashes.size(), req); - events.pop_front(); - // Store the first filled context block - EXPECT_EQ(std::get(events.front().data).blocks.size(), 1); - events.pop_front(); - // Remove a block for the decode phase - EXPECT_EQ(std::get(events.front().data).blockHashes.size(), 1); - events.pop_front(); - // Store the final context block and the decode block - EXPECT_EQ(std::get(events.front().data).blocks.size(), 1); - events.pop_front(); - EXPECT_EQ(std::get(events.front().data).blocks.size(), 1); - } - } - } - } - iter++; - } - EXPECT_EQ(done, true); - } -} - -TEST_F(GptExecutorTest, SingleRequestLora) -{ - bool streaming = true; - - SizeType32 beamWidth = 1; - auto executorConfig = ExecutorConfig(beamWidth); - auto trtEnginePath = GPT_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp1-pp1-cp1-gpu"; - auto executor = Executor(trtEnginePath, ModelType::kDECODER_ONLY, executorConfig); - - // Load lora weights, config - auto manager = tr::BufferManager(std::make_shared()); - auto loraWeightsTensor - = std::shared_ptr(tr::utils::loadNpy(manager, LORA_WEIGHTS_FILE.string(), tr::MemoryType::kCPU)); - auto loraConfigTensor - = std::shared_ptr(tr::utils::loadNpy(manager, LORA_CONFIG_FILE.string(), tr::MemoryType::kCPU)); - - // Create the request - SizeType32 maxNewTokens = 5; - VecTokens inputTokens{1, 2, 3, 4}; - auto request = Request(inputTokens, maxNewTokens, streaming, tensorrt_llm::executor::SamplingConfig()); - auto loraConfig = LoraConfig(0, detail::ofITensor(loraWeightsTensor), detail::ofITensor(loraConfigTensor)); - request.setLoraConfig(loraConfig); - - // Enqueue the request - auto requestId = executor.enqueueRequest(std::move(request)); - - // Get the new tokens - VecTokens tokens; - bool done = false; - int iter = 0; - std::chrono::milliseconds waitTime(1); - while (!done && iter < mMaxWaitMs) - { - auto responses = executor.awaitResponses(requestId, waitTime); - for (auto& response : responses) - { - if (response.hasError()) - { - // This request failed for some reason, get error msg - std::string errStr - = "Request id " + std::to_string(requestId) + " failed with err " + response.getErrorMsg(); - FAIL(); - } - else - { - auto result = response.getResult(); - done = result.isFinal; - // Append tokens - auto& newTokens = result.outputTokenIds.at(beamWidth - 1); - tokens.insert( - tokens.end(), std::make_move_iterator(newTokens.begin()), std::make_move_iterator(newTokens.end())); - } - } - ++iter; - } - EXPECT_LT(iter, mMaxWaitMs); - EXPECT_EQ(tokens.size(), maxNewTokens); -} - -TEST_P(GuidedDecodingParamsTest, All) -{ - auto const modelName = std::get<0>(GetParam()); - std::filesystem::path enginePath; - std::filesystem::path tokenizerInfoPath; - int tp_size = 1, pp_size = 1, cp_size = 1; - std::optional> deviceIds = std::nullopt; - - if (modelName == "gpt") - { - enginePath = GPT_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp1-pp1-cp1-gpu"; - tokenizerInfoPath = GPT_XGRAMMAR_TOKENIZER_INFO_PATH; - } - else if (modelName == "llama_tp1_pp1_cp1") - { - enginePath = LLAMA_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp1-pp1-cp1-gpu"; - tokenizerInfoPath = LLAMA_XGRAMMAR_TOKENIZER_INFO_PATH; - } - else if (modelName == "llama_tp4_pp1_cp1") - { - enginePath = LLAMA_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp4-pp1-cp1-gpu"; - tokenizerInfoPath = LLAMA_XGRAMMAR_TOKENIZER_INFO_PATH; - tp_size = 4; - } - else if (modelName == "llama_tp1_pp4_cp1") - { - enginePath = LLAMA_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp1-pp4-cp1-gpu"; - tokenizerInfoPath = LLAMA_XGRAMMAR_TOKENIZER_INFO_PATH; - pp_size = 4; - deviceIds = std::vector{3, 2, 1, 0}; - } - else if (modelName == "llama_tp2_pp2_cp1") - { - enginePath = LLAMA_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp2-pp2-cp1-gpu"; - tokenizerInfoPath = LLAMA_XGRAMMAR_TOKENIZER_INFO_PATH; - tp_size = 2; - pp_size = 2; - deviceIds = std::vector{2, 3, 0, 1}; - } - else - { - TLLM_THROW("Unrecognized modelName"); - } - - auto& comm = tensorrt_llm::mpi::MpiComm::world(); - auto const worldRank = comm.getRank(); - auto const worldSize = comm.getSize(); - - if (tp_size * pp_size * cp_size > 1) - { - // Run multi GPU test only when env variable is set - char const* val = getenv("RUN_LLAMA_MULTI_GPU"); - if (val == NULL) - { - GTEST_SKIP() << "Skipping multi-gpu guided decoding test"; - } - else - { - if (worldSize != 4) - { - FAIL() << "Leader mode and world size is not equal to 4"; - } - } - } - - bool streaming = false; - - SizeType32 beamWidth = 1; - auto executorConfig = ExecutorConfig(beamWidth); - - auto const tokenizerInfo = nlohmann::json::parse(std::ifstream{tokenizerInfoPath}); - auto const encodedVocab = tokenizerInfo["encoded_vocab"].template get>(); - auto const tokenizerStr = tokenizerInfo["tokenizer_str"].template get(); - auto const stopTokenIds = tokenizerInfo["stop_token_ids"].template get>(); - GuidedDecodingConfig guidedDecodingConfig( - GuidedDecodingConfig::GuidedDecodingBackend::kXGRAMMAR, encodedVocab, tokenizerStr, stopTokenIds); - executorConfig.setGuidedDecodingConfig(guidedDecodingConfig); - - if (deviceIds.has_value()) - { - auto parallelConfig = executorConfig.getParallelConfig().value_or(ParallelConfig()); - - parallelConfig.setDeviceIds(deviceIds.value()); - executorConfig.setParallelConfig(parallelConfig); - } - auto executor = Executor(enginePath, ModelType::kDECODER_ONLY, executorConfig); - - // Create the requests - VecTokens inputTokens; - if (modelName == "gpt") - { - inputTokens = {2061, 318, 352, 10, 16, 30, 23998, 39559, 287, 257, 8633, 287, 33918, 5794, 25, 220}; - } - else // llama - { - inputTokens = { - 128000, 62, 3923, 7037, 62, 16, 10, 16, 30, 62, 16533, 87710, 1265, 4404, 5356, 1265, 9643, 9132, 25, 62}; - } - SizeType32 maxNewTokens = 10; - SamplingConfig samplingConfig{}; - OutputConfig outputConfig{false, false, false, true}; - - std::vector requests; - requests.emplace_back(inputTokens, maxNewTokens, streaming, samplingConfig, outputConfig, stopTokenIds[0]); - - requests.emplace_back(inputTokens, maxNewTokens, streaming, samplingConfig, outputConfig, stopTokenIds[0]); - requests.back().setGuidedDecodingParams(GuidedDecodingParams(GuidedDecodingParams::GuideType::kJSON)); - - requests.emplace_back(inputTokens, maxNewTokens, streaming, samplingConfig, outputConfig, stopTokenIds[0]); - std::string jsonSchema{ - R"({"properties": {"answer": {"title": "Answer", "type": "integer"}}, "required": ["answer"], "title": "Answer", "type": "object"})"}; - requests.back().setGuidedDecodingParams( - GuidedDecodingParams(GuidedDecodingParams::GuideType::kJSON_SCHEMA, jsonSchema)); - - requests.emplace_back(inputTokens, maxNewTokens, streaming, samplingConfig, outputConfig, stopTokenIds[0]); - std::string regex{R"(\d+)"}; - requests.back().setGuidedDecodingParams(GuidedDecodingParams(GuidedDecodingParams::GuideType::kREGEX, regex)); - - requests.emplace_back(inputTokens, maxNewTokens, streaming, samplingConfig, outputConfig, stopTokenIds[0]); - std::string ebnfGrammar{R"(root ::= [0-9]+)"}; - requests.back().setGuidedDecodingParams( - GuidedDecodingParams(GuidedDecodingParams::GuideType::kEBNF_GRAMMAR, ebnfGrammar)); - - std::vector expectedOutputTokens; - if (modelName == "gpt") - { - expectedOutputTokens.push_back({1849, 7, 16, 10, 16, 8, 198, 16, 10, 16}); - expectedOutputTokens.push_back({90, 366, 3672, 1298, 366, 7554, 31780, 1600, 366, 12888}); - expectedOutputTokens.push_back({90, 366, 64, 77, 2032, 68, 81, 1, 1058, 352}); - expectedOutputTokens.push_back({25645, 25645, 25645, 25645, 25645, 25645, 25645, 25645, 25645, 25645}); - expectedOutputTokens.push_back({25645, 25645, 25645, 25645, 25645, 25645, 25645, 25645, 25645, 25645}); - } - else // llama - { - expectedOutputTokens.push_back({16, 10, 16, 28, 17, 198, 62, 3923, 7037, 62}); - expectedOutputTokens.push_back({5018, 16, 794, 330, 16, 498, 330, 17, 794, 330}); - expectedOutputTokens.push_back({5018, 9399, 794, 16, 92}); - expectedOutputTokens.push_back({16}); - expectedOutputTokens.push_back({16}); - } - - if (executor.canEnqueueRequests()) - { - // Enqueue the requests - auto reqIds = executor.enqueueRequests(std::move(requests)); - - // Get the responses - int numFinished = 0; - int iter = 0; - while (numFinished < 5 && iter < mMaxWaitMs) - { - std::chrono::milliseconds waitTime(1); - auto responses = executor.awaitResponses(waitTime); - for (auto& response : responses) - { - auto reqId = response.getRequestId(); - if (response.hasError()) - { - // This request failed for some reason, get error msg - std::string errStr - = "Request id " + std::to_string(reqId) + " failed with err " + response.getErrorMsg(); - FAIL(); - } - else - { - auto result = response.getResult(); - auto& newTokens = result.outputTokenIds.at(0); - - int reqIdx = std::find(reqIds.begin(), reqIds.end(), reqId) - reqIds.begin(); - EXPECT_THAT(newTokens, ::testing::ElementsAreArray(expectedOutputTokens[reqIdx])); - } - numFinished++; - } - } - EXPECT_LT(iter, mMaxWaitMs); - EXPECT_EQ(numFinished, 5); - } -} - -TEST_F(GptExecutorTest, GuidedDecodingFailure) -{ - bool streaming = false; - - SizeType32 beamWidth = 1; - auto executorConfig = ExecutorConfig(beamWidth); - - std::vector stopTokenIds{50256}; - auto trtEnginePath = GPT_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp1-pp1-cp1-gpu"; - auto executor = Executor(trtEnginePath, ModelType::kDECODER_ONLY, executorConfig); - - // Create the requests - SizeType32 maxNewTokens = 10; - SamplingConfig samplingConfig{}; - OutputConfig outputConfig{false, false, false, true}; - VecTokens inputTokens{2061, 318, 352, 10, 16, 30, 23998, 39559, 287, 257, 8633, 287, 33918, 5794, 25, 220}; - - std::vector requests; - requests.emplace_back(inputTokens, maxNewTokens, streaming, samplingConfig, outputConfig, stopTokenIds[0]); - requests.emplace_back(inputTokens, maxNewTokens, streaming, samplingConfig, outputConfig, stopTokenIds[0]); - requests.back().setGuidedDecodingParams(GuidedDecodingParams(GuidedDecodingParams::GuideType::kJSON)); - - // Enqueue the requests - auto reqIds = executor.enqueueRequests(std::move(requests)); - - // Get the responses - int numFinished = 0; - int iter = 0; - while (numFinished < 2 && iter < mMaxWaitMs) - { - std::chrono::milliseconds waitTime(1); - auto responses = executor.awaitResponses(waitTime); - for (auto& response : responses) - { - auto reqId = response.getRequestId(); - int reqIdx = std::find(reqIds.begin(), reqIds.end(), reqId) - reqIds.begin(); - if (reqIdx == 0) - { - EXPECT_FALSE(response.hasError()); - } - else - { - EXPECT_TRUE(response.hasError()); - } - numFinished++; - } - } - EXPECT_LT(iter, mMaxWaitMs); - EXPECT_EQ(numFinished, 2); -} - -TEST_P(ParamTest, SingleRequestCancelRequest) -{ - bool const streaming = std::get<0>(GetParam()); - bool const excludeInputFromOutput = std::get<1>(GetParam()); - auto const beamWidth = std::get<2>(GetParam()); - OutputConfig outConfig; - outConfig.excludeInputFromOutput = excludeInputFromOutput; - - auto executorConfig = ExecutorConfig(beamWidth); - auto trtEnginePath = GPT_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp1-pp1-cp1-gpu"; - auto executor = Executor(trtEnginePath, ModelType::kDECODER_ONLY, executorConfig); - - // Create the request - SizeType32 maxNewTokens = 300; - VecTokens inputTokens{1, 2, 3, 4}; - auto request - = Request(inputTokens, maxNewTokens, streaming, tensorrt_llm::executor::SamplingConfig(beamWidth), outConfig); - - auto requestId = executor.enqueueRequest(std::move(request)); - - std::this_thread::sleep_for(std::chrono::milliseconds(100)); - executor.cancelRequest(requestId); - - // Try to get the new tokens - bool done = false; - int iter = 0; - VecTokens tokens; - while (!done && iter < mMaxWaitMs) - { - std::chrono::milliseconds waitTime(1); - auto responses = executor.awaitResponses(requestId, waitTime); - for (auto& response : responses) - { - if (response.hasError()) - { - FAIL() << "Did not expect errors"; - } - else - { - auto result = response.getResult(); - done = result.isFinal; - // Append tokens - auto& newTokens = result.outputTokenIds.at(beamWidth - 1); - if (done) - { - for (SizeType32 beamIdx = 0; beamIdx < beamWidth; ++beamIdx) - { - EXPECT_EQ(result.finishReasons[beamIdx], FinishReason::kCANCELLED); - } - } - - if (streaming && beamWidth > 1) - { - tokens = newTokens; - } - else - { - tokens.insert(tokens.end(), newTokens.begin(), newTokens.end()); - } - } - } - ++iter; - } - EXPECT_EQ(done, true); - EXPECT_LT(iter, mMaxWaitMs); - auto expectedNumTokens - = streaming ? maxNewTokens : (excludeInputFromOutput ? 0 : inputTokens.size()) + maxNewTokens; - TLLM_LOG_INFO("num tokens: %d, expected %d", tokens.size(), expectedNumTokens); - EXPECT_LT(tokens.size(), expectedNumTokens); -} - -TEST_F(GptExecutorTest, orchModeFetchNewReqErr) -{ - SizeType32 beamWidth = 1; - auto executorConfig = ExecutorConfig(beamWidth); - - auto orchestratorConfig = OrchestratorConfig(true, PathUtil::EXECUTOR_WORKER_PATH()); - auto parallelConfig = ParallelConfig( - CommunicationType::kMPI, CommunicationMode::kORCHESTRATOR, std::nullopt, std::nullopt, orchestratorConfig); - executorConfig.setParallelConfig(parallelConfig); - - auto trtEnginePath = (GPT_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp1-pp1-cp1-gpu"); - auto executor = Executor(trtEnginePath, ModelType::kDECODER_ONLY, executorConfig); - - // Create a req with invalid parameters - SizeType32 maxNewTokens = 5; - // Create very long prompt which should result in error during request validate - VecTokens inputTokens(10000000); - - auto request = Request(inputTokens, maxNewTokens, false, tensorrt_llm::executor::SamplingConfig(beamWidth)); - auto requestId = executor.enqueueRequest(request); - auto requestId2 = executor.enqueueRequest(request); - - bool done = false; - int iter = 0; - while (!done && iter < mMaxWaitMs) - { - std::chrono::milliseconds waitTime(1); - auto responses = executor.awaitResponses(waitTime); - for (auto& response : responses) - { - if (response.hasError()) - { - auto err = response.getErrorMsg(); - EXPECT_THAT(err, testing::HasSubstr("exceeds maximum input length")); - EXPECT_THAT(err, testing::HasSubstr("Encountered an error when fetching new request:")); - done = true; - } - else - { - FAIL() << "Should get a response with error"; - } - } - ++iter; - } - EXPECT_LT(iter, mMaxWaitMs); -} - -TEST_F(GptExecutorTest, orchModeForwardError) -{ - SizeType32 constexpr maxBeamWidth{1}; - auto executorConfig = ExecutorConfig(maxBeamWidth); - - auto orchestratorConfig = OrchestratorConfig(true, PathUtil::EXECUTOR_WORKER_PATH()); - auto parallelConfig = ParallelConfig( - CommunicationType::kMPI, CommunicationMode::kORCHESTRATOR, std::nullopt, std::nullopt, orchestratorConfig); - executorConfig.setParallelConfig(parallelConfig); - - auto trtEnginePath = (GPT_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp1-pp1-cp1-gpu"); - auto executor = Executor(trtEnginePath, ModelType::kDECODER_ONLY, executorConfig); - - // Setting request beam width to 2 which should cause failure - SizeType32 constexpr beamWidth{2}; - SizeType32 constexpr maxNewTokens{5}; - VecTokens inputTokens{1, 2, 3, 4}; - - auto request = Request(inputTokens, maxNewTokens, false, tensorrt_llm::executor::SamplingConfig(beamWidth)); - auto requestId = executor.enqueueRequest(request); - auto requestId2 = executor.enqueueRequest(request); - - bool done = false; - int iter = 0; - while (!done && iter < mMaxWaitMs) - { - std::chrono::milliseconds waitTime(1); - auto responses = executor.awaitResponses(waitTime); - for (auto& response : responses) - { - if (response.hasError()) - { - auto err = response.getErrorMsg(); - std::cout << "err:" << err << std::endl; - EXPECT_THAT( - err, testing::HasSubstr("Requested beam width 2 is larger than configured max beam width 1")); - done = true; - } - else - { - FAIL() << "Should get a response with error"; - } - } - ++iter; - } - EXPECT_LT(iter, mMaxWaitMs); -} - -TEST_P(ParamCancelReqTest, MultipleRequestsMultiGpuCancelRequest) -{ - auto const useOrchestratorMode = std::get<0>(GetParam()); - auto const beamWidth = std::get<1>(GetParam()); - auto const modelName = std::get<2>(GetParam()); - - std::optional> deviceIds = std::nullopt; - - OutputConfig outConfig; - - auto executorConfig = ExecutorConfig(beamWidth); - std::filesystem::path modelPath; - if (modelName == "llama_tp4_pp1_cp1" || modelName == "llama_tp1_pp4_cp1" || modelName == "llama_tp2_pp2_cp1") - { - if (modelName == "llama_tp4_pp1_cp1") - { - modelPath = LLAMA_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp4-pp1-cp1-gpu"; - } - else if (modelName == "llama_tp1_pp4_cp1") - { - modelPath = LLAMA_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp1-pp4-cp1-gpu"; - deviceIds = std::vector{3, 2, 1, 0}; - } - else if (modelName == "llama_tp2_pp2_cp1") - { - modelPath = LLAMA_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp2-pp2-cp1-gpu"; - deviceIds = std::vector{2, 3, 0, 1}; - } - } - - // For llama model, only run for multiple GPUs - // This is detected by setting an env variable when running the test - char const* val = getenv("RUN_LLAMA_MULTI_GPU"); - if (val == NULL) - { - GTEST_SKIP() << "Skipping Llama test"; - } - else - { - // Check that it was launched with right number of MPI ranks - if (!useOrchestratorMode && COMM_SESSION.getSize() != 4) - { - // No orchestrator, need worldSize to match TP*PP - FAIL() << "Leader mode and world size is not equal to 4"; - } - else if (useOrchestratorMode && COMM_SESSION.getSize() != 1) - { - // No orchestrator, need worldSize to match TP*PP - FAIL() << "Orchestrator mode and World size is not equal to 1"; - } - } - - if (useOrchestratorMode) - { - auto orchestratorConfig = OrchestratorConfig(true, PathUtil::EXECUTOR_WORKER_PATH()); - auto parallelConfig = ParallelConfig(CommunicationType::kMPI, - useOrchestratorMode ? CommunicationMode::kORCHESTRATOR : CommunicationMode::kLEADER, std::nullopt, - std::nullopt, orchestratorConfig); - if (deviceIds.has_value()) - { - parallelConfig.setDeviceIds(deviceIds.value()); - } - executorConfig.setParallelConfig(parallelConfig); - } - else - { - if (deviceIds.has_value()) - { - auto parallelConfig = executorConfig.getParallelConfig().value_or(ParallelConfig()); - parallelConfig.setDeviceIds(deviceIds.value()); - executorConfig.setParallelConfig(parallelConfig); - } - } - - auto executor = Executor(modelPath, ModelType::kDECODER_ONLY, executorConfig); - - // Create the request - SizeType32 maxNewTokens = 50; - VecTokens inputTokens{1, 2, 3, 4}; - - std::vector requests; - for (auto streaming : {false, true}) - { - // Add two requests with numReturnSequences = 1 - auto samplingConfig = tensorrt_llm::executor::SamplingConfig(beamWidth); - requests.emplace_back(inputTokens, maxNewTokens, streaming, samplingConfig, outConfig); - requests.emplace_back(inputTokens, maxNewTokens, streaming, samplingConfig, outConfig); - // Add a request with numReturnSequences > 1 - auto samplingConfig2 = tensorrt_llm::executor::SamplingConfig(beamWidth); - auto constexpr numReturnSequences = 2; - samplingConfig2.setNumReturnSequences(numReturnSequences); - requests.emplace_back(inputTokens, maxNewTokens, streaming, samplingConfig2, outConfig); - } - std::vector cancelRequests{true, false, true, true, false, true}; - - if (executor.canEnqueueRequests()) - { - auto const requestIds = executor.enqueueRequests(requests); - - // Cancel the first and third requests - std::this_thread::sleep_for(std::chrono::milliseconds(50)); - for (SizeType32 i = 0; i < requests.size(); i++) - { - if (cancelRequests.at(i)) - { - executor.cancelRequest(requestIds.at(i)); - } - } - - std::unordered_map isStreaming; - std::unordered_map expectedNumTokens; - SizeType32 expectedNumResponses = 0; - for (SizeType32 i = 0; i < requests.size(); i++) - { - auto const& request = requests.at(i); - auto requestId = requestIds.at(i); - isStreaming[requestId] = request.getStreaming(); - expectedNumTokens[requestId] = (request.getStreaming() ? 0 : inputTokens.size()) + maxNewTokens; - auto const numResponses = request.getStreaming() ? expectedNumTokens[requestId] : 1; - auto const numReturnSequences = request.getSamplingConfig().getBeamWidth() > 1 - ? 1 - : request.getSamplingConfig().getNumReturnSequences().value_or(1); - expectedNumResponses += numResponses * numReturnSequences; - } - - std::unordered_map> tokens; - - // Get the new tokens for each requests - int32_t numFinished = 0; - int iter = 0; - SizeType32 numResponses = 0; - while (numFinished < requests.size() && iter < mMaxWaitMs) - { - std::chrono::milliseconds waitTime(1); - auto responses = executor.awaitResponses(waitTime); - for (auto& response : responses) - { - numResponses++; - if (!response.hasError()) - { - auto requestId = response.getRequestId(); - auto result = response.getResult(); - numFinished += result.isFinal; - auto seqIdx = result.sequenceIndex; - auto numSequences = result.outputTokenIds.size(); - auto& newTokens = result.outputTokenIds.at(numSequences - 1); - auto& reqResults = tokens[response.getRequestId()]; - auto& reqTokens = reqResults[seqIdx]; - if (isStreaming.at(requestId) && beamWidth > 1) - { - reqTokens = newTokens; - } - else - { - reqTokens.insert(reqTokens.end(), newTokens.begin(), newTokens.end()); - } - } - else - { - FAIL() << "Did not expect errors"; - } - } - ++iter; - } - - EXPECT_LE(numResponses, expectedNumResponses); - EXPECT_EQ(numFinished, requests.size()); - EXPECT_LT(iter, mMaxWaitMs); - - for (auto requestIdx = 0; requestIdx < requests.size(); requestIdx++) - { - auto const requestId = requestIds.at(requestIdx); - for (auto seqIdx = 0; seqIdx < tokens.at(requestId).size(); seqIdx++) - { - auto const& seqTokens = tokens.at(requestId).at(seqIdx); - if (cancelRequests.at(requestIdx)) - { - EXPECT_LT(seqTokens.size(), expectedNumTokens.at(requestId)); - } - else - { - EXPECT_EQ(seqTokens.size(), expectedNumTokens.at(requestId)); - } - } - } - } -} - -TEST_P(LeaderApiUsageTest, LeaderModeTest) -{ - auto const modelName = std::get<0>(GetParam()); - - SizeType32 beamWidth = 2; - OutputConfig outConfig; - std::optional> deviceIds = std::nullopt; - - auto executorConfig = ExecutorConfig(beamWidth); - std::filesystem::path modelPath; - if (modelName == "llama_tp4_pp1_cp1" || modelName == "llama_tp1_pp4_cp1" || modelName == "llama_tp2_pp2_cp1") - { - if (modelName == "llama_tp4_pp1_cp1") - { - modelPath = LLAMA_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp4-pp1-cp1-gpu"; - } - else if (modelName == "llama_tp1_pp4_cp1") - { - modelPath = LLAMA_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp1-pp4-cp1-gpu"; - deviceIds = std::vector{3, 2, 1, 0}; - } - else if (modelName == "llama_tp2_pp2_cp1") - { - modelPath = LLAMA_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp2-pp2-cp1-gpu"; - deviceIds = std::vector{2, 3, 0, 1}; - } - } - - // For llama model, only run for multiple GPUs - // This is detected by setting an env variable when running the test - char const* val = getenv("RUN_LLAMA_MULTI_GPU"); - if (val == NULL) - { - GTEST_SKIP() << "Skipping Llama test"; - } - else - { - // Check that it was launched with right number of MPI ranks - if (COMM_SESSION.getSize() != 4) - { - // No orchestrator, need worldSize to match TP*PP - FAIL() << "Leader mode and world size is not equal to 4"; - } - } - - if (deviceIds.has_value()) - { - auto parallelConfig = executorConfig.getParallelConfig().value_or(ParallelConfig()); - - parallelConfig.setDeviceIds(deviceIds.value()); - executorConfig.setParallelConfig(parallelConfig); - } - auto executor = Executor(modelPath, ModelType::kDECODER_ONLY, executorConfig); - - // Since this is leader mode, all ranks should participate - EXPECT_TRUE(executor.isParticipant()); - - // Create the request - SizeType32 maxNewTokens = 50; - VecTokens inputTokens{1, 2, 3, 4}; - auto request - = Request(inputTokens, maxNewTokens, false, tensorrt_llm::executor::SamplingConfig(beamWidth), outConfig); - auto requestStreaming - = Request(inputTokens, maxNewTokens, true, tensorrt_llm::executor::SamplingConfig(beamWidth), outConfig); - - // Leader enqueues requests and wait for responses - if (executor.canEnqueueRequests()) - { - auto requestId = executor.enqueueRequest(request); - auto requestId2 = executor.enqueueRequest(request); - auto requestId3 = executor.enqueueRequest(requestStreaming); - auto requestId4 = executor.enqueueRequest(requestStreaming); - - int32_t numFinished = 0; - int iter = 0; - SizeType32 numResponses = 0; - while (numFinished < 4 && iter < mMaxWaitMs) - { - std::chrono::milliseconds waitTime(1); - auto responses = executor.awaitResponses(waitTime); - for (auto& response : responses) - { - numResponses++; - if (!response.hasError()) - { - auto result = response.getResult(); - numFinished += result.isFinal; - } - else - { - FAIL() << "Did not expect errors"; - } - } - ++iter; - } - EXPECT_EQ(numFinished, 4); - EXPECT_LT(iter, mMaxWaitMs); - } - else - { - // Check that non-leader cannot enqueue requests - EXPECT_THROW({ auto reqId = executor.enqueueRequest(request); }, tensorrt_llm::common::TllmException); - EXPECT_THROW({ auto responses = executor.awaitResponses(); }, tensorrt_llm::common::TllmException); - EXPECT_THROW({ auto numResp = executor.getNumResponsesReady(); }, tensorrt_llm::common::TllmException); - EXPECT_THROW({ executor.cancelRequest(1); }, tensorrt_llm::common::TllmException); - EXPECT_THROW({ auto stats = executor.getLatestIterationStats(); }, tensorrt_llm::common::TllmException); - EXPECT_THROW({ auto stats = executor.getLatestRequestStats(); }, tensorrt_llm::common::TllmException); - } -} - -TEST_F(GptExecutorTest, validateParallelConfig) -{ - - auto trtEnginePath = (GPT_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp1-pp1-cp1-gpu"); - { - auto executorConfig = ExecutorConfig(); - auto executor = Executor(trtEnginePath, ModelType::kDECODER_ONLY, executorConfig); - } - - { - std::string expectedErrMsg = "OrchestratorConfig must be set"; - try - { - auto executorConfig = ExecutorConfig(); - auto parallelConfig = ParallelConfig(CommunicationType::kMPI, CommunicationMode::kORCHESTRATOR); - executorConfig.setParallelConfig(parallelConfig); - auto executor = Executor(trtEnginePath, ModelType::kDECODER_ONLY, executorConfig); - FAIL() << "Expected TllmException"; - } - catch (tc::TllmException& e) - { - EXPECT_THAT(e.what(), testing::HasSubstr(expectedErrMsg)); - } - catch (std::exception const& e) - { - FAIL() << "Expected TllmException"; - } - } -} - -TEST_P(TimeoutTest, TimeoutStreamingTest) -{ - auto const modelName = std::get<0>(GetParam()); - auto const useOrchestratorMode = std::get<1>(GetParam()); - auto const beamWidth = std::get<2>(GetParam()); - - auto executorConfig = ExecutorConfig(beamWidth); - std::filesystem::path modelPath; - bool isMultiGpu{false}; - std::optional> deviceIds = std::nullopt; - - if (modelName == "llama_tp4_pp1_cp1" || modelName == "llama_tp1_pp4_cp1" || modelName == "llama_tp2_pp2_cp1") - { - isMultiGpu = true; - if (modelName == "llama_tp4_pp1_cp1") - { - modelPath = LLAMA_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp4-pp1-cp1-gpu"; - } - else if (modelName == "llama_tp1_pp4_cp1") - { - modelPath = LLAMA_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp1-pp4-cp1-gpu"; - deviceIds = std::vector{3, 2, 1, 0}; - } - else if (modelName == "llama_tp2_pp2_cp1") - { - modelPath = LLAMA_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp2-pp2-cp1-gpu"; - deviceIds = std::vector{2, 3, 0, 1}; - } - } - if (modelName == "llama_tp1_pp1_cp1") - { - modelPath = LLAMA_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp1-pp1-cp1-gpu"; - } - // For llama model, only run for multiple GPUs - // This is detected by setting an env variable when running the test - char const* val = getenv("RUN_LLAMA_MULTI_GPU"); - if (val == NULL && isMultiGpu) - { - GTEST_SKIP() << "Skipping MultiGpu tests"; - } - if (val != NULL && !isMultiGpu) - { - GTEST_SKIP() << "Skipping SingleGpu tests"; - } - if (val != NULL && isMultiGpu) - { - // Check that it was launched with right number of MPI ranks - if (!useOrchestratorMode && COMM_SESSION.getSize() != 4) - { - // No orchestrator, need worldSize to match TP*PP - FAIL() << "Leader mode and world size is not equal to 4"; - } - if (useOrchestratorMode && COMM_SESSION.getSize() != 1) - { - // No orchestrator, need worldSize to match TP*PP - FAIL() << "Orchestrator mode and World size is not equal to 1"; - } - } - - if (useOrchestratorMode) - { - auto orchestratorConfig = OrchestratorConfig(true, PathUtil::EXECUTOR_WORKER_PATH()); - auto parallelConfig = ParallelConfig(CommunicationType::kMPI, - useOrchestratorMode ? CommunicationMode::kORCHESTRATOR : CommunicationMode::kLEADER, std::nullopt, - std::nullopt, orchestratorConfig); - executorConfig.setParallelConfig(parallelConfig); - if (deviceIds.has_value()) - { - parallelConfig.setDeviceIds(deviceIds.value()); - } - executorConfig.setParallelConfig(parallelConfig); - } - else - { - if (deviceIds.has_value()) - { - auto parallelConfig = executorConfig.getParallelConfig().value_or(ParallelConfig()); - parallelConfig.setDeviceIds(deviceIds.value()); - executorConfig.setParallelConfig(parallelConfig); - } - } - auto executor = Executor(modelPath, ModelType::kDECODER_ONLY, executorConfig); - - SizeType32 constexpr maxNewTokens = 10; - // create 1 request that times out immediately - // momentarily we don't cancel requests before forwardAsync so it will get scheduled for at least 1 forward - VecTokens immediateCancelTokens{1, 2, 3, 4}; - auto immediateCancelRequest - = Request(immediateCancelTokens, maxNewTokens, true, tensorrt_llm::executor::SamplingConfig(beamWidth)); - immediateCancelRequest.setReturnAllGeneratedTokens(true); - immediateCancelRequest.setAllottedTimeMs(std::chrono::milliseconds(0)); - SizeType32 constexpr immediateCancelMinLength = 0; - SizeType32 constexpr immediateCancelMaxLength = 1; - - // create 1 request that times out during the first forward - VecTokens oneForwardTokens{11, 12, 13, 14}; - auto oneForwardRequest - = Request(oneForwardTokens, maxNewTokens, true, tensorrt_llm::executor::SamplingConfig(beamWidth)); - oneForwardRequest.setReturnAllGeneratedTokens(true); - oneForwardRequest.setAllottedTimeMs(std::chrono::milliseconds(1)); - SizeType32 constexpr oneForwardlMinLength = 0; - SizeType32 constexpr oneForwardlMaxLength = 1; - - // Create the request that finishes by the number of tokens - VecTokens finishedTokens{101, 102, 103, 104}; - auto finishedRequest - = Request(finishedTokens, maxNewTokens, true, tensorrt_llm::executor::SamplingConfig(beamWidth)); - finishedRequest.setReturnAllGeneratedTokens(true); - finishedRequest.setAllottedTimeMs(std::chrono::milliseconds(5000)); - SizeType32 constexpr finishedMinLength = 5; - SizeType32 constexpr finishedMaxLength = maxNewTokens; - - std::vector referenceFinishReasons - = {FinishReason::kTIMED_OUT, FinishReason::kTIMED_OUT, FinishReason::kLENGTH}; - std::vector minLengths = {immediateCancelMinLength, oneForwardlMinLength, finishedMinLength}; - std::vector maxLengths = {immediateCancelMaxLength, oneForwardlMaxLength, finishedMaxLength}; - // workaround because the last response will be empty, but we want to have at least *some* responses surpass the - // minLength - std::vector achievedLength = {0, 0, 0}; - SizeType32 itNr{0}; - - if (executor.canEnqueueRequests()) - { - - std::vector requests = {immediateCancelRequest, oneForwardRequest, finishedRequest}; - auto requestIds = executor.enqueueRequests(requests); - - auto numFinished = 0; - - while (numFinished < static_cast(requests.size())) - { - itNr++; - std::chrono::milliseconds waitTime(mMaxWaitMs); - auto responses = executor.awaitResponses(requestIds, waitTime); - for (auto const& response : responses) - { - for (auto const& responseIt : response) - { - auto const reqId = responseIt.getRequestId(); - if (responseIt.hasError()) - { - // Allow response with error only if awaitResponse processed a terminated request id - std::string err - = "ReqId " + std::to_string(reqId) + " has already been processed and was terminated."; - if (responseIt.getErrorMsg() != err) - { - TLLM_THROW("Request id %lu encountered error: %s", reqId, responseIt.getErrorMsg().c_str()); - } - continue; - } - - auto const& result = responseIt.getResult(); - if (result.isFinal) - { - requestIds.erase(std::remove(requestIds.begin(), requestIds.end(), reqId), requestIds.end()); - numFinished++; - } - - auto const finishReason = result.finishReasons; - auto const actualResponse = result.outputTokenIds; - TLLM_LOG_DEBUG("reqId %d finished %d", reqId, result.isFinal); - TLLM_LOG_DEBUG("actual response:"); - - for (auto const& beam : actualResponse) - { - std::string tokenStr; - for (auto tok : beam) - { - tokenStr += std::to_string(tok) + " "; - } - TLLM_LOG_DEBUG("%s", tokenStr.c_str()); - } - - TLLM_LOG_DEBUG( - "beams' length must be in range [%d, %d]", minLengths[reqId - 1], maxLengths[reqId - 1]); - - if (result.isFinal) - { - TLLM_LOG_DEBUG("finishReason"); - std::string reasonStr; - for (auto const reason : finishReason) - { - // cast for easier visibility during debugging - EXPECT_EQ(static_cast(reason), static_cast(referenceFinishReasons[reqId - 1])); - reasonStr += std::to_string(static_cast(reason)) + " "; - } - TLLM_LOG_DEBUG("%s", reasonStr.c_str()); - } - - EXPECT_EQ(beamWidth, actualResponse.size()); - for (int beam = 0; beam < beamWidth; beam++) - { - EXPECT_LE(actualResponse.at(beam).size(), maxLengths[reqId - 1]) << "for request " << reqId; - achievedLength[reqId - 1] = std::max( - achievedLength[reqId - 1], static_cast(actualResponse.at(beam).size())); - } - } - } - } - - for (int reqIt = 0; reqIt < achievedLength.size(); ++reqIt) - { - EXPECT_GE(achievedLength[reqIt], minLengths[reqIt]) - << "request " << reqIt + 1 << " has not achieved min lengths"; - } - } -} - -TEST_P(TimeoutTest, TimeoutNonstreamingTest) -{ - auto const modelName = std::get<0>(GetParam()); - auto const useOrchestratorMode = std::get<1>(GetParam()); - auto const beamWidth = std::get<2>(GetParam()); - - std::optional> deviceIds = std::nullopt; - - auto executorConfig = ExecutorConfig(beamWidth); - std::filesystem::path modelPath; - bool isMultiGpu{false}; - if (modelName == "llama_tp4_pp1_cp1" || modelName == "llama_tp1_pp4_cp1" || modelName == "llama_tp2_pp2_cp1") - { - isMultiGpu = true; - if (modelName == "llama_tp4_pp1_cp1") - { - modelPath = LLAMA_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp4-pp1-cp1-gpu"; - } - else if (modelName == "llama_tp1_pp4_cp1") - { - modelPath = LLAMA_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp1-pp4-cp1-gpu"; - deviceIds = std::vector{3, 2, 1, 0}; - } - else if (modelName == "llama_tp2_pp2_cp1") - { - modelPath = LLAMA_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp2-pp2-cp1-gpu"; - deviceIds = std::vector{2, 3, 0, 1}; - } - } - if (modelName == "llama_tp1_pp1_cp1") - { - modelPath = LLAMA_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp1-pp1-cp1-gpu"; - } - // For llama model, only run for multiple GPUs - // This is detected by setting an env variable when running the test - char const* val = getenv("RUN_LLAMA_MULTI_GPU"); - if (val == NULL && isMultiGpu) - { - GTEST_SKIP() << "Skipping MultiGpu tests"; - } - if (val != NULL && !isMultiGpu) - { - GTEST_SKIP() << "Skipping SingleGpu tests"; - } - if (val != NULL && isMultiGpu) - { - // Check that it was launched with right number of MPI ranks - if (!useOrchestratorMode && COMM_SESSION.getSize() != 4) - { - // No orchestrator, need worldSize to match TP*PP - FAIL() << "Leader mode and world size is not equal to 4"; - } - if (useOrchestratorMode && COMM_SESSION.getSize() != 1) - { - // No orchestrator, need worldSize to match TP*PP - FAIL() << "Orchestrator mode and World size is not equal to 1"; - } - } - - if (useOrchestratorMode) - { - auto orchestratorConfig = OrchestratorConfig(true, PathUtil::EXECUTOR_WORKER_PATH()); - auto parallelConfig = ParallelConfig(CommunicationType::kMPI, - useOrchestratorMode ? CommunicationMode::kORCHESTRATOR : CommunicationMode::kLEADER, std::nullopt, - std::nullopt, orchestratorConfig); - executorConfig.setParallelConfig(parallelConfig); - if (deviceIds.has_value()) - { - parallelConfig.setDeviceIds(deviceIds.value()); - } - executorConfig.setParallelConfig(parallelConfig); - } - else - { - if (deviceIds.has_value()) - { - auto parallelConfig = executorConfig.getParallelConfig().value_or(ParallelConfig()); - parallelConfig.setDeviceIds(deviceIds.value()); - executorConfig.setParallelConfig(parallelConfig); - } - } - auto executor = Executor(modelPath, ModelType::kDECODER_ONLY, executorConfig); - - SizeType32 constexpr maxNewTokens = 5; - // create 1 request that times out immediately - // momentarily we don't cancel requests before forwardAsync so it will get scheduled for at least 1 forward - VecTokens immediateCancelTokens{1, 2, 3, 4}; - auto immediateCancelRequest - = Request(immediateCancelTokens, maxNewTokens, false, tensorrt_llm::executor::SamplingConfig(beamWidth)); - immediateCancelRequest.setAllottedTimeMs(std::chrono::milliseconds(0)); - std::vector> immediateCancelResponse = {immediateCancelTokens, immediateCancelTokens}; - - // create 1 request that times out during the first forward - VecTokens oneForwardTokens{11, 12, 13, 14}; - auto oneForwardRequest - = Request(oneForwardTokens, maxNewTokens, false, tensorrt_llm::executor::SamplingConfig(beamWidth)); - oneForwardRequest.setAllottedTimeMs(std::chrono::milliseconds(1)); - std::vector> oneForwardResponse = {oneForwardTokens, oneForwardTokens}; - - // Create the request that finishes by the number of tokens - VecTokens finishedTokens{101, 102, 103, 104}; - auto finishedRequest - = Request(finishedTokens, maxNewTokens, false, tensorrt_llm::executor::SamplingConfig(beamWidth)); - finishedRequest.setAllottedTimeMs(std::chrono::milliseconds(6000)); - std::vector> finishedReponse - = {{101, 102, 103, 104, 49849, 225, 49849, 232, 55742}, {101, 102, 103, 104, 49849, 225, 49849, 232, 29082}}; - - // assume responses will come in FIFO order - std::vector refResponses = {immediateCancelResponse, oneForwardResponse, finishedReponse}; - std::vector referenceFinishReasons - = {FinishReason::kTIMED_OUT, FinishReason::kTIMED_OUT, FinishReason::kLENGTH}; - if (executor.canEnqueueRequests()) - { - - std::vector requests = {immediateCancelRequest, oneForwardRequest, finishedRequest}; - auto requestIds = executor.enqueueRequests(requests); - - std::chrono::milliseconds waitTime(mMaxWaitMs); - auto responses = executor.awaitResponses(requestIds, waitTime); - for (auto const& response : responses) - { - for (auto const& responseIt : response) - { - auto const reqId = responseIt.getRequestId(); - if (responseIt.hasError()) - { - TLLM_THROW("Request id %lu encountered error: %s", reqId, responseIt.getErrorMsg().c_str()); - } - - auto const& result = responseIt.getResult(); - - auto const finishReason = result.finishReasons; - auto const actualResponse = result.outputTokenIds; - TLLM_LOG_DEBUG("reqId %d finished %d", reqId, result.isFinal); - TLLM_LOG_DEBUG("actual response:"); - - for (auto const& beam : actualResponse) - { - std::string tokenStr; - for (auto tok : beam) - { - tokenStr += std::to_string(tok) + " "; - } - TLLM_LOG_DEBUG("%s", tokenStr.c_str()); - } - - TLLM_LOG_DEBUG("reference:"); - auto referenceResponse = refResponses[reqId - 1]; - for (auto const& beam : referenceResponse) - { - std::string tokenStr; - for (auto tok : beam) - { - tokenStr += std::to_string(tok) + " "; - } - TLLM_LOG_DEBUG("%s", tokenStr.c_str()); - } - - if (result.isFinal) - { - TLLM_LOG_DEBUG("finishReason"); - std::string reasonStr; - for (auto const reason : finishReason) - { - // cast for easier visibility during debugging - EXPECT_EQ(static_cast(reason), static_cast(referenceFinishReasons[reqId - 1])); - reasonStr += std::to_string(static_cast(reason)) + " "; - } - TLLM_LOG_DEBUG("%s", reasonStr.c_str()); - } - - EXPECT_EQ(beamWidth, actualResponse.size()); - for (int beam = 0; beam < beamWidth; beam++) - { - EXPECT_EQ(referenceResponse.at(beam).size(), actualResponse.at(beam).size()); - EXPECT_THAT(actualResponse.at(beam), testing::ElementsAreArray(referenceResponse.at(beam))); - } - } - } - } -} - -INSTANTIATE_TEST_SUITE_P(GptExecutorTest, ParamTest, - testing::Combine( // - testing::Values(false, true), // streaming - testing::Values(false, true), // excludeInputFromOutput - testing::Values(1, 2) // beamWidth - ), - generateTestName); - -INSTANTIATE_TEST_SUITE_P(GptExecutorTest, ParamStatsTest, - testing::Combine( // - testing::Values(0, 1000), // iterStatsMaxIterations - testing::Values(false, true) // useOrchestratorMode - ), - generateTestNameStats); - -INSTANTIATE_TEST_SUITE_P(LlamaExecutorTest, ParamCancelReqTest, - testing::Combine( // - testing::Values(false, true), // useOrchestratorMode - testing::Values(1, 2), // beamWidth - testing::Values("llama_tp1_pp4_cp1", "llama_tp4_pp1_cp1", "llama_tp2_pp2_cp1") // modelName - ), - generateTestNameCancelReq); - -INSTANTIATE_TEST_SUITE_P(LlamaExecutorTest, TimeoutTest, - testing::Combine( // - testing::Values("llama_tp1_pp4_cp1", "llama_tp4_pp1_cp1", "llama_tp1_pp1_cp1"), // modelName - testing::Values(false, true), // useOrchestratorMode - testing::Values(2) // beamWidth - ), - generateTestNameTimeoutTest); - -INSTANTIATE_TEST_SUITE_P(LlamaExecutorTest, LeaderApiUsageTest, - testing::Combine( // - testing::Values("llama_tp1_pp4_cp1", "llama_tp4_pp1_cp1", "llama_tp2_pp2_cp1") // modelName - ), - generateTestNameLeaderApiUsage); - -INSTANTIATE_TEST_SUITE_P(GptExecutorTest, AllParamsTest, - testing::Combine( // - testing::Values(false, true), // streaming - testing::Values(1, 2), // beamWidth - testing::Values(true), // computeLogProbs - testing::Values(false, true), // excludeInputInOutput - testing::Values(true), // returnContextLogits - testing::Values(true), // returnGenerationLogits - testing::Values("gpt"), // modelName - testing::Values(false, true), // useOrchestratorMode - testing::Values(false, true), // returnAllGeneratedTokens - testing::Values(1, 2) // numReturnSequences - ), - generateTestNameAllParams); - -INSTANTIATE_TEST_SUITE_P(LlamaExecutorTest, AllParamsTest, - testing::Combine( // - testing::Values(false, true), // streaming - testing::Values(1, 2), // beamWidth - testing::Values(true), // computeLogProbs - testing::Values(false, true), // excludeInputInOutput - testing::Values(false), // returnContextLogits - testing::Values(true), // returnGenerationLogits - testing::Values("llama_tp1_pp4_cp1", "llama_tp4_pp1_cp1", "llama_tp2_pp2_cp1"), // modelName - testing::Values(false, true), // useOrchestratorMode - testing::Values(false), // returnAllGeneratedTokens - testing::Values(1) // numReturnSequences - ), - generateTestNameAllParams); - -INSTANTIATE_TEST_SUITE_P(LlamaMultiExecutorTest, AllParamsTest, - testing::Combine( // - testing::Values(false, true), // streaming - testing::Values(1, 2), // beamWidth - testing::Values(false), // computeLogProbs - testing::Values(false, true), // excludeInputInOutput - testing::Values(false), // returnContextLogits - testing::Values(false), // returnGenerationLogits - testing::Values("llama_tp1_pp2_cp1"), // modelName - testing::Values(false), // useOrchestratorMode - testing::Values(false), // returnAllGeneratedTokens - testing::Values(1) // numReturnSequences - ), - generateTestNameAllParams); - -INSTANTIATE_TEST_SUITE_P(MedusaExecutorTest, AllParamsTest, - testing::Combine( // - testing::Values(false, true), // streaming - testing::Values(1), // beamWidth - testing::Values(false), // computeLogProbs - testing::Values(false, true), // excludeInputInOutput - testing::Values(false), // returnContextLogits - testing::Values(false), // returnGenerationLogits - testing::Values("medusa"), // modelName - testing::Values(false, true), // useOrchestratorMode - testing::Values(false), // returnAllGeneratedTokens - testing::Values(1) // numReturnSequences - ), - generateTestNameAllParams); - -// Disable some of ChatGLM's tests since they are the same as gpt's. -INSTANTIATE_TEST_SUITE_P(ChatGlmExecutorTest, AllParamsTest, - testing::Combine( // - testing::Values(false), // streaming - testing::Values(1, 2), // beamWidth - testing::Values(false), // computeLogProbs - testing::Values(false), // excludeInputInOutput - testing::Values(false), // returnContextLogits - testing::Values(false), // returnGenerationLogits - testing::Values("chatglm"), // modelName - testing::Values(false), // useOrchestratorMode - testing::Values(false), // returnAllGeneratedTokens - testing::Values(1, 2) // numReturnSequences - ), - generateTestNameAllParams); - -// ChatGlm0 Test is for glm-10b. -INSTANTIATE_TEST_SUITE_P(ChatGlm0ExecutorTest, AllParamsTest, - testing::Combine( // - testing::Values(false), // streaming - testing::Values(1), // beamWidth - testing::Values(false), // computeLogProbs - testing::Values(false), // excludeInputInOutput - testing::Values(false), // returnContextLogits - testing::Values(false), // returnGenerationLogits - testing::Values("glm"), // modelName - testing::Values(false), // useOrchestratorMode - testing::Values(false), // returnAllGeneratedTokens - testing::Values(1) // numReturnSequences - ), - generateTestNameAllParams); - -INSTANTIATE_TEST_SUITE_P(ChatGlm2ExecutorTest, AllParamsTest, - testing::Combine( // - testing::Values(false), // streaming - testing::Values(1), // beamWidth - testing::Values(false), // computeLogProbs - testing::Values(false), // excludeInputInOutput - testing::Values(false), // returnContextLogits - testing::Values(false), // returnGenerationLogits - testing::Values("chatglm2"), // modelName - testing::Values(false), // useOrchestratorMode - testing::Values(false), // returnAllGeneratedTokens - testing::Values(1) // numReturnSequences - ), - generateTestNameAllParams); - -INSTANTIATE_TEST_SUITE_P(ChatGlm3ExecutorTest, AllParamsTest, - testing::Combine( // - testing::Values(false), // streaming - testing::Values(1), // beamWidth - testing::Values(false), // computeLogProbs - testing::Values(false), // excludeInputInOutput - testing::Values(false), // returnContextLogits - testing::Values(false), // returnGenerationLogits - testing::Values("chatglm3"), // modelName - testing::Values(false), // useOrchestratorMode - testing::Values(false), // returnAllGeneratedTokens - testing::Values(1) // numReturnSequences - ), - generateTestNameAllParams); - -INSTANTIATE_TEST_SUITE_P(LlamaExecutorTest, LogitsProcParamsTest, - testing::Combine( // - testing::Values( - "llama_tp1_pp1_cp1", "llama_tp4_pp1_cp1", "llama_tp2_pp2_cp1", "llama_tp1_pp4_cp1"), // modelName - testing::Values(false, true), // batched - testing::Values(false, true) // replicated - ), - generateTestNameLogitsProc); - -INSTANTIATE_TEST_SUITE_P(GptExecutorGuidedDecodingTest, GuidedDecodingParamsTest, - testing::Combine(testing::Values("gpt")), generateTestNameGuidedDecoding); - -INSTANTIATE_TEST_SUITE_P(LlamaExecutorGuidedDecodingTest, GuidedDecodingParamsTest, - testing::Combine( - testing::Values("llama_tp1_pp1_cp1", "llama_tp4_pp1_cp1", "llama_tp2_pp2_cp1", "llama_tp1_pp4_cp1")), - generateTestNameGuidedDecoding); diff --git a/cpp/tests/e2e_tests/executor/executorTest.h b/cpp/tests/e2e_tests/executor/executorTest.h deleted file mode 100644 index 7866a6992266..000000000000 --- a/cpp/tests/e2e_tests/executor/executorTest.h +++ /dev/null @@ -1,58 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2023-2025 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. - */ - -#pragma once - -#include "tensorrt_llm/plugins/api/tllmPlugin.h" -#include "tensorrt_llm/runtime/tllmLogger.h" -#include "tests/utils/common.h" - -#include -#include - -#include - -namespace tensorrt_llm::testing -{ - -class GptExecutorTest : public ::testing::Test // NOLINT(cppcoreguidelines-pro-type-member-init) -{ -public: - using SizeType32 = tensorrt_llm::testing::SizeType32; - -protected: - void SetUp() override - { - mDeviceCount = tensorrt_llm::common::getDeviceCount(); - if (mDeviceCount == 0) - { - GTEST_SKIP() << "No GPUs found"; - } - - mLogger = std::make_shared(); - initTrtLlmPlugins(mLogger.get()); - } - - void TearDown() override {} - - int mDeviceCount{}; - std::shared_ptr mLogger{}; - SizeType32 mMaxWaitMs = 300000; - SizeType32 mTrigWarnMs = 10000; -}; - -} // namespace tensorrt_llm::testing diff --git a/cpp/tests/resources/scripts/build_chatglm_engines.py b/cpp/tests/resources/scripts/build_chatglm_engines.py deleted file mode 100644 index abe187307604..000000000000 --- a/cpp/tests/resources/scripts/build_chatglm_engines.py +++ /dev/null @@ -1,193 +0,0 @@ -#!/usr/bin/env python3 -# 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 argparse -import os -import platform -import shutil -import sys -import typing -from pathlib import Path -from typing import Optional - -from build_engines_utils import run_command, wincopy - -import tensorrt_llm.bindings as _tb -from tensorrt_llm.bindings.internal.testing import ModelSpec - -resources_dir = Path(__file__).parent.resolve().parent -model_dir = resources_dir / "models" -chatglm_example_dir = Path("examples/chatglm") -bCopyModel = True # "False" to remove redundant copy of model from model_cache - - -def convert_ckpt(model_dir: str, output_dir: str, world_size: int): - if os.path.exists(output_dir): - print('Skip ckpt convert - output already exists') - return - - convert_cmd = [ - sys.executable, - str(chatglm_example_dir / "convert_checkpoint.py"), "--dtype=float16", - f"--model_dir={model_dir}", f"--output_dir={output_dir}", - f"--tp_size={world_size}" - ] - run_command(convert_cmd) - - -def build_engine(ckpt_dir: str, - engine_dir: str, - is_ifb: bool = False, - is_chatglm_6b_or_glm_10b: bool = False): - if os.path.exists(engine_dir): - print('Skip engine build - output already exists') - return - - build_cmd = [ - "trtllm-build", - f"--checkpoint_dir={ckpt_dir}", - f"--output_dir={engine_dir}", - "--log_level=error", - "--max_batch_size=8", - "--max_beam_width=2", - "--max_input_len=256", - "--max_seq_len=384", - "--gpt_attention_plugin=float16", - "--gemm_plugin=float16", - ] - if is_ifb: - build_cmd.extend([ - "--remove_input_padding=enable", - "--paged_kv_cache=enable", - "--context_fmha=enable", - "--use_paged_context_fmha=enable", - ]) - else: - build_cmd.extend([ - "--remove_input_padding=disable", - "--paged_kv_cache=disable", - ]) - - if is_chatglm_6b_or_glm_10b: - print("Disable Context FMHA for ChatGLM-6B and GLM-10B") - build_cmd.extend(["--context_fmha=disable"]) - - run_command(build_cmd) - - -def build_engines(model_cache: typing.Optional[str] = None, - world_size: int = 1, - clean: Optional[bool] = False): - - for model_name in [ - "chatglm-6b", "chatglm2-6b", "chatglm3-6b", "glm-10b", "glm-4-9b", - "chatglm3-6b-32k" - ]: - is_chatglm_6b_or_glm_10b = model_name in ["chatglm-6b", "glm-10b"] - if model_cache and (Path(model_cache) / model_name).is_dir(): - model_cache_dir = Path(model_cache) / model_name - if bCopyModel or model_name == "chatglm-6b": - print("Copy model from model_cache") - hf_dir = model_dir / model_name - if platform.system() == "Windows": - wincopy(source=str(model_cache_dir), - dest=model_name, - isdir=True, - cwd=model_dir) - else: - run_command(["rsync", "-rlptD", - str(model_cache_dir), "."], - cwd=model_dir) - else: - print("Use model from model_cache directly except ChatGLM-6B") - hf_dir = Path(model_cache) - - else: - hf_dir = model_dir / model_name - if not hf_dir.is_dir(): - print("Clone model from HF") - run_command( - [ - "git", "clone", - f"https://huggingface.co/THUDM/{model_name}", model_name - ], - cwd=model_dir, - ) - - # Build engines - print(f"Building {model_name}") - ckpt_dir = Path(model_dir) / "c-model" / model_name - if clean: - print('clean up ckpt folder ', ckpt_dir) - if ckpt_dir.is_dir(): - shutil.rmtree(ckpt_dir, ignore_errors=True) - - # Fix HF error for ChatGLM-6B / GLM-4-9B / ChatGLM2-6B / ChatGLM3-6B-32K, hope to remove this in the future - if model_name in [ - "chatglm-6b", "glm-4-9b", "chatglm2-6b", "chatglm3-6b-32k" - ]: - shutil.copy( - chatglm_example_dir / f"{model_name}/tokenization_chatglm.py", - hf_dir, - ) - - convert_ckpt(hf_dir, ckpt_dir, world_size) - - model_spec_obj = ModelSpec('input_tokens.npy', _tb.DataType.HALF) - model_spec_obj.set_kv_cache_type(_tb.KVCacheType.CONTINUOUS) - model_spec_obj.use_gpt_plugin() - engine_dir = Path( - model_dir - ) / "rt_engine" / model_name / model_spec_obj.get_model_path( - ) / "tp1-pp1-cp1-gpu" - if clean: - print('clean up engine folder ', engine_dir) - if engine_dir.is_dir(): - shutil.rmtree(engine_dir, ignore_errors=True) - build_engine(ckpt_dir, engine_dir, False, is_chatglm_6b_or_glm_10b) - - model_spec_obj.use_packed_input() - model_spec_obj.set_kv_cache_type(_tb.KVCacheType.PAGED) - engine_dir = Path( - model_dir - ) / "rt_engine" / model_name / model_spec_obj.get_model_path( - ) / "tp1-pp1-cp1-gpu" - if clean: - print('clean up engine folder ', engine_dir) - if engine_dir.is_dir(): - shutil.rmtree(engine_dir, ignore_errors=True) - build_engine(ckpt_dir, engine_dir, True, is_chatglm_6b_or_glm_10b) - - print("Done") - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument("--model_cache", - type=str, - help="Directory where models are stored") - - parser.add_argument('--world_size', - type=int, - default=1, - help='world size, only support tensor parallelism now') - - parser.add_argument('--clean', - action='store_true', - default=False, - help='Clean target folders before building engines') - - build_engines(**vars(parser.parse_args())) diff --git a/cpp/tests/resources/scripts/build_eagle_engines.py b/cpp/tests/resources/scripts/build_eagle_engines.py deleted file mode 100755 index 8b10698a603b..000000000000 --- a/cpp/tests/resources/scripts/build_eagle_engines.py +++ /dev/null @@ -1,143 +0,0 @@ -#!/usr/bin/env python3 -# SPDX-FileCopyrightText: Copyright (c) 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 argparse as _arg -import pathlib as _pl -import platform as _pf -import sys as _sys - -from build_engines_utils import run_command, wincopy - -import tensorrt_llm.bindings as _tb -from tensorrt_llm.bindings.internal.testing import ModelSpec - - -def build_engine(base_model_dir: _pl.Path, eagle_model_dir: _pl.Path, - engine_dir: _pl.Path, build_base_model: bool, *args): - - if build_base_model: - checkpoint_path = "examples/models/core/llama/convert_checkpoint.py" - else: - checkpoint_path = "examples/eagle/convert_checkpoint.py" - - covert_cmd = [_sys.executable, checkpoint_path] + ( - ['--model_dir', str(base_model_dir)] if base_model_dir else []) + [ - '--output_dir', str(engine_dir), '--dtype=float16' - ] + list(args) - - if not build_base_model: - covert_cmd += [ - '--eagle_model_dir', - str(eagle_model_dir), '--num_eagle_layers=4', '--max_draft_len=63' - ] - - run_command(covert_cmd) - - build_args = ["trtllm-build"] + ( - ['--checkpoint_dir', str(engine_dir)] if engine_dir else []) + [ - '--output_dir', - str(engine_dir), - '--gemm_plugin=float16', - '--max_batch_size=8', - '--max_input_len=12', - '--max_seq_len=140', - '--log_level=error', - '--paged_kv_cache=enable', - '--remove_input_padding=enable', - '--use_paged_context_fmha=enable', - ] - - if not build_base_model: - build_args += ['--speculative_decoding_mode=eagle'] - - run_command(build_args) - - -def build_engines(model_cache: str): - resources_dir = _pl.Path(__file__).parent.resolve().parent - models_dir = resources_dir / 'models' - model_name = 'vicuna-7b-eagle' - base_model_name = 'vicuna-7b-v1.3' - eagle_model_name = 'EAGLE-Vicuna-7B-v1.3' - - if model_cache: - print(f"Copy model from {model_cache}") - base_model_cache_dir = _pl.Path(model_cache) / base_model_name - eagle_cache_dir = _pl.Path(model_cache) / eagle_model_name - assert base_model_cache_dir.is_dir(), base_model_cache_dir - assert eagle_cache_dir.is_dir(), eagle_cache_dir - - if _pf.system() == "Windows": - wincopy(source=str(base_model_cache_dir), - dest=base_model_name, - isdir=True, - cwd=models_dir) - wincopy(source=str(eagle_cache_dir), - dest=eagle_model_name, - isdir=True, - cwd=models_dir) - else: - run_command(["rsync", "-rlptD", - str(base_model_cache_dir), "."], - cwd=models_dir) - run_command(["rsync", "-rlptD", - str(eagle_cache_dir), "."], - cwd=models_dir) - - base_model_dir = models_dir / base_model_name - eagle_model_dir = models_dir / eagle_model_name - assert base_model_dir.is_dir() - assert eagle_model_dir.is_dir() - - eagle_engine_dir = models_dir / 'rt_engine' / model_name - base_engine_dir = models_dir / 'rt_engine' / base_model_name - - model_spec_obj = ModelSpec('input_tokens.npy', _tb.DataType.HALF) - model_spec_obj.use_gpt_plugin() - model_spec_obj.set_kv_cache_type(_tb.KVCacheType.PAGED) - model_spec_obj.use_packed_input() - - base_full_engine_path = base_engine_dir / model_spec_obj.get_model_path( - ) / 'tp1-pp1-cp1-gpu' - print(f"\nBuilding fp16 engine at {str(base_full_engine_path)}") - build_engine(base_model_dir, - eagle_model_dir, - base_full_engine_path, - build_base_model=True) - - model_spec_obj = ModelSpec('input_tokens.npy', _tb.DataType.HALF) - model_spec_obj.use_gpt_plugin() - model_spec_obj.set_kv_cache_type(_tb.KVCacheType.PAGED) - model_spec_obj.use_packed_input() - model_spec_obj.use_eagle() - eagle_full_engine_path = eagle_engine_dir / model_spec_obj.get_model_path( - ) / 'tp1-pp1-cp1-gpu' - print(f"\nBuilding fp16 engine at {str(eagle_full_engine_path)}") - build_engine(base_model_dir, - eagle_model_dir, - eagle_full_engine_path, - build_base_model=False) - - print("Done.") - - -if __name__ == "__main__": - parser = _arg.ArgumentParser() - parser.add_argument("--model_cache", - type=str, - help="Directory where models are stored") - - build_engines(**vars(parser.parse_args())) diff --git a/cpp/tests/resources/scripts/build_enc_dec_engines.py b/cpp/tests/resources/scripts/build_enc_dec_engines.py deleted file mode 100644 index 7079916267d9..000000000000 --- a/cpp/tests/resources/scripts/build_enc_dec_engines.py +++ /dev/null @@ -1,187 +0,0 @@ -import os.path -from argparse import ArgumentParser -from dataclasses import dataclass, fields -from subprocess import run -from sys import stderr, stdout -from typing import List, Literal, Union - -split = os.path.split -join = os.path.join -dirname = os.path.dirname - - -@dataclass -class Arguments: - download: bool = False - dtype: Literal['float16', 'float32', 'bfloat16'] = 'float16' - - hf_repo_name: Literal[ - 'facebook/bart-large-cnn', 't5-small', - 'language_adapter-enc_dec_language_adapter'] = 'facebook/bart-large-cnn' - - model_cache: str = '/llm-models' - - tp: int = 1 - pp: int = 1 - - beams: str = '1' - gpus_per_node: int = 4 - debug: bool = False - - rm_pad: bool = True - gemm: bool = True - - max_new_tokens: int = 64 - - @property - def beams_tuple(self): - return eval(f'tuple([{self.beams}])') - - @property - def max_beam(self): - return max(self.beams_tuple) - - @property - def ckpt(self): - return self.hf_repo_name.split('/')[-1] - - @property - def base_dir(self): - return dirname(dirname(__file__)) - - @property - def data_dir(self): - return join(self.base_dir, 'data/enc_dec') - - @property - def models_dir(self): - return join(self.base_dir, 'models/enc_dec') - - @property - def hf_models_dir(self): - return join(self.model_cache, self.ckpt) - - @property - def trt_models_dir(self): - return join(self.models_dir, 'trt_models', self.ckpt) - - @property - def engines_dir(self): - return join(self.models_dir, 'trt_engines', self.ckpt, - f'{self.tp * self.pp}-gpu', self.dtype) - - @property - def model_type(self): - return self.ckpt.split('-')[0] - - def __post_init__(self): - parser = ArgumentParser() - for k in fields(self): - k = k.name - v = getattr(self, k) - if isinstance(v, bool): - parser.add_argument(f'--{k}', action='store_true') - else: - parser.add_argument(f'--{k}', default=v, type=type(v)) - - args = parser.parse_args() - for k, v in args._get_kwargs(): - setattr(self, k, v) - - -@dataclass -class RunCMDMixin: - args: Arguments - - def command(self) -> Union[str, List[str]]: - raise NotImplementedError - - def run(self): - cmd = self.command() - if cmd: - cmd = ' '.join(cmd) if isinstance(cmd, list) else cmd - print('+ ' + cmd) - run(cmd, shell='bash', stdout=stdout, stderr=stderr, check=True) - - -class DownloadHF(RunCMDMixin): - - def command(self): - args = self.args - return [ - 'git', 'clone', f'https://huggingface.co/{args.hf_repo_name}', - args.hf_models_dir - ] if args.download and args.model_type != 'language_adapter' else '' - - -class Convert(RunCMDMixin): - - def command(self): - args = self.args - return [ - f'python examples/models/core/enc_dec/convert_checkpoint.py', - f'--model_type {args.model_type}', - f'--model_dir {args.hf_models_dir}', - f'--output_dir {args.trt_models_dir}', - f'--tp_size {args.tp} --pp_size {args.pp}' - ] - - -class Build(RunCMDMixin): - - def command(self): - args = self.args - engine_dir = args.engines_dir - weight_dir = args.trt_models_dir - encoder_build = [ - f"trtllm-build --checkpoint_dir {join(weight_dir, 'encoder')}", - f"--output_dir {join(engine_dir, 'encoder')}", - f'--paged_kv_cache disable', - f'--max_beam_width {args.max_beam}', - f'--max_batch_size 8', - f'--max_input_len 512', - f'--gemm_plugin {args.dtype}', - f'--bert_attention_plugin {args.dtype}', - f'--gpt_attention_plugin {args.dtype}', - f'--remove_input_padding enable', - ] - - decoder_build = [ - f"trtllm-build --checkpoint_dir {join(weight_dir, 'decoder')}", - f"--output_dir {join(engine_dir, 'decoder')}", - f'--paged_kv_cache enable', - f'--max_beam_width {args.max_beam}', - f'--max_batch_size 8', - f'--max_seq_len 201', - f'--max_encoder_input_len 512', - f'--gemm_plugin {args.dtype}', - f'--bert_attention_plugin {args.dtype}', - f'--gpt_attention_plugin {args.dtype}', - f'--remove_input_padding enable', - '--max_input_len 1', - ] - - # t5 model with relative attention cannot use context_fmha - encoder_build.append(f'--context_fmha disable') - decoder_build.append(f'--context_fmha disable') - - # language adapter plugin leverages MOE plugin for static expert selection - if args.model_type == 'language_adapter': - encoder_build.append(f'--moe_plugin auto') - decoder_build.append(f'--moe_plugin auto') - else: - encoder_build.append(f'--moe_plugin disable') - decoder_build.append(f'--moe_plugin disable') - - encoder_build = ' '.join(encoder_build) - decoder_build = ' '.join(decoder_build) - ret = ' && '.join((encoder_build, decoder_build)) - return ret - - -if __name__ == "__main__": - # TODO: add support for more models / setup - args = Arguments() - DownloadHF(args).run() - Convert(args).run() - Build(args).run() diff --git a/cpp/tests/resources/scripts/build_engines_utils.py b/cpp/tests/resources/scripts/build_engines_utils.py deleted file mode 100644 index ad8525217e3a..000000000000 --- a/cpp/tests/resources/scripts/build_engines_utils.py +++ /dev/null @@ -1,64 +0,0 @@ -#!/usr/bin/env python3 -# 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 as _log -import os as _os -import pathlib as _pl -import subprocess as _sp -import typing as _tp - - -def run_command(command: _tp.Sequence[str], - *, - cwd=None, - timeout=None, - **kwargs) -> None: - _log.info("Running: cd %s && %s", str(cwd), " ".join(command)) - override_timeout = int(_os.environ.get("CPP_TEST_TIMEOUT_OVERRIDDEN", "-1")) - if override_timeout > 0 and (timeout is None or override_timeout > timeout): - _log.info("Overriding the command timeout: %s (before) and %s (after)", - timeout, override_timeout) - timeout = override_timeout - _sp.check_call(command, cwd=cwd, timeout=timeout, **kwargs) - - -# We can't use run_command() because robocopy (Robust Copy, rsync equivalent on Windows) -# for some reason uses nonzero return codes even on *successful* copies, so we need to check it manually. -# Also, robocopy only accepts dirs, not individual files, so we need a separate command for the -# single-file case. -def wincopy(source: str, dest: str, isdir: bool, cwd=None) -> None: - if not isdir: # Single-file copy - run_command(["cmd", "/c", "copy", - str(_pl.Path(source)), f".\\{dest}"], - cwd=cwd) - else: # Directory sync - copy_cmd = ["robocopy", source, f"./{dest}", "/mir", "/e"] - print(f"Running: cd %s && %s" % - (str(cwd or _pl.Path.cwd()), " ".join(copy_cmd))) - - # Run the command from the specified directory - result = _sp.run(copy_cmd, cwd=cwd) - - # Check for valid exit code - if result.returncode < 8: - print("ROBOCOPY completed successfully.") - else: - print( - "ROBOCOPY failure. Displaying error. See https://ss64.com/nt/robocopy-exit.html for exit code info." - ) - raise _sp.CalledProcessError(returncode=result.returncode, - cmd=copy_cmd, - output=result.stderr) diff --git a/cpp/tests/resources/scripts/build_gpt_engines.py b/cpp/tests/resources/scripts/build_gpt_engines.py deleted file mode 100755 index fa089d773dc0..000000000000 --- a/cpp/tests/resources/scripts/build_gpt_engines.py +++ /dev/null @@ -1,292 +0,0 @@ -#!/usr/bin/env python3 -# 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 argparse -import os -import platform -import shutil -import sys -from pathlib import Path -from typing import Optional - -from build_engines_utils import run_command, wincopy - -import tensorrt_llm.bindings as _tb -from tensorrt_llm.bindings.internal.testing import ModelSpec, QuantMethod - - -def convert_ckpt(model_dir: str, - output_dir: str, - *args, - world_size: int = 1, - dtype: str = 'float16'): - convert_cmd = [ - sys.executable, "examples/models/core/gpt/convert_checkpoint.py", - f"--model_dir={model_dir}", f"--output_dir={output_dir}", - f"--dtype={dtype}", f"--tp_size={world_size}" - ] + list(args) - run_command(convert_cmd) - - -def build_engine( - checkpoint_dir: str, - engine_dir: str, - *args, - max_input_len: int = 256, - max_seq_len: int = 384, -): - - build_cmd = [ - "trtllm-build", - '--log_level=error', - f'--checkpoint_dir={checkpoint_dir}', - f'--output_dir={engine_dir}', - '--max_batch_size=64', - f'--max_input_len={max_input_len}', - f'--max_seq_len={max_seq_len}', - '--max_beam_width=2', - '--kv_cache_type=continuous', - ] - legacy_args = [ - "--gpt_attention_plugin=disable", - "--context_fmha=disable", - "--remove_input_padding=disable", - ] - build_cmd = build_cmd + legacy_args + list(args) - run_command(build_cmd) - - -def build_engines(model_cache: Optional[str] = None, - world_size: int = 1, - clean: Optional[bool] = False): - # TODO add support of Pipeline parallelism to GPT - tp_size = world_size - pp_size = 1 - cp_size = 1 - - resources_dir = Path(__file__).parent.resolve().parent - models_dir = resources_dir / 'models' - model_name = 'gpt2' - - # Clone or update the model directory without lfs - hf_dir = models_dir / model_name - if hf_dir.exists(): - assert hf_dir.is_dir() - run_command(["git", "pull"], cwd=hf_dir) - else: - if platform.system() == "Windows": - url_prefix = "" - else: - url_prefix = "file://" - - model_url = url_prefix + str( - Path(model_cache) / - model_name) if model_cache else "https://huggingface.co/gpt2" - run_command([ - "git", "clone", model_url, "--single-branch", "--no-local", - model_name - ], - cwd=hf_dir.parent, - env={ - **os.environ, "GIT_LFS_SKIP_SMUDGE": "1" - }) - - assert hf_dir.is_dir() - - # Download the model file - model_file_name = "pytorch_model.bin" - if model_cache: - if platform.system() == "Windows": - wincopy(source=str( - Path(model_cache) / model_name / model_file_name), - dest=model_file_name, - isdir=False, - cwd=hf_dir) - else: - run_command([ - "rsync", "-rlptD", - str(Path(model_cache) / model_name / model_file_name), "." - ], - cwd=hf_dir) - else: - run_command(["git", "lfs", "pull", "--include", model_file_name], - cwd=hf_dir) - - safetensor_file = hf_dir / "model.safetensors" - has_safetensor = safetensor_file.exists() - if has_safetensor: - safetensor_file.rename(str(safetensor_file) + ".bak") - - assert (hf_dir / model_file_name).is_file() - - ckpt_dir = models_dir / 'c-model' / model_name - engine_dir = models_dir / 'rt_engine' / model_name - - if clean: - target_dir = Path(engine_dir) - print('clean up target folder ', target_dir) - if target_dir.is_dir(): - shutil.rmtree(target_dir, ignore_errors=True) - - tp_pp_cp_dir = f"tp{tp_size}-pp{pp_size}-cp{cp_size}-gpu" - tp_dir = f"{world_size}-gpu" - - print("\nConverting to fp16") - fp16_ckpt_dir = ckpt_dir / 'fp16' / tp_dir - convert_ckpt(str(hf_dir), - str(fp16_ckpt_dir), - world_size=tp_size, - dtype='float16') - - print("\nBuilding fp16 engines") - - input_file = 'input_tokens.npy' - # this engine can be use for in-flight batching - ifb_base_args = [ - '--gpt_attention_plugin=float16', - '--remove_input_padding=enable', - '--context_fmha=enable', - '--max_num_tokens=10000', - '--use_paged_context_fmha=enable', - ] - - paged_kv_cache_args = ['--kv_cache_type=paged'] - - no_kv_cache_args = ['--kv_cache_type=disabled'] - - def get_ifb_args(kv_cache_type): - if kv_cache_type == _tb.KVCacheType.DISABLED: - return ifb_base_args + no_kv_cache_args - elif kv_cache_type == _tb.KVCacheType.PAGED: - return ifb_base_args + paged_kv_cache_args - else: - assert False, f"Unsupported kv_cache_type: {kv_cache_type}" - - model_spec_obj = ModelSpec(input_file, _tb.DataType.HALF) - model_spec_obj.use_gpt_plugin() - model_spec_obj.set_kv_cache_type(_tb.KVCacheType.PAGED) - model_spec_obj.use_packed_input() - - model_spec_current = model_spec_obj.__copy__() - - for kv_cache_type in [_tb.KVCacheType.DISABLED, _tb.KVCacheType.PAGED]: - model_spec_current.set_kv_cache_type(kv_cache_type) - build_engine( - str(fp16_ckpt_dir), - str(engine_dir / model_spec_current.get_model_path() / - tp_pp_cp_dir), *get_ifb_args(kv_cache_type)) - - model_spec_current = model_spec_obj.__copy__() - max_draft_tokens = 5 - model_spec_current.use_draft_tokens_external_decoding() - model_spec_current.set_draft_tokens(max_draft_tokens) - - build_engine( - str(fp16_ckpt_dir), - str(engine_dir / model_spec_current.get_model_path() / tp_pp_cp_dir), - f'--max_draft_len={max_draft_tokens}', - '--speculative_decoding_mode=draft_tokens_external', - *get_ifb_args(_tb.KVCacheType.PAGED)) - - model_spec_current = model_spec_obj.__copy__() - model_spec_current.use_multiple_profiles() - - build_engine( - str(fp16_ckpt_dir), - str(engine_dir / model_spec_current.get_model_path() / tp_pp_cp_dir), - '--multiple_profiles=enable', *get_ifb_args(_tb.KVCacheType.PAGED)) - - model_spec_current = model_spec_obj.__copy__() - max_input_len = 128 - model_spec_current.set_max_input_length(max_input_len) - - build_engine(str(fp16_ckpt_dir), - str(engine_dir / model_spec_current.get_model_path() / - tp_pp_cp_dir), - *get_ifb_args(_tb.KVCacheType.PAGED), - max_input_len=max_input_len) - - # We build almost the same engine twice. But this engine has gather_context_logits - # to extract logits from python runtime and uses context FMHA for generation to match draft model executions, - # which uses context FMHA for draft tokens prediction. - # Currently the gather_context_logits is not supported with target model of speculative decoding - model_spec_current = model_spec_obj.__copy__() - model_spec_current.gather_logits() - - build_engine( - str(fp16_ckpt_dir), - str(engine_dir / model_spec_current.get_model_path() / tp_pp_cp_dir), - '--gather_context_logits', *get_ifb_args(_tb.KVCacheType.PAGED)) - - # build engine with lora enabled - model_spec_current = model_spec_obj.__copy__() - model_spec_current.use_lora_plugin() - build_engine( - str(fp16_ckpt_dir), - str(engine_dir / model_spec_current.get_model_path() / tp_pp_cp_dir), - "--lora_target_modules=attn_qkv", '--lora_plugin=float16', - *get_ifb_args(_tb.KVCacheType.PAGED)) - - if model_cache: - llm_datasets_root = Path(model_cache) / "datasets" - calib_dataset = llm_datasets_root / "cimec/lambada/" - else: - calib_dataset = "lambada" - print("\nConverting to fp16 SQ") - fp16_sq_ckpt_dir = ckpt_dir / 'fp16-sq' / tp_dir - convert_ckpt(str(hf_dir), - str(fp16_sq_ckpt_dir), - "--smoothquant=0.5", - f"--calib_dataset={calib_dataset}", - world_size=tp_size, - dtype='float16') - - print("\nBuilding fp16 SQ engines") - model_spec_current = ModelSpec(input_file, _tb.DataType.HALF) - model_spec_current.use_gpt_plugin() - model_spec_current.use_packed_input() - model_spec_current.set_quant_method(QuantMethod.SMOOTH_QUANT) - - for kv_cache_type in [_tb.KVCacheType.DISABLED, _tb.KVCacheType.PAGED]: - model_spec_current.set_kv_cache_type(kv_cache_type) - build_engine( - str(fp16_sq_ckpt_dir), - str(engine_dir / model_spec_current.get_model_path() / - tp_pp_cp_dir), *get_ifb_args(kv_cache_type)) - - if has_safetensor: - Path(str(safetensor_file) + ".bak").rename(safetensor_file) - - print("Done.") - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument("--model_cache", - type=str, - help="Directory where models are stored") - - parser.add_argument('--world_size', - type=int, - default=1, - help='World size, only support tensor parallelism now') - - parser.add_argument('--clean', - action='store_true', - default=False, - help='Clean target folders before building engines') - - build_engines(**vars(parser.parse_args())) diff --git a/cpp/tests/resources/scripts/build_gptj_engines.py b/cpp/tests/resources/scripts/build_gptj_engines.py deleted file mode 100755 index bfab97e0ec11..000000000000 --- a/cpp/tests/resources/scripts/build_gptj_engines.py +++ /dev/null @@ -1,189 +0,0 @@ -#!/usr/bin/env python3 -# 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 argparse as _arg -import os as _os -import pathlib as _pl -import platform as _pf -import sys as _sys -import typing as _tp - -from build_engines_utils import run_command, wincopy - -import tensorrt_llm.bindings as _tb -from tensorrt_llm.bindings.internal.testing import ModelSpec - - -def get_ckpt_without_quatization(model_dir, output_dir): - build_args = [ - _sys.executable, "examples/models/contrib/gpt/convert_checkpoint.py" - ] + [ - '--model_dir={}'.format(model_dir), - '--output_dir={}'.format(output_dir), - ] - run_command(build_args) - - -def get_ckpt_with_modelopt_quant(model_dir, output_dir, model_cache): - build_args = [_sys.executable, "examples/quantization/quantize.py"] + [ - '--model_dir={}'.format(model_dir), - '--output_dir={}'.format(output_dir), '--qformat=fp8', - '--kv_cache_dtype=fp8', - f'--calib_dataset={model_cache}/datasets/cnn_dailymail' - ] - run_command(build_args) - - -def build_engine(checkpoint_dir: _pl.Path, engine_dir: _pl.Path, *args): - build_args = ["trtllm-build"] + ( - ['--checkpoint_dir', str(checkpoint_dir)] if checkpoint_dir else []) + [ - '--output_dir', - str(engine_dir), - '--logits_dtype=float16', - '--gemm_plugin=float16', - '--max_batch_size=32', - '--max_input_len=40', - '--max_seq_len=60', - '--max_beam_width=2', - '--log_level=error', - ] + list(args) - run_command(build_args) - - -def build_engines(model_cache: _tp.Optional[str] = None, only_fp8=False): - resources_dir = _pl.Path(__file__).parent.resolve().parent - models_dir = resources_dir / 'models' - model_name = 'gpt-j-6b' - - # Clone or update the model directory without lfs - hf_dir = models_dir / model_name - if hf_dir.exists(): - assert hf_dir.is_dir() - run_command(["git", "pull"], cwd=hf_dir) - else: - if _pf.system() == "Windows": - url_prefix = "" - else: - url_prefix = "file://" - model_url = url_prefix + str( - _pl.Path(model_cache) / model_name - ) if model_cache else "https://huggingface.co/EleutherAI/gpt-j-6b" - run_command([ - "git", "clone", model_url, "--single-branch", "--no-local", - model_name - ], - cwd=hf_dir.parent, - env={ - **_os.environ, "GIT_LFS_SKIP_SMUDGE": "1" - }) - - assert (hf_dir.is_dir()) - - # Download the model file - model_file_name = "pytorch_model.bin" - if model_cache: - if _pf.system() == "Windows": - wincopy(source=str( - _pl.Path(model_cache) / model_name / model_file_name), - dest=model_file_name, - isdir=False, - cwd=hf_dir) - else: - run_command([ - "rsync", "-rlptD", - str(_pl.Path(model_cache) / model_name / model_file_name), "." - ], - cwd=hf_dir) - else: - run_command(["git", "lfs", "pull", "--include", model_file_name], - cwd=hf_dir) - - assert ((hf_dir / model_file_name).is_file()) - - engine_dir = models_dir / 'rt_engine' / model_name - - # TODO add Tensor and Pipeline parallelism to GPT-J - tp_size = 1 - pp_size = 1 - cp_size = 1 - tp_pp_cp_dir = f"tp{tp_size}-pp{pp_size}-cp{cp_size}-gpu" - input_file = 'input_tokens.npy' - - if only_fp8: - # with ifb, new plugin - print( - "\nBuilding fp8-plugin engine using gpt_attention_plugin with inflight-batching, packed" - ) - # TODO: use dummy scales atm; to re-enable when data is uploaded to the model cache - # quantized_fp8_model_arg = '--quantized_fp8_model_path=' + \ - # str(_pl.Path(model_cache) / 'fp8-quantized-modelopt' / 'gptj_tp1_rank0.npz') - fp8_ckpt_path = engine_dir / 'fp8' / tp_pp_cp_dir - get_ckpt_with_modelopt_quant(hf_dir, fp8_ckpt_path, model_cache) - model_spec_obj = ModelSpec(input_file, _tb.DataType.FP8) - model_spec_obj.use_gpt_plugin() - model_spec_obj.set_kv_cache_type(_tb.KVCacheType.PAGED) - model_spec_obj.use_packed_input() - build_engine( - fp8_ckpt_path, - engine_dir / model_spec_obj.get_model_path() / tp_pp_cp_dir, - '--gpt_attention_plugin=float16', - '--paged_kv_cache=enable', - '--remove_input_padding=enable', - '--use_paged_context_fmha=enable', - ) - else: - fp16_ckpt_path = engine_dir / 'fp16' / tp_pp_cp_dir - get_ckpt_without_quatization(hf_dir, fp16_ckpt_path) - print("\nBuilding fp16-plugin engine") - model_spec_obj = ModelSpec(input_file, _tb.DataType.HALF) - model_spec_obj.use_gpt_plugin() - model_spec_obj.set_kv_cache_type(_tb.KVCacheType.CONTINUOUS) - - build_engine( - fp16_ckpt_path, - engine_dir / model_spec_obj.get_model_path() / tp_pp_cp_dir, - '--gpt_attention_plugin=float16', '--paged_kv_cache=disable', - '--remove_input_padding=disable', "--context_fmha=disable") - - print("\nBuilding fp16-plugin-packed engine") - model_spec_obj.use_packed_input() - build_engine( - fp16_ckpt_path, - engine_dir / model_spec_obj.get_model_path() / tp_pp_cp_dir, - '--gpt_attention_plugin=float16', '--paged_kv_cache=disable', - '--remove_input_padding=enable', "--context_fmha=disable") - - print("\nBuilding fp16-plugin-packed-paged engine") - model_spec_obj.set_kv_cache_type(_tb.KVCacheType.PAGED) - build_engine( - fp16_ckpt_path, - engine_dir / model_spec_obj.get_model_path() / tp_pp_cp_dir, - '--gpt_attention_plugin=float16', '--paged_kv_cache=enable', - '--remove_input_padding=enable', "--context_fmha=disable") - print("Done.") - - -if __name__ == "__main__": - parser = _arg.ArgumentParser() - parser.add_argument("--model_cache", - type=str, - help="Directory where models are stored") - parser.add_argument( - "--only_fp8", - action="store_true", - help="Build engines for only FP8 tests. Implemented for H100 runners.") - - build_engines(**vars(parser.parse_args())) diff --git a/cpp/tests/resources/scripts/build_llama_engines.py b/cpp/tests/resources/scripts/build_llama_engines.py deleted file mode 100644 index dbac12621c73..000000000000 --- a/cpp/tests/resources/scripts/build_llama_engines.py +++ /dev/null @@ -1,149 +0,0 @@ -#!/usr/bin/env python3 -# 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 argparse as _arg -import pathlib as _pl -import platform as _pf -import sys as _sys -import time - -from build_engines_utils import run_command, wincopy - -import tensorrt_llm.bindings as _tb -from tensorrt_llm.bindings.internal.testing import ModelSpec - - -def build_engine(weight_dir: _pl.Path, engine_dir: _pl.Path, convert_extra_args, - build_extra_args): - - ckpt_dir = engine_dir / 'ckpt' - - convert_cmd = [ - _sys.executable, "examples/models/core/llama/convert_checkpoint.py" - ] + ([f'--model_dir={weight_dir}'] if weight_dir else []) + [ - f'--output_dir={ckpt_dir}', - '--dtype=float16', - ] + convert_extra_args - - run_command(convert_cmd) - - build_args = [ - 'trtllm-build', - f'--checkpoint_dir={ckpt_dir}', - f'--output_dir={engine_dir}', - '--gpt_attention_plugin=float16', - '--gemm_plugin=float16', - '--max_batch_size=32', - '--max_input_len=40', - '--max_seq_len=60', - '--max_beam_width=2', - '--log_level=error', - '--paged_kv_cache=enable', - '--remove_input_padding=enable', - ] + build_extra_args - - run_command(build_args) - - -def build_engines(model_cache: str, only_multi_gpu: bool): - resources_dir = _pl.Path(__file__).parent.resolve().parent - models_dir = resources_dir / 'models' - model_name = 'Llama-3.2-1B' - - if model_cache: - print("Copy model from model_cache") - model_cache_dir = _pl.Path( - model_cache) / 'llama-3.2-models' / model_name - assert (model_cache_dir.is_dir()), model_cache_dir - - if _pf.system() == "Windows": - wincopy(source=str(model_cache_dir), - dest=model_name, - isdir=True, - cwd=models_dir) - else: - run_command(["rsync", "-rlptD", - str(model_cache_dir), "."], - cwd=models_dir) - - hf_dir = models_dir / model_name - assert hf_dir.is_dir(), f"testing {hf_dir}" - - engine_dir = models_dir / 'rt_engine' / model_name - - model_spec_obj = ModelSpec('input_tokens_llama.npy', _tb.DataType.HALF) - model_spec_obj.use_gpt_plugin() - model_spec_obj.set_kv_cache_type(_tb.KVCacheType.PAGED) - model_spec_obj.use_packed_input() - - tp_pp_cp_sizes = [(1, 1, 1)] - if only_multi_gpu: - tp_pp_cp_sizes = [(1, 4, 1), (4, 1, 1), (1, 2, 1), (2, 2, 1), (2, 1, 1), - (1, 1, 2), (2, 1, 2)] - for tp_size, pp_size, cp_size in tp_pp_cp_sizes: - print(f"\nBuilding fp16 tp{tp_size} pp{pp_size} cp{cp_size} engine") - start_time = time.time() - - tp_pp_cp_dir = f"tp{tp_size}-pp{pp_size}-cp{cp_size}-gpu" - model_spec_obj.use_tensor_parallelism(tp_size) - model_spec_obj.use_pipeline_parallelism(pp_size) - model_spec_obj.use_context_parallelism(cp_size) - - build_engine( - hf_dir, engine_dir / model_spec_obj.get_model_path() / tp_pp_cp_dir, - [ - f'--tp_size={tp_size}', f'--pp_size={pp_size}', - f'--cp_size={cp_size}' - ], ['--use_paged_context_fmha=disable']) - - duration = time.time() - start_time - print( - f"Building fp16 tp{tp_size} pp{pp_size} cp{cp_size} engine took {duration} seconds" - ) - - if not only_multi_gpu: - print(f"\nBuilding lookahead engine") - start_time = time.time() - - model_spec_obj.use_tensor_parallelism(1) - model_spec_obj.use_pipeline_parallelism(1) - model_spec_obj.use_context_parallelism(1) - model_spec_obj.use_lookahead_decoding() - build_engine( - hf_dir, - engine_dir / model_spec_obj.get_model_path() / 'tp1-pp1-cp1-gpu', - [], [ - '--max_draft_len=39', - '--speculative_decoding_mode=lookahead_decoding' - ]) - - duration = time.time() - start_time - print(f"Building lookahead engine took {duration} seconds") - - print("Done.") - - -if __name__ == "__main__": - parser = _arg.ArgumentParser() - parser.add_argument("--model_cache", - type=str, - help="Directory where models are stored") - parser.add_argument( - "--only_multi_gpu", - action="store_true", - help="Flag to build only for Tensor and Pipeline parallelism") - - build_engines(**vars(parser.parse_args())) diff --git a/cpp/tests/resources/scripts/build_mamba_engines.py b/cpp/tests/resources/scripts/build_mamba_engines.py deleted file mode 100644 index 6b10a5b03531..000000000000 --- a/cpp/tests/resources/scripts/build_mamba_engines.py +++ /dev/null @@ -1,154 +0,0 @@ -#!/usr/bin/env python3 -# 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 argparse as _arg -import os as _os -import pathlib as _pl -import platform as _pf -import sys as _sys -import typing as _tp - -from build_engines_utils import run_command, wincopy - -import tensorrt_llm.bindings as _tb -from tensorrt_llm.bindings.internal.testing import ModelSpec - - -def build_engine(weight_dir: _pl.Path, ckpt_dir: _pl.Path, engine_dir: _pl.Path, - *args): - convert_args = [ - _sys.executable, "examples/models/core/mamba/convert_checkpoint.py" - ] + (['--model_dir', str(weight_dir)] if weight_dir else []) + [ - '--output_dir', - str(ckpt_dir), - '--dtype=float16', - ] - run_command(convert_args) - build_args = ["trtllm-build"] + ['--checkpoint_dir', - str(ckpt_dir)] + [ - '--output_dir', - str(engine_dir), - '--gpt_attention_plugin=disable', - '--paged_kv_cache=disable', - '--gemm_plugin=disable', - '--max_batch_size=8', - '--max_input_len=924', - '--max_seq_len=1024', - '--max_beam_width=1', - ] + list(args) - run_command(build_args) - - -def build_engines(model_cache: _tp.Optional[str] = None): - resources_dir = _pl.Path(__file__).parent.resolve().parent - models_dir = resources_dir / 'models' - model_name = 'mamba-2.8b-hf' - - if model_cache: - print("Copy model from model_cache") - model_cache_dir = _pl.Path(model_cache) / 'mamba' / model_name - if _pf.system() == "Windows": - wincopy(source=str(model_cache_dir), - dest=model_name, - isdir=True, - cwd=models_dir) - else: - run_command(["rsync", "-rlptD", - str(model_cache_dir), "."], - cwd=models_dir) - else: - print("Clone model from HF") - hf_dir = _pl.Path(models_dir) / model_name - run_command( - [ - "git", "clone", - "https://huggingface.co/state-spaces/mamba-2.8b-hf", model_name - ], - cwd=models_dir, - ) - hf_dir = models_dir / model_name - assert (hf_dir.is_dir()) - - # Clone or update the tokenizer directory without lfs - tokenizer_name = 'gpt-neox-20b' - tokenizer_hf_dir = models_dir / tokenizer_name - if tokenizer_hf_dir.exists(): - assert tokenizer_hf_dir.is_dir() - run_command(["git", "pull"], cwd=tokenizer_hf_dir) - else: - if _pf.system() == "Windows": - url_prefix = "" - else: - url_prefix = "file://" - tokenizer_url = url_prefix + str( - _pl.Path(model_cache) / tokenizer_name - ) if model_cache else "https://huggingface.co/EleutherAI/gpt-neox-20b" - run_command([ - "git", "clone", tokenizer_url, "--single-branch", "--no-local", - tokenizer_name - ], - cwd=tokenizer_hf_dir.parent, - env={ - **_os.environ, "GIT_LFS_SKIP_SMUDGE": "1" - }) - - tp_size = 1 - pp_size = 1 - cp_size = 1 - tp_pp_cp_dir = f"tp{tp_size}-pp{pp_size}-cp{cp_size}-gpu" - - ckpt_dir = models_dir / 'rt_ckpt' / model_name - engine_dir = models_dir / 'rt_engine' / model_name - model_spec_obj = ModelSpec('input_tokens.npy', _tb.DataType.HALF) - model_spec_obj.set_kv_cache_type(_tb.KVCacheType.CONTINUOUS) - model_spec_obj.use_tensor_parallelism(tp_size) - model_spec_obj.use_pipeline_parallelism(pp_size) - model_spec_obj.use_context_parallelism(cp_size) - - print("\nBuilding fp16 engine") - build_engine(hf_dir, - ckpt_dir / model_spec_obj.get_model_path() / tp_pp_cp_dir, - engine_dir / model_spec_obj.get_model_path() / tp_pp_cp_dir, - '--remove_input_padding=disable', '--paged_state=disable', - '--mamba_conv1d_plugin=disable') - print("\nBuilding fp16-plugin engine") - model_spec_obj.use_mamba_plugin() - build_engine(hf_dir, - ckpt_dir / model_spec_obj.get_model_path() / tp_pp_cp_dir, - engine_dir / model_spec_obj.get_model_path() / tp_pp_cp_dir, - '--remove_input_padding=disable', '--paged_state=disable') - print("\nBuilding fp16-plugin-packed engine") - model_spec_obj.use_packed_input() - build_engine(hf_dir, - ckpt_dir / model_spec_obj.get_model_path() / tp_pp_cp_dir, - engine_dir / model_spec_obj.get_model_path() / tp_pp_cp_dir, - '--remove_input_padding=enable', '--paged_state=disable') - print("\nBuilding fp16-plugin-packed-paged engine") - model_spec_obj.set_kv_cache_type(_tb.KVCacheType.PAGED) - build_engine(hf_dir, - ckpt_dir / model_spec_obj.get_model_path() / tp_pp_cp_dir, - engine_dir / model_spec_obj.get_model_path() / tp_pp_cp_dir, - '--remove_input_padding=enable', '--paged_state=enable') - print("Done.") - - -if __name__ == "__main__": - parser = _arg.ArgumentParser() - parser.add_argument("--model_cache", - type=str, - help="Directory where models are stored") - - build_engines(**vars(parser.parse_args())) diff --git a/cpp/tests/resources/scripts/build_medusa_engines.py b/cpp/tests/resources/scripts/build_medusa_engines.py deleted file mode 100755 index cf9c74f8779f..000000000000 --- a/cpp/tests/resources/scripts/build_medusa_engines.py +++ /dev/null @@ -1,115 +0,0 @@ -#!/usr/bin/env python3 -# 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 argparse as _arg -import pathlib as _pl -import platform as _pf -import sys as _sys - -from build_engines_utils import run_command, wincopy - -import tensorrt_llm.bindings as _tb -from tensorrt_llm.bindings.internal.testing import ModelSpec - - -def build_engine(base_model_dir: _pl.Path, medusa_model_dir: _pl.Path, - engine_dir: _pl.Path, *args): - - covert_cmd = [_sys.executable, "examples/medusa/convert_checkpoint.py"] + ( - ['--model_dir', str(base_model_dir)] if base_model_dir else []) + [ - '--medusa_model_dir', str(medusa_model_dir), \ - '--output_dir', str(engine_dir), '--dtype=float16', '--num_medusa_heads=4' - ] + list(args) - - run_command(covert_cmd) - - build_args = ["trtllm-build"] + ( - ['--checkpoint_dir', str(engine_dir)] if engine_dir else []) + [ - '--output_dir', - str(engine_dir), - '--gemm_plugin=float16', - '--max_batch_size=8', - '--max_input_len=12', - '--max_seq_len=140', - '--log_level=error', - '--paged_kv_cache=enable', - '--use_paged_context_fmha=enable', - '--remove_input_padding=enable', - '--speculative_decoding_mode=medusa', - ] - - run_command(build_args) - - -def build_engines(model_cache: str): - resources_dir = _pl.Path(__file__).parent.resolve().parent - models_dir = resources_dir / 'models' - model_name = 'vicuna-7b-medusa' - base_model_name = 'vicuna-7b-v1.3' - medusa_model_name = 'medusa-vicuna-7b-v1.3' - - if model_cache: - print(f"Copy model from {model_cache}") - base_model_cache_dir = _pl.Path(model_cache) / base_model_name - medusa_head_cache_dir = _pl.Path(model_cache) / medusa_model_name - assert base_model_cache_dir.is_dir(), base_model_cache_dir - assert medusa_head_cache_dir.is_dir(), medusa_head_cache_dir - - if _pf.system() == "Windows": - wincopy(source=str(base_model_cache_dir), - dest=base_model_name, - isdir=True, - cwd=models_dir) - wincopy(source=str(medusa_head_cache_dir), - dest=medusa_model_name, - isdir=True, - cwd=models_dir) - else: - run_command(["rsync", "-rlptD", - str(base_model_cache_dir), "."], - cwd=models_dir) - run_command(["rsync", "-rlptD", - str(medusa_head_cache_dir), "."], - cwd=models_dir) - - base_model_dir = models_dir / base_model_name - medusa_model_dir = models_dir / medusa_model_name - assert base_model_dir.is_dir() - assert medusa_model_dir.is_dir() - - engine_dir = models_dir / 'rt_engine' / model_name - - model_spec_obj = ModelSpec('input_tokens.npy', _tb.DataType.HALF) - model_spec_obj.use_gpt_plugin() - model_spec_obj.set_kv_cache_type(_tb.KVCacheType.PAGED) - model_spec_obj.use_packed_input() - model_spec_obj.use_medusa() - - full_engine_path = engine_dir / model_spec_obj.get_model_path( - ) / 'tp1-pp1-cp1-gpu' - print(f"\nBuilding fp16 engine at {str(full_engine_path)}") - build_engine(base_model_dir, medusa_model_dir, full_engine_path) - - print("Done.") - - -if __name__ == "__main__": - parser = _arg.ArgumentParser() - parser.add_argument("--model_cache", - type=str, - help="Directory where models are stored") - - build_engines(**vars(parser.parse_args())) diff --git a/cpp/tests/resources/scripts/build_recurrentgemma_engines.py b/cpp/tests/resources/scripts/build_recurrentgemma_engines.py deleted file mode 100644 index 293aab101d38..000000000000 --- a/cpp/tests/resources/scripts/build_recurrentgemma_engines.py +++ /dev/null @@ -1,135 +0,0 @@ -#!/usr/bin/env python3 -# 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 argparse as _arg -import os as _os -import pathlib as _pl -import platform as _pf -import sys as _sys -import typing as _tp - -from build_engines_utils import run_command, wincopy - -import tensorrt_llm.bindings as _tb -from tensorrt_llm.bindings.internal.testing import ModelSpec - - -def build_engine(weight_dir: _pl.Path, ckpt_dir: _pl.Path, engine_dir: _pl.Path, - *args): - convert_args = [ - _sys.executable, - "examples/models/core/recurrentgemma/convert_checkpoint.py" - ] + (['--model_dir', str(weight_dir)] if weight_dir else []) + [ - '--output_dir', - str(ckpt_dir), - '--ckpt_type=hf', - '--dtype=float16', - ] - run_command(convert_args) - build_args = ["trtllm-build"] + ['--checkpoint_dir', - str(ckpt_dir)] + [ - '--output_dir', - str(engine_dir), - '--gpt_attention_plugin=float16', - '--paged_kv_cache=enable', - '--gemm_plugin=float16', - '--max_batch_size=8', - '--max_input_len=924', - '--max_seq_len=1024', - '--max_beam_width=1', - ] + list(args) - run_command(build_args) - - -def build_engines(model_cache: _tp.Optional[str] = None): - resources_dir = _pl.Path(__file__).parent.resolve().parent - models_dir = resources_dir / 'models' - model_name = 'recurrentgemma-2b' - hf_dir = models_dir / model_name - - # Clone or update the model directory without lfs - if model_cache: - print("Copy model from model_cache") - model_cache_dir = _pl.Path(model_cache) / 'recurrentgemma' / model_name - print(model_cache_dir) - assert (model_cache_dir.is_dir()) - if _pf.system() == "Windows": - wincopy(source=str(model_cache_dir), - dest=model_name, - isdir=True, - cwd=models_dir) - else: - run_command(["rsync", "-rlptD", - str(model_cache_dir), "."], - cwd=models_dir) - else: - if not hf_dir.is_dir(): - if _pf.system() == "Windows": - url_prefix = "" - else: - url_prefix = "file://" - model_url = "https://huggingface.co/google/recurrentgemma-2b" - run_command([ - "git", "clone", model_url, "--single-branch", "--no-local", - model_name - ], - cwd=models_dir, - env={ - **_os.environ, "GIT_LFS_SKIP_SMUDGE": "1" - }) - - assert (hf_dir.is_dir()) - - # Download the model file - model_file_name = "*" - if not model_cache: - run_command(["git", "lfs", "pull", "--include", model_file_name], - cwd=hf_dir) - - tp_size = 1 - pp_size = 1 - cp_size = 1 - tp_pp_cp_dir = f"tp{tp_size}-pp{pp_size}-cp{cp_size}-gpu" - - ckpt_dir = models_dir / 'rt_ckpt' / model_name - engine_dir = models_dir / 'rt_engine' / model_name - - python_exe = _sys.executable - run_command([python_exe, "-m", "pip", "install", "transformers>=4.40.0"], - env=_os.environ, - timeout=300) - input_file = 'input_tokens.npy' - model_spec_obj = ModelSpec(input_file, _tb.DataType.HALF) - model_spec_obj.use_gpt_plugin() - model_spec_obj.use_packed_input() - model_spec_obj.set_kv_cache_type(_tb.KVCacheType.PAGED) - - print("\nBuilding fp16-plugin-packed-paged engine") - build_engine(hf_dir, - ckpt_dir / model_spec_obj.get_model_path() / tp_pp_cp_dir, - engine_dir / model_spec_obj.get_model_path() / tp_pp_cp_dir, - '--remove_input_padding=enable', '--paged_state=enable') - - print("Done.") - - -if __name__ == "__main__": - parser = _arg.ArgumentParser() - parser.add_argument("--model_cache", - type=str, - help="Directory where models are stored") - - build_engines(**vars(parser.parse_args())) diff --git a/cpp/tests/resources/scripts/build_redrafter_engines.py b/cpp/tests/resources/scripts/build_redrafter_engines.py deleted file mode 100755 index cdf3e889ac35..000000000000 --- a/cpp/tests/resources/scripts/build_redrafter_engines.py +++ /dev/null @@ -1,126 +0,0 @@ -#!/usr/bin/env python3 -# 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 argparse as _arg -import pathlib as _pl -import platform as _pf -import sys as _sys - -from build_engines_utils import run_command, wincopy - -import tensorrt_llm.bindings as _tb -from tensorrt_llm.bindings.internal.testing import ModelSpec - - -def build_engine(base_model_dir: _pl.Path, drafter_model_dir: _pl.Path, - engine_dir: _pl.Path, *args): - - base_ckpt_dir = f'{base_model_dir}-ckpt' - covert_cmd_base = [ - _sys.executable, "examples/models/core/llama/convert_checkpoint.py" - ] + (['--model_dir', str(base_model_dir)] if base_model_dir else []) + [ - '--output_dir', str(base_ckpt_dir), '--dtype=float16' - ] + list(args) - - run_command(covert_cmd_base) - - covert_cmd = [ - _sys.executable, "examples/redrafter/convert_checkpoint.py"] + ( - ['--base_model_checkpoint_dir', str(base_ckpt_dir)] if base_model_dir else []) + [ - '--drafter_model_dir', str(drafter_model_dir), \ - '--output_dir', str(engine_dir), '--dtype=float16', - '--redrafter_num_beams=5', '--redrafter_draft_len_per_beam=5' - ] + list(args) - - run_command(covert_cmd) - - build_args = ["trtllm-build"] + ( - ['--checkpoint_dir', str(engine_dir)] if engine_dir else []) + [ - '--output_dir', - str(engine_dir), - '--gemm_plugin=float16', - '--max_batch_size=8', - '--max_input_len=64', - '--max_seq_len=1024', - '--log_level=error', - '--paged_kv_cache=enable', - '--use_paged_context_fmha=enable', - '--remove_input_padding=enable', - '--speculative_decoding_mode=explicit_draft_tokens', - ] - - run_command(build_args) - - -def build_engines(model_cache: str): - resources_dir = _pl.Path(__file__).parent.resolve().parent - models_dir = resources_dir / 'models' - model_name = 'vicuna-7b-redrafter' - base_model_name = 'vicuna-7b-v1.3' - drafter_model_name = 'redrafter-vicuna-7b-v1.3' - - if model_cache: - print(f"Copy model from {model_cache}") - base_model_cache_dir = _pl.Path(model_cache) / base_model_name - drafter_cache_dir = _pl.Path(model_cache) / drafter_model_name - assert base_model_cache_dir.is_dir(), base_model_cache_dir - assert drafter_cache_dir.is_dir(), drafter_cache_dir - - if _pf.system() == "Windows": - wincopy(source=str(base_model_cache_dir), - dest=base_model_name, - isdir=True, - cwd=models_dir) - wincopy(source=str(drafter_cache_dir), - dest=drafter_model_name, - isdir=True, - cwd=models_dir) - else: - run_command(["rsync", "-rlptD", - str(base_model_cache_dir), "."], - cwd=models_dir) - run_command(["rsync", "-rlptD", - str(drafter_cache_dir), "."], - cwd=models_dir) - - base_model_dir = models_dir / base_model_name - drafter_model_dir = models_dir / drafter_model_name - assert base_model_dir.is_dir() - assert drafter_model_dir.is_dir() - - engine_dir = models_dir / 'rt_engine' / model_name - - model_spec_obj = ModelSpec('input_tokens.npy', _tb.DataType.HALF) - model_spec_obj.use_gpt_plugin() - model_spec_obj.set_kv_cache_type(_tb.KVCacheType.PAGED) - model_spec_obj.use_packed_input() - model_spec_obj.use_explicit_draft_tokens_decoding() - - full_engine_path = engine_dir / model_spec_obj.get_model_path( - ) / 'tp1-pp1-cp1-gpu' - print(f"\nBuilding fp16 engine at {str(full_engine_path)}") - build_engine(base_model_dir, drafter_model_dir, full_engine_path) - - print("Done.") - - -if __name__ == "__main__": - parser = _arg.ArgumentParser() - parser.add_argument("--model_cache", - type=str, - help="Directory where models are stored") - - build_engines(**vars(parser.parse_args())) diff --git a/cpp/tests/resources/scripts/generate_expected_chatglm_output.py b/cpp/tests/resources/scripts/generate_expected_chatglm_output.py deleted file mode 100755 index 416f76938700..000000000000 --- a/cpp/tests/resources/scripts/generate_expected_chatglm_output.py +++ /dev/null @@ -1,126 +0,0 @@ -#!/usr/bin/env python3 -# 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 os -from pathlib import Path - -import numpy as np - -# isort: off -import run -# isort: on - -import tensorrt_llm.bindings as _tb -from tensorrt_llm.bindings.internal.testing import ModelSpec - -resources_dir = Path(__file__).parent.resolve().parent -model_path = resources_dir / "models" - - -def generate_output( - model_name: str = "", - num_beams: int = 1, - max_output_len: int = 8, - output_logits: bool = False, - output_cum_log_probs: bool = False, - output_log_probs: bool = False, -): - hf_path = model_path / model_name - tp_size = 1 - pp_size = 1 - cp_size = 1 - tp_pp_cp_dir = f"tp{tp_size}-pp{pp_size}-cp{cp_size}-gpu/" - input_file = f"input_tokens_{model_name}.npy" - - data_input_file_name = resources_dir / "data" / input_file - if num_beams == 1: - output_dir = resources_dir / "data" / model_name / "sampling" - else: - output_dir = resources_dir / "data" / model_name / f"beam_search_{num_beams}" - output_dir.mkdir(exist_ok=True, parents=True) - - model_spec_obj_list = [ - ModelSpec(input_file, - _tb.DataType.HALF).use_gpt_plugin().set_kv_cache_type( - _tb.KVCacheType.CONTINUOUS), - ModelSpec(input_file, _tb.DataType.HALF).use_gpt_plugin(). - use_packed_input().set_kv_cache_type(_tb.KVCacheType.PAGED), - ] - - for model_spec_obj in model_spec_obj_list: - engine_dir = model_path / 'rt_engine' / model_name / model_spec_obj.get_model_path( - ) / tp_pp_cp_dir - base_output_name = os.path.splitext( - model_spec_obj.get_results_file())[0] - output_npy_file_name = output_dir / f'{base_output_name}.npy' - output_csv_file_name = output_dir / f'{base_output_name}.csv' - - args_list = [ - '--engine_dir', - str(engine_dir), - '--tokenizer_dir', - str(hf_path), - '--input_file', - str(data_input_file_name), - '--output_npy', - str(output_npy_file_name), - '--output_csv', - str(output_csv_file_name), - '--max_output_len', - str(max_output_len), - '--num_beams', - str(num_beams), - '--use_py_session', - ] - - if output_logits: - file_name = str(output_npy_file_name)[:-4] + "_logits.npy" - args_list.extend(['--output_logits_npy', file_name]) - - if output_cum_log_probs: - file_name = str(output_npy_file_name)[:-4] + "_cum_log_probs.npy" - args_list.extend(['--output_cum_log_probs_npy', file_name]) - - if output_log_probs: - file_name = str(output_npy_file_name)[:-4] + "_log_probs.npy" - args_list.extend(['--output_log_probs_npy', file_name]) - - args = run.parse_arguments(args_list) - run.main(args) - - # Convert pad_id to end_id in .npy out put file - data = np.load(str(output_npy_file_name)) - if model_name == 'chatglm-6b': - data[data == 3] = 130005 - elif model_name == 'chatglm2-6b' or model_name == 'chatglm3-6b': - data[data == 0] = 2 - elif model_name == 'glm-10b': - data[data == 50256] = 50258 - else: - raise NameError('bad model name') - - np.save(str(output_npy_file_name), data) - - -if __name__ == '__main__': - generate_output(model_name='chatglm-6b', num_beams=1) - generate_output(model_name='chatglm-6b', num_beams=2) - generate_output(model_name='chatglm2-6b', num_beams=1) - generate_output(model_name='chatglm2-6b', num_beams=2) - generate_output(model_name='chatglm3-6b', num_beams=1) - generate_output(model_name='chatglm3-6b', num_beams=2) - generate_output(model_name='glm-10b', num_beams=1) - print("Done") diff --git a/cpp/tests/resources/scripts/generate_expected_eagle_output.py b/cpp/tests/resources/scripts/generate_expected_eagle_output.py deleted file mode 100755 index 253a98beaf4e..000000000000 --- a/cpp/tests/resources/scripts/generate_expected_eagle_output.py +++ /dev/null @@ -1,86 +0,0 @@ -#!/usr/bin/env python3 -# 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 argparse as _arg -import os -from pathlib import Path - -# isort: off -import run -# isort: on - -import tensorrt_llm.bindings as _tb -from tensorrt_llm.bindings.internal.testing import ModelSpec - - -def generate_output(engine: str, - model_spec_obj: ModelSpec, - max_output_len: int = 8): - - model = 'vicuna-7b-v1.3' - model_eagle = 'vicuna-7b-eagle' - hf_model = 'vicuna-7b-v1.3' - resources_dir = Path(__file__).parent.resolve().parent - models_dir = resources_dir / 'models' - hf_dir = models_dir / hf_model - tp_pp_cp_dir = 'tp1-pp1-cp1-gpu/' - engine_dir = models_dir / 'rt_engine' / model / engine / tp_pp_cp_dir - - data_dir = resources_dir / 'data' - input_file = data_dir / 'input_vicuna.npy' - model_data_dir = data_dir / model_eagle - output_dir = model_data_dir / 'sampling' - - base_output_name = os.path.splitext(model_spec_obj.get_results_file())[0] - - args = run.parse_arguments([ - '--engine_dir', - str(engine_dir), '--input_file', - str(input_file), '--tokenizer_dir', - str(hf_dir), '--output_npy', - str(output_dir / (base_output_name + '.npy')), '--output_csv', - str(output_dir / (base_output_name + '.csv')), '--max_output_len', - str(max_output_len), '--use_py_session', '--temperature', '1.0' - ]) - run.main(args) - print(f"Output saved at {str(output_dir / base_output_name)}.[npy|csv]") - - -def generate_outputs(): - print(f'Generating outputs for Vicuna 7B v1.3 FP16') - max_output_len = 128 - model_spec_obj = ModelSpec('input_tokens_long.npy', _tb.DataType.HALF) - model_spec_obj.use_gpt_plugin() - model_spec_obj.set_max_output_length(max_output_len) - model_spec_obj.use_packed_input() - model_spec_obj.set_kv_cache_type(_tb.KVCacheType.PAGED) - - generate_output(engine=model_spec_obj.get_model_path(), - model_spec_obj=model_spec_obj, - max_output_len=max_output_len) - - -if __name__ == '__main__': - parser = _arg.ArgumentParser() - parser.add_argument( - "--only_multi_gpu", - action="store_true", - help="Generate data with Pipeline and Tensor Parallelism") - - args = parser.parse_args() - - generate_outputs() - print("Done") diff --git a/cpp/tests/resources/scripts/generate_expected_enc_dec_output.py b/cpp/tests/resources/scripts/generate_expected_enc_dec_output.py deleted file mode 100644 index fc3dc615c918..000000000000 --- a/cpp/tests/resources/scripts/generate_expected_enc_dec_output.py +++ /dev/null @@ -1,30 +0,0 @@ -from build_enc_dec_engines import Arguments, RunCMDMixin - - -class Run(RunCMDMixin): - - def command(self): - args = self.args - world_size = args.tp * args.pp - mpi_run = f'mpirun --allow-run-as-root -np {world_size}' if world_size > 1 else '' - ret = [] - for beam in args.beams_tuple: - ret.append(( - mpi_run, - f'python3 examples/models/core/enc_dec/run.py --engine_dir {args.engines_dir}', - f'--engine_name {args.ckpt}', - f'--model_name "{args.hf_models_dir}"', - f'--max_new_tokens={args.max_new_tokens}', - f'--num_beams={beam}', - f'--compare_hf_fp32', - f'--output_npy={args.data_dir}', - "--debug_mode" if args.debug else "", - )) - ret = [' '.join(x) for x in ret] - ret = ' && '.join(ret) - return ret - - -if __name__ == '__main__': - args = Arguments() - Run(args).run() diff --git a/cpp/tests/resources/scripts/generate_expected_gpt_output.py b/cpp/tests/resources/scripts/generate_expected_gpt_output.py deleted file mode 100755 index 16fa5cc8db64..000000000000 --- a/cpp/tests/resources/scripts/generate_expected_gpt_output.py +++ /dev/null @@ -1,191 +0,0 @@ -#!/usr/bin/env python3 -# 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 argparse -from pathlib import Path - -# isort: off -import run -# isort: on - -import os -import shutil - -import tensorrt_llm.bindings as _tb -from tensorrt_llm.bindings.internal.testing import ModelSpec, QuantMethod - - -def get_model_data_dir(): - resources_dir = Path(__file__).parent.resolve().parent - data_dir = resources_dir / 'data' - return data_dir / 'gpt2' - - -def generate_output(engine: str, - num_beams: int, - input_name: str, - model_spec_obj: ModelSpec, - max_output_len: int = 8, - output_logits: bool = False, - output_cum_log_probs: bool = False, - output_log_probs: bool = False): - tp_size = 1 - pp_size = 1 - cp_size = 1 - model = 'gpt2' - resources_dir = Path(__file__).parent.resolve().parent - models_dir = resources_dir / 'models' - tp_pp_cp_dir = 'tp' + str(tp_size) + '-pp' + str(pp_size) + '-cp' + str( - cp_size) + '-gpu/' - engine_dir = models_dir / 'rt_engine' / model / engine / tp_pp_cp_dir - - data_dir = resources_dir / 'data' - input_file = data_dir / input_name - model_data_dir = get_model_data_dir() - if num_beams <= 1: - output_dir = model_data_dir / 'sampling' - else: - output_dir = model_data_dir / ('beam_search_' + str(num_beams)) - - model_spec_obj.use_tensor_parallelism(tp_size).use_pipeline_parallelism( - pp_size).use_context_parallelism(cp_size) - - base_output_name = os.path.splitext(model_spec_obj.get_results_file())[0] - - args_list = [ - f'--engine_dir={engine_dir}', - f'--input_file={input_file}', - f'--tokenizer_dir={models_dir / model}', - f'--output_npy={output_dir / (base_output_name + ".npy")}', - f'--output_csv={output_dir / (base_output_name + ".csv")}', - f'--max_output_len={max_output_len}', - f'--num_beams={num_beams}', - '--use_py_session', - ] - - if output_logits: - args_list.extend([ - f'--output_logits_npy={output_dir / (base_output_name + "_logits.npy")}', - '--output_generation_logits', - ]) - - # Generate context_fmha_fp32_acc enabled results for GptExecutorTest.GenerationLogitsEarlyStop - if model_spec_obj.get_enable_context_fmha_fp32_acc(): - args_list.extend(["--enable_context_fmha_fp32_acc"]) - - if output_cum_log_probs: - args_list.extend([ - f'--output_cum_log_probs_npy={output_dir / model_spec_obj.get_cum_log_probs_file()}' - ]) - - if output_log_probs: - args_list.extend([ - f'--output_log_probs_npy={output_dir / model_spec_obj.get_log_probs_file()}' - ]) - - args = run.parse_arguments(args_list) - run.main(args) - - -def generate_outputs(num_beams): - input_name = 'input_tokens.npy' - input_name_long = 'input_tokens_long.npy' - - print('Generating GPT2 FP16 outputs') - model_spec_obj = ModelSpec(input_name, _tb.DataType.HALF) - model_spec_obj.use_gpt_plugin() - model_spec_obj.use_packed_input() - model_spec_obj.set_kv_cache_type(_tb.KVCacheType.PAGED) - model_spec_obj.gather_logits() - generate_output(engine=model_spec_obj.get_model_path(), - num_beams=num_beams, - input_name=input_name, - model_spec_obj=model_spec_obj, - output_logits=True, - output_log_probs=True, - output_cum_log_probs=True) - # GptExecutorTest.GenerationLogitsEarlyStop and several tests require to use context_fmha_fp32_acc flag in runtime - model_spec_obj.enable_context_fmha_fp32_acc() - generate_output(engine=model_spec_obj.get_model_path(), - num_beams=num_beams, - input_name=input_name, - model_spec_obj=model_spec_obj, - output_logits=True, - output_log_probs=True, - output_cum_log_probs=True) - - model_spec_obj = ModelSpec(input_name, _tb.DataType.HALF) - model_spec_obj.use_gpt_plugin() - model_spec_obj.set_kv_cache_type(_tb.KVCacheType.PAGED) - model_spec_obj.use_packed_input() - generate_output(engine=model_spec_obj.get_model_path(), - num_beams=num_beams, - input_name=input_name, - model_spec_obj=model_spec_obj, - output_logits=False, - output_log_probs=True, - output_cum_log_probs=True) - model_spec_obj.enable_context_fmha_fp32_acc() - generate_output(engine=model_spec_obj.get_model_path(), - num_beams=num_beams, - input_name=input_name, - model_spec_obj=model_spec_obj, - output_logits=False, - output_log_probs=True, - output_cum_log_probs=True) - model_spec_obj.set_max_output_length(128) - generate_output(engine=model_spec_obj.get_model_path(), - num_beams=num_beams, - input_name=input_name, - model_spec_obj=model_spec_obj, - output_logits=False, - max_output_len=128) - - model_spec_obj = ModelSpec(input_name_long, _tb.DataType.HALF) - model_spec_obj.use_gpt_plugin() - model_spec_obj.use_packed_input() - model_spec_obj.set_kv_cache_type(_tb.KVCacheType.PAGED) - generate_output(engine=model_spec_obj.get_model_path(), - num_beams=num_beams, - input_name=input_name_long, - model_spec_obj=model_spec_obj, - output_logits=False) - - model_spec_obj = ModelSpec(input_name, _tb.DataType.HALF) - model_spec_obj.use_gpt_plugin() - model_spec_obj.use_packed_input() - model_spec_obj.set_kv_cache_type(_tb.KVCacheType.PAGED) - model_spec_obj.set_quant_method(QuantMethod.SMOOTH_QUANT) - generate_output(engine=model_spec_obj.get_model_path(), - num_beams=num_beams, - input_name=input_name, - model_spec_obj=model_spec_obj, - output_logits=False) - - -if __name__ == '__main__': - parser = argparse.ArgumentParser() - parser.add_argument('--clean', - action='store_true', - default=False, - help='Clean target folders before building engines') - args = parser.parse_args() - if args.clean: - model_data_dir = get_model_data_dir() - print(f'Cleaning target folder {model_data_dir}') - shutil.rmtree(model_data_dir, ignore_errors=True) - generate_outputs(num_beams=1) - generate_outputs(num_beams=2) diff --git a/cpp/tests/resources/scripts/generate_expected_gptj_output.py b/cpp/tests/resources/scripts/generate_expected_gptj_output.py deleted file mode 100755 index 8d650d6bfc31..000000000000 --- a/cpp/tests/resources/scripts/generate_expected_gptj_output.py +++ /dev/null @@ -1,108 +0,0 @@ -#!/usr/bin/env python3 -# 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 argparse as _arg -import os -from pathlib import Path - -# isort: off -import run -# isort: on - -import tensorrt_llm.bindings as _tb -from tensorrt_llm.bindings.internal.testing import ModelSpec - - -def generate_output(engine: str, - num_beams: int, - model_spec_obj: ModelSpec, - max_output_len: int = 4): - - tp_size = 1 - pp_size = 1 - cp_size = 1 - model = 'gpt-j-6b' - resources_dir = Path(__file__).parent.resolve().parent - models_dir = resources_dir / 'models' - hf_dir = models_dir / model - tp_pp_cp_dir = 'tp' + str(tp_size) + '-pp' + str(pp_size) + '-cp' + str( - cp_size) + '-gpu/' - engine_dir = models_dir / 'rt_engine' / model / engine / tp_pp_cp_dir - - data_dir = resources_dir / 'data' - input_file = data_dir / 'input_tokens.npy' - model_data_dir = data_dir / model - if num_beams <= 1: - output_dir = model_data_dir / 'sampling' - else: - output_dir = model_data_dir / ('beam_search_' + str(num_beams)) - - base_output_name = os.path.splitext(model_spec_obj.get_results_file())[0] - - args = run.parse_arguments([ - '--engine_dir', - str(engine_dir), '--input_file', - str(input_file), '--tokenizer_dir', - str(hf_dir), '--output_npy', - str(output_dir / (base_output_name + '.npy')), '--output_csv', - str(output_dir / (base_output_name + '.csv')), '--max_output_len', - str(max_output_len), '--num_beams', - str(num_beams), '--use_py_session' - ]) - run.main(args) - - -def generate_outputs(only_fp8, num_beams): - input_file = 'input_tokens.npy' - if only_fp8 and num_beams == 1: - model_spec_obj = ModelSpec(input_file, _tb.DataType.FP8) - model_spec_obj.use_gpt_plugin() - model_spec_obj.set_kv_cache_type(_tb.KVCacheType.PAGED) - model_spec_obj.use_packed_input() - - print('Generating GPT-J FP8-kv-cache outputs') - generate_output(engine=model_spec_obj.get_model_path(), - num_beams=num_beams, - model_spec_obj=model_spec_obj) - elif not only_fp8: - print('Generating GPT-J FP16 outputs') - model_spec_obj = ModelSpec(input_file, _tb.DataType.HALF) - model_spec_obj.use_gpt_plugin() - model_spec_obj.set_kv_cache_type(_tb.KVCacheType.CONTINUOUS) - generate_output(engine=model_spec_obj.get_model_path(), - num_beams=num_beams, - model_spec_obj=model_spec_obj) - - model_spec_obj.use_packed_input() - generate_output(engine=model_spec_obj.get_model_path(), - num_beams=num_beams, - model_spec_obj=model_spec_obj) - - model_spec_obj.set_kv_cache_type(_tb.KVCacheType.PAGED) - generate_output(engine=model_spec_obj.get_model_path(), - num_beams=num_beams, - model_spec_obj=model_spec_obj) - - -if __name__ == '__main__': - parser = _arg.ArgumentParser() - parser.add_argument( - "--only_fp8", - action="store_true", - help="Generate data for only FP8 tests. Implemented for H100 runners.") - - generate_outputs(**vars(parser.parse_args()), num_beams=1) - generate_outputs(**vars(parser.parse_args()), num_beams=2) diff --git a/cpp/tests/resources/scripts/generate_expected_llama_output.py b/cpp/tests/resources/scripts/generate_expected_llama_output.py deleted file mode 100644 index 74916e77d053..000000000000 --- a/cpp/tests/resources/scripts/generate_expected_llama_output.py +++ /dev/null @@ -1,151 +0,0 @@ -#!/usr/bin/env python3 -# 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 argparse as _arg -import os -import time -from pathlib import Path - -from mpi4py.MPI import COMM_WORLD - -# isort: off -import run -# isort: on - -import tensorrt_llm.bindings as _tb -from tensorrt_llm.bindings.internal.testing import ModelSpec - - -def generate_output(engine: str, - num_beams: int, - model_spec_obj: ModelSpec, - tp_size: int = 1, - pp_size: int = 1, - cp_size: int = 1, - max_output_len: int = 8, - output_logits: bool = False, - output_cum_log_probs: bool = False, - output_log_probs: bool = False): - - model = 'Llama-3.2-1B' - resources_dir = Path(__file__).parent.resolve().parent - models_dir = resources_dir / 'models' - tp_pp_cp_dir = 'tp' + str(tp_size) + '-pp' + str(pp_size) + '-cp' + str( - cp_size) + '-gpu/' - engine_dir = models_dir / 'rt_engine' / model / engine / tp_pp_cp_dir - - data_dir = resources_dir / 'data' - input_file = data_dir / 'input_tokens_llama.npy' - model_data_dir = data_dir / model - if num_beams <= 1: - output_dir = model_data_dir / 'sampling' - else: - output_dir = model_data_dir / ('beam_search_' + str(num_beams)) - - base_output_name = os.path.splitext(model_spec_obj.get_results_file())[0] - - args_list = [ - f'--engine_dir={engine_dir}', - f'--input_file={input_file}', - f'--tokenizer_dir={models_dir / model}', - f'--output_npy={output_dir / (base_output_name + ".npy")}', - f'--output_csv={output_dir / (base_output_name + ".csv")}', - f'--max_output_len={max_output_len}', - f'--num_beams={num_beams}', - '--use_py_session', - ] - - if output_logits: - args_list.extend([ - f'--output_logits_npy={output_dir / (base_output_name + "_logits.npy")}', - '--output_generation_logits', - ]) - - if output_cum_log_probs: - args_list.extend([ - f'--output_cum_log_probs_npy={output_dir / model_spec_obj.get_cum_log_probs_file()}' - ]) - - if output_log_probs: - args_list.extend([ - f'--output_log_probs_npy={output_dir / model_spec_obj.get_log_probs_file()}' - ]) - - args = run.parse_arguments(args_list) - run.main(args) - - -def generate_outputs(num_beams, only_multi_gpu=False): - if not only_multi_gpu: - tp_pp_cp_sizes = [(1, 1, 1)] - elif COMM_WORLD.size == 4: - tp_pp_cp_sizes = [(4, 1, 1), (2, 2, 1), (1, 4, 1)] - elif COMM_WORLD.size == 2: - tp_pp_cp_sizes = [(1, 2, 1), (2, 1, 1)] - else: - raise RuntimeError( - f"The world size of MPI {COMM_WORLD.size} is not equal to 1, 2, or 4." - ) - model_spec_obj = ModelSpec('input_tokens_llama.npy', _tb.DataType.HALF) - model_spec_obj.use_gpt_plugin() - model_spec_obj.set_kv_cache_type(_tb.KVCacheType.PAGED) - model_spec_obj.use_packed_input() - - for tp_size, pp_size, cp_size in tp_pp_cp_sizes: - print( - f'Generating outputs for Llama FP16 with TP={tp_size}, PP={pp_size}, CP={cp_size}, BW={num_beams}' - ) - start_time = time.time() - - output_logits = False - output_log_probs = False - output_cum_log_probs = False - if tp_size == 4 and pp_size == 1: - output_logits = True - output_log_probs = True - output_cum_log_probs = True - - model_spec_obj.use_tensor_parallelism(tp_size) - model_spec_obj.use_pipeline_parallelism(pp_size) - model_spec_obj.use_context_parallelism(cp_size) - generate_output(engine=model_spec_obj.get_model_path(), - num_beams=num_beams, - tp_size=tp_size, - pp_size=pp_size, - cp_size=cp_size, - model_spec_obj=model_spec_obj, - output_logits=output_logits, - output_log_probs=output_log_probs, - output_cum_log_probs=output_cum_log_probs) - - duration = time.time() - start_time - print( - f"Generating outputs for Llama FP16 with TP={tp_size}, PP={pp_size}, CP={cp_size}, BW={num_beams} took {duration} seconds" - ) - - -if __name__ == '__main__': - parser = _arg.ArgumentParser() - parser.add_argument( - "--only_multi_gpu", - action="store_true", - help="Generate data with Pipeline and Tensor Parallelism") - - args = parser.parse_args() - - generate_outputs(num_beams=1, only_multi_gpu=args.only_multi_gpu) - generate_outputs(num_beams=2, only_multi_gpu=args.only_multi_gpu) - print("Done") diff --git a/cpp/tests/resources/scripts/generate_expected_mamba_output.py b/cpp/tests/resources/scripts/generate_expected_mamba_output.py deleted file mode 100644 index 16779c434775..000000000000 --- a/cpp/tests/resources/scripts/generate_expected_mamba_output.py +++ /dev/null @@ -1,107 +0,0 @@ -#!/usr/bin/env python3 -# 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 os -from pathlib import Path - -# isort: off -import run -# isort: on - -import tensorrt_llm.bindings as _tb -from tensorrt_llm.bindings.internal.testing import ModelSpec - - -def generate_output(engine: str, - num_beams: int, - input_name: str, - model_spec_obj: ModelSpec, - max_output_len: int = 8, - output_logits: bool = False): - tp_size = 1 - pp_size = 1 - cp_size = 1 - model = 'mamba-2.8b-hf' - resources_dir = Path(__file__).parent.resolve().parent - models_dir = resources_dir / 'models' - tp_pp_cp_dir = 'tp' + str(tp_size) + '-pp' + str(pp_size) + '-cp' + str( - cp_size) + '-gpu/' - engine_dir = models_dir / 'rt_engine' / model / engine / tp_pp_cp_dir - - data_dir = resources_dir / 'data' - input_file = data_dir / input_name - model_data_dir = data_dir / model - if num_beams <= 1: - output_dir = model_data_dir / 'sampling' - else: - output_dir = model_data_dir / ('beam_search_' + str(num_beams)) - - base_output_name = os.path.splitext(model_spec_obj.get_results_file())[0] - - output_logits_npy = None - if output_logits: - output_logits_npy = str(output_dir / - (base_output_name + '_logits' + '.npy')) - - args = run.parse_arguments([ - '--engine_dir', - str(engine_dir), '--input_file', - str(input_file), '--tokenizer_dir', - str(models_dir / 'gpt-neox-20b'), '--output_npy', - str(output_dir / (base_output_name + '.npy')), '--output_csv', - str(output_dir / (base_output_name + '.csv')), '--max_output_len', - str(max_output_len), '--num_beams', - str(num_beams), '--output_logits_npy', - str(output_logits_npy), '--use_py_session' - ]) - run.main(args) - - -def generate_outputs(num_beams): - print('Generating Mamba FP16 outputs') - input_name = 'input_tokens.npy' - model_spec_obj = ModelSpec(input_name, _tb.DataType.HALF) - model_spec_obj.set_kv_cache_type(_tb.KVCacheType.CONTINUOUS) - - generate_output(engine=model_spec_obj.get_model_path(), - num_beams=num_beams, - input_name=input_name, - model_spec_obj=model_spec_obj) - - print('Generating Mamba FP16-plugin outputs') - model_spec_obj.use_gpt_plugin() - generate_output(engine=model_spec_obj.get_model_path(), - num_beams=num_beams, - input_name=input_name, - model_spec_obj=model_spec_obj) - - print('Generating Mamba FP16-plugin-packed outputs') - model_spec_obj.use_packed_input() - generate_output(engine=model_spec_obj.get_model_path(), - num_beams=num_beams, - input_name=input_name, - model_spec_obj=model_spec_obj) - - print('Generating Mamba FP16-plugin-packed-paged outputs') - model_spec_obj.set_kv_cache_type(_tb.KVCacheType.PAGED) - generate_output(engine=model_spec_obj.get_model_path(), - num_beams=num_beams, - input_name=input_name, - model_spec_obj=model_spec_obj) - - -if __name__ == '__main__': - generate_outputs(num_beams=1) diff --git a/cpp/tests/resources/scripts/generate_expected_medusa_output.py b/cpp/tests/resources/scripts/generate_expected_medusa_output.py deleted file mode 100755 index e1cbc20c051b..000000000000 --- a/cpp/tests/resources/scripts/generate_expected_medusa_output.py +++ /dev/null @@ -1,88 +0,0 @@ -#!/usr/bin/env python3 -# 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 argparse as _arg -import os -from pathlib import Path - -# isort: off -import run -# isort: on - -import tensorrt_llm.bindings as _tb -from tensorrt_llm.bindings.internal.testing import ModelSpec - - -def generate_output(engine: str, - model_spec_obj: ModelSpec, - max_output_len: int = 8): - - model = 'vicuna-7b-medusa' - hf_model = 'vicuna-7b-v1.3' - resources_dir = Path(__file__).parent.resolve().parent - models_dir = resources_dir / 'models' - hf_dir = models_dir / hf_model - tp_pp_dir = 'tp1-pp1-cp1-gpu/' - engine_dir = models_dir / 'rt_engine' / model / engine / tp_pp_dir - - data_dir = resources_dir / 'data' - input_file = data_dir / 'input_vicuna.npy' - model_data_dir = data_dir / model - output_dir = model_data_dir / 'sampling' - - base_output_name = os.path.splitext(model_spec_obj.get_results_file())[0] - - args = run.parse_arguments([ - '--engine_dir', - str(engine_dir), '--input_file', - str(input_file), '--tokenizer_dir', - str(hf_dir), '--output_npy', - str(output_dir / (base_output_name + '.npy')), '--output_csv', - str(output_dir / (base_output_name + '.csv')), '--max_output_len', - str(max_output_len), '--use_py_session', - '--medusa_choices=[[0], [0, 0], [1], [0, 1], [2], [0, 0, 0], [1, 0], [0, 2], [3], [0, 3], [4], [0, 4], [2, 0], [0, 5], [0, 0, 1], [5], [0, 6], [6], [0, 7], [0, 1, 0], [1, 1], [7], [0, 8], [0, 0, 2], [3, 0], [0, 9], [8], [9], [1, 0, 0], [0, 2, 0], [1, 2], [0, 0, 3], [4, 0], [2, 1], [0, 0, 4], [0, 0, 5], [0, 0, 0, 0], [0, 1, 1], [0, 0, 6], [0, 3, 0], [5, 0], [1, 3], [0, 0, 7], [0, 0, 8], [0, 0, 9], [6, 0], [0, 4, 0], [1, 4], [7, 0], [0, 1, 2], [2, 0, 0], [3, 1], [2, 2], [8, 0], [0, 5, 0], [1, 5], [1, 0, 1], [0, 2, 1], [9, 0], [0, 6, 0], [0, 0, 0, 1], [1, 6], [0, 7, 0]]', - '--temperature', '1.0' - ]) - run.main(args) - print(f"Output saved at {str(output_dir / base_output_name)}.[npy|csv]") - - -def generate_outputs(): - print(f'Generating outputs for Medusa FP16') - max_output_len = 128 - model_spec_obj = ModelSpec('input_tokens_long.npy', _tb.DataType.HALF) - model_spec_obj.use_gpt_plugin() - model_spec_obj.set_max_output_length(max_output_len) - model_spec_obj.use_packed_input() - model_spec_obj.set_kv_cache_type(_tb.KVCacheType.PAGED) - model_spec_obj.use_medusa() - - generate_output(engine=model_spec_obj.get_model_path(), - model_spec_obj=model_spec_obj, - max_output_len=max_output_len) - - -if __name__ == '__main__': - parser = _arg.ArgumentParser() - parser.add_argument( - "--only_multi_gpu", - action="store_true", - help="Generate data with Pipeline and Tensor Parallelism") - - args = parser.parse_args() - - generate_outputs() - print("Done") diff --git a/cpp/tests/resources/scripts/generate_expected_recurrentgemma_output.py b/cpp/tests/resources/scripts/generate_expected_recurrentgemma_output.py deleted file mode 100644 index 0ef4cc4509fd..000000000000 --- a/cpp/tests/resources/scripts/generate_expected_recurrentgemma_output.py +++ /dev/null @@ -1,88 +0,0 @@ -#!/usr/bin/env python3 -# 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 os -from pathlib import Path - -# isort: off -import run -# isort: on - -import tensorrt_llm.bindings as _tb -from tensorrt_llm.bindings.internal.testing import ModelSpec - - -def generate_output(engine: str, - num_beams: int, - input_name: str, - model_spec_obj: ModelSpec, - max_output_len: int = 8, - output_logits: bool = False): - tp_size = 1 - pp_size = 1 - cp_size = 1 - model = 'recurrentgemma-2b' - resources_dir = Path(__file__).parent.resolve().parent - models_dir = resources_dir / 'models' - tp_pp_cp_dir = 'tp' + str(tp_size) + '-pp' + str(pp_size) + '-cp' + str( - cp_size) + '-gpu/' - engine_dir = models_dir / 'rt_engine' / model / engine / tp_pp_cp_dir - - data_dir = resources_dir / 'data' - input_file = data_dir / input_name - model_data_dir = data_dir / model - if num_beams <= 1: - output_dir = model_data_dir / 'sampling' - else: - output_dir = model_data_dir / ('beam_search_' + str(num_beams)) - - base_output_name = os.path.splitext(model_spec_obj.get_results_file())[0] - - output_logits_npy = None - if output_logits: - output_logits_npy = str(output_dir / - (base_output_name + '_logits' + '.npy')) - - args = run.parse_arguments([ - '--engine_dir', - str(engine_dir), '--input_file', - str(input_file), '--tokenizer_dir', - str(models_dir / model), '--output_npy', - str(output_dir / (base_output_name + '.npy')), '--output_csv', - str(output_dir / (base_output_name + '.csv')), '--max_output_len', - str(max_output_len), '--num_beams', - str(num_beams), '--output_logits_npy', - str(output_logits_npy), '--use_py_session' - ]) - run.main(args) - - -def generate_outputs(num_beams): - input_file = 'input_tokens.npy' - model_spec_obj = ModelSpec(input_file, _tb.DataType.HALF) - model_spec_obj.use_gpt_plugin() - model_spec_obj.set_kv_cache_type(_tb.KVCacheType.PAGED) - model_spec_obj.use_packed_input() - - print('Generating RecurrentGemma FP16-plugin-packed-paged outputs') - generate_output(engine=model_spec_obj.get_model_path(), - num_beams=num_beams, - input_name=input_file, - model_spec_obj=model_spec_obj) - - -if __name__ == '__main__': - generate_outputs(num_beams=1) diff --git a/cpp/tests/resources/scripts/generate_expected_redrafter_output.py b/cpp/tests/resources/scripts/generate_expected_redrafter_output.py deleted file mode 100644 index 989e029a5ab1..000000000000 --- a/cpp/tests/resources/scripts/generate_expected_redrafter_output.py +++ /dev/null @@ -1,88 +0,0 @@ -#!/usr/bin/env python3 -# SPDX-FileCopyrightText: Copyright (c) 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 argparse as _arg -import os -from pathlib import Path - -# isort: off -import run -# isort: on - -import tensorrt_llm.bindings as _tb -from tensorrt_llm.bindings.internal.testing import ModelSpec - - -def generate_output(engine: str, - model_spec_obj: ModelSpec, - max_output_len: int = 8): - - model = 'vicuna-7b-redrafter' - hf_model = 'vicuna-7b-v1.3' - resources_dir = Path(__file__).parent.resolve().parent - models_dir = resources_dir / 'models' - hf_dir = models_dir / hf_model - tp_pp_dir = 'tp1-pp1-cp1-gpu/' - engine_dir = models_dir / 'rt_engine' / model / engine / tp_pp_dir - - data_dir = resources_dir / 'data' - input_filename = model_spec_obj.get_input_file() - input_file = data_dir / input_filename - model_data_dir = data_dir / model - output_dir = model_data_dir / 'sampling' - - base_output_name = os.path.splitext(model_spec_obj.get_results_file())[0] - - args = run.parse_arguments([ - '--engine_dir', - str(engine_dir), - '--input_file', - str(input_file), - '--tokenizer_dir', - str(hf_dir), - '--output_npy', - str(output_dir / (base_output_name + '.npy')), - '--output_csv', - str(output_dir / (base_output_name + '.csv')), - '--max_output_len', - str(max_output_len), - '--use_py_session', - ]) - run.main(args) - print(f"Output saved at {str(output_dir / base_output_name)}.[npy|csv]") - - -def generate_outputs(): - print(f'Generating outputs for ReDrafter FP16') - max_output_len = 128 - model_spec_obj = ModelSpec('input_vicuna.npy', _tb.DataType.HALF) - model_spec_obj.use_gpt_plugin() - model_spec_obj.set_max_output_length(max_output_len) - model_spec_obj.use_packed_input() - model_spec_obj.set_kv_cache_type(_tb.KVCacheType.PAGED) - model_spec_obj.use_explicit_draft_tokens_decoding() - - generate_output(engine=model_spec_obj.get_model_path(), - model_spec_obj=model_spec_obj, - max_output_len=max_output_len) - - -if __name__ == '__main__': - parser = _arg.ArgumentParser() - args = parser.parse_args() - - generate_outputs() - print("Done") diff --git a/cpp/tests/resources/scripts/generate_hf_gpt_output.py b/cpp/tests/resources/scripts/generate_hf_gpt_output.py deleted file mode 100755 index a40ada8cb455..000000000000 --- a/cpp/tests/resources/scripts/generate_hf_gpt_output.py +++ /dev/null @@ -1,51 +0,0 @@ -#!/usr/bin/env python3 -# 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. - -from pathlib import Path - -import run_hf - - -def generate_hf_output(data_type: str, - output_name: str, - max_output_len: int = 8): - - model = 'gpt2' - resources_dir = Path(__file__).parent.resolve().parent - models_dir = resources_dir / 'models' - model_dir = models_dir / model - - data_dir = resources_dir / 'data' - input_file = data_dir / 'input_tokens.npy' - output_dir = data_dir / model / 'huggingface' - - run_hf.generate(model_dir=str(model_dir), - data_type=data_type, - input_file=str(input_file), - output_npy=str(output_dir / (output_name + '.npy')), - output_csv=str(output_dir / (output_name + '.csv')), - max_output_len=max_output_len) - - -def generate_hf_outputs(): - generate_hf_output(data_type='fp32', - output_name='output_tokens_fp32_huggingface') - generate_hf_output(data_type='fp16', - output_name='output_tokens_fp16_huggingface') - - -if __name__ == '__main__': - generate_hf_outputs() diff --git a/cpp/tests/resources/scripts/io_converter.py b/cpp/tests/resources/scripts/io_converter.py deleted file mode 100755 index 0ed6413c5ac4..000000000000 --- a/cpp/tests/resources/scripts/io_converter.py +++ /dev/null @@ -1,74 +0,0 @@ -#!/usr/bin/env python3 -# 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 argparse -import csv -import os - -import numpy as np - - -def csv_to_npy(input_file, output_file, pad_id, verbose): - data = [] - with open(input_file, newline='') as csvfile: - csv_reader = csv.reader(csvfile, delimiter=',') - for line in csv_reader: - data.append([int(e) for e in line]) - max_input_length = max([len(x) for x in data]) - data = [row + [pad_id] * (max_input_length - len(row)) for row in data] - data = np.array(data, dtype='int32') - if (verbose): - print(data, data.dtype) - np.save(output_file, data) - - -def npy_to_csv(input_file, output_file, verbose): - data = np.load(input_file) - if (verbose): - print(data, data.dtype) - np.savetxt(output_file, data, delimiter=",", fmt='%i') - - -if __name__ == '__main__': - parser = argparse.ArgumentParser() - parser.add_argument( - 'input_file', - type=str, - help='Read token ids from this file. Must be csv or npy.') - parser.add_argument('output_file', - type=str, - help='Write token ids this file. Must be csv or npy.') - parser.add_argument( - '-p', - '--pad_id', - type=int, - help= - 'Token id used for padding csv input with different sequence lengths.', - default=-1) - parser.add_argument('-v', '--verbose', action="store_true") - args = parser.parse_args() - - _, input_ext = os.path.splitext(args.input_file) - _, output_ext = os.path.splitext(args.output_file) - - if (input_ext == '.csv' and output_ext == '.npy'): - print('Converting csv to npy') - csv_to_npy(args.input_file, args.output_file, args.pad_id, args.verbose) - elif (input_ext == '.npy' and output_ext == '.csv'): - print('Converting npy to csv') - npy_to_csv(args.input_file, args.output_file, args.verbose) - else: - print('unknown file extensions') diff --git a/cpp/tests/unit_tests/CMakeLists.txt b/cpp/tests/unit_tests/CMakeLists.txt index 034de457bb78..9d22bd03b52a 100644 --- a/cpp/tests/unit_tests/CMakeLists.txt +++ b/cpp/tests/unit_tests/CMakeLists.txt @@ -27,4 +27,3 @@ add_subdirectory(multi_gpu) add_subdirectory(layers) add_subdirectory(runtime) add_subdirectory(thop) -add_subdirectory(utils) diff --git a/cpp/tests/unit_tests/batch_manager/CMakeLists.txt b/cpp/tests/unit_tests/batch_manager/CMakeLists.txt index f4e5e9fb2be0..a8f6c7dd4e8f 100644 --- a/cpp/tests/unit_tests/batch_manager/CMakeLists.txt +++ b/cpp/tests/unit_tests/batch_manager/CMakeLists.txt @@ -29,7 +29,5 @@ add_gtest(microBatchSchedulerTest microBatchSchedulerTest.cpp) add_gtest(peftCacheManagerTest peftCacheManagerTest.cpp) add_gtest(staticThreadPoolTest staticThreadPoolTest.cpp) add_gtest(rnnCacheFormatterTest rnnCacheFormatterTest.cpp) -add_gtest(cudaGraphExecutorCacheTest cudaGraphExecutorCacheTest.cpp) add_gtest(agentTreeTest agentTreeTest.cpp) add_gtest(truncateBlocksTest truncateBlocksTest.cpp) -add_gtest(encDecBeamSearchTest encDecBeamSearchTest.cpp) diff --git a/cpp/tests/unit_tests/batch_manager/bufferIndexHolderTest.cpp b/cpp/tests/unit_tests/batch_manager/bufferIndexHolderTest.cpp index 00664f4aeb4e..31400831a5bb 100644 --- a/cpp/tests/unit_tests/batch_manager/bufferIndexHolderTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/bufferIndexHolderTest.cpp @@ -18,7 +18,7 @@ #include "tensorrt_llm/batch_manager/baseTransBuffer.h" #include "tensorrt_llm/batch_manager/cacheTransBuffer.h" #include "tensorrt_llm/batch_manager/kvCacheManager.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include #include #include @@ -95,7 +95,7 @@ class BufferIndexHolderLifecycleTest : public ::testing::TestWithParam(numLayers, numHeads, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, maxBeamWidth, std::vector{kvMaxNumTokens}, - nvinfer1::DataType::kFLOAT, sinkTokenLength, stream, kvMaxNumTokens, kvMaxNumTokens, + tensorrt_llm::DataType::kFLOAT, sinkTokenLength, stream, kvMaxNumTokens, kvMaxNumTokens, /*enableBlockReuse=*/true, CacheType::kSELF, std::nullopt, nullptr, true); mKv->allocatePools(false); mTrans = std::make_unique(mKv.get(), std::optional{kvMaxNumTokens}); diff --git a/cpp/tests/unit_tests/batch_manager/cacheTransBufferTest.cpp b/cpp/tests/unit_tests/batch_manager/cacheTransBufferTest.cpp index 8150c6fa5406..2fa0477d2352 100644 --- a/cpp/tests/unit_tests/batch_manager/cacheTransBufferTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/cacheTransBufferTest.cpp @@ -18,6 +18,7 @@ #include "tensorrt_llm/batch_manager/cacheTransBuffer.h" #include "tensorrt_llm/batch_manager/kvCacheManager.h" #include "tensorrt_llm/common/envUtils.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/executor/executor.h" #include "tensorrt_llm/runtime/bufferManager.h" #include "tensorrt_llm/runtime/iTensor.h" @@ -51,7 +52,7 @@ class CacheTransBufferTest : public ::testing::Test auto constexpr blocksInSecondaryPool = 0; auto constexpr enableBlockReuse = true; - auto constexpr dataType = nvinfer1::DataType::kFLOAT; + auto constexpr dataType = tensorrt_llm::DataType::kFLOAT; using BlocksPerWindow = std::map>; auto const blocksPerWindow = BlocksPerWindow{{maxAttentionWindow, {totalNumBlocks, blocksInSecondaryPool}}}; diff --git a/cpp/tests/unit_tests/batch_manager/capacitySchedulerTest.cpp b/cpp/tests/unit_tests/batch_manager/capacitySchedulerTest.cpp index 7bb87c91e361..2c9757b847b2 100644 --- a/cpp/tests/unit_tests/batch_manager/capacitySchedulerTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/capacitySchedulerTest.cpp @@ -31,7 +31,7 @@ #include "tensorrt_llm/executor/types.h" #include "tensorrt_llm/testing/kvCacheManagerTestUtil.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include #include @@ -132,7 +132,7 @@ class CapacitySchedulerTest : public ::testing::Test // NOLINT(cppcoreguidelines auto const nbKvHeads = 10; auto constexpr sizePerHead = 1; auto const maxNumBlocks = tc::divUp(maxNumTokens, tokensPerBlock); - auto const kvDtype = nvinfer1::DataType::kHALF; + auto const kvDtype = tensorrt_llm::DataType::kHALF; CudaStreamPtr streamPtr = std::make_shared(); using BlocksPerWindow = std::map>; diff --git a/cpp/tests/unit_tests/batch_manager/cudaGraphExecutorCacheTest.cpp b/cpp/tests/unit_tests/batch_manager/cudaGraphExecutorCacheTest.cpp deleted file mode 100644 index f8ae0a9db64f..000000000000 --- a/cpp/tests/unit_tests/batch_manager/cudaGraphExecutorCacheTest.cpp +++ /dev/null @@ -1,161 +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. - */ - -#include "tensorrt_llm/batch_manager/common.h" -#include "tensorrt_llm/batch_manager/utils/inflightBatchingUtils.h" - -#include - -#include - -namespace tb = tensorrt_llm::batch_manager; -namespace tbu = tensorrt_llm::batch_manager::utils; -using SizeType32 = tensorrt_llm::runtime::SizeType32; - -namespace -{ -// A default-constructed CudaGraphExecutor holds mInstance == nullptr, so its destructor -// is a no-op and these tests do not require an active CUDA context. -std::shared_ptr makeDummyExecutor() -{ - return std::make_shared(); -} - -tb::BatchState makeBatchState(SizeType32 numTokens) -{ - return tb::BatchState{/*numCtxRequests=*/0, /*numGenRequests=*/1, numTokens, /*maxKvCacheLength=*/256}; -} -} // namespace - -class CudaGraphExecutorCacheTest : public ::testing::Test // NOLINT(cppcoreguidelines-pro-type-member-init) -{ -}; - -TEST_F(CudaGraphExecutorCacheTest, EmptyByDefault) -{ - tbu::CudaGraphExecutorCache cache(/*capacity=*/4); - EXPECT_EQ(cache.size(), 0); - EXPECT_FALSE(cache.get(makeBatchState(1)).has_value()); -} - -TEST_F(CudaGraphExecutorCacheTest, PutAndGetReturnsSameInstance) -{ - tbu::CudaGraphExecutorCache cache(/*capacity=*/4); - - auto bs = makeBatchState(1); - auto exec = makeDummyExecutor(); - cache.put(bs, exec); - - ASSERT_EQ(cache.size(), 1); - auto got = cache.get(bs); - ASSERT_TRUE(got.has_value()); - EXPECT_EQ(got->get(), exec.get()); -} - -TEST_F(CudaGraphExecutorCacheTest, PutWithExistingKeyReplaces) -{ - tbu::CudaGraphExecutorCache cache(/*capacity=*/4); - - auto bs = makeBatchState(1); - auto first = makeDummyExecutor(); - auto second = makeDummyExecutor(); - - cache.put(bs, first); - cache.put(bs, second); - - EXPECT_EQ(cache.size(), 1); - auto got = cache.get(bs); - ASSERT_TRUE(got.has_value()); - EXPECT_EQ(got->get(), second.get()); -} - -TEST_F(CudaGraphExecutorCacheTest, EvictsLeastRecentlyUsedAtCapacity) -{ - tbu::CudaGraphExecutorCache cache(/*capacity=*/2); - - auto bsA = makeBatchState(1); - auto bsB = makeBatchState(2); - auto bsC = makeBatchState(3); - - auto execA = makeDummyExecutor(); - auto execB = makeDummyExecutor(); - auto execC = makeDummyExecutor(); - - cache.put(bsA, execA); - cache.put(bsB, execB); - ASSERT_EQ(cache.size(), 2); - - // Access A so that B becomes the LRU entry. - EXPECT_TRUE(cache.get(bsA).has_value()); - - // Inserting C must evict B (the LRU), not A (just touched). - cache.put(bsC, execC); - EXPECT_EQ(cache.size(), 2); - EXPECT_TRUE(cache.get(bsA).has_value()); - EXPECT_FALSE(cache.get(bsB).has_value()); - EXPECT_TRUE(cache.get(bsC).has_value()); -} - -TEST_F(CudaGraphExecutorCacheTest, ClearDropsAllEntries) -{ - tbu::CudaGraphExecutorCache cache(/*capacity=*/4); - - auto bsA = makeBatchState(1); - auto bsB = makeBatchState(2); - auto bsC = makeBatchState(3); - - cache.put(bsA, makeDummyExecutor()); - cache.put(bsB, makeDummyExecutor()); - cache.put(bsC, makeDummyExecutor()); - ASSERT_EQ(cache.size(), 3); - - cache.clear(); - - EXPECT_EQ(cache.size(), 0); - EXPECT_FALSE(cache.get(bsA).has_value()); - EXPECT_FALSE(cache.get(bsB).has_value()); - EXPECT_FALSE(cache.get(bsC).has_value()); - - // After clearing, the cache must remain functional (i.e. clear() must not - // leave it in a broken state). - auto execA2 = makeDummyExecutor(); - cache.put(bsA, execA2); - EXPECT_EQ(cache.size(), 1); - auto got = cache.get(bsA); - ASSERT_TRUE(got.has_value()); - EXPECT_EQ(got->get(), execA2.get()); -} - -TEST_F(CudaGraphExecutorCacheTest, ClearReleasesExecutorOwnership) -{ - tbu::CudaGraphExecutorCache cache(/*capacity=*/4); - - auto exec = makeDummyExecutor(); - std::weak_ptr weak = exec; - - cache.put(makeBatchState(1), exec); - exec.reset(); - // The cache still owns one strong reference at this point. - ASSERT_FALSE(weak.expired()); - - cache.clear(); - - // After clear(), no strong references should remain. This guarantees that - // ~CudaGraphExecutor (which calls cudaGraphExecDestroy) actually runs for - // every cached entry - exactly what changeBeamWidth() relies on. - EXPECT_TRUE(weak.expired()); -} diff --git a/cpp/tests/unit_tests/batch_manager/encDecBeamSearchTest.cpp b/cpp/tests/unit_tests/batch_manager/encDecBeamSearchTest.cpp deleted file mode 100644 index 8e96c6a52ee3..000000000000 --- a/cpp/tests/unit_tests/batch_manager/encDecBeamSearchTest.cpp +++ /dev/null @@ -1,154 +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. - */ - -#include "tensorrt_llm/batch_manager/kvCacheManager.h" -#include "tensorrt_llm/batch_manager/llmRequest.h" -#include "tensorrt_llm/batch_manager/runtimeBuffers.h" -#include "tensorrt_llm/batch_manager/utils/inflightBatchingUtils.h" -#include "tensorrt_llm/common/memoryUtils.h" -#include "tensorrt_llm/kernels/kvCacheIndex.h" -#include "tensorrt_llm/runtime/bufferManager.h" -#include "tensorrt_llm/runtime/cudaStream.h" -#include "tensorrt_llm/runtime/iTensor.h" -#include "tensorrt_llm/runtime/samplingConfig.h" -#include "gtest/gtest.h" -#include - -using namespace tensorrt_llm::batch_manager; -using namespace tensorrt_llm::batch_manager::kv_cache_manager; -namespace tr = tensorrt_llm::runtime; -namespace tc = tensorrt_llm::common; -namespace tk = tensorrt_llm::kernels; -using SizeType32 = tr::SizeType32; - -// Verify that copyGenerationLogits correctly assembles the host logits buffer -// using the real kernel merge path, and that two back-to-back calls (simulating -// two requests flushing in the same batch) use distinct fragmentPointerDevice -// slots so their pointer arrays do not clobber each other. -TEST(CopyGenerationLogitsTest, KernelMergePathProducesCorrectHostLayoutAndSlotsAreIsolated) -{ - SizeType32 constexpr beamWidth = 2; - SizeType32 constexpr numSteps = RuntimeBuffers::GenerationLogitsCache::kCACHE_LENGTH; // full flush - SizeType32 constexpr vocabSize = 8; - SizeType32 constexpr promptLen = 1; - SizeType32 constexpr maxBatchSize = 4; // must be >= 2 to test slot isolation - - auto stream = std::make_shared(); - tr::BufferManager bufferMgr{stream}; - - // Build a real GenerationLogitsCache so that transposedLogits, - // fragmentPointerDevice and fragmentPointerHost are all properly allocated. - // cache.logits uses pinned memory so the test can fill it from the CPU while - // the GPU kernel can still read from it via DMA. - RuntimeBuffers::GenerationLogitsCache cache; - cache.logits = tr::BufferManager::pinnedPool( - tr::ITensor::makeShape({numSteps, maxBatchSize * beamWidth, vocabSize}), nvinfer1::DataType::kFLOAT); - cache.transposedLogits - = bufferMgr.gpu(tr::ITensor::makeShape({beamWidth, numSteps, vocabSize}), nvinfer1::DataType::kFLOAT); - cache.fragmentPointerDevice - = bufferMgr.gpu(tr::ITensor::makeShape({maxBatchSize, numSteps}), nvinfer1::DataType::kINT64); - cache.fragmentPointerHost - = tr::BufferManager::pinnedPool(tr::ITensor::makeShape({maxBatchSize, numSteps}), nvinfer1::DataType::kINT64); - - // Helper: build one LlmRequest that has numSteps fragments pointing into - // cache.logits[0..numSteps-1][logitsIndex:logitsIndex+beamWidth]. - // Each fragment is filled with sentinel value (step*100 + beam + reqOffset). - auto makeRequest = [&](RequestIdType reqId, SizeType32 logitsIndex, float reqOffset) -> std::shared_ptr - { - auto tokens = std::make_shared(promptLen, 0); - tr::SamplingConfig sc{beamWidth}; - auto req = std::make_shared(reqId, numSteps, tokens, sc, false); - - LlmRequest::BeamTokens gen(beamWidth, VecTokens(numSteps, 1)); - req->setGeneratedTokens(gen); - req->allocGenerationLogitsHost(vocabSize, nvinfer1::DataType::kFLOAT); - - // Write known values into the logits cache slots for this request and - // create matching fragment slice views. - for (SizeType32 step = 0; step < numSteps; ++step) - { - // cache.logits shape: [numSteps, maxBatchSize*beamWidth, vocabSize] - // Slice to [1, maxBS*bw, vocab], squeeze to [maxBS*bw, vocab]. - tr::ITensor::SharedPtr slot = tr::ITensor::slice(cache.logits, step, 1); - slot->squeeze(0); // [maxBS*bw, vocab] - auto* slotPtr = tr::bufferCast(*slot); - for (SizeType32 beam = 0; beam < beamWidth; ++beam) - { - float const val = reqOffset + static_cast(step * 100 + beam); - for (SizeType32 v = 0; v < vocabSize; ++v) - { - slotPtr[(logitsIndex + beam) * vocabSize + v] = val; - } - } - - // Fragment matches HandleGenerationLogits: slice [logitsIndex:logitsIndex+beamWidth] - // from the step slot, then unsqueeze(0) → [1, beamWidth, vocab]. - tr::ITensor::SharedPtr fragView = tr::ITensor::slice(slot, logitsIndex, beamWidth); - fragView->unsqueeze(0); // [1, beamWidth, vocab] - req->addGenerationLogitsFragment(fragView); - } - return req; - }; - - // Request 0 occupies logitsIndex=0 in the batch slot. - auto req0 = makeRequest(1, /*logitsIndex=*/0, /*reqOffset=*/0.0f); - // Request 1 occupies logitsIndex=beamWidth in the batch slot. - auto req1 = makeRequest(2, /*logitsIndex=*/beamWidth, /*reqOffset=*/1000.0f); - - // Flush request 0 — uses workIdx=0. - utils::copyGenerationLogits(cache, bufferMgr, *req0, /*beforeDecoder=*/false, {}); - // Flush request 1 — uses workIdx=1 (different slot → no pointer clobbering). - utils::copyGenerationLogits(cache, bufferMgr, *req1, /*beforeDecoder=*/false, {}); - - ASSERT_EQ(cudaStreamSynchronize(stream->get()), cudaSuccess); - - // Verify req0 host buffer: host[beam, step, v] == step*100 + beam - auto const* host0 = tr::bufferCast(*req0->getGenerationLogitsHost()); - for (SizeType32 beam = 0; beam < beamWidth; ++beam) - { - for (SizeType32 step = 0; step < numSteps; ++step) - { - float const expected = static_cast(step * 100 + beam); - for (SizeType32 v = 0; v < vocabSize; ++v) - { - SizeType32 const idx = (beam * numSteps + step) * vocabSize + v; - EXPECT_FLOAT_EQ(host0[idx], expected) << "req0 host[beam=" << beam << ",step=" << step << ",v=" << v - << "]=" << host0[idx] << " expected " << expected; - } - } - } - - // Verify req1 host buffer: host[beam, step, v] == 1000 + step*100 + beam - auto const* host1 = tr::bufferCast(*req1->getGenerationLogitsHost()); - for (SizeType32 beam = 0; beam < beamWidth; ++beam) - { - for (SizeType32 step = 0; step < numSteps; ++step) - { - float const expected = 1000.0f + static_cast(step * 100 + beam); - for (SizeType32 v = 0; v < vocabSize; ++v) - { - SizeType32 const idx = (beam * numSteps + step) * vocabSize + v; - EXPECT_FLOAT_EQ(host1[idx], expected) << "req1 host[beam=" << beam << ",step=" << step << ",v=" << v - << "]=" << host1[idx] << " expected " << expected; - } - } - } - - // Both requests must have had their fragments cleared. - EXPECT_EQ(req0->getGenerationLogitsFragmentsSize(), 0); - EXPECT_EQ(req1->getGenerationLogitsFragmentsSize(), 0); -} diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerFabricMemoryTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerFabricMemoryTest.cpp index d23cbabe4144..4d41d751476c 100644 --- a/cpp/tests/unit_tests/batch_manager/kvCacheManagerFabricMemoryTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerFabricMemoryTest.cpp @@ -20,6 +20,7 @@ #include "tensorrt_llm/batch_manager/kvCacheManager.h" #include "tensorrt_llm/batch_manager/llmRequest.h" #include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/runtime/iTensor.h" #include "tensorrt_llm/runtime/samplingConfig.h" #include "tensorrt_llm/testing/kvCacheManagerTestUtil.h" @@ -108,7 +109,7 @@ TEST_F(KVCacheManagerFabricMemoryTest, AllocatePoolsFallbackWhenFabricUnsupporte BlockManager blockManager(std::vector(numLayers, numKvHeads), sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, stream, maxAttentionWindow, beamWidth, - std::vector{maxAttentionWindow}, nvinfer1::DataType::kHALF, 0, 0); + std::vector{maxAttentionWindow}, tensorrt_llm::DataType::kHALF, 0, 0); blockManager.allocatePools(false); EXPECT_EQ(blockManager.getTokensPerBlock(), tokensPerBlock); @@ -147,7 +148,7 @@ TEST_F(KVCacheManagerFabricMemoryTest, AllocatePoolsWithFabricMemory) BlockManager blockManager(std::vector(numLayers, numKvHeads), sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, stream, maxAttentionWindow, beamWidth, - std::vector{maxAttentionWindow}, nvinfer1::DataType::kHALF, 0, 0); + std::vector{maxAttentionWindow}, tensorrt_llm::DataType::kHALF, 0, 0); blockManager.allocatePools(false); EXPECT_EQ(blockManager.getMaxNumBlocks(), blocksInPrimaryPool); @@ -191,7 +192,7 @@ TEST_F(KVCacheManagerFabricMemoryTest, OffloadOnboardRoundTripWithFabricPrimary) BlockManager blockManager(std::vector(numLayers, numKvHeads), sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, stream, maxAttentionWindow, beamWidth, - std::vector{maxAttentionWindow}, nvinfer1::DataType::kHALF, 0, 0); + std::vector{maxAttentionWindow}, tensorrt_llm::DataType::kHALF, 0, 0); blockManager.allocatePools(false); auto primaryPoolPtr = blockManager.getPrimaryPool(0); @@ -292,7 +293,7 @@ TEST_F(KVCacheManagerFabricMemoryTest, ReleasePoolsClearsFabricMemory) BlockManager blockManager(std::vector(numLayers, numKvHeads), sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, stream, maxAttentionWindow, beamWidth, - std::vector{maxAttentionWindow}, nvinfer1::DataType::kHALF, 0, 0); + std::vector{maxAttentionWindow}, tensorrt_llm::DataType::kHALF, 0, 0); size_t freeBefore = 0; size_t freeAfterAlloc = 0; diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp index 34cf8bc82def..98c2232b4062 100644 --- a/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp @@ -26,6 +26,7 @@ #include "tensorrt_llm/common/assert.h" #include "tensorrt_llm/common/cudaUtils.h" #include "tensorrt_llm/common/memoryUtils.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/executor/transferAgent.h" #include "tensorrt_llm/executor/types.h" #include "tensorrt_llm/kernels/kvCacheIndex.h" @@ -174,7 +175,8 @@ TEST_F(KVCacheManagerTest, BlockManagerTest) BlockManager blockManager(std::vector(numLayers, numKvHeads), sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, stream, maxAttentionWindow, beamWidth, - std::vector{maxAttentionWindow}, nvinfer1::DataType::kHALF, 0, maxAttentionWindow); + std::vector{maxAttentionWindow}, tensorrt_llm::DataType::kHALF, 0, + maxAttentionWindow); blockManager.allocatePools(false); EXPECT_EQ(blockManager.getTokensPerBlock(), tokensPerBlock); @@ -310,7 +312,7 @@ void writePatternToOffloadedBlocksGDS( } } -template +template void runPartialCopyTest() { auto constexpr numLayers = 12; @@ -521,59 +523,59 @@ void runPartialCopyTest() TEST_F(KVCacheManagerTest, BlockManagerTestPartialCopyINT64) { - runPartialCopyTest(); - runPartialCopyTest(); + runPartialCopyTest(); + runPartialCopyTest(); } TEST_F(KVCacheManagerTest, BlockManagerTestPartialCopyINT32) { - runPartialCopyTest(); - runPartialCopyTest(); + runPartialCopyTest(); + runPartialCopyTest(); } TEST_F(KVCacheManagerTest, BlockManagerTestPartialCopyFLOAT) { - runPartialCopyTest(); - runPartialCopyTest(); + runPartialCopyTest(); + runPartialCopyTest(); } #ifdef ENABLE_BF16 TEST_F(KVCacheManagerTest, BlockManagerTestPartialCopyBF16) { - runPartialCopyTest(); - runPartialCopyTest(); + runPartialCopyTest(); + runPartialCopyTest(); } #endif TEST_F(KVCacheManagerTest, BlockManagerTestPartialCopyHALF) { - runPartialCopyTest(); - runPartialCopyTest(); + runPartialCopyTest(); + runPartialCopyTest(); } TEST_F(KVCacheManagerTest, BlockManagerTestPartialCopyBOOL) { - runPartialCopyTest(); - runPartialCopyTest(); + runPartialCopyTest(); + runPartialCopyTest(); } TEST_F(KVCacheManagerTest, BlockManagerTestPartialCopyUINT8) { - runPartialCopyTest(); - runPartialCopyTest(); + runPartialCopyTest(); + runPartialCopyTest(); } TEST_F(KVCacheManagerTest, BlockManagerTestPartialCopyINT8) { - runPartialCopyTest(); - runPartialCopyTest(); + runPartialCopyTest(); + runPartialCopyTest(); } #ifdef ENABLE_FP8 TEST_F(KVCacheManagerTest, BlockManagerTestPartialCopyFP8) { - runPartialCopyTest(); - runPartialCopyTest(); + runPartialCopyTest(); + runPartialCopyTest(); } #endif @@ -731,8 +733,8 @@ TEST_F(KVCacheManagerTest, FindBlocksInReuseTreeByBlockKeysTest) auto const blocksPerWindow = BlocksPerWindow{{maxAttentionWindow, {blocksInPrimaryPool, blocksInSecondaryPool}}}; KVCacheManager kvCacheManager(numLayers, numKvHeads, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, - beamWidth, std::vector{maxAttentionWindow}, nvinfer1::DataType::kHALF, false, stream, - maxAttentionWindow, maxAttentionWindow, true); + beamWidth, std::vector{maxAttentionWindow}, tensorrt_llm::DataType::kHALF, false, + stream, maxAttentionWindow, maxAttentionWindow, true); // Add sequence [0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16] (17 tokens, three blocks) auto inputTokens = std::make_shared(VecTokens{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}); @@ -783,8 +785,8 @@ TEST_F(KVCacheManagerTest, FP4BlockScaleManagementTest) auto const blocksPerWindow = BlocksPerWindow{{maxAttentionWindow, {blocksInPrimaryPool, blocksInSecondaryPool}}}; KVCacheManager kvCacheManager(numLayers, numHeads, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, - beamWidth, std::vector{maxAttentionWindow}, nvinfer1::DataType::kFP4, false, stream, - maxAttentionWindow, maxAttentionWindow, true); + beamWidth, std::vector{maxAttentionWindow}, tensorrt_llm::DataType::kFP4, false, + stream, maxAttentionWindow, maxAttentionWindow, true); kvCacheManager.allocatePools(/*useUvm=*/false); @@ -823,14 +825,14 @@ TEST_F(KVCacheManagerTest, FP4AttentionWithHalfRecurrentStatesPoolTest) {maxAttentionWindow, {blocksInPrimaryPool, blocksInSecondaryPool}}, }; auto const poolConfigurations = std::vector{ - {recurrentStatesWindow, sizePerHead, nvinfer1::DataType::kHALF}, - {maxAttentionWindow, sizePerHead, nvinfer1::DataType::kFP4}, + {recurrentStatesWindow, sizePerHead, tensorrt_llm::DataType::kHALF}, + {maxAttentionWindow, sizePerHead, tensorrt_llm::DataType::kFP4}, }; auto const stream = std::make_shared(); KVCacheManager kvCacheManager(std::vector{0, numKvHeads}, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, maxBeamWidth, std::vector{recurrentStatesWindow, maxAttentionWindow}, - nvinfer1::DataType::kFP4, + tensorrt_llm::DataType::kFP4, /*sinkTokenLength=*/0, stream, maxAttentionWindow, /*chunkSize=*/0, /*enableBlockReuse=*/false, CacheType::kSELF, std::nullopt, nullptr, /*enablePartialReuse=*/false, /*copyOnPartialReuse=*/true, nullptr, /*enableIndexerKCache=*/false, /*indexerKCacheQuantBlockSize=*/128, /*indexerKCacheIndexHeadDim=*/0, @@ -838,13 +840,13 @@ TEST_F(KVCacheManagerTest, FP4AttentionWithHalfRecurrentStatesPoolTest) kvCacheManager.allocatePools(/*useUvm=*/false); auto const& blockManager = kvCacheManager.getBlockManager(); - EXPECT_EQ(blockManager.getDataTypeForWindow(recurrentStatesWindow), nvinfer1::DataType::kHALF); - EXPECT_EQ(blockManager.getDataTypeForWindow(maxAttentionWindow), nvinfer1::DataType::kFP4); + EXPECT_EQ(blockManager.getDataTypeForWindow(recurrentStatesWindow), tensorrt_llm::DataType::kHALF); + EXPECT_EQ(blockManager.getDataTypeForWindow(maxAttentionWindow), tensorrt_llm::DataType::kFP4); auto const& recurrentStatesPool = blockManager.getRecurrentStatesPool(); ASSERT_NE(recurrentStatesPool.primaryPtr, nullptr); - EXPECT_EQ(recurrentStatesPool.primaryPtr->getDataType(), nvinfer1::DataType::kHALF); - auto const recurrentStatesElementsPerBlock = recurrentStatesBytes / tc::getDTypeSize(nvinfer1::DataType::kHALF); + EXPECT_EQ(recurrentStatesPool.primaryPtr->getDataType(), tensorrt_llm::DataType::kHALF); + auto const recurrentStatesElementsPerBlock = recurrentStatesBytes / tc::getDTypeSize(tensorrt_llm::DataType::kHALF); EXPECT_EQ(recurrentStatesPool.blockSize, recurrentStatesElementsPerBlock); SizeType32 numRecurrentScalePools = 0; @@ -905,7 +907,8 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseTest) BlockManager blockManager(std::vector(numLayers, numKvHeads), sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, stream, maxAttentionWindow, beamWidth, - std::vector{maxAttentionWindow}, nvinfer1::DataType::kHALF, 0, maxAttentionWindow); + std::vector{maxAttentionWindow}, tensorrt_llm::DataType::kHALF, 0, + maxAttentionWindow); blockManager.allocatePools(false); EXPECT_EQ(blockManager.getTokensPerBlock(), tokensPerBlock); @@ -1240,7 +1243,8 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithExtraIdTest) BlockManager blockManager(std::vector(numLayers, numKvHeads), sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, stream, maxAttentionWindow, beamWidth, - std::vector{maxAttentionWindow}, nvinfer1::DataType::kHALF, 0, maxAttentionWindow); + std::vector{maxAttentionWindow}, tensorrt_llm::DataType::kHALF, 0, + maxAttentionWindow); blockManager.allocatePools(false); EXPECT_EQ(blockManager.getTokensPerBlock(), tokensPerBlock); @@ -1475,7 +1479,8 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithMultimodalHashTest) BlockManager blockManager(std::vector(numLayers, numKvHeads), sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, stream, maxAttentionWindow, beamWidth, - std::vector{maxAttentionWindow}, nvinfer1::DataType::kHALF, 0, maxAttentionWindow); + std::vector{maxAttentionWindow}, tensorrt_llm::DataType::kHALF, 0, + maxAttentionWindow); blockManager.allocatePools(false); EXPECT_EQ(blockManager.getTokensPerBlock(), tokensPerBlock); @@ -1679,7 +1684,8 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithLoraTaskIdTest) BlockManager blockManager(std::vector(numLayers, numKvHeads), sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, stream, maxAttentionWindow, beamWidth, - std::vector{maxAttentionWindow}, nvinfer1::DataType::kHALF, 0, maxAttentionWindow); + std::vector{maxAttentionWindow}, tensorrt_llm::DataType::kHALF, 0, + maxAttentionWindow); blockManager.allocatePools(false); EXPECT_EQ(blockManager.getTokensPerBlock(), tokensPerBlock); @@ -1972,7 +1978,8 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithExtraIdAndLoraTaskIdTest) BlockManager blockManager(std::vector(numLayers, numKvHeads), sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, stream, maxAttentionWindow, beamWidth, - std::vector{maxAttentionWindow}, nvinfer1::DataType::kHALF, 0, maxAttentionWindow); + std::vector{maxAttentionWindow}, tensorrt_llm::DataType::kHALF, 0, + maxAttentionWindow); blockManager.allocatePools(false); EXPECT_EQ(blockManager.getTokensPerBlock(), tokensPerBlock); @@ -2241,7 +2248,8 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithCacheSaltTest) BlockManager blockManager(std::vector(numLayers, numKvHeads), sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, stream, maxAttentionWindow, beamWidth, - std::vector{maxAttentionWindow}, nvinfer1::DataType::kHALF, 0, maxAttentionWindow); + std::vector{maxAttentionWindow}, tensorrt_llm::DataType::kHALF, 0, + maxAttentionWindow); blockManager.allocatePools(false); EXPECT_EQ(blockManager.getTokensPerBlock(), tokensPerBlock); @@ -2473,7 +2481,7 @@ TEST_F(KVCacheManagerTest, KVCacheManagerPerRequestStatsTest) auto const blocksPerWindow = BlocksPerWindow{{maxAttentionWindow, {blocksInPrimaryPool, blocksInSecondaryPool}}}; KVCacheManager kvCacheManager(numLayers, numHeads, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, - beamWidth, std::vector{maxAttentionWindow}, nvinfer1::DataType::kHALF, 0, stream, + beamWidth, std::vector{maxAttentionWindow}, tensorrt_llm::DataType::kHALF, 0, stream, maxSequenceLength, maxSequenceLength, true); kvCacheManager.allocatePools(false); @@ -2530,7 +2538,8 @@ TEST_F(KVCacheManagerTest, BlockManagerBlockPriorityTest) BlockManager blockManager(std::vector(numLayers, numKvHeads), sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, stream, maxAttentionWindow, beamWidth, - std::vector{maxAttentionWindow}, nvinfer1::DataType::kHALF, 0, maxAttentionWindow); + std::vector{maxAttentionWindow}, tensorrt_llm::DataType::kHALF, 0, + maxAttentionWindow); blockManager.allocatePools(false); EXPECT_EQ(blockManager.getTokensPerBlock(), tokensPerBlock); @@ -2668,7 +2677,7 @@ TEST_F(KVCacheManagerTest, KVCacheManagerDecodeBlockPriorityTest) auto const blocksPerWindow = BlocksPerWindow{{maxAttentionWindow, {blocksInPrimaryPool, blocksInSecondaryPool}}}; KVCacheManager kvCacheManager(numLayers, numHeads, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, - beamWidth, std::vector{maxAttentionWindow}, nvinfer1::DataType::kHALF, 0, stream, + beamWidth, std::vector{maxAttentionWindow}, tensorrt_llm::DataType::kHALF, 0, stream, maxSequenceLength, maxSequenceLength, true); kvCacheManager.allocatePools(false); @@ -2775,7 +2784,7 @@ TEST_F(KVCacheManagerTest, KVCacheManagerTimedEvictionTest) auto const blocksPerWindow = BlocksPerWindow{{maxAttentionWindow, {blocksInPrimaryPool, blocksInSecondaryPool}}}; KVCacheManager kvCacheManager(numLayers, numHeads, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, - beamWidth, std::vector{maxAttentionWindow}, nvinfer1::DataType::kHALF, 0, stream, + beamWidth, std::vector{maxAttentionWindow}, tensorrt_llm::DataType::kHALF, 0, stream, maxSequenceLength, maxSequenceLength, true); kvCacheManager.allocatePools(false); @@ -2847,7 +2856,7 @@ TEST_F(KVCacheManagerTest, KVCacheManagerDecodeTimedEvictionTest) auto const blocksPerWindow = BlocksPerWindow{{maxAttentionWindow, {blocksInPrimaryPool, blocksInSecondaryPool}}}; KVCacheManager kvCacheManager(numLayers, numHeads, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, - beamWidth, std::vector{maxAttentionWindow}, nvinfer1::DataType::kHALF, 0, stream, + beamWidth, std::vector{maxAttentionWindow}, tensorrt_llm::DataType::kHALF, 0, stream, maxSequenceLength, maxSequenceLength, true); kvCacheManager.allocatePools(false); { @@ -2940,7 +2949,7 @@ TEST_F(KVCacheManagerTest, KVCacheManagerSecondaryBlockPrimaryChildTest) auto const blocksPerWindow = BlocksPerWindow{{maxAttentionWindow, {blocksInPrimaryPool, blocksInSecondaryPool}}}; KVCacheManager kvCacheManager(numLayers, numHeads, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, - beamWidth, std::vector{maxAttentionWindow}, nvinfer1::DataType::kHALF, 0, stream, + beamWidth, std::vector{maxAttentionWindow}, tensorrt_llm::DataType::kHALF, 0, stream, maxSequenceLength, maxSequenceLength, true); kvCacheManager.allocatePools(false); @@ -3037,7 +3046,7 @@ TEST_F(KVCacheManagerTest, KVCacheManagerStoreContextBlocksUsesMaterializedConte auto const blocksPerWindow = BlocksPerWindow{{maxAttentionWindow, {blocksInPrimaryPool, blocksInSecondaryPool}}}; KVCacheManager kvCacheManager(numLayers, numHeads, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, - beamWidth, std::vector{maxAttentionWindow}, nvinfer1::DataType::kHALF, 0, stream, + beamWidth, std::vector{maxAttentionWindow}, tensorrt_llm::DataType::kHALF, 0, stream, maxSequenceLength, /*chunkSize*/ 0, true); kvCacheManager.allocatePools(false); @@ -3081,7 +3090,7 @@ TEST_F(KVCacheManagerTest, KVCacheManagerReleaseBlocksUsesMaterializedContextExt auto const blocksPerWindow = BlocksPerWindow{{maxAttentionWindow, {blocksInPrimaryPool, blocksInSecondaryPool}}}; KVCacheManager kvCacheManager(numLayers, numHeads, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, - beamWidth, std::vector{maxAttentionWindow}, nvinfer1::DataType::kHALF, 0, stream, + beamWidth, std::vector{maxAttentionWindow}, tensorrt_llm::DataType::kHALF, 0, stream, maxSequenceLength, /*chunkSize*/ 0, true); kvCacheManager.allocatePools(false); @@ -3123,7 +3132,7 @@ TEST_F(KVCacheManagerTest, KVCacheManagerLeafBlockTest) auto const blocksPerWindow = BlocksPerWindow{{maxAttentionWindow, {blocksInPrimaryPool, blocksInSecondaryPool}}}; KVCacheManager kvCacheManager(numLayers, numHeads, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, - beamWidth, std::vector{maxAttentionWindow}, nvinfer1::DataType::kHALF, 0, stream, + beamWidth, std::vector{maxAttentionWindow}, tensorrt_llm::DataType::kHALF, 0, stream, maxSequenceLength, maxSequenceLength, true); kvCacheManager.allocatePools(false); @@ -3208,7 +3217,7 @@ TEST_F(KVCacheManagerTest, KVCacheManagerLeafBlockWithDependentTest) auto const blocksPerWindow = BlocksPerWindow{{maxAttentionWindow, {blocksInPrimaryPool, blocksInSecondaryPool}}}; KVCacheManager kvCacheManager(numLayers, numHeads, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, - beamWidth, std::vector{maxAttentionWindow}, nvinfer1::DataType::kHALF, 0, stream, + beamWidth, std::vector{maxAttentionWindow}, tensorrt_llm::DataType::kHALF, 0, stream, maxSequenceLength, maxSequenceLength, true); kvCacheManager.allocatePools(false); @@ -3312,7 +3321,7 @@ TEST_P(KVCacheManagerTest, DISABLED_KVCacheManagerAllocationTest) auto constexpr maxNumSequences = 8; auto constexpr maxBeamWidth = 4; auto constexpr sinkTokenLength = 0; - auto constexpr dtype = nvinfer1::DataType::kHALF; + auto constexpr dtype = tensorrt_llm::DataType::kHALF; auto const stream = std::make_shared(); auto constexpr maxSequenceLength = tokensPerBlock * maxBlocksPerSeq; @@ -3335,11 +3344,12 @@ TEST_P(KVCacheManagerTest, DISABLED_KVCacheManagerAllocationTest) KVCacheManager kvCacheManager = homogeneousLayers ? KVCacheManager(numLayers, numHeads, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, - maxBeamWidth, std::vector{maxAttentionWindow}, nvinfer1::DataType::kHALF, + maxBeamWidth, std::vector{maxAttentionWindow}, tensorrt_llm::DataType::kHALF, sinkTokenLength, stream, maxSequenceLength, maxSequenceLength, enableBlockReuse) : KVCacheManager(std::vector(numLayers, numHeads), sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, maxBeamWidth, std::vector{maxAttentionWindow}, - nvinfer1::DataType::kHALF, sinkTokenLength, stream, maxSequenceLength, maxSequenceLength, enableBlockReuse); + tensorrt_llm::DataType::kHALF, sinkTokenLength, stream, maxSequenceLength, maxSequenceLength, + enableBlockReuse); auto const& blockManager = kvCacheManager.getBlockManager(); auto const& bufferManager = blockManager.getBufferManager(theOnlyWindowSize(kvCacheManager)); @@ -3410,10 +3420,10 @@ TEST_P(KVCacheManagerTest, KVCacheManagerTest) KVCacheManager kvCacheManager = homogeneousLayers ? KVCacheManager(numLayers, numHeads, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, - maxBeamWidth, std::vector{maxAttentionWindow}, nvinfer1::DataType::kHALF, + maxBeamWidth, std::vector{maxAttentionWindow}, tensorrt_llm::DataType::kHALF, sinkTokenLength, stream, maxSequenceLength, maxSequenceLength, enableBlockReuse) : KVCacheManager(numHeadsPerLayer, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, maxBeamWidth, - std::vector{maxAttentionWindow}, nvinfer1::DataType::kHALF, sinkTokenLength, + std::vector{maxAttentionWindow}, tensorrt_llm::DataType::kHALF, sinkTokenLength, stream, maxSequenceLength, maxSequenceLength, enableBlockReuse); kvCacheManager.allocatePools(false); @@ -3580,11 +3590,12 @@ TEST_P(KVCacheManagerTest, KVCacheManagerRewindTokensTest) KVCacheManager kvCacheManager = homogeneousLayers ? KVCacheManager(numLayers, numHeads, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, - maxBeamWidth, std::vector{maxAttentionWindow}, nvinfer1::DataType::kHALF, + maxBeamWidth, std::vector{maxAttentionWindow}, tensorrt_llm::DataType::kHALF, sinkTokenLength, stream, maxSequenceLength, maxSequenceLength, enableBlockReuse) : KVCacheManager(std::vector(numLayers, numHeads), sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, maxBeamWidth, std::vector{maxAttentionWindow}, - nvinfer1::DataType::kHALF, sinkTokenLength, stream, maxSequenceLength, maxSequenceLength, enableBlockReuse); + tensorrt_llm::DataType::kHALF, sinkTokenLength, stream, maxSequenceLength, maxSequenceLength, + enableBlockReuse); kvCacheManager.allocatePools(false); EXPECT_EQ(kvCacheManager.getTokensPerBlock(), tokensPerBlock); @@ -3688,10 +3699,10 @@ TEST_P(KVCacheManagerTest, KVCacheManagerMaxAttentionWindowTest) KVCacheManager kvCacheManager = homogeneousLayers ? KVCacheManager(numLayers, numHeads, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, - maxBeamWidth, std::vector{maxAttentionWindow}, nvinfer1::DataType::kHALF, + maxBeamWidth, std::vector{maxAttentionWindow}, tensorrt_llm::DataType::kHALF, sinkTokenLength, stream, maxSequenceLength, maxSequenceLength, enableBlockReuse) : KVCacheManager(numHeadsPerLayer, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, maxBeamWidth, - std::vector{maxAttentionWindow}, nvinfer1::DataType::kHALF, sinkTokenLength, + std::vector{maxAttentionWindow}, tensorrt_llm::DataType::kHALF, sinkTokenLength, stream, maxSequenceLength, maxSequenceLength, enableBlockReuse); kvCacheManager.allocatePools(false); @@ -3812,7 +3823,7 @@ TEST_F(KVCacheManagerTest, KVCacheManagerMaxAttentionWindowSmallerThanBlockSizeT auto const blocksPerWindow = BlocksPerWindow{{maxAttentionWindow, {blocksInPrimaryPool, blocksInSecondaryPool}}}; KVCacheManager kvCacheManager(numLayers, numHeads, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, - maxBeamWidth, std::vector{maxAttentionWindow}, nvinfer1::DataType::kHALF, + maxBeamWidth, std::vector{maxAttentionWindow}, tensorrt_llm::DataType::kHALF, sinkTokenLength, stream, maxSequenceLength, maxSequenceLength, enableBlockReuse); kvCacheManager.allocatePools(false); @@ -3906,7 +3917,7 @@ TEST_F(KVCacheManagerTest, KVCacheManagerEventStream) auto constexpr blocksInPrimaryPool = 8; auto constexpr blocksInSecondaryPool = 2; - auto constexpr dtype = nvinfer1::DataType::kHALF; + auto constexpr dtype = tensorrt_llm::DataType::kHALF; auto const stream = std::make_shared(); auto constexpr beamWidth = 1; @@ -4090,7 +4101,7 @@ TEST_F(KVCacheManagerTest, KVCacheManagerMaxAttentionWindowWithReuseTest) auto const blocksPerWindow = BlocksPerWindow{{maxAttentionWindow, {blocksInPrimaryPool, blocksInSecondaryPool}}}; KVCacheManager kvCacheManager(numLayers, numHeads, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, - maxBeamWidth, std::vector{maxAttentionWindow}, nvinfer1::DataType::kHALF, + maxBeamWidth, std::vector{maxAttentionWindow}, tensorrt_llm::DataType::kHALF, sinkTokenLength, stream, maxSequenceLength, /*chunkSize=*/tokensPerBlock, enableBlockReuse); kvCacheManager.allocatePools(false); @@ -4217,7 +4228,7 @@ TEST_F(KVCacheManagerTest, KVCacheManagerSWAInvalidateReuseTest) auto constexpr maxNumSequences = 8; auto constexpr maxBeamWidth = 1; auto constexpr sinkTokenLength = 0; - auto constexpr dtype = nvinfer1::DataType::kHALF; + auto constexpr dtype = tensorrt_llm::DataType::kHALF; auto const stream = std::make_shared(); auto constexpr maxSequenceLength = 128; SizeType32 constexpr maxNewTokens = 40; @@ -4300,7 +4311,7 @@ TEST_F(KVCacheManagerTest, KVCacheManagerVariableWindowAttentionWithReuseTest) auto constexpr maxNumSequences = 8; auto constexpr maxBeamWidth = 1; auto constexpr sinkTokenLength = 0; - auto constexpr dtype = nvinfer1::DataType::kHALF; + auto constexpr dtype = tensorrt_llm::DataType::kHALF; auto const stream = std::make_shared(); auto constexpr maxSequenceLength = 128; @@ -4426,7 +4437,7 @@ TEST_F(KVCacheManagerTest, KVCacheManagerEventStreamOverflow) auto constexpr blocksInPrimaryPool = 8; auto constexpr blocksInSecondaryPool = 2; - auto constexpr dtype = nvinfer1::DataType::kHALF; + auto constexpr dtype = tensorrt_llm::DataType::kHALF; auto const stream = std::make_shared(); auto constexpr beamWidth = 1; @@ -4486,7 +4497,7 @@ TEST_F(KVCacheManagerTest, KVCacheManagerEventStreamPriority) auto constexpr blocksInPrimaryPool = 8; auto constexpr blocksInSecondaryPool = 2; - auto constexpr dtype = nvinfer1::DataType::kHALF; + auto constexpr dtype = tensorrt_llm::DataType::kHALF; auto const stream = std::make_shared(); auto constexpr beamWidth = 1; @@ -4563,7 +4574,7 @@ TEST_F(KVCacheManagerTest, GetPriorityByBlockId) auto constexpr maxAttentionWindow = 32; auto constexpr maxNumSequences = 4; auto constexpr beamWidth = 1; - auto constexpr dtype = nvinfer1::DataType::kHALF; + auto constexpr dtype = tensorrt_llm::DataType::kHALF; auto const stream = std::make_shared(); SizeType32 constexpr maxNewTokens = 4; tr::SamplingConfig const samplingConfig{beamWidth}; @@ -4629,7 +4640,7 @@ TEST_F(KVCacheManagerTest, CommitAndGetBlockHashesForRequest) auto constexpr maxNumSequences = 4; auto constexpr beamWidth = 1; auto constexpr beamIdx = 0; - auto constexpr dtype = nvinfer1::DataType::kHALF; + auto constexpr dtype = tensorrt_llm::DataType::kHALF; auto const stream = std::make_shared(); SizeType32 constexpr maxNewTokens = 8; tr::SamplingConfig const samplingConfig{beamWidth}; @@ -4755,7 +4766,7 @@ TEST_F(KVCacheManagerTest, CommitAndGetBlockHashesFrontRunsTrailingFullBlock) auto constexpr maxNumSequences = 4; auto constexpr beamWidth = 1; auto constexpr beamIdx = 0; - auto constexpr dtype = nvinfer1::DataType::kHALF; + auto constexpr dtype = tensorrt_llm::DataType::kHALF; auto const stream = std::make_shared(); SizeType32 constexpr maxNewTokens = 8; tr::SamplingConfig const samplingConfig{beamWidth}; @@ -4871,7 +4882,7 @@ TEST_F(KVCacheManagerTest, PinAndUnpinBlocksById) BlocksPerWindow const blocksPerWindow{{maxAttentionWindow, {blocksInPrimaryPool, blocksInSecondaryPool}}}; KVCacheManager kvCacheManager(numLayers, numKvHeads, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, - beamWidth, std::vector{maxAttentionWindow}, nvinfer1::DataType::kHALF, 0, stream, + beamWidth, std::vector{maxAttentionWindow}, tensorrt_llm::DataType::kHALF, 0, stream, maxAttentionWindow, maxAttentionWindow, true); kvCacheManager.allocatePools(false); @@ -4923,7 +4934,7 @@ TEST_F(KVCacheManagerTest, StoreBlocksForReuseWithPinDoesNotCreateGhostFreeBlock BlocksPerWindow const blocksPerWindow{{maxAttentionWindow, {blocksInPrimaryPool, blocksInSecondaryPool}}}; KVCacheManager kvCacheManager(numLayers, numKvHeads, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, - beamWidth, std::vector{maxAttentionWindow}, nvinfer1::DataType::kHALF, 0, stream, + beamWidth, std::vector{maxAttentionWindow}, tensorrt_llm::DataType::kHALF, 0, stream, maxAttentionWindow, maxAttentionWindow, true /* enableBlockReuse */); kvCacheManager.allocatePools(false); @@ -4995,7 +5006,7 @@ TEST_F(KVCacheManagerTest, KVCacheManagerEventStreamBlocking) auto constexpr blocksInPrimaryPool = 8; auto constexpr blocksInSecondaryPool = 2; - auto constexpr dtype = nvinfer1::DataType::kHALF; + auto constexpr dtype = tensorrt_llm::DataType::kHALF; auto const stream = std::make_shared(); auto constexpr beamWidth = 1; @@ -5015,7 +5026,7 @@ TEST_F(KVCacheManagerTest, KVCacheManagerEventStreamBlocking) EXPECT_EQ(getEvents(kvCacheManagerTest).size(), 0); KVCacheManager kvCacheManager(numLayers, numHeads, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, - beamWidth, std::vector{maxAttentionWindow}, nvinfer1::DataType::kHALF, 0, stream, + beamWidth, std::vector{maxAttentionWindow}, tensorrt_llm::DataType::kHALF, 0, stream, maxSequenceLength, maxSequenceLength, true, CacheType::kSELF, std::nullopt, std::make_unique(1024)); @@ -5050,7 +5061,7 @@ TEST_F(KVCacheManagerTest, KVCacheManagerEventStreamWindowSize) auto blocksInPool = std::vector{8, 2}; auto blocksInSlidingWindowPool = std::vector{4, 2}; - auto constexpr dtype = nvinfer1::DataType::kHALF; + auto constexpr dtype = tensorrt_llm::DataType::kHALF; auto const stream = std::make_shared(); auto constexpr beamWidth = 1; @@ -5109,10 +5120,11 @@ TEST_F(KVCacheManagerTest, KVCacheTransferManagerConcurrencyTest) auto pool = KVCacheBlockPool(0, 2, 0, 0, 0); - pool.primaryPtr = bufferManager.gpu(tr::ITensor::makeShape({1, blockSize}), nvinfer1::DataType::kFLOAT); + pool.primaryPtr = bufferManager.gpu(tr::ITensor::makeShape({1, blockSize}), tensorrt_llm::DataType::kFLOAT); bufferManager.setZero(*pool.primaryPtr); - pool.secondaryPtr = tr::BufferManager::pinned(tr::ITensor::makeShape({1, blockSize}), nvinfer1::DataType::kFLOAT); + pool.secondaryPtr + = tr::BufferManager::pinned(tr::ITensor::makeShape({1, blockSize}), tensorrt_llm::DataType::kFLOAT); // Write some specific data into the cpu blocks. for (int i = 0; i < blockSize; i++) @@ -5149,11 +5161,11 @@ TEST_F(KVCacheManagerTest, KVCacheTransferManagerPendingTransfersDistinguishPrim auto pool = KVCacheBlockPool(0, 2, 0, 0, 0); pool.primaryPtr - = bufferManager.gpu(tr::ITensor::makeShape({kNumSlotsPerPool, kBlockSize}), nvinfer1::DataType::kFLOAT); + = bufferManager.gpu(tr::ITensor::makeShape({kNumSlotsPerPool, kBlockSize}), tensorrt_llm::DataType::kFLOAT); bufferManager.setZero(*pool.primaryPtr); - pool.secondaryPtr - = tr::BufferManager::pinned(tr::ITensor::makeShape({kNumSlotsPerPool, kBlockSize}), nvinfer1::DataType::kFLOAT); + pool.secondaryPtr = tr::BufferManager::pinned( + tr::ITensor::makeShape({kNumSlotsPerPool, kBlockSize}), tensorrt_llm::DataType::kFLOAT); auto primarySlot0 = std::make_shared(0, tk::KVCacheIndex(0, false)); auto primarySlot1 = std::make_shared(1, tk::KVCacheIndex(1, false)); @@ -5240,10 +5252,10 @@ TEST_P(KVCacheManagerTest, DISABLED_KVCacheManagerSinkTokenLengthTest) auto const maxSequenceLength = tokensPerBlock * maxBlocksPerSeq; KVCacheManager kvCacheManager = homogeneousLayers ? KVCacheManager(numLayers, numHeads, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, - maxBeamWidth, std::vector{maxAttentionWindow}, nvinfer1::DataType::kHALF, + maxBeamWidth, std::vector{maxAttentionWindow}, tensorrt_llm::DataType::kHALF, sinkTokenLength, stream, maxSequenceLength, maxSequenceLength, enableBlockReuse) : KVCacheManager(numHeadsPerLayer, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, maxBeamWidth, - std::vector{maxAttentionWindow}, nvinfer1::DataType::kHALF, sinkTokenLength, + std::vector{maxAttentionWindow}, tensorrt_llm::DataType::kHALF, sinkTokenLength, stream, maxSequenceLength, maxSequenceLength, enableBlockReuse); kvCacheManager.allocatePools(false); @@ -5402,10 +5414,10 @@ TEST_P(KVCacheManagerTest, KVCacheManagerBatchTest) KVCacheManager kvCacheManager = homogeneousLayers ? KVCacheManager(numLayers, numHeads, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, - maxBeamWidth, std::vector{maxAttentionWindow}, nvinfer1::DataType::kHALF, + maxBeamWidth, std::vector{maxAttentionWindow}, tensorrt_llm::DataType::kHALF, sinkTokenLength, stream, maxSequenceLength, maxSequenceLength, enableBlockReuse) : KVCacheManager(numHeadsPerLayer, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, maxBeamWidth, - std::vector{maxAttentionWindow}, nvinfer1::DataType::kHALF, sinkTokenLength, + std::vector{maxAttentionWindow}, tensorrt_llm::DataType::kHALF, sinkTokenLength, stream, maxSequenceLength, maxSequenceLength, enableBlockReuse); kvCacheManager.allocatePools(false); @@ -5542,12 +5554,12 @@ void testNeededBlocksOneStep(bool kv_cache_block_reuse, int beamWidth, int draft KVCacheManager kvCacheManager = homogeneousLayers ? KVCacheManager(numLayers, numHeads, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, - maxBeamWidth, std::vector{maxAttentionWindow}, nvinfer1::DataType::kHALF, - sinkTokenLength, stream, maxSequenceLength, + maxBeamWidth, std::vector{maxAttentionWindow}, + tensorrt_llm::DataType::kHALF, sinkTokenLength, stream, maxSequenceLength, /*chunkSize=*/tokensPerBlock, kv_cache_block_reuse) : KVCacheManager(numHeadsPerLayer, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, - maxBeamWidth, std::vector{maxAttentionWindow}, nvinfer1::DataType::kHALF, - sinkTokenLength, stream, maxSequenceLength, + maxBeamWidth, std::vector{maxAttentionWindow}, + tensorrt_llm::DataType::kHALF, sinkTokenLength, stream, maxSequenceLength, /*chunkSize=*/tokensPerBlock, kv_cache_block_reuse); kvCacheManager.allocatePools(false); @@ -5750,7 +5762,7 @@ struct KvCacheManagerInstantiationParameters SizeType32 maxNumTokens; bool kvCacheBlockReuse; std::vector maxAttentionWindowVec = {maxAttentionWindow}; - nvinfer1::DataType dtype = nvinfer1::DataType::kFLOAT; + tensorrt_llm::DataType dtype = tensorrt_llm::DataType::kFLOAT; }; BlocksPerWindow blocksAndWindow(SizeType32 numPrimaryBlocks, SizeType32 windowSize) @@ -7246,7 +7258,7 @@ TEST_F(KVCacheManagerTest, KVCacheManagerEventRemovedBatchedWithinWindow) auto constexpr maxNumSequences = 4; auto constexpr maxAttentionWindow = 32; auto constexpr beamWidth = 1; - auto constexpr dtype = nvinfer1::DataType::kHALF; + auto constexpr dtype = tensorrt_llm::DataType::kHALF; auto const stream = std::make_shared(); SizeType32 constexpr maxNewTokens{0}; tr::SamplingConfig const samplingConfig{beamWidth}; @@ -7324,7 +7336,7 @@ TEST_F(KVCacheManagerTest, KVCacheManagerEventRemovedOrderedBeforeStore) auto constexpr maxNumSequences = 4; auto constexpr maxAttentionWindow = 32; auto constexpr beamWidth = 1; - auto constexpr dtype = nvinfer1::DataType::kHALF; + auto constexpr dtype = tensorrt_llm::DataType::kHALF; auto const stream = std::make_shared(); SizeType32 constexpr maxNewTokens{0}; tr::SamplingConfig const samplingConfig{beamWidth}; @@ -7418,7 +7430,7 @@ TEST_F(KVCacheManagerTest, KVCacheManagerEventStoreForDifferentWindowDoesNotFlus auto constexpr blocksInSecondaryPool = 0; auto constexpr maxNumSequences = 4; auto constexpr beamWidth = 1; - auto constexpr dtype = nvinfer1::DataType::kHALF; + auto constexpr dtype = tensorrt_llm::DataType::kHALF; auto const stream = std::make_shared(); SizeType32 constexpr maxNewTokens{0}; tr::SamplingConfig const samplingConfig{beamWidth}; @@ -7531,7 +7543,8 @@ void testBlockManagerLinearAttention_ContextNoReuse(int beamWidth, int numTokens BlockManager blockManager(std::vector(numLayers, numKvHeads), sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, stream, maxAttentionWindow, beamWidth, - std::vector{linearWindowSizeCode, maxAttentionWindow}, nvinfer1::DataType::kHALF, 0, + std::vector{linearWindowSizeCode, maxAttentionWindow}, tensorrt_llm::DataType::kHALF, + 0, /*chunkSize*/ 0, CacheType::kSELF, std::nullopt, nullptr, false, true, nullptr, std::nullopt, false, 128, 0, false, linearAttentionMetadata); blockManager.allocatePools(false); @@ -7676,7 +7689,8 @@ void testBlockManagerLinearAttention_ContextReuse(int beamWidth, int numTokens0, BlockManager blockManager(std::vector(numLayers, numKvHeads), sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, stream, maxAttentionWindow, beamWidth, - std::vector{linearWindowSizeCode, maxAttentionWindow}, nvinfer1::DataType::kHALF, 0, + std::vector{linearWindowSizeCode, maxAttentionWindow}, tensorrt_llm::DataType::kHALF, + 0, /*chunkSize*/ 0, CacheType::kSELF, std::nullopt, nullptr, false, true, nullptr, std::nullopt, false, 128, 0, false, linearAttentionMetadata); blockManager.allocatePools(false); @@ -7900,7 +7914,7 @@ void testKVCacheManagerLinearAttention_DecodingBlockGrowth( {linearWindowSizeCode, {blocksInPrimaryPool, blocksInSecondaryPool}}}; KVCacheManager kvCacheManager(numLayers, numKvHeads, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, beamWidth, std::vector{linearWindowSizeCode}, - /*dtype*/ nvinfer1::DataType::kHALF, + /*dtype*/ tensorrt_llm::DataType::kHALF, /*sinkTokenLen*/ sinkTokenLen, /*stream*/ stream, /*maxSequenceLength*/ maxAttentionWindow, @@ -8012,7 +8026,7 @@ void testKVCacheManagerLinearAttention_BlockCopying( {linearWindowSizeCode, {blocksInPrimaryPool, blocksInSecondaryPool}}}; KVCacheManager kvCacheManager(numLayers, numKvHeads, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, beamWidth, std::vector{linearWindowSizeCode, maxAttentionWindow}, - nvinfer1::DataType::kHALF, sinkTokenLen, stream, maxAttentionWindow, /*chunkSize*/ 0, enableContextReuse, + tensorrt_llm::DataType::kHALF, sinkTokenLen, stream, maxAttentionWindow, /*chunkSize*/ 0, enableContextReuse, CacheType::kSELF, std::nullopt, nullptr, false, true, nullptr, false, 128, 0, false, linearAttentionMetadata); kvCacheManager.allocatePools(false); @@ -8338,7 +8352,7 @@ TEST_F(KVCacheManagerTest, StaticLinearHybridAllocationTest) // Static-hybrid path requires block reuse to be disabled. tle::KvCacheConfig const kvCacheConfigDisabledReuse{/*enableBlockReuse=*/false}; auto const blocksPerWindow - = KVCacheManager::calculateMaxNumBlocks(kvCacheConfigDisabledReuse, nvinfer1::DataType::kHALF, + = KVCacheManager::calculateMaxNumBlocks(kvCacheConfigDisabledReuse, tensorrt_llm::DataType::kHALF, numKvHeadsPerLayer, sizePerHead, tokensPerBlock, worldConfig, windowSizeToLayers, allottedPrimaryMemBytes, allottedSecondaryMemBytes, extraCostMemory, kvFactor, maxBatchSize, linearAttentionMetadata); @@ -8357,7 +8371,7 @@ TEST_F(KVCacheManagerTest, StaticLinearHybridAllocationTest) // so the linear pool falls back to memory-budget-based sizing rather than maxBatchSize. tle::KvCacheConfig const kvCacheConfigEnabledReuse{/*enableBlockReuse=*/true}; auto const dynamicBlocksPerWindow - = KVCacheManager::calculateMaxNumBlocks(kvCacheConfigEnabledReuse, nvinfer1::DataType::kHALF, + = KVCacheManager::calculateMaxNumBlocks(kvCacheConfigEnabledReuse, tensorrt_llm::DataType::kHALF, numKvHeadsPerLayer, sizePerHead, tokensPerBlock, worldConfig, windowSizeToLayers, allottedPrimaryMemBytes, allottedSecondaryMemBytes, extraCostMemory, kvFactor, maxBatchSize, linearAttentionMetadata); EXPECT_NE(std::get<0>(dynamicBlocksPerWindow.at(linearWindowSizeCode)), maxBatchSize); @@ -8390,7 +8404,7 @@ static auto makeBatchTestKVCacheManager(std::shared_ptr(numLayers, numKvHeads, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, beamWidth, std::vector{maxAttentionWindow}, - nvinfer1::DataType::kHALF, 0, stream, maxAttentionWindow, /*chunkSize=*/maxAttentionWindow, + tensorrt_llm::DataType::kHALF, 0, stream, maxAttentionWindow, /*chunkSize=*/maxAttentionWindow, /*enableBlockReuse=*/true, CacheType::kSELF, /*secondaryOffloadMinPriority=*/std::nullopt, /*eventManager=*/nullptr, @@ -8719,7 +8733,7 @@ TEST_F(KVCacheManagerTest, BatchAddSequence_NonLeafCopySourceTightPool) auto const blocksPerWindow = BlocksPerWindow{{maxAttentionWindow, {blocksInPrimaryPool, blocksInSecondaryPool}}}; KVCacheManager kvCacheManager(numLayers, numKvHeads, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, - beamWidth, std::vector{maxAttentionWindow}, nvinfer1::DataType::kHALF, 0, stream, + beamWidth, std::vector{maxAttentionWindow}, tensorrt_llm::DataType::kHALF, 0, stream, maxAttentionWindow, /*chunkSize=*/maxAttentionWindow, /*enableBlockReuse=*/true, CacheType::kSELF, /*secondaryOffloadMinPriority=*/std::nullopt, /*eventManager=*/nullptr, @@ -8876,7 +8890,7 @@ std::unique_ptr makePriorityEvictionManager( auto const blocksPerWindow = BlocksPerWindow{{maxAttentionWindow, {blocksInPrimaryPool, 0}}}; auto mgr = std::make_unique(kPE_NUM_LAYERS, kPE_NUM_HEADS, kPE_SIZE_PER_HEAD, kPE_TOKENS_PER_BLOCK, blocksPerWindow, kPE_MAX_NUM_SEQUENCES, kPE_BEAM_WIDTH, - std::vector{maxAttentionWindow}, nvinfer1::DataType::kHALF, 0, stream, + std::vector{maxAttentionWindow}, tensorrt_llm::DataType::kHALF, 0, stream, maxAttentionWindow, /*chunkSize=*/maxAttentionWindow, /*enableBlockReuse=*/true); mgr->allocatePools(false); return mgr; @@ -9210,7 +9224,7 @@ std::unique_ptr makeVSWAManager( { auto const blocksPerWindow = BlocksPerWindow{{kVSWA_ATTENTION_WINDOW, {blocksInPrimaryPool, 0}}}; auto mgr = std::make_unique(2, 2, 64, kVSWA_TOKENS_PER_BLOCK, blocksPerWindow, 8, kVSWA_BEAM_WIDTH, - std::vector{kVSWA_ATTENTION_WINDOW}, nvinfer1::DataType::kHALF, 0, stream, + std::vector{kVSWA_ATTENTION_WINDOW}, tensorrt_llm::DataType::kHALF, 0, stream, kVSWA_MAX_SEQUENCE_LENGTH, /*chunkSize=*/kVSWA_MAX_SEQUENCE_LENGTH, enableBlockReuse); mgr->allocatePools(false); return mgr; @@ -9229,7 +9243,7 @@ std::unique_ptr makeSmallWindowManager( SizeType32 constexpr kSmallMaxSeqLen = 128; auto const blocksPerWindow = BlocksPerWindow{{kSmallWindow, {blocksInPrimaryPool, 0}}}; auto mgr = std::make_unique(2, 2, 64, kSmallTpb, blocksPerWindow, 8, kVSWA_BEAM_WIDTH, - std::vector{kSmallWindow}, nvinfer1::DataType::kHALF, 0, stream, kSmallMaxSeqLen, + std::vector{kSmallWindow}, tensorrt_llm::DataType::kHALF, 0, stream, kSmallMaxSeqLen, /*chunkSize=*/kSmallMaxSeqLen, /*enableBlockReuse=*/true); mgr->allocatePools(false); return mgr; @@ -9952,7 +9966,7 @@ TEST_F(KVCacheManagerTest, VSWAEvictedPlaceholderAnchorAllowsTrailingReuse) auto const blocksPerWindow = BlocksPerWindow{{window, {blocksInPrimaryPool, 0}}}; KVCacheManager kvCacheManager(2, 2, 64, tpb, blocksPerWindow, 8, kVSWA_BEAM_WIDTH, std::vector{window}, - nvinfer1::DataType::kHALF, 0, stream, + tensorrt_llm::DataType::kHALF, 0, stream, /*maxSequenceLength=*/128, /*chunkSize=*/128, /*enableBlockReuse=*/true); kvCacheManager.allocatePools(false); auto const& blockManager = kvCacheManager.getBlockManager(); @@ -10088,7 +10102,7 @@ std::unique_ptr makeConnectorTestKVCacheManager( auto mgr = std::make_unique(std::vector(numLayers, numKvHeads), sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, beamWidth, std::vector{maxAttentionWindow}, - /*dtype*/ nvinfer1::DataType::kHALF, + /*dtype*/ tensorrt_llm::DataType::kHALF, /*sinkTokenLength*/ 0, stream, /*maxSequenceLength*/ maxAttentionWindow, /*chunkSize*/ maxAttentionWindow, @@ -10282,7 +10296,7 @@ TEST_F(KVCacheManagerTest, BlockManagerTestPerWindowFallback) auto constexpr maxBeamWidth = 1; auto constexpr smallWindow = 1024; auto constexpr largeWindow = 4096; - auto constexpr scalarDtype = nvinfer1::DataType::kHALF; + auto constexpr scalarDtype = tensorrt_llm::DataType::kHALF; auto const stream = std::make_shared(); auto const maxAttentionWindowVec = std::vector{smallWindow, largeWindow}; auto const blocksPerWindow = BlocksPerWindow{ @@ -10330,7 +10344,7 @@ TEST(BaseKVCacheManagerCalculateMaxNumBlocks, PerWindowOverrideDivergesByteBudge uint64_t const allottedPrimaryMemBytes = static_cast(1) << 30; // 1 GiB uint64_t const allottedSecondaryMemBytes = static_cast(1) << 30; size_t const extraCostMemory = 0; - auto const dtype = nvinfer1::DataType::kHALF; + auto const dtype = tensorrt_llm::DataType::kHALF; tensorrt_llm::executor::KvCacheConfig const config{}; tensorrt_llm::runtime::WorldConfig const worldConfig{}; @@ -10384,7 +10398,7 @@ TEST_F(KVCacheManagerTest, KVCacheManagerSWAEvictionCountPerWindow) auto constexpr maxNumSequences = 4; auto constexpr maxBeamWidth = 1; auto constexpr sinkTokenLength = 0; - auto constexpr dtype = nvinfer1::DataType::kHALF; + auto constexpr dtype = tensorrt_llm::DataType::kHALF; auto const stream = std::make_shared(); auto constexpr maxSequenceLength = 128; auto constexpr maxNewTokens = 40; @@ -10474,7 +10488,7 @@ TEST_F(KVCacheManagerTest, GenerationRequestClearCacheBlocksPerWindowResetsOnlyT BlockManager blockManager(std::vector(numLayers, numKvHeads), sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, stream, /*maxSequenceLength=*/fullWindow, maxBeamWidth, maxAttentionWindowVec, - nvinfer1::DataType::kHALF, /*sinkBubbleLength=*/0, /*chunkSize=*/0); + tensorrt_llm::DataType::kHALF, /*sinkBubbleLength=*/0, /*chunkSize=*/0); blockManager.allocatePools(/*useUvm=*/false); auto constexpr requestId = 7; @@ -10527,7 +10541,7 @@ TEST_F(KVCacheManagerTest, VswaMixedHeadDimReuseSmoke) auto constexpr smallSizePerHead = 256; auto constexpr largeSizePerHead = 512; auto constexpr sinkTokenLength = 0; - auto constexpr dtype = nvinfer1::DataType::kHALF; + auto constexpr dtype = tensorrt_llm::DataType::kHALF; auto constexpr maxSequenceLength = 64; auto constexpr maxNewTokens = 0; auto const stream = std::make_shared(); @@ -10633,11 +10647,12 @@ TEST_F(KVCacheManagerTest, VswaDisaggDtypeMismatchTriggersGuard) auto const maxAttentionWindowVec = std::vector{smallWindow, largeWindow}; auto const blocksPerWindow = BlocksPerWindow{ {smallWindow, {blocksInPrimary, blocksInSecondary}}, {largeWindow, {blocksInPrimary, blocksInSecondary}}}; - auto const poolConfigurations = std::vector{ - {smallWindow, sizePerHead, nvinfer1::DataType::kHALF}, {largeWindow, sizePerHead, nvinfer1::DataType::kBF16}}; + auto const poolConfigurations + = std::vector{{smallWindow, sizePerHead, tensorrt_llm::DataType::kHALF}, + {largeWindow, sizePerHead, tensorrt_llm::DataType::kBF16}}; auto kvCacheManager = std::make_unique(numLayers, numKvHeads, sizePerHead, tokensPerBlock, - blocksPerWindow, maxNumSequences, maxBeamWidth, maxAttentionWindowVec, /*dtype=*/nvinfer1::DataType::kHALF, + blocksPerWindow, maxNumSequences, maxBeamWidth, maxAttentionWindowVec, /*dtype=*/tensorrt_llm::DataType::kHALF, sinkTokenLength, stream, maxSequenceLength, /*chunkSize=*/0, /*enableBlockReuse=*/false, CacheType::kSELF, /*secondaryOffloadMinPriority=*/std::nullopt, /*eventManager=*/nullptr, diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheUtilsTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheUtilsTest.cpp index 2047a885d98c..a7b15e5ae8e0 100644 --- a/cpp/tests/unit_tests/batch_manager/kvCacheUtilsTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/kvCacheUtilsTest.cpp @@ -21,6 +21,7 @@ #include #include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/common/tllmDataType.h" namespace tc = tensorrt_llm::common; namespace tr = tensorrt_llm::runtime; @@ -48,7 +49,7 @@ TEST_F(BlockIteratorTest, BasicTest) auto constexpr mNumLayers = 5; auto constexpr mBlockSize = 32; auto const cacheShape = tr::ITensor::makeShape({mNumPrimaryBlocks, mNumLayers, 2, mBlockSize}); - constexpr nvinfer1::DataType dtype{tr::TRTDataType::value}; + constexpr tensorrt_llm::DataType dtype{tr::TRTDataType::value}; tr::ITensor::SharedPtr pool = tr::BufferManager::cpu(cacheShape, dtype); std::vector blockIds(mNumPrimaryBlocks); std::iota(blockIds.begin(), blockIds.end(), 0); @@ -75,7 +76,7 @@ TEST_F(BlockIteratorTest, BasicTest) TEST_F(BlockIteratorTest, CacheManagerTest) { - auto constexpr dataType = nvinfer1::DataType::kFLOAT; + auto constexpr dataType = tensorrt_llm::DataType::kFLOAT; auto constexpr numLayers = 12; auto constexpr numKvHeads = 6; auto constexpr sizePerHead = 16; diff --git a/cpp/tests/unit_tests/batch_manager/llmRequestTest.cpp b/cpp/tests/unit_tests/batch_manager/llmRequestTest.cpp index 120263b01af1..418178ff5023 100644 --- a/cpp/tests/unit_tests/batch_manager/llmRequestTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/llmRequestTest.cpp @@ -16,6 +16,7 @@ */ #include "tensorrt_llm/batch_manager/llmRequest.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/executor/executor.h" #include "tensorrt_llm/executor/types.h" @@ -97,7 +98,7 @@ TEST_F(LlmRequestTest, fromExecutorRequest) EXPECT_TRUE(llmReq.getStopWordsList().has_value()); { auto badWordsTensor = llmReq.getBadWordsList().value(); - EXPECT_EQ(badWordsTensor->getDataType(), nvinfer1::DataType::kINT32); + EXPECT_EQ(badWordsTensor->getDataType(), tensorrt_llm::DataType::kINT32); EXPECT_EQ(badWordsTensor->getShape().nbDims, 3); EXPECT_EQ(badWordsTensor->getShape().d[0], 1); EXPECT_EQ(badWordsTensor->getShape().d[1], 2); @@ -119,7 +120,7 @@ TEST_F(LlmRequestTest, fromExecutorRequest) { auto stopWordsTensor = llmReq.getStopWordsList().value(); - EXPECT_EQ(stopWordsTensor->getDataType(), nvinfer1::DataType::kINT32); + EXPECT_EQ(stopWordsTensor->getDataType(), tensorrt_llm::DataType::kINT32); EXPECT_EQ(stopWordsTensor->getShape().nbDims, 3); EXPECT_EQ(stopWordsTensor->getShape().d[0], 1); EXPECT_EQ(stopWordsTensor->getShape().d[1], 2); @@ -151,7 +152,7 @@ TEST_F(LlmRequestTest, fromExecutorRequest) EXPECT_EQ(llmReq.getPromptEmbeddingTable().value()->getShape().d[0], 1); EXPECT_EQ(llmReq.getPromptEmbeddingTable().value()->getShape().d[1], vocabSize); EXPECT_EQ(llmReq.getPromptEmbeddingTable().value()->getShape().d[2], hiddenSize); - EXPECT_EQ(llmReq.getPromptEmbeddingTable().value()->getDataType(), nvinfer1::DataType::kFLOAT); + EXPECT_EQ(llmReq.getPromptEmbeddingTable().value()->getDataType(), tensorrt_llm::DataType::kFLOAT); EXPECT_EQ(llmReq.getPromptVocabSize().value(), vocabSize); VecUniqueTokens uniqueTokens; for (size_t i = 0; i < inputTokens.size(); ++i) @@ -373,7 +374,7 @@ TEST_F(LlmRequestTest, testAllocateLogitsBuffer) EXPECT_EQ(llmReq.mPromptLen, 5); SizeType32 vocabSizePadded = 32000; - nvinfer1::DataType logitsDataType = nvinfer1::DataType::kFLOAT; + tensorrt_llm::DataType logitsDataType = tensorrt_llm::DataType::kFLOAT; // Test the allocation of context logits EXPECT_EQ(llmReq.getContextLogitsHost(), nullptr); @@ -462,7 +463,7 @@ TEST_F(LlmRequestTest, testCreateRequests) SizeType32 maxNewTokens{60}; tb::LlmRequest::RequestIdType requestId{77}; SizeType32 vocabSize{32}; - nvinfer1::DataType dtype{nvinfer1::DataType::kHALF}; + tensorrt_llm::DataType dtype{tensorrt_llm::DataType::kHALF}; tr::SamplingConfig samplingConfig(1); samplingConfig.randomSeed = std::vector{7}; diff --git a/cpp/tests/unit_tests/batch_manager/microBatchSchedulerTest.cpp b/cpp/tests/unit_tests/batch_manager/microBatchSchedulerTest.cpp index c2c308b37dd9..b2ecc787d6c6 100644 --- a/cpp/tests/unit_tests/batch_manager/microBatchSchedulerTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/microBatchSchedulerTest.cpp @@ -23,6 +23,7 @@ #include "tensorrt_llm/batch_manager/kvCacheManager.h" #include "tensorrt_llm/batch_manager/llmRequest.h" #include "tensorrt_llm/batch_manager/microBatchScheduler.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/testing/kvCacheManagerTestUtil.h" #include @@ -59,7 +60,7 @@ class MicroBatchSchedulerTest : public ::testing::Test // NOLINT(cppcoreguidelin { draftTokens = std::make_shared>(draftTokensLen, 2); draftLogits = BufferManager::cpu( - ITensor::makeShape({draftTokensLen, /* vocabSizePadded*/ 42}), nvinfer1::DataType::kFLOAT); + ITensor::makeShape({draftTokensLen, /* vocabSizePadded*/ 42}), tensorrt_llm::DataType::kFLOAT); } return std::make_shared(reqId, maxNewTokens, inputTokens, samplingConfig, /*isStreaming=*/false, @@ -1246,7 +1247,7 @@ class CombinedSchedulerTest : public ::testing::Test return std::make_shared( /*numLayers=*/10, /*nbKvHeads=*/10, /*sizePerHead=*/1, tokensPerBlock, blocksPerWindow, maxNumRequests, - /*maxBeamWidth=*/1, std::vector{maxNumTokensPerSeq}, nvinfer1::DataType::kHALF, + /*maxBeamWidth=*/1, std::vector{maxNumTokensPerSeq}, tensorrt_llm::DataType::kHALF, /*sinkTokenLength=*/0, stream, maxNumTokensPerSeq, /*chunkSize=*/maxNumTokensPerSeq, enableReuse); } diff --git a/cpp/tests/unit_tests/batch_manager/peftCacheManagerTest.cpp b/cpp/tests/unit_tests/batch_manager/peftCacheManagerTest.cpp index 49adfe5a6cb6..ef513894bfb6 100644 --- a/cpp/tests/unit_tests/batch_manager/peftCacheManagerTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/peftCacheManagerTest.cpp @@ -32,7 +32,7 @@ #include "tensorrt_llm/runtime/utils/numpyUtils.h" #include "tensorrt_llm/runtime/worldConfig.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include #include @@ -78,7 +78,7 @@ class PeftCacheManagerTest : public ::testing::Test // NOLINT(cppcoreguidelines- void SetUp() override { - mModelConfig = std::make_unique(0, 2, 2, 0, 1, 16, nvinfer1::DataType::kFLOAT); + mModelConfig = std::make_unique(0, 2, 2, 0, 1, 16, tensorrt_llm::DataType::kFLOAT); mModelConfig->setMlpHiddenSize(32); mWorldConfig = std::make_unique(2, 1, 1, 0); std::vector modules{ @@ -285,7 +285,7 @@ TEST_F(PeftCacheManagerTest, gptManagerSim) auto peftManager = std::make_unique(config, *mModelConfig, *mWorldConfig, *mManager); auto pageConfig = LoraCachePageManagerConfig( - runtime::MemoryType::kCPU, nvinfer1::DataType::kFLOAT, 128, 128, 2 * 8 * 64, 4 * 16, 1); + runtime::MemoryType::kCPU, tensorrt_llm::DataType::kFLOAT, 128, 128, 2 * 8 * 64, 4 * 16, 1); auto loraCache = std::make_unique(pageConfig, *mModelConfig, *mWorldConfig, *mManager); std::map> loras; @@ -505,7 +505,7 @@ TEST_F(PeftCacheManagerTest, getMaxNumSlots) config.numHostModuleLayer = 8192 * 8; config.numDeviceModuleLayer = 8292 * 2; auto [hostSlots, deviceSlots] - = PeftCacheManager::getMaxNumSlots(config, nvinfer1::DataType::kHALF, 256, 4 * 256, *mManager); + = PeftCacheManager::getMaxNumSlots(config, tensorrt_llm::DataType::kHALF, 256, 4 * 256, *mManager); EXPECT_EQ(262144, hostSlots); EXPECT_EQ(66336, deviceSlots); @@ -516,13 +516,13 @@ TEST_F(PeftCacheManagerTest, getMaxNumSlots) config.maxPagesPerBlockDevice = 8; std::tie(hostSlots, deviceSlots) - = PeftCacheManager::getMaxNumSlots(config, nvinfer1::DataType::kHALF, 256, 4 * 256, *mManager); + = PeftCacheManager::getMaxNumSlots(config, tensorrt_llm::DataType::kHALF, 256, 4 * 256, *mManager); EXPECT_EQ(195, hostSlots); EXPECT_EQ(66336, deviceSlots); std::tie(hostSlots, deviceSlots) - = PeftCacheManager::getMaxNumSlots(config, nvinfer1::DataType::kFLOAT, 384, 4 * 1024, *mManager); + = PeftCacheManager::getMaxNumSlots(config, tensorrt_llm::DataType::kFLOAT, 384, 4 * 1024, *mManager); config.hostCacheSize = 100000000; config.numHostModuleLayer = 8291 * 2; @@ -539,7 +539,7 @@ TEST_F(PeftCacheManagerTest, getPageManagerConfig) auto [hostCfg, deviceCfg] = PeftCacheManager::getPageManagerConfig(config, *mModelConfig, *mWorldConfig, *mManager); EXPECT_EQ(runtime::MemoryType::kCPU, hostCfg.getMemoryType()); - EXPECT_EQ(nvinfer1::DataType::kFLOAT, hostCfg.getDataType()); + EXPECT_EQ(tensorrt_llm::DataType::kFLOAT, hostCfg.getDataType()); EXPECT_EQ(456, hostCfg.getTotalNumPages()); EXPECT_EQ(24, hostCfg.getMaxPagesPerBlock()); EXPECT_EQ(288, hostCfg.getSlotsPerPage()); @@ -547,7 +547,7 @@ TEST_F(PeftCacheManagerTest, getPageManagerConfig) EXPECT_FALSE(hostCfg.getInitToZero()); EXPECT_EQ(runtime::MemoryType::kGPU, deviceCfg.getMemoryType()); - EXPECT_EQ(nvinfer1::DataType::kFLOAT, deviceCfg.getDataType()); + EXPECT_EQ(tensorrt_llm::DataType::kFLOAT, deviceCfg.getDataType()); EXPECT_EQ(116, deviceCfg.getTotalNumPages()); EXPECT_EQ(8, deviceCfg.getMaxPagesPerBlock()); EXPECT_EQ(288, deviceCfg.getSlotsPerPage()); @@ -563,7 +563,7 @@ TEST_F(PeftCacheManagerTest, getPageManagerConfig) = PeftCacheManager::getPageManagerConfig(config, *mModelConfig, *mWorldConfig, *mManager); EXPECT_EQ(runtime::MemoryType::kCPU, hostCfg.getMemoryType()); - EXPECT_EQ(nvinfer1::DataType::kFLOAT, hostCfg.getDataType()); + EXPECT_EQ(tensorrt_llm::DataType::kFLOAT, hostCfg.getDataType()); EXPECT_EQ(3617, hostCfg.getTotalNumPages()); EXPECT_EQ(4, hostCfg.getMaxPagesPerBlock()); EXPECT_EQ(288, hostCfg.getSlotsPerPage()); @@ -571,7 +571,7 @@ TEST_F(PeftCacheManagerTest, getPageManagerConfig) EXPECT_FALSE(hostCfg.getInitToZero()); EXPECT_EQ(runtime::MemoryType::kGPU, deviceCfg.getMemoryType()); - EXPECT_EQ(nvinfer1::DataType::kFLOAT, deviceCfg.getDataType()); + EXPECT_EQ(tensorrt_llm::DataType::kFLOAT, deviceCfg.getDataType()); EXPECT_EQ(116, deviceCfg.getTotalNumPages()); EXPECT_EQ(8, deviceCfg.getMaxPagesPerBlock()); EXPECT_EQ(288, deviceCfg.getSlotsPerPage()); @@ -586,7 +586,7 @@ class PeftCacheManagerPrefetchTest : public ::testing::Test // NOLINT(cppcoregui void SetUp() override { - mModelConfig = std::make_unique(0, 2, 2, 0, 1, 16, nvinfer1::DataType::kFLOAT); + mModelConfig = std::make_unique(0, 2, 2, 0, 1, 16, tensorrt_llm::DataType::kFLOAT); mModelConfig->setMlpHiddenSize(32); mWorldConfig = std::make_unique(2, 1, 1, 0); std::vector modules{ diff --git a/cpp/tests/unit_tests/batch_manager/rnnCacheFormatterTest.cpp b/cpp/tests/unit_tests/batch_manager/rnnCacheFormatterTest.cpp index 8840130df1ea..7dc4cea6e9af 100644 --- a/cpp/tests/unit_tests/batch_manager/rnnCacheFormatterTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/rnnCacheFormatterTest.cpp @@ -7,6 +7,7 @@ #include "tensorrt_llm/batch_manager/cacheFormatter.h" #include "tensorrt_llm/batch_manager/rnnCacheFormatter.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/executor/dataTransceiverState.h" #include @@ -27,12 +28,12 @@ class RnnTargetIRanksTest : public ::testing::Test std::vector kvLayersPerPP(pp, 0); // No attention layers auto state = texec::kv_cache::CacheState( /*nbAttentionLayers=*/0, /*nbKvHeads=*/1, /*sizePerHead=*/64, /*tokensPerBlock=*/32, tp, pp, - /*contextParallelism=*/1, kvLayersPerPP, nvinfer1::DataType::kFLOAT); + /*contextParallelism=*/1, kvLayersPerPP, tensorrt_llm::DataType::kFLOAT); texec::kv_cache::CacheState::RnnModelConfig rnnModelConfig{/*mDState=*/16, /*mDConv=*/4, /*mHiddenSize=*/256, /*mHeadDim=*/64, /*mConvDimSize=*/128, /*mNGroups=*/1, /*mNumLayers=*/numLayers, /*mNumHeads=*/4}; - state.setRnnConfig(rnnModelConfig, layersPerPP, nvinfer1::DataType::kFLOAT, nvinfer1::DataType::kFLOAT); + state.setRnnConfig(rnnModelConfig, layersPerPP, tensorrt_llm::DataType::kFLOAT, tensorrt_llm::DataType::kFLOAT); return state; } }; @@ -216,7 +217,7 @@ class HybridModelCounterpartsTest : public ::testing::Test SizeType32 tokensPerBlock = 32) { return texec::kv_cache::CacheState(numLayers, numHeads, sizePerHead, tokensPerBlock, tp, pp, - /*contextParallelism=*/1, layersPerPP, nvinfer1::DataType::kFLOAT, + /*contextParallelism=*/1, layersPerPP, tensorrt_llm::DataType::kFLOAT, texec::kv_cache::CacheState::AttentionType::kDEFAULT, /*kvFactor=*/2, /*enableAttentionDP=*/false, /*DPrank=*/0, /*DPsize=*/1); } @@ -228,7 +229,8 @@ class HybridModelCounterpartsTest : public ::testing::Test { auto state = makeKvCacheState(kvNumLayers, tp, pp, kvLayersPerPP, numHeads, sizePerHead, tokensPerBlock); texec::kv_cache::CacheState::RnnModelConfig rnnModelConfig{16, 4, 256, 64, 128, 1, rnnNumLayers, 4}; - state.setRnnConfig(rnnModelConfig, rnnLayersPerPP, nvinfer1::DataType::kFLOAT, nvinfer1::DataType::kFLOAT); + state.setRnnConfig( + rnnModelConfig, rnnLayersPerPP, tensorrt_llm::DataType::kFLOAT, tensorrt_llm::DataType::kFLOAT); return state; } @@ -471,7 +473,7 @@ class AttentionOnlyModelTest : public ::testing::Test SizeType32 tokensPerBlock = 32) { return texec::kv_cache::CacheState(numLayers, numHeads, sizePerHead, tokensPerBlock, tp, pp, - /*contextParallelism=*/1, layersPerPP, nvinfer1::DataType::kFLOAT, + /*contextParallelism=*/1, layersPerPP, tensorrt_llm::DataType::kFLOAT, texec::kv_cache::CacheState::AttentionType::kDEFAULT, /*kvFactor=*/2, /*enableAttentionDP=*/false, /*DPrank=*/0, /*DPsize=*/1); } diff --git a/cpp/tests/unit_tests/batch_manager/truncateBlocksTest.cpp b/cpp/tests/unit_tests/batch_manager/truncateBlocksTest.cpp index d0dce8eb71c5..85d4ac112244 100644 --- a/cpp/tests/unit_tests/batch_manager/truncateBlocksTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/truncateBlocksTest.cpp @@ -13,6 +13,7 @@ #include "tensorrt_llm/batch_manager/kvCacheManager.h" #include "tensorrt_llm/batch_manager/llmRequest.h" #include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/runtime/samplingConfig.h" #include "tensorrt_llm/testing/kvCacheManagerTestUtil.h" @@ -89,7 +90,7 @@ TEST_F(TruncateBlocksTest, MultiTurnConversationTruncation) // Create KVCacheManager with block reuse enabled KVCacheManager kvCacheManager(numLayers, numHeads, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, - beamWidth, std::vector{maxAttentionWindow}, nvinfer1::DataType::kHALF, 0, stream, + beamWidth, std::vector{maxAttentionWindow}, tensorrt_llm::DataType::kHALF, 0, stream, maxSequenceLength, maxSequenceLength /* chunkSize */, true /* enableBlockReuse */); kvCacheManager.allocatePools(false); @@ -219,7 +220,7 @@ TEST_F(TruncateBlocksTest, SharedPrefixTruncation) auto const blocksPerWindow = BlocksPerWindow{{maxAttentionWindow, {blocksInPrimaryPool, blocksInSecondaryPool}}}; KVCacheManager kvCacheManager(numLayers, numHeads, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, - beamWidth, std::vector{maxAttentionWindow}, nvinfer1::DataType::kHALF, 0, stream, + beamWidth, std::vector{maxAttentionWindow}, tensorrt_llm::DataType::kHALF, 0, stream, maxSequenceLength, maxSequenceLength /* chunkSize */, true /* enableBlockReuse */); kvCacheManager.allocatePools(false); @@ -317,7 +318,7 @@ TEST_F(TruncateBlocksTest, CompleteTruncation) auto const blocksPerWindow = BlocksPerWindow{{maxAttentionWindow, {blocksInPrimaryPool, blocksInSecondaryPool}}}; KVCacheManager kvCacheManager(numLayers, numHeads, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, - beamWidth, std::vector{maxAttentionWindow}, nvinfer1::DataType::kHALF, 0, stream, + beamWidth, std::vector{maxAttentionWindow}, tensorrt_llm::DataType::kHALF, 0, stream, maxSequenceLength, maxSequenceLength /* chunkSize */, true /* enableBlockReuse */); kvCacheManager.allocatePools(false); @@ -377,7 +378,7 @@ TEST_F(TruncateBlocksTest, NonExistentTokensTruncation) auto const blocksPerWindow = BlocksPerWindow{{maxAttentionWindow, {blocksInPrimaryPool, blocksInSecondaryPool}}}; KVCacheManager kvCacheManager(numLayers, numHeads, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, - beamWidth, std::vector{maxAttentionWindow}, nvinfer1::DataType::kHALF, 0, stream, + beamWidth, std::vector{maxAttentionWindow}, tensorrt_llm::DataType::kHALF, 0, stream, maxSequenceLength, maxSequenceLength /* chunkSize */, true /* enableBlockReuse */); kvCacheManager.allocatePools(false); @@ -453,7 +454,7 @@ TEST_F(TruncateBlocksTest, ComplexMultiTurnConversationTruncation) // Create KVCacheManager with block reuse enabled KVCacheManager kvCacheManager(numLayers, numHeads, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, - beamWidth, std::vector{maxAttentionWindow}, nvinfer1::DataType::kHALF, 0, stream, + beamWidth, std::vector{maxAttentionWindow}, tensorrt_llm::DataType::kHALF, 0, stream, maxSequenceLength, maxSequenceLength /* chunkSize */, true /* enableBlockReuse */); kvCacheManager.allocatePools(false); diff --git a/cpp/tests/unit_tests/common/loggerTest.cpp b/cpp/tests/unit_tests/common/loggerTest.cpp index 8aaea6863fb9..6fa28a69b5f8 100644 --- a/cpp/tests/unit_tests/common/loggerTest.cpp +++ b/cpp/tests/unit_tests/common/loggerTest.cpp @@ -26,7 +26,7 @@ using namespace tensorrt_llm::common; TEST(LoggerModuleTest, FormatModuleNoTrailingSpaces) { for (auto const* raw : {"batch_manager", "common", "cutlass_extensions", "deep_ep", "deep_gemm", "executor", - "executor_worker", "flash_mla", "kernels", "layers", "nanobind", "plugins", "runtime", "testing", "thop"}) + "flash_mla", "kernels", "layers", "nanobind", "runtime", "testing", "thop"}) { auto const fmt = formatModule(raw); EXPECT_FALSE(fmt.empty()); diff --git a/cpp/tests/unit_tests/executor/CMakeLists.txt b/cpp/tests/unit_tests/executor/CMakeLists.txt index a51baa6ed00f..b66727ca677d 100644 --- a/cpp/tests/unit_tests/executor/CMakeLists.txt +++ b/cpp/tests/unit_tests/executor/CMakeLists.txt @@ -19,21 +19,11 @@ add_gtest(decodingConfigTest decodingConfigTest.cpp) add_gtest(requestTest requestTest.cpp) add_gtest(responseTest responseTest.cpp) -add_gtest(executorTestSmall executorTestSmall.cpp) -target_link_libraries(executorTestSmall PRIVATE testingUtils) - -add_gtest(executorTestSmallArbitraryOutputTensors - executorTestSmallArbitraryOutputTensors.cpp) -target_link_libraries(executorTestSmallArbitraryOutputTensors - PRIVATE testingUtils) - add_gtest(executorConfigTest executorConfigTest.cpp) add_gtest(executorTensorTest tensorTest.cpp) add_gtest(serializeUtilsTest serializeUtilsTest.cpp) add_gtest(requestWithIdTest requestWithIdTest.cpp) add_gtest(loraConfigTest loraConfigTest.cpp) -add_gtest(intervalSetTest intervalSetTest.cpp) -add_gtest(dynamicBatchTunerTest dynamicBatchTunerTest.cpp) add_gtest(genUniqueAgentNameTest genUniqueAgentNameTest.cpp) target_link_libraries(genUniqueAgentNameTest PRIVATE ${Python3_LIBRARIES}) add_gtest(ucxCommTest ucxCommTest.cpp) diff --git a/cpp/tests/unit_tests/executor/agentCommTest.cpp b/cpp/tests/unit_tests/executor/agentCommTest.cpp index 194d5267c6b3..89488ba373e8 100644 --- a/cpp/tests/unit_tests/executor/agentCommTest.cpp +++ b/cpp/tests/unit_tests/executor/agentCommTest.cpp @@ -15,6 +15,7 @@ * limitations under the License. */ +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/executor/cache_transmission/agent_utils/connection.h" #include @@ -120,7 +121,7 @@ class AgentCommTest : public ::testing::TestWithParam auto constexpr blocksInSecondaryPool = 0; auto constexpr enableBlockReuse = true; - auto constexpr dataType = nvinfer1::DataType::kFLOAT; + auto constexpr dataType = tensorrt_llm::DataType::kFLOAT; using BlocksPerWindow = std::map>; BlocksPerWindow const blocksPerWindow diff --git a/cpp/tests/unit_tests/executor/dynamicBatchTunerTest.cpp b/cpp/tests/unit_tests/executor/dynamicBatchTunerTest.cpp deleted file mode 100644 index 04ce10393b67..000000000000 --- a/cpp/tests/unit_tests/executor/dynamicBatchTunerTest.cpp +++ /dev/null @@ -1,99 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2023-2025 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. - */ - -#include "tensorrt_llm/executor/dynamicBatchTuner.h" -#include "tensorrt_llm/common/tllmException.h" -#include "tensorrt_llm/executor/executor.h" -#include "tensorrt_llm/executor/types.h" -#include -#include - -using ::testing::_; -using ::testing::Invoke; - -using namespace tensorrt_llm::executor; -using namespace tensorrt_llm::common; - -TEST(DynamicBatchTunerTest, Stats) -{ - // moving average window size is 3 - DynamicBatchConfig dynamicBatchConfig(true, true, 3); - DynamicBatchTuner dynamicBatchTuner(dynamicBatchConfig); - - // check no division by zero issue - EXPECT_EQ(dynamicBatchTuner.getAverageInputLength(), 0); - EXPECT_EQ(dynamicBatchTuner.getAverageOutputLength(), 0); - - dynamicBatchTuner.updateStats(1, 2); - EXPECT_EQ(dynamicBatchTuner.getAverageInputLength(), 1); - EXPECT_EQ(dynamicBatchTuner.getAverageOutputLength(), 2); - - dynamicBatchTuner.updateStats(2, 3); - EXPECT_EQ(dynamicBatchTuner.getAverageInputLength(), 1.5); - EXPECT_EQ(dynamicBatchTuner.getAverageOutputLength(), 2.5); - - dynamicBatchTuner.updateStats(3, 4); - EXPECT_EQ(dynamicBatchTuner.getAverageInputLength(), 2); - EXPECT_EQ(dynamicBatchTuner.getAverageOutputLength(), 3); - - // check that the first element is removed from the moving average window - dynamicBatchTuner.updateStats(4, 5); - EXPECT_EQ(dynamicBatchTuner.getAverageInputLength(), 3); - EXPECT_EQ(dynamicBatchTuner.getAverageOutputLength(), 4); -} - -TEST(DynamicBatchConfig, RuntimeBatchSize) -{ - // moving average window size is 3 - DynamicBatchConfig dynamicBatchConfig(true, true, 3); - DynamicBatchTuner dynamicBatchTuner(dynamicBatchConfig); - // check runtime batch size computation - EXPECT_EQ(dynamicBatchTuner.getRuntimeBatchSize(143), 128); - EXPECT_EQ(dynamicBatchTuner.getRuntimeBatchSize(335), 256); - EXPECT_EQ(dynamicBatchTuner.getRuntimeBatchSize(671), 512); - EXPECT_EQ(dynamicBatchTuner.getRuntimeBatchSize(831), 768); - EXPECT_EQ(dynamicBatchTuner.getRuntimeBatchSize(1279), 1024); - EXPECT_EQ(dynamicBatchTuner.getRuntimeBatchSize(1663), 1536); - // fall back - EXPECT_EQ(dynamicBatchTuner.getRuntimeBatchSize(2049), 2048); - EXPECT_EQ(dynamicBatchTuner.getRuntimeBatchSize(1665), 1665); -} - -TEST(DynamicBatchConfig, RuntimeMaxNumTokens) -{ - // moving average window size is 1 - DynamicBatchConfig dynamicBatchConfig(true, true, 1); - DynamicBatchTuner dynamicBatchTuner(dynamicBatchConfig); - - // context heavy - dynamicBatchTuner.updateStats(100, 2); - EXPECT_EQ(dynamicBatchTuner.getRuntimeMaxNumTokens(1), 8192); - // context heavy fall back - EXPECT_EQ(dynamicBatchTuner.getRuntimeMaxNumTokens(256), 16384); - - // balanced - dynamicBatchTuner.updateStats(100, 100); - EXPECT_EQ(dynamicBatchTuner.getRuntimeMaxNumTokens(1), 4096); - // balanced fall back - EXPECT_EQ(dynamicBatchTuner.getRuntimeMaxNumTokens(4000), 8192); - - // gen heavy - dynamicBatchTuner.updateStats(2, 256); - EXPECT_EQ(dynamicBatchTuner.getRuntimeMaxNumTokens(1), 2048); - // gen heavy fall back - EXPECT_EQ(dynamicBatchTuner.getRuntimeMaxNumTokens(4000), 4096); -} diff --git a/cpp/tests/unit_tests/executor/executorTestSmall.cpp b/cpp/tests/unit_tests/executor/executorTestSmall.cpp deleted file mode 100644 index 2987509f16ac..000000000000 --- a/cpp/tests/unit_tests/executor/executorTestSmall.cpp +++ /dev/null @@ -1,289 +0,0 @@ -#include "tensorrt_llm/batch_manager/trtGptModelInflightBatching.h" -#include "tensorrt_llm/executor/executor.h" -#include "tensorrt_llm/runtime/common.h" -#include "tensorrt_llm/runtime/rawEngine.h" -#include "tensorrt_llm/runtime/tllmLogger.h" -#include "tests/utils/common.h" -#include "tests/utils/engines.h" -#include "tests/utils/executorUtils.h" - -#include "gtest/gtest.h" - -#include -#include -#include - -namespace tensorrt_llm::testing -{ - -struct TrivialConstantDecoderTestParameters -{ - using TupleT = std::tuple; - runtime::SizeType32 randomSeed; - runtime::SizeType32 vocabSize; - runtime::SizeType32 maxNumTokens; - runtime::SizeType32 maxBeamWidth; - runtime::SizeType32 maxBatchSize; - runtime::SizeType32 numRequests; - runtime::SizeType32 promptLength; - runtime::SizeType32 maxOutputLength; - - // Constructor that takes a tuple - TrivialConstantDecoderTestParameters( // NOLINT: implicit to allow gtest to convert from tuple generated by - // 'combine' - TupleT t) - : randomSeed(std::get<0>(t)) - , vocabSize(std::get<1>(t)) - , maxNumTokens(std::get<2>(t)) - , maxBeamWidth(std::get<3>(t)) - , maxBatchSize(std::get<4>(t)) - , numRequests(std::get<5>(t)) - , promptLength(std::get<6>(t)) - , maxOutputLength(std::get<7>(t)) - { - } -}; - -template -struct DecoderTestShared -{ - static constexpr runtime::SizeType32 kNumTokensPerBlock = 64; - static constexpr runtime::SizeType32 kKvCacheMaxTokens = 2048 * 8; - - DecoderTestShared(std::shared_ptr logger, std::mt19937 rng, - std::shared_ptr executor, std::vector randomLogits) - : logger(std::move(logger)) - , rng(rng) - , executor(std::move(executor)) - , randomLogits(std::move(randomLogits)){}; - std::shared_ptr logger; - std::mt19937 rng; - std::shared_ptr executor; - std::vector randomLogits; -}; - -template -std::unique_ptr> SetupDecoderTest(TrivialConstantDecoderTestParameters const& params) -{ - auto logger = std::make_shared(); - auto rng = std::mt19937(params.randomSeed); - auto randomLogits = tensorrt_llm::testing::randomLogits(params.vocabSize, &rng); - auto const decoderParameters = tensorrt_llm::testing::utils::engines::ConstantTrivialDecoderParameters{ - tensorrt_llm::testing::utils::engines::TrivialDecoderParameters{params.vocabSize, params.maxBatchSize, - params.maxNumTokens, DecoderTestShared::kNumTokensPerBlock, params.maxBeamWidth, false}, - randomLogits}; - auto engineHostMemory - = tensorrt_llm::testing::utils::engines::createConstantTrivialDecoder(decoderParameters, logger); - auto const engine = runtime::RawEngine(engineHostMemory.release()); - auto const dtype = runtime::TRTDataType::value; - auto modelConfig = runtime::ModelConfig(params.vocabSize, 1, 1, 0, 1, 1, dtype); - modelConfig.useGptAttentionPlugin(true); - modelConfig.setModelVariant(runtime::ModelConfig::ModelVariant::kGpt); - modelConfig.usePackedInput(true); - modelConfig.setKVCacheType(runtime::ModelConfig::KVCacheType::kPAGED); - modelConfig.setMaxNumTokens(params.maxNumTokens); - modelConfig.setMaxBatchSize(params.maxBatchSize); - modelConfig.setMaxBeamWidth(params.maxBeamWidth); - modelConfig.setMaxSequenceLen(params.maxNumTokens); - modelConfig.setMaxInputLen(params.maxNumTokens); - modelConfig.setLayerTypes({runtime::ModelConfig::LayerType::kATTENTION}); - modelConfig.setTokensPerBlock(DecoderTestShared::kNumTokensPerBlock); - modelConfig.setPagedContextFMHA(true); - - auto const worldConfig = runtime::WorldConfig(); - auto kvCacheConfig = executor::KvCacheConfig{}; - kvCacheConfig.setMaxTokens(DecoderTestShared::kKvCacheMaxTokens); - - auto const executorConfig - = tensorrt_llm::executor::ExecutorConfig(params.maxBeamWidth, executor::SchedulerConfig(), kvCacheConfig, true, - true, 1, 1, executor::BatchingType::kINFLIGHT, params.maxBatchSize, params.maxNumTokens, std::nullopt, - std::nullopt, std::nullopt, std::nullopt, false, 1, std::nullopt, executor::ExtendedRuntimePerfKnobConfig(), - std::nullopt, 0, executor::ExecutorConfig::kDefaultMaxSeqIdleMicroseconds, std::nullopt, std::nullopt); - - auto model = std::make_shared( - logger, modelConfig, worldConfig, engine, false, executorConfig, false); - - return std::make_unique>( - logger, rng, std::make_shared(model, executorConfig), randomLogits); -} - -template -class DecoderTest : public ::testing::Test, public ::testing::WithParamInterface -{ -protected: - std::unique_ptr> state; - - DecoderTest() - { - auto const params = GetParam(); - state = SetupDecoderTest(params); - } - - void runDecoderTest(TrivialConstantDecoderTestParameters const& parameters) - { - auto const requestTokens = createConsecutiveTokenSequence(parameters.promptLength, parameters.vocabSize, 0); - auto requests = std::vector{}; - requests.reserve(static_cast(parameters.numRequests)); - for (auto i = 0; i < parameters.numRequests; i++) - { - requests.emplace_back(requestTokens, parameters.maxOutputLength, false, executor::SamplingConfig{}, - executor::OutputConfig{false, false, false, true, false, false}); - } - auto const accumulatedResponses - = runThroughRequests(*state->executor, requests, std::chrono::duration(3600000)); - ASSERT_EQ(accumulatedResponses.size(), parameters.numRequests); - - std::sort(state->randomLogits.begin(), state->randomLogits.end()); - std::reverse(state->randomLogits.begin(), state->randomLogits.end()); - for (auto const& [requestId, responses] : accumulatedResponses) - { - for (auto const& response : responses) - { - ASSERT_FALSE(response.hasError()); - auto const& tokensByBeam = response.getResult().outputTokenIds; - ASSERT_EQ(tokensByBeam.size(), 1); - for (auto const& tokensForBeam : tokensByBeam) - { - ASSERT_EQ(tokensForBeam.size(), parameters.maxOutputLength); - } - } - } - } -}; - -namespace -{ -constexpr runtime::SizeType32 kRandomSeed1 = 45; -auto const randomSeeds = ::testing::Values(kRandomSeed1); - -constexpr runtime::SizeType32 kMinVocabSize = 16; -auto const vocabSizes = ::testing::Values(kMinVocabSize); - -constexpr runtime::SizeType32 kMinMaxNumTokens = 2048; -auto const maxNumTokenses = ::testing::Values(kMinMaxNumTokens); - -constexpr runtime::SizeType32 kMinBeamWidth = 1; -auto const beamWidths = ::testing::Values(kMinBeamWidth); - -constexpr runtime::SizeType32 kMinMaxBatchSize = 2048; -auto const maxBatchSizes = ::testing::Values(kMinMaxBatchSize); - -constexpr runtime::SizeType32 kMinNumRequests = 64; -auto const numRequestses = ::testing::Values(kMinNumRequests); - -constexpr runtime::SizeType32 kMinPromptLength = 32; -auto const promptLengths = ::testing::Values(kMinPromptLength); - -constexpr runtime::SizeType32 kMinMaxOutputLength = 16; -auto const maxOutputLengths = ::testing::Values(kMinMaxOutputLength); - -auto const paramGenerator - = ::testing::ConvertGenerator(::testing::Combine(randomSeeds, - vocabSizes, maxNumTokenses, beamWidths, maxBatchSizes, numRequestses, promptLengths, maxOutputLengths)); -} // namespace - -using DecoderFloatTest = DecoderTest; - -TEST_P(DecoderFloatTest, TestSizeAndValues) -{ - runDecoderTest(GetParam()); -} - -INSTANTIATE_TEST_SUITE_P(Float, DecoderFloatTest, paramGenerator, - [](::testing::TestParamInfo const& info) -> std::string - { - std::stringstream nameStringStream; - nameStringStream << "_maxBatchSize_" << info.param.maxBatchSize << "_vocabSize_" << info.param.vocabSize - << "_maxBeamWidth_" << info.param.maxBeamWidth << "_maxNumTokens_" << info.param.maxNumTokens - << "_maxOutputLength_" << info.param.maxOutputLength << "_numRequests_" - << info.param.numRequests << "_promptLength_" << info.param.promptLength << "_randomSeed_" - << info.param.randomSeed; - return nameStringStream.str(); - }); - -// Helper function to test calculateCacheSizePerToken with given parameters. -std::map calculateCacheSizePerTokenHelper( - std::vector const& maxAttentionWindowVec, runtime::SizeType32 kvFactor = 2, - runtime::SizeType32 vocabSize = 32, runtime::SizeType32 nbLayers = 4, runtime::SizeType32 nbAttentionLayers = 4, - runtime::SizeType32 nbRnnLayers = 0, runtime::SizeType32 nbHeads = 8, runtime::SizeType32 hiddenSize = 512, - bool isCrossAttention = false) -{ - // Create minimal ModelConfig for testing. - auto modelConfig = runtime::ModelConfig( - vocabSize, nbLayers, nbAttentionLayers, nbRnnLayers, nbHeads, hiddenSize, nvinfer1::DataType::kFLOAT); - modelConfig.useGptAttentionPlugin(true); - modelConfig.setModelVariant(runtime::ModelConfig::ModelVariant::kGpt); - modelConfig.setKVCacheType(runtime::ModelConfig::KVCacheType::kPAGED); - - auto const worldConfig = runtime::WorldConfig(); - - return batch_manager::TrtGptModelInflightBatching::calculateCacheSizePerTokenForDisagg( - modelConfig, worldConfig, maxAttentionWindowVec, isCrossAttention, kvFactor); -} - -// Test for TrtGptModelInflightBatching::calculateCacheSizePerToken function with different layer types. -TEST(TrtInflightBatchingTest, CalculateCacheSizePerTokenForDisagg) -{ - // Common parameters. - constexpr runtime::SizeType32 nbLayers = 5; - constexpr runtime::SizeType32 hiddenSize = 512; - constexpr runtime::SizeType32 kvFactor = 2; - constexpr runtime::SizeType32 vocabSize = 32; - constexpr runtime::SizeType32 nbHeads = 8; - // Test case 1: Single attention window size - attention layers only. - { - std::vector maxAttentionWindowVec = {128}; - constexpr runtime::SizeType32 nbAttentionLayers = 5; - constexpr runtime::SizeType32 numBytesPerFloatElement = 4; - constexpr runtime::SizeType32 nbRnnLayers = 0; - auto result = calculateCacheSizePerTokenHelper(maxAttentionWindowVec, kvFactor, vocabSize, nbLayers, - nbAttentionLayers, nbRnnLayers, nbHeads, hiddenSize, false); - EXPECT_EQ(result.size(), 1); - EXPECT_EQ(result.at(128), nbAttentionLayers * kvFactor * hiddenSize * numBytesPerFloatElement); - } - - // Test case 2: Multiple attention window sizes - attention layers only. - { - std::vector maxAttentionWindowVec = {128, 256}; - constexpr runtime::SizeType32 nbAttentionLayers = 5; - constexpr runtime::SizeType32 numBytesPerFloatElement = 4; - constexpr runtime::SizeType32 nbRnnLayers = 0; - auto result = calculateCacheSizePerTokenHelper(maxAttentionWindowVec, kvFactor, vocabSize, nbLayers, - nbAttentionLayers, nbRnnLayers, nbHeads, hiddenSize, false); - EXPECT_EQ(result.size(), 2); - auto const nbAttentionLayersIn128Window = 3; - auto const nbAttentionLayersIn256Window = 2; - EXPECT_EQ(result.at(128), nbAttentionLayersIn128Window * kvFactor * hiddenSize * numBytesPerFloatElement); - EXPECT_EQ(result.at(256), nbAttentionLayersIn256Window * kvFactor * hiddenSize * numBytesPerFloatElement); - } - - // Test case 3: Single attention window size - attention and rnn layers. - { - std::vector maxAttentionWindowVec = {128}; - constexpr runtime::SizeType32 nbAttentionLayers = 3; - constexpr runtime::SizeType32 numBytesPerFloatElement = 4; - constexpr runtime::SizeType32 nbRnnLayers = 2; - auto result = calculateCacheSizePerTokenHelper(maxAttentionWindowVec, kvFactor, vocabSize, nbLayers, - nbAttentionLayers, nbRnnLayers, nbHeads, hiddenSize, false); - EXPECT_EQ(result.size(), 1); - EXPECT_EQ(result.at(128), nbAttentionLayers * kvFactor * hiddenSize * numBytesPerFloatElement); - } - - // Test case 4: Multiple attention window sizes - attention and rnn layers. - { - std::vector maxAttentionWindowVec = {128, 256}; - constexpr runtime::SizeType32 nbAttentionLayers = 3; - constexpr runtime::SizeType32 numBytesPerFloatElement = 4; - constexpr runtime::SizeType32 nbRnnLayers = 2; - auto result = calculateCacheSizePerTokenHelper(maxAttentionWindowVec, kvFactor, vocabSize, nbLayers, - nbAttentionLayers, nbRnnLayers, nbHeads, hiddenSize, false); - EXPECT_EQ(result.size(), 2); - auto const nbAttentionLayersIn128Window = 2; - auto const nbAttentionLayersIn256Window = 1; - EXPECT_EQ(result.at(128), nbAttentionLayersIn128Window * kvFactor * hiddenSize * numBytesPerFloatElement); - EXPECT_EQ(result.at(256), nbAttentionLayersIn256Window * kvFactor * hiddenSize * numBytesPerFloatElement); - } -} - -} // namespace tensorrt_llm::testing diff --git a/cpp/tests/unit_tests/executor/executorTestSmallArbitraryOutputTensors.cpp b/cpp/tests/unit_tests/executor/executorTestSmallArbitraryOutputTensors.cpp deleted file mode 100644 index b64bd775fe30..000000000000 --- a/cpp/tests/unit_tests/executor/executorTestSmallArbitraryOutputTensors.cpp +++ /dev/null @@ -1,491 +0,0 @@ -#include "include/tensorrt_llm/executor/executor.h" -#include "tensorrt_llm/batch_manager/trtGptModelInflightBatching.h" -#include "tensorrt_llm/executor/types.h" -#include "tensorrt_llm/runtime/common.h" -#include "tensorrt_llm/runtime/iBuffer.h" -#include "tensorrt_llm/runtime/modelConfig.h" -#include "tensorrt_llm/runtime/rawEngine.h" -#include "tensorrt_llm/runtime/tllmLogger.h" -#include "tensorrt_llm/runtime/worldConfig.h" -#include "tests/utils/common.h" -#include "tests/utils/engines.h" -#include "tests/utils/executorUtils.h" - -#include "gtest/gtest.h" -#include -#include -#include -#include - -#include -#include -#include -#include -#include -#include -#include -#include -#include - -namespace tensorrt_llm::testing -{ - -struct TrivialConstantDecoderWithTopKLogitsTestParameters -{ - using TupleT = std::tuple; - runtime::SizeType32 randomSeed; - runtime::SizeType32 vocabSize; - runtime::SizeType32 maxNumTokens; - runtime::SizeType32 maxBeamWidth; - runtime::SizeType32 maxBatchSize; - runtime::SizeType32 numTopKLogits; - runtime::SizeType32 numRequests; - runtime::SizeType32 promptLength; - runtime::SizeType32 maxOutputLength; - bool gatherContext; - - // Constructor that takes a tuple - TrivialConstantDecoderWithTopKLogitsTestParameters( // NOLINT: implicit to allow gtest to convert from tuple - // generated by 'combine' - TupleT t) - : randomSeed(std::get<0>(t)) - , vocabSize(std::get<1>(t)) - , maxNumTokens(std::get<2>(t)) - , maxBeamWidth(std::get<3>(t)) - , maxBatchSize(std::get<4>(t)) - , numTopKLogits(std::get<5>(t)) - , numRequests(std::get<6>(t)) - , promptLength(std::get<7>(t)) - , maxOutputLength(std::get<8>(t)) - , gatherContext(std::get<9>(t)) - { - } -}; - -template -struct DecoderTestShared -{ - static constexpr runtime::SizeType32 kNumTokensPerBlock = 64; - static constexpr runtime::SizeType32 kKvCacheMaxTokens = 2048 * 8; - static constexpr auto kTopKTensorName = "topKLogits"; - - DecoderTestShared(std::shared_ptr logger, std::mt19937 rng, - std::shared_ptr executor, std::vector randomLogits) - : logger(std::move(logger)) - , rng(rng) - , executor(std::move(executor)) - , randomLogits(std::move(randomLogits)){}; - std::shared_ptr logger; - std::mt19937 rng; - std::shared_ptr executor; - std::vector randomLogits; -}; - -template -std::unique_ptr> SetupDecoderTest( - TrivialConstantDecoderWithTopKLogitsTestParameters const& params) -{ - auto logger = std::make_shared(); - auto rng = std::mt19937(params.randomSeed); - auto randomLogits = tensorrt_llm::testing::randomLogits(params.vocabSize, &rng); - auto const decoderParameters = tensorrt_llm::testing::utils::engines::ConstantTrivialDecoderParameters{ - tensorrt_llm::testing::utils::engines::TrivialDecoderParameters{params.vocabSize, params.maxBatchSize, - params.maxNumTokens, DecoderTestShared::kNumTokensPerBlock, params.maxBeamWidth, - params.gatherContext}, - randomLogits}; - auto engineHostMemory = tensorrt_llm::testing::utils::engines::createConstantTrivialDecoderWithTopKLogits( - decoderParameters, params.numTopKLogits, DecoderTestShared::kTopKTensorName, logger); - auto const engine = runtime::RawEngine(engineHostMemory.release()); - - auto const dtype = runtime::TRTDataType::value; - auto modelConfig = runtime::ModelConfig(params.vocabSize, 1, 1, 0, 1, 1, dtype); - modelConfig.useGptAttentionPlugin(true); - modelConfig.setModelVariant(runtime::ModelConfig::ModelVariant::kGpt); - modelConfig.usePackedInput(true); - modelConfig.setKVCacheType(runtime::ModelConfig::KVCacheType::kPAGED); - modelConfig.setMaxNumTokens(params.maxNumTokens); - modelConfig.setMaxBatchSize(params.maxBatchSize); - modelConfig.setMaxBeamWidth(params.maxBeamWidth); - modelConfig.setMaxSequenceLen(params.maxNumTokens); - modelConfig.setMaxInputLen(params.maxNumTokens); - modelConfig.setLayerTypes({runtime::ModelConfig::LayerType::kATTENTION}); - modelConfig.setTokensPerBlock(DecoderTestShared::kNumTokensPerBlock); - modelConfig.setPagedContextFMHA(true); - modelConfig.computeContextLogits(params.gatherContext); - - auto const worldConfig = runtime::WorldConfig(); - - auto kvCacheConfig = executor::KvCacheConfig{}; - kvCacheConfig.setMaxTokens(DecoderTestShared::kKvCacheMaxTokens); - - auto const executorConfig - = executor::ExecutorConfig(params.maxBeamWidth, executor::SchedulerConfig(), kvCacheConfig, true, true, 1, 1, - executor::BatchingType::kINFLIGHT, params.maxBatchSize, params.maxNumTokens, std::nullopt, std::nullopt, - std::nullopt, std::nullopt, false, 1, std::nullopt, executor::ExtendedRuntimePerfKnobConfig(), std::nullopt, - 0, executor::ExecutorConfig::kDefaultMaxSeqIdleMicroseconds, std::nullopt, std::nullopt, - std::vector{ - executor::AdditionalModelOutput{DecoderTestShared::kTopKTensorName, params.gatherContext}}); - - auto model = std::make_shared( - logger, modelConfig, worldConfig, engine, false, executorConfig, false); - - return std::make_unique>( - logger, rng, std::make_shared(model, executorConfig), randomLogits); -} - -template -class DecoderTopKGenerationLogitsTest - : public ::testing::Test, - public ::testing::WithParamInterface -{ -protected: - std::unique_ptr> state; - - DecoderTopKGenerationLogitsTest() - { - auto const params = GetParam(); - state = SetupDecoderTest(params); - } - - void runTopKGenerationLogitsTest(TrivialConstantDecoderWithTopKLogitsTestParameters const& parameters) - { - auto const requestTokens = createConsecutiveTokenSequence(parameters.promptLength, parameters.vocabSize, 0); - auto requests = std::vector{}; - requests.reserve(static_cast(parameters.numRequests)); - for (auto i = 0; i < parameters.numRequests; i++) - { - std::vector additionalOutputs{ - executor::AdditionalModelOutput{DecoderTestShared::kTopKTensorName}}; - requests.emplace_back(requestTokens, parameters.maxOutputLength, false, executor::SamplingConfig{}, - executor::OutputConfig{false, false, false, true, false, false, additionalOutputs}); - } - auto const accumulatedResponses - = runThroughRequests(*state->executor, requests, std::chrono::duration(100000)); - ASSERT_EQ(accumulatedResponses.size(), parameters.numRequests); - - std::sort(state->randomLogits.begin(), state->randomLogits.end()); - std::reverse(state->randomLogits.begin(), state->randomLogits.end()); - for (auto const& [requestId, responses] : accumulatedResponses) - { - for (auto const& response : responses) - { - ASSERT_FALSE(response.hasError()); - auto const& tokensByBeam = response.getResult().outputTokenIds; - auto const& additionalOutputs = response.getResult().additionalOutputs; - ASSERT_EQ(additionalOutputs.size(), 1); - auto const& topKLogits = additionalOutputs.front(); - auto const expectedOutputSize = parameters.maxOutputLength * parameters.numTopKLogits; - ASSERT_EQ(topKLogits.output.getSize(), expectedOutputSize); - auto const* topKLogitsData = reinterpret_cast(topKLogits.output.getData()); - for (auto i = 0; i < parameters.numTopKLogits; i++) - { - EXPECT_TRUE(almostEqual(topKLogitsData[i], state->randomLogits[i], 1e-5)) - << "requestId " << requestId << " i " << i << ": " << topKLogitsData[i] - << " != " << state->randomLogits[i]; - } - ASSERT_EQ(tokensByBeam.size(), 1); - for (auto const& tokensForBeam : tokensByBeam) - { - ASSERT_EQ(tokensForBeam.size(), parameters.maxOutputLength); - } - } - } - } -}; - -template -class DecoderTopKGenerationLogitsStreamingTest - : public ::testing::Test, - public ::testing::WithParamInterface -{ -protected: - std::unique_ptr> state; - - DecoderTopKGenerationLogitsStreamingTest() - { - auto const params = GetParam(); - state = SetupDecoderTest(params); - } - - void runTopKGenerationLogitsStreamingTest(TrivialConstantDecoderWithTopKLogitsTestParameters const& parameters) - { - auto const requestTokens = createConsecutiveTokenSequence(parameters.promptLength, parameters.vocabSize, 0); - auto requests = std::vector{}; - requests.reserve(static_cast(parameters.numRequests)); - for (auto i = 0; i < parameters.numRequests; i++) - { - std::vector additionalOutputs{ - executor::AdditionalModelOutput{DecoderTestShared::kTopKTensorName}}; - requests.emplace_back(requestTokens, parameters.maxOutputLength, true, executor::SamplingConfig{}, - executor::OutputConfig{false, false, false, true, false, false, additionalOutputs}); - } - auto const accumulatedResponses - = runThroughRequests(*state->executor, requests, std::chrono::duration(100000)); - ASSERT_EQ(accumulatedResponses.size(), parameters.numRequests); - - std::sort(state->randomLogits.begin(), state->randomLogits.end()); - std::reverse(state->randomLogits.begin(), state->randomLogits.end()); - for (auto const& idResponsesKvp : accumulatedResponses) - { - auto const& [requestId, responses] = idResponsesKvp; - auto numTokensForRequest = 0; - for (auto const& response : responses) - { - ASSERT_FALSE(response.hasError()); - auto const& tokensByBeam = response.getResult().outputTokenIds; - auto const& additionalOutputs = response.getResult().additionalOutputs; - ASSERT_EQ(additionalOutputs.size(), 1); - auto const& topKLogits = additionalOutputs.front(); - auto const expectedOutputSize = parameters.maxOutputLength * parameters.numTopKLogits; - ASSERT_EQ(topKLogits.output.getSize(), expectedOutputSize); - auto const* topKLogitsData = reinterpret_cast(topKLogits.output.getData()); - for (auto i = 0; i < parameters.numTopKLogits; i++) - { - EXPECT_TRUE(almostEqual(topKLogitsData[i], state->randomLogits[i], 1e-5)) - << "requestId " << requestId << " i " << i << ": " << topKLogitsData[i] - << " != " << state->randomLogits[i]; - } - ASSERT_EQ(tokensByBeam.size(), 1); - for (auto const& tokensForBeam : tokensByBeam) - { - numTokensForRequest += tokensForBeam.size(); - } - } - ASSERT_EQ(numTokensForRequest, parameters.maxOutputLength); - } - } -}; - -template -class DecoderTopKContextLogitsStreamingTest - : public ::testing::Test, - public ::testing::WithParamInterface -{ -protected: - std::unique_ptr> state; - - DecoderTopKContextLogitsStreamingTest() - { - auto const params = GetParam(); - state = SetupDecoderTest(params); - } - - void runTopKContextLogitsTest(TrivialConstantDecoderWithTopKLogitsTestParameters const& parameters) - { - auto requests = std::vector{}; - requests.reserve(static_cast(parameters.numRequests)); - for (auto i = 0; i < parameters.numRequests; i++) - { - // create different sequence for each request to avoid KV cache reuse - auto const requestTokens = createConsecutiveTokenSequence(parameters.promptLength, parameters.vocabSize, i); - std::vector additionalOutputs{ - executor::AdditionalModelOutput{DecoderTestShared::kTopKTensorName, true}}; - requests.emplace_back(requestTokens, parameters.maxOutputLength, true, executor::SamplingConfig{}, - executor::OutputConfig{false, false, false, true, false, false, additionalOutputs}); - } - auto const& accumulatedResponses - = runThroughRequests(*state->executor, requests, std::chrono::duration(100000)); - ASSERT_EQ(accumulatedResponses.size(), parameters.numRequests); - - std::sort(state->randomLogits.begin(), state->randomLogits.end()); - std::reverse(state->randomLogits.begin(), state->randomLogits.end()); - std::string const expectedAdditionalOutputName - = std::string("context_") + DecoderTestShared::kTopKTensorName; - for (auto const& idResponsesKvp : accumulatedResponses) - { - auto const& [requestId, responses] = idResponsesKvp; - std::size_t numTokensForRequest{0}; - for (auto const& response : responses) - { - ASSERT_FALSE(response.hasError()); - auto const& tokensByBeam = response.getResult().outputTokenIds; - auto const& additionalOutputs = response.getResult().additionalOutputs; - ASSERT_EQ(additionalOutputs.size(), 2); - auto const contextTopKLogitsPtr = std::find_if(additionalOutputs.cbegin(), additionalOutputs.cend(), - [&expectedAdditionalOutputName](auto const& ao) - { return ao.name == expectedAdditionalOutputName; }); - auto const expectedOutputSize = parameters.promptLength * parameters.numTopKLogits; - ASSERT_EQ(contextTopKLogitsPtr->output.getSize(), expectedOutputSize); - auto const* topKLogitsData = reinterpret_cast(contextTopKLogitsPtr->output.getData()); - for (auto i = 0; i < parameters.numTopKLogits; i++) - { - EXPECT_TRUE(almostEqual(topKLogitsData[i], state->randomLogits[i], 1e-5)) - << "requestId " << requestId << " i " << i << ": " << topKLogitsData[i] - << " != " << state->randomLogits[i]; - } - ASSERT_EQ(tokensByBeam.size(), 1); - for (auto const& tokensForBeam : tokensByBeam) - { - numTokensForRequest += static_cast(tokensForBeam.size()); - } - } - ASSERT_EQ(numTokensForRequest, parameters.maxOutputLength); - } - } -}; - -template -class DecoderTopKContextLogitsTest - : public ::testing::Test, - public ::testing::WithParamInterface -{ -protected: - std::unique_ptr> state; - - DecoderTopKContextLogitsTest() - { - auto const params = GetParam(); - state = SetupDecoderTest(params); - } - - void runTopKContextLogitsTest(TrivialConstantDecoderWithTopKLogitsTestParameters const& parameters) - { - auto requests = std::vector{}; - requests.reserve(static_cast(parameters.numRequests)); - for (auto i = 0; i < parameters.numRequests; i++) - { - // create different sequence for each request to avoid KV cache reuse - auto const requestTokens = createConsecutiveTokenSequence(parameters.promptLength, parameters.vocabSize, i); - std::vector additionalOutputs{ - executor::AdditionalModelOutput{DecoderTestShared::kTopKTensorName, true}}; - requests.emplace_back(requestTokens, parameters.maxOutputLength, false, executor::SamplingConfig{}, - executor::OutputConfig{false, false, false, true, false, false, additionalOutputs}); - } - auto const accumulatedResponses - = runThroughRequests(*state->executor, requests, std::chrono::duration(100000)); - ASSERT_EQ(accumulatedResponses.size(), parameters.numRequests); - - std::sort(state->randomLogits.begin(), state->randomLogits.end()); - std::reverse(state->randomLogits.begin(), state->randomLogits.end()); - std::string const expectedAdditionalOutputName - = std::string("context_") + DecoderTestShared::kTopKTensorName; - for (auto const& idResponsesKvp : accumulatedResponses) - { - auto const& [requestId, responses] = idResponsesKvp; - for (auto const& response : responses) - { - ASSERT_FALSE(response.hasError()); - auto const& tokensByBeam = response.getResult().outputTokenIds; - auto const& additionalOutputs = response.getResult().additionalOutputs; - ASSERT_EQ(additionalOutputs.size(), 2); - auto const contextTopKLogitsPtr = std::find_if(additionalOutputs.cbegin(), additionalOutputs.cend(), - [&expectedAdditionalOutputName](auto const& ao) - { return ao.name == expectedAdditionalOutputName; }); - auto const expectedOutputSize = parameters.promptLength * parameters.numTopKLogits; - ASSERT_EQ(contextTopKLogitsPtr->output.getSize(), expectedOutputSize); - auto const* topKLogitsData = reinterpret_cast(contextTopKLogitsPtr->output.getData()); - for (auto i = 0; i < parameters.numTopKLogits; i++) - { - EXPECT_TRUE(almostEqual(topKLogitsData[i], state->randomLogits[i], 1e-5)) - << "requestId " << requestId << " i " << i << ": " << topKLogitsData[i] - << " != " << state->randomLogits[i]; - } - ASSERT_EQ(tokensByBeam.size(), 1); - for (auto const& tokensForBeam : tokensByBeam) - { - ASSERT_EQ(tokensForBeam.size(), parameters.maxOutputLength); - } - } - } - } -}; - -namespace -{ -constexpr runtime::SizeType32 kRandomSeed1 = 45; -auto const randomSeeds = ::testing::Values(kRandomSeed1); - -constexpr runtime::SizeType32 kMinVocabSize = 64; -constexpr runtime::SizeType32 kMaxVocabSize = 2048; -auto const vocabSizes = ::testing::Values(kMinVocabSize); - -constexpr runtime::SizeType32 kMinMaxNumTokens = 2048; -auto const maxNumTokenses = ::testing::Values(kMinMaxNumTokens); - -constexpr runtime::SizeType32 kMinBeamWidth = 1; -auto const beamWidths = ::testing::Values(kMinBeamWidth); - -constexpr runtime::SizeType32 kMinMaxBatchSize = 2048; -auto const batchSizes = ::testing::Values(kMinMaxBatchSize); - -constexpr runtime::SizeType32 kMinNumTopKLogits = 4; -constexpr runtime::SizeType32 kMaxNumTopKLogits = 32; -auto const numTopKLogitses = ::testing::Values(kMinNumTopKLogits, kMaxNumTopKLogits); - -constexpr runtime::SizeType32 kMinNumRequests = 16; -constexpr runtime::SizeType32 kMaxNumRequests = 2048; -auto const numRequestses = ::testing::Values(kMinNumRequests); - -constexpr runtime::SizeType32 kMinPromptLength = 4; -constexpr runtime::SizeType32 kMaxPromptLength = 512; -auto const promptLengths = ::testing::Values(kMinPromptLength, kMaxPromptLength); - -constexpr runtime::SizeType32 kMinMaxOutputLength = 4; -constexpr runtime::SizeType32 kMaxMaxOutputLength = 256; -auto const maxOutputLengths = ::testing::Values(kMinMaxOutputLength, kMaxMaxOutputLength); - -auto const gatherContext = ::testing::Values(false, true); -auto const alwaysGatherContext = ::testing::Values(true); - -auto const paramGenerator = ::testing::ConvertGenerator( - ::testing::Combine(randomSeeds, vocabSizes, maxNumTokenses, beamWidths, batchSizes, numTopKLogitses, numRequestses, - promptLengths, maxOutputLengths, gatherContext)); - -auto const paramGeneratorGatherContext - = ::testing::ConvertGenerator( - ::testing::Combine(randomSeeds, vocabSizes, maxNumTokenses, beamWidths, batchSizes, numTopKLogitses, - numRequestses, promptLengths, maxOutputLengths, alwaysGatherContext)); - -auto const nameSuffixGenerator - = [](::testing::TestParamInfo const& info) -> std::string -{ - std::stringstream nameStringStream; - nameStringStream << "gatherContext_" << info.param.gatherContext << "_maxBatchSize_" << info.param.maxBatchSize - << "_vocabSize_" << info.param.vocabSize << "_maxBeamWidth_" << info.param.maxBeamWidth - << "_maxNumTokens_" << info.param.maxNumTokens << "_maxOutputLength_" << info.param.maxOutputLength - << "_numRequests_" << info.param.numRequests << "_numTopKLogits_" << info.param.numTopKLogits - << "_promptLength_" << info.param.promptLength << "_randomSeed_" << info.param.randomSeed; - return nameStringStream.str(); -}; - -} // namespace - -using DecoderTopKGenerationLogitsFloatTest = DecoderTopKGenerationLogitsTest; - -TEST_P(DecoderTopKGenerationLogitsFloatTest, TestSizeAndValues) -{ - runTopKGenerationLogitsTest(GetParam()); -} - -INSTANTIATE_TEST_SUITE_P(Float, DecoderTopKGenerationLogitsFloatTest, paramGenerator, nameSuffixGenerator); - -using DecoderTopKGenerationLogitsStreamingFloatTest = DecoderTopKGenerationLogitsStreamingTest; - -TEST_P(DecoderTopKGenerationLogitsStreamingFloatTest, TestSizeAndValues) -{ - runTopKGenerationLogitsStreamingTest(GetParam()); -} - -INSTANTIATE_TEST_SUITE_P(Float, DecoderTopKGenerationLogitsStreamingFloatTest, paramGenerator, nameSuffixGenerator); - -using DecoderTopKContextLogitsStreamingFloatTest = DecoderTopKContextLogitsStreamingTest; - -TEST_P(DecoderTopKContextLogitsStreamingFloatTest, TestSizeAndValues) -{ - runTopKContextLogitsTest(GetParam()); -} - -INSTANTIATE_TEST_SUITE_P( - Float, DecoderTopKContextLogitsStreamingFloatTest, paramGeneratorGatherContext, nameSuffixGenerator); - -using DecoderTopKContextLogitsFloatTest = DecoderTopKContextLogitsTest; - -TEST_P(DecoderTopKContextLogitsFloatTest, TestSizeAndValues) -{ - runTopKContextLogitsTest(GetParam()); -} - -INSTANTIATE_TEST_SUITE_P(Float, DecoderTopKContextLogitsFloatTest, paramGeneratorGatherContext, nameSuffixGenerator); - -} // namespace tensorrt_llm::testing diff --git a/cpp/tests/unit_tests/executor/intervalSetTest.cpp b/cpp/tests/unit_tests/executor/intervalSetTest.cpp deleted file mode 100644 index a2bb0a8f7532..000000000000 --- a/cpp/tests/unit_tests/executor/intervalSetTest.cpp +++ /dev/null @@ -1,224 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2023-2025 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. - */ - -#include "tensorrt_llm/executor/intervalSet.h" - -#include - -using tensorrt_llm::executor::IntervalSet; -using tensorrt_llm::executor::IdType; - -class IntervalSetTest : public ::testing::Test // NOLINT(cppcoreguidelines-pro-type-member-init) -{ -protected: - void SetUp() override {} - - void TearDown() override {} - - IntervalSet mIntervalSet; -}; - -namespace -{ - -TEST_F(IntervalSetTest, testPublicAPI) -{ - EXPECT_FALSE(mIntervalSet.contains(0)); - EXPECT_EQ(mIntervalSet.getNumElements(), 0); - mIntervalSet.insert(0); - mIntervalSet.insert(1); - mIntervalSet.insert(4); - mIntervalSet.insert(6); - EXPECT_TRUE(mIntervalSet.contains(0)); - EXPECT_TRUE(mIntervalSet.contains(4)); - EXPECT_FALSE(mIntervalSet.contains(2125)); - EXPECT_EQ(mIntervalSet.getNumElements(), 4); - mIntervalSet.insert(6); - mIntervalSet.insert(4); - mIntervalSet.insert(1); - mIntervalSet.insert(0); - EXPECT_EQ(mIntervalSet.getNumElements(), 4); - EXPECT_TRUE(mIntervalSet.contains(0)); - EXPECT_TRUE(mIntervalSet.contains(4)); - EXPECT_FALSE(mIntervalSet.contains(9)); - EXPECT_FALSE(mIntervalSet.contains(3)); - EXPECT_FALSE(mIntervalSet.contains(11)); - mIntervalSet.clear(); - EXPECT_EQ(mIntervalSet.getNumElements(), 0); -} - -TEST_F(IntervalSetTest, testClear) -{ - for (int i = 0; i < 100; i++) - { - if (i % 2 == 0) - { - EXPECT_FALSE(mIntervalSet.contains(i)); - mIntervalSet.insert(i); - EXPECT_TRUE(mIntervalSet.contains(i)); - mIntervalSet.clear(); - EXPECT_EQ(mIntervalSet.getNumElements(), 0); - } - EXPECT_FALSE(mIntervalSet.contains(i)); - } - for (int i = 0; i < 100; i++) - { - EXPECT_FALSE(mIntervalSet.contains(i)); - } - mIntervalSet.clear(); - EXPECT_EQ(mIntervalSet.getNumElements(), 0); - mIntervalSet.clear(); - EXPECT_EQ(mIntervalSet.getNumElements(), 0); - - for (int i = 19; i >= 10; i--) - { - if (i % 2 == 0) - { - EXPECT_FALSE(mIntervalSet.contains(i)); - mIntervalSet.insert(i); - EXPECT_TRUE(mIntervalSet.contains(i)); - } - else - { - EXPECT_FALSE(mIntervalSet.contains(i)); - } - } - EXPECT_EQ(mIntervalSet.getNumElements(), 5); - mIntervalSet.clear(); - EXPECT_EQ(mIntervalSet.getNumElements(), 0); - mIntervalSet.clear(); - EXPECT_EQ(mIntervalSet.getNumElements(), 0); -} - -TEST_F(IntervalSetTest, testRandomInsert) -{ - mIntervalSet.insert(4); - mIntervalSet.insert(8); - EXPECT_EQ(mIntervalSet.getNumElements(), 2); - std::vector idToAdd{9, 7, 5, 1, 6, 0, 2}; - for (auto id : idToAdd) - { - mIntervalSet.insert(id); - } - for (int i = 0; i < 10; i++) - { - if (i != 3) - { - EXPECT_TRUE(mIntervalSet.contains(i)); - } - else - { - EXPECT_FALSE(mIntervalSet.contains(i)); - } - } - EXPECT_EQ(mIntervalSet.getNumElements(), 9); - mIntervalSet.insert(3); - for (int i = 0; i < 10; i++) - { - EXPECT_TRUE(mIntervalSet.contains(i)); - } - EXPECT_EQ(mIntervalSet.getNumElements(), 10); - mIntervalSet.clear(); - EXPECT_EQ(mIntervalSet.getNumElements(), 0); -} - -TEST_F(IntervalSetTest, testTerminatedReqIdIntervals) -{ - mIntervalSet.insert(4); - mIntervalSet.insert(8); - EXPECT_EQ(mIntervalSet.getNumElements(), 2); - // terminatedReqIdIntervals is [[4, 4], [8, 8]] - EXPECT_EQ(mIntervalSet.getIntervals().size(), 2); - mIntervalSet.insert(3); - mIntervalSet.insert(5); - // terminatedReqIdIntervals is [[3, 5], [8, 8]] - EXPECT_EQ(mIntervalSet.getNumElements(), 4); - EXPECT_EQ(mIntervalSet.getIntervals().size(), 2); - mIntervalSet.insert(9); - mIntervalSet.insert(7); - // terminatedReqIdIntervals is [[3, 5], [7, 9]] - EXPECT_EQ(mIntervalSet.getNumElements(), 6); - EXPECT_EQ(mIntervalSet.getIntervals().size(), 2); - mIntervalSet.insert(6); - // terminatedReqIdIntervals is [[3, 9]] - EXPECT_EQ(mIntervalSet.getNumElements(), 7); - EXPECT_EQ(mIntervalSet.getIntervals().size(), 1); - mIntervalSet.insert(1); - // terminatedReqIdIntervals is [[1, 1], [3, 9]] - EXPECT_EQ(mIntervalSet.getNumElements(), 8); - EXPECT_EQ(mIntervalSet.getIntervals().size(), 2); - mIntervalSet.insert(0); - // terminatedReqIdIntervals is [[0, 1], [3, 9]] - EXPECT_EQ(mIntervalSet.getNumElements(), 9); - EXPECT_EQ(mIntervalSet.getIntervals().size(), 2); - mIntervalSet.insert(2); - // terminatedReqIdIntervals is [[0, 9]] - EXPECT_EQ(mIntervalSet.getNumElements(), 10); - EXPECT_EQ(mIntervalSet.getIntervals().size(), 1); - for (int i = 0; i < 10; i++) - { - mIntervalSet.insert(i); - // terminatedReqIdIntervals is always [[0, 9]] - EXPECT_EQ(mIntervalSet.getNumElements(), 10); - EXPECT_EQ(mIntervalSet.getIntervals().size(), 1); - } - mIntervalSet.clear(); - EXPECT_EQ(mIntervalSet.getNumElements(), 0); - EXPECT_EQ(mIntervalSet.getIntervals().size(), 0); - - // Insert continuous decreasing numbers. Interval size is always one. - for (int i = 19; i >= 10; i--) - { - mIntervalSet.insert(i); - EXPECT_TRUE(mIntervalSet.contains(i)); - EXPECT_EQ(mIntervalSet.getIntervals().size(), 1); - } - mIntervalSet.clear(); - EXPECT_EQ(mIntervalSet.getNumElements(), 0); - EXPECT_EQ(mIntervalSet.getIntervals().size(), 0); - - // Insert 50 disjoint even numbers - for (int i = 0; i < 100; i++) - { - if (i % 2 == 0) - { - mIntervalSet.insert(i); - EXPECT_EQ(mIntervalSet.getNumElements(), (i / 2) + 1); - EXPECT_EQ(mIntervalSet.getIntervals().size(), (i / 2) + 1); - } - } - - // Insert 50 disjoint odd numbers. Interval size should go down as the intervals are merged. - for (int i = 0; i < 100; i++) - { - if (i % 2 != 0) - { - mIntervalSet.insert(i); - EXPECT_EQ(mIntervalSet.getNumElements(), 50 + (i + 1) / 2); - if (i != 99) - { - EXPECT_EQ(mIntervalSet.getIntervals().size(), 50 - (i + 1) / 2); - } - else - { - EXPECT_EQ(mIntervalSet.getIntervals().size(), 1); - } - } - } -} - -} // namespace diff --git a/cpp/tests/unit_tests/executor/requestTest.cpp b/cpp/tests/unit_tests/executor/requestTest.cpp index f44f55b5141d..ec43d088657f 100644 --- a/cpp/tests/unit_tests/executor/requestTest.cpp +++ b/cpp/tests/unit_tests/executor/requestTest.cpp @@ -26,7 +26,9 @@ using ::testing::_; using ::testing::Invoke; using namespace tensorrt_llm::executor; -using namespace tensorrt_llm::common; +// Not a namespace-wide import: common also exports DataType/Dims, which would +// make the unqualified DataType (= executor::DataType) below ambiguous. +using tensorrt_llm::common::TllmException; TEST(RequestTest, validInputs) { diff --git a/cpp/tests/unit_tests/executor/serializeUtilsTest.cpp b/cpp/tests/unit_tests/executor/serializeUtilsTest.cpp index a39756bf7243..3c78659c87c4 100644 --- a/cpp/tests/unit_tests/executor/serializeUtilsTest.cpp +++ b/cpp/tests/unit_tests/executor/serializeUtilsTest.cpp @@ -18,6 +18,7 @@ #include "tensorrt_llm/executor/serializeUtils.h" #include "tensorrt_llm/batch_manager/kvCacheManager.h" #include "tensorrt_llm/common/logger.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/executor/cache_transmission/agent_utils/connection.h" #include "tensorrt_llm/executor/dataTransceiverState.h" #include "tensorrt_llm/executor/executor.h" @@ -751,7 +752,8 @@ TEST(SerializeUtilsTest, ContextPhaseParams) { auto state = std::make_unique(); state->setCommState(texec::kv_cache::CommState{12, "127.0.0.1"}); - state->setCacheState(texec::kv_cache::CacheState{10, 12, 128, 128, 8, 8, 8, {4}, nvinfer1::DataType::kFLOAT}); + state->setCacheState( + texec::kv_cache::CacheState{10, 12, 128, 128, 8, 8, 8, {4}, tensorrt_llm::DataType::kFLOAT}); auto stats = texec::ContextPhaseParams({10, 20, 30, 40, 50, 60}, 0, state.release(), VecTokens{10, 20}); auto stats2 = serializeDeserialize(stats); EXPECT_EQ(stats, stats2); @@ -1553,7 +1555,7 @@ TEST(SerializeUtilsTest, CacheStateIndexerKCache) texec::SizeType32 pp = 1; texec::SizeType32 cp = 1; std::vector attentionLayerNumPerPP{static_cast(nbKvHeadsPerLayer.size())}; - auto dataType = nvinfer1::DataType::kFLOAT; + auto dataType = tensorrt_llm::DataType::kFLOAT; auto attentionType = CacheState::AttentionType::kDEFAULT; int kvFactor = 2; bool enableAttentionDP = false; diff --git a/cpp/tests/unit_tests/executor/ucxCommTest.cpp b/cpp/tests/unit_tests/executor/ucxCommTest.cpp index 51d1d84d2676..5d2bfc6e772f 100644 --- a/cpp/tests/unit_tests/executor/ucxCommTest.cpp +++ b/cpp/tests/unit_tests/executor/ucxCommTest.cpp @@ -36,6 +36,7 @@ #include "tensorrt_llm/common/assert.h" #include "tensorrt_llm/common/cudaUtils.h" #include "tensorrt_llm/common/envUtils.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/executor/cache_transmission/mpi_utils/connection.h" #include "tensorrt_llm/executor/dataTransceiverState.h" #include "tensorrt_llm/executor/executor.h" @@ -132,11 +133,11 @@ TEST_F(UcxCommTest, Basic) tensorrt_llm::runtime::BufferManager bufferManager{std::make_shared()}; // Create and fill source CUDA buffer with random data - auto srcBuffer = bufferManager.gpu(buffer.size(), nvinfer1::DataType::kINT8); + auto srcBuffer = bufferManager.gpu(buffer.size(), tensorrt_llm::DataType::kINT8); bufferManager.copy(buffer.data(), *srcBuffer); bufferManager.getStream().synchronize(); - auto dstBuffer = bufferManager.gpu(buffer.size(), nvinfer1::DataType::kINT8); + auto dstBuffer = bufferManager.gpu(buffer.size(), tensorrt_llm::DataType::kINT8); // Send CUDA buffer using connection1 connection1->send(DataContext{0x75}, srcBuffer->data(), srcBuffer->getSizeInBytes()); @@ -204,14 +205,14 @@ TEST_F(UcxCommTest, multiSend) tensorrt_llm::runtime::BufferManager bufferManager{std::make_shared()}; - auto srcBuffer1 = bufferManager.gpu(buffer1.size(), nvinfer1::DataType::kINT8); - auto srcBuffer2 = bufferManager.gpu(buffer2.size(), nvinfer1::DataType::kINT8); + auto srcBuffer1 = bufferManager.gpu(buffer1.size(), tensorrt_llm::DataType::kINT8); + auto srcBuffer2 = bufferManager.gpu(buffer2.size(), tensorrt_llm::DataType::kINT8); bufferManager.copy(buffer1.data(), *srcBuffer1); bufferManager.copy(buffer2.data(), *srcBuffer2); bufferManager.getStream().synchronize(); - auto dstBuffer1 = bufferManager.gpu(buffer1.size(), nvinfer1::DataType::kINT8); - auto dstBuffer2 = bufferManager.gpu(buffer2.size(), nvinfer1::DataType::kINT8); + auto dstBuffer1 = bufferManager.gpu(buffer1.size(), tensorrt_llm::DataType::kINT8); + auto dstBuffer2 = bufferManager.gpu(buffer2.size(), tensorrt_llm::DataType::kINT8); connection1Peer->send(DataContext{0x75}, srcBuffer1->data(), srcBuffer1->getSizeInBytes()); connection2Peer->send(DataContext{0x75}, srcBuffer2->data(), srcBuffer2->getSizeInBytes()); diff --git a/cpp/tests/unit_tests/kernels/CMakeLists.txt b/cpp/tests/unit_tests/kernels/CMakeLists.txt index e593ce4b76b8..5f121d2406a9 100644 --- a/cpp/tests/unit_tests/kernels/CMakeLists.txt +++ b/cpp/tests/unit_tests/kernels/CMakeLists.txt @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2023-2025 NVIDIA CORPORATION & +# SPDX-FileCopyrightText: Copyright (c) 2023-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 @@ -35,9 +35,14 @@ if(USING_OSS_CUTLASS_MOE_GEMM) target_compile_definitions(mixtureOfExpertsTest PUBLIC USING_OSS_CUTLASS_MOE_GEMM) - add_gtest(mixtureOfExpertsInternalTest mixtureOfExpertsTest.cu) - remove_compile_definition(mixtureOfExpertsInternalTest - USING_OSS_CUTLASS_MOE_GEMM) + # The internal-path variant includes headers (quantization.h) that only ship + # with the internal cutlass kernels sources; the prebuilt tarball provides the + # library and a reduced header set only. + if(INTERNAL_CUTLASS_KERNELS_PATH) + add_gtest(mixtureOfExpertsInternalTest mixtureOfExpertsTest.cu) + remove_compile_definition(mixtureOfExpertsInternalTest + USING_OSS_CUTLASS_MOE_GEMM) + endif() endif() add_gtest(ropeTest ropeTest.cu) diff --git a/cpp/tests/unit_tests/kernels/banRepeatNGramsKernelsTest.cpp b/cpp/tests/unit_tests/kernels/banRepeatNGramsKernelsTest.cpp index 567cb95e8e44..83e94ea2d6ad 100644 --- a/cpp/tests/unit_tests/kernels/banRepeatNGramsKernelsTest.cpp +++ b/cpp/tests/unit_tests/kernels/banRepeatNGramsKernelsTest.cpp @@ -15,6 +15,7 @@ */ #include "tensorrt_llm/common/memoryUtils.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/kernels/banRepeatNgram.h" #include "tensorrt_llm/runtime/bufferManager.h" #include "tensorrt_llm/runtime/runtimeKernels.h" @@ -52,24 +53,25 @@ class BanRepeatNgramKernelsTest : public testing::Test SizeType32 const batchSize = outputIds.size(); auto const maxBatchSize = 2 * batchSize; - mLogits = BufferManager::pinned(ITensor::makeShape({batchSize, mVocabSizePadded}), nvinfer1::DataType::kFLOAT); + mLogits + = BufferManager::pinned(ITensor::makeShape({batchSize, mVocabSizePadded}), tensorrt_llm::DataType::kFLOAT); mSequenceLengths - = BufferManager::pinned(ITensor::makeShape({maxBatchSize, mBeamWidth}), nvinfer1::DataType::kINT32); + = BufferManager::pinned(ITensor::makeShape({maxBatchSize, mBeamWidth}), tensorrt_llm::DataType::kINT32); mFinished = BufferManager::pinned( ITensor::makeShape({maxBatchSize, mBeamWidth}), TRTDataType::value); mOutputIds = BufferManager::pinned( - ITensor::makeShape({maxBatchSize, mBeamWidth, mMaxSeqLen}), nvinfer1::DataType::kINT32); + ITensor::makeShape({maxBatchSize, mBeamWidth, mMaxSeqLen}), tensorrt_llm::DataType::kINT32); mOutputIdsPtr = BufferManager::pinned(ITensor::makeShape({maxBatchSize, mBeamWidth}), ptrType); mParentIds = BufferManager::pinned( - ITensor::makeShape({maxBatchSize, mBeamWidth, mMaxSeqLen}), nvinfer1::DataType::kINT32); + ITensor::makeShape({maxBatchSize, mBeamWidth, mMaxSeqLen}), tensorrt_llm::DataType::kINT32); mParentIdsPtr = BufferManager::pinned(ITensor::makeShape({maxBatchSize, mBeamWidth}), ptrType); - mNGramSizes = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kINT32); + mNGramSizes = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kINT32); - mBatchSlots = BufferManager::pinned(ITensor::makeShape({batchSize}), nvinfer1::DataType::kINT32); + mBatchSlots = BufferManager::pinned(ITensor::makeShape({batchSize}), tensorrt_llm::DataType::kINT32); auto batchSlotsPtr = bufferCast(*mBatchSlots); for (SizeType32 bi = 0; bi < batchSize; ++bi) diff --git a/cpp/tests/unit_tests/kernels/cudaCoreGemm/cudaCoreGemmKernelTest.cpp b/cpp/tests/unit_tests/kernels/cudaCoreGemm/cudaCoreGemmKernelTest.cpp index 05247e2d27a4..f7dabf98a93e 100644 --- a/cpp/tests/unit_tests/kernels/cudaCoreGemm/cudaCoreGemmKernelTest.cpp +++ b/cpp/tests/unit_tests/kernels/cudaCoreGemm/cudaCoreGemmKernelTest.cpp @@ -1,4 +1,3 @@ -#include #include #include #include diff --git a/cpp/tests/unit_tests/kernels/decodingKernelTest.cpp b/cpp/tests/unit_tests/kernels/decodingKernelTest.cpp index 4b94e67cb5a1..a458f783ef65 100644 --- a/cpp/tests/unit_tests/kernels/decodingKernelTest.cpp +++ b/cpp/tests/unit_tests/kernels/decodingKernelTest.cpp @@ -15,6 +15,7 @@ */ #include "tensorrt_llm/common/memoryUtils.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/kernels/decodingCommon.h" #include "tensorrt_llm/kernels/decodingKernels.h" #include "tensorrt_llm/kernels/speculativeDecoding/externalDraftTokensKernels.h" @@ -216,13 +217,15 @@ class TestBeamHypothesesCopy : public ::testing::Test srcBeams.empty(*mBufferManager); srcBeams.reshape(batchSize, beamWidth, maxSeqLen); - mSrcCumLogProbs = mBufferManager->gpu(ITensor::makeShape({batchSize, beamWidth}), nvinfer1::DataType::kFLOAT); + mSrcCumLogProbs + = mBufferManager->gpu(ITensor::makeShape({batchSize, beamWidth}), tensorrt_llm::DataType::kFLOAT); setBuffers(srcBeams, mSrcCumLogProbs, 2); dstBeams.empty(*mBufferManager); dstBeams.reshape(batchSize, beamWidth, maxSeqLen); - mDstCumLogProbs = mBufferManager->gpu(ITensor::makeShape({batchSize, beamWidth}), nvinfer1::DataType::kFLOAT); + mDstCumLogProbs + = mBufferManager->gpu(ITensor::makeShape({batchSize, beamWidth}), tensorrt_llm::DataType::kFLOAT); setBuffers(dstBeams, mDstCumLogProbs, 1); } @@ -544,7 +547,7 @@ class TestGatherTree : public ::testing::Test SizeType32 constexpr nbRnnLayers{0}; SizeType32 constexpr nbHeads{16}; SizeType32 constexpr hiddenSize{1024}; - nvinfer1::DataType constexpr dtype{nvinfer1::DataType::kFLOAT}; + tensorrt_llm::DataType constexpr dtype{tensorrt_llm::DataType::kFLOAT}; ModelConfig modelConfig{ vocabSize, nbAttentionLayers + nbRnnLayers, nbAttentionLayers, nbRnnLayers, nbHeads, hiddenSize, dtype}; @@ -1139,32 +1142,35 @@ class DecodingKernelsTest : public testing::Test auto const ptrType = TRTDataType::value; mDraftTokens = mBufferManager->pinnedPool( - ITensor::makeShape({mMaxBatchSize, mMaxDraftSeqlen}), nvinfer1::DataType::kINT32); + ITensor::makeShape({mMaxBatchSize, mMaxDraftSeqlen}), tensorrt_llm::DataType::kINT32); mTargetTokens = mBufferManager->pinnedPool( - ITensor::makeShape({mMaxBatchSize, mMaxTargetSeqlen}), nvinfer1::DataType::kINT32); - mOutputTokens - = mBufferManager->pinnedPool(ITensor::makeShape({mMaxBatchSize, mMaxSeqLen}), nvinfer1::DataType::kINT32); + ITensor::makeShape({mMaxBatchSize, mMaxTargetSeqlen}), tensorrt_llm::DataType::kINT32); + mOutputTokens = mBufferManager->pinnedPool( + ITensor::makeShape({mMaxBatchSize, mMaxSeqLen}), tensorrt_llm::DataType::kINT32); mNumsDraftTokens = mBufferManager->pinnedPool( - ITensor::makeShape({mMaxBatchSize, mMaxDraftSeqPerStep}), nvinfer1::DataType::kINT32); - mSequenceLengths = mBufferManager->pinnedPool(ITensor::makeShape({mMaxBatchSize}), nvinfer1::DataType::kINT32); - mAcceptedLengths = mBufferManager->pinnedPool(ITensor::makeShape({mMaxBatchSize}), nvinfer1::DataType::kINT32); - mContextLengths = mBufferManager->pinnedPool(ITensor::makeShape({mMaxBatchSize}), nvinfer1::DataType::kINT32); + ITensor::makeShape({mMaxBatchSize, mMaxDraftSeqPerStep}), tensorrt_llm::DataType::kINT32); + mSequenceLengths + = mBufferManager->pinnedPool(ITensor::makeShape({mMaxBatchSize}), tensorrt_llm::DataType::kINT32); + mAcceptedLengths + = mBufferManager->pinnedPool(ITensor::makeShape({mMaxBatchSize}), tensorrt_llm::DataType::kINT32); + mContextLengths + = mBufferManager->pinnedPool(ITensor::makeShape({mMaxBatchSize}), tensorrt_llm::DataType::kINT32); mFinishedSteps = mBufferManager->pinnedPool(ITensor::makeShape({mMaxDraftTokens + 1, mMaxBatchSize}), TRTDataType::value); mFinishedFinal = mBufferManager->pinnedPool( ITensor::makeShape({mMaxBatchSize}), TRTDataType::value); - mFinishedSum = mBufferManager->pinnedPool(ITensor::makeShape({mMaxBatchSize}), nvinfer1::DataType::kINT32); + mFinishedSum = mBufferManager->pinnedPool(ITensor::makeShape({mMaxBatchSize}), tensorrt_llm::DataType::kINT32); mPaths = mBufferManager->pinnedPool( - ITensor::makeShape({mMaxBatchSize, mMaxDraftSeqPerStep, mMaxDraftTokens}), nvinfer1::DataType::kINT32); - mEndIds = mBufferManager->pinnedPool(ITensor::makeShape({mMaxBatchSize}), nvinfer1::DataType::kINT32); + ITensor::makeShape({mMaxBatchSize, mMaxDraftSeqPerStep, mMaxDraftTokens}), tensorrt_llm::DataType::kINT32); + mEndIds = mBufferManager->pinnedPool(ITensor::makeShape({mMaxBatchSize}), tensorrt_llm::DataType::kINT32); - mBatchSlots = mBufferManager->pinnedPool(ITensor::makeShape({mMaxBatchSize}), nvinfer1::DataType::kINT32); + mBatchSlots = mBufferManager->pinnedPool(ITensor::makeShape({mMaxBatchSize}), tensorrt_llm::DataType::kINT32); auto batchSlotsRange = BufferRange(*mBatchSlots); std::iota(batchSlotsRange.begin(), batchSlotsRange.end(), 0); mCurandStates = mBufferManager->gpu( - ITensor::makeShape({mMaxBatchSize, sizeof(curandState_t)}), nvinfer1::DataType::kINT8); + ITensor::makeShape({mMaxBatchSize, sizeof(curandState_t)}), tensorrt_llm::DataType::kINT8); mAcceptedLen.resize(mMaxBatchSize); mOutputLen.resize(mMaxBatchSize); @@ -1194,8 +1200,9 @@ class DecodingKernelsTest : public testing::Test mMedusaInputLogitsPtrs = mBufferManager->pinnedPool(ITensor::makeShape({mMaxBatchSize, mMaxNumHeads}), ptrType); mTokensPerStep - = mBufferManager->pinnedPool(ITensor::makeShape({mMaxBatchSize}), nvinfer1::DataType::kINT32); - mBestPaths = mBufferManager->pinnedPool(ITensor::makeShape({mMaxBatchSize}), nvinfer1::DataType::kINT32); + = mBufferManager->pinnedPool(ITensor::makeShape({mMaxBatchSize}), tensorrt_llm::DataType::kINT32); + mBestPaths + = mBufferManager->pinnedPool(ITensor::makeShape({mMaxBatchSize}), tensorrt_llm::DataType::kINT32); } } diff --git a/cpp/tests/unit_tests/kernels/eaglePackDataTest.cpp b/cpp/tests/unit_tests/kernels/eaglePackDataTest.cpp index bdf74efb59be..8ce24f813e65 100644 --- a/cpp/tests/unit_tests/kernels/eaglePackDataTest.cpp +++ b/cpp/tests/unit_tests/kernels/eaglePackDataTest.cpp @@ -25,9 +25,8 @@ #include "tensorrt_llm/runtime/iBuffer.h" #include "tensorrt_llm/runtime/iTensor.h" #include "tensorrt_llm/runtime/runtimeKernels.h" -#include "tensorrt_llm/runtime/tllmLogger.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include #include @@ -132,81 +131,81 @@ class EaglePackDataTest : public ::testing::Test { // inputs mBatchSlots = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getBatchSize()}), nvinfer1::DataType::kINT32); + ITensor::makeShape({mSamplingParams.getBatchSize()}), tensorrt_llm::DataType::kINT32); mInputTemperatures = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), nvinfer1::DataType::kFLOAT); + ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), tensorrt_llm::DataType::kFLOAT); mInputRandomDataSample = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), nvinfer1::DataType::kFLOAT); + ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), tensorrt_llm::DataType::kFLOAT); mInputRandomDataValidation = BufferManager::pinnedPool( ITensor::makeShape({mSamplingParams.getMaxBatchSize(), mSamplingParams.getMaxDecodingTokens()}), - nvinfer1::DataType::kFLOAT); + tensorrt_llm::DataType::kFLOAT); mInputNextDraftTokens = BufferManager::pinnedPool( ITensor::makeShape({mSamplingParams.getMaxBatchSize(), mSamplingParams.getMaxDecodingDraftTokens()}), - nvinfer1::DataType::kINT32); + tensorrt_llm::DataType::kINT32); mInputNextDraftPaths = BufferManager::pinnedPool(ITensor::makeShape({mSamplingParams.getMaxBatchSize(), mSamplingParams.getMaxDecodingTokens(), mSamplingParams.getMaxPathLen()}), - nvinfer1::DataType::kINT32); + tensorrt_llm::DataType::kINT32); mInputSpecDecodingGenerationLengths = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), nvinfer1::DataType::kINT32); + ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), tensorrt_llm::DataType::kINT32); mInputSpecDecodingPositionOffsets = BufferManager::pinnedPool( ITensor::makeShape({mSamplingParams.getMaxBatchSize(), mSamplingParams.getMaxDecodingTokens()}), - nvinfer1::DataType::kINT32); + tensorrt_llm::DataType::kINT32); auto const numPackedMasks = static_cast(tensorrt_llm::common::divUp(mSamplingParams.getMaxDecodingTokens(), 32)); mInputSpecDecodingPackedMasks = BufferManager::pinnedPool( ITensor::makeShape( {mSamplingParams.getMaxBatchSize(), mSamplingParams.getMaxDecodingTokens(), numPackedMasks}), - nvinfer1::DataType::kINT32); + tensorrt_llm::DataType::kINT32); // outputs mOutputTemperatures = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getBatchSize()}), nvinfer1::DataType::kFLOAT); + ITensor::makeShape({mSamplingParams.getBatchSize()}), tensorrt_llm::DataType::kFLOAT); mOutputRandomDataSample = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getBatchSize()}), nvinfer1::DataType::kFLOAT); + ITensor::makeShape({mSamplingParams.getBatchSize()}), tensorrt_llm::DataType::kFLOAT); mOutputRandomDataValidation = BufferManager::pinnedPool( ITensor::makeShape({mSamplingParams.getBatchSize(), mSamplingParams.getMaxDecodingTokens()}), - nvinfer1::DataType::kFLOAT); + tensorrt_llm::DataType::kFLOAT); mOutputNextDraftTokens = BufferManager::pinnedPool( ITensor::makeShape({mSamplingParams.getBatchSize(), mSamplingParams.getMaxDecodingDraftTokens()}), - nvinfer1::DataType::kINT32); + tensorrt_llm::DataType::kINT32); mOutputNextDraftLens = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getBatchSize()}), nvinfer1::DataType::kINT32); + ITensor::makeShape({mSamplingParams.getBatchSize()}), tensorrt_llm::DataType::kINT32); mOutputNextDraftPaths = BufferManager::pinnedPool(ITensor::makeShape({mSamplingParams.getBatchSize(), mSamplingParams.getMaxDecodingTokens(), mSamplingParams.getMaxPathLen()}), - nvinfer1::DataType::kINT32); + tensorrt_llm::DataType::kINT32); mOutputSpecDecodingGenerationLengths = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getBatchSize()}), nvinfer1::DataType::kINT32); + ITensor::makeShape({mSamplingParams.getBatchSize()}), tensorrt_llm::DataType::kINT32); mOutputSpecDecodingPositionOffsets = BufferManager::pinnedPool( ITensor::makeShape({mSamplingParams.getBatchSize(), mSamplingParams.getMaxDecodingTokens()}), - nvinfer1::DataType::kINT32); + tensorrt_llm::DataType::kINT32); mOutputSpecDecodingPackedMasks = BufferManager::pinnedPool( ITensor::makeShape( {mSamplingParams.getBatchSize(), mSamplingParams.getMaxDecodingTokens(), numPackedMasks}), - nvinfer1::DataType::kINT32); + tensorrt_llm::DataType::kINT32); // workspace - mMaxGenerationLength = BufferManager::pinnedPool(ITensor::makeShape({1}), nvinfer1::DataType::kINT32); + mMaxGenerationLength = BufferManager::pinnedPool(ITensor::makeShape({1}), tensorrt_llm::DataType::kINT32); mCumSumGenerationLengths = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getBatchSize() + 1}), nvinfer1::DataType::kINT32); + ITensor::makeShape({mSamplingParams.getBatchSize() + 1}), tensorrt_llm::DataType::kINT32); mScanReduceTempStorageBytes = tksd::invokeScanReduceGenerationLengths( mSamplingParams.getBatchSize(), nullptr, nullptr, 0, nullptr, nullptr, mStream->get()); diff --git a/cpp/tests/unit_tests/kernels/mixtureOfExpertsTest.cu b/cpp/tests/unit_tests/kernels/mixtureOfExpertsTest.cu index 01cd1c4d792d..61c182dfd1ea 100644 --- a/cpp/tests/unit_tests/kernels/mixtureOfExpertsTest.cu +++ b/cpp/tests/unit_tests/kernels/mixtureOfExpertsTest.cu @@ -32,6 +32,7 @@ #endif #include "tensorrt_llm/kernels/cutlass_kernels/include/cutlass_kernel_selector.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/runtime/bufferManager.h" #include @@ -2508,31 +2509,31 @@ constexpr static auto typeToDtypeID() { if constexpr (std::is_same_v) { - return nvinfer1::DataType::kFP8; + return tensorrt_llm::DataType::kFP8; } else if constexpr (std::is_same_v) { - return nvinfer1::DataType::kFP4; + return tensorrt_llm::DataType::kFP4; } else if constexpr (std::is_same_v) { - return nvinfer1::DataType::kINT8; + return tensorrt_llm::DataType::kINT8; } else if constexpr (std::is_same_v) { - return nvinfer1::DataType::kINT4; + return tensorrt_llm::DataType::kINT4; } else if constexpr (std::is_same_v) { - return nvinfer1::DataType::kBF16; + return tensorrt_llm::DataType::kBF16; } else if constexpr (std::is_same_v) { - return nvinfer1::DataType::kHALF; + return tensorrt_llm::DataType::kHALF; } else if constexpr (std::is_same_v) { - return nvinfer1::DataType::kFLOAT; + return tensorrt_llm::DataType::kFLOAT; } else { @@ -2602,14 +2603,16 @@ TEST_F(MixtureOfExpertsProfilerTest, TestGeneratedProfilerDistribution) for (int ep : {1, 4, 8}) { #ifdef USING_OSS_CUTLASS_MOE_GEMM - backend.init(this->mMoERunner, GemmProfilerBackend::GemmToProfile::GEMM_1, nvinfer1::DataType::kHALF, - nvinfer1::DataType::kHALF, nvinfer1::DataType::kHALF, num_experts, k, 1024, 1024, 4096, mGroupSize, {}, - false, mUseLora, /*min_latency_mode=*/false, /*need_weights=*/true, MOEParallelismConfig{1, 0, ep, 0}, + backend.init(this->mMoERunner, GemmProfilerBackend::GemmToProfile::GEMM_1, tensorrt_llm::DataType::kHALF, + tensorrt_llm::DataType::kHALF, tensorrt_llm::DataType::kHALF, num_experts, k, 1024, 1024, 4096, + mGroupSize, {}, false, mUseLora, /*min_latency_mode=*/false, /*need_weights=*/true, + MOEParallelismConfig{1, 0, ep, 0}, /*enable_alltoall=*/false); #else - backend.init(this->mMoERunner, GemmProfilerBackend::GemmToProfile::GEMM_1, nvinfer1::DataType::kHALF, - nvinfer1::DataType::kHALF, nvinfer1::DataType::kHALF, num_experts, k, 1024, 4096, mGroupSize, {}, false, - mUseLora, /*min_latency_mode=*/false, /*need_weights=*/true, MOEParallelismConfig{1, 0, ep, ep - 1}); + backend.init(this->mMoERunner, GemmProfilerBackend::GemmToProfile::GEMM_1, tensorrt_llm::DataType::kHALF, + tensorrt_llm::DataType::kHALF, tensorrt_llm::DataType::kHALF, num_experts, k, 1024, 4096, mGroupSize, + {}, false, mUseLora, /*min_latency_mode=*/false, /*need_weights=*/true, + MOEParallelismConfig{1, 0, ep, ep - 1}); #endif auto ws_size = backend.getWorkspaceSize(num_tokens); diff --git a/cpp/tests/unit_tests/kernels/mlaChunkedPrefillTest.cu b/cpp/tests/unit_tests/kernels/mlaChunkedPrefillTest.cu index 3e4e9a1da0a3..c1fa77239729 100644 --- a/cpp/tests/unit_tests/kernels/mlaChunkedPrefillTest.cu +++ b/cpp/tests/unit_tests/kernels/mlaChunkedPrefillTest.cu @@ -8,6 +8,7 @@ #include "tensorrt_llm/kernels/kvCacheUtils.h" #include "tensorrt_llm/runtime/bufferManager.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/kernels/mlaChunkedPrefill.cuh" #include "tensorrt_llm/runtime/cudaStream.h" #include @@ -429,18 +430,18 @@ protected: using tensorrt_llm::runtime::ITensor; using tensorrt_llm::runtime::bufferCast; - auto dtype = nvinfer1::DataType::kHALF; + auto dtype = tensorrt_llm::DataType::kHALF; if constexpr (std::is_same_v) { - dtype = nvinfer1::DataType::kFLOAT; + dtype = tensorrt_llm::DataType::kFLOAT; } else if constexpr (std::is_same_v) { - dtype = nvinfer1::DataType::kHALF; + dtype = tensorrt_llm::DataType::kHALF; } else if constexpr (std::is_same_v) { - dtype = nvinfer1::DataType::kBF16; + dtype = tensorrt_llm::DataType::kBF16; } else { @@ -449,11 +450,11 @@ protected: auto cacheType = dtype; if constexpr (std::is_same_v) { - cacheType = nvinfer1::DataType::kFP8; + cacheType = tensorrt_llm::DataType::kFP8; this->h_kv_scale_quant_orig - = tensorrt_llm::runtime::BufferManager::pinned(ITensor::makeShape({1}), nvinfer1::DataType::kFLOAT); - this->d_kv_scale_quant_orig - = tensorrt_llm::runtime::BufferManager::gpuSync(ITensor::makeShape({1}), nvinfer1::DataType::kFLOAT); + = tensorrt_llm::runtime::BufferManager::pinned(ITensor::makeShape({1}), tensorrt_llm::DataType::kFLOAT); + this->d_kv_scale_quant_orig = tensorrt_llm::runtime::BufferManager::gpuSync( + ITensor::makeShape({1}), tensorrt_llm::DataType::kFLOAT); auto* kv_scale_quant_orig_ptr = bufferCast(*(this->h_kv_scale_quant_orig)); float kv_scale_orig_quant = 2.0F; kv_scale_quant_orig_ptr[0] = 1.0 / kv_scale_orig_quant; @@ -463,13 +464,13 @@ protected: // cu lens this->h_cu_kv_seq_lens = tensorrt_llm::runtime::BufferManager::pinned( - ITensor::makeShape({this->mBatchSize + 1}), nvinfer1::DataType::kINT64); + ITensor::makeShape({this->mBatchSize + 1}), tensorrt_llm::DataType::kINT64); this->h_cu_q_seq_lens = tensorrt_llm::runtime::BufferManager::pinned( - ITensor::makeShape({this->mBatchSize + 1}), nvinfer1::DataType::kINT64); + ITensor::makeShape({this->mBatchSize + 1}), tensorrt_llm::DataType::kINT64); this->d_cu_kv_seq_lens = tensorrt_llm::runtime::BufferManager::gpuSync( - this->h_cu_kv_seq_lens->getShape(), nvinfer1::DataType::kINT64); + this->h_cu_kv_seq_lens->getShape(), tensorrt_llm::DataType::kINT64); this->d_cu_q_seq_lens = tensorrt_llm::runtime::BufferManager::gpuSync( - this->h_cu_q_seq_lens->getShape(), nvinfer1::DataType::kINT64); + this->h_cu_q_seq_lens->getShape(), tensorrt_llm::DataType::kINT64); { this->mMaxSeqLen = 0; this->mMaxQSeqLen = 0; @@ -512,14 +513,14 @@ protected: int const total_cached_kv_len = this->mTotalKVLen - this->mTotalQLen; int const chunked_loop_num = (total_cached_kv_len + total_chunk_size - 1) / total_chunk_size; this->h_cu_chunk_lens = tensorrt_llm::runtime::BufferManager::pinned( - ITensor::makeShape({chunked_loop_num + 1, this->mBatchSize + 1}), nvinfer1::DataType::kINT64); + ITensor::makeShape({chunked_loop_num + 1, this->mBatchSize + 1}), tensorrt_llm::DataType::kINT64); this->h_chunked_ld_global_offset = tensorrt_llm::runtime::BufferManager::pinned( - ITensor::makeShape({chunked_loop_num + 1, this->mBatchSize}), nvinfer1::DataType::kINT64); + ITensor::makeShape({chunked_loop_num + 1, this->mBatchSize}), tensorrt_llm::DataType::kINT64); this->memsetZeroHost(this->h_chunked_ld_global_offset); this->d_cu_chunk_lens = tensorrt_llm::runtime::BufferManager::gpuSync( - this->h_cu_chunk_lens->getShape(), nvinfer1::DataType::kINT64); + this->h_cu_chunk_lens->getShape(), tensorrt_llm::DataType::kINT64); this->d_chunked_ld_global_offset = tensorrt_llm::runtime::BufferManager::gpuSync( - this->h_chunked_ld_global_offset->getShape(), nvinfer1::DataType::kINT64); + this->h_chunked_ld_global_offset->getShape(), tensorrt_llm::DataType::kINT64); // kv cache this->mMaxBlockPerSeq = (this->mMaxSeqLen + this->mTokensPerBlock - 1) / this->mTokensPerBlock; @@ -539,13 +540,13 @@ protected: this->mLoraSize + this->mRopeSize}), cacheType); this->h_compressed_offset_tensor = tensorrt_llm::runtime::BufferManager::pinned( - ITensor::makeShape({this->mBatchSize, 2, this->mMaxBlockPerSeq + 1}), nvinfer1::DataType::kINT32); + ITensor::makeShape({this->mBatchSize, 2, this->mMaxBlockPerSeq + 1}), tensorrt_llm::DataType::kINT32); this->d_kv_cache_tensor = tensorrt_llm::runtime::BufferManager::gpuSync(this->h_kv_cache_tensor->getShape(), dtype); this->d_compressed_kv_cache_tensor = tensorrt_llm::runtime::BufferManager::gpuSync(this->h_compressed_kv_cache_tensor->getShape(), cacheType); this->d_compressed_offset_tensor = tensorrt_llm::runtime::BufferManager::gpuSync( - this->h_compressed_offset_tensor->getShape(), nvinfer1::DataType::kINT32); + this->h_compressed_offset_tensor->getShape(), tensorrt_llm::DataType::kINT32); { auto* compressed_kv_cache_ptr = bufferCast(*(this->h_compressed_kv_cache_tensor)); @@ -601,15 +602,15 @@ protected: this->m_h_output_tensor = tensorrt_llm::runtime::BufferManager::pinned( ITensor::makeShape({this->mTotalQLen, this->mNumHeads, this->mNopeSize}), dtype); this->m_h_softmax_sum_tensor = tensorrt_llm::runtime::BufferManager::pinned( - ITensor::makeShape({2, this->mTotalQLen, this->mNumHeads}), nvinfer1::DataType::kFLOAT); + ITensor::makeShape({2, this->mTotalQLen, this->mNumHeads}), tensorrt_llm::DataType::kFLOAT); this->m_h_softmax_sum_accum_tensor = tensorrt_llm::runtime::BufferManager::pinned( - ITensor::makeShape({2, this->mTotalQLen, this->mNumHeads}), nvinfer1::DataType::kFLOAT); + ITensor::makeShape({2, this->mTotalQLen, this->mNumHeads}), tensorrt_llm::DataType::kFLOAT); this->m_h_output_tensor_ref = tensorrt_llm::runtime::BufferManager::pinned( ITensor::makeShape({this->mTotalQLen, this->mNumHeads, this->mNopeSize}), dtype); this->m_h_output_tensor_accum = tensorrt_llm::runtime::BufferManager::pinned( ITensor::makeShape({this->mTotalQLen, this->mNumHeads, this->mNopeSize}), dtype); this->m_h_merge_op = tensorrt_llm::runtime::BufferManager::pinned( - ITensor::makeShape({chunked_loop_num + 1, this->mBatchSize}), nvinfer1::DataType::kINT64); + ITensor::makeShape({chunked_loop_num + 1, this->mBatchSize}), tensorrt_llm::DataType::kINT64); this->m_d_q_tensor = tensorrt_llm::runtime::BufferManager::gpuSync(this->m_h_q_tensor->getShape(), dtype); this->m_d_kv_full_tensor = tensorrt_llm::runtime::BufferManager::gpuSync(this->m_h_kv_full_tensor->getShape(), dtype); @@ -618,13 +619,13 @@ protected: this->m_d_output_tensor = tensorrt_llm::runtime::BufferManager::gpuSync(this->m_h_output_tensor->getShape(), dtype); this->m_d_softmax_sum_tensor = tensorrt_llm::runtime::BufferManager::gpuSync( - this->m_h_softmax_sum_tensor->getShape(), nvinfer1::DataType::kFLOAT); + this->m_h_softmax_sum_tensor->getShape(), tensorrt_llm::DataType::kFLOAT); this->m_d_softmax_sum_accum_tensor = tensorrt_llm::runtime::BufferManager::gpuSync( - this->m_h_softmax_sum_accum_tensor->getShape(), nvinfer1::DataType::kFLOAT); + this->m_h_softmax_sum_accum_tensor->getShape(), tensorrt_llm::DataType::kFLOAT); this->m_d_output_tensor_accum = tensorrt_llm::runtime::BufferManager::gpuSync(this->m_h_output_tensor_accum->getShape(), dtype); - this->m_d_merge_op - = tensorrt_llm::runtime::BufferManager::gpuSync(this->m_h_merge_op->getShape(), nvinfer1::DataType::kINT64); + this->m_d_merge_op = tensorrt_llm::runtime::BufferManager::gpuSync( + this->m_h_merge_op->getShape(), tensorrt_llm::DataType::kINT64); { auto* q_ptr = bufferCast(*(this->m_h_q_tensor)); diff --git a/cpp/tests/unit_tests/kernels/mlaPreprocessTest.cu b/cpp/tests/unit_tests/kernels/mlaPreprocessTest.cu index f2c0863779bc..3fc249a2f24a 100644 --- a/cpp/tests/unit_tests/kernels/mlaPreprocessTest.cu +++ b/cpp/tests/unit_tests/kernels/mlaPreprocessTest.cu @@ -23,6 +23,7 @@ #include "tensorrt_llm/kernels/kvCacheUtils.h" #include "tensorrt_llm/runtime/bufferManager.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/kernels/mlaKernels.h" #include @@ -232,18 +233,18 @@ protected: using tensorrt_llm::runtime::ITensor; using tensorrt_llm::runtime::bufferCast; - auto dtype = nvinfer1::DataType::kHALF; + auto dtype = tensorrt_llm::DataType::kHALF; if constexpr (std::is_same_v) { - dtype = nvinfer1::DataType::kFLOAT; + dtype = tensorrt_llm::DataType::kFLOAT; } else if constexpr (std::is_same_v) { - dtype = nvinfer1::DataType::kHALF; + dtype = tensorrt_llm::DataType::kHALF; } else if constexpr (std::is_same_v) { - dtype = nvinfer1::DataType::kBF16; + dtype = tensorrt_llm::DataType::kBF16; } else { @@ -252,15 +253,15 @@ protected: auto cache_dtype = dtype; if constexpr (std::is_same_v) { - cache_dtype = nvinfer1::DataType::kFP8; + cache_dtype = tensorrt_llm::DataType::kFP8; this->h_kv_scale_orig_quant - = tensorrt_llm::runtime::BufferManager::pinned(ITensor::makeShape({1}), nvinfer1::DataType::kFLOAT); - this->d_kv_scale_orig_quant - = tensorrt_llm::runtime::BufferManager::gpuSync(ITensor::makeShape({1}), nvinfer1::DataType::kFLOAT); + = tensorrt_llm::runtime::BufferManager::pinned(ITensor::makeShape({1}), tensorrt_llm::DataType::kFLOAT); + this->d_kv_scale_orig_quant = tensorrt_llm::runtime::BufferManager::gpuSync( + ITensor::makeShape({1}), tensorrt_llm::DataType::kFLOAT); this->h_kv_scale_quant_orig - = tensorrt_llm::runtime::BufferManager::pinned(ITensor::makeShape({1}), nvinfer1::DataType::kFLOAT); - this->d_kv_scale_quant_orig - = tensorrt_llm::runtime::BufferManager::gpuSync(ITensor::makeShape({1}), nvinfer1::DataType::kFLOAT); + = tensorrt_llm::runtime::BufferManager::pinned(ITensor::makeShape({1}), tensorrt_llm::DataType::kFLOAT); + this->d_kv_scale_quant_orig = tensorrt_llm::runtime::BufferManager::gpuSync( + ITensor::makeShape({1}), tensorrt_llm::DataType::kFLOAT); auto* kv_scale_orig_quant_ptr = bufferCast(*(this->h_kv_scale_orig_quant)); auto* kv_scale_quant_orig_ptr = bufferCast(*(this->h_kv_scale_quant_orig)); float kv_scale_orig_quant = 2.0f; @@ -276,13 +277,13 @@ protected: static_assert(std::is_same_v, "TCache must be the same type as DataType"); } this->h_cu_seq_lens = tensorrt_llm::runtime::BufferManager::pinned( - ITensor::makeShape({this->mNumRequests + 1}), nvinfer1::DataType::kINT64); + ITensor::makeShape({this->mNumRequests + 1}), tensorrt_llm::DataType::kINT64); this->h_cu_ctx_cached_kv_lens = tensorrt_llm::runtime::BufferManager::pinned( - ITensor::makeShape({this->mNumRequests + 1}), nvinfer1::DataType::kINT64); + ITensor::makeShape({this->mNumRequests + 1}), tensorrt_llm::DataType::kINT64); this->d_cu_seq_lens = tensorrt_llm::runtime::BufferManager::gpuSync( - ITensor::makeShape({this->mNumRequests + 1}), nvinfer1::DataType::kINT64); + ITensor::makeShape({this->mNumRequests + 1}), tensorrt_llm::DataType::kINT64); this->d_cu_ctx_cached_kv_lens = tensorrt_llm::runtime::BufferManager::gpuSync( - ITensor::makeShape({this->mNumRequests + 1}), nvinfer1::DataType::kINT64); + ITensor::makeShape({this->mNumRequests + 1}), tensorrt_llm::DataType::kINT64); { // set random sequence length auto* cu_seq_lens_temp_ptr = bufferCast(*(this->h_cu_seq_lens)); @@ -333,9 +334,9 @@ protected: this->mTokensPerBlock, this->mLoraSize + this->mRopeSize}), cache_dtype); this->h_offset_tensor = tensorrt_llm::runtime::BufferManager::pinned( - ITensor::makeShape({this->mNumRequests, 2, this->mMaxBlockPerSeq}), nvinfer1::DataType::kINT32); + ITensor::makeShape({this->mNumRequests, 2, this->mMaxBlockPerSeq}), tensorrt_llm::DataType::kINT32); this->h_compressed_offset_tensor = tensorrt_llm::runtime::BufferManager::pinned( - ITensor::makeShape({this->mNumRequests, 2, this->mMaxBlockPerSeq}), nvinfer1::DataType::kINT32); + ITensor::makeShape({this->mNumRequests, 2, this->mMaxBlockPerSeq}), tensorrt_llm::DataType::kINT32); this->d_kv_cache_tensor = tensorrt_llm::runtime::BufferManager::gpuSync( ITensor::makeShape({this->mNumRequests, 2, this->mMaxBlockPerSeq, this->mNumHeadsUncompressed, this->mTokensPerBlock, this->mUncompressedHeadSize + this->mRopeSize}), @@ -349,9 +350,9 @@ protected: this->mTokensPerBlock, this->mLoraSize + this->mRopeSize}), cache_dtype); this->d_offset_tensor = tensorrt_llm::runtime::BufferManager::gpuSync( - ITensor::makeShape({this->mNumRequests, 2, this->mMaxBlockPerSeq}), nvinfer1::DataType::kINT32); + ITensor::makeShape({this->mNumRequests, 2, this->mMaxBlockPerSeq}), tensorrt_llm::DataType::kINT32); this->d_compressed_offset_tensor = tensorrt_llm::runtime::BufferManager::gpuSync( - ITensor::makeShape({this->mNumRequests, 2, this->mMaxBlockPerSeq}), nvinfer1::DataType::kINT32); + ITensor::makeShape({this->mNumRequests, 2, this->mMaxBlockPerSeq}), tensorrt_llm::DataType::kINT32); { auto* kv_cache_ptr = bufferCast(*(this->h_kv_cache_tensor)); auto* kv_cache_ref_ptr = bufferCast(*(this->h_kv_cache_tensor_ref)); diff --git a/cpp/tests/unit_tests/kernels/prepareCustomMaskTest.cpp b/cpp/tests/unit_tests/kernels/prepareCustomMaskTest.cpp index 61617934f236..6616c9c47668 100644 --- a/cpp/tests/unit_tests/kernels/prepareCustomMaskTest.cpp +++ b/cpp/tests/unit_tests/kernels/prepareCustomMaskTest.cpp @@ -25,6 +25,7 @@ #include "tensorrt_llm/common/assert.h" #include "tensorrt_llm/common/cudaUtils.h" #include "tensorrt_llm/common/memoryUtils.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h" #include "tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaRunnerParams.h" #include "tensorrt_llm/kernels/trtllmGenKernels/fmha/prepareCustomMask.h" @@ -293,8 +294,8 @@ class PrepareCustomMaskTest : public ::testing::Test int64_t totalMaskSize = static_cast(batchSize) * maxNumTilesQ * maxNumCustomMaskTilesKv * numInstsQ * numInstsKv * (tileSizeQ * tileSizeKvPadded) / 32; - auto customMaskOffsetsDevice = mBufferManager->gpu(batchSize, nvinfer1::DataType::kINT64); - auto customMaskDevice = mBufferManager->gpu(totalMaskSize, nvinfer1::DataType::kINT32); + auto customMaskOffsetsDevice = mBufferManager->gpu(batchSize, tensorrt_llm::DataType::kINT64); + auto customMaskDevice = mBufferManager->gpu(totalMaskSize, tensorrt_llm::DataType::kINT32); // Clear GPU buffers to ensure no stale data from previous tests cudaMemsetAsync(bufferCast(*customMaskOffsetsDevice), 0, batchSize * sizeof(int64_t), mStream->get()); diff --git a/cpp/tests/unit_tests/kernels/ropeTest.cu b/cpp/tests/unit_tests/kernels/ropeTest.cu index 517b006e4fde..36c91481722f 100644 --- a/cpp/tests/unit_tests/kernels/ropeTest.cu +++ b/cpp/tests/unit_tests/kernels/ropeTest.cu @@ -15,11 +15,13 @@ */ #include +#include "tensorrt_llm/common/cudaUtils.h" #include "tensorrt_llm/common/quantization.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/kernels/decodingCommon.h" +#include "tensorrt_llm/kernels/gptKernels.h" #include "tensorrt_llm/kernels/kvCacheUtils.h" #include "tensorrt_llm/kernels/unfusedAttentionKernels.h" -#include "tensorrt_llm/plugins/gptAttentionCommon/gptAttentionCommon.h" #include "tensorrt_llm/runtime/bufferManager.h" #include @@ -29,6 +31,7 @@ #include #endif +using namespace tensorrt_llm::common; using namespace tensorrt_llm::runtime; using namespace tensorrt_llm::kernels; @@ -502,26 +505,27 @@ protected: { auto const cu_seqlens_size = batch_size + 1; - cu_q_seqlens_tensor = mBufferManager->pinned(ITensor::makeShape({cu_seqlens_size}), nvinfer1::DataType::kINT32); + cu_q_seqlens_tensor + = mBufferManager->pinned(ITensor::makeShape({cu_seqlens_size}), tensorrt_llm::DataType::kINT32); cu_kv_seqlens_tensor - = mBufferManager->pinned(ITensor::makeShape({cu_seqlens_size}), nvinfer1::DataType::kINT32); - padding_offset_tensor - = mBufferManager->pinned(ITensor::makeShape({batch_size, input_seq_length}), nvinfer1::DataType::kINT32); - encoder_padding_offset_tensor - = mBufferManager->pinned(ITensor::makeShape({batch_size, cross_qkv_length}), nvinfer1::DataType::kINT32); + = mBufferManager->pinned(ITensor::makeShape({cu_seqlens_size}), tensorrt_llm::DataType::kINT32); + padding_offset_tensor = mBufferManager->pinned( + ITensor::makeShape({batch_size, input_seq_length}), tensorrt_llm::DataType::kINT32); + encoder_padding_offset_tensor = mBufferManager->pinned( + ITensor::makeShape({batch_size, cross_qkv_length}), tensorrt_llm::DataType::kINT32); fmha_tile_counter_ptr_tensor - = mBufferManager->pinned(ITensor::makeShape({mEnableContextFMHA ? 1 : 0}), nvinfer1::DataType::kINT32); + = mBufferManager->pinned(ITensor::makeShape({mEnableContextFMHA ? 1 : 0}), tensorrt_llm::DataType::kINT32); rotary_inv_freq_buf_tensor = mBufferManager->pinned( - ITensor::makeShape({batch_size, mRotaryEmbeddingDim / 2}), nvinfer1::DataType::kFLOAT); + ITensor::makeShape({batch_size, mRotaryEmbeddingDim / 2}), tensorrt_llm::DataType::kFLOAT); int const max_num_tokens = batch_size * input_seq_length; tokens_info_tensor - = mBufferManager->pinned(ITensor::makeShape({max_num_tokens, 2}), nvinfer1::DataType::kINT32); + = mBufferManager->pinned(ITensor::makeShape({max_num_tokens, 2}), tensorrt_llm::DataType::kINT32); #ifdef ENABLE_FP4 if constexpr (std::is_same_v) { - global_scale_tensor = mBufferManager->pinned(ITensor::makeShape({2}), nvinfer1::DataType::kFLOAT); + global_scale_tensor = mBufferManager->pinned(ITensor::makeShape({2}), tensorrt_llm::DataType::kFLOAT); } #endif } @@ -584,7 +588,7 @@ protected: // // Rotary cos sin cache buffer to avoid re-computing. SizeType32 maxOutputSize{generateRandomSizeSmallerThan(1024)}; rotary_cos_sin_tensor = this->mBufferManager->pinned( - ITensor::makeShape({mRotaryEmbeddingMaxPositions, mRotaryEmbeddingDim}), nvinfer1::DataType::kFLOAT); + ITensor::makeShape({mRotaryEmbeddingMaxPositions, mRotaryEmbeddingDim}), tensorrt_llm::DataType::kFLOAT); rotary_fill_help = bufferCast(*(rotary_cos_sin_tensor)); // createCosSinBuf(rotary_fill_help, mRotaryEmbeddingMaxPositions, mRotaryEmbeddingDim); //currently broken // fillWithOnesAndZerosInterleaved(rotary_fill_help, mRotaryEmbeddingMaxPositions* @@ -594,7 +598,7 @@ protected: batch_size = generateRandomSizeSmallerThan(12); - q_seq_lengths_tensor = mBufferManager->pinned(ITensor::makeShape({batch_size}), nvinfer1::DataType::kINT32); + q_seq_lengths_tensor = mBufferManager->pinned(ITensor::makeShape({batch_size}), tensorrt_llm::DataType::kINT32); q_seq_lengths = bufferCast(*(q_seq_lengths_tensor)); for (SizeType32 ii = 0; ii < batch_size; ++ii) diff --git a/cpp/tests/unit_tests/kernels/routing/routingDeepSeekTest.cpp b/cpp/tests/unit_tests/kernels/routing/routingDeepSeekTest.cpp index 78598fa4c417..be5c8f3c48d1 100644 --- a/cpp/tests/unit_tests/kernels/routing/routingDeepSeekTest.cpp +++ b/cpp/tests/unit_tests/kernels/routing/routingDeepSeekTest.cpp @@ -14,6 +14,7 @@ * limitations under the License. */ +#include "tensorrt_llm/common/tllmDataType.h" #include "tests/unit_tests/kernels/routing/routingTest.h" namespace tk = tensorrt_llm::kernels; @@ -156,8 +157,8 @@ class RoutingDeepSeekKernelTest : public RoutingKernelTest { RoutingKernelTest::allocateBuffers(param); int64_t scoresSize = param.numTokens * param.numExperts; - this->mPtrScoresHost = mBufferManager->pinned(ITensor::makeShape({scoresSize}), nvinfer1::DataType::kFLOAT); - this->mPtrScoresDevice = mBufferManager->gpu(ITensor::makeShape({scoresSize}), nvinfer1::DataType::kFLOAT); + this->mPtrScoresHost = mBufferManager->pinned(ITensor::makeShape({scoresSize}), tensorrt_llm::DataType::kFLOAT); + this->mPtrScoresDevice = mBufferManager->gpu(ITensor::makeShape({scoresSize}), tensorrt_llm::DataType::kFLOAT); this->mPtrRoutingBiasHost = mBufferManager->pinned(ITensor::makeShape({param.numExperts}), TRTDataType::value); @@ -430,9 +431,9 @@ TYPED_TEST(RoutingDeepSeekKernelTest, ClusterLevelWithFloat32Bias) // the GPU kernel (using fp32 bias) and the host reference (using T-typed bias) // observe numerically equivalent inputs. auto float32BiasHost - = this->mBufferManager->pinned(ITensor::makeShape({param.numExperts}), nvinfer1::DataType::kFLOAT); + = this->mBufferManager->pinned(ITensor::makeShape({param.numExperts}), tensorrt_llm::DataType::kFLOAT); auto float32BiasDevice - = this->mBufferManager->gpu(ITensor::makeShape({param.numExperts}), nvinfer1::DataType::kFLOAT); + = this->mBufferManager->gpu(ITensor::makeShape({param.numExperts}), tensorrt_llm::DataType::kFLOAT); auto fp32BiasPtr = bufferCast(*float32BiasHost); auto tBiasPtr = bufferCast(*this->mPtrRoutingBiasHost); for (int i = 0; i < param.numExperts; i++) diff --git a/cpp/tests/unit_tests/kernels/routing/routingTest.cpp b/cpp/tests/unit_tests/kernels/routing/routingTest.cpp index ba5c020ade9e..c71510dc319c 100644 --- a/cpp/tests/unit_tests/kernels/routing/routingTest.cpp +++ b/cpp/tests/unit_tests/kernels/routing/routingTest.cpp @@ -14,6 +14,7 @@ * limitations under the License. */ #include "tests/unit_tests/kernels/routing/routingTest.h" +#include "tensorrt_llm/common/tllmDataType.h" namespace tensorrt_llm::tests::kernels::routing { @@ -54,25 +55,25 @@ void RoutingKernelTest::allocateBuffers(RoutingKernelTestParam const& param) { countsSize = 2 * 256; } - mPtrExpertCountsHost = mBufferManager->pinned(ITensor::makeShape({countsSize}), nvinfer1::DataType::kINT32); - mPtrExpertCountsDevice = mBufferManager->gpu(ITensor::makeShape({countsSize}), nvinfer1::DataType::kINT32); + mPtrExpertCountsHost = mBufferManager->pinned(ITensor::makeShape({countsSize}), tensorrt_llm::DataType::kINT32); + mPtrExpertCountsDevice = mBufferManager->gpu(ITensor::makeShape({countsSize}), tensorrt_llm::DataType::kINT32); int64_t permIdxSize = 1; - mPtrPermutedIdxSizeHost = mBufferManager->pinned(ITensor::makeShape({permIdxSize}), nvinfer1::DataType::kINT32); - mPtrPermutedIdxSizeDevice = mBufferManager->gpu(ITensor::makeShape({permIdxSize}), nvinfer1::DataType::kINT32); + mPtrPermutedIdxSizeHost = mBufferManager->pinned(ITensor::makeShape({permIdxSize}), tensorrt_llm::DataType::kINT32); + mPtrPermutedIdxSizeDevice = mBufferManager->gpu(ITensor::makeShape({permIdxSize}), tensorrt_llm::DataType::kINT32); int64_t expIdxToPermIdxSize = numTokens * topK; mPtrExpandedIdxToPermutedIdxHost - = mBufferManager->pinned(ITensor::makeShape({expIdxToPermIdxSize}), nvinfer1::DataType::kINT32); + = mBufferManager->pinned(ITensor::makeShape({expIdxToPermIdxSize}), tensorrt_llm::DataType::kINT32); mPtrExpandedIdxToPermutedIdxDevice - = mBufferManager->gpu(ITensor::makeShape({expIdxToPermIdxSize}), nvinfer1::DataType::kINT32); + = mBufferManager->gpu(ITensor::makeShape({expIdxToPermIdxSize}), tensorrt_llm::DataType::kINT32); // int64_t permIdxToTokenIdxSize = (numTokens * topK + (numExperts << paddingLog2) - numExperts); int64_t permIdxToTokenIdxSize = (numTokens * topK + (numExperts * tileTokensDim) - numExperts); mPtrPermutedIdxToTokenIdxHost - = mBufferManager->pinned(ITensor::makeShape({permIdxToTokenIdxSize}), nvinfer1::DataType::kINT32); + = mBufferManager->pinned(ITensor::makeShape({permIdxToTokenIdxSize}), tensorrt_llm::DataType::kINT32); mPtrPermutedIdxToTokenIdxDevice - = mBufferManager->gpu(ITensor::makeShape({permIdxToTokenIdxSize}), nvinfer1::DataType::kINT32); + = mBufferManager->gpu(ITensor::makeShape({permIdxToTokenIdxSize}), tensorrt_llm::DataType::kINT32); int64_t expWeightsSize = numTokens * topK; mPtrTopKWeightsHost = mBufferManager->pinned(ITensor::makeShape({expWeightsSize}), TRTDataType::value); @@ -81,8 +82,8 @@ void RoutingKernelTest::allocateBuffers(RoutingKernelTestParam const& param) if (useTopKAsInput) { int64_t topKIdsSize = numTokens * topK; - mPtrTopKIdsHost = mBufferManager->pinned(ITensor::makeShape({topKIdsSize}), nvinfer1::DataType::kINT32); - mPtrTopKIdsDevice = mBufferManager->gpu(ITensor::makeShape({topKIdsSize}), nvinfer1::DataType::kINT32); + mPtrTopKIdsHost = mBufferManager->pinned(ITensor::makeShape({topKIdsSize}), tensorrt_llm::DataType::kINT32); + mPtrTopKIdsDevice = mBufferManager->gpu(ITensor::makeShape({topKIdsSize}), tensorrt_llm::DataType::kINT32); } else { @@ -91,23 +92,26 @@ void RoutingKernelTest::allocateBuffers(RoutingKernelTestParam const& param) } int64_t ctaIdxSize = numTokens * topK; - mPtrCtaIdxXyToBatchIdxHost = mBufferManager->pinned(ITensor::makeShape({ctaIdxSize}), nvinfer1::DataType::kINT32); - mPtrCtaIdxXyToBatchIdxDevice = mBufferManager->gpu(ITensor::makeShape({ctaIdxSize}), nvinfer1::DataType::kINT32); + mPtrCtaIdxXyToBatchIdxHost + = mBufferManager->pinned(ITensor::makeShape({ctaIdxSize}), tensorrt_llm::DataType::kINT32); + mPtrCtaIdxXyToBatchIdxDevice + = mBufferManager->gpu(ITensor::makeShape({ctaIdxSize}), tensorrt_llm::DataType::kINT32); - mPtrCtaIdxXyToMnLimitHost = mBufferManager->pinned(ITensor::makeShape({ctaIdxSize}), nvinfer1::DataType::kINT32); - mPtrCtaIdxXyToMnLimitDevice = mBufferManager->gpu(ITensor::makeShape({ctaIdxSize}), nvinfer1::DataType::kINT32); + mPtrCtaIdxXyToMnLimitHost + = mBufferManager->pinned(ITensor::makeShape({ctaIdxSize}), tensorrt_llm::DataType::kINT32); + mPtrCtaIdxXyToMnLimitDevice = mBufferManager->gpu(ITensor::makeShape({ctaIdxSize}), tensorrt_llm::DataType::kINT32); int64_t numNonExitingCtasSize = 1; mPtrNumNonExitingCtasHost - = mBufferManager->pinned(ITensor::makeShape({numNonExitingCtasSize}), nvinfer1::DataType::kINT32); + = mBufferManager->pinned(ITensor::makeShape({numNonExitingCtasSize}), tensorrt_llm::DataType::kINT32); mPtrNumNonExitingCtasDevice - = mBufferManager->gpu(ITensor::makeShape({numNonExitingCtasSize}), nvinfer1::DataType::kINT32); + = mBufferManager->gpu(ITensor::makeShape({numNonExitingCtasSize}), tensorrt_llm::DataType::kINT32); int64_t idxSize = numTokens * topK * sizeof(PackedType); - mPtrTopKPackedHost = mBufferManager->pinned(ITensor::makeShape({idxSize}), nvinfer1::DataType::kINT8); - mPtrTopKPackedDevice = mBufferManager->gpu(ITensor::makeShape({idxSize}), nvinfer1::DataType::kINT8); + mPtrTopKPackedHost = mBufferManager->pinned(ITensor::makeShape({idxSize}), tensorrt_llm::DataType::kINT8); + mPtrTopKPackedDevice = mBufferManager->gpu(ITensor::makeShape({idxSize}), tensorrt_llm::DataType::kINT8); mCurandStatesDevice - = mBufferManager->gpu(ITensor::makeShape({numTokens, sizeof(curandState_t)}), nvinfer1::DataType::kINT8); + = mBufferManager->gpu(ITensor::makeShape({numTokens, sizeof(curandState_t)}), tensorrt_llm::DataType::kINT8); } template @@ -127,19 +131,19 @@ void RoutingKernelTest::computePermutation(RoutingKernelTestParam const& para PackedType* expIdxHostPtr = reinterpret_cast(bufferCast(*this->mPtrTopKPackedHost)); auto tokenToExpertHost - = mBufferManager->pinned(ITensor::makeShape({param.numTokens * param.topK}), nvinfer1::DataType::kINT32); + = mBufferManager->pinned(ITensor::makeShape({param.numTokens * param.topK}), tensorrt_llm::DataType::kINT32); auto tokenToExpertHostPtr = bufferCast(*tokenToExpertHost); auto tokenToIdxInExpertHost - = mBufferManager->pinned(ITensor::makeShape({param.numTokens * param.topK}), nvinfer1::DataType::kINT32); + = mBufferManager->pinned(ITensor::makeShape({param.numTokens * param.topK}), tensorrt_llm::DataType::kINT32); auto tokenToIdxInExpertHostPtr = bufferCast(*tokenToIdxInExpertHost); auto expertScanCountsHost - = mBufferManager->pinned(ITensor::makeShape({param.numExperts + 1}), nvinfer1::DataType::kINT32); + = mBufferManager->pinned(ITensor::makeShape({param.numExperts + 1}), tensorrt_llm::DataType::kINT32); auto expertScanCountsHostPtr = bufferCast(*expertScanCountsHost); auto ctaScanCountsHost - = mBufferManager->pinned(ITensor::makeShape({param.numExperts + 1}), nvinfer1::DataType::kINT32); + = mBufferManager->pinned(ITensor::makeShape({param.numExperts + 1}), tensorrt_llm::DataType::kINT32); auto ctaScanCountsHostPtr = bufferCast(*ctaScanCountsHost); for (int ie = 0; ie < param.numExperts + 1; ++ie) @@ -407,7 +411,7 @@ void RoutingKernelTest::runTest(RoutingKernelTestParam const& param) // Retrieve the workspace size of the routing kernel. auto const workspaceSize = getDeviceWorkspaceSize(param); TensorPtr workspaceDevice - = mBufferManager->gpu(ITensor::makeShape({static_cast(workspaceSize)}), nvinfer1::DataType::kINT8); + = mBufferManager->gpu(ITensor::makeShape({static_cast(workspaceSize)}), tensorrt_llm::DataType::kINT8); // Call tested function routing callTestedFunction(param, workspaceDevice); // Verify results diff --git a/cpp/tests/unit_tests/kernels/routing/routingTest.h b/cpp/tests/unit_tests/kernels/routing/routingTest.h index 630cd72a5fcb..8b24ee3aa24e 100644 --- a/cpp/tests/unit_tests/kernels/routing/routingTest.h +++ b/cpp/tests/unit_tests/kernels/routing/routingTest.h @@ -23,7 +23,6 @@ #include "tensorrt_llm/runtime/cudaStream.h" #include "tensorrt_llm/runtime/iBuffer.h" #include "tensorrt_llm/runtime/runtimeKernels.h" -#include "tensorrt_llm/runtime/tllmLogger.h" #include #include #include //@todo check the usage of this diff --git a/cpp/tests/unit_tests/kernels/sampling/samplingPenaltyTest.cpp b/cpp/tests/unit_tests/kernels/sampling/samplingPenaltyTest.cpp index 8896dd005cf7..a188abf6bc31 100644 --- a/cpp/tests/unit_tests/kernels/sampling/samplingPenaltyTest.cpp +++ b/cpp/tests/unit_tests/kernels/sampling/samplingPenaltyTest.cpp @@ -14,6 +14,7 @@ * limitations under the License. */ +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/kernels/penaltyTypes.h" #include "tests/unit_tests/kernels/sampling/samplingTest.h" @@ -161,14 +162,14 @@ class TemperaturePenaltyTest : public SamplingKernelTest mLogitsPtrs = BufferManager::pinned(ITensor::makeShape({mBatchSize}), ptrType); mPenaltyWorkspaceDevice = mBufferManager->gpu( - ITensor::makeShape({mMaxBatchSize, mMaxTokensPerStep, mVocabSize * 2}), nvinfer1::DataType::kINT32); + ITensor::makeShape({mMaxBatchSize, mMaxTokensPerStep, mVocabSize * 2}), tensorrt_llm::DataType::kINT32); - mTokensPerStep = BufferManager::pinned(ITensor::makeShape({mMaxBatchSize}), nvinfer1::DataType::kINT32); + mTokensPerStep = BufferManager::pinned(ITensor::makeShape({mMaxBatchSize}), tensorrt_llm::DataType::kINT32); mBiasHost = BufferManager::pinned(ITensor::makeShape({mMaxBatchSize, mVocabSizePadded}), dataType); mBiasDevice = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize, mVocabSizePadded}), dataType); - mBatchSlots = BufferManager::pinned(ITensor::makeShape({mBatchSize}), nvinfer1::DataType::kINT32); + mBatchSlots = BufferManager::pinned(ITensor::makeShape({mBatchSize}), tensorrt_llm::DataType::kINT32); trk::invokeFill(*mLogitsRefHost, T{0.0f}, *mStream); trk::invokeFill(*mOutLogitsDevice, T{0.0f}, *mStream); @@ -204,7 +205,7 @@ class TemperaturePenaltyTest : public SamplingKernelTest ASSERT_EQ(param.temperaturesSize, mMaxBatchSize) << "Invalid test configuration."; mTemperaturesDevice - = mBufferManager->gpu(ITensor::makeShape({param.temperaturesSize}), nvinfer1::DataType::kFLOAT); + = mBufferManager->gpu(ITensor::makeShape({param.temperaturesSize}), tensorrt_llm::DataType::kFLOAT); mBufferManager->copy(*param.temperatures, *mTemperaturesDevice); } @@ -281,7 +282,8 @@ TYPED_TEST(TemperaturePenaltyTest, NoPenalty) { int32_t batchSize = 6; int32_t maxBatchSize = 2 * batchSize; - TensorPtr temperaturesHost = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kFLOAT); + TensorPtr temperaturesHost + = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kFLOAT); for (int32_t i = 0; i < maxBatchSize; ++i) { bufferCast(*temperaturesHost)[i] = 1.0f; @@ -297,7 +299,8 @@ TYPED_TEST(TemperaturePenaltyTest, LessThanOne) { int32_t batchSize = 6; int32_t maxBatchSize = 2 * batchSize; - TensorPtr temperaturesHost = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kFLOAT); + TensorPtr temperaturesHost + = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kFLOAT); for (int32_t i = 0; i < maxBatchSize; ++i) { bufferCast(*temperaturesHost)[i] = 0.53f; @@ -313,7 +316,8 @@ TYPED_TEST(TemperaturePenaltyTest, GreaterThaneOne) { int32_t batchSize = 6; int32_t maxBatchSize = 2 * batchSize; - TensorPtr temperaturesHost = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kFLOAT); + TensorPtr temperaturesHost + = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kFLOAT); for (int32_t i = 0; i < maxBatchSize; ++i) { bufferCast(*temperaturesHost)[i] = 2.01f; @@ -329,7 +333,8 @@ TYPED_TEST(TemperaturePenaltyTest, Mixed) { int32_t batchSize = 6; int32_t maxBatchSize = 2 * batchSize; - TensorPtr temperaturesHost = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kFLOAT); + TensorPtr temperaturesHost + = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kFLOAT); for (int32_t i = 0; i < maxBatchSize; ++i) { bufferCast(*temperaturesHost)[i] = 0.53f + 0.2f * i; @@ -345,7 +350,8 @@ TYPED_TEST(TemperaturePenaltyTest, LargeVocab) { int32_t batchSize = 6; int32_t maxBatchSize = 2 * batchSize; - TensorPtr temperaturesHost = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kFLOAT); + TensorPtr temperaturesHost + = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kFLOAT); for (int32_t i = 0; i < maxBatchSize; ++i) { bufferCast(*temperaturesHost)[i] = 0.53f + 0.2f * i; @@ -361,7 +367,8 @@ TYPED_TEST(TemperaturePenaltyTest, LargeVocabTokensPerStep) { int32_t batchSize = 6; int32_t maxBatchSize = 2 * batchSize; - TensorPtr temperaturesHost = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kFLOAT); + TensorPtr temperaturesHost + = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kFLOAT); for (int32_t i = 0; i < maxBatchSize; ++i) { bufferCast(*temperaturesHost)[i] = 1.f; // 0.53f + 0.2f * i; @@ -541,25 +548,25 @@ class RepetitionPenaltyTest : public SamplingKernelTest mLogitsPtrs = BufferManager::pinned(ITensor::makeShape({mBatchSize}), ptrType); mPenaltyWorkspaceDevice = mBufferManager->gpu( - ITensor::makeShape({mBatchSize, mMaxTokensPerStep, mVocabSize * 2}), nvinfer1::DataType::kINT32); + ITensor::makeShape({mBatchSize, mMaxTokensPerStep, mVocabSize * 2}), tensorrt_llm::DataType::kINT32); - mTokensPerStep = BufferManager::pinned(ITensor::makeShape({mMaxBatchSize}), nvinfer1::DataType::kINT32); + mTokensPerStep = BufferManager::pinned(ITensor::makeShape({mMaxBatchSize}), tensorrt_llm::DataType::kINT32); - mOutputIdsHost - = BufferManager::pinned(ITensor::makeShape({mMaxBatchSize, mSequenceLength}), nvinfer1::DataType::kINT32); + mOutputIdsHost = BufferManager::pinned( + ITensor::makeShape({mMaxBatchSize, mSequenceLength}), tensorrt_llm::DataType::kINT32); mOutputIdsDevice - = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize, mSequenceLength}), nvinfer1::DataType::kINT32); + = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize, mSequenceLength}), tensorrt_llm::DataType::kINT32); - mSeqLengthHost = BufferManager::pinned(ITensor::makeShape({mMaxBatchSize}), nvinfer1::DataType::kINT32); - mSeqLengthDevice = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize}), nvinfer1::DataType::kINT32); + mSeqLengthHost = BufferManager::pinned(ITensor::makeShape({mMaxBatchSize}), tensorrt_llm::DataType::kINT32); + mSeqLengthDevice = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize}), tensorrt_llm::DataType::kINT32); - mContextLengthHost = BufferManager::pinned(ITensor::makeShape({mMaxBatchSize}), nvinfer1::DataType::kINT32); - mContextLengthDevice = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize}), nvinfer1::DataType::kINT32); + mContextLengthHost = BufferManager::pinned(ITensor::makeShape({mMaxBatchSize}), tensorrt_llm::DataType::kINT32); + mContextLengthDevice = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize}), tensorrt_llm::DataType::kINT32); mIdsPtrHost = BufferManager::pinned(ITensor::makeShape({mMaxBatchSize}), ptrType); mIdsPtrDevice = BufferManager::pinned(ITensor::makeShape({mMaxBatchSize}), ptrType); - mBatchSlots = BufferManager::pinned(ITensor::makeShape({mBatchSize}), nvinfer1::DataType::kINT32); + mBatchSlots = BufferManager::pinned(ITensor::makeShape({mBatchSize}), tensorrt_llm::DataType::kINT32); auto batchSlotsPtr = bufferCast(*mBatchSlots); for (SizeType32 bi = 0; bi < mBatchSize; ++bi) @@ -606,13 +613,13 @@ class RepetitionPenaltyTest : public SamplingKernelTest ASSERT_EQ(param.frequencyPenaltiesSize, mMaxBatchSize) << "Invalid test configuration."; ASSERT_EQ(param.promptIgnoreLengthsSize, mMaxBatchSize) << "Invalid test configuration."; mRepetitionPenaltiesDevice - = mBufferManager->gpu(ITensor::makeShape({param.repetitionPenaltiesSize}), nvinfer1::DataType::kFLOAT); + = mBufferManager->gpu(ITensor::makeShape({param.repetitionPenaltiesSize}), tensorrt_llm::DataType::kFLOAT); mPresencePenaltiesDevice - = mBufferManager->gpu(ITensor::makeShape({param.presencePenaltiesSize}), nvinfer1::DataType::kFLOAT); + = mBufferManager->gpu(ITensor::makeShape({param.presencePenaltiesSize}), tensorrt_llm::DataType::kFLOAT); mFrequencyPenaltiesDevice - = mBufferManager->gpu(ITensor::makeShape({param.frequencyPenaltiesSize}), nvinfer1::DataType::kFLOAT); + = mBufferManager->gpu(ITensor::makeShape({param.frequencyPenaltiesSize}), tensorrt_llm::DataType::kFLOAT); mPromptIgnoreLengthsDevice - = mBufferManager->gpu(ITensor::makeShape({param.promptIgnoreLengthsSize}), nvinfer1::DataType::kINT32); + = mBufferManager->gpu(ITensor::makeShape({param.promptIgnoreLengthsSize}), tensorrt_llm::DataType::kINT32); mBufferManager->copy(*param.repetitionPenalties, *mRepetitionPenaltiesDevice); mBufferManager->copy(*param.presencePenalties, *mPresencePenaltiesDevice); mBufferManager->copy(*param.frequencyPenalties, *mFrequencyPenaltiesDevice); @@ -740,13 +747,13 @@ TYPED_TEST(RepetitionPenaltyTest, BatchNoPenalty) int32_t batchSize = 6; int32_t maxBatchSize = 2 * batchSize; TensorPtr repetitionPenaltyHost - = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kFLOAT); + = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kFLOAT); TensorPtr presencePenaltyHost - = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kFLOAT); + = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kFLOAT); TensorPtr frequencyPenaltyHost - = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kFLOAT); + = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kFLOAT); TensorPtr promptIgnoreLengthsHost - = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kINT32); + = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kINT32); for (int32_t i = 0; i < maxBatchSize; ++i) { bufferCast(*repetitionPenaltyHost)[i] = 1.0f; @@ -773,13 +780,13 @@ TYPED_TEST(RepetitionPenaltyTest, BatchRepetitionLessThanOne) int32_t batchSize = 6; int32_t maxBatchSize = 2 * batchSize; TensorPtr repetitionPenaltyHost - = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kFLOAT); + = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kFLOAT); TensorPtr presencePenaltyHost - = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kFLOAT); + = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kFLOAT); TensorPtr frequencyPenaltyHost - = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kFLOAT); + = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kFLOAT); TensorPtr promptIgnoreLengthsHost - = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kINT32); + = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kINT32); for (int32_t i = 0; i < maxBatchSize; ++i) { bufferCast(*repetitionPenaltyHost)[i] = 0.53f; @@ -806,13 +813,13 @@ TYPED_TEST(RepetitionPenaltyTest, BatchRepetitionGreaterThaneOne) int32_t batchSize = 6; int32_t maxBatchSize = 2 * batchSize; TensorPtr repetitionPenaltyHost - = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kFLOAT); + = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kFLOAT); TensorPtr presencePenaltyHost - = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kFLOAT); + = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kFLOAT); TensorPtr frequencyPenaltyHost - = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kFLOAT); + = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kFLOAT); TensorPtr promptIgnoreLengthsHost - = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kINT32); + = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kINT32); for (int32_t i = 0; i < maxBatchSize; ++i) { bufferCast(*repetitionPenaltyHost)[i] = 2.01f; @@ -839,13 +846,13 @@ TYPED_TEST(RepetitionPenaltyTest, BatchRepetitionMixed) int32_t batchSize = 6; int32_t maxBatchSize = 2 * batchSize; TensorPtr repetitionPenaltyHost - = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kFLOAT); + = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kFLOAT); TensorPtr presencePenaltyHost - = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kFLOAT); + = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kFLOAT); TensorPtr frequencyPenaltyHost - = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kFLOAT); + = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kFLOAT); TensorPtr promptIgnoreLengthsHost - = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kINT32); + = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kINT32); for (int32_t i = 0; i < maxBatchSize; ++i) { bufferCast(*repetitionPenaltyHost)[i] = 0.53 + i * 0.2f; @@ -872,13 +879,13 @@ TYPED_TEST(RepetitionPenaltyTest, BatchPresenceMixed) int32_t batchSize = 6; int32_t maxBatchSize = 2 * batchSize; TensorPtr repetitionPenaltyHost - = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kFLOAT); + = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kFLOAT); TensorPtr presencePenaltyHost - = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kFLOAT); + = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kFLOAT); TensorPtr frequencyPenaltyHost - = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kFLOAT); + = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kFLOAT); TensorPtr promptIgnoreLengthsHost - = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kINT32); + = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kINT32); for (int32_t i = 0; i < maxBatchSize; ++i) { bufferCast(*repetitionPenaltyHost)[i] = 1.0f; @@ -905,13 +912,13 @@ TYPED_TEST(RepetitionPenaltyTest, BatchPresenceHasDefaultValueZero2) int32_t batchSize = 6; int32_t maxBatchSize = 2 * batchSize; TensorPtr repetitionPenaltyHost - = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kFLOAT); + = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kFLOAT); TensorPtr presencePenaltyHost - = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kFLOAT); + = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kFLOAT); TensorPtr frequencyPenaltyHost - = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kFLOAT); + = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kFLOAT); TensorPtr promptIgnoreLengthsHost - = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kINT32); + = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kINT32); for (int32_t i = 0; i < maxBatchSize; ++i) { bufferCast(*repetitionPenaltyHost)[i] = 1.0f; @@ -938,13 +945,13 @@ TYPED_TEST(RepetitionPenaltyTest, BatchFrequencyMixed) int32_t batchSize = 6; int32_t maxBatchSize = 2 * batchSize; TensorPtr repetitionPenaltyHost - = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kFLOAT); + = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kFLOAT); TensorPtr presencePenaltyHost - = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kFLOAT); + = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kFLOAT); TensorPtr frequencyPenaltyHost - = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kFLOAT); + = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kFLOAT); TensorPtr promptIgnoreLengthsHost - = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kINT32); + = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kINT32); for (int32_t i = 0; i < maxBatchSize; ++i) { bufferCast(*repetitionPenaltyHost)[i] = 1.0f; @@ -971,13 +978,13 @@ TYPED_TEST(RepetitionPenaltyTest, BatchFrequencyHasDefaultValueZero2) int32_t batchSize = 6; int32_t maxBatchSize = 2 * batchSize; TensorPtr repetitionPenaltyHost - = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kFLOAT); + = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kFLOAT); TensorPtr presencePenaltyHost - = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kFLOAT); + = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kFLOAT); TensorPtr frequencyPenaltyHost - = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kFLOAT); + = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kFLOAT); TensorPtr promptIgnoreLengthsHost - = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kINT32); + = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kINT32); for (int32_t i = 0; i < maxBatchSize; ++i) { bufferCast(*repetitionPenaltyHost)[i] = 1.0f; @@ -1004,13 +1011,13 @@ TYPED_TEST(RepetitionPenaltyTest, PenaltyTypeRepetitionPresence) int32_t batchSize = 6; int32_t maxBatchSize = 2 * batchSize; TensorPtr repetitionPenaltyHost - = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kFLOAT); + = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kFLOAT); TensorPtr presencePenaltyHost - = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kFLOAT); + = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kFLOAT); TensorPtr frequencyPenaltyHost - = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kFLOAT); + = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kFLOAT); TensorPtr promptIgnoreLengthsHost - = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kINT32); + = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kINT32); for (int32_t i = 0; i < maxBatchSize; ++i) { bufferCast(*repetitionPenaltyHost)[i] = 0.53 + i * 0.2f; @@ -1037,13 +1044,13 @@ TYPED_TEST(RepetitionPenaltyTest, PenaltyTypeRepetitionFrequency) int32_t batchSize = 6; int32_t maxBatchSize = 2 * batchSize; TensorPtr repetitionPenaltyHost - = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kFLOAT); + = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kFLOAT); TensorPtr presencePenaltyHost - = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kFLOAT); + = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kFLOAT); TensorPtr frequencyPenaltyHost - = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kFLOAT); + = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kFLOAT); TensorPtr promptIgnoreLengthsHost - = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kINT32); + = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kINT32); for (int32_t i = 0; i < maxBatchSize; ++i) { bufferCast(*repetitionPenaltyHost)[i] = 0.53 + i * 0.2f; @@ -1070,13 +1077,13 @@ TYPED_TEST(RepetitionPenaltyTest, PenaltyTypePresenceFrequency) int32_t batchSize = 6; int32_t maxBatchSize = 2 * batchSize; TensorPtr repetitionPenaltyHost - = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kFLOAT); + = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kFLOAT); TensorPtr presencePenaltyHost - = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kFLOAT); + = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kFLOAT); TensorPtr frequencyPenaltyHost - = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kFLOAT); + = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kFLOAT); TensorPtr promptIgnoreLengthsHost - = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kINT32); + = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kINT32); for (int32_t i = 0; i < maxBatchSize; ++i) { bufferCast(*repetitionPenaltyHost)[i] = 1.0f; @@ -1103,13 +1110,13 @@ TYPED_TEST(RepetitionPenaltyTest, PenaltyTypeFull) int32_t batchSize = 6; int32_t maxBatchSize = 2 * batchSize; TensorPtr repetitionPenaltyHost - = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kFLOAT); + = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kFLOAT); TensorPtr presencePenaltyHost - = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kFLOAT); + = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kFLOAT); TensorPtr frequencyPenaltyHost - = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kFLOAT); + = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kFLOAT); TensorPtr promptIgnoreLengthsHost - = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kINT32); + = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kINT32); for (int32_t i = 0; i < maxBatchSize; ++i) { bufferCast(*repetitionPenaltyHost)[i] = 0.53 + i * 0.2f; @@ -1136,13 +1143,13 @@ TYPED_TEST(RepetitionPenaltyTest, PenaltyTypeFullTokensPerStep) int32_t batchSize = 6; int32_t maxBatchSize = 2 * batchSize; TensorPtr repetitionPenaltyHost - = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kFLOAT); + = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kFLOAT); TensorPtr presencePenaltyHost - = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kFLOAT); + = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kFLOAT); TensorPtr frequencyPenaltyHost - = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kFLOAT); + = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kFLOAT); TensorPtr promptIgnoreLengthsHost - = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kINT32); + = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kINT32); for (int32_t i = 0; i < maxBatchSize; ++i) { bufferCast(*repetitionPenaltyHost)[i] = 0.53 + i * 0.2f; @@ -1170,13 +1177,13 @@ TYPED_TEST(RepetitionPenaltyTest, PenaltyTypeFullWithPartialPromptIgnore) int32_t batchSize = 6; int32_t maxBatchSize = 2 * batchSize; TensorPtr repetitionPenaltyHost - = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kFLOAT); + = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kFLOAT); TensorPtr presencePenaltyHost - = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kFLOAT); + = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kFLOAT); TensorPtr frequencyPenaltyHost - = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kFLOAT); + = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kFLOAT); TensorPtr promptIgnoreLengthsHost - = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kINT32); + = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kINT32); for (int32_t i = 0; i < maxBatchSize; ++i) { bufferCast(*repetitionPenaltyHost)[i] = 0.53 + i * 0.2f; @@ -1203,13 +1210,13 @@ TYPED_TEST(RepetitionPenaltyTest, PenaltyTypeFullTokensPerStepWithFullPromptIgno int32_t batchSize = 6; int32_t maxBatchSize = 2 * batchSize; TensorPtr repetitionPenaltyHost - = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kFLOAT); + = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kFLOAT); TensorPtr presencePenaltyHost - = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kFLOAT); + = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kFLOAT); TensorPtr frequencyPenaltyHost - = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kFLOAT); + = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kFLOAT); TensorPtr promptIgnoreLengthsHost - = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kINT32); + = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kINT32); for (int32_t i = 0; i < maxBatchSize; ++i) { bufferCast(*repetitionPenaltyHost)[i] = 0.53 + i * 0.2f; @@ -1333,23 +1340,23 @@ class MinLengthPenaltyTest : public SamplingKernelTest mLogitsPtrs = BufferManager::pinned(ITensor::makeShape({mBatchSize}), ptrType); mPenaltyWorkspaceDevice = mBufferManager->gpu( - ITensor::makeShape({mBatchSize, mMaxTokensPerStep, mVocabSize}), nvinfer1::DataType::kINT32); + ITensor::makeShape({mBatchSize, mMaxTokensPerStep, mVocabSize}), tensorrt_llm::DataType::kINT32); - mTokensPerStep = BufferManager::pinned(ITensor::makeShape({mMaxBatchSize}), nvinfer1::DataType::kINT32); + mTokensPerStep = BufferManager::pinned(ITensor::makeShape({mMaxBatchSize}), tensorrt_llm::DataType::kINT32); - mSeqLengthHost = BufferManager::pinned(ITensor::makeShape({mMaxBatchSize}), nvinfer1::DataType::kINT32); - mSeqLengthDevice = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize}), nvinfer1::DataType::kINT32); + mSeqLengthHost = BufferManager::pinned(ITensor::makeShape({mMaxBatchSize}), tensorrt_llm::DataType::kINT32); + mSeqLengthDevice = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize}), tensorrt_llm::DataType::kINT32); - mContextLengthHost = BufferManager::pinned(ITensor::makeShape({mMaxBatchSize}), nvinfer1::DataType::kINT32); - mContextLengthDevice = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize}), nvinfer1::DataType::kINT32); + mContextLengthHost = BufferManager::pinned(ITensor::makeShape({mMaxBatchSize}), tensorrt_llm::DataType::kINT32); + mContextLengthDevice = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize}), tensorrt_llm::DataType::kINT32); - mMinLengthHost = BufferManager::pinned(ITensor::makeShape({mMaxBatchSize}), nvinfer1::DataType::kINT32); - mMinLengthDevice = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize}), nvinfer1::DataType::kINT32); + mMinLengthHost = BufferManager::pinned(ITensor::makeShape({mMaxBatchSize}), tensorrt_llm::DataType::kINT32); + mMinLengthDevice = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize}), tensorrt_llm::DataType::kINT32); - mEndIdsHost = BufferManager::pinned(ITensor::makeShape({mMaxBatchSize}), nvinfer1::DataType::kINT32); - mEndIdsDevice = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize}), nvinfer1::DataType::kINT32); + mEndIdsHost = BufferManager::pinned(ITensor::makeShape({mMaxBatchSize}), tensorrt_llm::DataType::kINT32); + mEndIdsDevice = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize}), tensorrt_llm::DataType::kINT32); - mBatchSlots = BufferManager::pinned(ITensor::makeShape({mBatchSize}), nvinfer1::DataType::kINT32); + mBatchSlots = BufferManager::pinned(ITensor::makeShape({mBatchSize}), tensorrt_llm::DataType::kINT32); auto batchSlotsPtr = bufferCast(*mBatchSlots); for (SizeType32 bi = 0; bi < mBatchSize; ++bi) @@ -1551,30 +1558,30 @@ class MinLengthPenaltyOOBSafetyTest : public SamplingKernelTest } // Defines currentStep. - mSeqLengthHost = BufferManager::pinned(ITensor::makeShape({mMaxBatchSize}), nvinfer1::DataType::kINT32); + mSeqLengthHost = BufferManager::pinned(ITensor::makeShape({mMaxBatchSize}), tensorrt_llm::DataType::kINT32); initConstant(bufferCast(*mSeqLengthHost), mMaxBatchSize, 3); - mSeqLengthDevice = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize}), nvinfer1::DataType::kINT32); + mSeqLengthDevice = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize}), tensorrt_llm::DataType::kINT32); mBufferManager->copy(*mSeqLengthHost, *mSeqLengthDevice); // Defines inputLength. - mContextLengthHost = BufferManager::pinned(ITensor::makeShape({mMaxBatchSize}), nvinfer1::DataType::kINT32); + mContextLengthHost = BufferManager::pinned(ITensor::makeShape({mMaxBatchSize}), tensorrt_llm::DataType::kINT32); initConstant(bufferCast(*mContextLengthHost), mMaxBatchSize, 2); - mContextLengthDevice = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize}), nvinfer1::DataType::kINT32); + mContextLengthDevice = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize}), tensorrt_llm::DataType::kINT32); mBufferManager->copy(*mContextLengthHost, *mContextLengthDevice); // Defines minLength. - mMinLengthHost = BufferManager::pinned(ITensor::makeShape({mMaxBatchSize}), nvinfer1::DataType::kINT32); + mMinLengthHost = BufferManager::pinned(ITensor::makeShape({mMaxBatchSize}), tensorrt_llm::DataType::kINT32); initConstant(bufferCast(*mMinLengthHost), mMaxBatchSize, 10); - mMinLengthDevice = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize}), nvinfer1::DataType::kINT32); + mMinLengthDevice = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize}), tensorrt_llm::DataType::kINT32); mBufferManager->copy(*mMinLengthHost, *mMinLengthDevice); // Defines endIds. - mEndIdsHost = BufferManager::pinned(ITensor::makeShape({mMaxBatchSize}), nvinfer1::DataType::kINT32); + mEndIdsHost = BufferManager::pinned(ITensor::makeShape({mMaxBatchSize}), tensorrt_llm::DataType::kINT32); initConstant(bufferCast(*mEndIdsHost), mMaxBatchSize, -1); - mEndIdsDevice = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize}), nvinfer1::DataType::kINT32); + mEndIdsDevice = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize}), tensorrt_llm::DataType::kINT32); mBufferManager->copy(*mEndIdsHost, *mEndIdsDevice); - mBatchSlots = BufferManager::pinned(ITensor::makeShape({mBatchSize}), nvinfer1::DataType::kINT32); + mBatchSlots = BufferManager::pinned(ITensor::makeShape({mBatchSize}), tensorrt_llm::DataType::kINT32); auto batchSlotsPtr = bufferCast(*mBatchSlots); for (SizeType32 bi = 0; bi < mBatchSize; ++bi) { diff --git a/cpp/tests/unit_tests/kernels/sampling/samplingTest.cpp b/cpp/tests/unit_tests/kernels/sampling/samplingTest.cpp index 90da247f5203..1d4f19c45ad2 100644 --- a/cpp/tests/unit_tests/kernels/sampling/samplingTest.cpp +++ b/cpp/tests/unit_tests/kernels/sampling/samplingTest.cpp @@ -14,6 +14,7 @@ * limitations under the License. */ #include "tests/unit_tests/kernels/sampling/samplingTest.h" +#include "tensorrt_llm/common/tllmDataType.h" namespace tensorrt_llm::tests::kernels::sampling { @@ -51,75 +52,78 @@ void SamplingKernelTest::allocateBuffers(SamplingKernelTestParam const& param auto const ptrType = TRTDataType::value; // Allocate GPU data - mSeqLengthsHost = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kINT32); - mSeqLengthsDevice = mBufferManager->gpu(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kINT32); + mSeqLengthsHost = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kINT32); + mSeqLengthsDevice = mBufferManager->gpu(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kINT32); mFinishedHost = BufferManager::pinned( ITensor::makeShape({maxBatchSize}), TRTDataType::value); mFinishedDevice = mBufferManager->gpu( ITensor::makeShape({maxBatchSize}), TRTDataType::value); - mOutputIdsHost = BufferManager::pinned(ITensor::makeShape({maxBatchSize, mMaxSeqLen}), nvinfer1::DataType::kINT32); - mOutputIdsDevice = mBufferManager->gpu(ITensor::makeShape({maxBatchSize, mMaxSeqLen}), nvinfer1::DataType::kINT32); + mOutputIdsHost + = BufferManager::pinned(ITensor::makeShape({maxBatchSize, mMaxSeqLen}), tensorrt_llm::DataType::kINT32); + mOutputIdsDevice + = mBufferManager->gpu(ITensor::makeShape({maxBatchSize, mMaxSeqLen}), tensorrt_llm::DataType::kINT32); mProbsHost = BufferManager::pinned(ITensor::makeShape({batchSize, maxTokensPerStep, vocabSize}), dataType); mProbsDevice = mBufferManager->gpu(ITensor::makeShape({batchSize, maxTokensPerStep, vocabSize}), dataType); mProbsPtrsDevice - = BufferManager::pinned(ITensor::makeShape({batchSize, maxTokensPerStep}), nvinfer1::DataType::kINT64); + = BufferManager::pinned(ITensor::makeShape({batchSize, maxTokensPerStep}), tensorrt_llm::DataType::kINT64); - mCumLogProbsDevice = mBufferManager->gpu(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kFLOAT); + mCumLogProbsDevice = mBufferManager->gpu(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kFLOAT); if (param.returnAllSelectedTokens) { SizeType32 maxTopK = param.topK == 0 ? vocabSize : param.topK; mOutputLogProbsDevice - = mBufferManager->gpu(ITensor::makeShape({maxBatchSize, maxTopK}), nvinfer1::DataType::kFLOAT); + = mBufferManager->gpu(ITensor::makeShape({maxBatchSize, maxTopK}), tensorrt_llm::DataType::kFLOAT); } else { mOutputLogProbsDevice - = mBufferManager->gpu(ITensor::makeShape({mMaxSeqLen, maxBatchSize}), nvinfer1::DataType::kFLOAT); + = mBufferManager->gpu(ITensor::makeShape({mMaxSeqLen, maxBatchSize}), tensorrt_llm::DataType::kFLOAT); } mZeroParentIdsDevice - = mBufferManager->gpu(ITensor::makeShape({maxBatchSize, maxTokensPerStep}), nvinfer1::DataType::kINT32); + = mBufferManager->gpu(ITensor::makeShape({maxBatchSize, maxTokensPerStep}), tensorrt_llm::DataType::kINT32); mLogitsHost = BufferManager::pinned(ITensor::makeShape({batchSize, maxTokensPerStep, vocabSize}), dataType); mLogProbsHost = BufferManager::pinned(ITensor::makeShape({batchSize, maxTokensPerStep, vocabSize}), dataType); mIdsPtrHost = BufferManager::pinned(ITensor::makeShape({2 * maxBatchSize}), ptrType); - mEndIdsHost = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kINT32); - mEndIdsDevice = mBufferManager->gpu(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kINT32); + mEndIdsHost = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kINT32); + mEndIdsDevice = mBufferManager->gpu(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kINT32); - mTopPsHost = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kFLOAT); - mTopPsDevice = mBufferManager->gpu(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kFLOAT); + mTopPsHost = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kFLOAT); + mTopPsDevice = mBufferManager->gpu(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kFLOAT); - mTopKsHost = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kINT32); - mTopKsDevice = mBufferManager->gpu(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kINT32); + mTopKsHost = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kINT32); + mTopKsDevice = mBufferManager->gpu(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kINT32); - mSkipDecodeHost = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kBOOL); - mSkipDecodeDevice = mBufferManager->gpu(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kBOOL); + mSkipDecodeHost = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kBOOL); + mSkipDecodeDevice = mBufferManager->gpu(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kBOOL); - mTokensPerStep = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kINT32); + mTokensPerStep = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kINT32); - mBatchSlots = BufferManager::pinned(ITensor::makeShape({batchSize}), nvinfer1::DataType::kINT32); + mBatchSlots = BufferManager::pinned(ITensor::makeShape({batchSize}), tensorrt_llm::DataType::kINT32); - mExpectedCumLogProbsHost = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kFLOAT); + mExpectedCumLogProbsHost + = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kFLOAT); if (param.returnAllSelectedTokens) { SizeType32 maxTopK = param.topK == 0 ? vocabSize : param.topK; mExpectedLogProbsHost - = BufferManager::pinned(ITensor::makeShape({maxBatchSize, maxTopK}), nvinfer1::DataType::kFLOAT); + = BufferManager::pinned(ITensor::makeShape({maxBatchSize, maxTopK}), tensorrt_llm::DataType::kFLOAT); } else { mExpectedLogProbsHost - = BufferManager::pinned(ITensor::makeShape({mMaxSeqLen, maxBatchSize}), nvinfer1::DataType::kFLOAT); + = BufferManager::pinned(ITensor::makeShape({mMaxSeqLen, maxBatchSize}), tensorrt_llm::DataType::kFLOAT); } mCurandStatesDevice - = mBufferManager->gpu(ITensor::makeShape({maxBatchSize, sizeof(curandState_t)}), nvinfer1::DataType::kINT8); + = mBufferManager->gpu(ITensor::makeShape({maxBatchSize, sizeof(curandState_t)}), tensorrt_llm::DataType::kINT8); } template @@ -494,7 +498,7 @@ void SamplingKernelTest::runTest(SamplingKernelTestParam const& param) // Retrieve the workspace size of the sampling kernel. auto const workspaceSize = getWorkspaceSize(param); TensorPtr workspaceDevice - = mBufferManager->gpu(ITensor::makeShape({static_cast(workspaceSize)}), nvinfer1::DataType::kINT8); + = mBufferManager->gpu(ITensor::makeShape({static_cast(workspaceSize)}), tensorrt_llm::DataType::kINT8); // Call tested function sampling callTestedFunction(param, workspaceDevice); diff --git a/cpp/tests/unit_tests/kernels/sampling/samplingTest.h b/cpp/tests/unit_tests/kernels/sampling/samplingTest.h index 0c7f52ba369e..74268bee092c 100644 --- a/cpp/tests/unit_tests/kernels/sampling/samplingTest.h +++ b/cpp/tests/unit_tests/kernels/sampling/samplingTest.h @@ -27,7 +27,6 @@ #include "tensorrt_llm/runtime/bufferManager.h" #include "tensorrt_llm/runtime/cudaStream.h" #include "tensorrt_llm/runtime/runtimeKernels.h" -#include "tensorrt_llm/runtime/tllmLogger.h" namespace tensorrt_llm::tests::kernels::sampling { diff --git a/cpp/tests/unit_tests/kernels/sampling/samplingUtilsTest.cu b/cpp/tests/unit_tests/kernels/sampling/samplingUtilsTest.cu index 71a6d767171c..f0f7690257e7 100644 --- a/cpp/tests/unit_tests/kernels/sampling/samplingUtilsTest.cu +++ b/cpp/tests/unit_tests/kernels/sampling/samplingUtilsTest.cu @@ -14,6 +14,7 @@ * limitations under the License. */ +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/runtime/gptDecoder.h" #include "tests/unit_tests/kernels/sampling/samplingTest.h" #include @@ -57,7 +58,8 @@ TEST_F(SamplingUtilsKernelTest, CurandInitialize) sync_check_cuda_error(this->mStream->get()); // Generate random numbers using initialized curand states.MemoryType - auto randValsDevice = this->mBufferManager->gpu(ITensor::makeShape({batchSize}), nvinfer1::DataType::kINT32); + auto randValsDevice + = this->mBufferManager->gpu(ITensor::makeShape({batchSize}), tensorrt_llm::DataType::kINT32); generateRandomNumber<<<1, batchSize, 0, this->mStream->get()>>>( bufferCast(*randValsDevice), batchSlotsPtr, curandStates, batchSize); auto randValsHost = this->mBufferManager->copyFrom(*randValsDevice, MemoryType::kCPU); @@ -97,7 +99,7 @@ TEST_F(SamplingUtilsKernelTest, CurandBatchInitialize) curandState_t* curandStates; cudaMalloc(&curandStates, sizeof(curandState_t) * 2 * batchSize); - auto randomSeedsHost = mBufferManager->pinnedPool(ITensor::makeShape({batchSize}), nvinfer1::DataType::kINT64); + auto randomSeedsHost = mBufferManager->pinnedPool(ITensor::makeShape({batchSize}), tensorrt_llm::DataType::kINT64); auto randomSeedsHostPtr = bufferCast(*randomSeedsHost); size_t const periodSize = 3; for (size_t i = 0; i < batchSize; ++i) @@ -106,7 +108,7 @@ TEST_F(SamplingUtilsKernelTest, CurandBatchInitialize) } auto randomSeedsDevice = mBufferManager->copyFrom(*randomSeedsHost, MemoryType::kGPU); - auto batchSlots = mBufferManager->pinnedPool(ITensor::makeShape({batchSize}), nvinfer1::DataType::kINT32); + auto batchSlots = mBufferManager->pinnedPool(ITensor::makeShape({batchSize}), tensorrt_llm::DataType::kINT32); auto batchSlotsPtr = bufferCast(*batchSlots); for (SizeType32 bi = 0; bi < batchSize; ++bi) @@ -120,7 +122,7 @@ TEST_F(SamplingUtilsKernelTest, CurandBatchInitialize) sync_check_cuda_error(mStream->get()); // Generate random numbers using initialized curand states. - auto randValsDevice = mBufferManager->gpu(ITensor::makeShape({batchSize}), nvinfer1::DataType::kINT32); + auto randValsDevice = mBufferManager->gpu(ITensor::makeShape({batchSize}), tensorrt_llm::DataType::kINT32); generateRandomNumber<<<1, batchSize, 0, this->mStream->get()>>>( bufferCast(*randValsDevice), batchSlotsPtr, curandStates, batchSize); auto const randValsHost = mBufferManager->copyFrom(*randValsDevice, MemoryType::kCPU); @@ -166,25 +168,26 @@ public: ITensor::makeShape({batchSize, maxBeamWidth, vocabSizePadded}), dataType); ITensor::SharedPtr logitsHostPtrs = this->mBufferManager->pinnedPool(ITensor::makeShape({batchSize}), ptrType); auto refLogitsHost = this->mBufferManager->pinnedPool( - ITensor::makeShape({batchSize, maxBeamWidth, vocabSizePadded}), nvinfer1::DataType::kFLOAT); + ITensor::makeShape({batchSize, maxBeamWidth, vocabSizePadded}), tensorrt_llm::DataType::kFLOAT); auto refEntropyHost = this->mBufferManager->pinnedPool( - ITensor::makeShape({maxBatchSize, maxBeamWidth}), nvinfer1::DataType::kFLOAT); - auto entropyDevice - = this->mBufferManager->gpu(ITensor::makeShape({maxBatchSize, maxBeamWidth}), nvinfer1::DataType::kFLOAT); + ITensor::makeShape({maxBatchSize, maxBeamWidth}), tensorrt_llm::DataType::kFLOAT); + auto entropyDevice = this->mBufferManager->gpu( + ITensor::makeShape({maxBatchSize, maxBeamWidth}), tensorrt_llm::DataType::kFLOAT); auto biasHost = this->mBufferManager->pinnedPool(ITensor::makeShape({vocabSize}), dataType); auto temperatureHost - = this->mBufferManager->pinnedPool(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kFLOAT); + = this->mBufferManager->pinnedPool(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kFLOAT); auto endIdsHost - = this->mBufferManager->pinnedPool(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kINT32); + = this->mBufferManager->pinnedPool(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kINT32); auto beamWidthsHost - = this->mBufferManager->pinnedPool(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kINT32); + = this->mBufferManager->pinnedPool(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kINT32); ITensor::SharedPtr finishedHost = this->mBufferManager->pinnedPool( ITensor::makeShape({maxBeamWidth, maxBatchSize}), TRTDataType::value); - auto batchSlots = this->mBufferManager->pinnedPool(ITensor::makeShape({batchSize}), nvinfer1::DataType::kINT32); + auto batchSlots + = this->mBufferManager->pinnedPool(ITensor::makeShape({batchSize}), tensorrt_llm::DataType::kINT32); auto batchSlotsPtr = bufferCast(*batchSlots); auto beamWidthsHostPtr = bufferCast(*beamWidthsHost); diff --git a/cpp/tests/unit_tests/kernels/shiftKCacheKernelTest.cu b/cpp/tests/unit_tests/kernels/shiftKCacheKernelTest.cu index b1b3bd6234bf..75fede666434 100644 --- a/cpp/tests/unit_tests/kernels/shiftKCacheKernelTest.cu +++ b/cpp/tests/unit_tests/kernels/shiftKCacheKernelTest.cu @@ -2,6 +2,7 @@ #include "tensorrt_llm/common/cudaUtils.h" #include "tensorrt_llm/common/memoryUtils.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/kernels/decoderMaskedMultiheadAttentionUtils.h" #include "tensorrt_llm/kernels/gptKernels.h" #include "tensorrt_llm/kernels/kvCacheUtils.h" @@ -197,37 +198,37 @@ public: std::vector const& tokenSeqIdxs) { // allocate buffer - mSeqLengthsHost = mBufferManager->pinned(ITensor::makeShape({batchSize}), nvinfer1::DataType::kINT32); - mSeqLengthsDevice = mBufferManager->gpu(ITensor::makeShape({batchSize}), nvinfer1::DataType::kINT32); + mSeqLengthsHost = mBufferManager->pinned(ITensor::makeShape({batchSize}), tensorrt_llm::DataType::kINT32); + mSeqLengthsDevice = mBufferManager->gpu(ITensor::makeShape({batchSize}), tensorrt_llm::DataType::kINT32); - mInputLengthsHost = mBufferManager->pinned(ITensor::makeShape({batchSize}), nvinfer1::DataType::kINT32); - mInputLengthsDevice = mBufferManager->gpu(ITensor::makeShape({batchSize}), nvinfer1::DataType::kINT32); + mInputLengthsHost = mBufferManager->pinned(ITensor::makeShape({batchSize}), tensorrt_llm::DataType::kINT32); + mInputLengthsDevice = mBufferManager->gpu(ITensor::makeShape({batchSize}), tensorrt_llm::DataType::kINT32); - mKScaleQuantOrigDevice = mBufferManager->gpu(ITensor::makeShape({1}), nvinfer1::DataType::kFLOAT); + mKScaleQuantOrigDevice = mBufferManager->gpu(ITensor::makeShape({1}), tensorrt_llm::DataType::kFLOAT); mTokenReadIdxsHost = mBufferManager->pinned( - ITensor::makeShape({static_cast(tokenReadIdxs.size())}), nvinfer1::DataType::kINT32); + ITensor::makeShape({static_cast(tokenReadIdxs.size())}), tensorrt_llm::DataType::kINT32); mTokenReadIdxsDevice = mBufferManager->gpu( - ITensor::makeShape({static_cast(tokenReadIdxs.size())}), nvinfer1::DataType::kINT32); + ITensor::makeShape({static_cast(tokenReadIdxs.size())}), tensorrt_llm::DataType::kINT32); mTokenWriteIdxsHost = mBufferManager->pinned( - ITensor::makeShape({static_cast(tokenWriteIdxs.size())}), nvinfer1::DataType::kINT32); + ITensor::makeShape({static_cast(tokenWriteIdxs.size())}), tensorrt_llm::DataType::kINT32); mTokenWriteIdxsDevice = mBufferManager->gpu( - ITensor::makeShape({static_cast(tokenWriteIdxs.size())}), nvinfer1::DataType::kINT32); + ITensor::makeShape({static_cast(tokenWriteIdxs.size())}), tensorrt_llm::DataType::kINT32); mTokenPosIdxsHost = mBufferManager->pinned( - ITensor::makeShape({static_cast(tokenPosIdxs.size())}), nvinfer1::DataType::kINT32); + ITensor::makeShape({static_cast(tokenPosIdxs.size())}), tensorrt_llm::DataType::kINT32); mTokenPosIdxsDevice = mBufferManager->gpu( - ITensor::makeShape({static_cast(tokenPosIdxs.size())}), nvinfer1::DataType::kINT32); + ITensor::makeShape({static_cast(tokenPosIdxs.size())}), tensorrt_llm::DataType::kINT32); mTokenSeqIdxsHost = mBufferManager->pinned( - ITensor::makeShape({static_cast(tokenSeqIdxs.size())}), nvinfer1::DataType::kINT32); + ITensor::makeShape({static_cast(tokenSeqIdxs.size())}), tensorrt_llm::DataType::kINT32); mTokenSeqIdxsDevice = mBufferManager->gpu( - ITensor::makeShape({static_cast(tokenSeqIdxs.size())}), nvinfer1::DataType::kINT32); + ITensor::makeShape({static_cast(tokenSeqIdxs.size())}), tensorrt_llm::DataType::kINT32); - // nvinfer1::DataType dataType = nvinfer1::DataType::kHALF - // nvinfer1::DataType::kHALF - // nvinfer1::DataType::kBF16 + // tensorrt_llm::DataType dataType = tensorrt_llm::DataType::kHALF + // tensorrt_llm::DataType::kHALF + // tensorrt_llm::DataType::kBF16 int32_t batchBeam = batchSize * beamWidth; if (pagedKvCache) { diff --git a/cpp/tests/unit_tests/kernels/sparseAttentionKernelsTest.cpp b/cpp/tests/unit_tests/kernels/sparseAttentionKernelsTest.cpp index 508d157e1f61..14aa58f04df5 100644 --- a/cpp/tests/unit_tests/kernels/sparseAttentionKernelsTest.cpp +++ b/cpp/tests/unit_tests/kernels/sparseAttentionKernelsTest.cpp @@ -1,5 +1,6 @@ #include +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/kernels/sparseAttentionKernels.h" #include "tensorrt_llm/runtime/bufferManager.h" #include "tensorrt_llm/runtime/cudaStream.h" @@ -40,28 +41,29 @@ TEST_F(sparseAttentionKernelsTest, GatherKvPageOffsetsKernelTest) constexpr int total_sparse_tokens = 14; // Create input buffers - auto kv_page_offsets - = mBufferManager->gpu(ITensor::makeShape({batch_size, 2, max_num_pages_per_seq}), nvinfer1::DataType::kINT32); - auto seq_lengths = mBufferManager->gpu(ITensor::makeShape({batch_size}), nvinfer1::DataType::kINT32); + auto kv_page_offsets = mBufferManager->gpu( + ITensor::makeShape({batch_size, 2, max_num_pages_per_seq}), tensorrt_llm::DataType::kINT32); + auto seq_lengths = mBufferManager->gpu(ITensor::makeShape({batch_size}), tensorrt_llm::DataType::kINT32); // Shape: [num_head_kv, total_sparse_tokens] - flattened across all batches auto sparse_indices - = mBufferManager->gpu(ITensor::makeShape({num_head_kv, total_sparse_tokens}), nvinfer1::DataType::kINT32); - auto sparse_indices_offsets = mBufferManager->gpu(ITensor::makeShape({batch_size + 1}), nvinfer1::DataType::kINT32); + = mBufferManager->gpu(ITensor::makeShape({num_head_kv, total_sparse_tokens}), tensorrt_llm::DataType::kINT32); + auto sparse_indices_offsets + = mBufferManager->gpu(ITensor::makeShape({batch_size + 1}), tensorrt_llm::DataType::kINT32); // Create output buffers auto output_kv_page_offsets = mBufferManager->gpu( - ITensor::makeShape({num_head_kv, batch_size, 2, max_num_pages_per_seq}), nvinfer1::DataType::kINT32); + ITensor::makeShape({num_head_kv, batch_size, 2, max_num_pages_per_seq}), tensorrt_llm::DataType::kINT32); auto output_seq_lengths - = mBufferManager->gpu(ITensor::makeShape({num_head_kv, batch_size}), nvinfer1::DataType::kINT32); + = mBufferManager->gpu(ITensor::makeShape({num_head_kv, batch_size}), tensorrt_llm::DataType::kINT32); // Create pinned host buffers for data initialization auto kv_page_offsets_host = mBufferManager->pinned( - ITensor::makeShape({batch_size, 2, max_num_pages_per_seq}), nvinfer1::DataType::kINT32); - auto seq_lengths_host = mBufferManager->pinned(ITensor::makeShape({batch_size}), nvinfer1::DataType::kINT32); - auto sparse_indices_host - = mBufferManager->pinned(ITensor::makeShape({num_head_kv, total_sparse_tokens}), nvinfer1::DataType::kINT32); + ITensor::makeShape({batch_size, 2, max_num_pages_per_seq}), tensorrt_llm::DataType::kINT32); + auto seq_lengths_host = mBufferManager->pinned(ITensor::makeShape({batch_size}), tensorrt_llm::DataType::kINT32); + auto sparse_indices_host = mBufferManager->pinned( + ITensor::makeShape({num_head_kv, total_sparse_tokens}), tensorrt_llm::DataType::kINT32); auto sparse_indices_offsets_host - = mBufferManager->pinned(ITensor::makeShape({batch_size + 1}), nvinfer1::DataType::kINT32); + = mBufferManager->pinned(ITensor::makeShape({batch_size + 1}), tensorrt_llm::DataType::kINT32); // Initialize test data auto kv_page_offsets_ptr = bufferCast(*kv_page_offsets_host); @@ -143,9 +145,9 @@ TEST_F(sparseAttentionKernelsTest, GatherKvPageOffsetsKernelTest) // Copy results back to host for verification auto output_kv_page_offsets_host = mBufferManager->pinned( - ITensor::makeShape({num_head_kv, batch_size, 2, max_num_pages_per_seq}), nvinfer1::DataType::kINT32); + ITensor::makeShape({num_head_kv, batch_size, 2, max_num_pages_per_seq}), tensorrt_llm::DataType::kINT32); auto output_seq_lengths_host - = mBufferManager->pinned(ITensor::makeShape({num_head_kv, batch_size}), nvinfer1::DataType::kINT32); + = mBufferManager->pinned(ITensor::makeShape({num_head_kv, batch_size}), tensorrt_llm::DataType::kINT32); mBufferManager->copy(*output_kv_page_offsets, *output_kv_page_offsets_host); mBufferManager->copy(*output_seq_lengths, *output_seq_lengths_host); diff --git a/cpp/tests/unit_tests/kernels/stopCriteriaKernelsTest.cpp b/cpp/tests/unit_tests/kernels/stopCriteriaKernelsTest.cpp index 2fefae39552e..9fe4737e82c0 100644 --- a/cpp/tests/unit_tests/kernels/stopCriteriaKernelsTest.cpp +++ b/cpp/tests/unit_tests/kernels/stopCriteriaKernelsTest.cpp @@ -16,6 +16,7 @@ #include "tensorrt_llm/kernels/stopCriteriaKernels.h" #include "tensorrt_llm/common/memoryUtils.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/kernels/decodingCommon.h" #include "tensorrt_llm/runtime/bufferManager.h" @@ -60,34 +61,35 @@ class StopCriteriaKernelsTest : public testing::Test std::uniform_int_distribution tokensPerStepDistr(1, mMaxTokensPerStep); mSequenceLengths - = BufferManager::pinned(ITensor::makeShape({maxBatchSize, beamWidth}), nvinfer1::DataType::kINT32); - mSequenceLengthLimits = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kINT32); + = BufferManager::pinned(ITensor::makeShape({maxBatchSize, beamWidth}), tensorrt_llm::DataType::kINT32); + mSequenceLengthLimits + = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kINT32); mFinished = BufferManager::pinned( ITensor::makeShape({maxBatchSize, beamWidth}), TRTDataType::value); - mFinishedSum = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kINT32); + mFinishedSum = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kINT32); mOutputIds = BufferManager::pinned( - ITensor::makeShape({maxBatchSize, beamWidth, mMaxSeqLen}), nvinfer1::DataType::kINT32); + ITensor::makeShape({maxBatchSize, beamWidth, mMaxSeqLen}), tensorrt_llm::DataType::kINT32); mOutputIdsPtr - = BufferManager::pinned(ITensor::makeShape({maxBatchSize, beamWidth}), nvinfer1::DataType::kINT64); + = BufferManager::pinned(ITensor::makeShape({maxBatchSize, beamWidth}), tensorrt_llm::DataType::kINT64); mParentIds = BufferManager::pinned( - ITensor::makeShape({maxBatchSize, beamWidth, mMaxSeqLen}), nvinfer1::DataType::kINT32); + ITensor::makeShape({maxBatchSize, beamWidth, mMaxSeqLen}), tensorrt_llm::DataType::kINT32); mParentIdsPtr - = BufferManager::pinned(ITensor::makeShape({maxBatchSize, beamWidth}), nvinfer1::DataType::kINT64); + = BufferManager::pinned(ITensor::makeShape({maxBatchSize, beamWidth}), tensorrt_llm::DataType::kINT64); mRefOutputIds = BufferManager::pinned( - ITensor::makeShape({maxBatchSize, beamWidth, mMaxSeqLen}), nvinfer1::DataType::kINT32); + ITensor::makeShape({maxBatchSize, beamWidth, mMaxSeqLen}), tensorrt_llm::DataType::kINT32); - mStopWords - = BufferManager::pinned(ITensor::makeShape({maxBatchSize, 2, maxStopWordsLen}), nvinfer1::DataType::kINT32); - mStopWordsPtr = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kINT64); - mStopWordsLen = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kINT32); + mStopWords = BufferManager::pinned( + ITensor::makeShape({maxBatchSize, 2, maxStopWordsLen}), tensorrt_llm::DataType::kINT32); + mStopWordsPtr = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kINT64); + mStopWordsLen = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kINT32); - mBatchSlots = BufferManager::pinned(ITensor::makeShape({batchSize}), nvinfer1::DataType::kINT32); + mBatchSlots = BufferManager::pinned(ITensor::makeShape({batchSize}), tensorrt_llm::DataType::kINT32); - mEndIds = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kINT32); - mTokensPerStep = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kINT32); + mEndIds = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kINT32); + mTokensPerStep = BufferManager::pinned(ITensor::makeShape({maxBatchSize}), tensorrt_llm::DataType::kINT32); auto batchSlotsPtr = bufferCast(*mBatchSlots); for (SizeType32 bi = 0; bi < batchSize; ++bi) diff --git a/cpp/tests/unit_tests/layers/baseSamplingLayerTest.cpp b/cpp/tests/unit_tests/layers/baseSamplingLayerTest.cpp index 8bdacb2e9f6c..7886a6b54e2b 100644 --- a/cpp/tests/unit_tests/layers/baseSamplingLayerTest.cpp +++ b/cpp/tests/unit_tests/layers/baseSamplingLayerTest.cpp @@ -15,6 +15,7 @@ */ #include "tests/unit_tests/layers/baseSamplingLayerTest.h" +#include "tensorrt_llm/common/tllmDataType.h" namespace tensorrt_llm::tests::layers::sampling { @@ -67,26 +68,26 @@ void BaseSamplingLayerTest::setup(uint64_t seed, TestSamplingParams const& pa BaseSamplingLayerTest::mMaxOutputLen * params.beamWidth, mVocabSize); } - mSeqLengthsDevice = mBufferManager->gpu(ITensor::makeShape({maxBatchSize()}), nvinfer1::DataType::kINT32); - mContextLengthDevice = mBufferManager->gpu(ITensor::makeShape({maxBatchSize()}), nvinfer1::DataType::kINT32); + mSeqLengthsDevice = mBufferManager->gpu(ITensor::makeShape({maxBatchSize()}), tensorrt_llm::DataType::kINT32); + mContextLengthDevice = mBufferManager->gpu(ITensor::makeShape({maxBatchSize()}), tensorrt_llm::DataType::kINT32); mFinishedDevice = params.isExternalDraftTokensLayerTest ? mBufferManager->gpu(ITensor::makeShape({mMaxTokensPerEngineStep, maxBatchSize()}), TRTDataType::value) : mBufferManager->gpu( ITensor::makeShape({maxBatchSize()}), TRTDataType::value); - mOutputIdsDevice - = mBufferManager->gpu(ITensor::makeShape({maxBatchSize(), mBeamWidth, mMaxSeqLen}), nvinfer1::DataType::kINT32); - mEndIdsDevice = mBufferManager->gpu(ITensor::makeShape({maxBatchSize()}), nvinfer1::DataType::kINT32); + mOutputIdsDevice = mBufferManager->gpu( + ITensor::makeShape({maxBatchSize(), mBeamWidth, mMaxSeqLen}), tensorrt_llm::DataType::kINT32); + mEndIdsDevice = mBufferManager->gpu(ITensor::makeShape({maxBatchSize()}), tensorrt_llm::DataType::kINT32); mIdsPtrHost = mBufferManager->pinned(ITensor::makeShape({maxBatchSize()}), ptrType); - mCumLogProbsDevice = mBufferManager->gpu(ITensor::makeShape({maxBatchSize()}), nvinfer1::DataType::kFLOAT); + mCumLogProbsDevice = mBufferManager->gpu(ITensor::makeShape({maxBatchSize()}), tensorrt_llm::DataType::kFLOAT); mOutputLogProbsDevice - = mBufferManager->gpu(ITensor::makeShape({maxBatchSize(), mMaxSeqLen}), nvinfer1::DataType::kFLOAT); + = mBufferManager->gpu(ITensor::makeShape({maxBatchSize(), mMaxSeqLen}), tensorrt_llm::DataType::kFLOAT); mBatchSlots - = mBufferManager->pinned(ITensor::makeShape({mBatchSize + mBatchSizeBadPad}), nvinfer1::DataType::kINT32); - mCurandStatesDevice - = mBufferManager->gpu(ITensor::makeShape({maxBatchSize(), sizeof(curandState_t)}), nvinfer1::DataType::kINT8); + = mBufferManager->pinned(ITensor::makeShape({mBatchSize + mBatchSizeBadPad}), tensorrt_llm::DataType::kINT32); + mCurandStatesDevice = mBufferManager->gpu( + ITensor::makeShape({maxBatchSize(), sizeof(curandState_t)}), tensorrt_llm::DataType::kINT8); auto const workspaceSize = mSamplingLayer->getWorkspaceSize(); @@ -152,11 +153,11 @@ void BaseSamplingLayerTest::setup(uint64_t seed, TestSamplingParams const& pa setupParams = samplingSetupParams; mSrcCacheIndirection = mBufferManager->gpu( - ITensor::makeShape({maxBatchSize(), mBeamWidth, mMaxSeqLen}), nvinfer1::DataType::kINT32); + ITensor::makeShape({maxBatchSize(), mBeamWidth, mMaxSeqLen}), tensorrt_llm::DataType::kINT32); mTgtCacheIndirection = mBufferManager->gpu( - ITensor::makeShape({maxBatchSize(), mBeamWidth, mMaxSeqLen}), nvinfer1::DataType::kINT32); + ITensor::makeShape({maxBatchSize(), mBeamWidth, mMaxSeqLen}), tensorrt_llm::DataType::kINT32); mParentIds = mBufferManager->gpu( - ITensor::makeShape({maxBatchSize(), mBeamWidth, mMaxSeqLen}), nvinfer1::DataType::kINT32); + ITensor::makeShape({maxBatchSize(), mBeamWidth, mMaxSeqLen}), tensorrt_llm::DataType::kINT32); auto constexpr nvTokenIdType = TRTDataType::value; auto constexpr nvSizeType = TRTDataType::value; diff --git a/cpp/tests/unit_tests/layers/baseSamplingLayerTest.h b/cpp/tests/unit_tests/layers/baseSamplingLayerTest.h index 5a375000a5e8..2bf782a5f5e8 100644 --- a/cpp/tests/unit_tests/layers/baseSamplingLayerTest.h +++ b/cpp/tests/unit_tests/layers/baseSamplingLayerTest.h @@ -34,7 +34,6 @@ #include "tensorrt_llm/runtime/bufferManager.h" #include "tensorrt_llm/runtime/cudaStream.h" #include "tensorrt_llm/runtime/runtimeKernels.h" -#include "tensorrt_llm/runtime/tllmLogger.h" #include "tensorrt_llm/common/tllmException.h" diff --git a/cpp/tests/unit_tests/layers/dynamicDecodeLayerTest.cpp b/cpp/tests/unit_tests/layers/dynamicDecodeLayerTest.cpp index a3c2d56de16e..cc3b9c411f5b 100644 --- a/cpp/tests/unit_tests/layers/dynamicDecodeLayerTest.cpp +++ b/cpp/tests/unit_tests/layers/dynamicDecodeLayerTest.cpp @@ -15,6 +15,7 @@ */ #include "tests/unit_tests/layers/dynamicDecodeLayerTest.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/executor/types.h" #include "tensorrt_llm/runtime/runtimeKernels.h" #include @@ -150,43 +151,43 @@ void DynamicDecodeLayerTest::allocateData(TestSamplingParams const& params, T mRuntimeLogitsHost = BufferManager::pinned(ITensor::makeShape({mBatchSize, mBeamWidth, mVocabSizePadded}), dataType); - mSeqLengthsDevice = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize}), nvinfer1::DataType::kINT32); - mContextLengthDevice = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize}), nvinfer1::DataType::kINT32); + mSeqLengthsDevice = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize}), tensorrt_llm::DataType::kINT32); + mContextLengthDevice = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize}), tensorrt_llm::DataType::kINT32); mFinishedDevice = mBufferManager->gpu( ITensor::makeShape({mMaxBatchSize}), TRTDataType::value); - mFinishedSumDevice = BufferManager::pinned(ITensor::makeShape({1}), nvinfer1::DataType::kFLOAT); - mOutputIdsDevice - = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize, mBeamWidth, mMaxSeqLen}), nvinfer1::DataType::kINT32); + mFinishedSumDevice = BufferManager::pinned(ITensor::makeShape({1}), tensorrt_llm::DataType::kFLOAT); + mOutputIdsDevice = mBufferManager->gpu( + ITensor::makeShape({mMaxBatchSize, mBeamWidth, mMaxSeqLen}), tensorrt_llm::DataType::kINT32); mNewTokens - = BufferManager::pinned(ITensor::makeShape({mMaxTokensPerStep, mMaxBatchSize}), nvinfer1::DataType::kINT32); - mEndIdsDevice = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize}), nvinfer1::DataType::kINT32); + = BufferManager::pinned(ITensor::makeShape({mMaxTokensPerStep, mMaxBatchSize}), tensorrt_llm::DataType::kINT32); + mEndIdsDevice = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize}), tensorrt_llm::DataType::kINT32); mEmbeddingBiasHost = BufferManager::pinned(ITensor::makeShape({mMaxBatchSize, mVocabSizePadded}), dataType); mEmbeddingBiasDevice = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize, mVocabSizePadded}), dataType); mRefLogProbsHost - = BufferManager::pinned(ITensor::makeShape({mMaxBatchSize, mMaxSeqLen}), nvinfer1::DataType::kFLOAT); + = BufferManager::pinned(ITensor::makeShape({mMaxBatchSize, mMaxSeqLen}), tensorrt_llm::DataType::kFLOAT); mOutputLogProbsDevice - = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize, mMaxSeqLen}), nvinfer1::DataType::kFLOAT); + = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize, mMaxSeqLen}), tensorrt_llm::DataType::kFLOAT); mOutputLogProbsTiledDevice - = mBufferManager->gpu(ITensor::makeShape({mMaxSeqLen, mMaxBatchSize}), nvinfer1::DataType::kFLOAT); + = mBufferManager->gpu(ITensor::makeShape({mMaxSeqLen, mMaxBatchSize}), tensorrt_llm::DataType::kFLOAT); - mCumLogProbsDevice = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize}), nvinfer1::DataType::kFLOAT); + mCumLogProbsDevice = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize}), tensorrt_llm::DataType::kFLOAT); mMaxBadWordsLen = getMaxWordsLen(params.badWords); mMaxStopWordsLen = getMaxWordsLen(params.stopWords); - mBadWords - = BufferManager::pinned(ITensor::makeShape({mMaxBatchSize, 2, mMaxBadWordsLen}), nvinfer1::DataType::kINT32); - mBadWordsLens = BufferManager::pinned(ITensor::makeShape({mMaxBatchSize}), nvinfer1::DataType::kINT32); - mBadWordsPtrs = BufferManager::pinned(ITensor::makeShape({mMaxBatchSize}), nvinfer1::DataType::kINT64); + mBadWords = BufferManager::pinned( + ITensor::makeShape({mMaxBatchSize, 2, mMaxBadWordsLen}), tensorrt_llm::DataType::kINT32); + mBadWordsLens = BufferManager::pinned(ITensor::makeShape({mMaxBatchSize}), tensorrt_llm::DataType::kINT32); + mBadWordsPtrs = BufferManager::pinned(ITensor::makeShape({mMaxBatchSize}), tensorrt_llm::DataType::kINT64); - mStopWords - = BufferManager::pinned(ITensor::makeShape({mMaxBatchSize, 2, mMaxStopWordsLen}), nvinfer1::DataType::kINT32); - mStopWordsLens = BufferManager::pinned(ITensor::makeShape({mMaxBatchSize}), nvinfer1::DataType::kINT32); - mStopWordsPtrs = BufferManager::pinned(ITensor::makeShape({mMaxBatchSize}), nvinfer1::DataType::kINT64); + mStopWords = BufferManager::pinned( + ITensor::makeShape({mMaxBatchSize, 2, mMaxStopWordsLen}), tensorrt_llm::DataType::kINT32); + mStopWordsLens = BufferManager::pinned(ITensor::makeShape({mMaxBatchSize}), tensorrt_llm::DataType::kINT32); + mStopWordsPtrs = BufferManager::pinned(ITensor::makeShape({mMaxBatchSize}), tensorrt_llm::DataType::kINT64); - mBatchSlots = BufferManager::pinned(ITensor::makeShape({mBatchSize}), nvinfer1::DataType::kINT32); + mBatchSlots = BufferManager::pinned(ITensor::makeShape({mBatchSize}), tensorrt_llm::DataType::kINT32); if (mDecodingMode.isMedusa()) { @@ -204,19 +205,19 @@ void DynamicDecodeLayerTest::allocateMedusaData(TestSamplingParams const& par auto const dataType = TRTDataType::value; mMaxMedusaHeads = params.maxNumMedusaHeads.value(); mPathsDevice = mBufferManager->gpu( - ITensor::makeShape({mMaxBatchSize, mMaxTokensPerStep, mMaxMedusaHeads + 1}), nvinfer1::DataType::kINT32); - mAcceptedLengths = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize}), nvinfer1::DataType::kINT32); + ITensor::makeShape({mMaxBatchSize, mMaxTokensPerStep, mMaxMedusaHeads + 1}), tensorrt_llm::DataType::kINT32); + mAcceptedLengths = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize}), tensorrt_llm::DataType::kINT32); mMedusaLogitsDevice = BufferManager::pinned( ITensor::makeShape({mMaxMedusaHeads, mMaxBatchSize, mMaxTokensPerStep, mVocabSizePadded}), dataType); - mNextDraftTokensDevice - = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize, mMaxTokensPerStep - 1}), nvinfer1::DataType::kINT32); - mTokensPerStepDevice = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize}), nvinfer1::DataType::kINT32); - mTreeIdsDevice - = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize, mMaxTokensPerStep - 1}), nvinfer1::DataType::kINT32); + mNextDraftTokensDevice = mBufferManager->gpu( + ITensor::makeShape({mMaxBatchSize, mMaxTokensPerStep - 1}), tensorrt_llm::DataType::kINT32); + mTokensPerStepDevice = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize}), tensorrt_llm::DataType::kINT32); + mTreeIdsDevice = mBufferManager->gpu( + ITensor::makeShape({mMaxBatchSize, mMaxTokensPerStep - 1}), tensorrt_llm::DataType::kINT32); mAcceptedLengthCumSumDevice - = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize + 1}), nvinfer1::DataType::kINT32); + = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize + 1}), tensorrt_llm::DataType::kINT32); mPackedPathsDevice - = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize * mMaxMedusaHeads}), nvinfer1::DataType::kINT32); + = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize * mMaxMedusaHeads}), tensorrt_llm::DataType::kINT32); } template @@ -604,7 +605,7 @@ template void DynamicDecodeLayerTest::batchCopy(SizeType32 step) { auto const logitsHost = ITensor::wrap(mTestLogitsInit.data() + step * mVocabSizePadded, - std::is_same_v ? nvinfer1::DataType::kFLOAT : nvinfer1::DataType::kHALF, + std::is_same_v ? tensorrt_llm::DataType::kFLOAT : tensorrt_llm::DataType::kHALF, ITensor::makeShape({mMaxTokensPerStep, mVocabSizePadded})); for (SizeType32 bi = 0; bi < mBatchSize; ++bi) { diff --git a/cpp/tests/unit_tests/layers/eagleLayerTest.cpp b/cpp/tests/unit_tests/layers/eagleLayerTest.cpp index bdb53f15618f..47bca93b47b4 100644 --- a/cpp/tests/unit_tests/layers/eagleLayerTest.cpp +++ b/cpp/tests/unit_tests/layers/eagleLayerTest.cpp @@ -22,9 +22,8 @@ #include "tensorrt_llm/runtime/iTensor.h" #include "tensorrt_llm/runtime/runtimeKernels.h" #include "tensorrt_llm/runtime/speculativeDecodingModule.h" -#include "tensorrt_llm/runtime/tllmLogger.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include #include @@ -555,126 +554,126 @@ void EagleDecodingLayerTest::allocateBuffers() // outputs mOutputIds = BufferManager::pinnedPool( ITensor::makeShape({mSamplingParams.getMaxBatchSize(), mSamplingParams.getMaxSeqLen()}), - nvinfer1::DataType::kINT32); + tensorrt_llm::DataType::kINT32); mSeqLengths = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), nvinfer1::DataType::kINT32); + ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), tensorrt_llm::DataType::kINT32); mOutputNextDraftTokens = BufferManager::pinnedPool( ITensor::makeShape({mSamplingParams.getMaxBatchSize(), mSamplingParams.getMaxDecodingDraftTokens()}), - nvinfer1::DataType::kINT32); + tensorrt_llm::DataType::kINT32); mOutputUnpackedNextDraftTokens = BufferManager::pinnedPool( ITensor::makeShape({mSamplingParams.getMaxBatchSize(), mSamplingParams.getMaxDecodingDraftTokens()}), - nvinfer1::DataType::kINT32); + tensorrt_llm::DataType::kINT32); mAcceptedLengths = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), nvinfer1::DataType::kINT32); + ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), tensorrt_llm::DataType::kINT32); mNextPosIds = BufferManager::pinnedPool( ITensor::makeShape({mSamplingParams.getMaxBatchSize(), mSamplingParams.getMaxDecodingTokens()}), - nvinfer1::DataType::kINT32); + tensorrt_llm::DataType::kINT32); mPrevDraftLengths = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), nvinfer1::DataType::kINT32); + ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), tensorrt_llm::DataType::kINT32); mNextDraftLengths = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), nvinfer1::DataType::kINT32); + ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), tensorrt_llm::DataType::kINT32); mNextGenerationLengths - = mBufferManager->gpu(ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), nvinfer1::DataType::kINT32); + = mBufferManager->gpu(ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), tensorrt_llm::DataType::kINT32); mNextGenerationLengthsHost = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), nvinfer1::DataType::kINT32); + ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), tensorrt_llm::DataType::kINT32); mAcceptedLengthCumSum = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize() + 1}), nvinfer1::DataType::kINT32); + ITensor::makeShape({mSamplingParams.getMaxBatchSize() + 1}), tensorrt_llm::DataType::kINT32); mPathsOffsets = BufferManager::pinnedPool( ITensor::makeShape({mSamplingParams.getMaxBatchSize() * mSamplingParams.getMaxDraftPathLen()}), - nvinfer1::DataType::kINT32); + tensorrt_llm::DataType::kINT32); mPackedMasks = BufferManager::pinnedPool( ITensor::makeShape({mSamplingParams.getMaxBatchSize(), mSamplingParams.getMaxDecodingTokens(), static_cast(divUp(mSamplingParams.getMaxDecodingTokens(), 32))}), - nvinfer1::DataType::kINT32); + tensorrt_llm::DataType::kINT32); mRandomDataSample = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), nvinfer1::DataType::kFLOAT); + ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), tensorrt_llm::DataType::kFLOAT); mRandomDataValidation = BufferManager::pinnedPool( ITensor::makeShape({mSamplingParams.getMaxBatchSize(), mSamplingParams.getMaxDecodingTokens()}), - nvinfer1::DataType::kFLOAT); + tensorrt_llm::DataType::kFLOAT); mOutputTemperatures = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), nvinfer1::DataType::kFLOAT); + ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), tensorrt_llm::DataType::kFLOAT); mOutputNextDraftPaths = BufferManager::pinnedPool(ITensor::makeShape({mSamplingParams.getMaxBatchSize(), mSamplingParams.getMaxDecodingTokens(), mSamplingParams.getMaxPathLen()}), - nvinfer1::DataType::kINT32); + tensorrt_llm::DataType::kINT32); mEagleNetCtxRequestTypesHost = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), nvinfer1::DataType::kINT32); + ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), tensorrt_llm::DataType::kINT32); mEagleNetCtxContextLengthsHost = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), nvinfer1::DataType::kINT32); + ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), tensorrt_llm::DataType::kINT32); mEagleNetCtxPastKeyValueLengthsHost = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), nvinfer1::DataType::kINT32); + ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), tensorrt_llm::DataType::kINT32); mEagleNetGenRequestTypesHost = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), nvinfer1::DataType::kINT32); + ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), tensorrt_llm::DataType::kINT32); mEagleNetGenContextLengthsHost = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), nvinfer1::DataType::kINT32); + ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), tensorrt_llm::DataType::kINT32); mEagleNetGenPastKeyValueLengthsHost = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), nvinfer1::DataType::kINT32); + ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), tensorrt_llm::DataType::kINT32); // inputs - mBatchSlots - = BufferManager::pinnedPool(ITensor::makeShape({mSamplingParams.getBatchSize()}), nvinfer1::DataType::kINT32); + mBatchSlots = BufferManager::pinnedPool( + ITensor::makeShape({mSamplingParams.getBatchSize()}), tensorrt_llm::DataType::kINT32); - mEndIds - = BufferManager::pinnedPool(ITensor::makeShape({mSamplingParams.getBatchSize()}), nvinfer1::DataType::kINT32); + mEndIds = BufferManager::pinnedPool( + ITensor::makeShape({mSamplingParams.getBatchSize()}), tensorrt_llm::DataType::kINT32); mInputNextDraftTokens = BufferManager::pinnedPool( ITensor::makeShape({mSamplingParams.getBatchSize(), mSamplingParams.getMaxDecodingDraftTokens()}), - nvinfer1::DataType::kINT32); + tensorrt_llm::DataType::kINT32); - mInputNextDraftLens - = BufferManager::pinnedPool(ITensor::makeShape({mSamplingParams.getBatchSize()}), nvinfer1::DataType::kINT32); + mInputNextDraftLens = BufferManager::pinnedPool( + ITensor::makeShape({mSamplingParams.getBatchSize()}), tensorrt_llm::DataType::kINT32); mInputNextDraftPaths = BufferManager::pinnedPool( ITensor::makeShape( {mSamplingParams.getBatchSize(), mSamplingParams.getMaxDecodingTokens(), mSamplingParams.getMaxPathLen()}), - nvinfer1::DataType::kINT32); + tensorrt_llm::DataType::kINT32); mInputLastDraftTokens = BufferManager::pinnedPool( ITensor::makeShape({mSamplingParams.getBatchSize(), mSamplingParams.getMaxDecodingDraftTokens()}), - nvinfer1::DataType::kINT32); + tensorrt_llm::DataType::kINT32); - mInputLastDraftLens - = BufferManager::pinnedPool(ITensor::makeShape({mSamplingParams.getBatchSize()}), nvinfer1::DataType::kINT32); + mInputLastDraftLens = BufferManager::pinnedPool( + ITensor::makeShape({mSamplingParams.getBatchSize()}), tensorrt_llm::DataType::kINT32); mInputLastDraftPaths = BufferManager::pinnedPool( ITensor::makeShape( {mSamplingParams.getBatchSize(), mSamplingParams.getMaxDecodingTokens(), mSamplingParams.getMaxPathLen()}), - nvinfer1::DataType::kINT32); + tensorrt_llm::DataType::kINT32); mInputAcceptedTokens = BufferManager::pinnedPool( ITensor::makeShape({mSamplingParams.getBatchSize(), mSamplingParams.getMaxPathLen()}), - nvinfer1::DataType::kINT32); + tensorrt_llm::DataType::kINT32); - mInputAcceptedLens - = BufferManager::pinnedPool(ITensor::makeShape({mSamplingParams.getBatchSize()}), nvinfer1::DataType::kINT32); + mInputAcceptedLens = BufferManager::pinnedPool( + ITensor::makeShape({mSamplingParams.getBatchSize()}), tensorrt_llm::DataType::kINT32); - mInputAcceptedPathIds - = BufferManager::pinnedPool(ITensor::makeShape({mSamplingParams.getBatchSize()}), nvinfer1::DataType::kINT32); + mInputAcceptedPathIds = BufferManager::pinnedPool( + ITensor::makeShape({mSamplingParams.getBatchSize()}), tensorrt_llm::DataType::kINT32); - mChunkedContextNextTokens - = BufferManager::pinnedPool(ITensor::makeShape({mSamplingParams.getBatchSize()}), nvinfer1::DataType::kINT32); + mChunkedContextNextTokens = BufferManager::pinnedPool( + ITensor::makeShape({mSamplingParams.getBatchSize()}), tensorrt_llm::DataType::kINT32); mDecodingWorkspace = std::make_shared(mBufferManager, decodingDomain, TRTDataType::value, mSamplingParams.getMaxBatchSize() * sizeof(curandState_t)); diff --git a/cpp/tests/unit_tests/layers/explicitDraftTokensLayerTest.cpp b/cpp/tests/unit_tests/layers/explicitDraftTokensLayerTest.cpp index e7831b57f77f..04d05e0d16a9 100644 --- a/cpp/tests/unit_tests/layers/explicitDraftTokensLayerTest.cpp +++ b/cpp/tests/unit_tests/layers/explicitDraftTokensLayerTest.cpp @@ -22,9 +22,8 @@ #include "tensorrt_llm/runtime/iTensor.h" #include "tensorrt_llm/runtime/runtimeKernels.h" #include "tensorrt_llm/runtime/speculativeDecodingModule.h" -#include "tensorrt_llm/runtime/tllmLogger.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include #include @@ -655,29 +654,29 @@ void ExplicitDraftTokensLayerTest::allocateBuffers() // outputs mOutputIds = BufferManager::pinnedPool( ITensor::makeShape({mSamplingParams.getMaxBatchSize(), mSamplingParams.getMaxSeqLen()}), - nvinfer1::DataType::kINT32); + tensorrt_llm::DataType::kINT32); mSeqLengths = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), nvinfer1::DataType::kINT32); + ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), tensorrt_llm::DataType::kINT32); mAcceptedLengths = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), nvinfer1::DataType::kINT32); + ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), tensorrt_llm::DataType::kINT32); mNextDraftLengths = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), nvinfer1::DataType::kINT32); + ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), tensorrt_llm::DataType::kINT32); mPrevDraftLengths = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), nvinfer1::DataType::kINT32); + ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), tensorrt_llm::DataType::kINT32); mAcceptedLengthCumSum = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize() + 1}), nvinfer1::DataType::kINT32); + ITensor::makeShape({mSamplingParams.getMaxBatchSize() + 1}), tensorrt_llm::DataType::kINT32); mOutputNextDraftTokens = BufferManager::pinnedPool( ITensor::makeShape({mSamplingParams.getMaxBatchSize(), mSamplingParams.getMaxDecodingDraftTokens()}), - nvinfer1::DataType::kINT32); + tensorrt_llm::DataType::kINT32); mOutputPositionIdsBase = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), nvinfer1::DataType::kINT32); + ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), tensorrt_llm::DataType::kINT32); mRandomDataSample = BufferManager::pinnedPool(ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), dataType); @@ -689,21 +688,21 @@ void ExplicitDraftTokensLayerTest::allocateBuffers() mPackedMasks = BufferManager::pinnedPool( ITensor::makeShape({mSamplingParams.getMaxBatchSize(), mSamplingParams.getMaxDecodingTokens(), static_cast(divUp(mSamplingParams.getMaxDecodingTokens(), 32))}), - nvinfer1::DataType::kINT32); + tensorrt_llm::DataType::kINT32); mNextPosIds = BufferManager::pinnedPool( ITensor::makeShape({mSamplingParams.getMaxBatchSize(), mSamplingParams.getMaxDecodingTokens()}), - nvinfer1::DataType::kINT32); + tensorrt_llm::DataType::kINT32); mOutputUnpackedNextDraftTokens = BufferManager::pinnedPool( ITensor::makeShape( {mSamplingParams.getMaxBatchSize(), mSamplingParams.getMaxNumPaths(), mSamplingParams.getMaxPathLen()}), - nvinfer1::DataType::kINT32); + tensorrt_llm::DataType::kINT32); mOutputUnpackedNextDraftIndices = BufferManager::pinnedPool( ITensor::makeShape( {mSamplingParams.getMaxBatchSize(), mSamplingParams.getMaxNumPaths(), mSamplingParams.getMaxPathLen()}), - nvinfer1::DataType::kINT32); + tensorrt_llm::DataType::kINT32); mOutputDraftProbs = BufferManager::pinnedPool( ITensor::makeShape({mSamplingParams.getMaxBatchSize(), mSamplingParams.getMaxNumPaths(), @@ -713,68 +712,68 @@ void ExplicitDraftTokensLayerTest::allocateBuffers() mOutputTemperatures = BufferManager::pinnedPool(ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), dataType); mOutputGenerationLengths = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), nvinfer1::DataType::kINT32); + ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), tensorrt_llm::DataType::kINT32); mOutputGenerationLengthsHost = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), nvinfer1::DataType::kINT32); + ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), tensorrt_llm::DataType::kINT32); - mMaxGenLengthHost = BufferManager::pinnedPool(ITensor::makeShape({1}), nvinfer1::DataType::kINT32); + mMaxGenLengthHost = BufferManager::pinnedPool(ITensor::makeShape({1}), tensorrt_llm::DataType::kINT32); // inputs - mBatchSlots - = BufferManager::pinnedPool(ITensor::makeShape({mSamplingParams.getBatchSize()}), nvinfer1::DataType::kINT32); + mBatchSlots = BufferManager::pinnedPool( + ITensor::makeShape({mSamplingParams.getBatchSize()}), tensorrt_llm::DataType::kINT32); mTokensPerStep = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), nvinfer1::DataType::kINT32); + ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), tensorrt_llm::DataType::kINT32); mPathsOffsets = BufferManager::pinnedPool( ITensor::makeShape({mSamplingParams.getMaxBatchSize() * mSamplingParams.getMaxDraftPathLen()}), - nvinfer1::DataType::kINT32); + tensorrt_llm::DataType::kINT32); mMasks = BufferManager::pinnedPool( ITensor::makeShape({mSamplingParams.getMaxBatchSize(), mSamplingParams.getMaxDecodingTokens(), mSamplingParams.getMaxDecodingTokens()}), - nvinfer1::DataType::kBOOL); + tensorrt_llm::DataType::kBOOL); mInputNextDraftTokens = BufferManager::pinnedPool( ITensor::makeShape( {mSamplingParams.getMaxBatchSize(), mSamplingParams.getMaxNumPaths(), mSamplingParams.getMaxPathLen()}), - nvinfer1::DataType::kINT32); + tensorrt_llm::DataType::kINT32); mLastDraftTokens = BufferManager::pinnedPool( ITensor::makeShape( {mSamplingParams.getBatchSize(), mSamplingParams.getMaxNumPaths(), mSamplingParams.getMaxPathLen()}), - nvinfer1::DataType::kINT32); + tensorrt_llm::DataType::kINT32); mPackedPosIds = BufferManager::pinnedPool( ITensor::makeShape({mSamplingParams.getMaxBatchSize(), mSamplingParams.getMaxDecodingTokens()}), - nvinfer1::DataType::kINT32); + tensorrt_llm::DataType::kINT32); mBestPathLengths = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), nvinfer1::DataType::kINT32); + ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), tensorrt_llm::DataType::kINT32); mBestPathIndices = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), nvinfer1::DataType::kINT32); + ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), tensorrt_llm::DataType::kINT32); mSpecDecodingGenerationLengths = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), nvinfer1::DataType::kINT32); + ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), tensorrt_llm::DataType::kINT32); mNextFlatTokens = BufferManager::pinnedPool( ITensor::makeShape({mSamplingParams.getMaxBatchSize() * mSamplingParams.getMaxDecodingTokens()}), - nvinfer1::DataType::kINT32); + tensorrt_llm::DataType::kINT32); mInputPositionIdsBase = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), nvinfer1::DataType::kINT32); + ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), tensorrt_llm::DataType::kINT32); mNextDraftIndices = BufferManager::pinnedPool( ITensor::makeShape( {mSamplingParams.getMaxBatchSize(), mSamplingParams.getMaxNumPaths(), mSamplingParams.getMaxPathLen()}), - nvinfer1::DataType::kINT32); + tensorrt_llm::DataType::kINT32); mLastDraftIndices = BufferManager::pinnedPool( ITensor::makeShape( {mSamplingParams.getMaxBatchSize(), mSamplingParams.getMaxNumPaths(), mSamplingParams.getMaxPathLen()}), - nvinfer1::DataType::kINT32); + tensorrt_llm::DataType::kINT32); mNextDraftProbs = BufferManager::pinnedPool( ITensor::makeShape({mSamplingParams.getMaxBatchSize(), mSamplingParams.getMaxNumPaths(), @@ -782,20 +781,20 @@ void ExplicitDraftTokensLayerTest::allocateBuffers() dataType); mEndIds = BufferManager::pinnedPool( - ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), nvinfer1::DataType::kINT32); + ITensor::makeShape({mSamplingParams.getMaxBatchSize()}), tensorrt_llm::DataType::kINT32); - mMaxGenLengthDevice = BufferManager::pinnedPool(ITensor::makeShape({1}), nvinfer1::DataType::kINT32); + mMaxGenLengthDevice = BufferManager::pinnedPool(ITensor::makeShape({1}), tensorrt_llm::DataType::kINT32); // Packed inputs - mMaxGenerationLength = BufferManager::pinnedPool(ITensor::makeShape({1}), nvinfer1::DataType::kINT32); - mCumSumGenerationLengths - = BufferManager::pinnedPool(ITensor::makeShape({mSamplingParams.getBatchSize()}), nvinfer1::DataType::kINT32); + mMaxGenerationLength = BufferManager::pinnedPool(ITensor::makeShape({1}), tensorrt_llm::DataType::kINT32); + mCumSumGenerationLengths = BufferManager::pinnedPool( + ITensor::makeShape({mSamplingParams.getBatchSize()}), tensorrt_llm::DataType::kINT32); // Packed outputs - mPackedPositionIdsBase - = BufferManager::pinnedPool(ITensor::makeShape({mSamplingParams.getBatchSize()}), nvinfer1::DataType::kINT32); - mPackedGenerationLengths - = BufferManager::pinnedPool(ITensor::makeShape({mSamplingParams.getBatchSize()}), nvinfer1::DataType::kINT32); + mPackedPositionIdsBase = BufferManager::pinnedPool( + ITensor::makeShape({mSamplingParams.getBatchSize()}), tensorrt_llm::DataType::kINT32); + mPackedGenerationLengths = BufferManager::pinnedPool( + ITensor::makeShape({mSamplingParams.getBatchSize()}), tensorrt_llm::DataType::kINT32); mPackedRandomDataSample = BufferManager::pinnedPool(ITensor::makeShape({mSamplingParams.getBatchSize()}), dataType); mPackedRandomDataVerification = BufferManager::pinnedPool( ITensor::makeShape( @@ -804,21 +803,21 @@ void ExplicitDraftTokensLayerTest::allocateBuffers() mPackedNextDraftTokens = BufferManager::pinnedPool( ITensor::makeShape( {mSamplingParams.getBatchSize(), mSamplingParams.getMaxNumPaths(), mSamplingParams.getMaxPathLen()}), - nvinfer1::DataType::kINT32); + tensorrt_llm::DataType::kINT32); mPackedNextDraftIndices = BufferManager::pinnedPool( ITensor::makeShape( {mSamplingParams.getBatchSize(), mSamplingParams.getMaxNumPaths(), mSamplingParams.getMaxPathLen()}), - nvinfer1::DataType::kINT32); + tensorrt_llm::DataType::kINT32); mPackedPackedMasks = BufferManager::pinnedPool( ITensor::makeShape({mSamplingParams.getBatchSize(), mSamplingParams.getMaxDecodingTokens(), static_cast(divUp(mSamplingParams.getMaxDecodingTokens(), 32))}), - nvinfer1::DataType::kINT32); + tensorrt_llm::DataType::kINT32); mPackedPositionOffsets = BufferManager::pinnedPool( ITensor::makeShape({mSamplingParams.getBatchSize(), mSamplingParams.getMaxDecodingTokens()}), - nvinfer1::DataType::kINT32); + tensorrt_llm::DataType::kINT32); mPackedPackedPosIds = BufferManager::pinnedPool( ITensor::makeShape({mSamplingParams.getBatchSize(), mSamplingParams.getMaxDecodingTokens()}), - nvinfer1::DataType::kINT32); + tensorrt_llm::DataType::kINT32); mPackedDraftProbs = BufferManager::pinnedPool( ITensor::makeShape({mSamplingParams.getBatchSize(), mSamplingParams.getMaxNumPaths(), mSamplingParams.getMaxDraftPathLen(), mSamplingParams.getVocabSize()}), @@ -1564,7 +1563,6 @@ class FillRandDataTest : public ::testing::Test // NOLINT(cppcoreguidelines-pro- void SetUp() override { - mLogger = std::make_shared(); mStream = std::make_shared(); mBufferManager = std::make_shared(mStream); } @@ -1576,12 +1574,12 @@ class FillRandDataTest : public ::testing::Test // NOLINT(cppcoreguidelines-pro- { SizeType32* batchSlotsPtr{nullptr}; - auto curandState = mBufferManager->gpu(ITensor::makeShape({batchSize, 48}), nvinfer1::DataType::kUINT8); + auto curandState = mBufferManager->gpu(ITensor::makeShape({batchSize, 48}), tensorrt_llm::DataType::kUINT8); auto* curandStatePtr = reinterpret_cast(bufferCast(*curandState)); if (batchInit) { - auto randomSeeds = mBufferManager->gpu(ITensor::makeShape({batchSize}), nvinfer1::DataType::kINT64); + auto randomSeeds = mBufferManager->gpu(ITensor::makeShape({batchSize}), tensorrt_llm::DataType::kINT64); trk::invokeFill(*randomSeeds, static_cast(randomSeed), *mStream); auto* randomSeedsPtr = bufferCast(*randomSeeds); tk::invokeCurandBatchInitialize(curandStatePtr, batchSlotsPtr, batchSize, randomSeedsPtr, mStream->get()); @@ -1625,7 +1623,6 @@ class FillRandDataTest : public ::testing::Test // NOLINT(cppcoreguidelines-pro- } private: - std::shared_ptr mLogger; std::shared_ptr mStream; std::shared_ptr mBufferManager; }; diff --git a/cpp/tests/unit_tests/layers/externalDraftTokensLayerTest.cpp b/cpp/tests/unit_tests/layers/externalDraftTokensLayerTest.cpp index c216a76bc2b8..cdc913d3b3aa 100644 --- a/cpp/tests/unit_tests/layers/externalDraftTokensLayerTest.cpp +++ b/cpp/tests/unit_tests/layers/externalDraftTokensLayerTest.cpp @@ -15,6 +15,7 @@ */ #include "tensorrt_llm/common/assert.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/executor/types.h" #include "tensorrt_llm/runtime/iBuffer.h" #include "tests/unit_tests/layers/baseSamplingLayerTest.h" @@ -88,7 +89,7 @@ class ExternalDraftTokensLayerTest : public BaseSamplingLayerTest dataType); mDraftTokenIds = this->mBufferManager->gpu( - ITensor::makeShape({this->maxBatchSize(), mMaxDraftLen}), nvinfer1::DataType::kINT32); + ITensor::makeShape({this->maxBatchSize(), mMaxDraftLen}), tensorrt_llm::DataType::kINT32); mUseDraftLogits = this->mBufferManager->gpu(ITensor::makeShape({this->maxBatchSize()}), TRTDataType::value); mUseDraftLogitsHost diff --git a/cpp/tests/unit_tests/layers/lookaheadAlgorithmTest.cpp b/cpp/tests/unit_tests/layers/lookaheadAlgorithmTest.cpp index cd1dc4799e6c..7f4791a85ba4 100644 --- a/cpp/tests/unit_tests/layers/lookaheadAlgorithmTest.cpp +++ b/cpp/tests/unit_tests/layers/lookaheadAlgorithmTest.cpp @@ -17,6 +17,7 @@ #include #include "tensorrt_llm/common/logger.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/executor/executor.h" #include "tensorrt_llm/layers/lookaheadAlgorithm.h" #include "tensorrt_llm/layers/lookaheadDecodingUtils.h" @@ -87,24 +88,25 @@ TEST_P(LookaheadAlgorithmTest, predict) auto shape = ITensor::makeShape({maxTokensPerStep}); auto shape2d = ITensor::makeShape({maxTokensPerStep, maxTokensPerStep}); auto shapeSingle = ITensor::makeShape({1}); - TensorPtr posidMax = BufferManager::cpu(shape, nvinfer1::DataType::kINT32); - TensorPtr attentionMaskMax = BufferManager::cpu(shape2d, nvinfer1::DataType::kBOOL); - TensorPtr inputLengthPtr = BufferManager::cpu(shapeSingle, nvinfer1::DataType::kINT32); + TensorPtr posidMax = BufferManager::cpu(shape, tensorrt_llm::DataType::kINT32); + TensorPtr attentionMaskMax = BufferManager::cpu(shape2d, tensorrt_llm::DataType::kBOOL); + TensorPtr inputLengthPtr = BufferManager::cpu(shapeSingle, tensorrt_llm::DataType::kINT32); auto& inputLength(*BufferRange(*inputLengthPtr).begin()); - TensorPtr outputMax = BufferManager::cpu(shape, nvinfer1::DataType::kINT32); - TensorPtr endIdPtr = BufferManager::cpu(shapeSingle, nvinfer1::DataType::kINT32); + TensorPtr outputMax = BufferManager::cpu(shape, tensorrt_llm::DataType::kINT32); + TensorPtr endIdPtr = BufferManager::cpu(shapeSingle, tensorrt_llm::DataType::kINT32); auto& endId(*BufferRange(*endIdPtr).begin()); endId = ascii->getEndToken(); - TensorPtr acceptedMax = BufferManager::cpu(shape, nvinfer1::DataType::kINT32); - TensorPtr acceptedOffsetsMax = BufferManager::cpu(shape, nvinfer1::DataType::kINT32); - TensorPtr acceptedLengthPtr = BufferManager::cpu(shapeSingle, nvinfer1::DataType::kINT32); + TensorPtr acceptedMax = BufferManager::cpu(shape, tensorrt_llm::DataType::kINT32); + TensorPtr acceptedOffsetsMax = BufferManager::cpu(shape, tensorrt_llm::DataType::kINT32); + TensorPtr acceptedLengthPtr = BufferManager::cpu(shapeSingle, tensorrt_llm::DataType::kINT32); auto& acceptedLength(*BufferRange(*acceptedLengthPtr).begin()); - TensorPtr sequence = BufferManager::cpu(ITensor::makeShape({maxSeqLen + maxDraftLen}), nvinfer1::DataType::kINT32); + TensorPtr sequence + = BufferManager::cpu(ITensor::makeShape({maxSeqLen + maxDraftLen}), tensorrt_llm::DataType::kINT32); BufferRange sequenceRange(*sequence); - TensorPtr sequenceLengthPtr = BufferManager::cpu(shapeSingle, nvinfer1::DataType::kINT32); + TensorPtr sequenceLengthPtr = BufferManager::cpu(shapeSingle, tensorrt_llm::DataType::kINT32); auto& sequenceLength(*bufferCast(*sequenceLengthPtr)); std::copy(promptRange.begin(), promptRange.end(), sequenceRange.begin()); @@ -224,13 +226,13 @@ TEST(LookaheadAlgorithmTest, treeEncodeTest) auto shape = inputTokens->getShape(); auto shape2d = ITensor::makeShape({shape.d[0], shape.d[0]}); - TensorPtr inputMasks = BufferManager::cpu(shape2d, nvinfer1::DataType::kBOOL); + TensorPtr inputMasks = BufferManager::cpu(shape2d, tensorrt_llm::DataType::kBOOL); LookaheadAlgorithm::posIdsToMask(inputMasks, inputPosIds); - TensorPtr outputTokens = BufferManager::cpu(shape, nvinfer1::DataType::kINT32); - TensorPtr outputPosIds = BufferManager::cpu(shape, nvinfer1::DataType::kINT32); - TensorPtr encodeMap = BufferManager::cpu(shape, nvinfer1::DataType::kINT32); - TensorPtr outputMasks = BufferManager::cpu(shape2d, nvinfer1::DataType::kBOOL); + TensorPtr outputTokens = BufferManager::cpu(shape, tensorrt_llm::DataType::kINT32); + TensorPtr outputPosIds = BufferManager::cpu(shape, tensorrt_llm::DataType::kINT32); + TensorPtr encodeMap = BufferManager::cpu(shape, tensorrt_llm::DataType::kINT32); + TensorPtr outputMasks = BufferManager::cpu(shape2d, tensorrt_llm::DataType::kBOOL); // auto len = LookaheadAlgorithm::treeEncode(outputTokens, outputPosIds, outputMasks, inputTokens, inputPosIds, // inputMasks, '$', 9); diff --git a/cpp/tests/unit_tests/layers/lookaheadDecodingLayerTest.cpp b/cpp/tests/unit_tests/layers/lookaheadDecodingLayerTest.cpp index 414e6f101743..917f6dbdca55 100644 --- a/cpp/tests/unit_tests/layers/lookaheadDecodingLayerTest.cpp +++ b/cpp/tests/unit_tests/layers/lookaheadDecodingLayerTest.cpp @@ -23,6 +23,7 @@ #include "tensorrt_llm/common/assert.h" #include "tensorrt_llm/common/logger.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/executor/executor.h" #include "tensorrt_llm/layers/decodingParams.h" #include "tensorrt_llm/layers/lookaheadDecodingLayer.h" @@ -320,7 +321,7 @@ void LookaheadDecodingLayerTest::allocateBuffers() mLlm[gbi] = std::make_shared(mAscii, mOracle[gbi], gbi); mScoreBoard[gbi] = std::ostringstream(); - mHistogram[gbi] = BufferManager::cpu(ITensor::makeShape({mTestParam.n + 1}), nvinfer1::DataType::kINT32); + mHistogram[gbi] = BufferManager::cpu(ITensor::makeShape({mTestParam.n + 1}), tensorrt_llm::DataType::kINT32); } switch (mTestParam.batchType) { @@ -348,48 +349,50 @@ void LookaheadDecodingLayerTest::allocateBuffers() auto maxBatchShape1D = ITensor::makeShape({maxBatchSize}); - mAlgoConfigBatch = BufferManager::pinnedPool(ITensor::makeShape({maxBatchSize, 3}), nvinfer1::DataType::kINT32); + mAlgoConfigBatch = BufferManager::pinnedPool(ITensor::makeShape({maxBatchSize, 3}), tensorrt_llm::DataType::kINT32); - mEndIds = BufferManager::pinnedPool(maxBatchShape1D, nvinfer1::DataType::kINT32); - mTokensPerStep = BufferManager::pinnedPool(maxBatchShape1D, nvinfer1::DataType::kINT32); + mEndIds = BufferManager::pinnedPool(maxBatchShape1D, tensorrt_llm::DataType::kINT32); + mTokensPerStep = BufferManager::pinnedPool(maxBatchShape1D, tensorrt_llm::DataType::kINT32); - mOutputIds = BufferManager::pinnedPool( - ITensor::makeShape({maxBatchSize, maxBeamSize, mMaxSeqLen + mMaxTokensPerStep}), nvinfer1::DataType::kINT32); - mSequenceLengths = BufferManager::pinnedPool(maxBatchShape1D, nvinfer1::DataType::kINT32); + mOutputIds + = BufferManager::pinnedPool(ITensor::makeShape({maxBatchSize, maxBeamSize, mMaxSeqLen + mMaxTokensPerStep}), + tensorrt_llm::DataType::kINT32); + mSequenceLengths = BufferManager::pinnedPool(maxBatchShape1D, tensorrt_llm::DataType::kINT32); mProbs = BufferManager::pinnedPool( - ITensor::makeShape({maxBatchSize, mMaxTokensPerStep, vocabSize}), nvinfer1::DataType::kFLOAT); + ITensor::makeShape({maxBatchSize, mMaxTokensPerStep, vocabSize}), tensorrt_llm::DataType::kFLOAT); mGoldenSampledTokens - = BufferManager::cpu(ITensor::makeShape({maxBatchSize, mMaxTokensPerStep}), nvinfer1::DataType::kINT32); - mInputTokensBatch - = BufferManager::pinnedPool(ITensor::makeShape({maxBatchSize, mMaxTokensPerStep}), nvinfer1::DataType::kINT32); - mPositionIdsBatch - = BufferManager::pinnedPool(ITensor::makeShape({maxBatchSize, mMaxTokensPerStep}), nvinfer1::DataType::kINT32); + = BufferManager::cpu(ITensor::makeShape({maxBatchSize, mMaxTokensPerStep}), tensorrt_llm::DataType::kINT32); + mInputTokensBatch = BufferManager::pinnedPool( + ITensor::makeShape({maxBatchSize, mMaxTokensPerStep}), tensorrt_llm::DataType::kINT32); + mPositionIdsBatch = BufferManager::pinnedPool( + ITensor::makeShape({maxBatchSize, mMaxTokensPerStep}), tensorrt_llm::DataType::kINT32); mNewTokens = BufferManager::pinnedPool( - ITensor::makeShape({mMaxTokensPerStep, maxBatchSize, 1}), nvinfer1::DataType::kINT32); - mNumNewTokens = BufferManager::pinnedPool(maxBatchShape1D, nvinfer1::DataType::kINT32); - mDraftLengths = BufferManager::pinnedPool(maxBatchShape1D, nvinfer1::DataType::kINT32); - mPrevDraftLengths = BufferManager::pinnedPool(maxBatchShape1D, nvinfer1::DataType::kINT32); + ITensor::makeShape({mMaxTokensPerStep, maxBatchSize, 1}), tensorrt_llm::DataType::kINT32); + mNumNewTokens = BufferManager::pinnedPool(maxBatchShape1D, tensorrt_llm::DataType::kINT32); + mDraftLengths = BufferManager::pinnedPool(maxBatchShape1D, tensorrt_llm::DataType::kINT32); + mPrevDraftLengths = BufferManager::pinnedPool(maxBatchShape1D, tensorrt_llm::DataType::kINT32); mDraftTokens - = BufferManager::pinnedPool(ITensor::makeShape({maxBatchSize, maxDraftLen}), nvinfer1::DataType::kINT32); + = BufferManager::pinnedPool(ITensor::makeShape({maxBatchSize, maxDraftLen}), tensorrt_llm::DataType::kINT32); auto packedMaskShape = ITensor::makeShape( {maxBatchSize, mMaxTokensPerStep, static_cast(common::divUp(mMaxTokensPerStep, 32))}); - mPackedMasks = BufferManager::pinnedPool(packedMaskShape, nvinfer1::DataType::kINT32); + mPackedMasks = BufferManager::pinnedPool(packedMaskShape, tensorrt_llm::DataType::kINT32); mPackedMasksBool = BufferManager::pinnedPool( - ITensor::makeShape({maxBatchSize, mMaxTokensPerStep, mMaxTokensPerStep}), nvinfer1::DataType::kBOOL); - mNumNewTokensCumSum = BufferManager::pinnedPool(ITensor::makeShape({maxBatchSize + 1}), nvinfer1::DataType::kINT32); + ITensor::makeShape({maxBatchSize, mMaxTokensPerStep, mMaxTokensPerStep}), tensorrt_llm::DataType::kBOOL); + mNumNewTokensCumSum + = BufferManager::pinnedPool(ITensor::makeShape({maxBatchSize + 1}), tensorrt_llm::DataType::kINT32); mPathsOffsets = BufferManager::pinnedPool( - ITensor::makeShape({maxBatchSize, maxAcceptedDraftLen}), nvinfer1::DataType::kINT32); - mGenerationLengths = BufferManager::pinnedPool(maxBatchShape1D, nvinfer1::DataType::kINT32); - mPositionOffsets - = BufferManager::pinnedPool(ITensor::makeShape({maxBatchSize, mMaxTokensPerStep}), nvinfer1::DataType::kINT32); - mPositionIds - = BufferManager::pinnedPool(ITensor::makeShape({maxBatchSize, mMaxTokensPerStep}), nvinfer1::DataType::kINT32); - mAttentionPackedMask = BufferManager::pinnedPool(packedMaskShape, nvinfer1::DataType::kINT32); - - mBatchSlotsMax = BufferManager::pinnedPool(maxBatchShape1D, nvinfer1::DataType::kINT32); + ITensor::makeShape({maxBatchSize, maxAcceptedDraftLen}), tensorrt_llm::DataType::kINT32); + mGenerationLengths = BufferManager::pinnedPool(maxBatchShape1D, tensorrt_llm::DataType::kINT32); + mPositionOffsets = BufferManager::pinnedPool( + ITensor::makeShape({maxBatchSize, mMaxTokensPerStep}), tensorrt_llm::DataType::kINT32); + mPositionIds = BufferManager::pinnedPool( + ITensor::makeShape({maxBatchSize, mMaxTokensPerStep}), tensorrt_llm::DataType::kINT32); + mAttentionPackedMask = BufferManager::pinnedPool(packedMaskShape, tensorrt_llm::DataType::kINT32); + + mBatchSlotsMax = BufferManager::pinnedPool(maxBatchShape1D, tensorrt_llm::DataType::kINT32); auto const batchSize = 0; auto batchShape1D = ITensor::makeShape({batchSize}); diff --git a/cpp/tests/unit_tests/layers/lookaheadRandomLlmTest.cpp b/cpp/tests/unit_tests/layers/lookaheadRandomLlmTest.cpp index 2cc2523d08cb..4ca7206c3bc8 100644 --- a/cpp/tests/unit_tests/layers/lookaheadRandomLlmTest.cpp +++ b/cpp/tests/unit_tests/layers/lookaheadRandomLlmTest.cpp @@ -15,6 +15,7 @@ */ #include +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/kernels/samplingTopKKernels.h" #include "tensorrt_llm/layers/lookaheadAlgorithm.h" #include "tensorrt_llm/layers/lookaheadDecodingUtils.h" @@ -48,7 +49,7 @@ TEST(LookaheadRandomllm, forward) std::string str("hello world!"); TensorPtr logits = BufferManager::cpu(ITensor::makeShape({static_cast(str.size()), ascii->getVocabSize()}), - nvinfer1::DataType::kFLOAT); + tensorrt_llm::DataType::kFLOAT); ascii->stringToLogits(logits, str); auto result = ascii->logitsToString(logits); EXPECT_EQ(result, str); @@ -67,7 +68,7 @@ TEST(LookaheadRandomllm, forward) std::vector positionIdVec({22, 23, 24, 23, 24, 25, 24, 25, 26, 25, 26, 27, 26, 27, 28}); TensorPtr positionIds = ITensor::wrap(positionIdVec, ITensor::makeShape({len})); TensorPtr outputLogits - = BufferManager::cpu(ITensor::makeShape({len, ascii->getVocabSize()}), nvinfer1::DataType::kFLOAT); + = BufferManager::cpu(ITensor::makeShape({len, ascii->getVocabSize()}), tensorrt_llm::DataType::kFLOAT); llm.forward(outputLogits, inputTokens, positionIds); @@ -123,29 +124,29 @@ TEST(LookaheadRandomllm, gpuSampling) SizeType32 workspaceSize = tensorrt_llm::kernels::getTopKWorkspaceSize(maxBatchSize, maxTokensPerStep, mMaxTopK, vocabSizePadded); - TensorPtr workspaceDevice - = mBufferManager->pinned(ITensor::makeShape({static_cast(workspaceSize)}), nvinfer1::DataType::kINT8); + TensorPtr workspaceDevice = mBufferManager->pinned( + ITensor::makeShape({static_cast(workspaceSize)}), tensorrt_llm::DataType::kINT8); auto const dataType = TRTDataType::value; auto const ptrType = TRTDataType::value; // Allocate GPU data - TensorPtr mSeqLengths = BufferManager::pinned(maxBatchShape1D, nvinfer1::DataType::kINT32); + TensorPtr mSeqLengths = BufferManager::pinned(maxBatchShape1D, tensorrt_llm::DataType::kINT32); TensorPtr mFinished = BufferManager::pinned(maxBatchShape1D, TRTDataType::value); - TensorPtr mEndIds = BufferManager::pinned(maxBatchShape1D, nvinfer1::DataType::kINT32); - TensorPtr mTopPs = BufferManager::pinned(maxBatchShape1D, nvinfer1::DataType::kFLOAT); - TensorPtr mTopKs = BufferManager::pinned(maxBatchShape1D, nvinfer1::DataType::kINT32); - TensorPtr mSkipDecode = BufferManager::pinned(maxBatchShape1D, nvinfer1::DataType::kBOOL); - TensorPtr mTokensPerStep = BufferManager::pinned(maxBatchShape1D, nvinfer1::DataType::kINT32); - - TensorPtr mCurandStates - = BufferManager::pinned(ITensor::makeShape({maxBatchSize, sizeof(curandState_t)}), nvinfer1::DataType::kINT8); + TensorPtr mEndIds = BufferManager::pinned(maxBatchShape1D, tensorrt_llm::DataType::kINT32); + TensorPtr mTopPs = BufferManager::pinned(maxBatchShape1D, tensorrt_llm::DataType::kFLOAT); + TensorPtr mTopKs = BufferManager::pinned(maxBatchShape1D, tensorrt_llm::DataType::kINT32); + TensorPtr mSkipDecode = BufferManager::pinned(maxBatchShape1D, tensorrt_llm::DataType::kBOOL); + TensorPtr mTokensPerStep = BufferManager::pinned(maxBatchShape1D, tensorrt_llm::DataType::kINT32); + + TensorPtr mCurandStates = BufferManager::pinned( + ITensor::makeShape({maxBatchSize, sizeof(curandState_t)}), tensorrt_llm::DataType::kINT8); TensorPtr mOutputIds - = BufferManager::pinned(ITensor::makeShape({maxBatchSize, mMaxSeqLen}), nvinfer1::DataType::kINT32); + = BufferManager::pinned(ITensor::makeShape({maxBatchSize, mMaxSeqLen}), tensorrt_llm::DataType::kINT32); TensorPtr mProbs = BufferManager::pinned(maxBatchShape3D, dataType); - TensorPtr mBatchSlots = BufferManager::pinned(batchShape1D, nvinfer1::DataType::kINT32); + TensorPtr mBatchSlots = BufferManager::pinned(batchShape1D, tensorrt_llm::DataType::kINT32); ///////////////////////////////////// std::copy(batchSlotsVec.begin(), batchSlotsVec.end(), BufferRange(*mBatchSlots).begin()); diff --git a/cpp/tests/unit_tests/layers/medusaDecodeLayerTest.cpp b/cpp/tests/unit_tests/layers/medusaDecodeLayerTest.cpp index 37a479eb4214..a93955d02c9a 100644 --- a/cpp/tests/unit_tests/layers/medusaDecodeLayerTest.cpp +++ b/cpp/tests/unit_tests/layers/medusaDecodeLayerTest.cpp @@ -15,6 +15,7 @@ */ #include "medusaDecodeLayerTest.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/kernels/decodingCommon.h" #include "tensorrt_llm/runtime/medusaModule.h" #include "tensorrt_llm/runtime/runtimeKernels.h" @@ -138,35 +139,36 @@ void MedusaDecodingLayerTest::allocateBuffers() mFinishedDevice = mBufferManager->gpu( ITensor::makeShape({mMaxBatchSize}), TRTDataType::value); - mOutputIdsDevice = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize, mMaxSeqLen}), nvinfer1::DataType::kINT32); + mOutputIdsDevice + = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize, mMaxSeqLen}), tensorrt_llm::DataType::kINT32); - mBatchSlots = BufferManager::pinned(ITensor::makeShape({mBatchSize}), nvinfer1::DataType::kINT32); + mBatchSlots = BufferManager::pinned(ITensor::makeShape({mBatchSize}), tensorrt_llm::DataType::kINT32); - mEndIdsDevice = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize}), nvinfer1::DataType::kINT32); + mEndIdsDevice = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize}), tensorrt_llm::DataType::kINT32); mPathsDevice = mBufferManager->gpu( - ITensor::makeShape({mMaxBatchSize, mMaxDecodingTokens, mMaxDraftPathLen + 1}), nvinfer1::DataType::kINT32); + ITensor::makeShape({mMaxBatchSize, mMaxDecodingTokens, mMaxDraftPathLen + 1}), tensorrt_llm::DataType::kINT32); - mSeqLengthsDevice = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize}), nvinfer1::DataType::kINT32); + mSeqLengthsDevice = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize}), tensorrt_llm::DataType::kINT32); - mAcceptedLengths = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize}), nvinfer1::DataType::kINT32); + mAcceptedLengths = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize}), tensorrt_llm::DataType::kINT32); - mTreeIdsDevice - = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize, mMaxDecodingTokens - 1}), nvinfer1::DataType::kINT32); + mTreeIdsDevice = mBufferManager->gpu( + ITensor::makeShape({mMaxBatchSize, mMaxDecodingTokens - 1}), tensorrt_llm::DataType::kINT32); mMedusaLogitsDevice = mBufferManager->gpu( ITensor::makeShape({mMaxDraftPathLen, mMaxBatchSize, mMaxDecodingTokens, mVocabSizePadded}), dataType); - mNextDraftTokensDevice - = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize, mMaxDecodingTokens - 1}), nvinfer1::DataType::kINT32); + mNextDraftTokensDevice = mBufferManager->gpu( + ITensor::makeShape({mMaxBatchSize, mMaxDecodingTokens - 1}), tensorrt_llm::DataType::kINT32); - mTokensPerStepDevice = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize}), nvinfer1::DataType::kINT32); + mTokensPerStepDevice = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize}), tensorrt_llm::DataType::kINT32); mAcceptedLengthCumSumDevice - = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize + 1}), nvinfer1::DataType::kINT32); + = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize + 1}), tensorrt_llm::DataType::kINT32); mPackedPathsDevice - = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize * mMaxDraftPathLen}), nvinfer1::DataType::kINT32); + = mBufferManager->gpu(ITensor::makeShape({mMaxBatchSize * mMaxDraftPathLen}), tensorrt_llm::DataType::kINT32); for (int32_t bi = 0; bi < mBatchSize; ++bi) { @@ -215,7 +217,7 @@ void MedusaDecodingLayerTest::setup(SamplingParams& params) for (SizeType32 bi = 0; bi < mBatchSize; ++bi) { auto const draftIdsHost = ITensor::wrap(reinterpret_cast(params.draftIds[bi].data()), - nvinfer1::DataType::kINT32, ITensor::makeShape({1, mMaxDecodingTokens - 1})); + tensorrt_llm::DataType::kINT32, ITensor::makeShape({1, mMaxDecodingTokens - 1})); auto draftIdsDeviceSlice = ITensor::slice(mNextDraftTokensDevice, batchSlotsPtr[bi], 1); mBufferManager->copy(*draftIdsHost, *draftIdsDeviceSlice); } @@ -224,7 +226,7 @@ void MedusaDecodingLayerTest::setup(SamplingParams& params) { auto& path = params.paths[bi]; auto const numPaths = static_cast(params.paths[bi].size() / (mMaxDraftPathLen + 1)); - auto const pathsHost = ITensor::wrap(reinterpret_cast(path.data()), nvinfer1::DataType::kINT32, + auto const pathsHost = ITensor::wrap(reinterpret_cast(path.data()), tensorrt_llm::DataType::kINT32, ITensor::makeShape({1, numPaths, mMaxDraftPathLen + 1})); TensorPtr pathsDeviceSlice = ITensor::slice(mPathsDevice, batchSlotsPtr[bi], 1); pathsDeviceSlice->squeeze(0); diff --git a/cpp/tests/unit_tests/layers/randomLlm.cpp b/cpp/tests/unit_tests/layers/randomLlm.cpp index 632cf8765c20..63aa85eaad16 100644 --- a/cpp/tests/unit_tests/layers/randomLlm.cpp +++ b/cpp/tests/unit_tests/layers/randomLlm.cpp @@ -16,6 +16,7 @@ #include "tests/unit_tests/layers/randomLlm.h" #include "tensorrt_llm/common/assert.h" #include "tensorrt_llm/common/logger.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/layers/lookaheadDecodingUtils.h" #include "tensorrt_llm/runtime/bufferManager.h" #include "tensorrt_llm/runtime/common.h" @@ -34,7 +35,7 @@ TensorPtr initTensor(std::string str, std::optional shape) { TLLM_CHECK(ITensor::volume(shape1d) == ITensor::volume(shape.value())); } - TensorPtr tensor = BufferManager::cpu(shape.value_or(shape1d), nvinfer1::DataType::kINT32); + TensorPtr tensor = BufferManager::cpu(shape.value_or(shape1d), tensorrt_llm::DataType::kINT32); auto tensorRange = BufferRange(*tensor); std::copy(str.begin(), str.end(), tensorRange.begin()); return tensor; @@ -42,7 +43,7 @@ TensorPtr initTensor(std::string str, std::optional shape) TensorConstPtr RandomTokenLogits::tokenToLogits(TokenIdType token) const { - TensorPtr logits = BufferManager::cpu(mVocabulary->getShape(), nvinfer1::DataType::kFLOAT); + TensorPtr logits = BufferManager::cpu(mVocabulary->getShape(), tensorrt_llm::DataType::kFLOAT); tokenToLogits(logits, token); return logits; } @@ -152,7 +153,7 @@ void RandomTokenLogits::logitsToTensor(TensorPtr const& tokens, TensorConstPtr c TensorConstPtr RandomTokenLogits::logitsToTensor(TensorConstPtr const& logits) const { auto len = logits->getShape().d[0]; - TensorPtr result = BufferManager::cpu(ITensor::makeShape({len}), nvinfer1::DataType::kINT32); + TensorPtr result = BufferManager::cpu(ITensor::makeShape({len}), tensorrt_llm::DataType::kINT32); logitsToTensor(result, logits); return result; } @@ -209,7 +210,7 @@ bool RandomLlm::verify(SizeType32 const offset, TensorConstPtr const& script) co void RandomLlm::forward(TensorPtr const& output, runtime::SizeType32 startId, TensorConstPtr const& input, TensorConstPtr const& offsets, TensorConstPtr const mask) const { - TensorPtr posIds = BufferManager::cpu(input->getShape(), nvinfer1::DataType::kINT32); + TensorPtr posIds = BufferManager::cpu(input->getShape(), tensorrt_llm::DataType::kINT32); BufferRange idRange(*posIds); BufferRange offsetRange(*offsets); for (auto i = 0; i < idRange.size(); i++) @@ -226,7 +227,7 @@ void RandomLlm::forward(TensorPtr const& output, TensorConstPtr const& input, Te TLLM_CHECK(ITensor::volume(input->getShape()) == ITensor::volume(position->getShape())); TLLM_CHECK(ITensor::volume(output->getShape()) == ITensor::volume(input->getShape()) * mTable->getVocabSize()); - TensorPtr tokens = BufferManager::cpu(input->getShape(), nvinfer1::DataType::kINT32); + TensorPtr tokens = BufferManager::cpu(input->getShape(), tensorrt_llm::DataType::kINT32); foretell(tokens, input, position, mask); // foretellOld(tokens, input, position); mTable->tensorToLogits(output, tokens); @@ -247,7 +248,7 @@ void LookaheadRandomLlm::foretell(TensorPtr const& output, TensorConstPtr const& TLLM_CHECK(mask->getShape().d[1] >= len); } - TensorPtr maskRebuilt = BufferManager::cpu(ITensor::makeShape({len, len}), nvinfer1::DataType::kBOOL); + TensorPtr maskRebuilt = BufferManager::cpu(ITensor::makeShape({len, len}), tensorrt_llm::DataType::kBOOL); posIdsToMask(maskRebuilt, position); auto outputRange = BufferRange(*output); diff --git a/cpp/tests/unit_tests/layers/randomLlm.h b/cpp/tests/unit_tests/layers/randomLlm.h index a6e898baedf5..b0f0564944dc 100644 --- a/cpp/tests/unit_tests/layers/randomLlm.h +++ b/cpp/tests/unit_tests/layers/randomLlm.h @@ -18,6 +18,7 @@ #include #include +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/layers/lookaheadDecodingUtils.h" #include "tensorrt_llm/runtime/common.h" #include "tensorrt_llm/runtime/runtimeKernels.h" @@ -77,7 +78,7 @@ class AsciiRandomTokenLogits : public RandomTokenLogits : RandomTokenLogits( []() { - auto vocab = BufferManager::cpu(ITensor::makeShape({128}), nvinfer1::DataType::kINT32); + auto vocab = BufferManager::cpu(ITensor::makeShape({128}), tensorrt_llm::DataType::kINT32); auto vocabRange = BufferRange(*vocab); TokenIdType token{0}; std::for_each(vocabRange.begin(), vocabRange.end(), [&token](auto& v) { v = token++; }); diff --git a/cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp b/cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp index 401cdbfd5d44..e294cdff9d49 100644 --- a/cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp +++ b/cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp @@ -35,6 +35,7 @@ #include "tensorrt_llm/common/assert.h" #include "tensorrt_llm/common/cudaUtils.h" #include "tensorrt_llm/common/envUtils.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/executor/cache_transmission/agent_utils/connection.h" #include "tensorrt_llm/executor/cache_transmission/mpi_utils/connection.h" #include "tensorrt_llm/executor/dataTransceiverState.h" @@ -105,7 +106,7 @@ TEST_F(RequestInfoTest, Basic) } auto state = std::make_unique(); state->setCommState(texec::kv_cache::CommState{12, "127.0.0.1"}); - state->setCacheState(texec::kv_cache::CacheState{10, 12, 128, 128, 8, 8, 8, {10}, nvinfer1::DataType::kFLOAT}); + state->setCacheState(texec::kv_cache::CacheState{10, 12, 128, 128, 8, 8, 8, {10}, tensorrt_llm::DataType::kFLOAT}); RequestInfo info{1, *state}; auto info2 = serializeDeserialize(info); EXPECT_EQ(info, info2); @@ -135,7 +136,7 @@ TEST_F(CacheConfigTest, EqualTo) constexpr SizeType32 nbRnnLayers{2}; constexpr SizeType32 nbHeads{12}; constexpr SizeType32 hiddenSize{768}; - constexpr nvinfer1::DataType dtype{nvinfer1::DataType::kFLOAT}; + constexpr tensorrt_llm::DataType dtype{tensorrt_llm::DataType::kFLOAT}; constexpr SizeType32 tokensPerBlock{64}; constexpr SizeType32 tensorParallelism{8}; constexpr SizeType32 pipelineParallelism{2}; @@ -216,7 +217,7 @@ class SymmetricalCacheTest : public ::testing::Test // NOLINT(cppcoreguidelines- auto constexpr blocksInSecondaryPool = 0; auto constexpr enableBlockReuse = false; - auto constexpr dataType = nvinfer1::DataType::kFLOAT; + auto constexpr dataType = tensorrt_llm::DataType::kFLOAT; using BlocksPerWindow = std::map>; auto const blocksPerWindow = BlocksPerWindow{{maxAttentionWindow, {totalNumBlocks, blocksInSecondaryPool}}}; @@ -428,8 +429,8 @@ TEST_F(SymmetricalCacheTest, SimpleTest) #if ENABLE_MULTI_DEVICE -using AsymmetricTestParam = std::tuple; +using AsymmetricTestParam = std::tuple; // CPMetaData struct to hold CP-specific information struct CPMetaData @@ -579,7 +580,7 @@ class AsymmetricalCacheTest : public ::testing::TestWithParam generateExpectedValue(size_t initial, int windowSize, - int tokenId, int layerId, int headId, int hiddenId, bool key, nvinfer1::DataType dataType) + int tokenId, int layerId, int headId, int hiddenId, bool key, tensorrt_llm::DataType dataType) { size_t seed = 0; std::size_t hashValue = std::hash{}(initial); @@ -1324,7 +1325,7 @@ TEST_P(AsymmetricalCacheTest, TestCase) int numHeads = std::get<7>(param); int sizePerHead = std::get<8>(param); int tokensPerBlock = std::get<9>(param); - nvinfer1::DataType dataType = std::get<10>(param); + tensorrt_llm::DataType dataType = std::get<10>(param); int kvFactor = std::get<11>(param); bool isMLA = std::get<12>(param); @@ -1439,7 +1440,7 @@ TEST_P(AsymmetricalCacheTestWithDP, TestCase) int numHeads = std::get<7>(param); int sizePerHead = std::get<8>(param); int tokensPerBlock = std::get<9>(param); - nvinfer1::DataType dataType = std::get<10>(param); + tensorrt_llm::DataType dataType = std::get<10>(param); int kvFactor = std::get<11>(param); bool isMLA = std::get<12>(param); @@ -1576,7 +1577,7 @@ TEST_P(UnexpectedTerminationRaceTest, UnexpectedTerminationRaceTest) int numHeads = std::get<7>(param); int sizePerHead = std::get<8>(param); int tokensPerBlock = std::get<9>(param); - nvinfer1::DataType dataType = std::get<10>(param); + tensorrt_llm::DataType dataType = std::get<10>(param); int kvFactor = std::get<11>(param); bool isMLA = std::get<12>(param); @@ -1723,87 +1724,87 @@ TEST_P(UnexpectedTerminationRaceTest, UnexpectedTerminationRaceTest) INSTANTIATE_TEST_CASE_P(UnexpectedTerminationRaceTest, UnexpectedTerminationRaceTest, testing::Combine(testing::Values(2), testing::Values(1), testing::Values(1), testing::Values(4), testing::Values(1), testing::Values(1), testing::Values(4), testing::Values(4), testing::Values(4), testing::Values(16), - testing::Values(nvinfer1::DataType::kFLOAT), testing::Values(2), testing::Values(false), testing::Values(false), - testing::Values(false), testing::Values(false), testing::Values(false), testing::Values(0), - testing::Values(128))); + testing::Values(tensorrt_llm::DataType::kFLOAT), testing::Values(2), testing::Values(false), + testing::Values(false), testing::Values(false), testing::Values(false), testing::Values(false), + testing::Values(0), testing::Values(128))); // Waive off isWindow test for now INSTANTIATE_TEST_CASE_P(AsymmetricCaseTest0, AsymmetricalCacheTest, testing::Combine(testing::Values(1, 2), testing::Values(1, 2), testing::Values(1), testing::Values(1, 2), testing::Values(1, 2), testing::Values(1), testing::Values(4), testing::Values(4), testing::Values(4), - testing::Values(16), testing::Values(nvinfer1::DataType::kFLOAT, nvinfer1::DataType::kINT8), testing::Values(2), - testing::Values(false), testing::Values(false), testing::Values(false), testing::Values(/*true,*/ false), - testing::Values(false), testing::Values(0), testing::Values(128))); + testing::Values(16), testing::Values(tensorrt_llm::DataType::kFLOAT, tensorrt_llm::DataType::kINT8), + testing::Values(2), testing::Values(false), testing::Values(false), testing::Values(false), + testing::Values(/*true,*/ false), testing::Values(false), testing::Values(0), testing::Values(128))); INSTANTIATE_TEST_CASE_P(AsymmetricCaseTestWithWindow, AsymmetricalCacheTest, testing::Combine(testing::Values(1), testing::Values(1), testing::Values(1), testing::Values(1), testing::Values(1), testing::Values(1), testing::Values(5), testing::Values(4), testing::Values(4), testing::Values(8), - testing::Values(nvinfer1::DataType::kFLOAT, nvinfer1::DataType::kINT8), testing::Values(2), + testing::Values(tensorrt_llm::DataType::kFLOAT, tensorrt_llm::DataType::kINT8), testing::Values(2), testing::Values(false), testing::Values(false), testing::Values(false), testing::Values(true), testing::Values(false), testing::Values(0), testing::Values(128))); INSTANTIATE_TEST_CASE_P(AsymmetricCaseTest1, AsymmetricalCacheTest, testing::Combine(testing::Values(4), testing::Values(1), testing::Values(1), testing::Values(1), testing::Values(4), testing::Values(1), testing::Values(8), testing::Values(4), testing::Values(4), testing::Values(8), - testing::Values(nvinfer1::DataType::kFLOAT, nvinfer1::DataType::kINT8), testing::Values(2), + testing::Values(tensorrt_llm::DataType::kFLOAT, tensorrt_llm::DataType::kINT8), testing::Values(2), testing::Values(false), testing::Values(false), testing::Values(false), testing::Values(false /*, true*/), testing::Values(false), testing::Values(0), testing::Values(128))); INSTANTIATE_TEST_CASE_P(AsymmetricCaseTest1EvenLayer, AsymmetricalCacheTest, testing::Combine(testing::Values(1), testing::Values(4), testing::Values(1), testing::Values(1), testing::Values(4), testing::Values(1), testing::Values(10), testing::Values(4), testing::Values(4), testing::Values(8), - testing::Values(nvinfer1::DataType::kFLOAT), testing::Values(2), testing::Values(false), testing::Values(false), - testing::Values(false), testing::Values(false), testing::Values(false), testing::Values(0), - testing::Values(128))); + testing::Values(tensorrt_llm::DataType::kFLOAT), testing::Values(2), testing::Values(false), + testing::Values(false), testing::Values(false), testing::Values(false), testing::Values(false), + testing::Values(0), testing::Values(128))); INSTANTIATE_TEST_CASE_P(AsymmetricCaseTest2EvenLayer, AsymmetricalCacheTest, testing::Combine(testing::Values(4), testing::Values(1), testing::Values(1), testing::Values(1), testing::Values(4), testing::Values(1), testing::Values(10), testing::Values(4), testing::Values(4), testing::Values(8), - testing::Values(nvinfer1::DataType::kFLOAT), testing::Values(2), testing::Values(false), testing::Values(false), - testing::Values(false), testing::Values(false), testing::Values(false), testing::Values(0), - testing::Values(128))); + testing::Values(tensorrt_llm::DataType::kFLOAT), testing::Values(2), testing::Values(false), + testing::Values(false), testing::Values(false), testing::Values(false), testing::Values(false), + testing::Values(0), testing::Values(128))); INSTANTIATE_TEST_CASE_P(AsymmetricCaseTest2, AsymmetricalCacheTest, testing::Combine(testing::Values(1), testing::Values(2), testing::Values(1), testing::Values(1), testing::Values(1, 4), testing::Values(1), testing::Values(16), testing::Values(16), testing::Values(4), - testing::Values(8), testing::Values(nvinfer1::DataType::kFLOAT), testing::Values(2), testing::Values(false), + testing::Values(8), testing::Values(tensorrt_llm::DataType::kFLOAT), testing::Values(2), testing::Values(false), testing::Values(false), testing::Values(false), testing::Values(false), testing::Values(false), testing::Values(0), testing::Values(128))); INSTANTIATE_TEST_CASE_P(AsymmetricCaseTest0ForMLA, AsymmetricalCacheTest, testing::Combine(testing::Values(1, 2), testing::Values(1, 2), testing::Values(1), testing::Values(1, 2), testing::Values(1, 2), testing::Values(1), testing::Values(4), testing::Values(1), testing::Values(4), - testing::Values(16), testing::Values(nvinfer1::DataType::kFLOAT, nvinfer1::DataType::kINT8), testing::Values(1), - testing::Values(true), testing::Values(false), testing::Values(false), testing::Values(false), - testing::Values(false), testing::Values(0), testing::Values(128))); + testing::Values(16), testing::Values(tensorrt_llm::DataType::kFLOAT, tensorrt_llm::DataType::kINT8), + testing::Values(1), testing::Values(true), testing::Values(false), testing::Values(false), + testing::Values(false), testing::Values(false), testing::Values(0), testing::Values(128))); INSTANTIATE_TEST_CASE_P(AsymmetricCaseTest1ForMLA, AsymmetricalCacheTest, testing::Combine(testing::Values(4), testing::Values(1), testing::Values(1), testing::Values(1), testing::Values(4), testing::Values(1), testing::Values(4), testing::Values(1), testing::Values(4), testing::Values(8), - testing::Values(nvinfer1::DataType::kFLOAT, nvinfer1::DataType::kINT8), testing::Values(1), + testing::Values(tensorrt_llm::DataType::kFLOAT, tensorrt_llm::DataType::kINT8), testing::Values(1), testing::Values(true), testing::Values(false), testing::Values(false), testing::Values(false), testing::Values(false), testing::Values(0), testing::Values(128))); INSTANTIATE_TEST_CASE_P(AsymmetricCaseTest1ForMLAEvenLayer, AsymmetricalCacheTestWithDP, testing::Combine(testing::Values(1), testing::Values(4), testing::Values(1), testing::Values(4), testing::Values(1), testing::Values(1), testing::Values(10), testing::Values(1), testing::Values(4), testing::Values(8), - testing::Values(nvinfer1::DataType::kFLOAT, nvinfer1::DataType::kINT8), testing::Values(1), + testing::Values(tensorrt_llm::DataType::kFLOAT, tensorrt_llm::DataType::kINT8), testing::Values(1), testing::Values(true), testing::Values(false), testing::Values(false, true), testing::Values(false), testing::Values(false), testing::Values(0), testing::Values(128))); INSTANTIATE_TEST_CASE_P(AsymmetricCaseTest2ForMLAEvenLayer, AsymmetricalCacheTestWithDP, testing::Combine(testing::Values(4), testing::Values(1), testing::Values(1), testing::Values(1), testing::Values(4), testing::Values(1), testing::Values(10), testing::Values(1), testing::Values(4), testing::Values(8), - testing::Values(nvinfer1::DataType::kFLOAT, nvinfer1::DataType::kINT8), testing::Values(1), + testing::Values(tensorrt_llm::DataType::kFLOAT, tensorrt_llm::DataType::kINT8), testing::Values(1), testing::Values(true), testing::Values(false), testing::Values(false, true), testing::Values(false), testing::Values(false), testing::Values(0), testing::Values(128))); INSTANTIATE_TEST_CASE_P(AsymmetricCaseTest0ForMLAWithIndexerKCache, AsymmetricalCacheTest, testing::Combine(testing::Values(1, 2), testing::Values(1, 2), testing::Values(1), testing::Values(1, 2), testing::Values(1, 2), testing::Values(1), testing::Values(4), testing::Values(1), testing::Values(4), - testing::Values(16), testing::Values(nvinfer1::DataType::kFLOAT, nvinfer1::DataType::kINT8), testing::Values(1), - testing::Values(true), testing::Values(false), testing::Values(false), testing::Values(false), - testing::Values(true), testing::Values(256), testing::Values(128))); + testing::Values(16), testing::Values(tensorrt_llm::DataType::kFLOAT, tensorrt_llm::DataType::kINT8), + testing::Values(1), testing::Values(true), testing::Values(false), testing::Values(false), + testing::Values(false), testing::Values(true), testing::Values(256), testing::Values(128))); // Tests cases where there's non-trivial TP and PP on context side but only CP on gen side. INSTANTIATE_TEST_CASE_P(AsymmetricCaseTest0WithCPForMLA, AsymmetricalCacheTest, @@ -1817,7 +1818,7 @@ INSTANTIATE_TEST_CASE_P(AsymmetricCaseTest0WithCPForMLA, AsymmetricalCacheTest, /*numHeads*/ testing::Values(1), /*sizePerHead*/ testing::Values(4), /*tokensPerBlock*/ testing::Values(8), - /*dataType*/ testing::Values(nvinfer1::DataType::kFLOAT, nvinfer1::DataType::kINT8), + /*dataType*/ testing::Values(tensorrt_llm::DataType::kFLOAT, tensorrt_llm::DataType::kINT8), /*kvFactor*/ testing::Values(1), /*isMLA*/ testing::Values(true), /*contextDP*/ testing::Values(false), @@ -1839,7 +1840,7 @@ INSTANTIATE_TEST_CASE_P(AsymmetricCaseTest1WithCPForMLA, AsymmetricalCacheTest, /*numHeads*/ testing::Values(1), /*sizePerHead*/ testing::Values(4), /*tokensPerBlock*/ testing::Values(8), - /*dataType*/ testing::Values(nvinfer1::DataType::kFLOAT, nvinfer1::DataType::kINT8), + /*dataType*/ testing::Values(tensorrt_llm::DataType::kFLOAT, tensorrt_llm::DataType::kINT8), /*kvFactor*/ testing::Values(1), /*isMLA*/ testing::Values(true), /*contextDP*/ testing::Values(false), @@ -1861,7 +1862,7 @@ INSTANTIATE_TEST_CASE_P(AsymmetricCaseTest0WithCPForGQA, AsymmetricalCacheTest, /*numHeads*/ testing::Values(4), /*sizePerHead*/ testing::Values(4), /*tokensPerBlock*/ testing::Values(8), - /*dataType*/ testing::Values(nvinfer1::DataType::kFLOAT, nvinfer1::DataType::kINT8), + /*dataType*/ testing::Values(tensorrt_llm::DataType::kFLOAT, tensorrt_llm::DataType::kINT8), /*kvFactor*/ testing::Values(2), /*isMLA*/ testing::Values(false), /*contextDP*/ testing::Values(false), @@ -1880,7 +1881,7 @@ INSTANTIATE_TEST_CASE_P(AsymmetricCaseTest1WithCPForGQA, AsymmetricalCacheTest, /*numHeads*/ testing::Values(4), /*sizePerHead*/ testing::Values(4), /*tokensPerBlock*/ testing::Values(8), - /*dataType*/ testing::Values(nvinfer1::DataType::kFLOAT, nvinfer1::DataType::kINT8), + /*dataType*/ testing::Values(tensorrt_llm::DataType::kFLOAT, tensorrt_llm::DataType::kINT8), /*kvFactor*/ testing::Values(2), /*isMLA*/ testing::Values(false), /*contextDP*/ testing::Values(false), @@ -1899,7 +1900,7 @@ INSTANTIATE_TEST_CASE_P(AsymmetricCaseTest0WithCPForMLAUnevenLayer, Asymmetrical /*numHeads*/ testing::Values(1), /*sizePerHead*/ testing::Values(4), /*tokensPerBlock*/ testing::Values(8), - /*dataType*/ testing::Values(nvinfer1::DataType::kFLOAT, nvinfer1::DataType::kINT8), + /*dataType*/ testing::Values(tensorrt_llm::DataType::kFLOAT, tensorrt_llm::DataType::kINT8), /*kvFactor*/ testing::Values(1), /*isMLA*/ testing::Values(true), /*contextDP*/ testing::Values(false), @@ -1921,7 +1922,7 @@ INSTANTIATE_TEST_CASE_P(AsymmetricCaseTest1WithCPForMLAUnevenLayer, Asymmetrical /*numHeads*/ testing::Values(1), /*sizePerHead*/ testing::Values(4), /*tokensPerBlock*/ testing::Values(8), - /*dataType*/ testing::Values(nvinfer1::DataType::kFLOAT, nvinfer1::DataType::kINT8), + /*dataType*/ testing::Values(tensorrt_llm::DataType::kFLOAT, tensorrt_llm::DataType::kINT8), /*kvFactor*/ testing::Values(1), /*isMLA*/ testing::Values(true), /*contextDP*/ testing::Values(false), @@ -1943,7 +1944,7 @@ INSTANTIATE_TEST_CASE_P(AsymmetricCaseTest2WithCPForMLAUnevenLayer, Asymmetrical /*numHeads*/ testing::Values(1), /*sizePerHead*/ testing::Values(4), /*tokensPerBlock*/ testing::Values(8), - /*dataType*/ testing::Values(nvinfer1::DataType::kFLOAT, nvinfer1::DataType::kINT8), + /*dataType*/ testing::Values(tensorrt_llm::DataType::kFLOAT, tensorrt_llm::DataType::kINT8), /*kvFactor*/ testing::Values(1), /*isMLA*/ testing::Values(true), /*contextDP*/ testing::Values(false), @@ -1965,7 +1966,7 @@ INSTANTIATE_TEST_CASE_P(AsymmetricCaseTestWithCPAndDPForMLA0, AsymmetricalCacheT /*numHeads*/ testing::Values(1), /*sizePerHead*/ testing::Values(4), /*tokensPerBlock*/ testing::Values(8), - /*dataType*/ testing::Values(nvinfer1::DataType::kFLOAT, nvinfer1::DataType::kINT8), + /*dataType*/ testing::Values(tensorrt_llm::DataType::kFLOAT, tensorrt_llm::DataType::kINT8), /*kvFactor*/ testing::Values(1), /*isMLA*/ testing::Values(true), /*contextDP*/ testing::Values(false), @@ -1987,7 +1988,7 @@ INSTANTIATE_TEST_CASE_P(AsymmetricCaseTestWithCPAndDPForMLA1, AsymmetricalCacheT /*numHeads*/ testing::Values(1), /*sizePerHead*/ testing::Values(4), /*tokensPerBlock*/ testing::Values(8), - /*dataType*/ testing::Values(nvinfer1::DataType::kFLOAT, nvinfer1::DataType::kINT8), + /*dataType*/ testing::Values(tensorrt_llm::DataType::kFLOAT, tensorrt_llm::DataType::kINT8), /*kvFactor*/ testing::Values(1), /*isMLA*/ testing::Values(true), /*contextDP*/ testing::Values(true), @@ -2009,7 +2010,7 @@ INSTANTIATE_TEST_CASE_P(AsymmetricCaseTestWithCPAndDPForGQA0, AsymmetricalCacheT /*numHeads*/ testing::Values(4), /*sizePerHead*/ testing::Values(4), /*tokensPerBlock*/ testing::Values(8), - /*dataType*/ testing::Values(nvinfer1::DataType::kFLOAT, nvinfer1::DataType::kINT8), + /*dataType*/ testing::Values(tensorrt_llm::DataType::kFLOAT, tensorrt_llm::DataType::kINT8), /*kvFactor*/ testing::Values(2), /*isMLA*/ testing::Values(false), /*contextDP*/ testing::Values(false), @@ -2028,7 +2029,7 @@ INSTANTIATE_TEST_CASE_P(AsymmetricCaseTestWithCPAndDPForGQA1, AsymmetricalCacheT /*numHeads*/ testing::Values(4), /*sizePerHead*/ testing::Values(4), /*tokensPerBlock*/ testing::Values(8), - /*dataType*/ testing::Values(nvinfer1::DataType::kFLOAT, nvinfer1::DataType::kINT8), + /*dataType*/ testing::Values(tensorrt_llm::DataType::kFLOAT, tensorrt_llm::DataType::kINT8), /*kvFactor*/ testing::Values(2), /*isMLA*/ testing::Values(false), /*contextDP*/ testing::Values(true), @@ -2038,97 +2039,97 @@ INSTANTIATE_TEST_CASE_P(AsymmetricCaseTestWithCPAndDPForGQA1, AsymmetricalCacheT INSTANTIATE_TEST_CASE_P(AsymmetricCaseTestWithDPForMLA1, AsymmetricalCacheTestWithDP, testing::Combine(testing::Values(1, 2), testing::Values(1, 2), testing::Values(1), testing::Values(1, 2), testing::Values(1, 2), testing::Values(1), testing::Values(4), testing::Values(1), testing::Values(4), - testing::Values(16), testing::Values(nvinfer1::DataType::kFLOAT, nvinfer1::DataType::kINT8), testing::Values(1), - testing::Values(true), testing::Values(true), testing::Values(true), testing::Values(false), + testing::Values(16), testing::Values(tensorrt_llm::DataType::kFLOAT, tensorrt_llm::DataType::kINT8), + testing::Values(1), testing::Values(true), testing::Values(true), testing::Values(true), testing::Values(false), testing::Values(false), testing::Values(0), testing::Values(128))); INSTANTIATE_TEST_CASE_P(AsymmetricCaseTestWithDPForMLA2, AsymmetricalCacheTestWithDP, testing::Combine(testing::Values(1, 2), testing::Values(1, 2), testing::Values(1), testing::Values(1, 2), testing::Values(1, 2), testing::Values(1), testing::Values(4), testing::Values(1), testing::Values(4), - testing::Values(16), testing::Values(nvinfer1::DataType::kFLOAT, nvinfer1::DataType::kINT8), testing::Values(1), - testing::Values(true), testing::Values(true), testing::Values(false), testing::Values(false), - testing::Values(false), testing::Values(0), testing::Values(128))); + testing::Values(16), testing::Values(tensorrt_llm::DataType::kFLOAT, tensorrt_llm::DataType::kINT8), + testing::Values(1), testing::Values(true), testing::Values(true), testing::Values(false), + testing::Values(false), testing::Values(false), testing::Values(0), testing::Values(128))); INSTANTIATE_TEST_CASE_P(AsymmetricCaseTestWithDPForMLA3, AsymmetricalCacheTestWithDP, testing::Combine(testing::Values(1, 2), testing::Values(1, 2), testing::Values(1), testing::Values(1, 2), testing::Values(1, 2), testing::Values(1), testing::Values(4), testing::Values(1), testing::Values(4), - testing::Values(16), testing::Values(nvinfer1::DataType::kFLOAT, nvinfer1::DataType::kINT8), testing::Values(1), - testing::Values(true), testing::Values(false), testing::Values(true), testing::Values(false), - testing::Values(false), testing::Values(0), testing::Values(128))); + testing::Values(16), testing::Values(tensorrt_llm::DataType::kFLOAT, tensorrt_llm::DataType::kINT8), + testing::Values(1), testing::Values(true), testing::Values(false), testing::Values(true), + testing::Values(false), testing::Values(false), testing::Values(0), testing::Values(128))); INSTANTIATE_TEST_CASE_P(AsymmetricCaseTestWithDPForMLA4, AsymmetricalCacheTestWithDP, testing::Combine(testing::Values(2), testing::Values(1), testing::Values(1), testing::Values(4), testing::Values(1), testing::Values(1), testing::Values(4), testing::Values(1), testing::Values(4), testing::Values(16), - testing::Values(nvinfer1::DataType::kFLOAT, nvinfer1::DataType::kINT8), testing::Values(1), + testing::Values(tensorrt_llm::DataType::kFLOAT, tensorrt_llm::DataType::kINT8), testing::Values(1), testing::Values(true), testing::Values(false), testing::Values(true), testing::Values(false), testing::Values(false), testing::Values(0), testing::Values(128))); INSTANTIATE_TEST_CASE_P(AsymmetricCaseTestWithDPForMLA5, AsymmetricalCacheTestWithDP, testing::Combine(testing::Values(4), testing::Values(1), testing::Values(1), testing::Values(2), testing::Values(1), testing::Values(1), testing::Values(4), testing::Values(1), testing::Values(4), testing::Values(16), - testing::Values(nvinfer1::DataType::kFLOAT, nvinfer1::DataType::kINT8), testing::Values(1), + testing::Values(tensorrt_llm::DataType::kFLOAT, tensorrt_llm::DataType::kINT8), testing::Values(1), testing::Values(true), testing::Values(false), testing::Values(true), testing::Values(false), testing::Values(false), testing::Values(0), testing::Values(128))); INSTANTIATE_TEST_CASE_P(AsymmetricCaseTestWithDPForNoMLA, AsymmetricalCacheTestWithDP, testing::Combine(testing::Values(1, 2), testing::Values(1, 2), testing::Values(1), testing::Values(1, 2), testing::Values(1, 2), testing::Values(1), testing::Values(4), testing::Values(4), testing::Values(4), - testing::Values(16), testing::Values(nvinfer1::DataType::kFLOAT, nvinfer1::DataType::kINT8), testing::Values(2), - testing::Values(false), testing::Values(true), testing::Values(true), testing::Values(false), - testing::Values(false), testing::Values(0), testing::Values(128))); + testing::Values(16), testing::Values(tensorrt_llm::DataType::kFLOAT, tensorrt_llm::DataType::kINT8), + testing::Values(2), testing::Values(false), testing::Values(true), testing::Values(true), + testing::Values(false), testing::Values(false), testing::Values(0), testing::Values(128))); INSTANTIATE_TEST_CASE_P(AsymmetricCaseTestWithDPForNoMLA1, AsymmetricalCacheTestWithDP, testing::Combine(testing::Values(1, 2), testing::Values(1, 2), testing::Values(1), testing::Values(1, 2), testing::Values(1, 2), testing::Values(1), testing::Values(4), testing::Values(4), testing::Values(4), - testing::Values(16), testing::Values(nvinfer1::DataType::kFLOAT, nvinfer1::DataType::kINT8), testing::Values(2), - testing::Values(false), testing::Values(true), testing::Values(false), testing::Values(false), - testing::Values(false), testing::Values(0), testing::Values(128))); + testing::Values(16), testing::Values(tensorrt_llm::DataType::kFLOAT, tensorrt_llm::DataType::kINT8), + testing::Values(2), testing::Values(false), testing::Values(true), testing::Values(false), + testing::Values(false), testing::Values(false), testing::Values(0), testing::Values(128))); INSTANTIATE_TEST_CASE_P(AsymmetricCaseTestWithDPForNoMLA2, AsymmetricalCacheTestWithDP, testing::Combine(testing::Values(1, 2), testing::Values(1, 2), testing::Values(1), testing::Values(1, 2), testing::Values(1, 2), testing::Values(1), testing::Values(4), testing::Values(4), testing::Values(4), - testing::Values(16), testing::Values(nvinfer1::DataType::kFLOAT, nvinfer1::DataType::kINT8), testing::Values(2), - testing::Values(false), testing::Values(false), testing::Values(true), testing::Values(false), - testing::Values(false), testing::Values(0), testing::Values(128))); + testing::Values(16), testing::Values(tensorrt_llm::DataType::kFLOAT, tensorrt_llm::DataType::kINT8), + testing::Values(2), testing::Values(false), testing::Values(false), testing::Values(true), + testing::Values(false), testing::Values(false), testing::Values(0), testing::Values(128))); INSTANTIATE_TEST_CASE_P(AsymmetricCaseTestWithDPForNoMLADuplicate0, AsymmetricalCacheTestWithDP, testing::Combine(testing::Values(1, 2), testing::Values(1, 2), testing::Values(1), testing::Values(4), testing::Values(1), testing::Values(1), testing::Values(4), testing::Values(2), testing::Values(4), - testing::Values(16), testing::Values(nvinfer1::DataType::kFLOAT, nvinfer1::DataType::kINT8), testing::Values(2), - testing::Values(false), testing::Values(true, false), testing::Values(false), testing::Values(false), - testing::Values(false), testing::Values(0), testing::Values(128))); + testing::Values(16), testing::Values(tensorrt_llm::DataType::kFLOAT, tensorrt_llm::DataType::kINT8), + testing::Values(2), testing::Values(false), testing::Values(true, false), testing::Values(false), + testing::Values(false), testing::Values(false), testing::Values(0), testing::Values(128))); INSTANTIATE_TEST_CASE_P(AsymmetricCaseTestWithDPForNoMLADuplicate0EvenLayer, AsymmetricalCacheTestWithDP, testing::Combine(testing::Values(1), testing::Values(4), testing::Values(1), testing::Values(4), testing::Values(1), testing::Values(1), testing::Values(5), testing::Values(2), testing::Values(4), testing::Values(16), - testing::Values(nvinfer1::DataType::kFLOAT, nvinfer1::DataType::kINT8), testing::Values(2), + testing::Values(tensorrt_llm::DataType::kFLOAT, tensorrt_llm::DataType::kINT8), testing::Values(2), testing::Values(false), testing::Values(true, false), testing::Values(false), testing::Values(false), testing::Values(false), testing::Values(0), testing::Values(128))); INSTANTIATE_TEST_CASE_P(AsymmetricCaseTestWithDPForNoMLADuplicate1, AsymmetricalCacheTestWithDP, testing::Combine(testing::Values(1, 2), testing::Values(1, 2), testing::Values(1), testing::Values(2), testing::Values(2), testing::Values(1), testing::Values(4), testing::Values(1), testing::Values(4), - testing::Values(16), testing::Values(nvinfer1::DataType::kFLOAT, nvinfer1::DataType::kINT8), testing::Values(2), - testing::Values(false), testing::Values(true, false), testing::Values(false), testing::Values(false), - testing::Values(false), testing::Values(0), testing::Values(128))); + testing::Values(16), testing::Values(tensorrt_llm::DataType::kFLOAT, tensorrt_llm::DataType::kINT8), + testing::Values(2), testing::Values(false), testing::Values(true, false), testing::Values(false), + testing::Values(false), testing::Values(false), testing::Values(0), testing::Values(128))); INSTANTIATE_TEST_CASE_P(AsymmetricCaseTestWithDPForNoMLADuplicate2, AsymmetricalCacheTestWithDP, testing::Combine(testing::Values(4), testing::Values(1), testing::Values(1), testing::Values(4, 2), testing::Values(1), testing::Values(1), testing::Values(4), testing::Values(2), testing::Values(4), - testing::Values(16), testing::Values(nvinfer1::DataType::kFLOAT, nvinfer1::DataType::kINT8), testing::Values(2), - testing::Values(false), testing::Values(false), testing::Values(false), testing::Values(false), - testing::Values(false), testing::Values(0), testing::Values(128))); + testing::Values(16), testing::Values(tensorrt_llm::DataType::kFLOAT, tensorrt_llm::DataType::kINT8), + testing::Values(2), testing::Values(false), testing::Values(false), testing::Values(false), + testing::Values(false), testing::Values(false), testing::Values(0), testing::Values(128))); INSTANTIATE_TEST_CASE_P(AsymmetricCaseTestWithDPForNoMLADuplicate3, AsymmetricalCacheTestWithDP, testing::Combine(testing::Values(2), testing::Values(1), testing::Values(1), testing::Values(4), testing::Values(1), testing::Values(1), testing::Values(4), testing::Values(2), testing::Values(4), testing::Values(16), - testing::Values(nvinfer1::DataType::kFLOAT, nvinfer1::DataType::kINT8), testing::Values(2), + testing::Values(tensorrt_llm::DataType::kFLOAT, tensorrt_llm::DataType::kINT8), testing::Values(2), testing::Values(false), testing::Values(false), testing::Values(true), testing::Values(false), testing::Values(false), testing::Values(0), testing::Values(128))); INSTANTIATE_TEST_CASE_P(AsymmetricCaseTestWithDPForNoMLADuplicate4, AsymmetricalCacheTestWithDP, testing::Combine(testing::Values(4), testing::Values(1), testing::Values(1), testing::Values(1, 2), testing::Values(2), testing::Values(1), testing::Values(4), testing::Values(1, 2), testing::Values(4), - testing::Values(16), testing::Values(nvinfer1::DataType::kFLOAT, nvinfer1::DataType::kINT8), testing::Values(2), - testing::Values(false), testing::Values(false), testing::Values(false), testing::Values(false), - testing::Values(false), testing::Values(0), testing::Values(128))); + testing::Values(16), testing::Values(tensorrt_llm::DataType::kFLOAT, tensorrt_llm::DataType::kINT8), + testing::Values(2), testing::Values(false), testing::Values(false), testing::Values(false), + testing::Values(false), testing::Values(false), testing::Values(0), testing::Values(128))); #endif @@ -2139,7 +2140,7 @@ TEST(targetTest, CacheStateNODP) int const numHeads = 2; int const sizePerHead = 64; int const tokensPerBlock = 64; - auto const dataType = nvinfer1::DataType::kFLOAT; + auto const dataType = tensorrt_llm::DataType::kFLOAT; bool const isMLA = true; int const kvFactor = 2; @@ -2429,7 +2430,7 @@ TEST(targetTest, CacheStateNODPForGQAWithCP) int const numHeads = 4; int const sizePerHead = 64; int const tokensPerBlock = 64; - auto const dataType = nvinfer1::DataType::kFLOAT; + auto const dataType = tensorrt_llm::DataType::kFLOAT; bool const isMLA = false; int const kvFactor = 2; @@ -2645,7 +2646,7 @@ TEST(targetTest, CacheStateContextDP) int const numHeads = 2; int const sizePerHead = 64; int const tokensPerBlock = 64; - auto const dataType = nvinfer1::DataType::kFLOAT; + auto const dataType = tensorrt_llm::DataType::kFLOAT; bool const isMLA = true; int const kvFactor = 2; diff --git a/cpp/tests/unit_tests/multi_gpu/kernels/allReduce/allReduceFusionTest.cu b/cpp/tests/unit_tests/multi_gpu/kernels/allReduce/allReduceFusionTest.cu index 4b9c7af29a46..e95c84d7d05a 100644 --- a/cpp/tests/unit_tests/multi_gpu/kernels/allReduce/allReduceFusionTest.cu +++ b/cpp/tests/unit_tests/multi_gpu/kernels/allReduce/allReduceFusionTest.cu @@ -25,6 +25,7 @@ #include #include +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/kernels/communicationKernels/allReduceFusionKernels.h" #include "tensorrt_llm/kernels/communicationKernels/allReduceWorkspace.h" #include "tensorrt_llm/kernels/quantization.h" @@ -235,28 +236,28 @@ template <> struct DTypeTraits { static constexpr ncclDataType_t kNCCLDataType = ncclFloat16; - static constexpr nvinfer1::DataType kTRTDataType = nvinfer1::DataType::kHALF; + static constexpr tensorrt_llm::DataType kTRTDataType = tensorrt_llm::DataType::kHALF; }; template <> struct DTypeTraits<__nv_bfloat16> { static constexpr ncclDataType_t kNCCLDataType = ncclBfloat16; - static constexpr nvinfer1::DataType kTRTDataType = nvinfer1::DataType::kBF16; + static constexpr tensorrt_llm::DataType kTRTDataType = tensorrt_llm::DataType::kBF16; }; template <> struct DTypeTraits { static constexpr ncclDataType_t kNCCLDataType = ncclFloat32; - static constexpr nvinfer1::DataType kTRTDataType = nvinfer1::DataType::kFLOAT; + static constexpr tensorrt_llm::DataType kTRTDataType = tensorrt_llm::DataType::kFLOAT; }; template class TestRunner { static constexpr ncclDataType_t kNCCLDataType = DTypeTraits::kNCCLDataType; - static constexpr nvinfer1::DataType kTRTDataType = DTypeTraits::kTRTDataType; + static constexpr tensorrt_llm::DataType kTRTDataType = DTypeTraits::kTRTDataType; static constexpr bool kFP4QuantOutSupport = !std::is_same_v; static_assert(kFP4QuantOutSupport || Pattern != ar_fusion::AllReduceFusionPattern::kARResidualRMSNormFP4Quant, "kARResidualRMSNormFP4Quant is not supported for float dtype"); diff --git a/cpp/tests/unit_tests/multi_gpu/kernels/allReduce/allReduceKernelTest.cu b/cpp/tests/unit_tests/multi_gpu/kernels/allReduce/allReduceKernelTest.cu index b3d120e7015c..78747373accf 100644 --- a/cpp/tests/unit_tests/multi_gpu/kernels/allReduce/allReduceKernelTest.cu +++ b/cpp/tests/unit_tests/multi_gpu/kernels/allReduce/allReduceKernelTest.cu @@ -32,6 +32,7 @@ #include #include +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/kernels/customAllReduceKernels.h" #include "tensorrt_llm/runtime/bufferManager.h" #include "tensorrt_llm/runtime/cudaStream.h" @@ -212,7 +213,7 @@ public: { } - void set_params(AllReduceParams& params, nvinfer1::DataType dataType, int token_num, int hidden_size, + void set_params(AllReduceParams& params, tensorrt_llm::DataType dataType, int token_num, int hidden_size, AllReduceFusionOp op) const { int world_size = world_config.getSize(); @@ -316,7 +317,7 @@ bool test(Workspace const& workspace, int token_num, int hidden_size, bool has_b in.copy_from(input_buffer.data()); AllReduceParams params; - workspace.set_params(params, nvinfer1::DataType::kHALF, token_num, hidden_size, fusion_op); + workspace.set_params(params, tensorrt_llm::DataType::kHALF, token_num, hidden_size, fusion_op); params.ranks_per_node = world_size; params.local_rank = rank; params.local_output_buffer_ptr = out.data(); @@ -334,21 +335,21 @@ bool test(Workspace const& workspace, int token_num, int hidden_size, bool has_b cudaEventCreate(&begin); cudaEventCreate(&end); lamportInitialize( - params.fusion_params.lamport_peer_comm_buffer_ptrs[rank], message_size, nvinfer1::DataType::kHALF, s); + params.fusion_params.lamport_peer_comm_buffer_ptrs[rank], message_size, tensorrt_llm::DataType::kHALF, s); lamportInitialize(params.fusion_params.lamport_peer_comm_buffer_ptrs[rank + MAX_RANKS_PER_NODE], message_size, - nvinfer1::DataType::kHALF, s); + tensorrt_llm::DataType::kHALF, s); lamportInitialize(params.fusion_params.lamport_peer_comm_buffer_ptrs[rank + MAX_RANKS_PER_NODE * 2], message_size, - nvinfer1::DataType::kHALF, s); + tensorrt_llm::DataType::kHALF, s); cudaDeviceSynchronize(); comm.barrier(); for (int i = 0; i < warmup; ++i) { - customAllReduce(params, nvinfer1::DataType::kHALF, runtime_strategy, config, fusion_op, s); + customAllReduce(params, tensorrt_llm::DataType::kHALF, runtime_strategy, config, fusion_op, s); } cudaEventRecord(begin, s); for (int i = 0; i < iter; ++i) { - customAllReduce(params, nvinfer1::DataType::kHALF, runtime_strategy, config, fusion_op, s); + customAllReduce(params, tensorrt_llm::DataType::kHALF, runtime_strategy, config, fusion_op, s); } cudaEventRecord(end, s); cudaEventSynchronize(end); @@ -462,7 +463,7 @@ bool test_prepostnorm(Workspace const& workspace, int token_num, int hidden_size in.copy_from(input_buffer.data()); AllReduceParams params; - workspace.set_params(params, nvinfer1::DataType::kHALF, token_num, hidden_size, fusion_op); + workspace.set_params(params, tensorrt_llm::DataType::kHALF, token_num, hidden_size, fusion_op); params.ranks_per_node = world_size; params.local_rank = rank; params.local_output_buffer_ptr = out.data(); @@ -484,12 +485,12 @@ bool test_prepostnorm(Workspace const& workspace, int token_num, int hidden_size comm.barrier(); for (int i = 0; i < warmup; ++i) { - customAllReduce(params, nvinfer1::DataType::kHALF, runtime_strategy, config, fusion_op, s); + customAllReduce(params, tensorrt_llm::DataType::kHALF, runtime_strategy, config, fusion_op, s); } cudaEventRecord(begin, s); for (int i = 0; i < iter; ++i) { - customAllReduce(params, nvinfer1::DataType::kHALF, runtime_strategy, config, fusion_op, s); + customAllReduce(params, tensorrt_llm::DataType::kHALF, runtime_strategy, config, fusion_op, s); } cudaEventRecord(end, s); cudaEventSynchronize(end); diff --git a/cpp/tests/unit_tests/multi_gpu/kernels/allReduce/gemmAllReduceTest.cu b/cpp/tests/unit_tests/multi_gpu/kernels/allReduce/gemmAllReduceTest.cu index 415a3c32da81..2ff2c8c37b9c 100644 --- a/cpp/tests/unit_tests/multi_gpu/kernels/allReduce/gemmAllReduceTest.cu +++ b/cpp/tests/unit_tests/multi_gpu/kernels/allReduce/gemmAllReduceTest.cu @@ -1,5 +1,5 @@ /* - * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION. All rights reserved. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -25,12 +25,11 @@ #else #include "allreduce_gemm_runner.h" #endif -#include "common.h" #include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/kernels/userbuffers/ub_interface.h" #include "tensorrt_llm/runtime/ipcNvlsMemory.h" #include "tensorrt_llm/runtime/utils/mpiUtils.h" -#include #include "cute/tensor.hpp" #include "cutlass/cutlass.h" @@ -53,8 +52,59 @@ #include "cutlass/util/reference/host/gett.hpp" #include "cutlass/util/reference/host/tensor_fill.h" +namespace tensorrt_llm::testing +{ + +/** + * GPU timer for recording the elapsed time across kernel(s) launched in GPU stream + */ +struct GpuTimer +{ + cudaStream_t _stream_id; + cudaEvent_t _start; + cudaEvent_t _stop; + + /// Constructor + GpuTimer() + : _stream_id(0) + { + TLLM_CUDA_CHECK(cudaEventCreate(&_start)); + TLLM_CUDA_CHECK(cudaEventCreate(&_stop)); + } + + /// Destructor + ~GpuTimer() + { + TLLM_CUDA_CHECK(cudaEventDestroy(_start)); + TLLM_CUDA_CHECK(cudaEventDestroy(_stop)); + } + + /// Start the timer for a given stream (defaults to the default stream) + void start(cudaStream_t stream_id = 0) + { + _stream_id = stream_id; + TLLM_CUDA_CHECK(cudaEventRecord(_start, _stream_id)); + } + + /// Stop the timer + void stop() + { + TLLM_CUDA_CHECK(cudaEventRecord(_stop, _stream_id)); + } + + /// Return the elapsed time (in milliseconds) + float elapsed_millis() + { + float elapsed = 0.0; + TLLM_CUDA_CHECK(cudaEventSynchronize(_stop)); + TLLM_CUDA_CHECK(cudaEventElapsedTime(&elapsed, _start, _stop)); + return elapsed; + } +}; + +} // namespace tensorrt_llm::testing + using namespace cutlass; -using namespace nvinfer1; using namespace tensorrt_llm::mpi; using namespace tensorrt_llm::runtime; using namespace tensorrt_llm::common; @@ -238,7 +288,7 @@ struct ToType template <> struct ToType { - nvinfer1::DataType trt_value = nvinfer1::DataType::kBF16; + tensorrt_llm::DataType trt_value = tensorrt_llm::DataType::kBF16; ncclDataType_t nccl_value = ncclBfloat16; char const* str_value = "bf16"; }; @@ -246,7 +296,7 @@ struct ToType template <> struct ToType { - nvinfer1::DataType trt_value = nvinfer1::DataType::kHALF; + tensorrt_llm::DataType trt_value = tensorrt_llm::DataType::kHALF; ncclDataType_t nccl_value = ncclFloat16; char const* str_value = "fp16"; }; @@ -254,7 +304,7 @@ struct ToType template <> struct ToType { - nvinfer1::DataType trt_value = nvinfer1::DataType::kFP8; + tensorrt_llm::DataType trt_value = tensorrt_llm::DataType::kFP8; ncclDataType_t nccl_value = ncclFloat8e4m3; char const* str_value = "fp8_e4m3"; }; diff --git a/cpp/tests/unit_tests/multi_gpu/kernels/allReduce/moeAllReduceFusionTest.cu b/cpp/tests/unit_tests/multi_gpu/kernels/allReduce/moeAllReduceFusionTest.cu index 8abccf214b7a..4b81a4133fa9 100644 --- a/cpp/tests/unit_tests/multi_gpu/kernels/allReduce/moeAllReduceFusionTest.cu +++ b/cpp/tests/unit_tests/multi_gpu/kernels/allReduce/moeAllReduceFusionTest.cu @@ -24,6 +24,7 @@ #include #include +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/kernels/communicationKernels/allReduceWorkspace.h" #include "tensorrt_llm/kernels/communicationKernels/moeAllReduceFusionKernels.h" #include "tensorrt_llm/kernels/quantization.h" @@ -392,8 +393,8 @@ class MoEARFuseTestRunner { static_assert(std::is_same_v || std::is_same_v); static constexpr ncclDataType_t kNCCLDataType = std::is_same_v ? ncclFloat16 : ncclBfloat16; - static constexpr nvinfer1::DataType kTRTDataType - = std::is_same_v ? nvinfer1::DataType::kHALF : nvinfer1::DataType::kBF16; + static constexpr tensorrt_llm::DataType kTRTDataType + = std::is_same_v ? tensorrt_llm::DataType::kHALF : tensorrt_llm::DataType::kBF16; public: MoEARFuseTestRunner(int max_token_num, int hidden_dim, int max_expert_num) diff --git a/cpp/tests/unit_tests/multi_gpu/mpiUtilsTest.cpp b/cpp/tests/unit_tests/multi_gpu/mpiUtilsTest.cpp index 221cd98b5f02..c8a04f658409 100644 --- a/cpp/tests/unit_tests/multi_gpu/mpiUtilsTest.cpp +++ b/cpp/tests/unit_tests/multi_gpu/mpiUtilsTest.cpp @@ -20,7 +20,6 @@ #include "tensorrt_llm/runtime/utils/mpiUtils.h" #if ENABLE_MULTI_DEVICE -#include "tensorrt_llm/plugins/common/plugin.h" #include #endif // ENABLE_MULTI_DEVICE @@ -92,11 +91,6 @@ TEST(MPIUtils, BroadcastNcclId) EXPECT_TRUE(std::any_of( id.internal, id.internal + sizeof(id.internal) / sizeof(id.internal[0]), [](auto x) { return x != 0; })); } - -TEST(MPIUtils, GlobalSessionHandle) -{ - EXPECT_EQ(tensorrt_llm::plugins::getCommSessionHandle(), &COMM_SESSION); -} #endif // ENABLE_MULTI_DEVICE template diff --git a/cpp/tests/unit_tests/runtime/CMakeLists.txt b/cpp/tests/unit_tests/runtime/CMakeLists.txt index c022ba31ebca..3a171ee39877 100644 --- a/cpp/tests/unit_tests/runtime/CMakeLists.txt +++ b/cpp/tests/unit_tests/runtime/CMakeLists.txt @@ -35,7 +35,6 @@ add_gtest(samplingConfigTest samplingConfigTest.cpp) add_gtest(samplingTest samplingTest.cpp) add_gtest(sanitizerTest sanitizerTest.cpp) add_gtest(tllmBuffersTest tllmBuffersTest.cpp) -add_gtest(tllmRuntimeTest tllmRuntimeTest.cpp) add_gtest(transposeKVKernelTest transposeKVKernelTest.cpp) add_gtest(utilsTest utilsTest.cpp) add_gtest(virtualMemoryTest virtualMemoryTest.cpp) diff --git a/cpp/tests/unit_tests/runtime/bufferManagerTest.cpp b/cpp/tests/unit_tests/runtime/bufferManagerTest.cpp index 194cf88b6765..8bdc6d0352a0 100644 --- a/cpp/tests/unit_tests/runtime/bufferManagerTest.cpp +++ b/cpp/tests/unit_tests/runtime/bufferManagerTest.cpp @@ -17,6 +17,7 @@ #include #include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/runtime/bufferManager.h" #include "tensorrt_llm/runtime/cudaMemPool.h" @@ -121,7 +122,7 @@ TEST_F(BufferManagerTest, Pointers) static_assert(std::is_same_v); static_assert(trtPointerType.isPointer()); static_assert(trtPointerType.getDataType() == TRTDataType::value); - static_assert(static_cast(trtPointerType) == BufferDataType::kTrtPointerType); + static_assert(static_cast(trtPointerType) == BufferDataType::kTrtPointerType); static_assert(trtPointerType == BufferDataType::kTrtPointerType); // uses implicit type conversion // The C++ type corresponding to the TensorRT type for storing pointers (int64_t) using cppStorageType = DataTypeTraits::type; diff --git a/cpp/tests/unit_tests/runtime/decodingLayerWorkspaceTest.cpp b/cpp/tests/unit_tests/runtime/decodingLayerWorkspaceTest.cpp index bb6ce6410ad5..74f8faa37c87 100644 --- a/cpp/tests/unit_tests/runtime/decodingLayerWorkspaceTest.cpp +++ b/cpp/tests/unit_tests/runtime/decodingLayerWorkspaceTest.cpp @@ -16,6 +16,7 @@ #include "tensorrt_llm/runtime/decodingLayerWorkspace.h" #include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/common/workspace.h" #include #include @@ -121,8 +122,9 @@ auto const tensorDataTypesTuples = testing::Combine(tensorDataTypes, tensorDataT auto const tensorShapeTuples = testing::Combine(tensorDimensions, tensorDimensions, tensorDimensions); auto const mirrorInWorkspaceParams = testing::Combine(tensorDataTypesTuples, tensorShapeTuples, randomSeeds); -using MirrorInWorkspaceParamType = std::tuple, - std::tuple, std::uint64_t>; +using MirrorInWorkspaceParamType + = std::tuple, + std::tuple, std::uint64_t>; class MirrorInWorkspaceTest : public testing::TestWithParam { diff --git a/cpp/tests/unit_tests/runtime/gptDecoderBatchedTest.cpp b/cpp/tests/unit_tests/runtime/gptDecoderBatchedTest.cpp index 15476899fce3..a979c9a4699f 100644 --- a/cpp/tests/unit_tests/runtime/gptDecoderBatchedTest.cpp +++ b/cpp/tests/unit_tests/runtime/gptDecoderBatchedTest.cpp @@ -17,10 +17,11 @@ #include "tensorrt_llm/runtime/gptDecoderBatched.h" #include "tensorrt_llm/batch_manager/createNewDecoderRequests.h" #include "tensorrt_llm/batch_manager/decoderBuffers.h" -#include "tensorrt_llm/batch_manager/makeDecodingBatchInputOutput.h" +#include "tensorrt_llm/batch_manager/llmRequest.h" #include "tensorrt_llm/common/assert.h" #include "tensorrt_llm/common/logger.h" #include "tensorrt_llm/common/memoryUtils.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/executor/types.h" #include "tensorrt_llm/runtime/bufferManager.h" #include "tensorrt_llm/runtime/common.h" @@ -47,6 +48,62 @@ using TensorPtr = ITensor::SharedPtr; namespace { +// Local copy of the former MakeDecodingBatchInputOutput::createDecoderBatchInputs +// helper, which was removed with the TensorRT-engine execution path. The decoder +// under test is backend-agnostic; this builds its step-batched inputs directly. +void createDecoderBatchInputs(tb::DecoderInputBuffers& inputBuffers, std::vector const& activeSlots, + decoder::DecoderState const& decoderState) +{ + auto const& numDecodingEngineTokens = decoderState.getNumDecodingEngineTokens(); + auto const& maxDecodingEngineTokens = decoderState.getMaxDecodingEngineTokens(); + auto const& maxDecodingDecoderTokens = decoderState.getMaxDecodingDecoderTokens(); + auto const maxDecoderSteps = tc::ceilDiv(maxDecodingEngineTokens, maxDecodingDecoderTokens); + + auto& batchSlots = inputBuffers.forwardBatchSlots; + auto& decoderLogits = inputBuffers.decoderLogits; + + for (SizeType32 step = 0; step < maxDecoderSteps; ++step) + { + batchSlots.at(step)->resize(activeSlots.size()); + } + + auto constexpr singleRequest = 1; + + std::vector batchSizes(maxDecoderSteps); + std::vector> batchLogits(maxDecoderSteps); + auto maxActiveDecoderSteps = 1; + for (size_t batchIdx = 0; batchIdx < activeSlots.size(); ++batchIdx) + { + auto const slot = activeSlots.at(batchIdx); + auto const& logits = decoderLogits.at(batchIdx); + + auto const numDecoderSteps = tc::ceilDiv(numDecodingEngineTokens.at(slot), maxDecodingDecoderTokens); + maxActiveDecoderSteps = std::max(maxActiveDecoderSteps, numDecoderSteps); + for (SizeType32 step = 0; step < numDecoderSteps; ++step) + { + auto batchSlotsRange = BufferRange(*batchSlots.at(step)); + batchSlotsRange[batchSizes[step]] = slot; + batchSizes[step]++; + auto logitsSlice = ITensor::slice(logits, step, singleRequest); + batchLogits[step].emplace_back(std::move(logitsSlice)); + } + } + + for (SizeType32 step = 0; step < maxDecoderSteps; ++step) + { + batchSlots.at(step)->resize(batchSizes[step]); + } + batchLogits.resize(maxActiveDecoderSteps); + + inputBuffers.maxDecoderSteps = maxActiveDecoderSteps; + inputBuffers.batchLogits = batchLogits; +} + +} // namespace + +namespace +{ + std::shared_ptr createLlmRequest(SizeType32 batchSlot, SizeType32 inputLengths, SizeType32 generatedTokensPerSteps, SizeType32 acceptedTokensPerStep, TokenIdType inputTokenId, TokenIdType expectedTokenId, SizeType32 maxNewTokens, SamplingConfig const& samplingConfig, TokenIdType endId) @@ -93,7 +150,7 @@ std::vector> createLlmRequests(std::vector> const& requests, TensorPtr const& batchSlots, - nvinfer1::DataType logitsType, ModelConfig const& modelConfig, WorldConfig const& worldConfig, + tensorrt_llm::DataType logitsType, ModelConfig const& modelConfig, WorldConfig const& worldConfig, tle::DecodingConfig const& decodingConfig, GptDecoderBatched& decoder, CudaStream const& runtimeStream, SizeType32 maxSequenceLength, tb::DecoderInputBuffers& inputBuffers, decoder::DecoderState& decoderState) { @@ -125,7 +182,7 @@ void newRequests(std::vector> const& requests, T } void createDecoderInputs(tb::DecoderInputBuffers& inputBuffers, SizeType32 batchSize, SizeType32 vocabSizePadded, - nvinfer1::DataType dataType, std::vector& samplingConfigs, + tensorrt_llm::DataType dataType, std::vector& samplingConfigs, std::vector const& generatedTokensPerSteps, bool computeLogProbs, BufferManager& manager) { auto& logits = inputBuffers.decoderLogits; @@ -242,8 +299,8 @@ void verifyResults(BufferManager& manager, decoder::DecoderState const& decoderS } } -void testDecoder(nvinfer1::DataType const dtype, std::vector& samplingConfigs, SizeType32 maxBeamWidth, - bool computeLogProbs) +void testDecoder(tensorrt_llm::DataType const dtype, std::vector& samplingConfigs, + SizeType32 maxBeamWidth, bool computeLogProbs) { TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); SizeType32 constexpr tensorParallelism{1}; @@ -345,7 +402,7 @@ void testDecoder(nvinfer1::DataType const dtype, std::vector& sa auto activeSlots = std::vector(batchSize); std::iota(activeSlots.begin(), activeSlots.end(), 0); - tb::MakeDecodingBatchInputOutput::createDecoderBatchInputs(inputBuffers, activeSlots, decoderState); + createDecoderBatchInputs(inputBuffers, activeSlots, decoderState); decoder.forward(decoderState, inputBuffers); checkSequenceLengths(*decoderState.getSequenceLengths(), expectedLengths, manager); @@ -375,7 +432,7 @@ void testDecoder(nvinfer1::DataType const dtype, std::vector& sa EXPECT_FALSE(getFinished(*decoderState.getFinishedSum(), samplingConfigs, manager)[0]); } -void testDecoderWavefront(nvinfer1::DataType const dtype, std::vector& samplingConfigs, +void testDecoderWavefront(tensorrt_llm::DataType const dtype, std::vector& samplingConfigs, SizeType32 maxBeamWidth, bool computeLogProbs) { TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); @@ -474,7 +531,7 @@ void testDecoderWavefront(nvinfer1::DataType const dtype, std::vector(batchIdx + 1); std::iota(activeSlots.begin(), activeSlots.end(), 0); - tb::MakeDecodingBatchInputOutput::createDecoderBatchInputs(inputBuffers, activeSlots, decoderState); + createDecoderBatchInputs(inputBuffers, activeSlots, decoderState); decoder.forward(decoderState, inputBuffers); advanceSequenceLengths( @@ -496,7 +553,7 @@ void testDecoderWavefront(nvinfer1::DataType const dtype, std::vector& samplingConfigs, +void testDecoderDraft(tensorrt_llm::DataType const dtype, std::vector& samplingConfigs, SizeType32 maxBeamWidth, std::vector const& generatedTokensPerSteps, std::vector const& acceptedTokensPerStep, SizeType32 maxGeneratedTokensPerStep) { @@ -631,7 +688,7 @@ void testDecoderDraft(nvinfer1::DataType const dtype, std::vector(batchSize); std::iota(activeSlots.begin(), activeSlots.end(), 0); - tb::MakeDecodingBatchInputOutput::createDecoderBatchInputs(inputBuffers, activeSlots, decoderState); + createDecoderBatchInputs(inputBuffers, activeSlots, decoderState); decoder.forward(decoderState, inputBuffers); checkSequenceLengths(*decoderState.getSequenceLengths(), expectedLengths, manager); EXPECT_THAT(getFinished(*decoderState.getFinishedSum(), samplingConfigs, manager), ::testing::Each(false)); @@ -648,11 +705,11 @@ struct BeamConfig std::vector beamWidths; }; -using ParamType = std::tuple; +using ParamType = std::tuple; std::string generateTestName(testing::TestParamInfo const& info) { - std::string name{std::get<0>(info.param) == nvinfer1::DataType::kFLOAT ? "Float" : "Half"}; + std::string name{std::get<0>(info.param) == tensorrt_llm::DataType::kFLOAT ? "Float" : "Half"}; BeamConfig const beamConfig = std::get<1>(info.param); name.append("MaxBeamWidth" + std::to_string(beamConfig.maxBeamWidth)); for (auto const beamWdith : beamConfig.beamWidths) @@ -673,7 +730,7 @@ class ParamTest : public ::testing::TestWithParam TEST_P(ParamTest, Test) { - nvinfer1::DataType const dtype{std::get<0>(GetParam())}; + tensorrt_llm::DataType const dtype{std::get<0>(GetParam())}; BeamConfig const beamConfig{std::get<1>(GetParam())}; bool const computeLogProbs{std::get<2>(GetParam())}; std::vector samplingConfigs; @@ -686,7 +743,7 @@ TEST_P(ParamTest, Test) } INSTANTIATE_TEST_SUITE_P(DecoderBwTest, ParamTest, - testing::Combine(testing::Values(nvinfer1::DataType::kFLOAT, nvinfer1::DataType::kHALF), + testing::Combine(testing::Values(tensorrt_llm::DataType::kFLOAT, tensorrt_llm::DataType::kHALF), testing::Values(BeamConfig{1, {1, 1, 1}}, BeamConfig{3, {3, 3, 3, 3}}, BeamConfig{4, {4, 4, 4}}, BeamConfig{10, {10, 10, 10}}), testing::Values(false, true)), @@ -698,7 +755,7 @@ class ParamWavefrontTest : public ::testing::TestWithParam TEST_P(ParamWavefrontTest, Test) { - nvinfer1::DataType const dtype{std::get<0>(GetParam())}; + tensorrt_llm::DataType const dtype{std::get<0>(GetParam())}; BeamConfig const beamConfig{std::get<1>(GetParam())}; bool const computeLogProbs{std::get<2>(GetParam())}; bool const normalizeLogProbs{true}; @@ -712,7 +769,7 @@ TEST_P(ParamWavefrontTest, Test) } INSTANTIATE_TEST_SUITE_P(DecoderBwTest, ParamWavefrontTest, - testing::Combine(testing::Values(nvinfer1::DataType::kFLOAT, nvinfer1::DataType::kHALF), + testing::Combine(testing::Values(tensorrt_llm::DataType::kFLOAT, tensorrt_llm::DataType::kHALF), testing::Values(BeamConfig{1, {1, 1, 1}}, BeamConfig{3, {3, 3, 3, 3}}, BeamConfig{4, {4, 4, 4}}, BeamConfig{10, {10, 10, 10}}), testing::Values(false, true)), @@ -725,7 +782,7 @@ struct DraftConfig std::vector acceptedTokensPerStep; }; -using DraftTestParamType = std::tuple; +using DraftTestParamType = std::tuple; class ParamDraftTest : public ::testing::TestWithParam { @@ -733,7 +790,7 @@ class ParamDraftTest : public ::testing::TestWithParam TEST_P(ParamDraftTest, Test) { - nvinfer1::DataType const dtype{std::get<0>(GetParam())}; + tensorrt_llm::DataType const dtype{std::get<0>(GetParam())}; BeamConfig const beamConfig{std::get<1>(GetParam())}; DraftConfig const draftConfig{std::get<2>(GetParam())}; @@ -751,7 +808,7 @@ TEST_P(ParamDraftTest, Test) } INSTANTIATE_TEST_SUITE_P(DecoderTest, ParamDraftTest, - testing::Combine(testing::Values(nvinfer1::DataType::kFLOAT, nvinfer1::DataType::kHALF), + testing::Combine(testing::Values(tensorrt_llm::DataType::kFLOAT, tensorrt_llm::DataType::kHALF), testing::Values(BeamConfig{1, {1, 1, 1}}), testing::Values( // DraftConfig{2, {1, 1, 1}, {0, 0, 0}}, DraftConfig{2, {2, 2, 2}, {1, 1, 1}}, @@ -760,7 +817,7 @@ INSTANTIATE_TEST_SUITE_P(DecoderTest, ParamDraftTest, )), [](testing::TestParamInfo const& info) { - std::string name{std::get<0>(info.param) == nvinfer1::DataType::kFLOAT ? "Float" : "Half"}; + std::string name{std::get<0>(info.param) == tensorrt_llm::DataType::kFLOAT ? "Float" : "Half"}; BeamConfig const beamConfig = std::get<1>(info.param); DraftConfig const draftConfig = std::get<2>(info.param); name.append("MaxBeamWidth" + std::to_string(beamConfig.maxBeamWidth)); diff --git a/cpp/tests/unit_tests/runtime/gptDecoderTest.cpp b/cpp/tests/unit_tests/runtime/gptDecoderTest.cpp index e1fed49293bd..5f620aa4d9c4 100644 --- a/cpp/tests/unit_tests/runtime/gptDecoderTest.cpp +++ b/cpp/tests/unit_tests/runtime/gptDecoderTest.cpp @@ -17,6 +17,7 @@ #include #include "tensorrt_llm/common/memoryUtils.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/executor/types.h" #include "tensorrt_llm/runtime/bufferManager.h" #include "tensorrt_llm/runtime/gptDecoder.h" @@ -68,7 +69,7 @@ bool forwardAndSync(std::unique_ptr const& decoder, DecodingOutput& } } -void testDecoder(nvinfer1::DataType const dtype, SamplingConfig const& samplingConfig) +void testDecoder(tensorrt_llm::DataType const dtype, SamplingConfig const& samplingConfig) { SizeType32 constexpr tensorParallelism{1}; SizeType32 constexpr pipelineParallelism{1}; @@ -140,21 +141,21 @@ void testDecoder(nvinfer1::DataType const dtype, SamplingConfig const& samplingC if (beamWidth > 1) { auto srcCacheIndirection = std::shared_ptr( - manager.gpu(ITensor::makeShape({batchSize, beamWidth, maxSeqLength}), nvinfer1::DataType::kINT32)); + manager.gpu(ITensor::makeShape({batchSize, beamWidth, maxSeqLength}), tensorrt_llm::DataType::kINT32)); manager.setZero(*srcCacheIndirection); inputs.cacheIndirection = srcCacheIndirection; } // set up outputs auto outputIds = std::shared_ptr( - manager.gpu(ITensor::makeShape({batchSize, beamWidth, maxSeqLength}), nvinfer1::DataType::kINT32)); + manager.gpu(ITensor::makeShape({batchSize, beamWidth, maxSeqLength}), tensorrt_llm::DataType::kINT32)); manager.setZero(*outputIds); auto gatheredOutputIds = std::shared_ptr( - manager.gpu(ITensor::makeShape({batchSize, beamWidth, maxSeqLength}), nvinfer1::DataType::kINT32)); + manager.gpu(ITensor::makeShape({batchSize, beamWidth, maxSeqLength}), tensorrt_llm::DataType::kINT32)); manager.setZero(*gatheredOutputIds); DecodingOutput outputs{outputIds, gatheredOutputIds}; auto newTokens - = std::shared_ptr(manager.gpu(ITensor::makeShape({batchSize, beamWidth}), nvinfer1::DataType::kINT32)); + = std::shared_ptr(manager.gpu(ITensor::makeShape({batchSize, beamWidth}), tensorrt_llm::DataType::kINT32)); manager.setZero(*newTokens); outputs.newTokens = newTokens; @@ -165,7 +166,7 @@ void testDecoder(nvinfer1::DataType const dtype, SamplingConfig const& samplingC TRTDataType::value); inputs.finishReasons = ITensor::view(outputs.finishReasons); manager.setZero(*outputs.finishReasons); - outputs.finishedSum = BufferManager::pinnedPool(ITensor::makeShape({batchSize}), nvinfer1::DataType::kINT32); + outputs.finishedSum = BufferManager::pinnedPool(ITensor::makeShape({batchSize}), tensorrt_llm::DataType::kINT32); auto finishedSumHost = bufferCast(*outputs.finishedSum); for (SizeType32 bi = 0; bi < batchSize; ++bi) { @@ -175,17 +176,17 @@ void testDecoder(nvinfer1::DataType const dtype, SamplingConfig const& samplingC if (beamWidth > 1) { auto tgtCacheIndirection = std::shared_ptr( - manager.gpu(ITensor::makeShape({batchSize, beamWidth, maxSeqLength}), nvinfer1::DataType::kINT32)); + manager.gpu(ITensor::makeShape({batchSize, beamWidth, maxSeqLength}), tensorrt_llm::DataType::kINT32)); manager.setZero(*tgtCacheIndirection); outputs.cacheIndirection = tgtCacheIndirection; auto cumLogProbs - = std::shared_ptr(manager.gpu(ITensor::makeShape({batchSize, beamWidth}), nvinfer1::DataType::kFLOAT)); + = std::shared_ptr(manager.gpu(ITensor::makeShape({batchSize, beamWidth}), tensorrt_llm::DataType::kFLOAT)); manager.setZero(*cumLogProbs); outputs.cumLogProbs = cumLogProbs; auto parentIds = std::shared_ptr( - manager.gpu(ITensor::makeShape({batchSize, beamWidth, maxSeqLength}), nvinfer1::DataType::kINT32)); + manager.gpu(ITensor::makeShape({batchSize, beamWidth, maxSeqLength}), tensorrt_llm::DataType::kINT32)); manager.setZero(*parentIds); outputs.parentIds = parentIds; } @@ -245,13 +246,13 @@ void testDecoder(nvinfer1::DataType const dtype, SamplingConfig const& samplingC } // namespace -class ParamTest : public ::testing::TestWithParam> +class ParamTest : public ::testing::TestWithParam> { }; TEST_P(ParamTest, Test) { - nvinfer1::DataType const dtype{std::get<0>(GetParam())}; + tensorrt_llm::DataType const dtype{std::get<0>(GetParam())}; SizeType32 const beamWidth{std::get<1>(GetParam())}; SamplingConfig const samplingConfig{beamWidth}; @@ -259,10 +260,11 @@ TEST_P(ParamTest, Test) } INSTANTIATE_TEST_SUITE_P(DecoderTest, ParamTest, - testing::Combine(testing::Values(nvinfer1::DataType::kFLOAT, nvinfer1::DataType::kHALF), testing::Values(1, 3)), + testing::Combine( + testing::Values(tensorrt_llm::DataType::kFLOAT, tensorrt_llm::DataType::kHALF), testing::Values(1, 3)), [](testing::TestParamInfo const& info) { - std::string name{std::get<0>(info.param) == nvinfer1::DataType::kFLOAT ? "Float" : "Half"}; + std::string name{std::get<0>(info.param) == tensorrt_llm::DataType::kFLOAT ? "Float" : "Half"}; auto const beamWidth = std::get<1>(info.param); name.append(beamWidth == 1 ? "Sampling" : "BeamWidth" + std::to_string(beamWidth)); return name; diff --git a/cpp/tests/unit_tests/runtime/iTensorTest.cpp b/cpp/tests/unit_tests/runtime/iTensorTest.cpp index 54ba8aa3beec..4637474f72c7 100644 --- a/cpp/tests/unit_tests/runtime/iTensorTest.cpp +++ b/cpp/tests/unit_tests/runtime/iTensorTest.cpp @@ -17,6 +17,7 @@ #include #include +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/runtime/bufferManager.h" #include "tensorrt_llm/runtime/iTensor.h" @@ -25,7 +26,7 @@ using namespace tensorrt_llm::runtime; TEST(ITensorTest, SqueezeTensor) { auto dims = ITensor::makeShape({16, 1, 4}); - auto constexpr dataType = nvinfer1::DataType::kFLOAT; + auto constexpr dataType = tensorrt_llm::DataType::kFLOAT; ITensor::SharedPtr tensor{BufferManager::cpu(dims, dataType)}; auto squeezeDim = 0; @@ -102,7 +103,7 @@ TEST(ITensorTest, UnsqueezeTensor) auto oldShape = ITensor::makeShape({2, 3, 4, 5}); { - auto tensor = BufferManager::cpu(oldShape, nvinfer1::DataType::kINT32); + auto tensor = BufferManager::cpu(oldShape, tensorrt_llm::DataType::kINT32); tensor->unsqueeze(0); auto shape = tensor->getShape(); @@ -114,7 +115,7 @@ TEST(ITensorTest, UnsqueezeTensor) EXPECT_EQ(shape.d[4], 5); } { - auto tensor = BufferManager::cpu(oldShape, nvinfer1::DataType::kINT32); + auto tensor = BufferManager::cpu(oldShape, tensorrt_llm::DataType::kINT32); tensor->unsqueeze(1); auto shape = tensor->getShape(); @@ -127,7 +128,7 @@ TEST(ITensorTest, UnsqueezeTensor) } { - auto tensor = BufferManager::cpu(oldShape, nvinfer1::DataType::kINT32); + auto tensor = BufferManager::cpu(oldShape, tensorrt_llm::DataType::kINT32); tensor->unsqueeze(4); auto shape = tensor->getShape(); @@ -144,7 +145,7 @@ TEST(ITensorTest, UnsqueezeTensor) { try { - auto tensor = BufferManager::cpu(oldShape, nvinfer1::DataType::kINT32); + auto tensor = BufferManager::cpu(oldShape, tensorrt_llm::DataType::kINT32); tensor->unsqueeze(invalidDim); FAIL() << "Expected failure"; } @@ -162,7 +163,7 @@ TEST(ITensorTest, UnsqueezeTensor) TEST(ITensorTest, TensorView) { auto const dims = ITensor::makeShape({16, 1, 4}); - auto constexpr dataType = nvinfer1::DataType::kFLOAT; + auto constexpr dataType = tensorrt_llm::DataType::kFLOAT; ITensor::SharedPtr tensor = BufferManager::cpu(dims, dataType); auto const viewDims = ITensor::makeShape({16, 1, 2}); @@ -180,7 +181,7 @@ TEST(ITensorTest, TensorView) TEST(ITensorTest, TensorSlice) { auto dims = ITensor::makeShape({16, 8, 4}); - auto constexpr dataType = nvinfer1::DataType::kFLOAT; + auto constexpr dataType = tensorrt_llm::DataType::kFLOAT; ITensor::SharedPtr tensor{BufferManager::cpu(dims, dataType)}; auto offset = dims.d[0] / 4; auto slice = ITensor::slice(tensor, offset); @@ -221,7 +222,7 @@ TEST(ITensorTest, TensorSlice) TEST(ITensorTest, TensorDimsSliceAtManual) { auto shape = ITensor::makeShape({5, 5, 5, 5, 5}); - auto constexpr dataType = nvinfer1::DataType::kFLOAT; + auto constexpr dataType = tensorrt_llm::DataType::kFLOAT; ITensor::SharedPtr tensor(BufferManager::cpu(shape, dataType)); auto offsetDims = ITensor::makeShape({4, 3, 3}); auto sizeDim = 2; @@ -282,7 +283,7 @@ TEST(ITensorTest, TensorDimsSliceAtManual) TEST(ITensorTest, TensorDimsSliceAtExtrame) { - auto constexpr dataType = nvinfer1::DataType::kFLOAT; + auto constexpr dataType = tensorrt_llm::DataType::kFLOAT; { auto shape = ITensor::makeShape({5, 5, 5, 5, 5}); ITensor::SharedPtr tensor(BufferManager::cpu(shape, dataType)); @@ -540,7 +541,7 @@ TEST(ShapeRange, test) TEST(ITensorTest, TensorDimsSliceAt) { auto shape = ITensor::makeShape({5, 5, 5, 5}); - auto constexpr dataType = nvinfer1::DataType::kFLOAT; + auto constexpr dataType = tensorrt_llm::DataType::kFLOAT; ITensor::SharedPtr tensor(BufferManager::cpu(shape, dataType)); auto verify = [&shape, &tensor, &dataType](ITensor::Shape const& index) @@ -657,7 +658,7 @@ TEST(ITensorTest, TensorDimsSliceAt) TEST(BufferRangeTest, ConstType) { auto shape = ITensor::makeShape({5, 5, 5, 5, 5}); - auto constexpr dataType = nvinfer1::DataType::kFLOAT; + auto constexpr dataType = tensorrt_llm::DataType::kFLOAT; ITensor::SharedPtr tensor(BufferManager::cpu(shape, dataType)); ITensor::SharedConstPtr tensorConst = tensor; @@ -694,7 +695,7 @@ TEST(BufferRangeTest, ConstType) TEST(ITensorTest, GetDimension) { auto shape = ITensor::makeShape({10, 11, 12}); - auto constexpr dataType = nvinfer1::DataType::kFLOAT; + auto constexpr dataType = tensorrt_llm::DataType::kFLOAT; ITensor::SharedPtr tensor(BufferManager::cpu(shape, dataType)); auto firstDimensionFromStart = tensor->getDimension<0>(); diff --git a/cpp/tests/unit_tests/runtime/loraCacheTest.cpp b/cpp/tests/unit_tests/runtime/loraCacheTest.cpp index 4d4dc86dc824..6a91d11df3db 100644 --- a/cpp/tests/unit_tests/runtime/loraCacheTest.cpp +++ b/cpp/tests/unit_tests/runtime/loraCacheTest.cpp @@ -28,7 +28,7 @@ #include "tensorrt_llm/runtime/utils/numpyUtils.h" #include "tensorrt_llm/runtime/worldConfig.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include #include @@ -78,7 +78,7 @@ class LoraCacheTest : public ::testing::Test, void SetUp() override { - mModelConfig = std::make_unique(0, 2, 2, 0, 1, 16, nvinfer1::DataType::kFLOAT); + mModelConfig = std::make_unique(0, 2, 2, 0, 1, 16, tensorrt_llm::DataType::kFLOAT); mModelConfig->setMlpHiddenSize(32); mWorldConfig = std::make_unique(2, 1, 1, 0); std::vector modules{ @@ -101,7 +101,7 @@ class LoraCacheTest : public ::testing::Test, mManager = std::make_unique(mStream); auto pageConfig = LoraCachePageManagerConfig( - runtime::MemoryType::kCPU, nvinfer1::DataType::kFLOAT, 2 * 8, 6, 64, 4 * 16, 1); + runtime::MemoryType::kCPU, tensorrt_llm::DataType::kFLOAT, 2 * 8, 6, 64, 4 * 16, 1); pageConfig.setInitToZero(true); auto pageConfig2 = pageConfig; pageConfig2.setInitToZero(true); @@ -125,7 +125,7 @@ TEST_F(LoraCacheTest, LoraCachePageManagerTest) auto pageShape = ITensor::makeShape({maxAdapterSize, maxAdapterWeights}); LoraCachePageManagerConfig config( - runtime::MemoryType::kCPU, nvinfer1::DataType::kFLOAT, 8, 6, maxAdapterSize, maxAdapterWeights, 1); + runtime::MemoryType::kCPU, tensorrt_llm::DataType::kFLOAT, 8, 6, maxAdapterSize, maxAdapterWeights, 1); LoraCachePageManager manager(config, *mManager); auto block0 = manager.blockPtr(0); @@ -182,11 +182,11 @@ TEST_F(LoraCacheTest, LoraCachePageManagerTest) TEST_F(LoraCacheTest, determineNumPages) { - ModelConfig modelConfig(0, 2, 2, 0, 1, 4, nvinfer1::DataType::kFLOAT); + ModelConfig modelConfig(0, 2, 2, 0, 1, 4, tensorrt_llm::DataType::kFLOAT); modelConfig.setLoraModules(LoraModule::createLoraModules({"attn_dense", "attn_qkv"}, 4, 4, 1, 1, 2, 2, 0)); WorldConfig worldConfig(1, 1, 1, 0); - LoraCachePageManagerConfig pageConfig(MemoryType::kCPU, nvinfer1::DataType::kFLOAT, 12393, 40, 80, 16, 1); + LoraCachePageManagerConfig pageConfig(MemoryType::kCPU, tensorrt_llm::DataType::kFLOAT, 12393, 40, 80, 16, 1); LoraCache cache(pageConfig, modelConfig, worldConfig, *mManager); @@ -374,7 +374,7 @@ TEST_F(LoraCacheTest, basicPutGet) TEST_F(LoraCacheTest, splitTransposeCpu) { - auto modelConfig = ModelConfig(0, 2, 2, 0, 1, 16, nvinfer1::DataType::kFLOAT); + auto modelConfig = ModelConfig(0, 2, 2, 0, 1, 16, tensorrt_llm::DataType::kFLOAT); auto worldConfig = WorldConfig(2, 1, 1, 0); SizeType32 const split{2}; @@ -391,8 +391,8 @@ TEST_F(LoraCacheTest, splitTransposeCpu) auto const outputShape = ITensor::makeShape({batchSize, inputLength / split}); auto inputTensor = mManager->copyFrom(input, inputShape, MemoryType::kCPU); - auto outputTensorRank0 = mManager->cpu(outputShape, nvinfer1::DataType::kINT32); - auto outputTensorRank1 = mManager->cpu(outputShape, nvinfer1::DataType::kINT32); + auto outputTensorRank0 = mManager->cpu(outputShape, tensorrt_llm::DataType::kINT32); + auto outputTensorRank1 = mManager->cpu(outputShape, tensorrt_llm::DataType::kINT32); mManager->setZero(*outputTensorRank0); mManager->setZero(*outputTensorRank1); @@ -416,8 +416,8 @@ TEST_F(LoraCacheTest, splitTransposeCpu) auto const outputShape = ITensor::makeShape({batchSize, inputLength / split}); auto inputTensor = mManager->copyFrom(input, inputShape, MemoryType::kCPU); - auto outputTensorRank0 = mManager->cpu(outputShape, nvinfer1::DataType::kINT32); - auto outputTensorRank1 = mManager->cpu(outputShape, nvinfer1::DataType::kINT32); + auto outputTensorRank0 = mManager->cpu(outputShape, tensorrt_llm::DataType::kINT32); + auto outputTensorRank1 = mManager->cpu(outputShape, tensorrt_llm::DataType::kINT32); mManager->setZero(*outputTensorRank0); mManager->setZero(*outputTensorRank1); @@ -438,7 +438,7 @@ TEST_F(LoraCacheTest, splitTransposeCpu) TEST_P(LoraCacheTest, copyToPages_tp1) { bool const isDora = GetParam(); - auto modelConfig = ModelConfig(0, 2, 2, 0, 1, 16, nvinfer1::DataType::kFLOAT); + auto modelConfig = ModelConfig(0, 2, 2, 0, 1, 16, tensorrt_llm::DataType::kFLOAT); modelConfig.setMlpHiddenSize(32); auto worldConfig = WorldConfig(1, 1, 1, 0); std::vector modules{ @@ -501,7 +501,7 @@ TEST_P(LoraCacheTest, copyToPages_tp1) TEST_P(LoraCacheTest, copyToPages_tp2_rank0) { bool const isDora = GetParam(); - auto modelConfig = ModelConfig(0, 2, 2, 0, 1, 16, nvinfer1::DataType::kFLOAT); + auto modelConfig = ModelConfig(0, 2, 2, 0, 1, 16, tensorrt_llm::DataType::kFLOAT); modelConfig.setMlpHiddenSize(32); auto worldConfig = WorldConfig(2, 1, 1, 0); std::vector modules{ @@ -562,7 +562,7 @@ TEST_P(LoraCacheTest, copyToPages_tp2_rank0) TEST_P(LoraCacheTest, copyToPages_tp2_rank1) { bool const isDora = GetParam(); - auto modelConfig = ModelConfig(0, 2, 2, 0, 1, 16, nvinfer1::DataType::kFLOAT); + auto modelConfig = ModelConfig(0, 2, 2, 0, 1, 16, tensorrt_llm::DataType::kFLOAT); modelConfig.setMlpHiddenSize(32); auto worldConfig = WorldConfig(2, 1, 1, 1); std::vector modules{ diff --git a/cpp/tests/unit_tests/runtime/loraManagerTest.cpp b/cpp/tests/unit_tests/runtime/loraManagerTest.cpp index 6910719da76f..11c19d22efb0 100644 --- a/cpp/tests/unit_tests/runtime/loraManagerTest.cpp +++ b/cpp/tests/unit_tests/runtime/loraManagerTest.cpp @@ -32,6 +32,7 @@ #include "tensorrt_llm/runtime/modelConfig.h" #include "tensorrt_llm/runtime/worldConfig.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/runtime/utils/numpyUtils.h" #include @@ -66,7 +67,7 @@ class LoraManagerTest { protected: LoraManagerTest() - : mModelConfig(1, 2, 2, 0, 1, 4, nvinfer1::DataType::kFLOAT) + : mModelConfig(1, 2, 2, 0, 1, 4, tensorrt_llm::DataType::kFLOAT) { } @@ -87,7 +88,7 @@ class LoraManagerTest PeftTable getPeftTable(SizeType32 tpRank = 0) { - auto modelConfig = ModelConfig(0, 2, 2, 0, 1, 16, nvinfer1::DataType::kFLOAT); + auto modelConfig = ModelConfig(0, 2, 2, 0, 1, 16, tensorrt_llm::DataType::kFLOAT); modelConfig.setMlpHiddenSize(32); auto worldConfig = WorldConfig(2, 2, 1, 3); std::vector modules{ @@ -102,7 +103,7 @@ class LoraManagerTest }; modelConfig.setLoraModules(modules); auto pageConfig = LoraCachePageManagerConfig( - runtime::MemoryType::kCPU, nvinfer1::DataType::kFLOAT, 2 * 8, 6, 64, 4 * 16, 1); + runtime::MemoryType::kCPU, tensorrt_llm::DataType::kFLOAT, 2 * 8, 6, 64, 4 * 16, 1); pageConfig.setInitToZero(true); LoraCache loraCache(pageConfig, modelConfig, worldConfig, *mManager); @@ -213,7 +214,7 @@ static void checkLoraTensors(LoraManager const& loraManager, std::vectorsecond; auto actualTensor = inputTensors.find(fieldName)->second; ITensor::shapeEquals(expectedTensor->getShape(), actualTensor->getShape()); - if (expectedTensor->getDataType() == nvinfer1::DataType::kINT64) + if (expectedTensor->getDataType() == tensorrt_llm::DataType::kINT64) { auto expT = bufferCast(*expectedTensor); auto actT = bufferCast(*actualTensor); @@ -308,7 +309,7 @@ TEST_P(LoraManagerTest, fillInputTensors) bool const isDora = GetParam(); LoraManager loraManager; - auto modelConfig = ModelConfig(0, 2, 2, 0, 1, 16, nvinfer1::DataType::kFLOAT); + auto modelConfig = ModelConfig(0, 2, 2, 0, 1, 16, tensorrt_llm::DataType::kFLOAT); modelConfig.setMlpHiddenSize(32); auto worldConfig = WorldConfig(1, 1, 1, 0); std::vector modules{ @@ -332,9 +333,9 @@ TEST_P(LoraManagerTest, fillInputTensors) auto numLayers = static_cast(modelConfig.getNbAttentionLayers()); SizeType32 numSeqs = 4; TensorPtr weightsPtrs - = mManager->cpu(ITensor::makeShape({numModules, numLayers, numSeqs, 3}), nvinfer1::DataType::kINT64); + = mManager->cpu(ITensor::makeShape({numModules, numLayers, numSeqs, 3}), tensorrt_llm::DataType::kINT64); TensorPtr adapterSizes - = mManager->cpu(ITensor::makeShape({numModules, numLayers, numSeqs}), nvinfer1::DataType::kINT32); + = mManager->cpu(ITensor::makeShape({numModules, numLayers, numSeqs}), tensorrt_llm::DataType::kINT32); mManager->setZero(*weightsPtrs); mManager->setZero(*adapterSizes); diff --git a/cpp/tests/unit_tests/runtime/loraUtilsTest.cpp b/cpp/tests/unit_tests/runtime/loraUtilsTest.cpp index a14fa7bb8c47..994a77acf818 100644 --- a/cpp/tests/unit_tests/runtime/loraUtilsTest.cpp +++ b/cpp/tests/unit_tests/runtime/loraUtilsTest.cpp @@ -25,7 +25,7 @@ #include "tensorrt_llm/runtime/modelConfig.h" #include "tensorrt_llm/runtime/worldConfig.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include #include @@ -53,7 +53,7 @@ class LoraUtilsTest : public ::testing::Test // NOLINT(cppcoreguidelines-pro-typ TEST_F(LoraUtilsTest, null_values) { std::optional optReqLoraWeights = std::nullopt; - std::optional optReqLoraConfig = mManager->emptyTensor(MemoryType::kCPU, nvinfer1::DataType::kHALF); + std::optional optReqLoraConfig = mManager->emptyTensor(MemoryType::kCPU, tensorrt_llm::DataType::kHALF); EXPECT_THAT([&]() { loraValidateRequestTensorDims(optReqLoraWeights, optReqLoraConfig); }, testing::Throws()); @@ -66,33 +66,35 @@ TEST_F(LoraUtilsTest, null_values) TEST_F(LoraUtilsTest, dims_mem_type) { - std::optional optReqLoraWeights = mManager->cpu(ITensor::makeShape({1, 2}), nvinfer1::DataType::kHALF); + std::optional optReqLoraWeights + = mManager->cpu(ITensor::makeShape({1, 2}), tensorrt_llm::DataType::kHALF); std::optional optReqLoraConfig - = mManager->cpu(ITensor::makeShape({1, 2, 3}), nvinfer1::DataType::kINT32); + = mManager->cpu(ITensor::makeShape({1, 2, 3}), tensorrt_llm::DataType::kINT32); EXPECT_THAT([&]() { loraValidateRequestTensorDims(optReqLoraWeights, optReqLoraConfig); }, testing::Throws()); - std::optional optGpuWeights = mManager->gpu(ITensor::makeShape({1, 2, 50}), nvinfer1::DataType::kHALF); + std::optional optGpuWeights + = mManager->gpu(ITensor::makeShape({1, 2, 50}), tensorrt_llm::DataType::kHALF); EXPECT_THAT([&]() { loraValidateRequestTensorDims(optGpuWeights, optReqLoraConfig); }, testing::Throws()); - optReqLoraWeights = mManager->cpu(ITensor::makeShape({1, 2, 50}), nvinfer1::DataType::kHALF); - optReqLoraConfig = mManager->cpu(ITensor::makeShape({1, 2, 3}), nvinfer1::DataType::kINT32); + optReqLoraWeights = mManager->cpu(ITensor::makeShape({1, 2, 50}), tensorrt_llm::DataType::kHALF); + optReqLoraConfig = mManager->cpu(ITensor::makeShape({1, 2, 3}), tensorrt_llm::DataType::kINT32); loraValidateRequestTensorDims(optReqLoraWeights, optReqLoraConfig); } TEST_F(LoraUtilsTest, loraValidateRequestTensors) { - auto modelConfig = ModelConfig(0, 2, 2, 0, 1, 4, nvinfer1::DataType::kFLOAT); + auto modelConfig = ModelConfig(0, 2, 2, 0, 1, 4, tensorrt_llm::DataType::kFLOAT); auto worldConfig = WorldConfig(); std::optional optReqLoraWeights - = mManager->cpu(ITensor::makeShape({1, 2, 32}), nvinfer1::DataType::kFLOAT); + = mManager->cpu(ITensor::makeShape({1, 2, 32}), tensorrt_llm::DataType::kFLOAT); std::optional optReqLoraConfig - = mManager->cpu(ITensor::makeShape({1, 2, 3}), nvinfer1::DataType::kINT32); + = mManager->cpu(ITensor::makeShape({1, 2, 3}), tensorrt_llm::DataType::kINT32); std::vector config{1, 0, 4, 1, 1, 4}; diff --git a/cpp/tests/unit_tests/runtime/medusaModuleTest.cpp b/cpp/tests/unit_tests/runtime/medusaModuleTest.cpp index 3cba1bf2994d..8e7ff6c75459 100644 --- a/cpp/tests/unit_tests/runtime/medusaModuleTest.cpp +++ b/cpp/tests/unit_tests/runtime/medusaModuleTest.cpp @@ -20,7 +20,7 @@ #include "tensorrt_llm/runtime/iTensor.h" #include "tensorrt_llm/runtime/utils/speculativeChoicesUtils.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include #include @@ -58,14 +58,16 @@ class MedusaModuleTest : public ::testing::Test // NOLINT(cppcoreguidelines-pro- auto const tokensPerStep = medusaModule.getMaxDecodingTokens(); // batch size = 1 here. - TensorPtr medusaGenerationLengthsHost = mManager->pinned(ITensor::makeShape({1}), nvinfer1::DataType::kINT32); + TensorPtr medusaGenerationLengthsHost + = mManager->pinned(ITensor::makeShape({1}), tensorrt_llm::DataType::kINT32); TensorPtr medusaPositionOffsetsHost - = mManager->pinned(ITensor::makeShape({tokensPerStep}), nvinfer1::DataType::kINT32); - TensorPtr medusaTreeIdsHost = mManager->pinned(ITensor::makeShape({tokensPerStep}), nvinfer1::DataType::kINT32); + = mManager->pinned(ITensor::makeShape({tokensPerStep}), tensorrt_llm::DataType::kINT32); + TensorPtr medusaTreeIdsHost + = mManager->pinned(ITensor::makeShape({tokensPerStep}), tensorrt_llm::DataType::kINT32); TensorPtr medusaPathsHost - = mManager->pinned(ITensor::makeShape({tokensPerStep, medusaHeads + 1}), nvinfer1::DataType::kINT32); + = mManager->pinned(ITensor::makeShape({tokensPerStep, medusaHeads + 1}), tensorrt_llm::DataType::kINT32); TensorPtr attentionPackedMaskHost - = mManager->pinned(ITensor::makeShape({tokensPerStep, numPackedMasks}), nvinfer1::DataType::kINT32); + = mManager->pinned(ITensor::makeShape({tokensPerStep, numPackedMasks}), tensorrt_llm::DataType::kINT32); std::vector topKs; utils::initTensorsFromChoices(medusaModule, choices, topKs, medusaGenerationLengthsHost, diff --git a/cpp/tests/unit_tests/runtime/runtimeKernelTest.cpp b/cpp/tests/unit_tests/runtime/runtimeKernelTest.cpp index dd517c5de5a4..58372a6bd479 100644 --- a/cpp/tests/unit_tests/runtime/runtimeKernelTest.cpp +++ b/cpp/tests/unit_tests/runtime/runtimeKernelTest.cpp @@ -23,7 +23,7 @@ #include "tensorrt_llm/runtime/iTensor.h" #include "tensorrt_llm/runtime/runtimeKernels.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include #include @@ -85,7 +85,7 @@ TEST_F(RuntimeKernelTest, FillBufferInt8) { for (auto size : {123LLU, 1025LLU, 1LLU << 32}) { - auto buffer = mManager->gpu(size, nvinfer1::DataType::kINT8); + auto buffer = mManager->gpu(size, tensorrt_llm::DataType::kINT8); testFill(*buffer, *mManager, *mStream); } } @@ -94,7 +94,7 @@ TEST_F(RuntimeKernelTest, FillTensorInt8) { for (auto size : {123, 1025, std::numeric_limits::max()}) { - auto tensor = mManager->gpu(tr::ITensor::makeShape({size, 2}), nvinfer1::DataType::kINT8); + auto tensor = mManager->gpu(tr::ITensor::makeShape({size, 2}), tensorrt_llm::DataType::kINT8); testFill(*tensor, *mManager, *mStream); } } @@ -111,7 +111,7 @@ TEST_F(RuntimeKernelTest, ScatterHalf) auto const outputShape = tr::ITensor::makeShape({batchSize * beamWidth, inputLength}); auto inputTensor = mManager->copyFrom(input, inputShape, tr::MemoryType::kGPU); - auto outputTensor = mManager->gpu(outputShape, nvinfer1::DataType::kHALF); + auto outputTensor = mManager->gpu(outputShape, tensorrt_llm::DataType::kHALF); mManager->setZero(*outputTensor); tr::kernels::scatterTensor(*outputTensor, *inputTensor, beamWidth, *mStream); @@ -174,7 +174,7 @@ TEST_F(RuntimeKernelTest, TileInt32) auto const outputShape = tr::ITensor::makeShape({batchSize * beamWidth, inputLength}); auto inputTensor = mManager->copyFrom(input, inputShape, tr::MemoryType::kGPU); - auto outputTensor = mManager->gpu(outputShape, nvinfer1::DataType::kINT32); + auto outputTensor = mManager->gpu(outputShape, tensorrt_llm::DataType::kINT32); tr::kernels::tileTensor(*outputTensor, *inputTensor, beamWidth, *mStream); @@ -194,7 +194,7 @@ TEST_F(RuntimeKernelTest, TileHalf) auto const outputShape = tr::ITensor::makeShape({batchSize * beamWidth, inputLength}); auto inputTensor = mManager->copyFrom(input, inputShape, tr::MemoryType::kGPU); - auto outputTensor = mManager->gpu(outputShape, nvinfer1::DataType::kHALF); + auto outputTensor = mManager->gpu(outputShape, tensorrt_llm::DataType::kHALF); tr::kernels::tileTensor(*outputTensor, *inputTensor, beamWidth, *mStream); @@ -228,11 +228,11 @@ TEST_F(RuntimeKernelTest, TileInt8Large) // Scope the allocated tensors to ensure they are de-allocated before the test ends. { - auto inputTensor = mManager->gpu(inputShape, nvinfer1::DataType::kINT8); + auto inputTensor = mManager->gpu(inputShape, tensorrt_llm::DataType::kINT8); tr::kernels::invokeFill(*inputTensor, value, *mStream); mStream->synchronize(); - auto outputTensor = mManager->gpu(outputShape, nvinfer1::DataType::kINT8); + auto outputTensor = mManager->gpu(outputShape, tensorrt_llm::DataType::kINT8); tr::kernels::tileTensor(*outputTensor, *inputTensor, beamWidth, *mStream); mStream->synchronize(); @@ -257,11 +257,11 @@ void testCopyBatch(tr::SizeType64 stride, tr::BufferManager& manager, tr::CudaSt auto const bufferShape = tr::ITensor::makeShape({rows, stride}); auto const indicesShape = tr::ITensor::makeShape({numIndices}); - auto srcBufferHost = tr::BufferManager::cpu(bufferShape, nvinfer1::DataType::kINT32); - auto dstBufferDevice = manager.gpu(bufferShape, nvinfer1::DataType::kINT32); - auto srcOffsets = tr::BufferManager::pinned(indicesShape, nvinfer1::DataType::kINT64); - auto dstOffsets = tr::BufferManager::pinned(indicesShape, nvinfer1::DataType::kINT64); - auto sizes = tr::BufferManager::pinned(indicesShape, nvinfer1::DataType::kINT64); + auto srcBufferHost = tr::BufferManager::cpu(bufferShape, tensorrt_llm::DataType::kINT32); + auto dstBufferDevice = manager.gpu(bufferShape, tensorrt_llm::DataType::kINT32); + auto srcOffsets = tr::BufferManager::pinned(indicesShape, tensorrt_llm::DataType::kINT64); + auto dstOffsets = tr::BufferManager::pinned(indicesShape, tensorrt_llm::DataType::kINT64); + auto sizes = tr::BufferManager::pinned(indicesShape, tensorrt_llm::DataType::kINT64); tr::kernels::invokeFill(*dstBufferDevice, 0, stream); auto* srcBufferHostPtr = tr::bufferCast(*srcBufferHost); diff --git a/cpp/tests/unit_tests/runtime/samplingTest.cpp b/cpp/tests/unit_tests/runtime/samplingTest.cpp index dad99323164e..bb93478cf8ff 100644 --- a/cpp/tests/unit_tests/runtime/samplingTest.cpp +++ b/cpp/tests/unit_tests/runtime/samplingTest.cpp @@ -14,12 +14,12 @@ * limitations under the License. */ +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/executor/types.h" #include "tensorrt_llm/layers/dynamicDecodeLayer.h" #include "tensorrt_llm/runtime/bufferManager.h" #include "tensorrt_llm/runtime/cudaStream.h" #include "tensorrt_llm/runtime/gptDecoder.h" -#include "tensorrt_llm/runtime/tllmLogger.h" #include @@ -39,14 +39,11 @@ class SamplingTest : public ::testing::Test // NOLINT(cppcoreguidelines-pro-type if (mDeviceCount == 0) GTEST_SKIP() << "No GPUs found"; - - mLogger = std::make_shared(); } void TearDown() override {} int mDeviceCount; - std::shared_ptr mLogger; }; std::shared_ptr dynamicDecodeTest(std::shared_ptr manager, size_t vocabSize, @@ -67,10 +64,10 @@ std::shared_ptr dynamicDecodeTest(std::shared_ptrgpu(ITensor::makeShape({signedBatchSize}), nvinfer1::DataType::kINT32); + ITensor::SharedPtr gpuEndIds = manager->gpu(ITensor::makeShape({signedBatchSize}), tensorrt_llm::DataType::kINT32); manager->copy(cpuEndIds.data(), *gpuEndIds, MemoryType::kCPU); ITensor::SharedPtr gpuOutputIds = manager->gpu( - ITensor::makeShape({signedBatchSize, signedBeamWidth, signedMaxSeqLength}), nvinfer1::DataType::kINT32); + ITensor::makeShape({signedBatchSize, signedBeamWidth, signedMaxSeqLength}), tensorrt_llm::DataType::kINT32); manager->copy(cpuOutputIds.data(), *gpuOutputIds, MemoryType::kCPU); auto const decodingMode = beamWidth == 1 ? tle::DecodingMode::TopKTopP() : tle::DecodingMode::BeamSearch(); @@ -92,7 +89,7 @@ std::shared_ptr dynamicDecodeTest(std::shared_ptr(gpuEndIds, batchSlots, step, ite, localBatchSize); auto logitsShape = ITensor::makeShape({signedBatchSize, static_cast(beamWidth), static_cast(vocabSizePadded)}); - ITensor::SharedPtr inputLogits = manager->gpu(logitsShape, nvinfer1::DataType::kFLOAT); + ITensor::SharedPtr inputLogits = manager->gpu(logitsShape, tensorrt_llm::DataType::kFLOAT); forwardParams->logits = inputLogits; manager->copy(cpuLogits.data(), *inputLogits, MemoryType::kCPU); @@ -101,10 +98,10 @@ std::shared_ptr dynamicDecodeTest(std::shared_ptrstopCriteriaInputs = std::make_shared(localBatchSize); auto outputParams = std::make_shared(gpuOutputIds); - outputParams->sequenceLength = manager->gpu(ITensor::makeShape({signedBatchSize}), nvinfer1::DataType::kINT32); + outputParams->sequenceLength = manager->gpu(ITensor::makeShape({signedBatchSize}), tensorrt_llm::DataType::kINT32); manager->copy(cpuSequenceLengths.data(), *outputParams->sequenceLength.value(), MemoryType::kCPU); outputParams->newTokens - = manager->gpu(ITensor::makeShape({signedBatchSize, signedBeamWidth}), nvinfer1::DataType::kINT32); + = manager->gpu(ITensor::makeShape({signedBatchSize, signedBeamWidth}), tensorrt_llm::DataType::kINT32); outputParams->finished = manager->gpu( ITensor::makeShape({signedBatchSize, signedBeamWidth}), TRTDataType::value); diff --git a/cpp/tests/unit_tests/runtime/tllmBuffersTest.cpp b/cpp/tests/unit_tests/runtime/tllmBuffersTest.cpp index c901061695bb..4080a0f29e9e 100644 --- a/cpp/tests/unit_tests/runtime/tllmBuffersTest.cpp +++ b/cpp/tests/unit_tests/runtime/tllmBuffersTest.cpp @@ -19,6 +19,7 @@ #include "tensorrt_llm/common/cudaUtils.h" #include "tensorrt_llm/common/stringUtils.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/runtime/bufferManager.h" #include "tensorrt_llm/runtime/cudaMemPool.h" #include "tensorrt_llm/runtime/iTensor.h" @@ -230,7 +231,7 @@ TEST_F(TllmBuffersTest, DeviceBuffer) { CudaAllocatorAsync allocator{mStream, mMemPool}; { - DeviceBuffer buffer{size, nvinfer1::DataType::kFLOAT, allocator}; + DeviceBuffer buffer{size, tensorrt_llm::DataType::kFLOAT, allocator}; testBuffer(buffer, sizeof(float)); } streamPtr->synchronize(); @@ -242,7 +243,7 @@ TEST_F(TllmBuffersTest, DeviceBuffer) { CudaAllocator allocator{}; { - StaticDeviceBuffer buffer{size, nvinfer1::DataType::kFLOAT, allocator}; + StaticDeviceBuffer buffer{size, tensorrt_llm::DataType::kFLOAT, allocator}; testBuffer(buffer, sizeof(float)); } streamPtr->synchronize(); @@ -263,10 +264,10 @@ TEST_F(TllmBuffersTest, DeviceTensor) GTEST_SKIP() << noPoolSkipReason; } auto streamPtr = std::make_shared(); - nvinfer1::Dims constexpr dims{3, 16, 8, 4}; + tensorrt_llm::Dims constexpr dims{3, 16, 8, 4}; CudaAllocatorAsync allocator{streamPtr, mMemPool}; { - DeviceTensor tensor{dims, nvinfer1::DataType::kFLOAT, allocator}; + DeviceTensor tensor{dims, tensorrt_llm::DataType::kFLOAT, allocator}; EXPECT_EQ(tensor.getSize(), ITensor::volume(dims)); testBuffer(tensor, sizeof(float)); EXPECT_EQ(tensor.getSize(), ITensor::volume(tensor.getShape())); @@ -281,7 +282,7 @@ TEST_F(TllmBuffersTest, BufferSlice) { auto constexpr size = 1024; HostAllocator allocator{}; - auto constexpr dataType = nvinfer1::DataType::kFLOAT; + auto constexpr dataType = tensorrt_llm::DataType::kFLOAT; auto buffer = std::make_shared(size, dataType, allocator); auto offset = size / 8; auto slice = IBuffer::slice(buffer, offset); @@ -319,7 +320,7 @@ TEST_F(TllmBuffersTest, BufferOutput) CudaAllocatorAsync allocator{streamPtr, mMemPool}; for (std::size_t size : {0, 16}) { - DeviceBuffer buffer{size, nvinfer1::DataType::kFLOAT, allocator}; + DeviceBuffer buffer{size, tensorrt_llm::DataType::kFLOAT, allocator}; TLLM_CUDA_CHECK(cudaMemsetAsync(buffer.data(), 0, buffer.getSizeInBytes(), streamPtr->get())); streamPtr->synchronize(); std::stringstream ss; @@ -343,11 +344,11 @@ TEST_F(TllmBuffersTest, TensorOutput) } auto streamPtr = std::make_shared(); - nvinfer1::Dims constexpr dims{3, 16, 8, 4}; + tensorrt_llm::Dims constexpr dims{3, 16, 8, 4}; CudaAllocatorAsync allocator{streamPtr, mMemPool}; - for (auto dataType : - {nvinfer1::DataType::kFLOAT, nvinfer1::DataType::kHALF, nvinfer1::DataType::kBOOL, nvinfer1::DataType::kINT8, - nvinfer1::DataType::kINT32, nvinfer1::DataType::kINT64, nvinfer1::DataType::kUINT8}) + for (auto dataType : {tensorrt_llm::DataType::kFLOAT, tensorrt_llm::DataType::kHALF, tensorrt_llm::DataType::kBOOL, + tensorrt_llm::DataType::kINT8, tensorrt_llm::DataType::kINT32, tensorrt_llm::DataType::kINT64, + tensorrt_llm::DataType::kUINT8}) { DeviceTensor tensor{dims, dataType, allocator}; TLLM_CUDA_CHECK(cudaMemsetAsync(tensor.data(), 0, tensor.getSizeInBytes(), streamPtr->get())); @@ -483,8 +484,8 @@ TEST_F(TllmBuffersTest, PinnedPoolAllocator) EXPECT_EQ(segments.size(), 0); { - auto a = BufferManager::pinnedPool(ITensor::makeShape({512, 4, 4}), nvinfer1::DataType::kFLOAT); - auto b = BufferManager::pinnedPool(ITensor::makeShape({512, 10}), nvinfer1::DataType::kHALF); + auto a = BufferManager::pinnedPool(ITensor::makeShape({512, 4, 4}), tensorrt_llm::DataType::kFLOAT); + auto b = BufferManager::pinnedPool(ITensor::makeShape({512, 10}), tensorrt_llm::DataType::kHALF); pool.logSegments(); auto it = std::begin(segments); EXPECT_NE(it->tag, nullptr); @@ -512,7 +513,7 @@ TEST_F(TllmBuffersTest, PinnedPoolAllocator) std::size_t secondChunkSize; { // Test creating a new chunk - auto c = BufferManager::pinnedPool(ITensor::makeShape({initChunkSize + 1}), nvinfer1::DataType::kUINT8); + auto c = BufferManager::pinnedPool(ITensor::makeShape({initChunkSize + 1}), tensorrt_llm::DataType::kUINT8); pool.logSegments(); auto it = std::begin(segments); EXPECT_EQ(it->tag, nullptr); diff --git a/cpp/tests/unit_tests/runtime/tllmRuntimeTest.cpp b/cpp/tests/unit_tests/runtime/tllmRuntimeTest.cpp deleted file mode 100644 index b3a6f99146c2..000000000000 --- a/cpp/tests/unit_tests/runtime/tllmRuntimeTest.cpp +++ /dev/null @@ -1,147 +0,0 @@ -/* - * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. - * - * 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. - */ - -#ifndef TOP_LEVEL_DIR -#error "Define TOP_LEVEL_DIR" -#endif - -#include -#include -#include - -#include "tensorrt_llm/common/cudaUtils.h" -#include "tensorrt_llm/runtime/rawEngine.h" -#include "tensorrt_llm/runtime/tllmLogger.h" -#include "tensorrt_llm/runtime/tllmRuntime.h" - -#include -#include -#include -#include -#include - -namespace fs = std::filesystem; -namespace trt = nvinfer1; - -namespace -{ -auto const TEST_RESOURCE_DIR = fs::path{TOP_LEVEL_DIR} / "cpp/tests/resources"; -auto const MNIST_MODEL_PATH = TEST_RESOURCE_DIR / "models/mnist.onnx"; - -template -std::unique_ptr makeUnique(T* ptr) -{ - EXPECT_NE(ptr, nullptr); - return std::unique_ptr(ptr); -} - -std::unique_ptr buildMnistEngine(trt::ILogger& logger) -{ - EXPECT_TRUE(fs::exists(MNIST_MODEL_PATH)); - auto builder = makeUnique(trt::createInferBuilder(logger)); - auto const explicitBatch = 1U << static_cast(trt::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH); - auto network = makeUnique(builder->createNetworkV2(explicitBatch)); - auto parser = makeUnique(nvonnxparser::createParser(*network, logger)); - auto const parsingSuccess = parser->parseFromFile( - MNIST_MODEL_PATH.string().c_str(), static_cast(trt::ILogger::Severity::kWARNING)); - EXPECT_TRUE(parsingSuccess); - auto config = makeUnique(builder->createBuilderConfig()); - return makeUnique(builder->buildSerializedNetwork(*network, *config)); -} -} // namespace - -using namespace tensorrt_llm::runtime; -namespace tc = tensorrt_llm::common; - -class TllmRuntimeTest : public ::testing::Test // NOLINT(cppcoreguidelines-pro-type-member-init) -{ -protected: - void SetUp() override - { - mDeviceCount = tc::getDeviceCount(); - - if (mDeviceCount == 0) - GTEST_SKIP(); - - mLogger.setLevel(trt::ILogger::Severity::kINFO); - mSerializedEngine = buildMnistEngine(mLogger); - ASSERT_NE(mSerializedEngine, nullptr); - } - - void TearDown() override {} - - int mDeviceCount; - TllmLogger mLogger{}; - std::unique_ptr mSerializedEngine; -}; - -TEST_F(TllmRuntimeTest, SinglePass) -{ - EXPECT_TRUE(mSerializedEngine); - TllmRuntime rt{RawEngine(mSerializedEngine.get()), &mLogger, false, 1.0F}; - auto& engine = rt.getEngine(); - EXPECT_FALSE(engine.hasImplicitBatchDimension()); - EXPECT_EQ(rt.getNbProfiles(), engine.getNbOptimizationProfiles()); - EXPECT_EQ(rt.getNbContexts(), 0); - auto const nbIoTensors = engine.getNbIOTensors(); - EXPECT_EQ(nbIoTensors, 2); - rt.addContext(0); - EXPECT_EQ(rt.getNbContexts(), 1); - - auto constexpr dataType = trt::DataType::kFLOAT; - - auto const inputName = engine.getIOTensorName(0); - EXPECT_EQ(engine.getTensorIOMode(inputName), trt::TensorIOMode::kINPUT); - auto const inputDims = engine.getTensorShape(inputName); - std::array constexpr inputDimsExpected = {1, 1, 28, 28}; - EXPECT_EQ(inputDims.nbDims, inputDimsExpected.size()); - for (int i = 0; i < inputDims.nbDims; ++i) - { - EXPECT_EQ(inputDims.d[i], inputDimsExpected[i]); - } - EXPECT_EQ(engine.getTensorDataType(inputName), dataType); - - auto const outputName = engine.getIOTensorName(1); - EXPECT_EQ(engine.getTensorIOMode(outputName), trt::TensorIOMode::kOUTPUT); - auto const outputDims = engine.getTensorShape(outputName); - std::array constexpr outputDimsExpected = {1, 10}; - EXPECT_EQ(outputDims.nbDims, outputDimsExpected.size()); - for (int i = 0; i < outputDims.nbDims; ++i) - { - EXPECT_EQ(outputDims.d[i], outputDimsExpected[i]); - } - EXPECT_EQ(engine.getTensorDataType(outputName), dataType); - - auto& allocator = rt.getBufferManager(); - TllmRuntime::TensorMap tensorMap{}; - auto inputBuffer = std::shared_ptr{allocator.gpu(inputDims, dataType)}; - allocator.setZero(*inputBuffer); - tensorMap.insert(std::make_pair(inputName, inputBuffer)); - rt.setInputTensors(0, tensorMap); - rt.setOutputTensors(0, tensorMap); - ASSERT_NE(tensorMap.find(outputName), tensorMap.end()); - auto outputBuffer = tensorMap.at(outputName); - allocator.setZero(*outputBuffer); - rt.executeContext(0); - - std::vector output(outputBuffer->getSize()); - allocator.copy(*outputBuffer, output.data()); - rt.getStream().synchronize(); - auto min = std::min_element(output.begin(), output.end()); - EXPECT_NEAR(*min, -0.126409f, 1e-5f); - auto max = std::max_element(output.begin(), output.end()); - EXPECT_NEAR(*max, 0.140218f, 1e-5f); -} diff --git a/cpp/tests/unit_tests/runtime/torchTest.cpp b/cpp/tests/unit_tests/runtime/torchTest.cpp index 4aa498de8d25..4ca36d875d44 100644 --- a/cpp/tests/unit_tests/runtime/torchTest.cpp +++ b/cpp/tests/unit_tests/runtime/torchTest.cpp @@ -17,6 +17,7 @@ #include #include +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/runtime/bufferManager.h" #include "tensorrt_llm/runtime/torch.h" #include "tensorrt_llm/runtime/torchView.h" @@ -52,7 +53,7 @@ class TorchTest : public ::testing::Test // NOLINT(cppcoreguidelines-pro-type-me namespace { -template +template void checkFilled(IBuffer& buffer, int fillValue) { if (DType == buffer.getDataType()) @@ -79,13 +80,13 @@ TEST_F(TorchTest, Aten) } auto constexpr fillValue = 1; - auto tensorHostBase = manager.allocate(MemoryType::kPINNED, shapeTllm, nvinfer1::DataType::kINT64); + auto tensorHostBase = manager.allocate(MemoryType::kPINNED, shapeTllm, tensorrt_llm::DataType::kINT64); for (auto memoryType : {MemoryType::kCPU, MemoryType::kGPU, MemoryType::kPINNED}) { - for (auto dtype : {nvinfer1::DataType::kFLOAT, nvinfer1::DataType::kHALF, nvinfer1::DataType::kINT8, - nvinfer1::DataType::kUINT8, nvinfer1::DataType::kINT32, nvinfer1::DataType::kINT64, - nvinfer1::DataType::kBF16, nvinfer1::DataType::kFP8, nvinfer1::DataType::kBOOL}) + for (auto dtype : {tensorrt_llm::DataType::kFLOAT, tensorrt_llm::DataType::kHALF, tensorrt_llm::DataType::kINT8, + tensorrt_llm::DataType::kUINT8, tensorrt_llm::DataType::kINT32, tensorrt_llm::DataType::kINT64, + tensorrt_llm::DataType::kBF16, tensorrt_llm::DataType::kFP8, tensorrt_llm::DataType::kBOOL}) { ITensor::SharedPtr tensorTllm{manager.allocate(memoryType, shapeTllm, dtype)}; @@ -98,20 +99,20 @@ TEST_F(TorchTest, Aten) EXPECT_THAT(tensorAten.sizes(), ::testing::ElementsAreArray(shapeAten)); EXPECT_EQ(tensorAten.data_ptr(), tensorTllm->data()); - if (dtype != nvinfer1::DataType::kFP8) + if (dtype != tensorrt_llm::DataType::kFP8) { tensorAten.fill_(c10::Scalar(fillValue)); auto tensorHost = ITensor::wrap(tensorHostBase->data(), dtype, shapeTllm); manager.copy(*tensorTllm, *tensorHost); mStream->synchronize(); - checkFilled(*tensorHost, fillValue); - checkFilled(*tensorHost, fillValue); - checkFilled(*tensorHost, fillValue); - checkFilled(*tensorHost, fillValue); - checkFilled(*tensorHost, fillValue); - checkFilled(*tensorHost, fillValue); - checkFilled(*tensorHost, fillValue); - checkFilled(*tensorHost, fillValue); + checkFilled(*tensorHost, fillValue); + checkFilled(*tensorHost, fillValue); + checkFilled(*tensorHost, fillValue); + checkFilled(*tensorHost, fillValue); + checkFilled(*tensorHost, fillValue); + checkFilled(*tensorHost, fillValue); + checkFilled(*tensorHost, fillValue); + checkFilled(*tensorHost, fillValue); } // Conversion back to TRT-LLM tensor diff --git a/cpp/tests/unit_tests/runtime/utilsTest.cpp b/cpp/tests/unit_tests/runtime/utilsTest.cpp index 8b69070c03d8..58882824a7dc 100644 --- a/cpp/tests/unit_tests/runtime/utilsTest.cpp +++ b/cpp/tests/unit_tests/runtime/utilsTest.cpp @@ -18,6 +18,7 @@ #error "Define TOP_LEVEL_DIR" #endif +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/runtime/bufferManager.h" #include "tensorrt_llm/runtime/iBuffer.h" #include "tensorrt_llm/runtime/iTensor.h" @@ -71,7 +72,7 @@ TEST_F(UtilsTest, LoadNpy) TEST_F(UtilsTest, LoadStoreNpy) { auto dims = ITensor::makeShape({2, 3, 4}); - auto constexpr dataType = nvinfer1::DataType::kFLOAT; + auto constexpr dataType = tensorrt_llm::DataType::kFLOAT; ITensor::SharedPtr tensor{BufferManager::cpu(dims, dataType)}; auto tensorRange = BufferRange(*tensor); std::iota(tensorRange.begin(), tensorRange.end(), 0); @@ -96,7 +97,7 @@ TEST_F(UtilsTest, LoadStoreNpy) TEST_F(UtilsTest, LoadStoreNpyGPU) { auto dims = ITensor::makeShape({2, 3, 4}); - auto constexpr dataType = nvinfer1::DataType::kFLOAT; + auto constexpr dataType = tensorrt_llm::DataType::kFLOAT; ITensor::SharedPtr tensor{BufferManager::cpu(dims, dataType)}; auto tensorRange = BufferRange(*tensor); std::iota(tensorRange.begin(), tensorRange.end(), 0); diff --git a/cpp/tests/unit_tests/runtime/virtualMemoryTest.cpp b/cpp/tests/unit_tests/runtime/virtualMemoryTest.cpp index f2045e7659d1..159a07770694 100644 --- a/cpp/tests/unit_tests/runtime/virtualMemoryTest.cpp +++ b/cpp/tests/unit_tests/runtime/virtualMemoryTest.cpp @@ -19,6 +19,7 @@ #include "tensorrt_llm/common/cudaUtils.h" #include "tensorrt_llm/common/nvmlWrapper.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/runtime/bufferManager.h" #include "tensorrt_llm/runtime/tllmBuffers.h" #include "tensorrt_llm/runtime/virtualMemory.h" @@ -1502,7 +1503,7 @@ TEST_F(VirtualMemoryManagerTest, TestCudaVirtualMemoryAllocator) // Create a buffer using the virtual address allocator auto buffer = std::make_unique( - size, nvinfer1::DataType::kINT8, CudaVirtualMemoryAllocator{config}); + size, tensorrt_llm::DataType::kINT8, CudaVirtualMemoryAllocator{config}); auto memoryAfterAllocation = getCurrentProcessMemoryInfo(); if (memoryInfoAvailable()) @@ -1513,7 +1514,7 @@ TEST_F(VirtualMemoryManagerTest, TestCudaVirtualMemoryAllocator) // Test that we can access the buffer data ASSERT_NE(buffer->data(), nullptr) << "Buffer data should not be null"; ASSERT_EQ(buffer->getSize(), size) << "Buffer size should match requested size"; - ASSERT_EQ(buffer->getDataType(), nvinfer1::DataType::kINT8) << "Buffer data type should be INT8"; + ASSERT_EQ(buffer->getDataType(), tensorrt_llm::DataType::kINT8) << "Buffer data type should be INT8"; ASSERT_EQ(buffer->getMemoryType(), MemoryType::kGPU) << "Buffer memory type should be GPU"; // Test memory access by setting memory to a known pattern @@ -1574,7 +1575,7 @@ TEST_F(VirtualMemoryManagerTest, TestCudaVirtualMemoryAllocatorUnalignedSize) // Create a buffer using the virtual address allocator auto buffer = std::make_unique( - size, nvinfer1::DataType::kINT8, CudaVirtualMemoryAllocator{config}); + size, tensorrt_llm::DataType::kINT8, CudaVirtualMemoryAllocator{config}); auto memoryAfterAllocation = getCurrentProcessMemoryInfo(); if (memoryInfoAvailable()) diff --git a/cpp/tests/unit_tests/utils/CMakeLists.txt b/cpp/tests/unit_tests/utils/CMakeLists.txt deleted file mode 100644 index 2d7e9145c817..000000000000 --- a/cpp/tests/unit_tests/utils/CMakeLists.txt +++ /dev/null @@ -1,16 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2023-2025 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. - -add_gtest(testUtilsTest utilsTest.cpp) diff --git a/cpp/tests/unit_tests/utils/utilsTest.cpp b/cpp/tests/unit_tests/utils/utilsTest.cpp deleted file mode 100644 index 38587120ef84..000000000000 --- a/cpp/tests/unit_tests/utils/utilsTest.cpp +++ /dev/null @@ -1,55 +0,0 @@ - -#include "common.h" -#include "tensorrt_llm/runtime/common.h" - -#include - -#include -#include - -struct RandomLogitsTestParameters -{ - using TupleT = std::tuple; - - int32_t randomSeed; - tensorrt_llm::runtime::SizeType32 vocabSize; - - // Constructor that takes a tuple - RandomLogitsTestParameters( // NOLINT: implicit to allow gtest to convert from tuple generated by - // 'combine' - TupleT t) - : randomSeed(std::get<0>(t)) - , vocabSize(std::get<1>(t)) - { - } -}; - -class RandomLogits : public ::testing::Test, public ::testing::WithParamInterface -{ -protected: - static constexpr int randomSeed = 2345; -}; - -namespace -{ -constexpr int32_t kRandomSeed1 = 45; -constexpr int32_t kRandomSeed2 = 567; -auto const randomSeeds = ::testing::Values(kRandomSeed1, kRandomSeed2); - -constexpr tensorrt_llm::runtime::SizeType32 kMinVocabSize = 16; -constexpr tensorrt_llm::runtime::SizeType32 kMaxVocabSize = 100000; -auto const vocabSizes = ::testing::Values(kMinVocabSize, kMaxVocabSize); - -auto const paramGenerator - = ::testing::ConvertGenerator(::testing::Combine(randomSeeds, vocabSizes)); -} // namespace - -TEST_P(RandomLogits, FloatSumToOne) -{ - auto rng = std::mt19937(randomSeed); - auto const randomLogits = tensorrt_llm::testing::randomLogits(456, &rng); - auto const sum = std::reduce(randomLogits.begin(), randomLogits.end()); - ASSERT_DOUBLE_EQ(sum, 1.0); -} - -INSTANTIATE_TEST_SUITE_P(Float, RandomLogits, paramGenerator); diff --git a/cpp/tests/utils/CMakeLists.txt b/cpp/tests/utils/CMakeLists.txt deleted file mode 100644 index 0123ac1c4600..000000000000 --- a/cpp/tests/utils/CMakeLists.txt +++ /dev/null @@ -1,19 +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. - -add_library(testingUtils common.cpp engines.cpp executorUtils.cpp) -target_link_libraries(testingUtils PUBLIC gtest_main ${SHARED_TARGET}) -target_include_directories(testingUtils PRIVATE ${MPI_C_INCLUDE_DIRS}) -target_compile_definitions(testingUtils PUBLIC TOP_LEVEL_DIR="${TOP_LEVEL_DIR}") diff --git a/cpp/tests/utils/common.cpp b/cpp/tests/utils/common.cpp deleted file mode 100644 index 5640cef7b495..000000000000 --- a/cpp/tests/utils/common.cpp +++ /dev/null @@ -1,684 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2024-2025 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. - */ - -#include "common.h" - -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/memoryUtils.h" -#include "tensorrt_llm/executor/executor.h" -#include "tensorrt_llm/executor/types.h" -#include "tensorrt_llm/runtime/iBuffer.h" -#include "tensorrt_llm/runtime/iTensor.h" -#include "tensorrt_llm/runtime/utils/numpyUtils.h" -#include "tensorrt_llm/testing/modelSpec.h" -#include "tests/utils/common.h" - -#include - -#include -#include - -namespace tensorrt_llm::testing -{ -namespace fs = std::filesystem; -namespace tr = tensorrt_llm::runtime; -namespace tc = tensorrt_llm::common; - -std::string PathUtil::FP16_GPT_ATTENTION_PACKED_DIR() -{ - return ModelSpec::getDefaultModelSpec().setKVCacheType(KVCacheType::kCONTINUOUS).getModelPath(); -} - -std::string PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() -{ - return ModelSpec::getDefaultModelSpec().getModelPath(); -} - -std::string PathUtil::FP16_GPT_LORA_DIR() -{ - return ModelSpec::getDefaultModelSpec().useLoraPlugin().getModelPath(); -} - -std::string PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DRAFT_TOKENS_DIR() -{ - return ModelSpec::getDefaultModelSpec().useDraftTokensExternalDecoding().getModelPath(); -} - -std::string PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_GATHER_DIR() -{ - return ModelSpec::getDefaultModelSpec().gatherLogits().getModelPath(); -} - -std::string PathUtil::FP16_PLUGIN_PACKED_PAGED_RESULT_FILE() -{ - return ModelSpec::getDefaultModelSpec().getResultsFile(); -} - -std::string PathUtil::FP16_PLUGIN_PACKED_PAGED_LONG_RESULT_FILE() -{ - return ModelSpec::getDefaultModelSpec().setInputFile("input_tokens_long.npy").getResultsFile(); -} - -std::string PathUtil::FP16_PLUGIN_PACKED_PAGED_GATHER_RESULT_FILE() -{ - return ModelSpec::getDefaultModelSpec().gatherLogits().getResultsFile(); -} - -std::string PathUtil::FP16_PLUGIN_PACKED_PAGED_GENERATION_LOGITS_FILE() -{ - return ModelSpec::getDefaultModelSpec().gatherLogits().getGenerationLogitsFile(); -} - -std::string PathUtil::FP16_PLUGIN_PACKED_PAGED_CONTEXT_LOGITS_FILE() -{ - return ModelSpec::getDefaultModelSpec().gatherLogits().getContextLogitsFile(); -} - -std::string PathUtil::FP16_PLUGIN_PACKED_PAGED_CUM_LOG_PROBS_FILE() -{ - return ModelSpec::getDefaultModelSpec().getCumLogProbsFile(); -} - -std::string PathUtil::FP16_PLUGIN_PACKED_PAGED_GATHER_CUM_LOG_PROBS_FILE() -{ - return ModelSpec::getDefaultModelSpec().gatherLogits().getCumLogProbsFile(); -} - -std::string PathUtil::FP16_PLUGIN_PACKED_PAGED_LOG_PROBS_FILE() -{ - return ModelSpec::getDefaultModelSpec().getLogProbsFile(); -} - -std::string PathUtil::FP16_PLUGIN_PACKED_PAGED_GATHER_LOG_PROBS_FILE() -{ - return ModelSpec::getDefaultModelSpec().gatherLogits().getLogProbsFile(); -} - -std::string PathUtil::FP16_PLUGIN_PACKED_PAGED_RESULT_TP1_PP1_FILE() -{ - return ModelSpec::getDefaultModelSpec().getResultsFile(); -} - -std::string PathUtil::FP16_PLUGIN_PACKED_PAGED_RESULT_TP4_PP1_FILE() -{ - return ModelSpec::getDefaultModelSpec().useTensorParallelism(4).getResultsFile(); -} - -std::string PathUtil::FP16_PLUGIN_PACKED_PAGED_RESULT_TP2_PP2_FILE() -{ - return ModelSpec::getDefaultModelSpec().useTensorParallelism(2).usePipelineParallelism(2).getResultsFile(); -} - -std::string PathUtil::FP16_PLUGIN_PACKED_PAGED_RESULT_TP1_PP4_FILE() -{ - return ModelSpec::getDefaultModelSpec().usePipelineParallelism(4).getResultsFile(); -} - -std::string PathUtil::FP16_PLUGIN_PACKED_PAGED_RESULT_TP1_PP2_FILE() -{ - return ModelSpec::getDefaultModelSpec().usePipelineParallelism(2).getResultsFile(); -} - -std::string PathUtil::FP16_PLUGIN_PACKED_PAGED_RESULT_TP2_PP1_FILE() -{ - return ModelSpec::getDefaultModelSpec().useTensorParallelism(2).getResultsFile(); -} - -std::string PathUtil::FP16_PLUGIN_PACKED_PAGED_CONTEXT_LOGITS_TP4_PP1_FILE() -{ - return ModelSpec::getDefaultModelSpec().useTensorParallelism(4).getContextLogitsFile(); -} - -std::string PathUtil::FP16_PLUGIN_PACKED_PAGED_GENERATION_LOGITS_TP4_PP1_FILE() -{ - return ModelSpec::getDefaultModelSpec().useTensorParallelism(4).getGenerationLogitsFile(); -} - -std::string PathUtil::FP16_PLUGIN_PACKED_PAGED_CUM_LOG_PROBS_TP4_PP1_FILE() -{ - return ModelSpec::getDefaultModelSpec().useTensorParallelism(4).getCumLogProbsFile(); -} - -std::string PathUtil::FP16_PLUGIN_PACKED_PAGED_LOG_PROBS_TP4_PP1_FILE() -{ - return ModelSpec::getDefaultModelSpec().useTensorParallelism(4).getLogProbsFile(); -} - -std::string PathUtil::FP16_PLUGIN_PACKED_PAGED_GATHER_CONTEXTFMHAFP32ACC_RESULT_FILE() -{ - return ModelSpec::getDefaultModelSpec().gatherLogits().enableContextFMHAFp32Acc().getResultsFile(); -} - -std::string PathUtil::FP16_PLUGIN_PACKED_PAGED_CONTEXTFMHAFP32ACC_GENERATION_LOGITS_FILE() -{ - return ModelSpec::getDefaultModelSpec().gatherLogits().enableContextFMHAFp32Acc().getGenerationLogitsFile(); -} - -std::string PathUtil::FP16_PLUGIN_PACKED_PAGED_CONTEXTFMHAFP32ACC_CONTEXT_LOGITS_FILE() -{ - return ModelSpec::getDefaultModelSpec().gatherLogits().enableContextFMHAFp32Acc().getContextLogitsFile(); -} - -void TestData::loadLogProbs( - fs::path const& cumLogProbsFile, fs::path const& logProbsFile, tr::BufferManager const& manager) -{ - TLLM_CHECK_WITH_INFO( - cumLogProbsFile != "", "Testing return log probs, but missing the expected cum log probs results file."); - auto expectedCumLogProbsPtr - = std::shared_ptr(tr::utils::loadNpy(manager, cumLogProbsFile.string(), MemoryType::kCPU)); - - TLLM_CHECK_WITH_INFO( - logProbsFile != "", "Testing return log probs, but missing the expected log probs results file."); - auto expectedLogProbsPtr = std::shared_ptr(tr::utils::loadNpy(manager, logProbsFile.string(), MemoryType::kCPU)); - - for (SizeType32 inputIdx = 0; inputIdx < nbGivenInputs; ++inputIdx) - { - for (SizeType32 beam = 0; beam < beamWidth; ++beam) - { - auto expectedCumLogProbsBatchSlice = std::shared_ptr(ITensor::slice(expectedCumLogProbsPtr, inputIdx, 1)); - expectedCumLogProbsBatchSlice->squeeze(0); // bs - expectedCumLogProbs[inputIdx] = expectedCumLogProbsBatchSlice; // shape: [beamWidth] - - auto expectedLogProbsBatchSlice = std::shared_ptr(ITensor::slice(expectedLogProbsPtr, inputIdx, 1)); - expectedLogProbsBatchSlice->squeeze(0); // bs - expectedLogProbs[inputIdx] = expectedLogProbsBatchSlice; // shape: [beamWidth, numOutputTokens] - } - } -} - -void TestData::loadContextLogits(fs::path const& contextLogitsFile, std::vector const& givenInputLengths, - tr::BufferManager const& manager) -{ - TLLM_CHECK_WITH_INFO(contextLogitsFile != "", - "Testing with gather or replace logits, but missing the expected context logits results file."); - auto expectedContextLogitsPtr - = std::shared_ptr(tr::utils::loadNpy(manager, contextLogitsFile.string(), MemoryType::kCPU)); - - int promptOffset = 0; - for (SizeType32 bi = 0; bi < nbGivenInputs; ++bi) - { - for (SizeType32 beam = 0; beam < beamWidth; ++beam) - { - auto expectedContextLogitBatchSlice - = std::shared_ptr(ITensor::slice(expectedContextLogitsPtr, promptOffset, givenInputLengths.at(bi))); - expectedContextLogits.at(bi) = expectedContextLogitBatchSlice; // shape: [prompt_length, vocab_size] - } - promptOffset += givenInputLengths.at(bi); - } -} - -void TestData::loadGenerationLogits(fs::path const& genLogitsFile, tr::BufferManager const& manager) -{ - TLLM_CHECK_WITH_INFO(genLogitsFile != "", - "Testing with gather or replace logits, but missing the expected generation logits results file."); - auto expectedGenerationLogitsPtr - = std::shared_ptr(tr::utils::loadNpy(manager, genLogitsFile.string(), MemoryType::kCPU)); - - for (SizeType32 bi = 0; bi < nbGivenInputs; ++bi) - { - for (SizeType32 beam = 0; beam < beamWidth; ++beam) - { - auto expectedGenerationLogitBatchSlice - = std::shared_ptr(ITensor::slice(expectedGenerationLogitsPtr, bi, 1)); - expectedGenerationLogitBatchSlice->squeeze(0); // bs - expectedGenerationLogitBatchSlice->squeeze(0); // beam - expectedGenerationLogits.at(bi) = expectedGenerationLogitBatchSlice; // shape: [max_output_len, vocab_size] - } - } -} - -void TestData::makeDraft(SizeType32 maxDraftTokens, bool acceptDraftByLogits, fs::path const& genLogitsFile, - std::vector const& givenInputLengths, tr::BufferManager const& manager) -{ - TLLM_CHECK(beamWidth == 1); - - ITensor::SharedPtr expectedGenerationLogitsPtr; - if (acceptDraftByLogits) - { - TLLM_CHECK_WITH_INFO( - genLogitsFile != "", "Testing Draft token, but missing the expected generation logits results file."); - expectedGenerationLogitsPtr - = std::shared_ptr(tr::utils::loadNpy(manager, genLogitsFile.string(), MemoryType::kCPU)); - } - - std::vector draftLengths(givenInputLengths.size()); - // first draft length stays 0 - std::transform(givenInputLengths.begin() + 1, givenInputLengths.end(), draftLengths.begin() + 1, - [this, &maxDraftTokens](auto inputLength) - { return std::rand() % std::min((maxSeqLen - (inputLength + 1)), maxDraftTokens) + 1; }); - - auto* const expectedOutputData = tr::bufferCast(*expectedOutputIds); - for (SizeType32 bi = 0; bi < nbGivenInputs; ++bi) - { - SizeType32 constexpr beamIdx{0}; - auto const endId = endIds.at(bi); - auto const draftLen = draftLengths.at(bi); - auto acceptedLen = draftLen > 0 ? std::rand() % draftLen : 0; - - if (acceptDraftByLogits && draftLen > 0) - { - auto expectedLogitBatchSlice = std::shared_ptr(ITensor::slice(expectedGenerationLogitsPtr, bi, 1)); - expectedLogitBatchSlice->squeeze(0); // bs - expectedLogitBatchSlice->squeeze(0); // beam - auto expectedLogitBatchStepSlice = std::shared_ptr(ITensor::slice(expectedLogitBatchSlice, 1, draftLen)); - auto expectedLogitBatchStepView = ITensor::view(expectedLogitBatchStepSlice, - ITensor::makeShape({draftLen, 1, 1, expectedLogitBatchStepSlice->getShape().d[1]})); - draftLogits.at(bi) = manager.copyFrom(*expectedLogitBatchStepView, MemoryType::kCPU); - } - - for (SizeType32 si = 0; si < draftLen; ++si) - { - auto const draftIndex - = tc::flat_index3(bi, beamIdx, givenInputLengths.at(bi) + si + 1, beamWidth, maxSeqLen); - auto draftToken = expectedOutputData[draftIndex]; - if (draftToken == endId) - { - acceptedLen = std::min(acceptedLen, si); - } - if (si >= acceptedLen) - { - draftToken = -1; - if (acceptDraftByLogits) - { - auto vocabSizePadded = expectedGenerationLogitsPtr->getShape().d[3]; - auto* draftLogitsPtr = tr::bufferCast(*draftLogits.at(bi)); - for (SizeType32 vi = 0; vi < vocabSizePadded; ++vi) - { - draftLogitsPtr[si * vocabSizePadded + vi] = 0.f; - } - } - } - draftTokens.at(bi).push_back(draftToken); - } - acceptedDraftTokensLengths.at(bi) = acceptedLen; - - auto const expectedLen = expectedOutputLengths.at(bi * beamWidth + beamIdx); - TLLM_CHECK(expectedLen > 0); - expectedOutputLengths[bi * beamWidth + beamIdx] - = draftLen > 0 ? std::min(expectedLen, (givenInputLengths.at(bi) + 1) + acceptedLen + 1) : expectedLen; - } -} - -template -bool invokeCompareLogits(ITensor const& groundTruthLogits, ITensor const& outputLogits, float atol, float rtol) -{ - bool allMatch = true; - T const* const gtLogitsPtr = tr::bufferCast(groundTruthLogits); - T const* const outputLogitsPtr = tr::bufferCast(outputLogits); - - size_t outputSize = outputLogits.getSize(); - int errorNumber = 0; - - for (size_t i = 0; i < outputSize; i++) - { - if (!almostEqual(outputLogitsPtr[i], gtLogitsPtr[i], atol, rtol)) - { - TLLM_LOG_DEBUG("Mismatch value. Position of logits: %d, expected value: %f, output value: %f", i, - gtLogitsPtr[i], outputLogitsPtr[i]); - allMatch = false; - errorNumber++; - if (errorNumber == 10) - { - break; - } - } - } - return allMatch; -} - -bool compareLogits(ITensor const& groundTruthLogits, ITensor const& outputLogits, float atol, float rtol) -{ - EXPECT_EQ(groundTruthLogits.getDataType(), outputLogits.getDataType()); - switch (groundTruthLogits.getDataType()) - { - case nvinfer1::DataType::kFLOAT: return invokeCompareLogits(groundTruthLogits, outputLogits, atol, rtol); - case nvinfer1::DataType::kHALF: return invokeCompareLogits(groundTruthLogits, outputLogits, atol, rtol); - default: TLLM_THROW("Unsupported data type"); - } -} - -std::tuple getRequestGivenInputIdxLength( - std::uint64_t requestId, SizeType32 nbGivenInputs, std::vector const& givenInputLengths) -{ - auto const givenInputIdx = requestId % nbGivenInputs; - auto const inputLength = givenInputLengths.at(givenInputIdx); - return {givenInputIdx, inputLength}; -} - -std::tuple, SizeType32, SizeType32> getGivenInputLengths( - ITensor const& givenInput, SizeType32 padId) -{ - auto const& inputShape = givenInput.getShape(); - auto const nbGivenInputs = static_cast(inputShape.d[0]); - auto const maxInputLength = static_cast(inputShape.d[1]); - auto const* const givenInputData = tr::bufferCast(givenInput); - - std::vector givenInputLengths(nbGivenInputs); - for (SizeType32 i = 0; i < nbGivenInputs; ++i) - { - auto const* const seqBegin = givenInputData + i * maxInputLength; - auto const* const it = std::find(seqBegin, seqBegin + maxInputLength, padId); - givenInputLengths[i] = std::distance(seqBegin, it); - } - - return {givenInputLengths, nbGivenInputs, maxInputLength}; -} - -std::vector createConsecutiveTokenSequence( - tr::SizeType32 length, tr::SizeType32 vocabSize, tr::TokenIdType firstTokenId) -{ - auto result = std::vector(static_cast(length), 0); - std::iota(result.begin(), result.end(), firstTokenId); - std::transform(result.begin(), result.end(), result.begin(), [&](auto const i) { return i % vocabSize; }); - return result; -} - -TestData TestData::loadTestData(BeamResult const& beamResults, ITensor const& givenInput, SizeType32 const maxBeamWidth, - tr::BufferManager& manager, executor::OutputConfig const& outConfig, ModelIds const& modelIds) -{ - auto const [givenInputLengths, nbGivenInputs, maxInputLength] = getGivenInputLengths(givenInput, modelIds.padId); - auto const& [beamWidth, resultsFile, contextLogitsFile, genLogitsFile, cumLogProbsFile, logProbsFile] = beamResults; - - TestData testData{nbGivenInputs, beamWidth}; - testData.expectedOutputIds = tr::utils::loadNpy(manager, resultsFile.string(), tr::MemoryType::kCPU); - - auto const& outputShape = testData.expectedOutputIds->getShape(); - EXPECT_EQ(outputShape.nbDims, 2); - EXPECT_EQ(nbGivenInputs * beamWidth, outputShape.d[0]); - testData.maxSeqLen = static_cast(outputShape.d[1]); - EXPECT_LE(maxInputLength, testData.maxSeqLen); - EXPECT_LE(beamWidth, maxBeamWidth); - - auto const maxNewTokens = testData.maxSeqLen - maxInputLength; - - testData.endIds.insert(testData.endIds.end(), nbGivenInputs, modelIds.endId); - - if (outConfig.returnContextLogits && beamWidth == 1) - { - testData.loadContextLogits(contextLogitsFile, givenInputLengths, manager); - } - if (outConfig.returnGenerationLogits && beamWidth == 1) - { - testData.loadGenerationLogits(genLogitsFile, manager); - } - if (outConfig.returnLogProbs && beamWidth == 1) - { - testData.loadLogProbs(cumLogProbsFile, logProbsFile, manager); - } - - for (SizeType32 inputIdx = 0; inputIdx < nbGivenInputs; ++inputIdx) - { - for (SizeType32 beam = 0; beam < beamWidth; ++beam) - { - SizeType32 expectedLen = givenInputLengths[inputIdx] + maxNewTokens; - testData.expectedOutputLengths[inputIdx * beamWidth + beam] = expectedLen; - } - } - - return testData; -} - -void TestData::verifyOutput(std::unordered_map> const& resultTokens, - std::vector const& givenInputLengths, bool streaming, bool excludeInputFromOutput, - FlakyTestInfo flakyTestInfo, bool isSpeculativeDecoding, SizeType32 reqBeamWidth, SizeType32 numReturnSequences, - bool isNonGreedySampling) -{ - for (auto const& [batchId, beamTokens] : resultTokens) - { - for (auto seqIdx = 0; seqIdx < numReturnSequences; seqIdx++) - { - auto const& tokens = beamTokens.at(seqIdx); - auto const inputLength = givenInputLengths.at(batchId); - SizeType32 const numReturnBeams = tokens.size(); - auto const* const expectedOutputData = tr::bufferCast(*this->expectedOutputIds); - auto const expectedOutputLengths = this->expectedOutputLengths; - auto const endId = this->endIds[batchId]; - auto const maxSeqLen = this->maxSeqLen; - - for (SizeType32 beam = 0; beam < numReturnBeams; ++beam) - { - bool isFlaky = flakyTestInfo.batchIdBeams.count(std::make_pair(batchId, beam)); - if (isFlaky) - { - TLLM_LOG_WARNING("Disabling token comparison for batchId %d beam %d, test if flaky", batchId, beam); - } - - auto const expectInputOutputLength - = expectedOutputLengths[batchId * reqBeamWidth + beam]; // Ground truth output length - auto expectedOutputLength - = expectInputOutputLength - inputLength; // Number of new generated output tokens - - bool inputNotIncluded = (streaming || excludeInputFromOutput); - bool anyMismatch = false; - auto predictedTokens = tokens.at(beam); - // Remove the prompt - if (!inputNotIncluded) - { - predictedTokens.erase(predictedTokens.begin(), predictedTokens.begin() + inputLength); - } - - if (!isNonGreedySampling) - { - EXPECT_EQ(predictedTokens.size(), expectedOutputLength) - << "b: " << batchId << " seq: " << seqIdx << " beam: " << beam; - } - - auto numPredTokens = static_cast(predictedTokens.size()); - - if (isSpeculativeDecoding) - { - // WAR to ensure bulk execution of spec decoding. - // We hope that no request in batch can finish 2x faster than any other request. - // For the cases when BS < 8, some predicted tokens are mismatched to reference data. - numPredTokens /= 2; - } - - for (auto i = 0; i < numPredTokens; ++i) - { - // Use the expected data for that beamWidth - auto const expectIndex = tc::flat_index3(batchId, beam, inputLength + i, reqBeamWidth, maxSeqLen); - auto const expectedToken = expectedOutputData[expectIndex]; - if (expectedToken == endId) - { - // TODO: can not find the error when (expectedToken == endId) && (predictedToken != endId) - break; - } - auto const predictedToken = predictedTokens.at(i); - if (!isFlaky && !isNonGreedySampling) - { - EXPECT_EQ(predictedToken, expectedToken) - << "b: " << batchId << " seq: " << seqIdx << " beam: " << beam << " i: " << i; - } - anyMismatch |= (predictedToken != expectedToken); - } - if (!isFlaky && !isNonGreedySampling) - { - EXPECT_FALSE(anyMismatch) << "b: " << batchId << " seq: " << seqIdx << " beam: " << beam; - } - else if (isNonGreedySampling) - { - EXPECT_TRUE(anyMismatch) << "b: " << batchId << " seq: " << seqIdx << " beam: " << beam; - } - } - } - } -} - -void TestData::verifyLogProbs(bool computeLogProbs, bool streaming, bool excludeInputFromOutput, SizeType32 inputLength, - SizeType32 beamWidth, executor::BeamTokens const& beamTokens, - std::optional const& cumLogProbs, - std::optional> const& logProbs, SizeType32 batchId, FlakyTestInfo flakyTestInfo) -{ - auto expectedCumLogProbs = this->expectedCumLogProbs[batchId]; - auto expectedLogProbs = this->expectedLogProbs[batchId]; - auto const expectedOutputLengths = this->expectedOutputLengths; - auto const numReturnBeams = beamTokens.size(); - - if (computeLogProbs) - { - EXPECT_TRUE(cumLogProbs.has_value()) << "bid: " << batchId; - EXPECT_TRUE(logProbs.has_value()) << "bid: " << batchId; - EXPECT_EQ(cumLogProbs.value().size(), numReturnBeams) << "bid: " << batchId; - EXPECT_EQ(logProbs.value().size(), numReturnBeams) << "bid: " << batchId; - - bool removeInput = !excludeInputFromOutput && !streaming; - - for (SizeType32 beam = 0; beam < numReturnBeams; ++beam) - { - bool isFlaky = flakyTestInfo.batchIdBeams.count(std::make_pair(batchId, beam)); - if (isFlaky) - { - TLLM_LOG_WARNING("Disabling token comparison for batchId %d beam %d, test if flaky", batchId, beam); - } - - auto expectedOutputLength = expectedOutputLengths[batchId * beamWidth + beam]; - expectedOutputLength -= inputLength; - - auto numPredTokens = logProbs.value().at(beam).size(); - // Check shape - EXPECT_EQ(numPredTokens, beamTokens.at(beam).size() - (removeInput ? inputLength : 0)) - << "bid: " << batchId << " beam: " << beam; - - // If beamWidth == 1, compare log probs against python runtime - if (beamWidth == 1) - { - auto* const reqExpectedCumLogProbs = tr::bufferCast(*expectedCumLogProbs); - // Only check cumLogProbs for the last generated token - if (numPredTokens == expectedOutputLength && !isFlaky) - { - EXPECT_TRUE(almostEqual(reqExpectedCumLogProbs[beam], cumLogProbs.value().at(beam), 2e-1, 5e-2)) - << "expectedCumLogProbs : " << reqExpectedCumLogProbs[beam] - << " cumlogProbs : " << cumLogProbs.value().at(beam); - } - - auto expectedLogProbsBeam = std::shared_ptr(tr::ITensor::slice(expectedLogProbs, beam, 1)); - expectedLogProbsBeam->squeeze(0); - auto* const reqExpectedLogProbs = tr::bufferCast(*expectedLogProbsBeam); - for (auto i = 0; i < numPredTokens; ++i) - { - if (!isFlaky) - { - EXPECT_TRUE( - almostEqual(reqExpectedLogProbs[inputLength + i], logProbs.value()[beam][i], 5e-2, 5e-2)) - << "expectedLogProbs : " << reqExpectedLogProbs[inputLength + i] - << " logProbs : " << logProbs.value()[beam][i]; - } - } - } - } - } - else - { - EXPECT_FALSE(cumLogProbs.has_value()) << "bid: " << batchId; - EXPECT_FALSE(logProbs.has_value()) << "bid: " << batchId; - } -} - -void TestData::validateContextLogits(bool getContextLogits, SizeType32 inputLength, SizeType32 beamWidth, - std::optional const& contextLogits, SizeType32 vocabSizePadded, SizeType32 batchId, float atol, - float rtol) -{ - if (getContextLogits) - { - EXPECT_TRUE(contextLogits.has_value()) << "bid: " << batchId; - EXPECT_EQ(contextLogits.value().getShape().size(), 2); - EXPECT_EQ(contextLogits.value().getShape()[0], inputLength); - EXPECT_EQ(contextLogits.value().getShape()[1], vocabSizePadded); - auto const expectedContextLogits = this->expectedContextLogits[batchId]; - - if (beamWidth == 1) - { - cudaDeviceSynchronize(); // Make sure the logits copy is complete. - EXPECT_TRUE(compareLogits( - *expectedContextLogits, *(executor::detail::toITensor(contextLogits.value())), atol, rtol)); - } - } - else - { - EXPECT_FALSE(contextLogits.has_value()) << "bid: " << batchId; - } -} - -void TestData::validateGenerationLogits(bool getGenLogits, bool isFinal, bool streaming, bool excludeInputFromOutput, - SizeType32 inputLength, SizeType32 maxOutputLen, SizeType32 beamWidth, executor::BeamTokens const& beamTokens, - std::optional const& genLogits, SizeType32 vocabSizePadded, SizeType32 batchId, - bool const returnAllGeneratedTokens, float atol, float rtol) -{ - auto const numReturnBeams = beamTokens.size(); - - if (getGenLogits) - { - EXPECT_TRUE(genLogits.has_value()) << "bid: " << batchId; - EXPECT_EQ(genLogits.value().getShape().size(), 3); - - // Expected generation logits - auto const& expectedGenerationLogits - = this->expectedGenerationLogits[batchId]; // [maxOutputLen, vocabSizePadded] - // Output generation logits - // 1. non-streaming: [beamWidth, maxOutputLen, vocabSizePadded] - // 2. streaming: [maxOutputLen (or 1), beamWidth, vocabSizePadded] - auto const& outputGenerationLogits = executor::detail::toITensor(genLogits.value()); - - if (streaming) - { - EXPECT_EQ(genLogits.value().getShape()[1], numReturnBeams); - EXPECT_EQ(beamWidth, 1); // Only support streaming && beamWidth == 1 - - SizeType32 const beamIdx = 0; - bool removeInput = !excludeInputFromOutput && !streaming; - // If returnAllGeneratedTokens, will contain duplicate tokens - auto const& numPredTokens = beamTokens.at(beamIdx).size() - (removeInput ? inputLength : 0); - - SizeType32 numGeneratedToken = genLogits.value().getShape()[0]; - if (returnAllGeneratedTokens) - { - EXPECT_EQ(numGeneratedToken, numPredTokens); - } - else - { - EXPECT_EQ(numGeneratedToken, 1); - } - SizeType32 sliceOffset = returnAllGeneratedTokens ? 0 : numPredTokens - 1; - - auto const& expectedGenerationLogitsSlice - = std::shared_ptr(ITensor::slice(expectedGenerationLogits, sliceOffset, - numGeneratedToken)); // [numGeneratedToken, vocabSizePadded] - - cudaDeviceSynchronize(); // Make sure the logits copy is complete. - EXPECT_TRUE(compareLogits(*expectedGenerationLogitsSlice, *outputGenerationLogits, atol, rtol)); - } - else - { - // Non-streaming - EXPECT_EQ(genLogits.value().getShape()[0], numReturnBeams); - EXPECT_EQ(genLogits.value().getShape()[1], maxOutputLen); - - if (isFinal && beamWidth == 1) - { - cudaDeviceSynchronize(); // Make sure the logits copy is complete. - EXPECT_TRUE(compareLogits(*expectedGenerationLogits, *outputGenerationLogits, atol, rtol)); - } - } - EXPECT_EQ(genLogits.value().getShape()[2], vocabSizePadded); - } - else - { - EXPECT_FALSE(genLogits.has_value()) << "bid: " << batchId; - } -} - -} // namespace tensorrt_llm::testing diff --git a/cpp/tests/utils/common.h b/cpp/tests/utils/common.h deleted file mode 100644 index f7b73a9acea4..000000000000 --- a/cpp/tests/utils/common.h +++ /dev/null @@ -1,352 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2024-2025 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. - */ - -#pragma once - -#ifndef TOP_LEVEL_DIR -#error "Define TOP_LEVEL_DIR" -#endif - -#include "tensorrt_llm/executor/executor.h" -#include "tensorrt_llm/executor/types.h" -#include "tensorrt_llm/runtime/bufferManager.h" -#include "tensorrt_llm/runtime/common.h" -#include "tensorrt_llm/runtime/iBuffer.h" -#include "tensorrt_llm/runtime/iTensor.h" - -#include -#include -#include -#include -#include -#include -#include - -namespace tensorrt_llm::testing -{ -namespace fs = std::filesystem; -namespace tr = tensorrt_llm::runtime; - -using tr::SizeType32; -using tr::TokenIdType; -using tr::ITensor; -using tr::MemoryType; - -auto const TEST_RESOURCE_PATH = fs::path{TOP_LEVEL_DIR} / "cpp/tests/resources"; - -auto const ENGINE_PATH = TEST_RESOURCE_PATH / "models/rt_engine"; -auto const GPT_MODEL_PATH = ENGINE_PATH / "gpt2"; -auto const LLAMA_MODEL_PATH = ENGINE_PATH / "Llama-3.2-1B"; -auto const MEDUSA_MODEL_PATH = ENGINE_PATH / "vicuna-7b-medusa"; -auto const CHATGLM_MODEL_PATH = ENGINE_PATH / "chatglm-6b"; -auto const CHATGLM2_MODEL_PATH = ENGINE_PATH / "chatglm2-6b"; -auto const CHATGLM3_MODEL_PATH = ENGINE_PATH / "chatglm3-6b"; -auto const GLM_MODEL_PATH = ENGINE_PATH / "glm-10b"; -auto const ENC_DEC_ENGINE_BASE = TEST_RESOURCE_PATH / "models/enc_dec/trt_engines"; - -auto const DATA_PATH = TEST_RESOURCE_PATH / "data"; -auto const GPT_DATA_PATH = DATA_PATH / "gpt2"; -auto const GPT_XGRAMMAR_TOKENIZER_INFO_PATH = GPT_DATA_PATH / "xgrammar_tokenizer_info.json"; -auto const LLAMA_DATA_PATH = DATA_PATH / "Llama-3.2-1B"; -auto const LLAMA_XGRAMMAR_TOKENIZER_INFO_PATH = LLAMA_DATA_PATH / "xgrammar_tokenizer_info.json"; -auto const MEDUSA_DATA_PATH = DATA_PATH / "vicuna-7b-medusa"; -auto const CHATGLM_DATA_PATH = DATA_PATH / "chatglm-6b"; -auto const CHATGLM2_DATA_PATH = DATA_PATH / "chatglm2-6b"; -auto const CHATGLM3_DATA_PATH = DATA_PATH / "chatglm3-6b"; -auto const GLM_DATA_PATH = DATA_PATH / "glm-10b"; -auto const ENC_DEC_DATA_BASE = DATA_PATH / "enc_dec"; - -auto constexpr T5_NAME = "t5-small"; -auto constexpr BART_NAME = "bart-large-cnn"; -auto constexpr LANGUAGE_ADAPTER_NAME = "language_adapter-enc_dec_language_adapter"; - -class PathUtil -{ -public: - static std::string EXECUTOR_WORKER_PATH() - { - return (std::filesystem::path{TOP_LEVEL_DIR} / "cpp/build/tensorrt_llm/executor_worker/executorWorker") - .string(); - } - - // model paths - static std::string FP16_GPT_ATTENTION_PACKED_DIR(); - static std::string FP16_GPT_ATTENTION_PACKED_PAGED_DIR(); - static std::string FP16_GPT_LORA_DIR(); - static std::string FP16_GPT_ATTENTION_PACKED_PAGED_DRAFT_TOKENS_DIR(); - static std::string FP16_GPT_ATTENTION_PACKED_PAGED_GATHER_DIR(); - static std::string FP16_PLUGIN_PACKED_PAGED_RESULT_FILE(); - static std::string FP16_PLUGIN_PACKED_PAGED_LONG_RESULT_FILE(); - static std::string FP16_PLUGIN_PACKED_PAGED_GATHER_RESULT_FILE(); - // logits - static std::string FP16_PLUGIN_PACKED_PAGED_GENERATION_LOGITS_FILE(); - static std::string FP16_PLUGIN_PACKED_PAGED_CONTEXT_LOGITS_FILE(); - static std::string FP16_PLUGIN_PACKED_PAGED_CUM_LOG_PROBS_FILE(); - static std::string FP16_PLUGIN_PACKED_PAGED_GATHER_CUM_LOG_PROBS_FILE(); - static std::string FP16_PLUGIN_PACKED_PAGED_LOG_PROBS_FILE(); - static std::string FP16_PLUGIN_PACKED_PAGED_GATHER_LOG_PROBS_FILE(); - // results - static std::string FP16_PLUGIN_PACKED_PAGED_RESULT_TP1_PP1_FILE(); - static std::string FP16_PLUGIN_PACKED_PAGED_RESULT_TP4_PP1_FILE(); - static std::string FP16_PLUGIN_PACKED_PAGED_RESULT_TP2_PP2_FILE(); - static std::string FP16_PLUGIN_PACKED_PAGED_RESULT_TP1_PP4_FILE(); - static std::string FP16_PLUGIN_PACKED_PAGED_RESULT_TP1_PP2_FILE(); - static std::string FP16_PLUGIN_PACKED_PAGED_RESULT_TP2_PP1_FILE(); - static std::string FP16_PLUGIN_PACKED_PAGED_CONTEXT_LOGITS_TP4_PP1_FILE(); - static std::string FP16_PLUGIN_PACKED_PAGED_GENERATION_LOGITS_TP4_PP1_FILE(); - static std::string FP16_PLUGIN_PACKED_PAGED_CUM_LOG_PROBS_TP4_PP1_FILE(); - static std::string FP16_PLUGIN_PACKED_PAGED_LOG_PROBS_TP4_PP1_FILE(); - // GptExecutorTest.GenerationLogitsEarlyStop requires to use context_fmha_fp32_acc flag in runtime for better - // accuracy - static std::string FP16_PLUGIN_PACKED_PAGED_GATHER_CONTEXTFMHAFP32ACC_RESULT_FILE(); - static std::string FP16_PLUGIN_PACKED_PAGED_CONTEXTFMHAFP32ACC_GENERATION_LOGITS_FILE(); - static std::string FP16_PLUGIN_PACKED_PAGED_CONTEXTFMHAFP32ACC_CONTEXT_LOGITS_FILE(); -}; - -class ModelIds -{ -public: - ModelIds() = default; - - constexpr ModelIds(TokenIdType endId, TokenIdType padId) - : endId{endId} - , padId{padId} - { - } - - TokenIdType endId{}; - TokenIdType padId{}; -}; - -class BeamResult -{ -public: - explicit BeamResult(SizeType32 beamWidth) - : beamWidth{beamWidth} {}; - - BeamResult(SizeType32 beamWidth, fs::path resultsFile, fs::path contextLogitsFile, fs::path genLogitsFile, - fs::path cumLogProbsFile, fs::path logProbsFile) - : beamWidth{beamWidth} - , resultsFile{std::move(resultsFile)} - , contextLogitsFile{std::move(contextLogitsFile)} - , genLogitsFile{std::move(genLogitsFile)} - , cumLogProbsFile{std::move(cumLogProbsFile)} - , logProbsFile{std::move(logProbsFile)} {}; - - SizeType32 beamWidth; - fs::path resultsFile; - - fs::path contextLogitsFile; - fs::path genLogitsFile; - - fs::path cumLogProbsFile; - fs::path logProbsFile; -}; - -using BeamResults = std::vector; - -struct FlakyTestInfo -{ - // Pair of batch ID + beam which are flaky - std::set> batchIdBeams; -}; - -class TestData -{ -public: - explicit TestData(SizeType32 nbGivenInputs, SizeType32 beamWidth) - : nbGivenInputs{nbGivenInputs} - , beamWidth{beamWidth} - { - expectedOutputLengths.resize(nbGivenInputs * beamWidth); - - draftTokens.resize(nbGivenInputs); - draftLogits.resize(nbGivenInputs); - acceptedDraftTokensLengths.resize(nbGivenInputs); - expectedGenerationLogits.resize(nbGivenInputs); - expectedContextLogits.resize(nbGivenInputs); - expectedCumLogProbs.resize(nbGivenInputs); - expectedLogProbs.resize(nbGivenInputs); - } - - void loadLogProbs(fs::path const& cumLogProbsFile, fs::path const& logProbsFile, tr::BufferManager const& manager); - - void loadContextLogits(fs::path const& contextLogitsFile, std::vector const& givenInputLengths, - tr::BufferManager const& manager); - void loadGenerationLogits(fs::path const& genLogitsFile, tr::BufferManager const& manager); - - void makeDraft(SizeType32 maxDraftTokens, bool acceptDraftByLogits, fs::path const& genLogitsFile, - std::vector const& givenInputLengths, tr::BufferManager const& manager); - - static TestData loadTestData(BeamResult const& beamResults, ITensor const& givenInput, SizeType32 maxBeamWidth, - tr::BufferManager& manager, executor::OutputConfig const& outConfig, ModelIds const& modelIds); - - void verifyOutput(std::unordered_map> const& resultTokens, - std::vector const& givenInputLengths, bool streaming, bool excludeInputFromOutput, - FlakyTestInfo flakyTestInfo, bool isSpeculativeDecoding, SizeType32 reqBeamWidth, SizeType32 numReturnSequences, - bool isNonGreedySampling); - - void verifyLogProbs(bool computeLogProbs, bool streaming, bool excludeInputFromOutput, SizeType32 inputLength, - SizeType32 beamWidth, executor::BeamTokens const& beamTokens, - std::optional const& cumLogProbs, - std::optional> const& logProbs, SizeType32 batchId, - FlakyTestInfo flakyTestInfo); - - void validateContextLogits(bool getContextLogits, SizeType32 inputLength, SizeType32 beamWidth, - std::optional const& contextLogits, SizeType32 vocabSizePadded, SizeType32 batchId, - float atol = 1e-2, float rtol = 1e-3); - - void validateGenerationLogits(bool getGenLogits, bool isFinal, bool streaming, bool excludeInputFromOutput, - SizeType32 inputLength, SizeType32 maxOutputLen, SizeType32 beamWidth, executor::BeamTokens const& beamTokens, - std::optional const& genLogits, SizeType32 vocabSizePadded, SizeType32 batchId, - bool returnAllGeneratedTokens, float atol = 1e-2, float rtol = 1e-3); - - SizeType32 nbGivenInputs{}; - SizeType32 beamWidth{}; - SizeType32 maxSeqLen{}; - ITensor::SharedPtr expectedOutputIds; - std::vector expectedOutputLengths; - std::vector endIds; - std::vector draftTokens; - std::vector draftLogits; - std::vector acceptedDraftTokensLengths; - std::vector expectedGenerationLogits; - std::vector expectedContextLogits; - std::vector expectedCumLogProbs; - std::vector expectedLogProbs; -}; - -inline bool almostEqual(float a, float b, float atol = 1e-2, float rtol = 1e-3) -{ - // Params: a = value to compare and b = reference - // This function follows implementation of numpy.isclose(), which checks - // abs(a - b) <= (atol + rtol * abs(b)). - // Note that the inequality above is asymmetric where b is considered as - // a reference value. To account into both absolute/relative errors, it - // uses absolute tolerance and relative tolerance at the same time. The - // default values of atol and rtol borrowed from numpy.isclose(). For the - // case of nan value, the result will be true. - if (std::isnan(a) && std::isnan(b)) - { - return true; - } - return fabs(a - b) <= (atol + rtol * fabs(b)); -} - -bool compareLogits(ITensor const& groundTruthLogits, ITensor const& outputLogits, float atol = 1e-2, float rtol = 1e-3); - -std::tuple getRequestGivenInputIdxLength( - std::uint64_t requestId, SizeType32 nbGivenInputs, std::vector const& givenInputLengths); - -std::tuple, SizeType32, SizeType32> getGivenInputLengths( - ITensor const& givenInput, SizeType32 padId); - -/// @brief Generates a vector of floating point values summing to 1, that can be used as logits. -/// -/// @tparam TEngine The type of the random engine. -/// @tparam TLogits The type of floating point values. -/// @param vocabSize The vocabulary size, i.e. the size of the vector. -/// @param engine A random engine. -/// @return std::vector A vector of floating point values, summing to 1. -template -std::vector randomLogits(runtime::SizeType32 vocabSize, TEngine* engine) -{ - if constexpr (std::disjunction_v, std::is_same>) - { - // This algorithm ensures the resulting values sum to 1 by: - // 1. Sampling in the interval 0..1 - // 2. Sorting the sampled values and adding a last value equal to 1 - // 3. Calculating the adjacent differences of the sorted values - // Since the values are sorted and the last value is 1, we get that all the differences are positive and must - // sum to 1. It can be proven recursively by seeing that the first value sums to itself, and the n-1 first - // values must sum to the value at n, minus the difference between the n-th and n-1-th values. - // It is also helpful to convince yourself of it with a quick drawing. - auto distribution = std::uniform_real_distribution(0, 1); - std::vector samples(vocabSize); - samples.back() = 1.0; - std::transform(samples.begin(), samples.end() - 1, samples.begin(), - [&](auto const /*i*/) { return distribution(*engine); }); - std::sort(samples.begin(), samples.end() - 1); - std::vector result(vocabSize); - std::adjacent_difference(samples.begin(), samples.end(), result.begin()); - if constexpr (std::is_same_v) - { - return result; - } - - if constexpr (std::is_same_v) - { - std::vector halfResults(vocabSize); - std::transform( - result.begin(), result.end(), halfResults.begin(), [&](auto const f) { return __float2half(f); }); - return halfResults; - } - } - TLLM_THROW("Unsupported logits type."); -} - -std::vector createConsecutiveTokenSequence( - tr::SizeType32 length, tr::SizeType32 vocabSize, tr::TokenIdType firstTokenId); - -/** - * GPU timer for recording the elapsed time across kernel(s) launched in GPU stream - */ -struct GpuTimer -{ - cudaStream_t _stream_id; - cudaEvent_t _start; - cudaEvent_t _stop; - - /// Construct`or - GpuTimer() - : _stream_id(0) - { - TLLM_CUDA_CHECK(cudaEventCreate(&_start)); - TLLM_CUDA_CHECK(cudaEventCreate(&_stop)); - } - - /// Destructor - ~GpuTimer() - { - TLLM_CUDA_CHECK(cudaEventDestroy(_start)); - TLLM_CUDA_CHECK(cudaEventDestroy(_stop)); - } - - /// Start the timer for a given stream (defaults to the default stream) - void start(cudaStream_t stream_id = 0) - { - _stream_id = stream_id; - TLLM_CUDA_CHECK(cudaEventRecord(_start, _stream_id)); - } - - /// Stop the timer - void stop() - { - TLLM_CUDA_CHECK(cudaEventRecord(_stop, _stream_id)); - } - - /// Return the elapsed time (in milliseconds) - float elapsed_millis() - { - float elapsed = 0.0; - TLLM_CUDA_CHECK(cudaEventSynchronize(_stop)); - TLLM_CUDA_CHECK(cudaEventElapsedTime(&elapsed, _start, _stop)); - return elapsed; - } -}; - -} // namespace tensorrt_llm::testing diff --git a/cpp/tests/utils/engines.cpp b/cpp/tests/utils/engines.cpp deleted file mode 100644 index 28f0d9b15935..000000000000 --- a/cpp/tests/utils/engines.cpp +++ /dev/null @@ -1,97 +0,0 @@ -#include "engines.h" - -#include "tensorrt_llm/batch_manager/transformerBuffers.h" -#include "tensorrt_llm/common/logger.h" -#include "tensorrt_llm/runtime/common.h" -#include "tensorrt_llm/runtime/iTensor.h" -#include "tensorrt_llm/runtime/tllmLogger.h" - -#include -#include -#include - -nvinfer1::ITensor& tensorrt_llm::testing::utils::engines::details::addInputIds( - EngineBuildState& buildState, runtime::SizeType32 maxNumTokens) -{ - auto* input_ids = buildState.networkDefinition->addInput(batch_manager::RuntimeBuffers::kInputIdsTensorName, - nvinfer1::DataType::kINT32, runtime::ITensor::makeShape({-1})); - buildState.tensors.push_back(input_ids); - buildState.profile->setDimensions(batch_manager::RuntimeBuffers::kInputIdsTensorName, - nvinfer1::OptProfileSelector::kMAX, runtime::ITensor::makeShape({maxNumTokens})); - buildState.profile->setDimensions(batch_manager::RuntimeBuffers::kInputIdsTensorName, - nvinfer1::OptProfileSelector::kOPT, runtime::ITensor::makeShape({maxNumTokens / 2})); - buildState.profile->setDimensions(batch_manager::RuntimeBuffers::kInputIdsTensorName, - nvinfer1::OptProfileSelector::kMIN, runtime::ITensor::makeShape({1})); - return *input_ids; -} - -nvinfer1::ITensor* tensorrt_llm::testing::utils::engines::details::addLastTokenIds( - EngineBuildState& buildState, runtime::SizeType32 maxBatchSize, runtime::SizeType32 maxBeamWidth) -{ - auto* last_token_ids - = buildState.networkDefinition->addInput(batch_manager::RuntimeBuffers::kLastTokenIdsTensorName, - nvinfer1::DataType::kINT32, runtime::ITensor::makeShape({-1})); - buildState.tensors.push_back(last_token_ids); - buildState.profile->setDimensions(batch_manager::RuntimeBuffers::kLastTokenIdsTensorName, - nvinfer1::OptProfileSelector::kMAX, runtime::ITensor::makeShape({maxBatchSize * maxBeamWidth})); - buildState.profile->setDimensions(batch_manager::RuntimeBuffers::kLastTokenIdsTensorName, - nvinfer1::OptProfileSelector::kOPT, runtime::ITensor::makeShape({maxBatchSize * maxBeamWidth / 2})); - buildState.profile->setDimensions(batch_manager::RuntimeBuffers::kLastTokenIdsTensorName, - nvinfer1::OptProfileSelector::kMIN, runtime::ITensor::makeShape({1})); - return last_token_ids; -} - -nvinfer1::ITensor& tensorrt_llm::testing::utils::engines::details::addKvCacheOffsets(EngineBuildState& buildState, - runtime::SizeType32 numPools, runtime::SizeType32 tokensPerBlock, runtime::SizeType32 maxBatchSize, - runtime::SizeType32 maxNumTokens, runtime::SizeType32 maxBeamWidth) -{ - auto* kvCacheOffsets = buildState.networkDefinition->addInput( - batch_manager::TransformerBuffers::kKvCacheBlockOffsetsTensorName, nvinfer1::DataType::kINT32, - runtime::ITensor::makeShape({numPools, -1, 2, -1})); // [numPools, maxBatch * maxBeamWidth, 2, maxBlocksPerSeq] - buildState.tensors.push_back(kvCacheOffsets); - auto const maxBlocksPerSeq = maxNumTokens / tokensPerBlock; - buildState.profile->setDimensions(batch_manager::TransformerBuffers::kKvCacheBlockOffsetsTensorName, - nvinfer1::OptProfileSelector::kMAX, - runtime::ITensor::makeShape({numPools, maxBatchSize * maxBeamWidth, 2, maxBlocksPerSeq})); - buildState.profile->setDimensions(batch_manager::TransformerBuffers::kKvCacheBlockOffsetsTensorName, - nvinfer1::OptProfileSelector::kOPT, - runtime::ITensor::makeShape({numPools, maxBatchSize * maxBeamWidth / 2, 2, maxBlocksPerSeq / 2})); - buildState.profile->setDimensions(batch_manager::TransformerBuffers::kKvCacheBlockOffsetsTensorName, - nvinfer1::OptProfileSelector::kMIN, runtime::ITensor::makeShape({numPools, 1, 2, 1})); - return *kvCacheOffsets; -} - -tensorrt_llm::testing::utils::engines::details::EngineBuildState -tensorrt_llm::testing::utils::engines::initializeEngineBuild(std::shared_ptr const& logger) -{ - auto* builder = nvinfer1::createInferBuilder(*logger); - auto* profile = builder->createOptimizationProfile(); - auto* network = builder->createNetworkV2( - 1U << static_cast(nvinfer1::NetworkDefinitionCreationFlag::kSTRONGLY_TYPED)); - nvinfer1::IBuilderConfig* config = builder->createBuilderConfig(); - return {builder, network, profile, config}; -} - -nvinfer1::ITensor& tensorrt_llm::testing::utils::engines::details::addSingleOutputLayer( - tensorrt_llm::testing::utils::engines::details::EngineBuildState& buildState, nvinfer1::ILayer* layer) -{ - buildState.layers.push_back(layer); - auto* output = layer->getOutput(0); - buildState.tensors.push_back(output); - TLLM_LOG_INFO("Adding layer %s with output shape %s.", layer->getName(), - tensorrt_llm::runtime::ITensor::toString(output->getDimensions()).c_str()); - - return *output; -} - -tensorrt_llm::common::OptionalRef tensorrt_llm::testing::utils::engines::details::getTensorByName( - tensorrt_llm::testing::utils::engines::details::EngineBuildState& buildState, std::string_view name) -{ - auto result = std::find_if(buildState.tensors.begin(), buildState.tensors.end(), - [name](auto const tensor) { return tensor->getName() == name; }); - if (result == buildState.tensors.end()) - { - return tensorrt_llm::common::OptionalRef{}; - } - return **result; -} diff --git a/cpp/tests/utils/engines.h b/cpp/tests/utils/engines.h deleted file mode 100644 index 54b53e5c6097..000000000000 --- a/cpp/tests/utils/engines.h +++ /dev/null @@ -1,356 +0,0 @@ -#ifndef CA1B91B5_DF64_4CF8_948F_5AFF243A2555 -#define CA1B91B5_DF64_4CF8_948F_5AFF243A2555 - -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/optionalRef.h" -#include "tensorrt_llm/runtime/common.h" -#include "tensorrt_llm/runtime/iBuffer.h" -#include "tensorrt_llm/runtime/iTensor.h" -#include "tensorrt_llm/runtime/tllmLogger.h" -#include -#include -#include -#include -#include -#include -#include -#include - -namespace tensorrt_llm::testing::utils::engines -{ - -namespace details -{ - -struct EngineBuildResource -{ - EngineBuildResource() = default; - virtual ~EngineBuildResource() = default; - EngineBuildResource(EngineBuildResource const& vector) = default; - EngineBuildResource& operator=(EngineBuildResource const& vector) = default; - EngineBuildResource(EngineBuildResource&& vector) noexcept = default; - EngineBuildResource& operator=(EngineBuildResource&& vector) noexcept = default; -}; - -template -struct Vector : public EngineBuildResource -{ - explicit Vector(std::vector values) - : values(std::move(values)){}; - Vector(Vector const& vector) = default; - Vector& operator=(Vector const& vector) = default; - Vector(Vector&& vector) noexcept = default; - Vector& operator=(Vector&& vector) noexcept = default; - ~Vector() override = default; - std::vector values; -}; - -template -struct Array : public EngineBuildResource -{ - explicit Array(std::array values) - : values(std::move(values)){}; - Array(Array const& vector) = default; - Array& operator=(Array const& vector) = default; - Array(Array&& vector) noexcept = default; - Array& operator=(Array&& vector) noexcept = default; - ~Array() override = default; - std::array values; -}; - -struct EngineBuildState -{ - EngineBuildState(nvinfer1::IBuilder* builder, nvinfer1::INetworkDefinition* networkDefinition, - nvinfer1::IOptimizationProfile* profile, nvinfer1::IBuilderConfig* builderConfig) - : builder(builder) - , networkDefinition(networkDefinition) - , profile(profile) - , builderConfig(builderConfig){}; - EngineBuildState(EngineBuildState const& vector) = delete; - EngineBuildState& operator=(EngineBuildState const& vector) = delete; - EngineBuildState(EngineBuildState&& vector) noexcept = default; - EngineBuildState& operator=(EngineBuildState&& vector) noexcept = default; - std::unique_ptr builder; - std::unique_ptr networkDefinition; - nvinfer1::IOptimizationProfile* profile; - - // While building the engine, one might need some data for weights and such. Turns out, TensorRT does not keep a - // copy of those, so if you create them as temporaries and pass them to the TRT APIs, you will get UB. So we need - // some place where we can keep those things. - std::unique_ptr builderConfig; - std::vector> resources; - std::vector tensors; - std::vector layers; - - ~EngineBuildState() - { - // Builder needs to be deleteds last. - networkDefinition.reset(); - builderConfig.reset(); - builder.reset(); - } -}; - -common::OptionalRef getTensorByName(EngineBuildState& buildState, std::string_view name); - -nvinfer1::ITensor& addSingleOutputLayer(EngineBuildState& buildState, nvinfer1::ILayer* layer); - -template -TResource& addResource(EngineBuildState& buildState, TResource resource) -{ - return *dynamic_cast( - buildState.resources.emplace_back(std::make_unique(std::move(resource))).get()); -} - -template -Vector& addSingleConstantVectorResource(EngineBuildState& buildState, TValue value, std::size_t length) -{ - std::vector weights(length); - std::fill(weights.begin(), weights.end(), value); - return addResource(buildState, Vector{weights}); -} - -template -Vector& addConstantVectorResource(EngineBuildState& buildState, std::vector values) -{ - return addResource(buildState, Vector{values}); -} - -template -Array& addConstantScalarResource(EngineBuildState& buildState, TValue value) -{ - return addResource(buildState, Array{{value}}); -} - -nvinfer1::ITensor& addInputIds(EngineBuildState& buildState, runtime::SizeType32 maxNumTokens); -nvinfer1::ITensor* addLastTokenIds( - EngineBuildState& buildState, runtime::SizeType32 maxBatchSize, runtime::SizeType32 maxBeamWidth); -nvinfer1::ITensor& addKvCacheOffsets(EngineBuildState& buildState, runtime::SizeType32 numPools, - runtime::SizeType32 tokensPerBlock, runtime::SizeType32 maxBatchSize, runtime::SizeType32 maxNumTokens, - runtime::SizeType32 maxBeamWidth); - -template -nvinfer1::ITensor& addSingleConstantVector(EngineBuildState& buildState, TValue value, runtime::SizeType32 length) -{ - auto& resourceWeights = addSingleConstantVectorResource(buildState, value, length); - auto const trtDatatype = runtime::TRTDataType::value; - auto* layer = buildState.networkDefinition->addConstant(runtime::ITensor::makeShape({length}), - {trtDatatype, resourceWeights.values.data(), static_cast(length)}); - return addSingleOutputLayer(buildState, layer); -} - -template -nvinfer1::ITensor& addSingleConstantTensor(EngineBuildState& buildState, TValue value, runtime::SizeType32 length) -{ - auto& resourceWeights = addSingleConstantVectorResource(buildState, value, length); - auto const trtDatatype = runtime::TRTDataType::value; - auto* layer = buildState.networkDefinition->addConstant(runtime::ITensor::makeShape({1, length}), - {trtDatatype, resourceWeights.values.data(), static_cast(length)}); - return addSingleOutputLayer(buildState, layer); -} - -template -nvinfer1::ITensor& addConstantVector(EngineBuildState& buildState, std::vector values) -{ - auto& resourceWeights = addConstantVectorResource(buildState, values); - auto const trtDatatype = runtime::TRTDataType::value; - auto const length = static_cast(values.size()); - auto* layer = buildState.networkDefinition->addConstant( - runtime::ITensor::makeShape({length}), {trtDatatype, resourceWeights.values.data(), length}); - return addSingleOutputLayer(buildState, layer); -} - -template -nvinfer1::ITensor& addConstantTensor( - EngineBuildState& buildState, std::vector values, runtime::ITensor::Shape shape) -{ - auto& resourceWeights = addConstantVectorResource(buildState, values); - auto const trtDatatype = runtime::TRTDataType::value; - auto const count = runtime::ITensor::volume(shape); - auto* layer = buildState.networkDefinition->addConstant(shape, {trtDatatype, resourceWeights.values.data(), count}); - return addSingleOutputLayer(buildState, layer); -} - -template -nvinfer1::ITensor& addSingleConstantTensor(EngineBuildState& buildState, TValue value, runtime::ITensor::Shape shape) -{ - auto const count = runtime::ITensor::volume(shape); - auto& resourceWeights = addSingleConstantVectorResource(buildState, value, count); - auto const trtDatatype = runtime::TRTDataType::value; - auto* layer = buildState.networkDefinition->addConstant(shape, {trtDatatype, resourceWeights.values.data(), count}); - return addSingleOutputLayer(buildState, layer); -} - -template -nvinfer1::ITensor& addConstantScalar(EngineBuildState& buildState, TValue value) -{ - auto& resourceWeights = addConstantScalarResource(buildState, value); - auto const trtDatatype = runtime::TRTDataType::value; - auto* layer = buildState.networkDefinition->addConstant( - runtime::ITensor::makeShape({}), {trtDatatype, resourceWeights.values.data(), 1}); - return addSingleOutputLayer(buildState, layer); -} - -template -nvinfer1::ITensor& oneHotEncode( - EngineBuildState& buildState, nvinfer1::ITensor& inputIds, runtime::SizeType32 vocabSize) -{ - auto const trtValueType = runtime::TRTDataType::value; - auto& oneHotValues = addConstantVector(buildState, {0, 1}); - auto& oneHotDepth = addConstantScalar(buildState, vocabSize); - auto* oneHotLayer = buildState.networkDefinition->addOneHot(inputIds, oneHotValues, oneHotDepth, 0); - return addSingleOutputLayer(buildState, oneHotLayer); -} -} // namespace details - -struct TrivialDecoderParameters -{ - TrivialDecoderParameters(runtime::SizeType32 vocabSize, runtime::SizeType32 maxBatchSize, - runtime::SizeType32 maxNumTokens, runtime::SizeType32 tokensPerBlock, runtime::SizeType32 maxBeamWidth, - bool gatherContextLogits) - : vocabSize(vocabSize) - , maxBatchSize(maxBatchSize) - , maxNumTokens(maxNumTokens) - , tokensPerBlock(tokensPerBlock) - , maxBeamWidth(maxBeamWidth) - , gatherContextLogits(gatherContextLogits){}; - runtime::SizeType32 vocabSize; - runtime::SizeType32 maxBatchSize; - runtime::SizeType32 maxNumTokens; - runtime::SizeType32 tokensPerBlock; - runtime::SizeType32 maxBeamWidth; - bool gatherContextLogits; -}; - -details::EngineBuildState initializeEngineBuild(std::shared_ptr const& logger); - -template -std::unique_ptr createTrivialDecoder( - TrivialDecoderParameters parameters, std::shared_ptr const& logger) -{ - auto const trtLogitsType = runtime::TRTDataType::value; - auto buildState = initializeEngineBuild(logger); - auto* builder = buildState.builder.get(); - auto* profile = buildState.profile; - auto* network = buildState.networkDefinition.get(); - auto& inputIds = details::addInputIds(buildState, parameters.maxNumTokens); - auto& kvCacheOffsets = details::addKvCacheOffsets(buildState, 1, parameters.tokensPerBlock, parameters.maxBatchSize, - parameters.maxNumTokens, parameters.maxBeamWidth); - - auto& oneHotLayerOutput = details::oneHotEncode(buildState, inputIds, parameters.vocabSize); - oneHotLayerOutput.setName(batch_manager::RuntimeBuffers::kLogitsTensorName); - network->markOutput(oneHotLayerOutput); - - buildState.builderConfig->addOptimizationProfile(profile); - buildState.builderConfig->setProfilingVerbosity(nvinfer1::ProfilingVerbosity::kDETAILED); - auto* engine = builder->buildSerializedNetwork(*network, *buildState.builderConfig); - return std::unique_ptr(engine); -} - -template -struct ConstantTrivialDecoderParameters -{ - ConstantTrivialDecoderParameters(TrivialDecoderParameters trivialDecoderParameters, std::vector logits) - : trivialDecoderParameters(trivialDecoderParameters) - , logits(logits) - { - auto const sizeTypeVocabSize = static_cast(trivialDecoderParameters.vocabSize); - auto const logitsSize = logits.size(); - TLLM_CHECK_WITH_INFO(static_cast(trivialDecoderParameters.vocabSize) == logits.size(), - "The size of the constant logits (%lu) has to be equal to the vocabulary size (%lu).", logitsSize, - sizeTypeVocabSize); - }; - - TrivialDecoderParameters trivialDecoderParameters; - std::vector logits; -}; - -template -details::EngineBuildState createConstantTrivialDecoderBase( - ConstantTrivialDecoderParameters parameters, std::shared_ptr const& logger) -{ - auto const trtLogitsType = runtime::TRTDataType::value; - auto buildState = initializeEngineBuild(logger); - auto* builder = buildState.builder.get(); - auto* profile = buildState.profile; - auto* network = buildState.networkDefinition.get(); - auto& inputIds = details::addInputIds(buildState, parameters.trivialDecoderParameters.maxNumTokens); - nvinfer1::ITensor* lastTokenIds = nullptr; - if (!parameters.trivialDecoderParameters.gatherContextLogits) - { - lastTokenIds = details::addLastTokenIds(buildState, parameters.trivialDecoderParameters.maxBatchSize, - parameters.trivialDecoderParameters.maxBeamWidth); - } - auto& kvCacheOffsets = details::addKvCacheOffsets(buildState, 1, parameters.trivialDecoderParameters.tokensPerBlock, - parameters.trivialDecoderParameters.maxBatchSize, parameters.trivialDecoderParameters.maxNumTokens, - parameters.trivialDecoderParameters.maxBeamWidth); - - auto const vocabSize = static_cast(parameters.logits.size()); - - auto& constantLogitsPerToken = details::addConstantTensor( - buildState, parameters.logits, runtime::ITensor::makeShape({vocabSize, 1})); - auto& oneHotLayerOutput - = details::oneHotEncode(buildState, inputIds, parameters.trivialDecoderParameters.vocabSize); - auto& ones = details::addSingleConstantTensor(buildState, 1, runtime::ITensor::makeShape({1, vocabSize})); - auto* intermediateLayer1 = network->addMatrixMultiply( - ones, nvinfer1::MatrixOperation::kNONE, oneHotLayerOutput, nvinfer1::MatrixOperation::kNONE); - auto* intermediateLayer1Output = intermediateLayer1->getOutput(0); - - nvinfer1::ITensor* gatherLayerOutput = nullptr; - if (!parameters.trivialDecoderParameters.gatherContextLogits) - { - auto& one = details::addSingleConstantTensor(buildState, 1, runtime::ITensor::makeShape({1})); - auto* lastTokenIdsMinus1Layer - = network->addElementWise(*lastTokenIds, one, nvinfer1::ElementWiseOperation::kSUB); - auto* gatherLayer = network->addGather(*intermediateLayer1Output, *lastTokenIdsMinus1Layer->getOutput(0), 1); - gatherLayerOutput = gatherLayer->getOutput(0); - } - else - { - gatherLayerOutput = intermediateLayer1Output; - } - - auto* constLogitsLayer = network->addMatrixMultiply(*gatherLayerOutput, nvinfer1::MatrixOperation::kTRANSPOSE, - constantLogitsPerToken, nvinfer1::MatrixOperation::kTRANSPOSE); - auto* outputLogits = constLogitsLayer->getOutput(0); - network->markOutput(*outputLogits); - outputLogits->setName(batch_manager::RuntimeBuffers::kLogitsTensorName); - buildState.tensors.push_back(outputLogits); - return buildState; -} - -template -std::unique_ptr createConstantTrivialDecoder( - ConstantTrivialDecoderParameters parameters, std::shared_ptr const& logger) -{ - auto buildState = createConstantTrivialDecoderBase(parameters, logger); - buildState.builderConfig->addOptimizationProfile(buildState.profile); - buildState.builderConfig->setProfilingVerbosity(nvinfer1::ProfilingVerbosity::kDETAILED); - auto* engine = buildState.builder->buildSerializedNetwork(*buildState.networkDefinition, *buildState.builderConfig); - return std::unique_ptr(engine); -} - -template -std::unique_ptr createConstantTrivialDecoderWithTopKLogits( - ConstantTrivialDecoderParameters parameters, runtime::SizeType32 numTopLogits, std::string_view outputName, - std::shared_ptr const& logger) -{ - auto buildState = createConstantTrivialDecoderBase(parameters, logger); - auto logits = details::getTensorByName(buildState, batch_manager::RuntimeBuffers::kLogitsTensorName); - TLLM_CHECK_WITH_INFO(static_cast(logits), - "You can only add topk logits on top of a network which contains a tensor named %s", - batch_manager::RuntimeBuffers::kLogitsTensorName); - auto* topKLayer = buildState.networkDefinition->addTopK( - logits.value(), nvinfer1::TopKOperation::kMAX, numTopLogits, 1UL << 1UL); - auto* topKLayerOutput = topKLayer->getOutput(0); - topKLayerOutput->setName(outputName.data()); - buildState.networkDefinition->markOutput(*topKLayerOutput); - auto* profile = buildState.profile; - buildState.builderConfig->addOptimizationProfile(profile); - buildState.builderConfig->setProfilingVerbosity(nvinfer1::ProfilingVerbosity::kDETAILED); - auto* engine = buildState.builder->buildSerializedNetwork(*buildState.networkDefinition, *buildState.builderConfig); - return std::unique_ptr(engine); -} -} // namespace tensorrt_llm::testing::utils::engines - -#endif /* CA1B91B5_DF64_4CF8_948F_5AFF243A2555 */ diff --git a/cpp/tests/utils/executorUtils.cpp b/cpp/tests/utils/executorUtils.cpp deleted file mode 100644 index 6dbd6223ef19..000000000000 --- a/cpp/tests/utils/executorUtils.cpp +++ /dev/null @@ -1,55 +0,0 @@ -#include "executorUtils.h" -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/logger.h" -#include -#include - -std::unordered_map> -tensorrt_llm::testing::runThroughRequests(executor::Executor& executor, std::vector const& requests, - std::chrono::duration timeout) -{ - std::unordered_map> accumulatedResponses; - auto responseReadFuture = std::async(std::launch::async, - [&]() -> std::optional - { - auto remainingRequests = requests.size(); - try - { - while (remainingRequests > 0) - { - auto const responses = executor.awaitResponses(); - for (auto const& response : responses) - { - auto const requestId = response.getRequestId(); - if (response.hasError()) - { - TLLM_LOG_ERROR("Error response received for request: %lu", requestId); - TLLM_THROW(response.getErrorMsg()); - } - auto const isFinal = response.hasError() || response.getResult().isFinal; - accumulatedResponses[requestId].emplace_back(response); - if (isFinal) - { - TLLM_LOG_DEBUG("Final response received for request: %lu", requestId); - --remainingRequests; - } - } - } - - return std::nullopt; - } - catch (std::exception const& e) - { - TLLM_LOG_EXCEPTION(e); - return e; - } - }); - auto const requestIds = executor.enqueueRequests(requests); - responseReadFuture.wait_for(timeout); - auto const readResult = responseReadFuture.get(); - if (readResult.has_value()) - { - throw std::exception(readResult.value()); - } - return accumulatedResponses; -} diff --git a/cpp/tests/utils/executorUtils.h b/cpp/tests/utils/executorUtils.h deleted file mode 100644 index 97f154ca275d..000000000000 --- a/cpp/tests/utils/executorUtils.h +++ /dev/null @@ -1,18 +0,0 @@ -#ifndef A073F2DA_315E_434B_B811_D420F0A59DF3 -#define A073F2DA_315E_434B_B811_D420F0A59DF3 - -#include "tensorrt_llm/batch_manager/common.h" -#include "tensorrt_llm/executor/executor.h" -#include -#include - -namespace tensorrt_llm::testing -{ - -std::unordered_map> runThroughRequests( - executor::Executor& executor, std::vector const& requests, - std::chrono::duration timeout); - -} // namespace tensorrt_llm::testing - -#endif /* A073F2DA_315E_434B_B811_D420F0A59DF3 */ diff --git a/docker/Dockerfile.multi b/docker/Dockerfile.multi index e72a9ad30a76..b25a24086af2 100644 --- a/docker/Dockerfile.multi +++ b/docker/Dockerfile.multi @@ -119,7 +119,7 @@ COPY .gitmodules setup.py requirements.txt requirements-dev.txt constraints.txt ENV CCACHE_DIR=/root/.cache/ccache # Build the TRT-LLM wheel ARG GITHUB_MIRROR="" -ARG BUILD_WHEEL_ARGS="--clean --benchmarks" +ARG BUILD_WHEEL_ARGS="--clean" ARG BUILD_WHEEL_SCRIPT="scripts/build_wheel.py" RUN --mount=type=cache,target=/root/.cache/pip --mount=type=cache,target=${CCACHE_DIR} \ GITHUB_MIRROR=$GITHUB_MIRROR python3 ${BUILD_WHEEL_SCRIPT} ${BUILD_WHEEL_ARGS} @@ -138,7 +138,6 @@ RUN --mount=type=bind,source=README.md,target=/mnt/ctx/README.md \ --mount=type=bind,source=examples,target=/mnt/ctx/examples \ --mount=type=bind,from=wheel,source=/src/tensorrt_llm/build,target=/mnt/wheel \ --mount=type=bind,from=wheel,source=/src/tensorrt_llm/benchmarks,target=/mnt/benchmarks \ - --mount=type=bind,from=wheel,source=/src/tensorrt_llm/cpp/build/benchmarks,target=/mnt/cpp_benchmarks \ # Copy build context files cp /mnt/ctx/README.md ./ && \ cp -r /mnt/ctx/docs ./docs && \ @@ -148,24 +147,12 @@ RUN --mount=type=bind,source=README.md,target=/mnt/ctx/README.md \ # Copy wheel stage outputs cp /mnt/wheel/tensorrt_llm*.whl ./ && \ cp -r /mnt/benchmarks ./benchmarks && \ - mkdir -p benchmarks/cpp && \ - cp /mnt/cpp_benchmarks/bertBenchmark \ - /mnt/cpp_benchmarks/gptManagerBenchmark \ - /mnt/cpp_benchmarks/disaggServerBenchmark \ - benchmarks/cpp/ && \ - rm -v \ - benchmarks/cpp/bertBenchmark.cpp \ - benchmarks/cpp/gptManagerBenchmark.cpp \ - benchmarks/cpp/disaggServerBenchmark.cpp \ - benchmarks/cpp/CMakeLists.txt && \ - # Create symlinks to installed package binaries and libraries - ln -sv $(python3 -c 'import site; print(f"{site.getsitepackages()[0]}/tensorrt_llm/bin")') bin && \ - test -f bin/executorWorker && \ + # Create a symlink to installed package libraries ln -sv $(python3 -c 'import site; print(f"{site.getsitepackages()[0]}/tensorrt_llm/libs")') lib && \ - test -f lib/libnvinfer_plugin_tensorrt_llm.so && \ + test -f lib/libtensorrt_llm.so && \ echo "/app/tensorrt_llm/lib" > /etc/ld.so.conf.d/tensorrt_llm.conf && \ ldconfig && \ - ! ( ldd -v bin/executorWorker | grep tensorrt_llm | grep -q "not found" ) && \ + ! ( ldd -v lib/libth_common.so | grep tensorrt_llm | grep -q "not found" ) && \ # Clean up caches and CVE workarounds rm -rf /root/.cache/uv/archive-v0 && \ # WAR against https://github.com/advisories/GHSA-58pv-8j8x-9vj2 diff --git a/docs/source/blogs/Best_perf_practice_on_DeepSeek-R1_in_TensorRT-LLM.md b/docs/source/blogs/Best_perf_practice_on_DeepSeek-R1_in_TensorRT-LLM.md index 22d4688c503d..91d89d70cacf 100644 --- a/docs/source/blogs/Best_perf_practice_on_DeepSeek-R1_in_TensorRT-LLM.md +++ b/docs/source/blogs/Best_perf_practice_on_DeepSeek-R1_in_TensorRT-LLM.md @@ -92,7 +92,7 @@ Here we set `LOCAL_USER=1` argument to set up the local user instead of root acc Here we compile the source inside the container: ``` bash -python3 ./scripts/build_wheel.py --trt_root /usr/local/tensorrt --benchmarks --cuda_architectures "90-real;100-real" --python_bindings --clean +python3 ./scripts/build_wheel.py --cuda_architectures "90-real;100-real" --clean ``` You can set the cuda_architectures to "100-real" if targeting Blackwell only, and "90-real" to target Hopper only to save some build time. diff --git a/docs/source/commands/trtllm-bench.rst b/docs/source/commands/trtllm-bench.rst index fee60a9ab70c..a55899a961d1 100644 --- a/docs/source/commands/trtllm-bench.rst +++ b/docs/source/commands/trtllm-bench.rst @@ -23,7 +23,7 @@ Dataset preparation prepare_dataset.py ^^^^^^^^^^^^^^^^^^ -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 is designed to work with the `prepare_dataset.py `_ script, which generates benchmark datasets in the required format. The prepare_dataset script supports: **Dataset Types:** diff --git a/docs/source/developer-guide/overview.md b/docs/source/developer-guide/overview.md index af7f44f139cf..d8fe31631612 100644 --- a/docs/source/developer-guide/overview.md +++ b/docs/source/developer-guide/overview.md @@ -101,7 +101,6 @@ Module names longer than 8 characters are abbreviated to fit the fixed-width tag | `deep_ep` | `deep_ep ` | | `deep_gemm` | `deepgemm` | | `executor` | `executor` | -| `executor_worker` | `exec_wkr` | | `flash_mla` | `flashmla` | | `kernels` | `kernels ` | | `layers` | `layers ` | diff --git a/docs/source/developer-guide/perf-benchmarking.md b/docs/source/developer-guide/perf-benchmarking.md index 4eb04eef4d13..5903c12e32ea 100644 --- a/docs/source/developer-guide/perf-benchmarking.md +++ b/docs/source/developer-guide/perf-benchmarking.md @@ -171,7 +171,7 @@ can simply read a line and assume a complete entry. When creating a dataset, be JSON entry is on every line. ``` -In order to prepare a synthetic dataset, you can use the provided script in the `benchmarks/cpp` +In order to prepare a synthetic dataset, you can use the provided script in the `benchmarks` directory. For example, to generate a synthetic dataset of 1000 requests with a uniform ISL/OSL of 128/128 for [meta-llama/Llama-3.1-8B](https://huggingface.co/meta-llama/Llama-3.1-8B), run: diff --git a/docs/source/developer-guide/perf-overview.md b/docs/source/developer-guide/perf-overview.md index 223ac5e8e92b..dbf999d4ca03 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/cpp/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 | `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` | | [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 @@ -288,11 +288,11 @@ Testing was performed using the PyTorch backend - this workflow does not require ### Preparing a Dataset -In order to prepare a dataset, you can use the provided [script](source:benchmarks/cpp/prepare_dataset.py). +In order to prepare a dataset, you can use the provided [script](source:benchmarks/prepare_dataset.py). To generate a synthetic dataset, run the following command: ```shell -python benchmarks/cpp/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 +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 ``` 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/legacy/performance/perf-benchmarking.md b/docs/source/legacy/performance/perf-benchmarking.md index 4fc460596f22..8d1567aee020 100644 --- a/docs/source/legacy/performance/perf-benchmarking.md +++ b/docs/source/legacy/performance/perf-benchmarking.md @@ -202,7 +202,7 @@ can simply read a line and assume a complete entry. When creating a dataset, be JSON entry is on every line. ``` -In order to prepare a synthetic dataset, you can use the provided script in the `benchmarks/cpp` +In order to prepare a synthetic dataset, you can use the provided script in the `benchmarks` directory. For example, to generate a synthetic dataset of 1000 requests with a uniform ISL/OSL of 128/128 for [meta-llama/Llama-3.1-8B](https://huggingface.co/meta-llama/Llama-3.1-8B), run: @@ -481,7 +481,7 @@ The PyTorch workflow supports benchmarking with LoRA (Low-Rank Adaptation) adapt Use `prepare_dataset.py` with LoRA-specific options to generate requests with LoRA metadata: ```shell -python3 benchmarks/cpp/prepare_dataset.py \ +python3 benchmarks/prepare_dataset.py \ --stdout \ --rand-task-id 0 1 \ --tokenizer /path/to/tokenizer \ @@ -555,7 +555,7 @@ To benchmark multi-modal models with PyTorch workflow, you can follow the simila First, prepare the dataset: ``` -python ./benchmarks/cpp/prepare_dataset.py \ +python ./benchmarks/prepare_dataset.py \ --tokenizer Qwen/Qwen2-VL-2B-Instruct \ --stdout \ dataset \ @@ -846,7 +846,7 @@ The following table summarizes the commands needed for running benchmarks: | Scenario | Phase | Command | | - | - | - | -| Dataset | Preparation | `python benchmarks/cpp/prepare_dataset.py --stdout --tokenizer $HF_MODEL token-norm-dist --input-mean $ISL --output-mean $OSL --input-stdev 0 --output-stdev 0 --num-requests $NUM_REQUESTS > $DATASET_PATH` | +| Dataset | Preparation | `python benchmarks/prepare_dataset.py --stdout --tokenizer $HF_MODEL token-norm-dist --input-mean $ISL --output-mean $OSL --input-stdev 0 --output-stdev 0 --num-requests $NUM_REQUESTS > $DATASET_PATH` | | Throughput | Build | `trtllm-bench --model $HF_MODEL build --dataset $DATASET_PATH` | | Throughput | Benchmark | `trtllm-bench --model $HF_MODEL throughput --dataset $DATASET_PATH --engine_dir $ENGINE_DIR` | | Latency | Build | See [section about building low latency engines](#low-latency-tensorrt-llm-engine-for-llama-3-70b) | diff --git a/docs/source/legacy/performance/performance-tuning-guide/benchmarking-default-performance.md b/docs/source/legacy/performance/performance-tuning-guide/benchmarking-default-performance.md index 17cb9aef45ba..7277b0afa5b0 100644 --- a/docs/source/legacy/performance/performance-tuning-guide/benchmarking-default-performance.md +++ b/docs/source/legacy/performance/performance-tuning-guide/benchmarking-default-performance.md @@ -86,7 +86,7 @@ The README in the examples folder for supported models walks through building en `trtllm-bench` expects to be passed in a dataset of requests to run through the model. This guide creates a dummy dataset of 1000 requests with every request having input and output sequence length of 2048. TensorRT-LLM provides the `prepare_dataset.py` script to produce the dataset. To use it clone the TensorRT-LLM Repo and run the following command: -`python benchmarks/cpp/prepare_dataset.py --stdout --tokenizer /path/to/hf/Llama-3.3-70B-Instruct/ token-norm-dist --input-mean 2048 --output-mean 2048 --input-stdev 0 --output-stdev 0 --num-requests 1000 > synthetic_2048_2048.txt` +`python benchmarks/prepare_dataset.py --stdout --tokenizer /path/to/hf/Llama-3.3-70B-Instruct/ token-norm-dist --input-mean 2048 --output-mean 2048 --input-stdev 0 --output-stdev 0 --num-requests 1000 > synthetic_2048_2048.txt` `trtllm-bench` can also take in real data, see [`trtllm-bench` documentation](../perf-benchmarking.md) for more details on the required format. diff --git a/examples/auto_deploy/paragraf/create_standalone_package.py b/examples/auto_deploy/paragraf/create_standalone_package.py index 9b48c23692f2..45f12c85dc57 100644 --- a/examples/auto_deploy/paragraf/create_standalone_package.py +++ b/examples/auto_deploy/paragraf/create_standalone_package.py @@ -557,9 +557,9 @@ def insert_optional_trtllm_guard() -> None: 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/cpp directory. + # but it does not ship TensorRT-LLM's source-tree benchmarks/ directory. content = content.replace( - ' script_dir = Path(root_dir, "benchmarks", "cpp")\n', + ' script_dir = Path(root_dir, "benchmarks")\n', " script_dir = Path(temp_dir)\n", ) diff --git a/examples/bindings/executor/README.md b/examples/bindings/executor/README.md deleted file mode 100644 index df44568fd4f1..000000000000 --- a/examples/bindings/executor/README.md +++ /dev/null @@ -1,76 +0,0 @@ -# Python Bindings Example - -This example shows how to use the python bindings interface to generate tokens -using a TensorRT engine. - -## Setup - -Build a TensorRT engine for one of the supported TensorRT LLM model following -instructions in the corresponding `examples` folder. - -## Usage - -### Basic example - -Run `example_basic.py`, passing in the directory where the TensorRT engine was generated. For example: - -``` -cd examples/bindings -python3 example_basic.py --model_path=../llama/tmp/7B/trt_engines/fp16/1-gpu/ -``` - -### Debug example - -This example shows how you can define which engine IO tensors should be kept or dumped to numpy files. -Run `example_debug.py`, passing in the directory where the TensorRT engine was generated. For example: - -``` -cd examples/bindings -python3 example_debug.py --model_path=../llama/tmp/7B/trt_engines/fp16/1-gpu/ -``` - -### Advanced example - -This example shows how you can use the python bindings to generate tokens for a larger number of requests concurrently and demonstrate how tokens can be returned in a streaming fashion. - -The full list of supported input parameters can be obtained with: -``` -pytho3 example_advanced.py -h -``` - -For example, assuming a CSV file named `input_tokens.csv` exist which contains the following input tokens: -``` -1, 2, 3, 4, 5, 6 -1, 2, 3, 4 -1, 2, 3, 4, 5, 6, 7, 8, 9, 10 -``` -one can generate output tokens for those 3 prompts with: -``` -python3 example_advanced.py --model_path --input_tokens_csv_file input_tokens.csv -``` -Upon successful completion, the output tokens will be written to file `output_tokens.csv`. - -### Multi-GPU Example - -To run the two examples for models requiring more than one gpu, you can run the example with MPI. - -For example, the basic example can be run as follows: -``` -mpirun -n 4 --allow-run-as-root python3 example_basic.py --model_path=../llama/tmp/7B/trt_engines/fp16/4gpu_tp4_pp1/ -``` - -The advanced example can also be run using the ORCHESTRATOR mode, where the additional processes needed for multi-GPU runs will automatically be spawned. -This can be done by running: -``` -python3 example_advanced.py --model_path=../llama/tmp/7B/trt_engines/fp16/4gpu_tp4_pp1/ --use_orchestrator_mode -``` - -### Logits post processor example - -This example shows how to generate JSON structured output using LogitsPostProcessor API. - -``` -python3 example_logits_processor.py -t -e --batch_size 8 -``` - -LogitsPostProcessorBatched, which fuses logits processing for all samples in a batch into a single callback, is enabled by `--lpp_batched` diff --git a/examples/bindings/executor/example_advanced.py b/examples/bindings/executor/example_advanced.py deleted file mode 100644 index 25f063a855e2..000000000000 --- a/examples/bindings/executor/example_advanced.py +++ /dev/null @@ -1,166 +0,0 @@ -import argparse -import csv -import datetime -from pathlib import Path - -import tensorrt_llm - -trtllm_package_dir = Path(tensorrt_llm.__file__).parent -executor_worker_path = trtllm_package_dir / 'bin' / 'executorWorker' - -import tensorrt_llm.bindings.executor as trtllm - - -# Read input tokens from csv file -def read_input_tokens(input_tokens_csv_file: str) -> list[int]: - - input_tokens = [] - with open(input_tokens_csv_file, mode='r') as file: - csvFile = csv.reader(file) - for lines in csvFile: - input_tokens.append([int(item) for item in lines]) - return input_tokens - - -# Prepare and enqueue the requests -def enqueue_requests(args: argparse.Namespace, - executor: trtllm.Executor) -> None: - - output_config = trtllm.OutputConfig() - output_config.exclude_input_from_output = args.exclude_input_from_output - sampling_config = trtllm.SamplingConfig(args.beam_width) - input_tokens = read_input_tokens(args.input_tokens_csv_file) - - request_ids = [] - for tokens in input_tokens: - req = trtllm.Request(input_token_ids=tokens, - max_tokens=args.max_tokens, - streaming=args.streaming, - sampling_config=sampling_config, - output_config=output_config) - req_id = executor.enqueue_request(req) - request_ids.append(req_id) - - return request_ids - - -# Wait for responses and store output tokens -def wait_for_responses(args: argparse.Namespace, request_ids: list[int], - executor: trtllm.Executor) -> dict[dict[list[int]]]: - - output_tokens = { - req_id: { - beam: [] - for beam in range(args.beam_width) - } - for req_id in request_ids - } - num_finished = 0 - iter = 0 - while (num_finished < len(request_ids) and iter < args.timeout_ms): - responses = executor.await_responses( - datetime.timedelta(milliseconds=args.timeout_ms)) - for response in responses: - req_id = response.request_id - if not response.has_error(): - result = response.result - num_finished += 1 if result.is_final else 0 - for beam, outTokens in enumerate(result.output_token_ids): - output_tokens[req_id][beam].extend(outTokens) - else: - raise RuntimeError( - str(req_id) + " encountered error:" + response.error_msg) - - return output_tokens - - -# Write the output tokens to file -def write_output_tokens(output_tokens_csv_file: str, request_ids: list[int], - output_tokens: dict[dict[list[int]]], - beam_width: int) -> None: - - with open(output_tokens_csv_file, 'w') as csvfile: - - writer = csv.writer(csvfile) - for req_id in request_ids: - out_tokens = output_tokens[req_id] - for beam in range(args.beam_width): - beam_tokens = out_tokens[beam] - writer.writerow(beam_tokens) - - print("Output tokens written to:", output_tokens_csv_file) - - -if __name__ == "__main__": - parser = argparse.ArgumentParser(description="Executor Bindings Example") - parser.add_argument("--model_path", - type=str, - required=True, - help="Directory containing model engine") - parser.add_argument("--input_tokens_csv_file", - type=str, - required=True, - help="CSV file containing the input tokens") - parser.add_argument("--output_tokens_csv_file", - type=str, - required=False, - default="output_tokens.csv", - help="CSV file where to write output tokens") - parser.add_argument("--beam_width", - type=int, - required=False, - default=1, - help="The beam width") - parser.add_argument("--streaming", - default=False, - action="store_true", - help="Operate in streaming mode") - - parser.add_argument("--use_orchestrator_mode", - default=False, - action="store_true", - help="Operate in orchestrator mode for multi-GPU runs") - - parser.add_argument( - "--exclude_input_from_output", - default=False, - action="store_true", - help= - "Exclude input token when writing output tokens. Only has effect for streaming=False since in streaming mode, input tokens are never included in output." - ) - parser.add_argument("--max_tokens", - type=int, - required=False, - default=10, - help="The max number of tokens to be generated") - parser.add_argument( - "--timeout_ms", - type=int, - required=False, - default=10000, - help="The maximum time to wait for all responses, in milliseconds") - - args = parser.parse_args() - executor_config = trtllm.ExecutorConfig(args.beam_width) - - if args.use_orchestrator_mode: - orchestrator_config = trtllm.OrchestratorConfig( - True, str(executor_worker_path)) - executor_config.parallel_config = trtllm.ParallelConfig( - trtllm.CommunicationType.MPI, trtllm.CommunicationMode.ORCHESTRATOR, - None, None, orchestrator_config) - - # Create the executor. - executor = trtllm.Executor(args.model_path, trtllm.ModelType.DECODER_ONLY, - executor_config) - - if executor.can_enqueue_requests(): - # Enqueue the requests - request_ids = enqueue_requests(args, executor) - - # Wait for the responses - output_tokens = wait_for_responses(args, request_ids, executor) - - # Write the output tokens - write_output_tokens(args.output_tokens_csv_file, request_ids, - output_tokens, args.beam_width) diff --git a/examples/bindings/executor/example_basic.py b/examples/bindings/executor/example_basic.py deleted file mode 100644 index 3c71bde594e1..000000000000 --- a/examples/bindings/executor/example_basic.py +++ /dev/null @@ -1,34 +0,0 @@ -import argparse - -import tensorrt_llm.bindings.executor as trtllm - -# This example hows to use the python bindings to create an executor, enqueue a -# request, and get the generated tokens. - -# First, follow the steps in README.md to generate the engines. - -if __name__ == "__main__": - parser = argparse.ArgumentParser(description="Executor Bindings Example") - parser.add_argument("--model_path", - type=str, - required=True, - help="Directory containing model engine") - args = parser.parse_args() - - # Create the executor. - executor = trtllm.Executor(args.model_path, trtllm.ModelType.DECODER_ONLY, - trtllm.ExecutorConfig(1)) - - if executor.can_enqueue_requests(): - # Create the request. - request = trtllm.Request(input_token_ids=[1, 2, 3, 4], max_tokens=10) - - # Enqueue the request. - request_id = executor.enqueue_request(request) - - # Wait for the new tokens. - responses = executor.await_responses(request_id) - output_tokens = responses[0].result.output_token_ids - - # Print tokens. - print(output_tokens) diff --git a/examples/bindings/executor/example_debug.py b/examples/bindings/executor/example_debug.py deleted file mode 100644 index f7c0669b1254..000000000000 --- a/examples/bindings/executor/example_debug.py +++ /dev/null @@ -1,67 +0,0 @@ -import argparse -import pathlib as pl - -import numpy as np - -import tensorrt_llm.bindings.executor as trtllm - -# This example hows to use the python bindings to create an executor, enqueue a -# request, and get the generated tokens. - -# First, follow the steps in README.md to generate the engines. - -if __name__ == "__main__": - parser = argparse.ArgumentParser(description="Executor Bindings Example") - parser.add_argument("--model_path", - type=str, - required=True, - help="Directory containing model engine") - parser.add_argument("--dump_tensors", - action="store_true", - help="Dump debug tensors to files") - args = parser.parse_args() - - max_tokens = 2 - - # Select which tensors should be kept or dumped - debug_config = trtllm.DebugConfig( - debug_tensor_names=["sequence_length"], - debug_tensors_max_iterations=0 if args.dump_tensors else max_tokens) - - # Create the executor. - executor = trtllm.Executor( - args.model_path, trtllm.ModelType.DECODER_ONLY, - trtllm.ExecutorConfig(1, debug_config=debug_config)) - - if executor.can_enqueue_requests(): - # Create the request. - request = trtllm.Request(input_token_ids=[1, 2, 3, 4], - max_tokens=max_tokens) - - # Enqueue the request. - request_id = executor.enqueue_request(request) - - # Wait for the new tokens. - responses = executor.await_responses(request_id) - output_tokens = responses[0].result.output_token_ids - - # Print tokens. - print(output_tokens) - - if args.dump_tensors: - print("debug tensors from files:") - debug_dir = pl.Path("/tmp/tllm_debug/PP_1/TP_1") - if debug_dir.is_dir(): - for iter_dir in [x for x in debug_dir.iterdir() if x.is_dir()]: - print(iter_dir.name) - for file in [x for x in iter_dir.iterdir() if x.is_file()]: - print(file.name, np.load(file)) - else: - print("debug dir not found") - else: - print("debug tensors from queue:") - debug_tensors = executor.get_latest_debug_tensors() - for debug_iter in debug_tensors: - print(f"iteration {debug_iter.iter}") - for [name, tensor] in debug_iter.debug_tensors.items(): - print(name, tensor) diff --git a/examples/bindings/executor/example_logits_processor.py b/examples/bindings/executor/example_logits_processor.py deleted file mode 100644 index 6cb1a751a6da..000000000000 --- a/examples/bindings/executor/example_logits_processor.py +++ /dev/null @@ -1,212 +0,0 @@ -import argparse -import datetime -import typing as _tp - -import torch as _tor -from lmformatenforcer import (JsonSchemaParser, TokenEnforcer, - TokenEnforcerTokenizerData) -from pydantic import BaseModel -from transformers import AutoTokenizer - -import tensorrt_llm.bindings.executor as trtllm - - -def _build_regular_tokens_list( - tokenizer) -> _tp.List[_tp.Tuple[int, str, bool]]: - token_0 = [tokenizer.encode("0")[-1]] - regular_tokens = [] - vocab_size = tokenizer.vocab_size - for token_idx in range(vocab_size): - if token_idx in tokenizer.all_special_ids: - continue - # We prepend token 0 and skip the first letter of the result to get a space if the token is a start word. - tensor_after_0 = _tor.tensor(token_0 + [token_idx], dtype=_tor.long) - decoded_after_0 = tokenizer.decode(tensor_after_0)[1:] - decoded_regular = tokenizer.decode(token_0) - is_word_start_token = len(decoded_after_0) > len(decoded_regular) - regular_tokens.append((token_idx, decoded_after_0, is_word_start_token)) - return regular_tokens - - -def build_token_enforcer(tokenizer, character_level_parser): - """ - Build logits processor for feeding it into generate function (use_py_session should be True) - """ - regular_tokens = _build_regular_tokens_list(tokenizer) - - def _decode(tokens: _tp.List[int]) -> str: - tensor = _tor.tensor(tokens, dtype=_tor.long) - return tokenizer.decode(tensor) - - tokenizer_data = TokenEnforcerTokenizerData(regular_tokens, _decode, - tokenizer.eos_token_id) - return TokenEnforcer(tokenizer_data, character_level_parser) - - -# Prepare and enqueue the requests -def enqueue_requests(args: argparse.Namespace, - executor: trtllm.Executor) -> None: - - sampling_config = trtllm.SamplingConfig(args.beam_width) - - request_ids = [] - for iter_id in range(args.batch_size): - # Create the request. - request = trtllm.Request(input_token_ids=prompt, - max_tokens=25, - end_id=tokenizer.eos_token_id, - sampling_config=sampling_config, - client_id=iter_id % 2) - request.logits_post_processor_name = request.BATCHED_POST_PROCESSOR_NAME if args.lpp_batched else "my_logits_pp" - - # Enqueue the request. - req_id = executor.enqueue_request(request) - request_ids.append(req_id) - - return request_ids - - -# Wait for responses and store output tokens -def wait_for_responses(args: argparse.Namespace, request_ids: list[int], - executor: trtllm.Executor) -> dict[dict[list[int]]]: - - output_tokens = { - req_id: { - beam: [] - for beam in range(args.beam_width) - } - for req_id in request_ids - } - num_finished = 0 - iter = 0 - while (num_finished < len(request_ids) and iter < args.timeout_ms): - responses = executor.await_responses( - datetime.timedelta(milliseconds=args.timeout_ms)) - for response in responses: - req_id = response.request_id - if not response.has_error(): - result = response.result - num_finished += 1 if result.is_final else 0 - for beam, outTokens in enumerate(result.output_token_ids): - output_tokens[req_id][beam].extend(outTokens) - else: - raise RuntimeError( - str(req_id) + " encountered error:" + response.error_msg) - - return output_tokens - - -if __name__ == "__main__": - parser = argparse.ArgumentParser(description="Executor Bindings Example") - parser.add_argument("--tokenizer_path", - "-t", - type=str, - required=True, - help="Directory containing model tokenizer") - parser.add_argument("--engine_path", - "-e", - type=str, - required=True, - help="Directory containing model engine") - parser.add_argument("--beam_width", - type=int, - required=False, - default=1, - help="The beam width") - parser.add_argument("--batch_size", - type=int, - required=False, - default=1, - help="The batch size") - parser.add_argument( - "--timeout_ms", - type=int, - required=False, - default=10000, - help="The maximum time to wait for all responses, in milliseconds") - parser.add_argument("--lpp_batched", - action="store_true", - default=False, - help="Enable batched logits post processor") - - args = parser.parse_args() - - tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_path) - - class AnswerFormat(BaseModel): - last_name: str - year_of_birth: int - - parser = JsonSchemaParser(AnswerFormat.model_json_schema()) - token_enforcer = build_token_enforcer(tokenizer, parser) - - def get_allowed_tokens(ids, client_id): - if client_id is None or client_id == 0: return [42] - - def _trim(ids): - return [x for x in ids if x != tokenizer.eos_token_id] - - allowed = token_enforcer.get_allowed_tokens(_trim(ids[0])) - return allowed - - def logits_post_processor(req_id: int, logits: _tor.Tensor, - ids: _tp.List[_tp.List[int]], stream_ptr: int, - client_id: _tp.Optional[int]): - mask = _tor.full_like(logits, fill_value=float("-inf"), device="cpu") - allowed = get_allowed_tokens(ids, client_id) - mask[:, :, allowed] = 0 - - with _tor.cuda.stream(_tor.cuda.ExternalStream(stream_ptr)): - mask = mask.to(logits.device, non_blocking=True) - logits += mask - - def logits_post_processor_batched( - req_ids_batch: _tp.List[int], logits_batch: _tp.List[_tor.Tensor], - ids_batch: _tp.List[_tp.List[_tp.List[int]]], stream_ptr, - client_ids_batch: _tp.List[_tp.Optional[int]]): - masks = [] - for req_id, logits, ids, client_id in zip(req_ids_batch, logits_batch, - ids_batch, client_ids_batch): - del req_id - mask = _tor.full_like(logits, - fill_value=float("-inf"), - device="cpu") - allowed = get_allowed_tokens(ids, client_id) - mask[:, :, allowed] = 0 - masks.append(mask) - - with _tor.cuda.stream(_tor.cuda.ExternalStream(stream_ptr)): - for logits, mask in zip(logits_batch, masks): - logits += mask.to(logits.device, non_blocking=True) - - # Create the executor. - executor_config = trtllm.ExecutorConfig(args.beam_width) - logits_proc_config = trtllm.LogitsPostProcessorConfig() - if not args.lpp_batched: - logits_proc_config.processor_map = { - "my_logits_pp": logits_post_processor - } - else: - logits_proc_config.processor_batched = logits_post_processor_batched - executor_config.logits_post_processor_config = logits_proc_config - executor = trtllm.Executor(args.engine_path, trtllm.ModelType.DECODER_ONLY, - executor_config) - - input = "Please give me information about Michael Jordan. You MUST answer using the following json schema: " - prompt = tokenizer.encode(input) - print(f"Input text: {input}\n") - - if executor.can_enqueue_requests(): - request_ids = enqueue_requests(args, executor) - output_tokens = wait_for_responses(args, request_ids, executor) - - # Print output - for req_id in request_ids: - for beam_id in range(args.beam_width): - result = tokenizer.decode( - output_tokens[req_id][beam_id][len(prompt):]) - generated_tokens = len( - output_tokens[req_id][beam_id]) - len(prompt) - print( - f"Request {req_id} Beam {beam_id} ({generated_tokens} tokens): {result}" - ) diff --git a/examples/cpp/executor/CMakeLists.txt b/examples/cpp/executor/CMakeLists.txt deleted file mode 100644 index b448667e8d75..000000000000 --- a/examples/cpp/executor/CMakeLists.txt +++ /dev/null @@ -1,161 +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. cmake needs this line - -cmake_minimum_required(VERSION 3.27) - -set(TRTLLM_DIR "${CMAKE_CURRENT_SOURCE_DIR}/../../..") -list(APPEND CMAKE_MODULE_PATH "${TRTLLM_DIR}/cpp/cmake/modules") - -if(NOT TRTLLM_BUILD_DIR) - set(TRTLLM_BUILD_DIR "${TRTLLM_DIR}/cpp/build") -endif() -set(TRTLLM_LIB_PATH "${TRTLLM_BUILD_DIR}/tensorrt_llm/libtensorrt_llm.so") -if(NOT EXISTS ${TRTLLM_LIB_PATH}) - message(FATAL_ERROR "Cannot find ${TRTLLM_LIB_PATH}") -endif() - -set(TRTLLM_PLUGIN_PATH - "${TRTLLM_BUILD_DIR}/tensorrt_llm/plugins/libnvinfer_plugin_tensorrt_llm.so" -) -set(TRTLLM_INCLUDE_DIR "${TRTLLM_DIR}/cpp/include") - -option( - ENABLE_MULTI_DEVICE - "Enable multi device/instance examples building (requires MPI headers/libs)" - ON) - -# Determine CXX11 ABI compatibility -execute_process( - COMMAND bash -c "nm -f posix -D ${TRTLLM_LIB_PATH} | grep __cxx11" - RESULT_VARIABLE GLIB_CXX11_FOUND - OUTPUT_QUIET) -if(GLIB_CXX11_FOUND EQUAL 0) - set(USE_CXX11_ABI 1) -else() - set(USE_CXX11_ABI 0) -endif() -message(STATUS "Use CXX11 ABI: ${USE_CXX11_ABI}") -add_compile_options("-D_GLIBCXX_USE_CXX11_ABI=${USE_CXX11_ABI}") - -set(CMAKE_CXX_STANDARD 17) -set(CMAKE_CXX_STANDARD_REQUIRED TRUE) -set(CMAKE_VERBOSE_MAKEFILE 1) - -# Define project name -project(executorExamples) - -# Compile options $ ? -if(ENABLE_MULTI_DEVICE) - set(EMD 1) -else() - set(EMD 0) -endif() -set(CMAKE_CXX_FLAGS "-Wall -pthread -lstdc++ -DENABLE_MULTI_DEVICE=${EMD} ") -set(CMAKE_CXX_FLAGS_RELEASE "-O3") -set(CMAKE_BUILD_TYPE release) - -find_package(CUDAToolkit REQUIRED COMPONENTS cuda_driver cudart_static nvml) -message(STATUS "CUDA library status:") -message(STATUS " version: ${CUDAToolkit_VERSION}") -message(STATUS " libraries: ${CUDAToolkit_LIBRARY_DIR}") -message(STATUS " include path: ${CUDAToolkit_INCLUDE_DIRS}") - -# TRT dependencies -find_package(TensorRT 10 REQUIRED) - -if(${CUDAToolkit_VERSION} VERSION_GREATER_EQUAL "11") - add_definitions("-DENABLE_BF16") - message( - STATUS - "CUDA_VERSION ${CUDA_VERSION} is greater or equal than 11.0, enable -DENABLE_BF16 flag" - ) -endif() - -if(${CUDAToolkit_VERSION} VERSION_GREATER_EQUAL "11.8") - add_definitions("-DENABLE_FP8") - message( - STATUS - "CUDA_VERSION ${CUDA_VERSION} is greater or equal than 11.8, enable -DENABLE_FP8 flag" - ) -endif() - -add_subdirectory(${TRTLLM_DIR}/3rdparty 3rdparty) -FetchContent_MakeAvailable(cxxopts) - -# tensorrt_llm shared lib -add_library(tensorrt_llm SHARED IMPORTED) -set_property(TARGET tensorrt_llm PROPERTY IMPORTED_LOCATION ${TRTLLM_LIB_PATH}) -set_property( - TARGET tensorrt_llm PROPERTY IMPORTED_LINK_INTERFACE_LIBRARIES - CUDA::cuda_driver CUDA::cudart_static CUDA::nvml) - -# nvinfer_plugin_tensorrt_llm shared lib -add_library(nvinfer_plugin_tensorrt_llm SHARED IMPORTED) -set_property(TARGET nvinfer_plugin_tensorrt_llm PROPERTY IMPORTED_LOCATION - ${TRTLLM_PLUGIN_PATH}) -set_property(TARGET nvinfer_plugin_tensorrt_llm - PROPERTY IMPORTED_LINK_INTERFACE_LIBRARIES tensorrt_llm) - -include_directories(${TRTLLM_INCLUDE_DIR} ${CUDAToolkit_INCLUDE_DIRS}) - -# Basic -add_executable(executorExampleBasic executorExampleBasic.cpp) -target_link_libraries(executorExampleBasic nvinfer_plugin_tensorrt_llm) - -add_executable(executorExampleDebug executorExampleDebug.cpp) -target_link_libraries(executorExampleDebug nvinfer_plugin_tensorrt_llm) - -add_executable(executorExampleKvEvents executorExampleKvEvents.cpp) -target_link_libraries(executorExampleKvEvents nvinfer_plugin_tensorrt_llm - cxxopts::cxxopts) - -add_executable(executorExampleLogitsProcessor - executorExampleLogitsProcessor.cpp) -target_link_libraries(executorExampleLogitsProcessor - nvinfer_plugin_tensorrt_llm) - -# Advanced -if(NOT TARGET cxxopts::cxxopts) - add_subdirectory(${CMAKE_BINARY_DIR}/_deps/cxxopts-src - ${CMAKE_CURRENT_BINARY_DIR}/cxxopts) -endif() - -add_executable(executorExampleAdvanced executorExampleAdvanced.cpp) -target_link_libraries(executorExampleAdvanced nvinfer_plugin_tensorrt_llm - cxxopts::cxxopts) - -# MultiInstance -if(ENABLE_MULTI_DEVICE) - find_package(MPI REQUIRED) - message(STATUS "Using MPI_C_INCLUDE_DIRS: ${MPI_C_INCLUDE_DIRS}") - message(STATUS "Using MPI_C_LIBRARIES: ${MPI_C_LIBRARIES}") - include_directories(${MPI_C_INCLUDE_DIRS}) - - add_executable(executorExampleAdvancedMultiInstances - executorExampleAdvancedMultiInstances.cpp) - target_link_libraries( - executorExampleAdvancedMultiInstances nvinfer_plugin_tensorrt_llm - cxxopts::cxxopts ${MPI_C_LIBRARIES}) - - # FastLogits - add_executable(executorExampleFastLogits executorExampleFastLogits.cpp) - target_link_libraries(executorExampleFastLogits nvinfer_plugin_tensorrt_llm - cxxopts::cxxopts ${MPI_C_LIBRARIES}) - - add_executable(executorExampleDisaggregated executorExampleDisaggregated.cpp) - target_link_libraries( - executorExampleDisaggregated nvinfer_plugin_tensorrt_llm cxxopts::cxxopts - ${MPI_C_LIBRARIES}) -endif() diff --git a/examples/cpp/executor/README.md b/examples/cpp/executor/README.md deleted file mode 100644 index 597a38effe01..000000000000 --- a/examples/cpp/executor/README.md +++ /dev/null @@ -1,135 +0,0 @@ -# Executor API examples - -This directory contains several examples that demonstrate how to use the `Executor` API: -- The example defined in `executorExampleBasic.cpp` shows how you can generate output tokens for a single prompt in only a few lines of code. -- The example defined in `executorExampleAdvanced.cpp` supports more options such as providing an arbitrary number of input requests with arbitrary tokens per request and running in streaming mode. -- The example defined in `executorExampleLogitsProcessor.cpp` shows how to use `LogitsPostProcessor` to control output tokens. -- The example defined in `executorExampleFastLogits.cpp` shows how to use `ExternalDraftTokensConfig` for speculative decoding and optionally use the fast logits feature. -- The example defined in `executorExampleKvEvents.cpp` shows how to use the KV cache event API. -- The example defined in `executorExampleDisaggregated.cpp` shows how to use the disaggregated executor API. - -## Building the examples - -To build the examples, you first need to build the TensorRT LLM C++ shared libraries (`libtensorrt_llm.so` and `libnvinfer_plugin_tensorrt_llm.so`) using the [`build_wheel.py`](source:scripts/build_wheel.py) script. Alternatively, if you have already build the TensorRT LLM libraries, you can modify the provided `CMakeLists.txt` such that the `libtensorrt_llm.so` and `libnvinfer_plugin_tensorrt_llm.so` are imported properly. - -Once the TensorRT LLM libraries are built, you can run - -``` -mkdir build -cd build -cmake .. -make -j -``` -from the `./examples/cpp/executor/` folder to build the basic and advanced examples. - -## Preparing the TensorRT LLM engine(s) - -Before you run the examples, please make sure that you have already built engine(s) using the TensorRT LLM API. - -Use `trtllm-build` to build the TRT-LLM engine. - -## Running the examples - -### executorExampleBasic - -From the `examples/cpp/executor/build` folder, you can get run the `executorExampleBasic` example with: - -``` -./executorExampleBasic -``` -where `` is the path to the directly containing the TensorRT engine files. - -### executorExampleDebug - -This example shows how you can define which engine IO tensors should be dumped to numpy files. -From the `examples/cpp/executor/build` folder, you can get run the `executorExampleDebug` example with: - -``` -./executorExampleDebug -``` -where `` is the path to the directly containing the TensorRT engine files. - -### executorExampleAdvanced - -From the `examples/cpp/executor/build` folder, you can also run the `executorExampleAdvanced` example. To get the full list of supported input arguments, type - -``` -./executorExampleAdvanced -h -``` - -For example, you can run: - -``` -./executorExampleAdvanced --engine_dir --input_tokens_csv_file ../inputTokens.csv -``` - -to run with the provided dummy input tokens from `inputTokens.csv`. Upon successful completion, you should see the following in the logs: -``` -[TensorRT-LLM][INFO] Creating request with 6 input tokens -[TensorRT-LLM][INFO] Creating request with 4 input tokens -[TensorRT-LLM][INFO] Creating request with 10 input tokens -[TensorRT-LLM][INFO] Got 20 tokens for beam 0 for requestId 3 -[TensorRT-LLM][INFO] Request id 3 is completed. -[TensorRT-LLM][INFO] Got 14 tokens for beam 0 for requestId 2 -[TensorRT-LLM][INFO] Request id 2 is completed. -[TensorRT-LLM][INFO] Got 16 tokens for beam 0 for requestId 1 -[TensorRT-LLM][INFO] Request id 1 is completed. -[TensorRT-LLM][INFO] Writing output tokens to outputTokens.csv -[TensorRT-LLM][INFO] Exiting. -``` - -#### Multi-GPU run - -To run the `executorExampleAdvanced` on models that require multiple GPUs, you can run the example using MPI as follows: - -``` -mpirun -n --allow-run-as-root ./executorExampleAdvanced --engine_dir --input_tokens_csv_file ../inputTokens.csv -``` -where `` must equal to `tp*pp` for the TensorRT engine. By default GPU device IDs `[0...(num_ranks-1)]` will be used. - -Alternatively, it's also possible to run multi-GPU model by using the so-called `Orchestrator` communication mode, where the `Executor` instance will automatically spawn additional processes to run the model on multiple GPUs. To use the `Orchestrator` communication mode, you can run the example with: - -``` -./executorExampleAdvanced --engine_dir --input_tokens_csv_file ../inputTokens.csv --use_orchestrator_mode --worker_executable_path -``` -where `` is the absolute path to the stand-alone executor worker executable, located at`cpp/build/tensorrt_llm/executor_worker/executorWorker` by default. - - -### executorExampleFastLogits - -To run the `executorExampleFastLogits`, you need two GPUs (one for the draft model and one for the target model). You can run it as follows: - -``` -mpirun -n 3 --allow-run-as-root ./executorExampleFastLogits --engine_dir --draft_engine_dir --num_draft_tokens=3 -``` - -The examples uses 3 MPI ranks (one for the orchestrator, one for the draft model and one for the target model). - -Use `--fast_logits=false` to disable the fast logits feature. - -### executorExampleKvEvents - -From the `examples/cpp/executor/build` folder, you can get run the `executorExampleKvEvents` example with: - -``` -./executorExampleKvEvents --engine_dir -``` -where `` is the path to the directly containing the TensorRT engine files. - -This example shows how the KV Cache Event API can be used to reconstruct the state of TRT-LLM's internal radix tree. This can be used in applications such as smart routing to route requests between multiple executor instances to maximize KV Cache reuse. Events are emitted when blocks are stored, removed, or updated in the radix tree. - -### executorExampleDisaggregated - -From the `examples/cpp/executor/build` folder, you can also run the `executorExampleDisaggregated` example. To get the full list of supported input arguments, type -``` -./executorExampleDisaggregated -h -``` -Note setting `TRTLLM_USE_UCX_KVCACHE=1` is required to run disaggregated executor. -For example, you can run : -``` -export TRTLLM_USE_UCX_KVCACHE=1 - -mpirun -n --allow-run-as-root --oversubscribe ./executorExampleDisaggregated --context_engine_dir --context_rank_size --generation_engine_dir --generation_rank_size --input_tokens_csv_file ../inputTokens.csv - -``` -where `` must equal to `tp*pp` for the context engine, and `` must equal to `tp*pp` for the generation engine,the context engine and generation engine can be heterogeneous in parallelism. `` must equal to `++1`, the additional rank is used as orchestrator process. diff --git a/examples/cpp/executor/executorExampleAdvanced.cpp b/examples/cpp/executor/executorExampleAdvanced.cpp deleted file mode 100644 index e2fdf489f63d..000000000000 --- a/examples/cpp/executor/executorExampleAdvanced.cpp +++ /dev/null @@ -1,368 +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. - */ - -#include -#include -#include -#include -#include - -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/logger.h" -#include "tensorrt_llm/executor/executor.h" -#include "tensorrt_llm/plugins/api/tllmPlugin.h" -#include - -namespace tle = tensorrt_llm::executor; - -namespace fs = std::filesystem; - -struct RuntimeOptions -{ - std::string trtEnginePath; - std::string inputTokensCsvFile; - std::string outputTokensCsvFile; - - bool streaming; - bool excludeInputFromOutput; - tle::SizeType32 maxNewTokens; - tle::SizeType32 beamWidth; - std::optional numReturnSequences; - tle::SizeType32 timeoutMs; - - bool useOrchestratorMode; - std::string workerExecutablePath; -}; - -// Utility function to parse input arguments -RuntimeOptions parseArgs(int argc, char* argv[]); - -// Function that enqueues requests -std::vector enqueueRequests(RuntimeOptions const& runtimeOpts, tle::Executor& executor); - -// Function that waits for responses and stores output tokens -std::unordered_map waitForResponses( - RuntimeOptions const& runtimeOpts, std::vector const& requestIds, tle::Executor& executor); - -// Utility function to read input tokens from csv file -std::vector readInputTokens(std::string const& path); - -// Utility function to write output tokens from csv file -void writeOutputTokens(std::string const& path, std::vector& requestIds, - std::unordered_map const& outputTokens, tle::SizeType32 beamWidth); - -tle::SizeType32 getNumSequencesPerRequest(RuntimeOptions const& runtimeOpts); - -// Main -int main(int argc, char* argv[]) -{ - // Register the TRT-LLM plugins - initTrtLlmPlugins(); - - auto runtimeOpts = parseArgs(argc, argv); - - // Create the executor for this engine - auto executorConfig = tle::ExecutorConfig(runtimeOpts.beamWidth); - - if (runtimeOpts.useOrchestratorMode) - { - auto orchestratorConfig = tle::OrchestratorConfig(true, runtimeOpts.workerExecutablePath); - auto parallelConfig = tle::ParallelConfig(tle::CommunicationType::kMPI, tle::CommunicationMode::kORCHESTRATOR, - std::nullopt, std::nullopt, orchestratorConfig); - executorConfig.setParallelConfig(parallelConfig); - } - - auto executor = tle::Executor(runtimeOpts.trtEnginePath, tle::ModelType::kDECODER_ONLY, executorConfig); - - if (executor.canEnqueueRequests()) - { - // Create the requests - auto requestIds = enqueueRequests(runtimeOpts, executor); - - // Wait for responses and store output tokens - auto outputTokens = waitForResponses(runtimeOpts, requestIds, executor); - - // Write output tokens csv file - TLLM_LOG_INFO("Writing output tokens to %s", runtimeOpts.outputTokensCsvFile.c_str()); - auto numSequences = getNumSequencesPerRequest(runtimeOpts); - writeOutputTokens(runtimeOpts.outputTokensCsvFile, requestIds, outputTokens, numSequences); - } - TLLM_LOG_INFO("Exiting."); - return 0; -} - -RuntimeOptions parseArgs(int argc, char* argv[]) -{ - RuntimeOptions runtimeOpts; - - cxxopts::Options options(argv[0], "Example that demonstrates how to use the Executor API"); - options.add_options()("h,help", "Print usage"); - options.add_options()("engine_dir", "Directory that store the engines.", cxxopts::value()); - options.add_options()("beam_width", "The beam width", cxxopts::value()->default_value("1")); - options.add_options()( - "num_return_sequences", "The number of return sequences per request.", cxxopts::value>()); - options.add_options()("streaming", "Operate in streaming mode", cxxopts::value()->default_value("false")); - options.add_options()("exclude_input_from_output", - "Exclude input tokens when writing output tokens. Only has effect for streaming = false. For streaming = true, " - "output tokens are not included.", - cxxopts::value()->default_value("false")); - options.add_options()( - "max_new_tokens", "The maximum number of tokens to generate", cxxopts::value()->default_value("10")); - options.add_options()( - "input_tokens_csv_file", "Path to a csv file that contains input tokens", cxxopts::value()); - options.add_options()("output_tokens_csv_file", "Path to a csv file that will contain the output tokens", - cxxopts::value()->default_value("outputTokens.csv")); - options.add_options()("timeout_ms", "The maximum time to wait for all responses, in milliseconds.", - cxxopts::value()->default_value("10000")); - options.add_options()("use_orchestrator_mode", "Use orchestrator communication mode.", - cxxopts::value()->default_value("false")); - options.add_options()("worker_executable_path", "The location of the worker executable.", - cxxopts::value()->default_value("")); - - auto parsedOptions = options.parse(argc, argv); - - // Argument: help - if (parsedOptions.count("help")) - { - TLLM_LOG_ERROR(options.help()); - exit(0); - } - - // Argument: Engine directory - if (!parsedOptions.count("engine_dir")) - { - TLLM_LOG_ERROR(options.help()); - TLLM_LOG_ERROR("Please specify engine directory."); - exit(1); - } - runtimeOpts.trtEnginePath = parsedOptions["engine_dir"].as(); - if (!fs::exists(runtimeOpts.trtEnginePath) || !fs::is_directory(runtimeOpts.trtEnginePath)) - { - TLLM_LOG_ERROR("Engine directory doesn't exist."); - exit(1); - } - - // Argument: Input tokens csv file - if (!parsedOptions.count("input_tokens_csv_file")) - { - TLLM_LOG_ERROR(options.help()); - TLLM_LOG_ERROR("Please specify input_tokens_csv_file"); - exit(1); - } - runtimeOpts.inputTokensCsvFile = parsedOptions["input_tokens_csv_file"].as(); - runtimeOpts.streaming = parsedOptions["streaming"].as(); - runtimeOpts.excludeInputFromOutput = parsedOptions["exclude_input_from_output"].as(); - runtimeOpts.maxNewTokens = parsedOptions["max_new_tokens"].as(); - runtimeOpts.beamWidth = parsedOptions["beam_width"].as(); - if (parsedOptions.count("num_return_sequences") > 0) - { - runtimeOpts.numReturnSequences = parsedOptions["num_return_sequences"].as>(); - } - runtimeOpts.timeoutMs = parsedOptions["timeout_ms"].as(); - runtimeOpts.outputTokensCsvFile = parsedOptions["output_tokens_csv_file"].as(); - - runtimeOpts.useOrchestratorMode = parsedOptions["use_orchestrator_mode"].as(); - runtimeOpts.workerExecutablePath = parsedOptions["worker_executable_path"].as(); - - return runtimeOpts; -} - -std::vector enqueueRequests(RuntimeOptions const& runtimeOpts, tle::Executor& executor) -{ - tle::OutputConfig outputConfig; - outputConfig.excludeInputFromOutput = runtimeOpts.excludeInputFromOutput; - tle::SamplingConfig samplingConfig(runtimeOpts.beamWidth); - if (runtimeOpts.numReturnSequences && runtimeOpts.beamWidth == 1) - { - samplingConfig.setTopP(0.9); - } - samplingConfig.setNumReturnSequences(runtimeOpts.numReturnSequences); - - TLLM_LOG_INFO("Reading input tokens from %s", runtimeOpts.inputTokensCsvFile.c_str()); - auto inputTokens = readInputTokens(runtimeOpts.inputTokensCsvFile); - TLLM_LOG_INFO("Number of requests: %d", inputTokens.size()); - - std::vector requests; - for (auto& tokens : inputTokens) - { - TLLM_LOG_INFO("Creating request with %d input tokens", tokens.size()); - requests.emplace_back( - std::move(tokens), runtimeOpts.maxNewTokens, runtimeOpts.streaming, samplingConfig, outputConfig); - } - - // Enqueue the requests - auto requestIds = executor.enqueueRequests(std::move(requests)); - - return requestIds; -} - -std::unordered_map waitForResponses( - RuntimeOptions const& runtimeOpts, std::vector const& requestIds, tle::Executor& executor) -{ - // Map that will be used to store output tokens for requests - std::unordered_map outputTokens; - auto numSequences = getNumSequencesPerRequest(runtimeOpts); - for (auto requestId : requestIds) - { - outputTokens[requestId] = tle::BeamTokens(numSequences); - } - - tle::SizeType32 numFinished{0}; - tle::SizeType32 iter{0}; - - // Get the new tokens for each request - while (numFinished < static_cast(requestIds.size()) && iter < runtimeOpts.timeoutMs) - { - std::chrono::milliseconds waitTime(1); - // Wait for any response - auto responses = executor.awaitResponses(waitTime); - - auto insertResponseTokens - = [&outputTokens](tle::IdType requestId, tle::SizeType32 seqIdx, tle::VecTokens const& respTokens) - { - TLLM_LOG_INFO("Got %d tokens for seqIdx %d for requestId %d", respTokens.size(), seqIdx, requestId); - - // Store the output tokens for that request id - auto& outTokens = outputTokens.at(requestId).at(seqIdx); - outTokens.insert(outTokens.end(), std::make_move_iterator(respTokens.begin()), - std::make_move_iterator(respTokens.end())); - }; - - // Loop over the responses - for (auto const& response : responses) - { - auto requestId = response.getRequestId(); - if (!response.hasError()) - { - auto result = response.getResult(); - numFinished += result.isFinal; - if (runtimeOpts.beamWidth > 1) - { - for (tle::SizeType32 beam = 0; beam < numSequences; ++beam) - { - insertResponseTokens(requestId, beam, result.outputTokenIds.at(beam)); - } - } - else - { - insertResponseTokens(requestId, result.sequenceIndex, result.outputTokenIds.at(0)); - } - if (result.isFinal) - { - TLLM_LOG_INFO("Request id %lu is completed.", requestId); - } - } - else - { - // Allow response with error only if awaitResponse processed a terminated request id - std::string err = "ReqId " + std::to_string(response.getRequestId()) - + " has already been processed and was terminated."; - if (response.getErrorMsg() != err) - { - TLLM_THROW("Request id %lu encountered error: %s", requestId, response.getErrorMsg().c_str()); - } - } - } - ++iter; - } - if (iter == runtimeOpts.timeoutMs) - { - TLLM_THROW("Timeout exceeded."); - } - - return outputTokens; -} - -std::vector readInputTokens(std::string const& path) -{ - std::vector data; - std::ifstream file(path); - - if (!file.is_open()) - { - auto const err = std::string{"Failed to open file: "} + path; - TLLM_LOG_ERROR(err); - TLLM_THROW(err); - } - - std::string line; - while (std::getline(file, line)) - { - std::vector row; - std::stringstream ss(line); - std::string token; - - while (std::getline(ss, token, ',')) - { - try - { - row.push_back(std::stoi(token)); - } - catch (std::invalid_argument const& e) - { - TLLM_LOG_ERROR("Invalid argument: %s", e.what()); - } - catch (std::out_of_range const& e) - { - TLLM_LOG_ERROR("Out of range: %s", e.what()); - } - } - - data.push_back(row); - } - - file.close(); - return data; -} - -void writeOutputTokens(std::string const& path, std::vector& requestIds, - std::unordered_map const& outputTokens, tle::SizeType32 beamWidth) -{ - std::ofstream file(path); - - if (!file.is_open()) - { - TLLM_LOG_ERROR("Failed to open file %s", path.c_str()); - return; - } - - for (auto requestId : requestIds) - { - auto const& outTokens = outputTokens.at(requestId); - for (tle::SizeType32 beam = 0; beam < beamWidth; ++beam) - { - auto const& beamTokens = outTokens.at(beam); - for (size_t i = 0; i < beamTokens.size(); ++i) - { - file << beamTokens[i]; - if (i < beamTokens.size() - 1) - { - file << ", "; - } - } - file << "\n"; - } - } - - file.close(); -} - -tle::SizeType32 getNumSequencesPerRequest(RuntimeOptions const& runtimeOpts) -{ - auto numReturnSequences = runtimeOpts.numReturnSequences.value_or(runtimeOpts.beamWidth); - return runtimeOpts.beamWidth > 1 ? std::min(numReturnSequences, runtimeOpts.beamWidth) : numReturnSequences; -} diff --git a/examples/cpp/executor/executorExampleAdvancedMultiInstances.cpp b/examples/cpp/executor/executorExampleAdvancedMultiInstances.cpp deleted file mode 100644 index f967661ccd27..000000000000 --- a/examples/cpp/executor/executorExampleAdvancedMultiInstances.cpp +++ /dev/null @@ -1,381 +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. - */ - -#include -#include -#include -#include -#include - -#include "tensorrt_llm/common/logger.h" -#include "tensorrt_llm/executor/executor.h" -#include "tensorrt_llm/plugins/api/tllmPlugin.h" -#include "tensorrt_llm/runtime/utils/mpiUtils.h" -#include - -namespace tle = tensorrt_llm::executor; - -namespace fs = std::filesystem; - -struct RuntimeOptions -{ - std::string trtEnginePath; - std::string inputTokensCsvFile; - std::string outputTokensCsvFile; - - bool streaming; - bool excludeInputFromOutput; - tle::SizeType32 maxNewTokens; - tle::SizeType32 beamWidth; - tle::SizeType32 timeoutMs; - - bool useOrchestratorMode; - std::string workerExecutablePath; - bool spawnProcesses; -}; - -// Utility function to parse input arguments -RuntimeOptions parseArgs(int argc, char* argv[]); - -// Function that enqueues requests -std::vector> enqueueRequests( - RuntimeOptions const& runtimeOpts, std::deque& executors); - -// Function that waits for responses and stores output tokens -std::map, tle::BeamTokens> waitForResponses(RuntimeOptions const& runtimeOpts, - std::vector> const& instanceRequestIds, std::deque& executors); - -// Utility function to read input tokens from csv file -std::vector readInputTokens(std::string const& path); - -// Utility function to write output tokens from csv file -void writeOutputTokens(std::string const& path, std::vector>& requestIds, - std::map, tle::BeamTokens> const& outputTokens, tle::SizeType32 beamWidth); - -// Main -int main(int argc, char* argv[]) -{ - // Register the TRT-LLM plugins - initTrtLlmPlugins(); - - auto runtimeOpts = parseArgs(argc, argv); - - // Create the executor for this engine - auto executorConfig = tle::ExecutorConfig(runtimeOpts.beamWidth); - - tle::KvCacheConfig kvCacheConfig{false, 10000}; - executorConfig.setKvCacheConfig(kvCacheConfig); - - bool isOrchestrator = true; - if (!runtimeOpts.spawnProcesses) - { - tensorrt_llm::mpi::initialize(tensorrt_llm::mpi::MpiThreadSupport::THREAD_MULTIPLE); - int myRank = tensorrt_llm::mpi::MpiComm::world().getRank(); - isOrchestrator = (myRank == 0); - } - - auto orchestratorConfig = tle::OrchestratorConfig( - isOrchestrator, runtimeOpts.workerExecutablePath, nullptr, runtimeOpts.spawnProcesses); - auto parallelConfig = tle::ParallelConfig(tle::CommunicationType::kMPI, tle::CommunicationMode::kORCHESTRATOR, - std::nullopt, std::nullopt, orchestratorConfig); - executorConfig.setParallelConfig(parallelConfig); - - int numInstances = 3; - if (!runtimeOpts.spawnProcesses) - { - // Keep one rank for orchestrator - numInstances = tensorrt_llm::mpi::MpiComm::world().getSize() - 1; - } - std::deque executors; - for (int instanceId = 0; instanceId < numInstances; ++instanceId) - { - auto executorConfigTmp = executorConfig; - // Set the rank id participating in each model instance - if (!runtimeOpts.spawnProcesses) - { - parallelConfig.setParticipantIds({instanceId + 1}); - } - executorConfigTmp.setParallelConfig(parallelConfig); - executors.emplace_back(runtimeOpts.trtEnginePath, tle::ModelType::kDECODER_ONLY, executorConfigTmp); - } - - // Only orchestrator rank (rank 0) will enter - if (isOrchestrator) - { - // Create the requests - auto instanceRequestIds = enqueueRequests(runtimeOpts, executors); - - // Wait for responses and store output tokens - auto outputTokens = waitForResponses(runtimeOpts, instanceRequestIds, executors); - - // Write output tokens csv file - TLLM_LOG_INFO("Writing output tokens to %s", runtimeOpts.outputTokensCsvFile.c_str()); - writeOutputTokens(runtimeOpts.outputTokensCsvFile, instanceRequestIds, outputTokens, runtimeOpts.beamWidth); - } - TLLM_LOG_INFO("Exiting."); - return 0; -} - -RuntimeOptions parseArgs(int argc, char* argv[]) -{ - RuntimeOptions runtimeOpts; - - cxxopts::Options options(argv[0], "Example that demonstrates how to use the Executor API"); - options.add_options()("h,help", "Print usage"); - options.add_options()("engine_dir", "Directory that store the engines.", cxxopts::value()); - options.add_options()("beam_width", "The beam width", cxxopts::value()->default_value("1")); - options.add_options()("streaming", "Operate in streaming mode", cxxopts::value()->default_value("false")); - options.add_options()("exclude_input_from_output", - "Exclude input tokens when writing output tokens. Only has effect for streaming = false. For streaming = true, " - "output tokens are not included.", - cxxopts::value()->default_value("false")); - options.add_options()( - "max_new_tokens", "The maximum number of tokens to generate", cxxopts::value()->default_value("10")); - options.add_options()( - "input_tokens_csv_file", "Path to a csv file that contains input tokens", cxxopts::value()); - options.add_options()("output_tokens_csv_file", "Path to a csv file that will contain the output tokens", - cxxopts::value()->default_value("outputTokens.csv")); - options.add_options()("timeout_ms", "The maximum time to wait for all responses, in milliseconds.", - cxxopts::value()->default_value("10000")); - options.add_options()("worker_executable_path", "The location of the worker executable.", - cxxopts::value()->default_value("")); - options.add_options()("spawn_processes", - "Flag that controls if MPI_Comm_spawn should be used to spawn worker processes, or if they have been launched " - "with mpi already.", - cxxopts::value()->default_value("true")); - - auto parsedOptions = options.parse(argc, argv); - - // Argument: help - if (parsedOptions.count("help")) - { - TLLM_LOG_ERROR(options.help()); - exit(0); - } - - // Argument: Engine directory - if (!parsedOptions.count("engine_dir")) - { - TLLM_LOG_ERROR(options.help()); - TLLM_LOG_ERROR("Please specify engine directory."); - exit(1); - } - runtimeOpts.trtEnginePath = parsedOptions["engine_dir"].as(); - if (!fs::exists(runtimeOpts.trtEnginePath) || !fs::is_directory(runtimeOpts.trtEnginePath)) - { - TLLM_LOG_ERROR("Engine directory doesn't exist."); - exit(1); - } - - // Argument: Input tokens csv file - if (!parsedOptions.count("input_tokens_csv_file")) - { - TLLM_LOG_ERROR(options.help()); - TLLM_LOG_ERROR("Please specify input_tokens_csv_file"); - exit(1); - } - runtimeOpts.inputTokensCsvFile = parsedOptions["input_tokens_csv_file"].as(); - runtimeOpts.streaming = parsedOptions["streaming"].as(); - runtimeOpts.excludeInputFromOutput = parsedOptions["exclude_input_from_output"].as(); - runtimeOpts.maxNewTokens = parsedOptions["max_new_tokens"].as(); - runtimeOpts.beamWidth = parsedOptions["beam_width"].as(); - runtimeOpts.timeoutMs = parsedOptions["timeout_ms"].as(); - runtimeOpts.outputTokensCsvFile = parsedOptions["output_tokens_csv_file"].as(); - - runtimeOpts.workerExecutablePath = parsedOptions["worker_executable_path"].as(); - runtimeOpts.spawnProcesses = parsedOptions["spawn_processes"].as(); - - return runtimeOpts; -} - -std::vector> enqueueRequests( - RuntimeOptions const& runtimeOpts, std::deque& executors) -{ - tle::OutputConfig outputConfig; - outputConfig.excludeInputFromOutput = runtimeOpts.excludeInputFromOutput; - tle::SamplingConfig samplingConfig(runtimeOpts.beamWidth); - - TLLM_LOG_INFO("Reading input tokens from %s", runtimeOpts.inputTokensCsvFile.c_str()); - auto inputTokens = readInputTokens(runtimeOpts.inputTokensCsvFile); - TLLM_LOG_INFO("Number of requests: %d", inputTokens.size()); - - std::vector requests; - for (auto& tokens : inputTokens) - { - TLLM_LOG_INFO("Creating request with %d input tokens", tokens.size()); - requests.emplace_back( - std::move(tokens), runtimeOpts.maxNewTokens, runtimeOpts.streaming, samplingConfig, outputConfig); - } - - // Enqueue the requests - // Round robin over instances - std::vector> instanceRequestIds; - for (size_t req = 0; req < requests.size(); ++req) - { - auto instanceId = req % executors.size(); - TLLM_LOG_INFO("Enqueuing request %d for instance %d", req, instanceId); - auto requestId = executors.at(instanceId).enqueueRequest(requests[req]); - instanceRequestIds.emplace_back(instanceId, requestId); - } - TLLM_LOG_INFO("Enqueued %d requests", instanceRequestIds.size()); - return instanceRequestIds; -} - -std::map, tle::BeamTokens> waitForResponses(RuntimeOptions const& runtimeOpts, - std::vector> const& instanceRequestIds, std::deque& executors) -{ - // Map that will be used to store output tokens for requests - int numRequests = 0; - std::map, tle::BeamTokens> outputTokens; - for (auto instanceRequestId : instanceRequestIds) - { - outputTokens[instanceRequestId] = tle::BeamTokens(runtimeOpts.beamWidth); - numRequests++; - } - - tle::SizeType32 numFinished{0}; - tle::SizeType32 iter{0}; - - // Get the new tokens for each request - while (numFinished < numRequests && iter < runtimeOpts.timeoutMs) - { - std::chrono::milliseconds waitTime(1); - for (size_t instanceId = 0; instanceId < executors.size(); ++instanceId) - { - // Wait for any response for given instance - auto responses = executors.at(instanceId).awaitResponses(waitTime); - // Loop over the responses - for (auto const& response : responses) - { - auto requestId = response.getRequestId(); - if (!response.hasError()) - { - auto result = response.getResult(); - numFinished += result.isFinal; - TLLM_LOG_INFO("Number of finished requests: %d", numFinished); - - for (tle::SizeType32 beam = 0; beam < runtimeOpts.beamWidth; ++beam) - { - auto& respTokens = result.outputTokenIds.at(beam); - - TLLM_LOG_INFO("Got %d tokens for beam %d for requestId %d", respTokens.size(), beam, requestId); - - // Store the output tokens for that request id - auto& outTokens = outputTokens.at(std::make_pair(instanceId, requestId)).at(beam); - outTokens.insert(outTokens.end(), std::make_move_iterator(respTokens.begin()), - std::make_move_iterator(respTokens.end())); - } - if (result.isFinal) - { - TLLM_LOG_INFO("Request id %lu is completed.", requestId); - } - } - else - { - // Allow response with error only if awaitResponse processed a terminated request id - std::string err = "ReqId " + std::to_string(response.getRequestId()) - + " has already been processed and was terminated."; - if (response.getErrorMsg() != err) - { - TLLM_THROW("Request id %lu encountered error: %s", requestId, response.getErrorMsg().c_str()); - } - } - } - } - ++iter; - } - if (iter == runtimeOpts.timeoutMs) - { - TLLM_THROW("Timeout exceeded."); - } - - return outputTokens; -} - -std::vector readInputTokens(std::string const& path) -{ - std::vector data; - std::ifstream file(path); - - if (!file.is_open()) - { - auto const err = std::string{"Failed to open file: "} + path; - TLLM_LOG_ERROR(err); - TLLM_THROW(err); - } - - std::string line; - while (std::getline(file, line)) - { - std::vector row; - std::stringstream ss(line); - std::string token; - - while (std::getline(ss, token, ',')) - { - try - { - row.push_back(std::stoi(token)); - } - catch (std::invalid_argument const& e) - { - TLLM_LOG_ERROR("Invalid argument: %s", e.what()); - } - catch (std::out_of_range const& e) - { - TLLM_LOG_ERROR("Out of range: %s", e.what()); - } - } - - data.push_back(row); - } - - file.close(); - return data; -} - -void writeOutputTokens(std::string const& path, std::vector>& instanceRequestIds, - std::map, tle::BeamTokens> const& outputTokens, tle::SizeType32 beamWidth) -{ - std::ofstream file(path); - - if (!file.is_open()) - { - TLLM_LOG_ERROR("Failed to open file %s", path.c_str()); - return; - } - - for (auto instanceRequestId : instanceRequestIds) - { - auto const& outTokens = outputTokens.at(instanceRequestId); - for (tle::SizeType32 beam = 0; beam < beamWidth; ++beam) - { - auto const& beamTokens = outTokens.at(beam); - for (size_t i = 0; i < beamTokens.size(); ++i) - { - file << beamTokens[i]; - if (i < beamTokens.size() - 1) - { - file << ", "; - } - } - file << "\n"; - } - } - - file.close(); -} diff --git a/examples/cpp/executor/executorExampleBasic.cpp b/examples/cpp/executor/executorExampleBasic.cpp deleted file mode 100644 index b3ae3328392c..000000000000 --- a/examples/cpp/executor/executorExampleBasic.cpp +++ /dev/null @@ -1,60 +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. - */ -#include - -#include "tensorrt_llm/common/logger.h" -#include "tensorrt_llm/executor/executor.h" -#include "tensorrt_llm/plugins/api/tllmPlugin.h" - -namespace tlc = tensorrt_llm::common; -namespace tle = tensorrt_llm::executor; - -int main(int argc, char* argv[]) -{ - // Register the TRT-LLM plugins - initTrtLlmPlugins(); - - if (argc != 2) - { - TLLM_LOG_ERROR("Usage: %s ", argv[0]); - return 1; - } - - // Create the executor for this engine - tle::SizeType32 beamWidth = 1; - auto executorConfig = tle::ExecutorConfig(beamWidth); - auto trtEnginePath = argv[1]; - auto executor = tle::Executor(trtEnginePath, tle::ModelType::kDECODER_ONLY, executorConfig); - - // Create the request - tle::SizeType32 maxNewTokens = 5; - tle::VecTokens inputTokens{1, 2, 3, 4}; - auto request = tle::Request(inputTokens, maxNewTokens); - - // Enqueue the request - auto requestId = executor.enqueueRequest(request); - - // Wait for the response - auto responses = executor.awaitResponses(requestId); - - // Get outputTokens - auto outputTokens = responses.at(0).getResult().outputTokenIds.at(beamWidth - 1); - - TLLM_LOG_INFO("Output tokens: %s", tlc::vec2str(outputTokens).c_str()); - - return 0; -} diff --git a/examples/cpp/executor/executorExampleDebug.cpp b/examples/cpp/executor/executorExampleDebug.cpp deleted file mode 100644 index d0af1a8140b5..000000000000 --- a/examples/cpp/executor/executorExampleDebug.cpp +++ /dev/null @@ -1,65 +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. - */ -#include - -#include "tensorrt_llm/common/logger.h" -#include "tensorrt_llm/executor/executor.h" -#include "tensorrt_llm/plugins/api/tllmPlugin.h" - -namespace tlc = tensorrt_llm::common; -namespace tle = tensorrt_llm::executor; - -int main(int argc, char* argv[]) -{ - // Register the TRT-LLM plugins - initTrtLlmPlugins(); - - if (argc != 2) - { - TLLM_LOG_ERROR("Usage: %s ", argv[0]); - return 1; - } - - // Create the executor for this engine - tle::SizeType32 beamWidth = 1; - auto executorConfig = tle::ExecutorConfig(beamWidth); - // Select which tensors should be dumped - auto debugConfig = tle::DebugConfig(); - debugConfig.setDebugTensorNames({"host_request_types"}); - executorConfig.setDebugConfig(debugConfig); - - auto trtEnginePath = argv[1]; - auto executor = tle::Executor(trtEnginePath, tle::ModelType::kDECODER_ONLY, executorConfig); - - // Create the request - tle::SizeType32 maxNewTokens = 2; - tle::VecTokens inputTokens{1, 2, 3, 4}; - auto request = tle::Request(inputTokens, maxNewTokens); - - // Enqueue the request - auto requestId = executor.enqueueRequest(request); - - // Wait for the response - auto responses = executor.awaitResponses(requestId); - - // Get outputTokens - auto outputTokens = responses.at(0).getResult().outputTokenIds.at(beamWidth - 1); - - TLLM_LOG_INFO("Output tokens: %s", tlc::vec2str(outputTokens).c_str()); - - return 0; -} diff --git a/examples/cpp/executor/executorExampleDisaggregated.cpp b/examples/cpp/executor/executorExampleDisaggregated.cpp deleted file mode 100644 index ef9ff85c5e3b..000000000000 --- a/examples/cpp/executor/executorExampleDisaggregated.cpp +++ /dev/null @@ -1,441 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2022-2025 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. - */ - -#include -#include -#include -#include -#include -#include - -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/logger.h" -#include "tensorrt_llm/executor/executor.h" -#include "tensorrt_llm/executor/types.h" -#include "tensorrt_llm/plugins/api/tllmPlugin.h" -#include "tensorrt_llm/runtime/utils/mpiUtils.h" - -#include - -namespace tle = tensorrt_llm::executor; - -namespace fs = std::filesystem; - -struct RuntimeOptions -{ - std::string trtContextEnginePath; - std::string trtGenerationEnginePath; - std::string inputTokensCsvFile; - std::string outputTokensCsvFile; - - bool streaming; - bool excludeInputFromOutput; - int contextRankSize; - int generationRankSize; - tle::SizeType32 maxNewTokens; - tle::SizeType32 beamWidth; - std::optional numReturnSequences; - tle::SizeType32 timeoutMs; -}; - -RuntimeOptions parseArgs(int argc, char* argv[]); - -// Function that enqueues requests into context executor and generation executor -std::unordered_map enqueueRequests( - RuntimeOptions const& runtimeOpts, tle::Executor& contextExecutor, tle::Executor& generationExecutor); - -// Function that waits for gen responses and stores output tokens -std::unordered_map waitForGenResponses(RuntimeOptions const& runtimeOpts, - std::unordered_map const& genRequestIdToContextRequestId, - tle::Executor& generationExecutor); - -// Utility function to read input tokens from csv file -std::vector readInputTokens(std::string const& path); - -// Utility function to write output tokens from csv file -void writeOutputTokens(std::string const& path, - std::unordered_map& genRequestIdToContextRequestId, - std::unordered_map const& outputTokens, tle::SizeType32 beamWidth); - -int main(int argc, char* argv[]) -{ - - // Register the TRT-LLM plugins - initTrtLlmPlugins(); - - auto runtimeOpts = parseArgs(argc, argv); - TLLM_CHECK_WITH_INFO(runtimeOpts.beamWidth == 1, "Only support beamWidth =1"); - TLLM_CHECK_WITH_INFO( - runtimeOpts.numReturnSequences.has_value() == false || runtimeOpts.numReturnSequences.value() == 1, - "Only support numReturnSequences =1"); - // Create the executor for this engine - auto contextExecutorConfig = tle::ExecutorConfig(runtimeOpts.beamWidth); - auto generationExecutorConfig = tle::ExecutorConfig(runtimeOpts.beamWidth); - bool isOrchestrator = (tensorrt_llm::mpi::MpiComm::world().getRank() == 0); - auto orchestratorConfig = tle::OrchestratorConfig(isOrchestrator, "", nullptr, false); - int contextRankSize = runtimeOpts.contextRankSize; - int generationRankSize = runtimeOpts.generationRankSize; - TLLM_CHECK_WITH_INFO(tensorrt_llm::mpi::MpiComm::world().getSize() >= contextRankSize + generationRankSize + 1, - " MPI should launch at least [contextRankSize+generationRankSize+1]: %d processes", - contextRankSize + generationRankSize + 1); - int deviceCount = -1; - TLLM_CHECK(cudaGetDeviceCount(&deviceCount) == cudaSuccess); - - std::vector contextRankIds(contextRankSize); - std::vector contextDeviceIds(contextRankSize); - std::vector generationRankIds(generationRankSize); - std::vector generationDeviceIds(generationRankSize); - for (int i = 0; i < contextRankSize; i++) - { - contextRankIds[i] = i + 1; - contextDeviceIds[i] = i % deviceCount; - TLLM_LOG_INFO("context Rank %d on device %d", contextRankIds[i], contextDeviceIds[i]); - } - tle::ParallelConfig contextParallelConfig{tensorrt_llm::executor::CommunicationType::kMPI, - tensorrt_llm::executor::CommunicationMode::kORCHESTRATOR, contextDeviceIds, contextRankIds, orchestratorConfig}; - - for (int i = 0; i < generationRankSize; i++) - { - generationRankIds[i] = i + 1 + contextRankSize; - generationDeviceIds[i] = (i + contextRankSize) % deviceCount; - TLLM_LOG_INFO("generation Rank %d on device %d", generationRankIds[i], generationDeviceIds[i]); - } - tle::ParallelConfig generationParallelConfig{tensorrt_llm::executor::CommunicationType::kMPI, - tensorrt_llm::executor::CommunicationMode::kORCHESTRATOR, generationDeviceIds, generationRankIds, - orchestratorConfig}; - - contextExecutorConfig.setParallelConfig(contextParallelConfig); - generationExecutorConfig.setParallelConfig(generationParallelConfig); - - auto contextExecutor - = tle::Executor(runtimeOpts.trtContextEnginePath, tle::ModelType::kDECODER_ONLY, contextExecutorConfig); - auto generationExecutor - = tle::Executor(runtimeOpts.trtGenerationEnginePath, tle::ModelType::kDECODER_ONLY, generationExecutorConfig); - tensorrt_llm::mpi::MpiComm::world().barrier(); - - if (tensorrt_llm::mpi::MpiComm::world().getRank() == 0) - { - - TLLM_CHECK_WITH_INFO(contextExecutor.canEnqueueRequests(), "contextExecutor can't enqueue requests"); - TLLM_CHECK_WITH_INFO(generationExecutor.canEnqueueRequests(), "generationExecutor can't enqueue requests"); - auto genRequestIdsToContextRequestIds = enqueueRequests(runtimeOpts, contextExecutor, generationExecutor); - auto outputTokens = waitForGenResponses(runtimeOpts, genRequestIdsToContextRequestIds, generationExecutor); - TLLM_LOG_INFO("Writing output tokens to %s", runtimeOpts.outputTokensCsvFile.c_str()); - writeOutputTokens( - runtimeOpts.outputTokensCsvFile, genRequestIdsToContextRequestIds, outputTokens, runtimeOpts.beamWidth); - } - tensorrt_llm::mpi::MpiComm::world().barrier(); - - TLLM_LOG_INFO("Exiting."); - return 0; -} - -RuntimeOptions parseArgs(int argc, char* argv[]) -{ - - RuntimeOptions runtimeOpts; - - cxxopts::Options options(argv[0], "Example that demonstrates how to use the Executor Disaggregated API"); - options.add_options()("h,help", "Print usage"); - options.add_options()( - "context_engine_dir", "Directory that store the context engine.", cxxopts::value()); - options.add_options()( - "generation_engine_dir", "Directory that store the generation engine.", cxxopts::value()); - options.add_options()( - "context_rank_size", "The number of ranks for the context engine", cxxopts::value()->default_value("1")); - options.add_options()("generation_rank_size", "The number of ranks for the generation engine", - cxxopts::value()->default_value("1")); - options.add_options()("beam_width", "The beam width", cxxopts::value()->default_value("1")); - options.add_options()( - "num_return_sequences", "The number of return sequences per request.", cxxopts::value>()); - options.add_options()("streaming", "Operate in streaming mode", cxxopts::value()->default_value("false")); - options.add_options()("exclude_input_from_output", - "Exclude input tokens when writing output tokens. Only has effect for streaming = false. For streaming = true, " - "output tokens are not included.", - cxxopts::value()->default_value("false")); - options.add_options()( - "max_new_tokens", "The maximum number of tokens to generate", cxxopts::value()->default_value("10")); - options.add_options()( - "input_tokens_csv_file", "Path to a csv file that contains input tokens", cxxopts::value()); - options.add_options()("output_tokens_csv_file", "Path to a csv file that will contain the output tokens", - cxxopts::value()->default_value("outputTokens.csv")); - options.add_options()("timeout_ms", "The maximum time to wait for all responses, in milliseconds.", - cxxopts::value()->default_value("10000")); - - auto parsedOptions = options.parse(argc, argv); - - // Argument: help - if (parsedOptions.count("help")) - { - TLLM_LOG_ERROR(options.help()); - exit(0); - } - - runtimeOpts.trtContextEnginePath = parsedOptions["context_engine_dir"].as(); - if (!fs::exists(runtimeOpts.trtContextEnginePath) || !fs::is_directory(runtimeOpts.trtContextEnginePath)) - { - TLLM_LOG_ERROR("Context engine directory doesn't exist."); - exit(1); - } - - runtimeOpts.trtGenerationEnginePath = parsedOptions["generation_engine_dir"].as(); - if (!fs::exists(runtimeOpts.trtGenerationEnginePath) || !fs::is_directory(runtimeOpts.trtGenerationEnginePath)) - { - TLLM_LOG_ERROR("Generation engine directory doesn't exist."); - exit(1); - } - // Argument: Input tokens csv file - if (!parsedOptions.count("input_tokens_csv_file")) - { - TLLM_LOG_ERROR(options.help()); - TLLM_LOG_ERROR("Please specify input_tokens_csv_file"); - exit(1); - } - - runtimeOpts.inputTokensCsvFile = parsedOptions["input_tokens_csv_file"].as(); - runtimeOpts.streaming = parsedOptions["streaming"].as(); - runtimeOpts.excludeInputFromOutput = parsedOptions["exclude_input_from_output"].as(); - runtimeOpts.maxNewTokens = parsedOptions["max_new_tokens"].as(); - runtimeOpts.beamWidth = parsedOptions["beam_width"].as(); - runtimeOpts.contextRankSize = parsedOptions["context_rank_size"].as(); - runtimeOpts.generationRankSize = parsedOptions["generation_rank_size"].as(); - if (parsedOptions.count("num_return_sequences") > 0) - { - runtimeOpts.numReturnSequences = parsedOptions["num_return_sequences"].as>(); - } - runtimeOpts.timeoutMs = parsedOptions["timeout_ms"].as(); - runtimeOpts.outputTokensCsvFile = parsedOptions["output_tokens_csv_file"].as(); - - return runtimeOpts; -} - -std::unordered_map enqueueRequests( - RuntimeOptions const& runtimeOpts, tle::Executor& contextExecutor, tle::Executor& generationExecutor) -{ - - tle::OutputConfig outputConfig; - outputConfig.excludeInputFromOutput = runtimeOpts.excludeInputFromOutput; - tle::SamplingConfig samplingConfig(runtimeOpts.beamWidth); - std::unordered_map genRequestIdToContextRequestId; - if (runtimeOpts.numReturnSequences && runtimeOpts.beamWidth == 1) - { - samplingConfig.setTopP(0.9); - } - samplingConfig.setNumReturnSequences(runtimeOpts.numReturnSequences); - - TLLM_LOG_INFO("Reading input tokens from %s", runtimeOpts.inputTokensCsvFile.c_str()); - auto inputTokens = readInputTokens(runtimeOpts.inputTokensCsvFile); - TLLM_LOG_INFO("Number of requests: %d", inputTokens.size()); - - std::vector requests; - for (auto& tokens : inputTokens) - { - TLLM_LOG_INFO("Creating request with %d input tokens", tokens.size()); - requests.emplace_back( - std::move(tokens), runtimeOpts.maxNewTokens, runtimeOpts.streaming, samplingConfig, outputConfig); - requests.back().setRequestType(tensorrt_llm::executor::RequestType::REQUEST_TYPE_CONTEXT_ONLY); - } - - auto contextRequestIds = contextExecutor.enqueueRequests(requests); - - for (size_t i = 0; i < requests.size(); i++) - { - - TLLM_LOG_INFO("waiting response for Context request id: %lu,", contextRequestIds[i]); - auto response = contextExecutor.awaitResponses(contextRequestIds[i]); - TLLM_LOG_INFO("response received for Context request id: %lu", contextRequestIds[i]); - TLLM_CHECK(response.size() == 1); - TLLM_CHECK(response.back().getResult().contextPhaseParams.has_value()); - requests.at(i).setContextPhaseParams(response.back().getResult().contextPhaseParams.value()); - requests.at(i).setRequestType(tensorrt_llm::executor::RequestType::REQUEST_TYPE_GENERATION_ONLY); - auto genRequestId = generationExecutor.enqueueRequest(requests.at(i)); - genRequestIdToContextRequestId[genRequestId] = contextRequestIds[i]; - - TLLM_LOG_INFO("enqueuing generation request for Context request id: %lu, generation request id: %lu", - contextRequestIds[i], genRequestId); - } - - return genRequestIdToContextRequestId; -} - -std::unordered_map waitForGenResponses(RuntimeOptions const& runtimeOpts, - std::unordered_map const& genRequestIdToContextRequestId, - tle::Executor& generationExecutor) -{ - - // Map that will be used to store output tokens for requests - std::unordered_map outputTokens; - std::vector contextRequestIds{}; - std::vector genRequestIds{}; - for (auto const& [key, value] : genRequestIdToContextRequestId) - { - genRequestIds.push_back(key); - contextRequestIds.push_back(value); - } - for (auto contextRequestId : contextRequestIds) - { - outputTokens[contextRequestId] = tle::BeamTokens(runtimeOpts.beamWidth); - } - - tle::SizeType32 numFinished{0}; - tle::SizeType32 iter{0}; - - // Get the new tokens for each request - while (numFinished < static_cast(genRequestIds.size()) && iter < runtimeOpts.timeoutMs) - { - std::chrono::milliseconds waitTime(1); - // Wait for any response - auto responses = generationExecutor.awaitResponses(waitTime); - - auto insertResponseTokens = [&outputTokens, &genRequestIdToContextRequestId](tle::IdType genRequestId, - tle::SizeType32 seqIdx, tle::VecTokens const& respTokens) - { - TLLM_LOG_INFO("Got %d tokens for seqIdx %d for genRequestId %d,contextRequestId %d", respTokens.size(), - seqIdx, genRequestId, genRequestIdToContextRequestId.at(genRequestId)); - - // Store the output tokens for that request id - auto& outTokens = outputTokens.at(genRequestIdToContextRequestId.at(genRequestId)).at(seqIdx); - outTokens.insert(outTokens.end(), std::make_move_iterator(respTokens.begin()), - std::make_move_iterator(respTokens.end())); - }; - - // Loop over the responses - for (auto const& response : responses) - { - auto genRequestId = response.getRequestId(); - if (!response.hasError()) - { - auto result = response.getResult(); - numFinished += result.isFinal; - if (runtimeOpts.beamWidth > 1) - { - for (tle::SizeType32 beam = 0; beam < runtimeOpts.beamWidth; ++beam) - { - insertResponseTokens(genRequestId, beam, result.outputTokenIds.at(beam)); - } - } - else - { - insertResponseTokens(genRequestId, result.sequenceIndex, result.outputTokenIds.at(0)); - } - if (result.isFinal) - { - TLLM_LOG_INFO("genRequest id %lu ,contextRequestId %lu is completed.", genRequestId, - genRequestIdToContextRequestId.at(genRequestId)); - } - } - else - { - // Allow response with error only if awaitResponse processed a terminated request id - std::string err = "genReqId " + std::to_string(response.getRequestId()) - + " has already been processed and was terminated."; - if (response.getErrorMsg() != err) - { - TLLM_THROW("GenRequest id %lu encountered error: %s", genRequestId, response.getErrorMsg().c_str()); - } - } - } - ++iter; - } - if (iter == runtimeOpts.timeoutMs) - { - TLLM_THROW("Timeout exceeded."); - } - - return outputTokens; -} - -std::vector readInputTokens(std::string const& path) -{ - std::vector data; - std::ifstream file(path); - - if (!file.is_open()) - { - auto const err = std::string{"Failed to open file: "} + path; - TLLM_LOG_ERROR(err); - TLLM_THROW(err); - } - - std::string line; - while (std::getline(file, line)) - { - std::vector row; - std::stringstream ss(line); - std::string token; - - while (std::getline(ss, token, ',')) - { - try - { - row.push_back(std::stoi(token)); - } - catch (std::invalid_argument const& e) - { - TLLM_LOG_ERROR("Invalid argument: %s", e.what()); - } - catch (std::out_of_range const& e) - { - TLLM_LOG_ERROR("Out of range: %s", e.what()); - } - } - - data.push_back(row); - } - - file.close(); - return data; -} - -void writeOutputTokens(std::string const& path, - std::unordered_map& genRequestIdToContextRequestId, - std::unordered_map const& outputTokens, tle::SizeType32 beamWidth) -{ - std::ofstream file(path); - - if (!file.is_open()) - { - TLLM_LOG_ERROR("Failed to open file %s", path.c_str()); - return; - } - std::vector requestIds; - for (auto const& [key, value] : genRequestIdToContextRequestId) - { - requestIds.push_back(value); - } - std::sort(requestIds.begin(), requestIds.end()); - - for (auto requestId : requestIds) - { - auto const& outTokens = outputTokens.at(requestId); - for (tle::SizeType32 beam = 0; beam < beamWidth; ++beam) - { - auto const& beamTokens = outTokens.at(beam); - for (size_t i = 0; i < beamTokens.size(); ++i) - { - file << beamTokens[i]; - if (i < beamTokens.size() - 1) - { - file << ", "; - } - } - file << "\n"; - } - } - - file.close(); -} diff --git a/examples/cpp/executor/executorExampleFastLogits.cpp b/examples/cpp/executor/executorExampleFastLogits.cpp deleted file mode 100644 index 3611a1bb73d7..000000000000 --- a/examples/cpp/executor/executorExampleFastLogits.cpp +++ /dev/null @@ -1,264 +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. - */ - -#include -#include -#include -#include -#include - -#include "tensorrt_llm/common/logger.h" -#include "tensorrt_llm/executor/executor.h" -#include "tensorrt_llm/plugins/api/tllmPlugin.h" -#include "tensorrt_llm/runtime/utils/mpiUtils.h" -#include - -namespace tlc = tensorrt_llm::common; -namespace tle = tensorrt_llm::executor; - -namespace fs = std::filesystem; - -struct RuntimeOptions -{ - std::string trtDraftEnginePath; - std::string trtEnginePath; - bool fastLogits; - tle::SizeType32 numDraftTokens; -}; - -// Utility function to parse input arguments -RuntimeOptions parseArgs(int argc, char* argv[]); - -// Runs a draft request -tle::Result executeDraftRequest(tle::Executor& executor, RuntimeOptions const& runtimeOpts); - -// Runs a target request -tle::Result executeTargetRequest( - tle::Executor& executor, tle::Result const& draftResult, RuntimeOptions const& runtimeOpts); - -// Main -int main(int argc, char* argv[]) -{ - // Register the TRT-LLM plugins - initTrtLlmPlugins(); - - auto runtimeOpts = parseArgs(argc, argv); - - // Create the executor for this engine - auto executorConfig = tle::ExecutorConfig(); - - tensorrt_llm::mpi::initialize(tensorrt_llm::mpi::MpiThreadSupport::THREAD_MULTIPLE); - int const myRank = tensorrt_llm::mpi::MpiComm::world().getRank(); - bool const isOrchestrator = (myRank == 0); - - auto kvCacheConfig = tle::KvCacheConfig(true /* enableBlockReuse */); - executorConfig.setKvCacheConfig(kvCacheConfig); - - auto orchestratorConfig - = tle::OrchestratorConfig(isOrchestrator, "" /* workerExecutablePath */, nullptr, false /* spawnPrcesses */); - auto parallelConfig = tle::ParallelConfig(tle::CommunicationType::kMPI, tle::CommunicationMode::kORCHESTRATOR, - std::nullopt, std::nullopt, orchestratorConfig); - executorConfig.setParallelConfig(parallelConfig); - - auto specDecConfig = tle::SpeculativeDecodingConfig(runtimeOpts.fastLogits); - executorConfig.setSpecDecConfig(specDecConfig); - - std::unique_ptr draftExecutor; - std::unique_ptr targetExecutor; - - if (isOrchestrator) - { - auto executorConfigDraft = executorConfig; - parallelConfig.setParticipantIds({1}); - executorConfigDraft.setParallelConfig(parallelConfig); - - draftExecutor = std::make_unique( - runtimeOpts.trtDraftEnginePath, tle::ModelType::kDECODER_ONLY, executorConfigDraft); - - parallelConfig.setParticipantIds({2}); - executorConfig.setParallelConfig(parallelConfig); - - targetExecutor - = std::make_unique(runtimeOpts.trtEnginePath, tle::ModelType::kDECODER_ONLY, executorConfig); - } - else if (myRank == 1) // draft model process - { - parallelConfig.setParticipantIds({1}); - parallelConfig.setDeviceIds({0}); - executorConfig.setParallelConfig(parallelConfig); - draftExecutor = std::make_unique( - runtimeOpts.trtDraftEnginePath, tle::ModelType::kDECODER_ONLY, executorConfig); - } - else if (myRank == 2) // target model process - { - parallelConfig.setParticipantIds({2}); - parallelConfig.setDeviceIds({1}); - executorConfig.setParallelConfig(parallelConfig); - targetExecutor - = std::make_unique(runtimeOpts.trtEnginePath, tle::ModelType::kDECODER_ONLY, executorConfig); - ; - } - - // Only orchestrator rank (rank 0) will enter - if (isOrchestrator) - { - auto draftResult = executeDraftRequest(*draftExecutor, runtimeOpts); - - executeTargetRequest(*targetExecutor, draftResult, runtimeOpts); - } - TLLM_LOG_INFO("Exiting."); - return 0; -} - -RuntimeOptions parseArgs(int argc, char* argv[]) -{ - RuntimeOptions runtimeOpts; - - cxxopts::Options options(argv[0], "Example that demonstrates how to use the Executor API"); - options.add_options()("h,help", "Print usage"); - options.add_options()("engine_dir", "Directory that store the engine.", cxxopts::value()); - options.add_options()("draft_engine_dir", "Directory that store the draft engine.", cxxopts::value()); - options.add_options()( - "fast_logits", "Use speculative decoding fast logits feature", cxxopts::value()->default_value("true")); - options.add_options()( - "num_draft_tokens", "Number of draft tokens to use", cxxopts::value()->default_value("5")); - - auto parsedOptions = options.parse(argc, argv); - - // Argument: help - if (parsedOptions.count("help")) - { - TLLM_LOG_ERROR(options.help()); - exit(0); - } - - // Argument: Engine directory - if (!parsedOptions.count("engine_dir")) - { - TLLM_LOG_ERROR(options.help()); - TLLM_LOG_ERROR("Please specify engine directory."); - exit(1); - } - runtimeOpts.trtEnginePath = parsedOptions["engine_dir"].as(); - if (!fs::exists(runtimeOpts.trtEnginePath) || !fs::is_directory(runtimeOpts.trtEnginePath)) - { - TLLM_LOG_ERROR("Engine directory doesn't exist."); - exit(1); - } - - // Argument: Draft engine directory - if (!parsedOptions.count("draft_engine_dir")) - { - TLLM_LOG_ERROR(options.help()); - TLLM_LOG_ERROR("Please specify draft engine directory."); - exit(1); - } - runtimeOpts.trtDraftEnginePath = parsedOptions["draft_engine_dir"].as(); - if (!fs::exists(runtimeOpts.trtDraftEnginePath) || !fs::is_directory(runtimeOpts.trtDraftEnginePath)) - { - TLLM_LOG_ERROR("Draft engine directory doesn't exist."); - exit(1); - } - - runtimeOpts.fastLogits = parsedOptions["fast_logits"].as(); - runtimeOpts.numDraftTokens = parsedOptions["num_draft_tokens"].as(); - - return runtimeOpts; -} - -tle::Result executeDraftRequest(tle::Executor& executor, RuntimeOptions const& runtimeOpts) -{ - tle::OutputConfig outputConfig; - outputConfig.returnGenerationLogits = true; - - // Create the request - tle::SizeType32 maxNewTokens = runtimeOpts.numDraftTokens; - tle::VecTokens inputTokens{1, 2, 3, 4}; - - tle::Request request{std::move(inputTokens), maxNewTokens}; - request.setOutputConfig(outputConfig); - - // Enqueue the request - auto requestId = executor.enqueueRequest(std::move(request)); - - // Wait for the response - auto responses = executor.awaitResponses(requestId); - - if (responses.at(0).hasError()) - { - TLLM_LOG_ERROR(responses.at(0).getErrorMsg()); - exit(1); - } - - auto outputTokens = responses.at(0).getResult().outputTokenIds.at(0); - - TLLM_LOG_INFO("[DRAFT] Output tokens: %s", tlc::vec2str(outputTokens).c_str()); - - return responses.at(0).getResult(); -} - -tle::Result executeTargetRequest( - tle::Executor& executor, tle::Result const& draftResult, RuntimeOptions const& runtimeOpts) -{ - // Create the request - tle::SizeType32 maxNewTokens = runtimeOpts.numDraftTokens + 1; - tle::VecTokens inputTokens{1, 2, 3, 4}; - - tle::Request request{std::move(inputTokens), maxNewTokens}; - - tle::VecTokens const& outputTokenIds = draftResult.outputTokenIds.at(0); - tle::VecTokens draftTokens(outputTokenIds.end() - runtimeOpts.numDraftTokens, outputTokenIds.end()); - TLLM_LOG_INFO("[DRAFT] Draft tokens: %s", tlc::vec2str(draftTokens).c_str()); - - tle::Tensor logitsTensor; - - if (runtimeOpts.fastLogits) - { - auto const& logitsInfo = draftResult.specDecFastLogitsInfo.value(); - logitsTensor = logitsInfo.toTensor(); - } - else - { - auto generationLogits = draftResult.generationLogits.value(); - auto logitsShape = generationLogits.getShape(); - TLLM_CHECK(logitsShape[0] == 1); - logitsTensor = tle::Tensor::cpu(generationLogits.getDataType(), {logitsShape[1], logitsShape[2]}); - std::memcpy(logitsTensor.getData(), generationLogits.getData(), generationLogits.getSizeInBytes()); - } - - tle::ExternalDraftTokensConfig draftTokensConfig( - std::move(draftTokens), logitsTensor, std::nullopt /* acceptance threshold */, runtimeOpts.fastLogits); - request.setExternalDraftTokensConfig(draftTokensConfig); - - // Enqueue the request - auto requestId = executor.enqueueRequest(std::move(request)); - - // Wait for the response - auto responses = executor.awaitResponses(requestId); - - if (responses.at(0).hasError()) - { - TLLM_LOG_ERROR(responses.at(0).getErrorMsg()); - exit(1); - } - - auto outputTokens = responses.at(0).getResult().outputTokenIds.at(0); - - TLLM_LOG_INFO("[TARGET] Output tokens: %s", tlc::vec2str(outputTokens).c_str()); - - return responses.at(0).getResult(); -} diff --git a/examples/cpp/executor/executorExampleKvEvents.cpp b/examples/cpp/executor/executorExampleKvEvents.cpp deleted file mode 100644 index ea1923294382..000000000000 --- a/examples/cpp/executor/executorExampleKvEvents.cpp +++ /dev/null @@ -1,341 +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. - */ -#include -#include -#include -#include -#include -#include - -#include "tensorrt_llm/common/logger.h" -#include "tensorrt_llm/executor/executor.h" -#include "tensorrt_llm/plugins/api/tllmPlugin.h" -#include - -namespace tlc = tensorrt_llm::common; -namespace tle = tensorrt_llm::executor; -namespace fs = std::filesystem; - -struct RuntimeOptions -{ - std::string trtEnginePath; - tle::SizeType32 numSysPrompts; - - tle::SizeType32 sysPromptTokens; - tle::SizeType32 contextTokens; - - tle::SizeType32 maxTokensMean; - tle::SizeType32 maxTokensStddev; - - tle::SizeType32 numRequests; - - size_t hostCacheSize; - size_t maxTokensInPagedKvCache; -}; - -struct KVCacheBlock -{ - KVCacheBlock(size_t hash, int cacheLevel, int priority, std::optional loraId = std::nullopt, - std::shared_ptr prevBlock = nullptr, std::optional cacheSalt = std::nullopt); - - size_t hash; - int cacheLevel; - int priority; - - std::optional loraId; - std::optional cacheSalt; - - std::shared_ptr prevBlock; - std::unordered_map> nextBlocks; -}; - -class RadixTree -{ -public: - explicit RadixTree(tle::Executor& executor); - // Check the executor for new events. - void pollEvents(); - -private: - std::shared_ptr mCacheEventManager; - // The root block of the radix tree - std::shared_ptr root; - // A table mapping block hashes to their pointers - std::unordered_map> blockTable; - // Event counter - size_t eventCounter; -}; - -// Utility function to parse input arguments -RuntimeOptions parseArgs(int argc, char* argv[]); - -// Create a tle::Request -tle::Request makeRequest(int sysPromptTokens, int contextTokens, std::uniform_int_distribution sysPromptSelector, - std::normal_distribution maxNumTokensSelector); - -std::default_random_engine gen; - -int main(int argc, char* argv[]) -{ - // Register the TRT-LLM plugins - initTrtLlmPlugins(); - - auto runtimeOpts = parseArgs(argc, argv); - - // Create the executor for this engine - auto executorConfig = tle::ExecutorConfig(1); // Beam width 1 is required for cache block reuse - auto kvCacheConfig = tle::KvCacheConfig(true, - runtimeOpts.maxTokensInPagedKvCache ? std::optional(runtimeOpts.maxTokensInPagedKvCache) - : std::nullopt); // Enable cache block reuse - kvCacheConfig.setHostCacheSize(runtimeOpts.hostCacheSize); - kvCacheConfig.setEventBufferMaxSize(32768); - executorConfig.setKvCacheConfig(kvCacheConfig); - - auto executor = tle::Executor(runtimeOpts.trtEnginePath, tle::ModelType::kDECODER_ONLY, executorConfig); - - auto radixTree = RadixTree(executor); - - auto activeRequests = runtimeOpts.numRequests; - - std::uniform_int_distribution sysPromptSelector( - 1, runtimeOpts.numSysPrompts); // Select a system prompt between 1 and `runtimeOpts.numSysPrompts` - std::normal_distribution maxNumTokensSelector(runtimeOpts.maxTokensMean, runtimeOpts.maxTokensStddev); - - // Create and enqueue the requests - for (int i = 0; i < runtimeOpts.numRequests; i++) - { - std::ignore = executor.enqueueRequest(makeRequest( - runtimeOpts.sysPromptTokens, runtimeOpts.contextTokens, sysPromptSelector, maxNumTokensSelector)); - } - - while (activeRequests > 0) - { - auto responses = executor.awaitResponses(std::chrono::milliseconds(20)); - for (auto const& response : responses) - { - if (response.getResult().isFinal) - activeRequests--; - } - // Only call pollEvents once every 20ms. Events are only added to the queue once per iteration, so no need to - // poll faster than this. - radixTree.pollEvents(); - } - - return 0; -} - -RuntimeOptions parseArgs(int argc, char* argv[]) -{ - RuntimeOptions runtimeOpts; - - cxxopts::Options options(argv[0], "Example that demonstrates how to use the ExecutorKVCacheManager API"); - options.add_options()("h,help", "Print usage"); - options.add_options()("engine_dir", "Directory that store the engines.", cxxopts::value()); - options.add_options()("num_sys_prompts", "Amount of unique simulated system prompts to use", - cxxopts::value()->default_value("10")); - options.add_options()( - "sys_prompt_tokens", "Size of the simulated system prompts", cxxopts::value()->default_value("256")); - options.add_options()("context_tokens", "Amount of varying context tokens coming after the system prompts", - cxxopts::value()->default_value("128")); - options.add_options()( - "max_tokens_mean", "Mean number of max output tokens", cxxopts::value()->default_value("128")); - options.add_options()( - "max_tokens_stddev", "Standard deviation of max output tokens", cxxopts::value()->default_value("32")); - options.add_options()( - "num_requests", "Amount of requests to send to the engine", cxxopts::value()->default_value("100")); - options.add_options()("host_cache_size", "Size of the KV Cache in host memory in bytes", - cxxopts::value()->default_value("0")); - options.add_options()("max_tokens_in_paged_kv_cache", "Amount of tokens in the kv cache", - cxxopts::value()->default_value("0")); - - auto parsedOptions = options.parse(argc, argv); - - // Argument: help - if (parsedOptions.count("help")) - { - TLLM_LOG_ERROR(options.help()); - exit(0); - } - - // Argument: Engine directory - if (!parsedOptions.count("engine_dir")) - { - TLLM_LOG_ERROR(options.help()); - TLLM_LOG_ERROR("Please specify engine directory."); - exit(1); - } - runtimeOpts.trtEnginePath = parsedOptions["engine_dir"].as(); - if (!fs::exists(runtimeOpts.trtEnginePath) || !fs::is_directory(runtimeOpts.trtEnginePath)) - { - TLLM_LOG_ERROR("Engine directory doesn't exist."); - exit(1); - } - - runtimeOpts.numSysPrompts = parsedOptions["num_sys_prompts"].as(); - runtimeOpts.sysPromptTokens = parsedOptions["sys_prompt_tokens"].as(); - runtimeOpts.contextTokens = parsedOptions["context_tokens"].as(); - runtimeOpts.maxTokensMean = parsedOptions["max_tokens_mean"].as(); - runtimeOpts.maxTokensStddev = parsedOptions["max_tokens_stddev"].as(); - runtimeOpts.numRequests = parsedOptions["num_requests"].as(); - runtimeOpts.hostCacheSize = parsedOptions["host_cache_size"].as(); - runtimeOpts.maxTokensInPagedKvCache = parsedOptions["max_tokens_in_paged_kv_cache"].as(); - - return runtimeOpts; -} - -KVCacheBlock::KVCacheBlock(size_t hash, int cacheLevel, int priority, std::optional loraId, - std::shared_ptr prevBlock, std::optional cacheSalt) - : hash{hash} - , cacheLevel{cacheLevel} - , priority{priority} - , loraId{loraId} - , cacheSalt{std::move(cacheSalt)} - , prevBlock{prevBlock} - , nextBlocks{} -{ -} - -RadixTree::RadixTree(tle::Executor& executor) - : mCacheEventManager(*executor.getKVCacheEventManager()) - , eventCounter{1} -{ - // Use id=-1 for the root block. Doesn't matter what exact id is used, just that it is unique. - root = std::make_shared(-1, -1, -1); - blockTable[-1] = root; - - // Wait for the `CREATED` event to be emitted. - while (true) - { - auto events = mCacheEventManager->getLatestEvents(); - if (events.size() == 1) - { - auto const& eventData = std::get(events.front().data); - TLLM_LOG_INFO("Event ID %d: KV Cache Manager initialized with blocks per level of: %s", - events.front().eventId, tlc::vec2str(eventData.numBlocksPerCacheLevel).c_str()); - break; - } - } -}; - -void RadixTree::pollEvents() -{ - auto events = mCacheEventManager->getLatestEvents(std::chrono::milliseconds(20)); - for (tle::KVCacheEvent const& event : events) - { - TLLM_CHECK(event.eventId == eventCounter++); - if (std::holds_alternative(event.data)) - { - // Blocks have been stored into the radix tree - auto const& eventData = std::get(event.data); - auto prevBlock = blockTable[eventData.parentHash.value_or(-1)]; - - // This block should be in the tree - TLLM_CHECK(blockTable.find(prevBlock->hash) != blockTable.end()); - - for (auto& block : eventData.blocks) - { - - TLLM_LOG_INFO("Event ID %d: Block %04x was inserted into the radix tree with parent %04x.", - event.eventId, block.blockHash, prevBlock->hash); - - // This block shouldn't already exist in the tree, and should have tokens associated with it - TLLM_CHECK(blockTable.find(block.blockHash) == blockTable.end()); - TLLM_CHECK(block.tokens.size() > 0); - - auto thisBlock = std::make_shared( - block.blockHash, block.cacheLevel, block.priority, block.loraId, prevBlock, block.cacheSalt); - - blockTable[block.blockHash] = thisBlock; - // Link the parent to the new block - prevBlock->nextBlocks[block.blockHash] = thisBlock; - - prevBlock = thisBlock; - } - } - else if (std::holds_alternative(event.data)) - { - auto const& eventData = std::get(event.data); - - for (auto const& hash : eventData.blockHashes) - { - - TLLM_LOG_INFO("Event ID %d: Block %04x was removed from the radix tree.", event.eventId, hash); - - // This block should exist in the tree - TLLM_CHECK(blockTable.find(hash) != blockTable.end()); - - auto& block = blockTable[hash]; - - // Check that the block has no children, and that the parent has the block listed as a child - TLLM_CHECK(block->nextBlocks.size() == 0); - TLLM_CHECK(block->prevBlock->nextBlocks.find(block->hash) != block->prevBlock->nextBlocks.end()); - - // Remove the block from it's parent, and remove the entry in the block table - block->prevBlock->nextBlocks.erase(block->hash); - blockTable.erase(hash); - } - } - else if (std::holds_alternative(event.data)) - { - auto const& eventData = std::get(event.data); - - if (eventData.priority.has_value()) - { - // The block priority was updated - TLLM_LOG_INFO("Event ID %d: Block %04x priority was changed from %d to %d", event.eventId, - eventData.blockHash, eventData.priority->oldValue, eventData.priority->newValue); - - TLLM_CHECK(blockTable[eventData.blockHash]->priority == eventData.priority->oldValue); - blockTable[eventData.blockHash]->priority = eventData.priority->newValue; - } - - if (eventData.cacheLevel.has_value()) - { - // The block cache level was updated - TLLM_LOG_INFO("Event ID %d: Block %04x cache level was changed from %d to %d", event.eventId, - eventData.blockHash, eventData.cacheLevel->oldValue, eventData.cacheLevel->newValue); - - TLLM_CHECK(blockTable[eventData.blockHash]->cacheLevel == eventData.cacheLevel->oldValue); - blockTable[eventData.blockHash]->cacheLevel = eventData.cacheLevel->newValue; - } - } - else - { - TLLM_LOG_ERROR("Unsupported event type. This shouldn't happen!"); - } - } -} - -tle::Request makeRequest(int sysPromptTokens, int contextTokens, std::uniform_int_distribution sysPromptSelector, - std::normal_distribution maxNumTokensSelector) -{ - int sysPromptVersion = sysPromptSelector(gen); - tle::VecTokens inputTokens; - - // Add `sysPromptTokens` tokens. Add the version to the token ids to create a unique system prompt - for (int i = 0; i < sysPromptTokens; i++) - { - inputTokens.emplace_back(sysPromptVersion + i); - } - // Add random context tokens - for (int i = 0; i < contextTokens; i++) - { - inputTokens.emplace_back(rand() % 1000); - } - - return tle::Request(inputTokens, maxNumTokensSelector(gen)); -} diff --git a/examples/cpp/executor/executorExampleLogitsProcessor.cpp b/examples/cpp/executor/executorExampleLogitsProcessor.cpp deleted file mode 100644 index 0913b77b1775..000000000000 --- a/examples/cpp/executor/executorExampleLogitsProcessor.cpp +++ /dev/null @@ -1,91 +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. - */ -#include - -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/logger.h" -#include "tensorrt_llm/executor/executor.h" -#include "tensorrt_llm/plugins/api/tllmPlugin.h" - -namespace tlc = tensorrt_llm::common; -namespace tle = tensorrt_llm::executor; - -int main(int argc, char* argv[]) -{ - // Register the TRT-LLM plugins - initTrtLlmPlugins(); - - if (argc != 2) - { - TLLM_LOG_ERROR("Usage: %s ", argv[0]); - return 1; - } - - int constexpr sentinels[] = {42, 29}; - int step = 0; - - auto logitsPostProcessorFn - = [&step, &sentinels](tle::IdType reqId, tle::Tensor& logits, tle::BeamTokens const& tokens, - tle::StreamPtr const& streamPtr, std::optional clientId) - { - auto logitsDataType = logits.getDataType(); - auto logitsCpu = tensorrt_llm::executor::Tensor::cpu(logitsDataType, logits.getShape()); - auto* dataPtr = logitsCpu.getData(); - auto* dataPtrFloat = static_cast(dataPtr); - for (size_t i = 0; i < logitsCpu.getSize(); ++i) - { - dataPtrFloat[i] = -1.0e20; - } - dataPtrFloat[sentinels[step]] = 0.0f; - - logits.setFrom(logitsCpu, streamPtr); - step = (1 - step); - }; - - std::string logitsPostProcessorName = "MyLogitsPP"; - - // Create the executor for this engine - tle::SizeType32 beamWidth = 1; - auto executorConfig = tle::ExecutorConfig(beamWidth); - - auto logitsProcConfig = tle::LogitsPostProcessorConfig(); - logitsProcConfig.setProcessorMap(std::unordered_map{ - {logitsPostProcessorName, logitsPostProcessorFn}}); - executorConfig.setLogitsPostProcessorConfig(logitsProcConfig); - - auto trtEnginePath = argv[1]; - auto executor = tle::Executor(trtEnginePath, tle::ModelType::kDECODER_ONLY, executorConfig); - - // Create the request - tle::SizeType32 maxNewTokens = 5; - tle::VecTokens inputTokens{1, 2, 3, 4}; - auto request = tle::Request(inputTokens, maxNewTokens); - request.setLogitsPostProcessorName(logitsPostProcessorName); - - // Enqueue the request - auto requestId = executor.enqueueRequest(std::move(request)); - - // Wait for the response - auto responses = executor.awaitResponses(requestId); - - // Get outputTokens - auto outputTokens = responses.at(0).getResult().outputTokenIds.at(beamWidth - 1); - - TLLM_LOG_INFO("Output tokens: %s", tlc::vec2str(outputTokens).c_str()); - - return 0; -} diff --git a/examples/cpp/executor/inputTokens.csv b/examples/cpp/executor/inputTokens.csv deleted file mode 100644 index 4cb3974a91b5..000000000000 --- a/examples/cpp/executor/inputTokens.csv +++ /dev/null @@ -1,3 +0,0 @@ -1, 2, 3, 4, 5, 6 -1, 2, 3, 4 -1, 2, 3, 4, 5, 6, 7, 8, 9, 10 diff --git a/examples/disaggregated/slurm/benchmark/disaggr_torch.slurm b/examples/disaggregated/slurm/benchmark/disaggr_torch.slurm index 275016c255dd..d73430ba5820 100644 --- a/examples/disaggregated/slurm/benchmark/disaggr_torch.slurm +++ b/examples/disaggregated/slurm/benchmark/disaggr_torch.slurm @@ -101,7 +101,7 @@ elif [ -d "${trtllm_repo}" ]; then if [ "${build_wheel}" = "true" ]; then echo "Building TensorRT-LLM wheel on one node..." - build_command="python3 ./scripts/build_wheel.py --trt_root /usr/local/tensorrt --benchmarks --use_ccache --clean" + build_command="python3 ./scripts/build_wheel.py --trt_root /usr/local/tensorrt --use_ccache --clean" if [ -n "${cuda_architectures:-}" ]; then build_command="${build_command} --cuda_architectures \"${cuda_architectures}\"" fi diff --git a/examples/disaggregated/slurm/benchmark/disaggr_torch_dwdp.slurm b/examples/disaggregated/slurm/benchmark/disaggr_torch_dwdp.slurm index 8835399d6c50..07015e6473a0 100644 --- a/examples/disaggregated/slurm/benchmark/disaggr_torch_dwdp.slurm +++ b/examples/disaggregated/slurm/benchmark/disaggr_torch_dwdp.slurm @@ -101,7 +101,7 @@ elif [ -d "${trtllm_repo}" ]; then if [ "${build_wheel}" = "true" ]; then echo "Building TensorRT-LLM wheel on one node..." - build_command="python3 ./scripts/build_wheel.py --trt_root /usr/local/tensorrt --benchmarks --use_ccache --clean" + build_command="python3 ./scripts/build_wheel.py --trt_root /usr/local/tensorrt --use_ccache --clean" if [ -n "${cuda_architectures:-}" ]; then build_command="${build_command} --cuda_architectures \"${cuda_architectures}\"" fi diff --git a/examples/layer_wise_benchmarks/sample_performance_alignment.sh b/examples/layer_wise_benchmarks/sample_performance_alignment.sh index 812fe46f0024..4b9bf883c3c9 100755 --- a/examples/layer_wise_benchmarks/sample_performance_alignment.sh +++ b/examples/layer_wise_benchmarks/sample_performance_alignment.sh @@ -14,7 +14,7 @@ export TLLM_AUTOTUNER_CACHE_PATH="$PROFILE_DIR/sample_performance_alignment_cach mkdir -p -- "$PROFILE_DIR" mkdir -p -- "$(dirname -- "$TLLM_AUTOTUNER_CACHE_PATH")" -python3 ../../benchmarks/cpp/prepare_dataset.py \ +python3 ../../benchmarks/prepare_dataset.py \ --tokenizer "$MODEL" \ --stdout \ --random-seed 42 \ diff --git a/jenkins/Build.groovy b/jenkins/Build.groovy index caadaca639b3..da963453cd52 100644 --- a/jenkins/Build.groovy +++ b/jenkins/Build.groovy @@ -421,24 +421,12 @@ def runLLMBuild(pipeline, buildFlags, tarName, is_linux_x86_64) def buildJobs = buildFlags[BUILD_JOBS_FOR_CONFIG] ?: BUILD_JOBS withCredentials([usernamePassword(credentialsId: "urm-artifactory-creds", usernameVariable: 'CONAN_LOGIN_USERNAME', passwordVariable: 'CONAN_PASSWORD')]) { - sh "cd ${LLM_ROOT} && python3 scripts/build_wheel.py --use_ccache -G Ninja -j ${buildJobs} -a '${buildFlags[WHEEL_ARCHS]}' ${buildFlags[WHEEL_EXTRA_ARGS]} --benchmarks" + sh "cd ${LLM_ROOT} && python3 scripts/build_wheel.py --use_ccache -G Ninja -j ${buildJobs} -a '${buildFlags[WHEEL_ARCHS]}' ${buildFlags[WHEEL_EXTRA_ARGS]}" } - if (is_linux_x86_64) { - sh "cd ${LLM_ROOT} && python3 scripts/build_cpp_examples.py" - } - // Step 3: packaging wheels into tarfile sh "cp ${LLM_ROOT}/build/tensorrt_llm-*.whl TensorRT-LLM/" - // Step 4: packaging benchmark and required cpp dependencies into tarfile - sh "mkdir -p TensorRT-LLM/benchmarks/cpp" - sh "cp ${LLM_ROOT}/cpp/build/benchmarks/bertBenchmark TensorRT-LLM/benchmarks/cpp" - sh "cp ${LLM_ROOT}/cpp/build/benchmarks/gptManagerBenchmark TensorRT-LLM/benchmarks/cpp" - sh "cp ${LLM_ROOT}/cpp/build/benchmarks/disaggServerBenchmark TensorRT-LLM/benchmarks/cpp" - sh "cp ${LLM_ROOT}/cpp/build/tensorrt_llm/libtensorrt_llm.so TensorRT-LLM/benchmarks/cpp" - sh "cp ${LLM_ROOT}/cpp/build/tensorrt_llm/plugins/libnvinfer_plugin_tensorrt_llm.so TensorRT-LLM/benchmarks/cpp" - - // Step 5: packaging attribution files into tarfile when they exist + // Step 4: packaging attribution files into tarfile when they exist sh "mkdir -p TensorRT-LLM/attribution" sh "cp ${LLM_ROOT}/cpp/build/attribution/missing_files.json TensorRT-LLM/attribution/ || true" sh "cp ${LLM_ROOT}/cpp/build/attribution/import_payload.json TensorRT-LLM/attribution/ || true" @@ -571,7 +559,7 @@ def launchStages(pipeline, cpu_arch, enableFailFast, globalVars) stage(key) { stage("[${key}] Run") { echoNodeAndGpuInfo(pipeline, key) - buildWheelInContainer(pipeline, [], X86_64_TRIPLE, false, false, "cp312", "-a '90-real' -b Debug --benchmarks --micro_benchmarks") + buildWheelInContainer(pipeline, [], X86_64_TRIPLE, false, false, "cp312", "-a '90-real' -b Debug --micro_benchmarks") } } }) diff --git a/jenkins/L0_MergeRequest.groovy b/jenkins/L0_MergeRequest.groovy index 7ecfc1274865..0d445268d2b5 100644 --- a/jenkins/L0_MergeRequest.groovy +++ b/jenkins/L0_MergeRequest.groovy @@ -954,7 +954,6 @@ def getMultiGpuFileChanged(pipeline, testFilter, globalVars) "cpp/include/tensorrt_llm/runtime/worldConfig.h", "cpp/tensorrt_llm/batch_manager/", "cpp/tensorrt_llm/executor/", - "cpp/tensorrt_llm/executor_worker/", "cpp/tensorrt_llm/kernels/communicationKernels/", "cpp/tensorrt_llm/kernels/customAllReduceKernels.cu", "cpp/tensorrt_llm/kernels/customAllReduceKernels.h", @@ -969,13 +968,6 @@ def getMultiGpuFileChanged(pipeline, testFilter, globalVars) "cpp/tensorrt_llm/kernels/userbuffers/", "cpp/tensorrt_llm/kernels/xqaDispatcher.cpp", "cpp/tensorrt_llm/kernels/xqaDispatcher.h", - "cpp/tensorrt_llm/plugins/cpSplitPlugin/cpSplitPlugin.cpp", - "cpp/tensorrt_llm/plugins/cpSplitPlugin/cpSplitPlugin.h", - "cpp/tensorrt_llm/plugins/gptAttentionCommon/gptAttentionCommon.cpp", - "cpp/tensorrt_llm/plugins/gptAttentionCommon/gptAttentionCommon.h", - "cpp/tensorrt_llm/plugins/gptAttentionPlugin/gptAttentionPlugin.cpp", - "cpp/tensorrt_llm/plugins/gptAttentionPlugin/gptAttentionPlugin.h", - "cpp/tensorrt_llm/plugins/ncclPlugin/", "cpp/tensorrt_llm/nanobind/", "cpp/tensorrt_llm/runtime/ipcUtils.cpp", "cpp/tensorrt_llm/runtime/ncclCommunicator.cpp", @@ -985,8 +977,6 @@ def getMultiGpuFileChanged(pipeline, testFilter, globalVars) "cpp/tensorrt_llm/thop/allgatherOp.cpp", "cpp/tensorrt_llm/thop/allreduceOp.cpp", "cpp/tensorrt_llm/thop/reducescatterOp.cpp", - "cpp/tests/e2e_tests/batch_manager/", - "cpp/tests/e2e_tests/executor/", "cpp/tests/unit_tests/multi_gpu/", "jenkins/L0_Test.groovy", "tensorrt_llm/_ipc_utils.py", diff --git a/jenkins/L0_Test.groovy b/jenkins/L0_Test.groovy index e2a43dcc92ac..14ba33b94f59 100644 --- a/jenkins/L0_Test.groovy +++ b/jenkins/L0_Test.groovy @@ -4732,7 +4732,6 @@ def launchTestJobs(pipeline, testFilter) "A10-TensorRT-1": ["a10", "l0_a10", 1, 1], "A30-PyTorch-1": ["a30", "l0_a30", 1, 2], "A30-PyTorch-2": ["a30", "l0_a30", 2, 2], - "A10-CPP-1": ["a10", "l0_a10", 1, 1], "A30-CPP-1": ["a30", "l0_a30", 1, 1], "A30-AutoDeploy-1": ["a30", "l0_a30", 1, 1], "A100X-PyTorch-1": ["a100x", "l0_a100", 1, 1], @@ -4775,7 +4774,6 @@ def launchTestJobs(pipeline, testFilter) // "L40S-TensorRT-Post-Merge-5": ["l40s", "l0_l40s", 5, 5], "L40S-FMHA-Post-Merge-1": ["l40s", "l0_l40s", 1, 1], "H100_PCIe-AutoDeploy-Post-Merge-1": ["h100-cr", "l0_h100", 1, 1], - "H100_PCIe-CPP-Post-Merge-1": ["h100-cr", "l0_h100", 1, 1], // "H100_PCIe-TensorRT-Post-Merge-1": ["h100-cr", "l0_h100", 1, 5], // "H100_PCIe-TensorRT-Post-Merge-2": ["h100-cr", "l0_h100", 2, 5], // "H100_PCIe-TensorRT-Post-Merge-3": ["h100-cr", "l0_h100", 3, 5], diff --git a/jenkins/scripts/perf/local/slurm_install.sh b/jenkins/scripts/perf/local/slurm_install.sh index 91bb7d664e62..2524c7c8ac79 100755 --- a/jenkins/scripts/perf/local/slurm_install.sh +++ b/jenkins/scripts/perf/local/slurm_install.sh @@ -21,7 +21,7 @@ slurm_build_wheel() { fi echo "Building wheel on node ${SLURM_NODEID:-0}, task ${SLURM_LOCALID:-0}" - retry_command bash -c "cd $llmSrcNode && rm -rf .venv-3.12 && python3 ./scripts/build_wheel.py --trt_root /usr/local/tensorrt --benchmarks --use_ccache --cuda_architectures '100-real' --clean -c" + retry_command bash -c "cd $llmSrcNode && rm -rf .venv-3.12 && python3 ./scripts/build_wheel.py --use_ccache --cuda_architectures '100-real' --clean -c" cd $jobWorkspace echo "(Writing build wheel lock) Lock file: $build_lock_file" diff --git a/legacy-files.txt b/legacy-files.txt index d6c8b23308aa..de8f7f433ec0 100644 --- a/legacy-files.txt +++ b/legacy-files.txt @@ -1,14 +1,14 @@ .devcontainer/make_env.py .github/scripts/label_community_user.py .github/scripts/pr_checklist_check.py -benchmarks/cpp/__init__.py -benchmarks/cpp/prepare_dataset.py -benchmarks/cpp/utils/__init__.py -benchmarks/cpp/utils/convert_nemo_dataset.py -benchmarks/cpp/utils/generate_rand_loras.py -benchmarks/cpp/utils/prepare_real_data.py -benchmarks/cpp/utils/prepare_synthetic_data.py -benchmarks/cpp/utils/utils.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 @@ -33,38 +33,11 @@ cpp/micro_benchmarks/gen-moe-benchmark-file.py cpp/tensorrt_llm/deep_ep/strip_nvshmem_helper.py cpp/tensorrt_llm/kernels/cutlass_kernels/python/generate_kernels.py cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/copy_cu.py -cpp/tests/resources/scripts/build_chatglm_engines.py -cpp/tests/resources/scripts/build_eagle_engines.py -cpp/tests/resources/scripts/build_enc_dec_engines.py -cpp/tests/resources/scripts/build_engines_utils.py -cpp/tests/resources/scripts/build_gpt_engines.py -cpp/tests/resources/scripts/build_gptj_engines.py -cpp/tests/resources/scripts/build_llama_engines.py -cpp/tests/resources/scripts/build_mamba_engines.py -cpp/tests/resources/scripts/build_medusa_engines.py -cpp/tests/resources/scripts/build_recurrentgemma_engines.py -cpp/tests/resources/scripts/build_redrafter_engines.py -cpp/tests/resources/scripts/generate_expected_chatglm_output.py -cpp/tests/resources/scripts/generate_expected_eagle_output.py -cpp/tests/resources/scripts/generate_expected_enc_dec_output.py -cpp/tests/resources/scripts/generate_expected_gpt_output.py -cpp/tests/resources/scripts/generate_expected_gptj_output.py -cpp/tests/resources/scripts/generate_expected_llama_output.py -cpp/tests/resources/scripts/generate_expected_mamba_output.py -cpp/tests/resources/scripts/generate_expected_medusa_output.py -cpp/tests/resources/scripts/generate_expected_recurrentgemma_output.py -cpp/tests/resources/scripts/generate_expected_redrafter_output.py -cpp/tests/resources/scripts/generate_hf_gpt_output.py cpp/tests/resources/scripts/generate_test_lora_weights.py -cpp/tests/resources/scripts/io_converter.py docs/source/conf.py docs/source/helper.py examples/apps/chat.py examples/apps/fastapi_server.py -examples/bindings/executor/example_advanced.py -examples/bindings/executor/example_basic.py -examples/bindings/executor/example_debug.py -examples/bindings/executor/example_logits_processor.py examples/disaggregated/clients/disagg_client.py examples/disaggregated/slurm/benchmark/submit.py examples/dora/normalize_weights.py @@ -238,7 +211,6 @@ examples/wide_ep/slurm_scripts/process_gen_iterlog.py jenkins/scripts/mergeWaiveList.py jenkins/scripts/open_search_db.py jenkins/scripts/test_rerun.py -scripts/build_cpp_examples.py scripts/build_wheel.py scripts/check_test_list.py scripts/dco_check.py @@ -843,11 +815,8 @@ tests/integration/defs/common.py tests/integration/defs/conftest.py tests/integration/defs/cpp/conftest.py tests/integration/defs/cpp/cpp_common.py -tests/integration/defs/cpp/test_e2e.py tests/integration/defs/cpp/test_multi_gpu.py tests/integration/defs/cpp/test_unit_tests.py -tests/integration/defs/deterministic/mixtral_deterministic.py -tests/integration/defs/deterministic/test_mixtral_deterministic.py tests/integration/defs/disaggregated/test_auto_scaling.py tests/integration/defs/disaggregated/test_disaggregated.py tests/integration/defs/disaggregated/test_disaggregated_etcd.py @@ -920,7 +889,6 @@ tests/integration/defs/test_e2e.py tests/integration/defs/test_fmha.py tests/integration/defs/test_list_parser.py tests/integration/defs/test_list_validation.py -tests/integration/defs/test_mlpf_results.py tests/integration/defs/test_sanity.py tests/integration/defs/test_unittests.py tests/integration/defs/triton_server/__init__.py @@ -1073,7 +1041,6 @@ tests/unittest/_torch/thop/serial/test_moe.py tests/unittest/_torch/thop/serial/test_moe_alltoall.py tests/unittest/api_stability/api_stability_core.py tests/unittest/api_stability/test_llm_api.py -tests/unittest/bindings/binding_test_utils.py tests/unittest/bindings/test_bindings_moe.py tests/unittest/bindings/test_bindings_ut.py tests/unittest/bindings/test_executor_bindings.py @@ -1161,7 +1128,6 @@ 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_leak.py tests/unittest/others/test_mapping.py tests/unittest/others/test_model_dtype.py tests/unittest/others/test_module.py diff --git a/pyproject.toml b/pyproject.toml index 8110c926318e..49c5584d5680 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -58,14 +58,14 @@ exclude = [ ".devcontainer/make_env.py", ".github/scripts/label_community_user.py", ".github/scripts/pr_checklist_check.py", - "benchmarks/cpp/__init__.py", - "benchmarks/cpp/prepare_dataset.py", - "benchmarks/cpp/utils/__init__.py", - "benchmarks/cpp/utils/convert_nemo_dataset.py", - "benchmarks/cpp/utils/generate_rand_loras.py", - "benchmarks/cpp/utils/prepare_real_data.py", - "benchmarks/cpp/utils/prepare_synthetic_data.py", - "benchmarks/cpp/utils/utils.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", @@ -90,38 +90,11 @@ exclude = [ "cpp/tensorrt_llm/deep_ep/strip_nvshmem_helper.py", "cpp/tensorrt_llm/kernels/cutlass_kernels/python/generate_kernels.py", "cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/copy_cu.py", - "cpp/tests/resources/scripts/build_chatglm_engines.py", - "cpp/tests/resources/scripts/build_eagle_engines.py", - "cpp/tests/resources/scripts/build_enc_dec_engines.py", - "cpp/tests/resources/scripts/build_engines_utils.py", - "cpp/tests/resources/scripts/build_gpt_engines.py", - "cpp/tests/resources/scripts/build_gptj_engines.py", - "cpp/tests/resources/scripts/build_llama_engines.py", - "cpp/tests/resources/scripts/build_mamba_engines.py", - "cpp/tests/resources/scripts/build_medusa_engines.py", - "cpp/tests/resources/scripts/build_recurrentgemma_engines.py", - "cpp/tests/resources/scripts/build_redrafter_engines.py", - "cpp/tests/resources/scripts/generate_expected_chatglm_output.py", - "cpp/tests/resources/scripts/generate_expected_eagle_output.py", - "cpp/tests/resources/scripts/generate_expected_enc_dec_output.py", - "cpp/tests/resources/scripts/generate_expected_gpt_output.py", - "cpp/tests/resources/scripts/generate_expected_gptj_output.py", - "cpp/tests/resources/scripts/generate_expected_llama_output.py", - "cpp/tests/resources/scripts/generate_expected_mamba_output.py", - "cpp/tests/resources/scripts/generate_expected_medusa_output.py", - "cpp/tests/resources/scripts/generate_expected_recurrentgemma_output.py", - "cpp/tests/resources/scripts/generate_expected_redrafter_output.py", - "cpp/tests/resources/scripts/generate_hf_gpt_output.py", "cpp/tests/resources/scripts/generate_test_lora_weights.py", - "cpp/tests/resources/scripts/io_converter.py", "docs/source/conf.py", "docs/source/helper.py", "examples/apps/chat.py", "examples/apps/fastapi_server.py", - "examples/bindings/executor/example_advanced.py", - "examples/bindings/executor/example_basic.py", - "examples/bindings/executor/example_debug.py", - "examples/bindings/executor/example_logits_processor.py", "examples/disaggregated/clients/disagg_client.py", "examples/disaggregated/slurm/benchmark/submit.py", "examples/dora/normalize_weights.py", @@ -295,7 +268,6 @@ exclude = [ "jenkins/scripts/mergeWaiveList.py", "jenkins/scripts/open_search_db.py", "jenkins/scripts/test_rerun.py", - "scripts/build_cpp_examples.py", "scripts/build_wheel.py", "scripts/check_test_list.py", "scripts/dco_check.py", @@ -900,11 +872,8 @@ exclude = [ "tests/integration/defs/conftest.py", "tests/integration/defs/cpp/conftest.py", "tests/integration/defs/cpp/cpp_common.py", - "tests/integration/defs/cpp/test_e2e.py", "tests/integration/defs/cpp/test_multi_gpu.py", "tests/integration/defs/cpp/test_unit_tests.py", - "tests/integration/defs/deterministic/mixtral_deterministic.py", - "tests/integration/defs/deterministic/test_mixtral_deterministic.py", "tests/integration/defs/disaggregated/test_auto_scaling.py", "tests/integration/defs/disaggregated/test_disaggregated.py", "tests/integration/defs/disaggregated/test_disaggregated_etcd.py", @@ -977,7 +946,6 @@ exclude = [ "tests/integration/defs/test_fmha.py", "tests/integration/defs/test_list_parser.py", "tests/integration/defs/test_list_validation.py", - "tests/integration/defs/test_mlpf_results.py", "tests/integration/defs/test_sanity.py", "tests/integration/defs/test_unittests.py", "tests/integration/defs/triton_server/__init__.py", @@ -1130,7 +1098,6 @@ exclude = [ "tests/unittest/_torch/thop/serial/test_moe_alltoall.py", "tests/unittest/api_stability/api_stability_core.py", "tests/unittest/api_stability/test_llm_api.py", - "tests/unittest/bindings/binding_test_utils.py", "tests/unittest/bindings/test_bindings_moe.py", "tests/unittest/bindings/test_bindings_ut.py", "tests/unittest/bindings/test_executor_bindings.py", @@ -1218,7 +1185,6 @@ exclude = [ "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_leak.py", "tests/unittest/others/test_mapping.py", "tests/unittest/others/test_model_dtype.py", "tests/unittest/others/test_module.py", diff --git a/requirements.txt b/requirements.txt index dab577614f57..99d03e56cbc2 100644 --- a/requirements.txt +++ b/requirements.txt @@ -22,7 +22,6 @@ pandas h5py==3.12.1 StrEnum sentencepiece>=0.1.99 -tensorrt~=10.16.1 # https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes/rel-26-05.html#rel-26-05 uses 2.12.0a0. torch>=2.11.0,<=2.13.0a0 torchvision diff --git a/ruff-legacy-baseline.json b/ruff-legacy-baseline.json index ee17a68362c2..9a7bf5d452b3 100644 --- a/ruff-legacy-baseline.json +++ b/ruff-legacy-baseline.json @@ -1,16 +1,16 @@ { "_meta": { "generated_by": "scripts/legacy_utils.py lint-update-violations", - "total_violations": 2917, - "total_files": 303 + "total_violations": 2358, + "total_files": 271 }, ".github/scripts/label_community_user.py": { "D212": 1 }, - "benchmarks/cpp/utils/convert_nemo_dataset.py": { + "benchmarks/utils/convert_nemo_dataset.py": { "E741": 1 }, - "benchmarks/cpp/utils/prepare_real_data.py": { + "benchmarks/utils/prepare_real_data.py": { "D202": 1, "D205": 1, "D410": 7, @@ -80,14 +80,6 @@ "D212": 1, "D300": 1 }, - "examples/bindings/executor/example_advanced.py": { - "E402": 1 - }, - "examples/bindings/executor/example_logits_processor.py": { - "D200": 1, - "D212": 1, - "D415": 1 - }, "examples/dora/normalize_weights.py": { "D200": 3, "D202": 1, @@ -257,9 +249,6 @@ "D212": 1, "E741": 1 }, - "scripts/check_test_list.py": { - "D212": 1 - }, "scripts/dco_check.py": { "D212": 1 }, @@ -292,39 +281,21 @@ "D212": 1, "E402": 20 }, - "tensorrt_llm/_torch/attention_backend/sparse/dsa.py": { - "F821": 2 - }, "tensorrt_llm/_torch/attention_backend/sparse/kernel.py": { "E731": 3 }, "tensorrt_llm/_torch/attention_backend/sparse/rocket.py": { "E712": 2 }, - "tensorrt_llm/_torch/attention_backend/sparse/utils.py": { - "F821": 4 - }, "tensorrt_llm/_torch/attention_backend/trtllm.py": { "E712": 1 }, - "tensorrt_llm/_torch/attention_backend/utils.py": { - "F821": 2 - }, - "tensorrt_llm/_torch/compilation/patterns/ar_residual_norm.py": { - "E402": 1 - }, "tensorrt_llm/_torch/compilation/patterns/residual_add_norm.py": { "E402": 1 }, "tensorrt_llm/_torch/compilation/piecewise_optimizer.py": { "E711": 2 }, - "tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py": { - "E741": 11 - }, - "tensorrt_llm/_torch/custom_ops/torch_custom_ops.py": { - "F821": 2 - }, "tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dense_blockscaled_gemm_persistent.py": { "E741": 3 }, @@ -332,10 +303,6 @@ "E712": 1, "E731": 1 }, - "tensorrt_llm/_torch/model_config.py": { - "F811": 1, - "F821": 4 - }, "tensorrt_llm/_torch/models/checkpoints/auto_mapper.py": { "F821": 1 }, @@ -366,12 +333,6 @@ "tensorrt_llm/_torch/models/modeling_nemotron.py": { "E731": 1 }, - "tensorrt_llm/_torch/models/modeling_qwen2vl.py": { - "F811": 1 - }, - "tensorrt_llm/_torch/models/modeling_qwen3_next.py": { - "F821": 1 - }, "tensorrt_llm/_torch/models/modeling_siglip.py": { "F811": 1 }, @@ -379,9 +340,6 @@ "E731": 1, "F821": 5 }, - "tensorrt_llm/_torch/modules/attention.py": { - "E731": 2 - }, "tensorrt_llm/_torch/modules/fused_moe/deep_ep_utils.py": { "F821": 2 }, @@ -415,18 +373,6 @@ "tensorrt_llm/_torch/modules/mamba/ssd_state_passing.py": { "E731": 1 }, - "tensorrt_llm/_torch/pyexecutor/_util.py": { - "F811": 1 - }, - "tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py": { - "F821": 1 - }, - "tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py": { - "F821": 1 - }, - "tensorrt_llm/_torch/pyexecutor/model_engine.py": { - "F821": 1 - }, "tensorrt_llm/_torch/pyexecutor/model_loader.py": { "F811": 1 }, @@ -434,24 +380,15 @@ "E712": 1, "E721": 1 }, - "tensorrt_llm/_torch/pyexecutor/resource_manager.py": { - "F821": 2 - }, "tensorrt_llm/_torch/pyexecutor/seq_slot_manager.py": { "F821": 1 }, "tensorrt_llm/_torch/speculative/auto_heuristic.py": { "F821": 1 }, - "tensorrt_llm/_torch/speculative/interface.py": { - "F821": 1 - }, "tensorrt_llm/_torch/speculative/ngram.py": { "E741": 1 }, - "tensorrt_llm/_torch/speculative/speculation_gate.py": { - "F821": 1 - }, "tensorrt_llm/_utils.py": { "E722": 1, "F821": 1 @@ -535,9 +472,7 @@ "D202": 1, "D205": 1, "D208": 3, - "D210": 1, "D212": 1, - "D300": 1, "E722": 1 }, "tensorrt_llm/executor/ray_executor.py": { @@ -578,7 +513,6 @@ }, "tensorrt_llm/executor/rpc_proxy.py": { "D205": 1, - "D212": 1, "D415": 1 }, "tensorrt_llm/executor/rpc_worker.py": { @@ -589,24 +523,6 @@ "D300": 6, "F811": 1 }, - "tensorrt_llm/functional.py": { - "D200": 42, - "D202": 6, - "D205": 20, - "D207": 1, - "D208": 6, - "D210": 2, - "D212": 142, - "D214": 3, - "D300": 144, - "D301": 2, - "D411": 8, - "D415": 9 - }, - "tensorrt_llm/inputs/evs.py": { - "D205": 1, - "D212": 2 - }, "tensorrt_llm/inputs/multimodal.py": { "D415": 1 }, @@ -657,11 +573,7 @@ }, "tensorrt_llm/llmapi/mpi_session.py": { "D200": 1, - "D205": 2, - "D210": 7, - "D212": 3, - "D300": 10, - "D411": 1, + "D205": 1, "D415": 1 }, "tensorrt_llm/llmapi/reasoning_parser.py": { @@ -703,25 +615,6 @@ "E731": 1, "F821": 1 }, - "tensorrt_llm/models/modeling_utils.py": { - "D200": 3, - "D202": 1, - "D205": 4, - "D208": 4, - "D210": 3, - "D212": 2, - "D300": 5, - "D415": 4, - "E721": 2, - "E722": 1 - }, - "tensorrt_llm/quantization/functional.py": { - "D205": 4, - "D212": 4, - "D300": 4, - "D411": 2, - "D415": 2 - }, "tensorrt_llm/quantization/quantize_by_modelopt.py": { "D200": 1, "D205": 1, @@ -797,9 +690,6 @@ "tensorrt_llm/scaffolding/task.py": { "F821": 1 }, - "tensorrt_llm/scaffolding/task_collection.py": { - "F821": 1 - }, "tensorrt_llm/scaffolding/worker.py": { "PLE0302": 1 }, @@ -818,11 +708,6 @@ "D210": 1, "F811": 1 }, - "tensorrt_llm/serve/openai_server.py": { - "D205": 1, - "D212": 2, - "F821": 3 - }, "tensorrt_llm/serve/responses_utils.py": { "D200": 2, "D205": 4, @@ -896,20 +781,10 @@ "D403": 9, "D415": 11 }, - "tests/integration/defs/accuracy/accuracy_core.py": { - "D205": 1, - "D212": 1 - }, "tests/integration/defs/accuracy/test_disaggregated_serving.py": { "D212": 1, "F601": 1 }, - "tests/integration/defs/accuracy/test_llm_api_autodeploy.py": { - "D205": 1, - "D209": 1, - "D415": 1, - "F811": 2 - }, "tests/integration/defs/accuracy/test_llm_api_pytorch.py": { "D212": 1, "D300": 3, @@ -1159,16 +1034,6 @@ "E741": 2, "F811": 1 }, - "tests/integration/defs/test_mlpf_results.py": { - "D200": 1, - "D205": 1, - "D208": 4, - "D212": 2, - "D415": 2 - }, - "tests/integration/defs/triton_server/common.py": { - "E402": 1 - }, "tests/integration/defs/triton_server/conftest.py": { "D200": 3, "D202": 1, @@ -1215,16 +1080,6 @@ "tests/integration/defs/utils/timeout_manager.py": { "D212": 9 }, - "tests/microbenchmarks/build_time_benchmark.py": { - "D200": 1, - "D300": 1, - "D415": 1 - }, - "tests/microbenchmarks/build_time_dashboard.py": { - "D200": 1, - "D300": 1, - "D415": 1 - }, "tests/scripts/allreduce_perf/allreduce_heuristic_code_gen.py": { "D212": 1, "E712": 1 @@ -1245,11 +1100,7 @@ "D415": 6 }, "tests/unittest/bindings/test_executor_bindings.py": { - "D202": 1, - "E712": 54, - "F403": 2, - "F405": 3, - "F811": 1 + "E712": 51 }, "tests/unittest/conftest.py": { "D200": 3, @@ -1329,10 +1180,6 @@ "D205": 1, "D212": 1 }, - "tests/unittest/llmapi/test_llm_utils.py": { - "F403": 2, - "F405": 14 - }, "tests/unittest/llmapi/test_mpi_session.py": { "D200": 1, "D212": 1, @@ -1345,12 +1192,6 @@ "tests/unittest/others/test_kv_cache_transceiver.py": { "D212": 1 }, - "tests/unittest/others/test_leak.py": { - "D200": 1, - "D210": 1, - "D300": 1, - "D415": 1 - }, "tests/unittest/others/test_time_breakdown.py": { "D212": 1, "D415": 1 diff --git a/ruff-legacy.toml b/ruff-legacy.toml index 2f9d1bd00745..3aadfee4149c 100644 --- a/ruff-legacy.toml +++ b/ruff-legacy.toml @@ -18,14 +18,14 @@ include = [ ".devcontainer/make_env.py", ".github/scripts/label_community_user.py", ".github/scripts/pr_checklist_check.py", - "benchmarks/cpp/__init__.py", - "benchmarks/cpp/prepare_dataset.py", - "benchmarks/cpp/utils/__init__.py", - "benchmarks/cpp/utils/convert_nemo_dataset.py", - "benchmarks/cpp/utils/generate_rand_loras.py", - "benchmarks/cpp/utils/prepare_real_data.py", - "benchmarks/cpp/utils/prepare_synthetic_data.py", - "benchmarks/cpp/utils/utils.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", @@ -50,38 +50,11 @@ include = [ "cpp/tensorrt_llm/deep_ep/strip_nvshmem_helper.py", "cpp/tensorrt_llm/kernels/cutlass_kernels/python/generate_kernels.py", "cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/copy_cu.py", - "cpp/tests/resources/scripts/build_chatglm_engines.py", - "cpp/tests/resources/scripts/build_eagle_engines.py", - "cpp/tests/resources/scripts/build_enc_dec_engines.py", - "cpp/tests/resources/scripts/build_engines_utils.py", - "cpp/tests/resources/scripts/build_gpt_engines.py", - "cpp/tests/resources/scripts/build_gptj_engines.py", - "cpp/tests/resources/scripts/build_llama_engines.py", - "cpp/tests/resources/scripts/build_mamba_engines.py", - "cpp/tests/resources/scripts/build_medusa_engines.py", - "cpp/tests/resources/scripts/build_recurrentgemma_engines.py", - "cpp/tests/resources/scripts/build_redrafter_engines.py", - "cpp/tests/resources/scripts/generate_expected_chatglm_output.py", - "cpp/tests/resources/scripts/generate_expected_eagle_output.py", - "cpp/tests/resources/scripts/generate_expected_enc_dec_output.py", - "cpp/tests/resources/scripts/generate_expected_gpt_output.py", - "cpp/tests/resources/scripts/generate_expected_gptj_output.py", - "cpp/tests/resources/scripts/generate_expected_llama_output.py", - "cpp/tests/resources/scripts/generate_expected_mamba_output.py", - "cpp/tests/resources/scripts/generate_expected_medusa_output.py", - "cpp/tests/resources/scripts/generate_expected_recurrentgemma_output.py", - "cpp/tests/resources/scripts/generate_expected_redrafter_output.py", - "cpp/tests/resources/scripts/generate_hf_gpt_output.py", "cpp/tests/resources/scripts/generate_test_lora_weights.py", - "cpp/tests/resources/scripts/io_converter.py", "docs/source/conf.py", "docs/source/helper.py", "examples/apps/chat.py", "examples/apps/fastapi_server.py", - "examples/bindings/executor/example_advanced.py", - "examples/bindings/executor/example_basic.py", - "examples/bindings/executor/example_debug.py", - "examples/bindings/executor/example_logits_processor.py", "examples/disaggregated/clients/disagg_client.py", "examples/disaggregated/slurm/benchmark/submit.py", "examples/dora/normalize_weights.py", @@ -255,7 +228,6 @@ include = [ "jenkins/scripts/mergeWaiveList.py", "jenkins/scripts/open_search_db.py", "jenkins/scripts/test_rerun.py", - "scripts/build_cpp_examples.py", "scripts/build_wheel.py", "scripts/check_test_list.py", "scripts/dco_check.py", @@ -860,11 +832,8 @@ include = [ "tests/integration/defs/conftest.py", "tests/integration/defs/cpp/conftest.py", "tests/integration/defs/cpp/cpp_common.py", - "tests/integration/defs/cpp/test_e2e.py", "tests/integration/defs/cpp/test_multi_gpu.py", "tests/integration/defs/cpp/test_unit_tests.py", - "tests/integration/defs/deterministic/mixtral_deterministic.py", - "tests/integration/defs/deterministic/test_mixtral_deterministic.py", "tests/integration/defs/disaggregated/test_auto_scaling.py", "tests/integration/defs/disaggregated/test_disaggregated.py", "tests/integration/defs/disaggregated/test_disaggregated_etcd.py", @@ -937,7 +906,6 @@ include = [ "tests/integration/defs/test_fmha.py", "tests/integration/defs/test_list_parser.py", "tests/integration/defs/test_list_validation.py", - "tests/integration/defs/test_mlpf_results.py", "tests/integration/defs/test_sanity.py", "tests/integration/defs/test_unittests.py", "tests/integration/defs/triton_server/__init__.py", @@ -1090,7 +1058,6 @@ include = [ "tests/unittest/_torch/thop/serial/test_moe_alltoall.py", "tests/unittest/api_stability/api_stability_core.py", "tests/unittest/api_stability/test_llm_api.py", - "tests/unittest/bindings/binding_test_utils.py", "tests/unittest/bindings/test_bindings_moe.py", "tests/unittest/bindings/test_bindings_ut.py", "tests/unittest/bindings/test_executor_bindings.py", @@ -1178,7 +1145,6 @@ include = [ "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_leak.py", "tests/unittest/others/test_mapping.py", "tests/unittest/others/test_model_dtype.py", "tests/unittest/others/test_module.py", diff --git a/scripts/build_cpp_examples.py b/scripts/build_cpp_examples.py deleted file mode 100644 index cb2591acfea2..000000000000 --- a/scripts/build_cpp_examples.py +++ /dev/null @@ -1,88 +0,0 @@ -import argparse -import contextlib -import logging -import os -import platform -import shutil -import subprocess -from os import PathLike -from pathlib import Path - - -@contextlib.contextmanager -def working_directory(path: PathLike): - """Changes working directory and returns to previous on exit.""" - prev_cwd = Path.cwd() - os.chdir(path) - try: - yield - finally: - os.chdir(prev_cwd) - - -def build_cpp_examples(build_dir: PathLike, trt_dir: PathLike, - enable_multi_device: str, loglevel: int) -> None: - logging.basicConfig(level=loglevel, - format='%(asctime)s - %(levelname)s - %(message)s') - # Convert input paths to pathlib.Path objects - build_dir = Path(build_dir) - trt_dir = Path(trt_dir) - - assert trt_dir.is_dir() - - def cmake_parse(path: PathLike) -> str: - return str(path).replace("\\", "/") - - # Remove the build directory if it exists - if build_dir.exists(): - logging.info(f"Removed directory: {build_dir}") - shutil.rmtree(build_dir) - - # Create the build directory - build_dir.mkdir(parents=True, exist_ok=True) - - # Change to the build directory - with working_directory(build_dir): - # Run CMake with the specified TensorRT directories - generator = ["-GNinja"] if platform.system() == "Windows" else [] - generate_command = [ - 'cmake', - '-S', - '..', - '-B', - '.', - f'-DTensorRT_ROOT={cmake_parse(trt_dir)}', - f'-DENABLE_MULTI_DEVICE={enable_multi_device}', - ] + generator - logging.info(f"Executing {generate_command}") - subprocess.run(generate_command, check=True) - - # Build the project using make - build_command = ["cmake", "--build", ".", "--config", "Release"] - logging.info(f"Executing {build_command}") - subprocess.run(build_command, check=True) - - -if __name__ == '__main__': - parser = argparse.ArgumentParser(description='Build C++ examples') - parser.add_argument('--build-dir', - default='examples/cpp/executor/build', - help='Build directory path') - parser.add_argument('--trt-dir', - default='/usr/local/tensorrt', - help='TensorRT directory path') - parser.add_argument('--enable-multi-device', - default='ON', - help='Enable multi device support (requires MPI)') - parser.add_argument('-v', - '--verbose', - help="verbose", - action="store_const", - dest="loglevel", - const=logging.DEBUG, - default=logging.INFO) - cli = parser.parse_args() - - args = vars(cli) - print(args) # Log on Jenkins instance. - build_cpp_examples(**args) diff --git a/scripts/build_wheel.py b/scripts/build_wheel.py index 263f061d14b6..7beb8b88f829 100755 --- a/scripts/build_wheel.py +++ b/scripts/build_wheel.py @@ -489,7 +489,7 @@ def main(*, job_count: int = None, extra_cmake_vars: Sequence[str] = tuple(), extra_make_targets: str = "", - trt_root: str = '/usr/local/tensorrt', + trt_root: str = None, nccl_root: str = None, nixl_root: str = None, mooncake_root: str = None, @@ -505,7 +505,6 @@ def main(*, install: bool = False, skip_building_wheel: bool = False, linking_install_binary: bool = False, - benchmarks: bool = False, micro_benchmarks: bool = False, nvtx: bool = False, skip_stubs: bool = False, @@ -542,21 +541,6 @@ def main(*, no_venv, yes=yes) - # Ensure base TRT is installed (check inside the venv) - try: - check_output([str(venv_python), "-m", "pip", "show", "tensorrt"]) - except CalledProcessError: - error_msg = "TensorRT was not installed properly." - if on_windows: - error_msg += ( - " Please download the TensorRT zip file manually," - " install it and relaunch build_wheel.py." - " See https://docs.nvidia.com/deeplearning/tensorrt/install-guide/index.html#installing-zip for more details." - ) - else: - error_msg += f" Please install tensorrt into the venv using \"`{venv_python}` -m pip install tensorrt\" and relaunch build_wheel.py" - raise RuntimeError(error_msg) - if cuda_architectures is not None: if "70-real" in cuda_architectures: raise RuntimeError("Volta architecture is deprecated support.") @@ -670,7 +654,7 @@ def main(*, "-- BOLT: Forcing NVRTC_DYNAMIC_LINKING=ON (static NVIDIA libs lack relocations)" ) - targets = ["tensorrt_llm", "nvinfer_plugin_tensorrt_llm"] + targets = ["tensorrt_llm"] if cpp_only: build_pyt = "OFF" @@ -687,9 +671,6 @@ def main(*, build_deep_gemm = "ON" build_flash_mla = "ON" - if benchmarks: - targets.append("benchmarks") - if micro_benchmarks: targets.append("micro_benchmarks") build_micro_benchmarks = "ON" @@ -698,9 +679,6 @@ def main(*, disable_nvtx = "OFF" if nvtx else "ON" - if not on_windows: - targets.append("executorWorker") - source_dir = get_source_dir() fmha_v2_cu_dir = project_dir / "cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fmha_v2_cu" @@ -932,18 +910,11 @@ def copy_resolving_symlink(src_path, dst_path): lib_dir / "tensorrt_llm.dll") install_file(build_dir / f"tensorrt_llm/thop/th_common.dll", lib_dir / "th_common.dll") - install_file( - build_dir / f"tensorrt_llm/plugins/nvinfer_plugin_tensorrt_llm.dll", - lib_dir / "nvinfer_plugin_tensorrt_llm.dll") else: install_file(build_dir / "tensorrt_llm/libtensorrt_llm.so", lib_dir / "libtensorrt_llm.so") install_file(build_dir / "tensorrt_llm/thop/libth_common.so", lib_dir / "libth_common.so") - install_file( - build_dir / - "tensorrt_llm/plugins/libnvinfer_plugin_tensorrt_llm.so", - lib_dir / "libnvinfer_plugin_tensorrt_llm.so") if os.path.exists( build_dir / "tensorrt_llm/executor/cache_transmission/ucx_utils/libtensorrt_llm_ucx_wrapper.so" @@ -1028,15 +999,6 @@ def copy_resolving_symlink(src_path, dst_path): clear_folder(deep_gemm_dir) deep_gemm_dir.rmdir() - bin_dir = pkg_dir / "bin" - if bin_dir.exists(): - clear_folder(bin_dir) - bin_dir.mkdir(parents=True, exist_ok=True) - - if not on_windows: - install_file(build_dir / "tensorrt_llm/executor_worker/executorWorker", - bin_dir / "executorWorker") - scripts_dir = pkg_dir / "scripts" if scripts_dir.exists(): clear_folder(scripts_dir) @@ -1276,7 +1238,9 @@ def add_arguments(parser: ArgumentParser): parser.add_argument( "--trt_root", default="/usr/local/tensorrt", - help="Directory containing TensorRT headers and libraries") + help="[DEPRECATED] No effect: TensorRT is no longer required to build. " + "Accepted for backward compatibility and will be removed in a future release." + ) parser.add_argument("--nccl_root", help="Directory containing NCCL headers and libraries") parser.add_argument("--nixl_root", @@ -1313,9 +1277,6 @@ def add_arguments(parser: ArgumentParser): help= "Install the built binary by creating symbolic links instead of copying files" ) - parser.add_argument("--benchmarks", - action="store_true", - help="Build the benchmarks for the C++ runtime") parser.add_argument("--micro_benchmarks", action="store_true", help="Build the micro benchmarks for C++ components") diff --git a/scripts/get_wheel_from_package.py b/scripts/get_wheel_from_package.py index cb604482c27b..f8dc16652361 100644 --- a/scripts/get_wheel_from_package.py +++ b/scripts/get_wheel_from_package.py @@ -78,27 +78,11 @@ def get_wheel_from_package(arch, artifact_path, timeout): build_dir = llm_root / "build" build_dir.mkdir(parents=True, exist_ok=True) - benchmarks_dir = llm_root / "cpp" / "build" / "benchmarks" - benchmarks_dir.mkdir(parents=True, exist_ok=True) - wheel_files = glob.glob(str(tmp_dir / "tensorrt_llm*.whl")) for wheel_file in wheel_files: shutil.move(wheel_file, str(build_dir)) print(f"Moved wheel file: {wheel_file} -> {build_dir}") - benchmark_files = [ - "bertBenchmark", "gptManagerBenchmark", "disaggServerBenchmark" - ] - - for benchmark in benchmark_files: - src_path = tmp_dir / "benchmarks" / "cpp" / benchmark - if src_path.exists(): - dst_path = benchmarks_dir / benchmark - shutil.copy2(src_path, dst_path) - print(f"Copied benchmark file: {src_path} -> {dst_path}") - else: - print(f"Warning: Benchmark file not found: {src_path}") - shutil.rmtree(tmp_dir) if os.path.exists(tarfile_name): diff --git a/setup.py b/setup.py index 42b130f73fb0..4724f6d86b8c 100644 --- a/setup.py +++ b/setup.py @@ -137,15 +137,13 @@ def has_ext_modules(self): if on_windows: package_data = [ - 'libs/th_common.dll', 'libs/tensorrt_llm.dll', - 'libs/nvinfer_plugin_tensorrt_llm.dll', 'bindings.*.pyd', "include/**/*" + 'libs/th_common.dll', 'libs/tensorrt_llm.dll', 'bindings.*.pyd', + "include/**/*" ] else: package_data = [ - 'bin/executorWorker', 'libs/libtensorrt_llm.so', 'libs/libth_common.so', - 'libs/libnvinfer_plugin_tensorrt_llm.so', 'libs/libtensorrt_llm_ucx_wrapper.so', 'libs/libdecoder_attention_0.so', 'libs/libtensorrt_llm_nixl_wrapper.so', diff --git a/tensorrt_llm/__init__.py b/tensorrt_llm/__init__.py index 91e5f98a500d..d707ca0e0694 100644 --- a/tensorrt_llm/__init__.py +++ b/tensorrt_llm/__init__.py @@ -104,24 +104,6 @@ def _setup_vendored_triton_kernels(): # ImportError: libc10.so: cannot open shared object file: No such file or directory import torch # noqa - -def _preload_tensorrt_libs(): - """Preload the TensorRT libraries needed by the bindings extension. - - The C++ runtime still links against the TensorRT libraries until it is - decoupled from TensorRT. Importing the tensorrt package loads libnvinfer - from the tensorrt_libs wheel for environments where it is not on the - system loader path; without it, importing tensorrt_llm.bindings raises - ImportError: libnvinfer.so.10: cannot open shared object file. - """ - try: - import tensorrt # noqa: F401 - except ImportError: - pass - - -_preload_tensorrt_libs() - import tensorrt_llm._torch.models as torch_models import tensorrt_llm.math_utils as math_utils import tensorrt_llm.models as models diff --git a/tensorrt_llm/_common.py b/tensorrt_llm/_common.py index 00bac35cc11b..d7feeededb9f 100644 --- a/tensorrt_llm/_common.py +++ b/tensorrt_llm/_common.py @@ -18,7 +18,6 @@ import platform import threading import time -from functools import wraps from pathlib import Path import torch @@ -89,27 +88,3 @@ def _print_stacks(): print_stacks_thread.start() logger.info("TensorRT LLM inited.") - - -# TODO: dead on the Python side (no remaining @_is_building users); the IS_BUILDING -# env var is only read by the C++ isBuilding() in the TensorRT plugins. Remove this -# together with that C++ half in the C++ decouple step. -class _BuildingFlag: - def __enter__(self): - os.environ["IS_BUILDING"] = "1" - - def __exit__(self, type, value, tb): - del os.environ["IS_BUILDING"] - - -def _is_building(f): - """Use this to decorate functions which are called during engine building/refitting process, - otherwise, the plugin registration will fail. - """ - - @wraps(f) - def decorated(*args, **kwargs): - with _BuildingFlag(): - return f(*args, **kwargs) - - return decorated diff --git a/tensorrt_llm/executor/base_worker.py b/tensorrt_llm/executor/base_worker.py index 086ff957d54b..0294bf64534d 100644 --- a/tensorrt_llm/executor/base_worker.py +++ b/tensorrt_llm/executor/base_worker.py @@ -207,24 +207,8 @@ def _create_py_executor(): self.max_seq_len = _executor.max_seq_len return _executor - def _create_engine(executor_config): - engine = self._engine - if executor_config is None: - executor_config = tllm.ExecutorConfig(1) - executor_config.logits_post_processor_config = tllm.LogitsPostProcessorConfig( - processor_batched=self._batched_logits_processor, - replicate=False) - comm_ranks, device_ids = self._get_comm_ranks_device_id() - executor_config.parallel_config = tllm.ParallelConfig( - participant_ids=comm_ranks, device_ids=device_ids) - - assert not hasattr(executor_config, "backend") - return tllm.Executor(engine, tllm.ModelType.DECODER_ONLY, - executor_config) - - self.engine = _create_py_executor( - ) if self.llm_args is not None else _create_engine( - self._executor_config) + assert self.llm_args is not None, "llm_args is required to set up the worker engine" + self.engine = _create_py_executor() self._lora_manager: Optional[LoraManager] = None self._prompt_adapter_manager: Optional[PromptAdapterManager] = None @@ -245,16 +229,10 @@ def await_responses(self, timeout: Optional[float] = None) -> list: seconds=timeout) if timeout is not None else None) def fetch_stats(self) -> list: - if isinstance(self.engine, tllm.Executor): - iter_stats = self.engine.get_latest_iteration_stats() - #TODO: Support req stats with TRT engine - # This would require ensuring iter and req stats have same size - return [(iter_stat, None, None) for iter_stat in iter_stats] - else: - return self.engine.get_latest_iteration_stats() + return self.engine.get_latest_iteration_stats() def fetch_kv_cache_capacity(self) -> dict: - if self.engine is None or isinstance(self.engine, tllm.Executor): + if self.engine is None: return {} from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor @@ -264,10 +242,7 @@ def fetch_kv_cache_capacity(self) -> dict: return {} def fetch_kv_cache_events(self) -> list: - if isinstance(self.engine, tllm.Executor): - return self.engine.get_latest_kv_cache_events() - else: - return self.engine.get_latest_kv_cache_events() + return self.engine.get_latest_kv_cache_events() def set_result_queue(self, queue): """In multi-gpu mode, result_queue will be set here to communicate between the proxy and the worker 0 process.""" diff --git a/tensorrt_llm/executor/worker.py b/tensorrt_llm/executor/worker.py index 5958d296e1c4..342168664ef1 100644 --- a/tensorrt_llm/executor/worker.py +++ b/tensorrt_llm/executor/worker.py @@ -153,11 +153,6 @@ def shutdown(self): def block_subordinates(self): if self.rank != 0: - if isinstance(self.engine, tllm.Executor): - self.shutdown() - raise self.WorkerExit( - "block_subordinates() should be used in a `with GenerationExecutorWorker() as ...:` block" - ) from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor if isinstance(self.engine, PyExecutor): self.engine.wait_shutdown() diff --git a/tests/integration/README.md b/tests/integration/README.md index 8653989af4db..1f2f745bc10f 100644 --- a/tests/integration/README.md +++ b/tests/integration/README.md @@ -29,8 +29,6 @@ SM version mapping: - `sm100` = Blackwell (e.g., B100, B200) - `sm103` = Blackwell-Ultra (e.g., B300, GB300) -- To run perf tests, you also need to first build the cpp benchmark by calling `build_wheel.py` with `--benchmarks` flag. - ## Run perf tests All the perf test names are in the form of `perf/test_perf.py::test_perf[...]` where the `...` part is the test parameters. diff --git a/tests/integration/defs/.test_durations b/tests/integration/defs/.test_durations index 1805229baa7e..8d65874b1e3b 100644 --- a/tests/integration/defs/.test_durations +++ b/tests/integration/defs/.test_durations @@ -525,43 +525,8 @@ "accuracy/test_llm_api_pytorch_multimodal.py::TestQwen2_5_VL_7B::test_auto_dtype": 828.9645, "accuracy/test_llm_api_pytorch_multimodal.py::TestQwen3VL::test_auto_dtype[forced_chunked_prefill]": 238.39, "accuracy/test_llm_api_pytorch_multimodal.py::TestVILA1_5_3B::test_auto_dtype": 356.457, - "cpp/test_e2e.py::test_benchmarks[t5-90]": 580.63, - "cpp/test_e2e.py::test_model[-enc_dec_language_adapter-90]": 237.6715, - "cpp/test_e2e.py::test_model[-gpt_executor-80]": 1174.5125, - "cpp/test_e2e.py::test_model[-gpt_tests-80]": 705.1885, - "cpp/test_e2e.py::test_model[-medusa-86]": 545.2025, - "cpp/test_e2e.py::test_model[-t5-90]": 183.012, - "cpp/test_e2e.py::test_model[fp8-llama-90]": 373.718, - "cpp/test_multi_gpu.py::TestDisagg::test_asymmetric_executor[llama-4proc-mpi_kvcache-90]": 854.0535, - "cpp/test_multi_gpu.py::TestDisagg::test_asymmetric_executor[llama-4proc-nixl_kvcache-90]": 12.3445, - "cpp/test_multi_gpu.py::TestDisagg::test_asymmetric_executor[llama-4proc-ucx_kvcache-90]": 12.0945, - "cpp/test_multi_gpu.py::TestDisagg::test_asymmetric_executor[llama-6proc-mpi_kvcache-90]": 27.0285, - "cpp/test_multi_gpu.py::TestDisagg::test_asymmetric_executor[llama-6proc-nixl_kvcache-90]": 26.472, - "cpp/test_multi_gpu.py::TestDisagg::test_asymmetric_executor[llama-6proc-ucx_kvcache-90]": 26.258, - "cpp/test_multi_gpu.py::TestDisagg::test_asymmetric_executor[llama-8proc-mpi_kvcache-90]": 14.722, - "cpp/test_multi_gpu.py::TestDisagg::test_asymmetric_executor[llama-8proc-nixl_kvcache-90]": 15.325, - "cpp/test_multi_gpu.py::TestDisagg::test_asymmetric_executor[llama-8proc-ucx_kvcache-90]": 14.7605, - "cpp/test_multi_gpu.py::TestDisagg::test_orchestrator_params[llama-mpi_kvcache-90]": 18.438, - "cpp/test_multi_gpu.py::TestDisagg::test_orchestrator_params[llama-ucx_kvcache-90]": 18.755, - "cpp/test_multi_gpu.py::TestDisagg::test_spawn_orchestrator[llama-nixl_kvcache-90]": 40.666, - "cpp/test_multi_gpu.py::TestDisagg::test_spawn_orchestrator[llama-ucx_kvcache-90]": 40.163, - "cpp/test_multi_gpu.py::TestDisagg::test_symmetric_executor[llama-2proc-mpi_kvcache-90]": 6.9815, - "cpp/test_multi_gpu.py::TestDisagg::test_symmetric_executor[llama-2proc-nixl_kvcache-90]": 7.0665, - "cpp/test_multi_gpu.py::TestDisagg::test_symmetric_executor[llama-2proc-ucx_kvcache-90]": 6.9995, - "cpp/test_multi_gpu.py::TestDisagg::test_symmetric_executor[llama-4proc-mpi_kvcache-90]": 11.5145, - "cpp/test_multi_gpu.py::TestDisagg::test_symmetric_executor[llama-4proc-nixl_kvcache-90]": 11.65, - "cpp/test_multi_gpu.py::TestDisagg::test_symmetric_executor[llama-4proc-ucx_kvcache-90]": 11.75, - "cpp/test_multi_gpu.py::TestDisagg::test_symmetric_executor[llama-8proc-mpi_kvcache-90]": 14.1805, - "cpp/test_multi_gpu.py::TestDisagg::test_symmetric_executor[llama-8proc-nixl_kvcache-90]": 14.174, - "cpp/test_multi_gpu.py::TestDisagg::test_symmetric_executor[llama-8proc-ucx_kvcache-90]": 14.1315, - "cpp/test_multi_gpu.py::test_enc_dec[t5-90]": 140.1025, "cpp/test_multi_gpu.py::test_fused_gemm_allreduce[4proc-90]": 21.5805, - "cpp/test_multi_gpu.py::test_llama_executor[llama-leader-90]": 619.3225, - "cpp/test_multi_gpu.py::test_llama_executor[llama-orchestrator-90]": 188.325, - "cpp/test_multi_gpu.py::test_llama_executor_guided_decoding[llama-90]": 15.448, - "cpp/test_multi_gpu.py::test_llama_executor_logits_proc[llama-90]": 46.4445, "cpp/test_multi_gpu.py::test_mpi_utils[90]": 47.964, - "cpp/test_multi_gpu.py::test_trt_gpt_real_decoder[llama-90]": 112.2465, "cpp/test_multi_gpu.py::test_user_buffer[2proc-90]": 4.2345, "cpp/test_unit_tests.py::test_unit_tests[common-80]": 257.673, "cpp/test_unit_tests.py::test_unit_tests[common-90]": 19.7405, @@ -573,8 +538,6 @@ "cpp/test_unit_tests.py::test_unit_tests[runtime-80]": 719.1235, "cpp/test_unit_tests.py::test_unit_tests[thop-80]": 3.441, "cpp/test_unit_tests.py::test_unit_tests[thop-90]": 3.1505, - "cpp/test_unit_tests.py::test_unit_tests[utils-80]": 4.5175, - "cpp/test_unit_tests.py::test_unit_tests[utils-90]": 4.088, "disaggregated/test_auto_scaling.py::test_disagg_server_restart[etcd-round_robin]": 115.3005, "disaggregated/test_auto_scaling.py::test_disagg_server_restart[http-round_robin]": 155.1545, "disaggregated/test_auto_scaling.py::test_minimal_instances[etcd-round_robin]": 118.212, diff --git a/tests/integration/defs/common.py b/tests/integration/defs/common.py index e40db78ac36c..03abc6fed72b 100644 --- a/tests/integration/defs/common.py +++ b/tests/integration/defs/common.py @@ -33,8 +33,8 @@ from tensorrt_llm.lora_manager import LoraConfig from tensorrt_llm.sampling_params import SamplingParams -from .trt_test_alternative import (check_call, check_output, exists, is_windows, - print_info, print_warning) +from .trt_test_alternative import (check_call, check_output, exists, print_info, + print_warning) def venv_check_call(venv, cmd, env=None, **kwargs): @@ -573,26 +573,6 @@ def generate_summary_cmd(example_root, *args, **kwargs): return summary_cmd -def generate_deterministic_cmd(example_root, *args, **kwargs): - "generate deterministic command" - deterministic_cmd = [ - f"{example_root}/mixtral_deterministic.py", - "--check_deterministic_accuracy" - ] - - for key, value in kwargs.items(): - if isinstance(value, bool): - if value: - deterministic_cmd.extend(f"--{key}") - else: - deterministic_cmd.extend([f"--{key}", f"{value}"]) - - for arg in args: - deterministic_cmd.append(f"--{arg}") - - return deterministic_cmd - - def quantize_data(llm_venv, example_root, model_dir, @@ -729,24 +709,6 @@ def run_and_check(llm_venv, run_cmd, valid_outputs, streaming=False): ]), f"output is: {output}" -def get_cpp_benchmark(cpp_benchmark_name, llm_root): - suffix = ".exe" if is_windows() else "" - cpp_benchmark_name += suffix - # In CI/CD, we copy the cpp binary into the same folder as cpp to avoid package sanity - ci_path = os.path.join(os.path.dirname(os.path.realpath(llm_root)), - "benchmarks", "cpp", cpp_benchmark_name) - if os.path.exists(ci_path): - return ci_path - # In QA, we keep the benchmark build at its original location - qa_path = os.path.join(llm_root, "cpp", "build", "benchmarks", - cpp_benchmark_name) - if os.path.exists(qa_path): - return qa_path - raise Exception( - f"Cannot find cpp benchmark binary in either {ci_path} or {qa_path}. Did you forget --benchmark in building TRT-LLM?" - ) - - def generate_dummy_loras( hf_model_dir, lora_output_dir, diff --git a/tests/integration/defs/conftest.py b/tests/integration/defs/conftest.py index 2ec7f4d81f11..43d50860bb5e 100644 --- a/tests/integration/defs/conftest.py +++ b/tests/integration/defs/conftest.py @@ -2587,15 +2587,6 @@ def pytest_runtest_protocol(item, nextitem): return ret -@pytest.fixture(scope="function") -def deterministic_test_root(llm_root, llm_venv): - "Get deterministic test root" - deterministic_root = os.path.join(llm_root, - "tests/integration/defs/deterministic") - - return deterministic_root - - @pytest.fixture(scope="function") def disaggregated_test_root(llm_root, llm_venv): "Get disaggregated test root" diff --git a/tests/integration/defs/cpp/conftest.py b/tests/integration/defs/cpp/conftest.py index cdcb8e7237e1..d9c54742e72e 100644 --- a/tests/integration/defs/cpp/conftest.py +++ b/tests/integration/defs/cpp/conftest.py @@ -3,7 +3,6 @@ import pathlib as _pl import shutil import sys as _sys -import time import defs.cpp.cpp_common as _cpp import pytest @@ -175,52 +174,6 @@ def build_google_tests(request, build_type): ) -@pytest.fixture(scope="session") -def build_benchmarks(build_google_tests, build_dir, build_type): - - make_benchmarks = [ - "cmake", - "--build", - ".", - "--config", - build_type, - "-j", - "--target", - "benchmarks", - ] - - _cpp.run_command(make_benchmarks, cwd=build_dir, timeout=300) - - -@pytest.fixture(scope="session") -def prepare_model( - root_dir, - cpp_resources_dir, - python_exe, - model_cache_arg, - install_additional_requirements, -): - - def _prepare(model_name: str, run_fp8=False): - install_additional_requirements(model_name) - - start_time = time.time() - - _cpp.prepare_model_tests( - model_name=model_name, - python_exe=python_exe, - root_dir=root_dir, - resources_dir=cpp_resources_dir, - model_cache_arg=model_cache_arg, - ) - - duration = time.time() - start_time - print(f"Built model: {model_name}") - print(f"Duration: {duration} seconds") - - return _prepare - - @pytest.fixture(scope="function", autouse=True) def keep_log_files(build_dir): """Backup previous cpp test results when run multiple ctest invocations.""" diff --git a/tests/integration/defs/cpp/cpp_common.py b/tests/integration/defs/cpp/cpp_common.py index e90c599d8003..b4017d2dd672 100755 --- a/tests/integration/defs/cpp/cpp_common.py +++ b/tests/integration/defs/cpp/cpp_common.py @@ -242,86 +242,3 @@ def produce_mpirun_command(*, global_commands, nranks, local_commands, l += ["-n", "1"] + local_commands + (leader_commands if rank == 0 else []) + [":"] return l[:-1] - - -def prepare_model_tests(model_name: str, - python_exe: str, - root_dir: _pl.Path, - resources_dir: _pl.Path, - model_cache_arg=[], - only_fp8_arg=[], - only_multi_gpu_arg=[]): - scripts_dir = resources_dir / "scripts" - - model_env = {**_os.environ, "PYTHONPATH": f"examples/{model_name}"} - enc_dec_model_name_arg = [] - beams_arg = [] - if model_name in ('bart', 't5', 'enc_dec_language_adapter'): - enc_dec_repo_name_dict = { - 'bart': 'facebook/bart-large-cnn', - 't5': 't5-small', - 'enc_dec_language_adapter': - 'language_adapter-enc_dec_language_adapter' - } - enc_dec_model_name_arg = [ - '--hf_repo_name', enc_dec_repo_name_dict[model_name] - ] - if model_name == 't5' and (not only_multi_gpu_arg): - beams_arg = ['--beams', '1,2'] - model_name = 'enc_dec' - - # share the same script for gpt related tests - if model_name == 'gpt_executor' or model_name == 'gpt_tests': - model_name = 'gpt' - - build_engines = [ - python_exe, - str(scripts_dir / f"build_{model_name}_engines.py") - ] + model_cache_arg + only_fp8_arg + only_multi_gpu_arg + enc_dec_model_name_arg + beams_arg - - if model_name in ['gpt']: - build_engines += ['--clean'] - run_command(build_engines, cwd=root_dir, env=model_env, timeout=1800) - - model_env["PYTHONPATH"] = "examples" - generate_expected_output = [ - python_exe, - str(scripts_dir / f"generate_expected_{model_name}_output.py") - ] + only_fp8_arg + only_multi_gpu_arg + enc_dec_model_name_arg - if "enc_dec" in model_name: - generate_expected_output += model_cache_arg - generate_expected_output += beams_arg - - if model_name in ['gpt']: - generate_expected_output += ['--clean'] - - if only_multi_gpu_arg and model_name != 'enc_dec': - for world_size in (2, 4): - generate_command = [ - "mpirun", "-n", - str(world_size), "--allow-run-as-root", "--timeout", "600" - ] + generate_expected_output - run_command(generate_command, - cwd=root_dir, - env=model_env, - timeout=600) - else: - run_command(generate_expected_output, - cwd=root_dir, - env=model_env, - timeout=600) - - if model_name in ['gpt', 'llama']: - if model_name == 'gpt': - script_model_name = 'gpt2' - elif model_name == 'llama': - script_model_name = 'Llama-3.2-1B' - generate_tokenizer_info = [ - python_exe, "examples/generate_xgrammar_tokenizer_info.py", - f"--model_dir={str(resources_dir / 'models' / script_model_name)}", - f"--output_dir={str(resources_dir / 'data' / script_model_name)}" - ] - run_command(generate_tokenizer_info, - cwd=root_dir, - env=model_env, - timeout=600) diff --git a/tests/integration/defs/cpp/test_e2e.py b/tests/integration/defs/cpp/test_e2e.py deleted file mode 100644 index ce583671cfd8..000000000000 --- a/tests/integration/defs/cpp/test_e2e.py +++ /dev/null @@ -1,328 +0,0 @@ -import copy -import logging as _logger -import os as _os -import pathlib as _pl -from dataclasses import dataclass -from typing import List, Optional - -import defs.cpp.cpp_common as _cpp -import pytest - - -@dataclass(frozen=True) -class DatasetConfig: - """Configuration for a benchmark dataset.""" - name: str - local_path: str - split: str - input_key: str - output_key: str - max_input_len: str - num_requests: str - config_name: Optional[str] = None - prompt: Optional[str] = None - prompt_key: Optional[str] = None - - @property - def token_file(self) -> str: - return "prepared_" + self.name.replace('/', '_') - - def get_dataset_args(self) -> dict[str, str]: - """Build the dataset args dict for prepare_dataset.py.""" - args = { - '--dataset-local-path': self.local_path, - '--dataset-split': self.split, - '--dataset-input-key': self.input_key, - '--dataset-output-key': self.output_key, - } - if self.config_name: - args['--dataset-config-name'] = self.config_name - if self.prompt: - args['--dataset-prompt'] = self.prompt - if self.prompt_key: - args['--dataset-prompt-key'] = self.prompt_key - return args - - -def get_benchmark_dataset_configs(model_cache: str) -> List[DatasetConfig]: - """Define dataset configurations for benchmark tests. - - To add a new dataset, add a new DatasetConfig entry to this list. - """ - datasets_dir = _pl.Path(model_cache) / "datasets" - - return [ - DatasetConfig( - name="ccdv/cnn_dailymail", - local_path=str(datasets_dir / "ccdv" / "cnn_dailymail"), - config_name="3.0.0", - split="validation", - input_key="article", - prompt="Summarize the following article:", - output_key="highlights", - max_input_len="256", - num_requests="50", - ), - DatasetConfig( - name="Open-Orca/1million-gpt-4", - local_path=str(datasets_dir / "Open-Orca" / "1million-gpt-4" / - "1M-GPT4-Augmented.parquet"), - split="train", - input_key="question", - prompt_key="system_prompt", - output_key="response", - max_input_len="20", - num_requests="10", - ), - ] - - -def run_single_gpu_tests(build_dir: _pl.Path, - test_list: List[str], - run_fp8=False, - timeout=3600): - - cpp_env = {**_os.environ} - tests_dir = build_dir / "tests" / "e2e_tests" - - included_tests = list(_cpp.generate_included_model_tests(test_list)) - - fname_list = list(_cpp.generate_result_file_name(test_list, - run_fp8=run_fp8)) - resultFileName = "-".join(fname_list) + ".xml" - - excluded_tests = ["FP8"] if not run_fp8 else [] - - excluded_tests.extend(list(_cpp.generate_excluded_test_list(test_list))) - - ctest = ["ctest", "--output-on-failure", "--output-junit", resultFileName] - - if included_tests: - ctest.extend(["-R", "|".join(included_tests)]) - if excluded_tests: - ctest.extend(["-E", "|".join(excluded_tests)]) - - parallel = _cpp.default_test_parallel - if parallel_override := _os.environ.get("LLM_TEST_PARALLEL_OVERRIDE", - None): - parallel = int(parallel_override) - - _cpp.parallel_run_ctest(ctest, - cwd=tests_dir, - env=cpp_env, - timeout=timeout, - parallel=parallel) - if "gpt" in test_list: - xml_output_file = build_dir / "results-single-gpu-disagg-executor_gpt.xml" - new_env = copy.copy(cpp_env) - new_env["TRTLLM_USE_MPI_KVCACHE"] = "1" - trt_model_test = _cpp.produce_mpirun_command( - global_commands=["mpirun", "--allow-run-as-root"], - nranks=2, - local_commands=[ - "executor/disaggExecutorTest", - "--gtest_filter=*GptSingleDeviceDisaggSymmetricExecutorTest*" - ], - leader_commands=[f"--gtest_output=xml:{xml_output_file}"]) - _cpp.run_command(trt_model_test, - cwd=tests_dir, - env=new_env, - timeout=timeout) - - run_spec_dec_tests(build_dir=build_dir) - - -def run_benchmarks( - model_name: str, - python_exe: str, - root_dir: _pl.Path, - build_dir: _pl.Path, - resources_dir: _pl.Path, - model_cache: str, - batching_types: list[str], - api_types: list[str], -): - benchmark_exe_dir = build_dir / "benchmarks" - if model_name == "gpt": - model_engine_dir = resources_dir / "models" / "rt_engine" / "gpt2" - tokenizer_dir = resources_dir / "models" / "gpt2" - elif model_name in ('bart', 't5'): - if model_name == "t5": - hf_repo_name = "t5-small" - elif model_name == "bart": - hf_repo_name = "bart-large-cnn" - model_engine_dir = resources_dir / "models" / "enc_dec" / "trt_engines" / hf_repo_name - tokenizer_dir = model_cache + "/" + hf_repo_name - model_engine_path = model_engine_dir / "1-gpu" / "float16" / "decoder" - encoder_model_engine_path = model_engine_dir / "1-gpu" / "float16" / "encoder" - model_name = "enc_dec" - else: - _logger.info( - f"run_benchmark test does not support {model_name}. Skipping benchmarks" - ) - return NotImplementedError - - if model_name == "gpt": - model_engine_path = model_engine_dir / "fp16_plugin_packed_paged" / "tp1-pp1-cp1-gpu" - - # WAR: Currently importing the bindings here causes a segfault in pybind 11 during shutdown - # As this just builds a path we hard-code for now to obviate the need for import of bindings - - # model_spec_obj = model_spec.ModelSpec(input_file, _tb.DataType.HALF) - # model_spec_obj.set_kv_cache_type(_tb.KVCacheType.PAGED) - # model_spec_obj.use_gpt_plugin() - # model_spec_obj.use_packed_input() - # model_engine_path = model_engine_dir / model_spec_obj.get_model_path( - # ) / "tp1-pp1-cp1-gpu" - - for config in get_benchmark_dataset_configs(model_cache): - benchmark_src_dir = _pl.Path("benchmarks") / "cpp" - data_dir = resources_dir / "data" - prepare_dataset = [ - python_exe, - str(benchmark_src_dir / "prepare_dataset.py"), "--tokenizer", - str(tokenizer_dir), "--output", - str(data_dir / config.token_file), "dataset", "--max-input-len", - config.max_input_len, "--num-requests", config.num_requests - ] - for k, v in config.get_dataset_args().items(): - prepare_dataset += [k, v] - - # Use environment variable to force HuggingFace to use offline cached dataset - offline_env = {**_os.environ, 'HF_DATASETS_OFFLINE': '1'} - _cpp.run_command(prepare_dataset, - cwd=root_dir, - timeout=300, - env=offline_env) - - for batching_type in batching_types: - for api_type in api_types: - benchmark = [ - str(benchmark_exe_dir / "gptManagerBenchmark"), - "--engine_dir", - str(model_engine_path), "--type", - str(batching_type), "--api", - str(api_type), "--dataset", - str(data_dir / config.token_file) - ] - if model_name == "enc_dec": - benchmark += [ - "--encoder_engine_dir", - str(encoder_model_engine_path) - ] - - _cpp.run_command(benchmark, cwd=root_dir, timeout=600) - req_rate_benchmark = benchmark + [ - "--request_rate", "100", "--enable_exp_delays" - ] - _cpp.run_command(req_rate_benchmark, cwd=root_dir, timeout=600) - concurrency_benchmark = benchmark + ["--concurrency", "30"] - _cpp.run_command(concurrency_benchmark, - cwd=root_dir, - timeout=600) - - if "IFB" in batching_types and "executor" in api_types: - # executor streaming test - benchmark = [ - str(benchmark_exe_dir / "gptManagerBenchmark"), "--engine_dir", - str(model_engine_path), "--type", "IFB", "--dataset", - str(data_dir / config.token_file), "--api", "executor", - "--streaming" - ] - if model_name == "enc_dec": - benchmark += [ - "--encoder_engine_dir", - str(encoder_model_engine_path) - ] - _cpp.run_command(benchmark, cwd=root_dir, timeout=600) - - -def run_spec_dec_tests(build_dir: _pl.Path): - xml_output_file = build_dir / "results-spec-dec-fast-logits.xml" - cpp_env = {**_os.environ} - tests_dir = build_dir / "tests" / "e2e_tests" - trt_model_test = _cpp.produce_mpirun_command( - global_commands=["mpirun", "--allow-run-as-root"], - nranks=3, - local_commands=[ - "executor/executorTest", "--gtest_filter=*SpecDecFastLogits*" - ], - leader_commands=[f"--gtest_output=xml:{xml_output_file}"]) - _cpp.run_command(trt_model_test, cwd=tests_dir, env=cpp_env, timeout=1500) - - -@pytest.fixture(scope="session") -def run_model_tests(build_dir, lora_setup): - - def _run(model_name: str, run_fp8: bool): - run_single_gpu_tests( - build_dir=build_dir, - test_list=[model_name], - timeout=_cpp.default_test_timeout, - run_fp8=run_fp8, - ) - - return _run - - -@pytest.fixture(scope="session") -def run_model_benchmarks(root_dir, build_dir, cpp_resources_dir, python_exe, - model_cache): - - def _run( - model_name: str, - batching_types: List[str], - api_types: List[str], - ): - - run_benchmarks( - model_name=model_name, - python_exe=python_exe, - root_dir=root_dir, - build_dir=build_dir, - resources_dir=cpp_resources_dir, - model_cache=model_cache, - batching_types=batching_types, - api_types=api_types, - ) - - return _run - - -@pytest.mark.parametrize("build_google_tests", ["80", "86", "89", "90"], - indirect=True) -@pytest.mark.parametrize("model", [ - "bart", "chatglm", "eagle", "encoder", "enc_dec_language_adapter", "gpt", - "gpt_executor", "gpt_tests", "llama", "mamba", "medusa", "recurrentgemma", - "redrafter", "t5" -]) -@pytest.mark.parametrize("run_fp8", [False, True], ids=["", "fp8"]) -def test_model(build_google_tests, model, prepare_model, run_model_tests, - run_fp8): - - if model == "recurrentgemma": - pytest.skip( - "TODO: fix recurrentgemma OOM with newest version of transformers") - return - - prepare_model(model, run_fp8) - - run_model_tests(model, run_fp8) - - -@pytest.mark.parametrize("build_google_tests", ["80", "86", "89", "90"], - indirect=True) -@pytest.mark.parametrize("model", ["bart", "gpt", "t5"]) -def test_benchmarks(build_benchmarks, model, prepare_model, - run_model_benchmarks): - - prepare_model(model) - - batching_types = ["IFB"] - api_types = ["executor"] - - run_model_benchmarks( - model_name=model, - batching_types=batching_types, - api_types=api_types, - ) diff --git a/tests/integration/defs/cpp/test_multi_gpu.py b/tests/integration/defs/cpp/test_multi_gpu.py index 1124178cccc6..12a0353fda10 100644 --- a/tests/integration/defs/cpp/test_multi_gpu.py +++ b/tests/integration/defs/cpp/test_multi_gpu.py @@ -1,9 +1,7 @@ import os as _os import pathlib as _pl import platform -import time from enum import Enum, auto -from typing import List, Optional import defs.cpp.cpp_common as _cpp import pytest @@ -149,344 +147,6 @@ def run_nccl_utils_tests(build_dir: _pl.Path, nprocs=2, timeout=300): timeout=timeout) -def run_llama_executor_leader_tests(build_dir: _pl.Path, timeout=1500): - tests_dir = build_dir / "tests" / "e2e_tests" - - mgpu_env = get_multi_gpu_env(llama_multi_gpu=True) - - #Executor test in leader mode - xml_output_file = build_dir / "results-multi-gpu-llama-exec-leader-mode.xml" - trt_model_test = _cpp.produce_mpirun_command( - global_commands=["mpirun", "--allow-run-as-root"], - nranks=4, - local_commands=[ - "executor/executorTest", - "--gtest_filter=*LlamaExecutorTest*LeaderMode*:*LlamaMultiExecutorTest*LeaderMode*" - ], - leader_commands=[f"--gtest_output=xml:{xml_output_file}"]) - - _cpp.run_command(trt_model_test, cwd=tests_dir, env=mgpu_env, timeout=1500) - - -def run_llama_executor_orchestrator_tests(build_dir: _pl.Path, timeout=1500): - tests_dir = build_dir / "tests" / "e2e_tests" - - mgpu_env = get_multi_gpu_env(llama_multi_gpu=True) - - #Executor test in orchestrator mode - xml_output_file = build_dir / "results-multi-gpu-llama-exec-orch-mode.xml" - trt_model_test = [ - "mpirun", "-n", "1", "--allow-run-as-root", "executor/executorTest", - "--gtest_filter=*LlamaExecutorTest*OrchMode*", - f"--gtest_output=xml:{xml_output_file}" - ] - _cpp.run_command(trt_model_test, cwd=tests_dir, env=mgpu_env, timeout=1500) - - -def run_llama_executor_logits_proc_tests(build_dir: _pl.Path, timeout=1500): - tests_dir = build_dir / "tests" / "e2e_tests" - - mgpu_env = get_multi_gpu_env(llama_multi_gpu=True) - - #Logits processor test in leader mode - xml_output_file = build_dir / "results-multi-gpu-logits-proc.xml" - - tp_pp_sizes = [(4, 1), (2, 2), (1, 4)] - gtest_filter = [ - f"LlamaExecutorTest/LogitsProcParamsTest*tp{tp}_pp{pp}*" - for tp, pp in tp_pp_sizes - ] - - gtest_filter = ":".join(gtest_filter) - - trt_model_test = _cpp.produce_mpirun_command( - global_commands=["mpirun", "--allow-run-as-root"], - nranks=4, - local_commands=[ - "executor/executorTest", f"--gtest_filter={gtest_filter}" - ], - leader_commands=[f"--gtest_output=xml:{xml_output_file}"]) - - _cpp.run_command(trt_model_test, cwd=tests_dir, env=mgpu_env, timeout=1500) - - -def run_llama_executor_guided_decoding_tests(build_dir: _pl.Path, timeout=1500): - tests_dir = build_dir / "tests" / "e2e_tests" - - mgpu_env = get_multi_gpu_env(llama_multi_gpu=True) - - #Guided decoding test in leader mode - xml_output_file = build_dir / "results-multi-gpu-guided-decoding.xml" - - tp_pp_sizes = [(4, 1), (2, 2), (1, 4)] - gtest_filter = [ - f"LlamaExecutorGuidedDecodingTest/GuidedDecodingParamsTest*tp{tp}_pp{pp}*" - for tp, pp in tp_pp_sizes - ] - - gtest_filter = ":".join(gtest_filter) - - trt_model_test = _cpp.produce_mpirun_command( - global_commands=["mpirun", "--allow-run-as-root"], - nranks=4, - local_commands=[ - "executor/executorTest", f"--gtest_filter={gtest_filter}" - ], - leader_commands=[f"--gtest_output=xml:{xml_output_file}"]) - - _cpp.run_command(trt_model_test, cwd=tests_dir, env=mgpu_env, timeout=1500) - - -def run_enc_dec_multi_gpu_tests(build_dir: _pl.Path, timeout=1500): - tests_dir = build_dir / "tests" / "e2e_tests" - cpp_env = {**_os.environ} - - #EncDec test in leader mode - xml_output_file = build_dir / "results-multi-gpu-t5-exec-leader-mode.xml" - trt_model_test = _cpp.produce_mpirun_command( - global_commands=["mpirun", "--allow-run-as-root"], - nranks=4, - local_commands=[ - "executor/encDecTest", - "--gtest_filter=T5MultiGPUTest/EncDecParamsTest.Forward*" - ], - leader_commands=[f"--gtest_output=xml:{xml_output_file}"], - ) - _cpp.run_command(trt_model_test, cwd=tests_dir, env=cpp_env, timeout=1500) - - -def run_trt_gpt_model_real_decoder_multi_gpu_tests(build_dir: _pl.Path, - timeout=1500): - tests_dir = build_dir / "tests" / "e2e_tests" - cpp_env = {**_os.environ} - - xml_output_file = build_dir / "results-multi-gpu-real-decoder.xml" - trt_model_test = _cpp.produce_mpirun_command( - global_commands=["mpirun", "--allow-run-as-root"], - nranks=4, - local_commands=[ - "batch_manager/trtGptModelRealDecoderTest", - "--gtest_filter=*TP*:*PP*" - ], - leader_commands=[f"--gtest_output=xml:{xml_output_file}"]) - _cpp.run_command(trt_model_test, - cwd=tests_dir, - env=cpp_env, - timeout=timeout) # expecting ~ 1200s - - -def run_disagg_symmetric_executor_tests(build_dir: _pl.Path, - model: str, - nprocs=2, - kvcache_type=KVCacheType.MPI, - timeout=1500): - tests_dir = build_dir / "tests" / "e2e_tests" - - prefix = get_model_test_filter_prefix(model) - - mgpu_env = get_multi_gpu_env(kv_cache_type=kvcache_type, - llama_multi_gpu=True) - - xml_output_file = build_dir / f"results-multi-gpu-disagg-executor-{nprocs}-process.xml" - trt_model_test = _cpp.produce_mpirun_command( - global_commands=["mpirun", "--allow-run-as-root"], - nranks=nprocs, - local_commands=[ - "executor/disaggExecutorTest", - f"--gtest_filter=*{prefix}*DisaggSymmetricExecutorTest*" - ], - leader_commands=[f"--gtest_output=xml:{xml_output_file}"]) - - _cpp.run_command(trt_model_test, - cwd=tests_dir, - env=mgpu_env, - timeout=timeout) - - -def run_disagg_asymmetric_executor_tests(build_dir: _pl.Path, - model: str, - nprocs=4, - kvcache_type=KVCacheType.MPI, - timeout=1500): - - tests_dir = build_dir / "tests" / "e2e_tests" - - prefix = get_model_test_filter_prefix(model) - - mgpu_env = get_multi_gpu_env(kv_cache_type=kvcache_type, - llama_multi_gpu=True) - - xml_output_file = build_dir / f"results-multi-gpu-disagg-asymmetric-executor-{nprocs}-process.xml" - - trt_model_test = _cpp.produce_mpirun_command( - global_commands=["mpirun", "--allow-run-as-root"], - nranks=nprocs, - local_commands=[ - "executor/disaggExecutorTest", - f"--gtest_filter=*{prefix}*DisaggAsymmetricExecutorTest*" - ], - leader_commands=[f"--gtest_output=xml:{xml_output_file}"]) - - _cpp.run_command(trt_model_test, - cwd=tests_dir, - env=mgpu_env, - timeout=timeout) - - -def run_disagg_orchestrator_params_tests(build_dir: _pl.Path, - model: str, - kvcache_type=KVCacheType.MPI, - timeout=1500): - - tests_dir = build_dir / "tests" / "e2e_tests" - - prefix = get_model_test_filter_prefix(model) - - mgpu_env = get_multi_gpu_env(kv_cache_type=kvcache_type, - llama_multi_gpu=True) - - xml_output_file = build_dir / "results-multi-gpu-disagg-asymmetric-orchestrator-executor-7-process.xml" - trt_model_test = _cpp.produce_mpirun_command( - global_commands=["mpirun", "--allow-run-as-root"], - nranks=7, - local_commands=[ - "executor/disaggExecutorTest", - f"--gtest_filter=*{prefix}*DisaggOrchestratorParamsTest*" - ], - leader_commands=[f"--gtest_output=xml:{xml_output_file}"]) - _cpp.run_command(trt_model_test, - cwd=tests_dir, - env=mgpu_env, - timeout=timeout) - - -def run_disagg_spawn_orchestrator_tests(build_dir: _pl.Path, - model: str, - kvcache_type=False, - timeout=1500): - - tests_dir = build_dir / "tests" / "e2e_tests" - - prefix = get_model_test_filter_prefix(model) - - mgpu_env = get_multi_gpu_env(kv_cache_type=kvcache_type, - llama_multi_gpu=True) - - xml_output_file = build_dir / "results-multi-gpu-disagg-spawn-asymmetric-orchestrator-executor-1-process.xml" - - comms = [ - "executor/disaggExecutorTest", - f"--gtest_filter=*{prefix}*DisaggSpawnOrchestrator*", - f"--gtest_output=xml:{xml_output_file}" - ] - _cpp.run_command(comms, cwd=tests_dir, env=mgpu_env, timeout=timeout) - - -def prepare_multi_gpu_model_tests(test_list: List[str], - python_exe: str, - root_dir: _pl.Path, - resources_dir: _pl.Path, - model_cache: Optional[str] = None): - - model_cache_arg = ["--model_cache", model_cache] if model_cache else [] - - if "llama" in test_list: - _cpp.prepare_model_tests(model_name="llama", - python_exe=python_exe, - root_dir=root_dir, - resources_dir=resources_dir, - model_cache_arg=model_cache_arg, - only_multi_gpu_arg=["--only_multi_gpu"]) - - if "t5" in test_list: - _cpp.prepare_model_tests(model_name="t5", - python_exe=python_exe, - root_dir=root_dir, - resources_dir=resources_dir, - model_cache_arg=model_cache_arg, - only_multi_gpu_arg=['--tp', '4', '--pp', '1']) - - -@pytest.fixture(scope="session") -def prepare_model_multi_gpu(python_exe, root_dir, cpp_resources_dir, - model_cache): - - def _prepare(model_name: str): - if platform.system() != "Windows": - - start_time = time.time() - - prepare_multi_gpu_model_tests( - test_list=[model_name], - python_exe=python_exe, - root_dir=root_dir, - resources_dir=cpp_resources_dir, - model_cache=model_cache, - ) - - duration = time.time() - start_time - print(f"Built multi-GPU model: {model_name}") - print(f"Duration: {duration} seconds") - - return _prepare - - -@pytest.fixture(scope="session") -def gpt_single_gpu_model(prepare_model): - prepare_model("gpt") - return "gpt" - - -@pytest.fixture(scope="session") -def llama_single_gpu_model(prepare_model): - prepare_model("llama") - return "llama" - - -@pytest.fixture(scope="session") -def llama_multi_gpu_model(prepare_model_multi_gpu): - prepare_model_multi_gpu("llama") - return "llama" - - -# Allow us to dynamically choose a fixture at runtime -# Combined with session scope fixtures above to ensure -# that the model is built only once per pytest session -@pytest.fixture -def prepare_models_disagg(request): - - def _prepare(model_name: str): - if model_name == "llama": - fixture_names = [ - "llama_single_gpu_model", - "llama_multi_gpu_model", - ] - elif model_name == "gpt": - fixture_names = [ - "gpt_single_gpu_model", - ] - else: - raise ValueError(f"Disagg tests don't support model: {model_name}") - - print(f"Preparing models for disagg tests: {fixture_names}") - # Run the fixtures - for fixture_name in fixture_names: - request.getfixturevalue(fixture_name) - - return _prepare - - -# Use indirect parameterization to ensure that the model is built -# only once per pytest session -@pytest.fixture(scope="session") -def multi_gpu_model(request, prepare_model_multi_gpu): - - model_name = request.param - prepare_model_multi_gpu(model_name) - - return model_name - - @pytest.mark.parametrize("build_google_tests", ["80", "86", "89", "90"], indirect=True) def test_mpi_utils(build_google_tests, build_dir): @@ -535,138 +195,3 @@ def test_nccl_utils(build_google_tests, nprocs, build_dir): if platform.system() != "Windows": run_nccl_utils_tests(build_dir=build_dir, nprocs=nprocs, timeout=300) - - -@pytest.mark.parametrize("build_google_tests", ["80", "86", "89", "90"], - indirect=True) -@pytest.mark.parametrize("multi_gpu_model", ["t5"], indirect=True) -def test_enc_dec(build_google_tests, multi_gpu_model, build_dir): - - if platform.system() != "Windows": - run_enc_dec_multi_gpu_tests(build_dir=build_dir, - timeout=_cpp.default_test_timeout) - - -@pytest.mark.parametrize("build_google_tests", ["80", "86", "89", "90"], - indirect=True) -@pytest.mark.parametrize("mode", ["orchestrator", "leader"]) -@pytest.mark.parametrize("multi_gpu_model", ["llama"], indirect=True) -def test_llama_executor(build_google_tests, multi_gpu_model, mode, lora_setup, - build_dir): - - if platform.system() == "Windows": - return - - if mode == "orchestrator": - run_llama_executor_orchestrator_tests(build_dir=build_dir, - timeout=_cpp.default_test_timeout) - elif mode == "leader": - run_llama_executor_leader_tests(build_dir=build_dir, - timeout=_cpp.default_test_timeout) - else: - raise ValueError(f"Unsupported mode: {mode}") - - -@pytest.mark.parametrize("build_google_tests", ["80", "86", "89", "90"], - indirect=True) -@pytest.mark.parametrize("multi_gpu_model", ["llama"], indirect=True) -def test_llama_executor_logits_proc(build_google_tests, multi_gpu_model, - lora_setup, build_dir): - - if platform.system() != "Windows": - run_llama_executor_logits_proc_tests(build_dir=build_dir, - timeout=_cpp.default_test_timeout) - - -@pytest.mark.parametrize("build_google_tests", ["80", "86", "89", "90"], - indirect=True) -@pytest.mark.parametrize("multi_gpu_model", ["llama"], indirect=True) -def test_llama_executor_guided_decoding(build_google_tests, multi_gpu_model, - lora_setup, build_dir): - - if platform.system() != "Windows": - run_llama_executor_guided_decoding_tests( - build_dir=build_dir, timeout=_cpp.default_test_timeout) - - -@pytest.mark.parametrize("build_google_tests", ["80", "86", "89", "90"], - indirect=True) -@pytest.mark.parametrize("multi_gpu_model", ["llama"], indirect=True) -def test_trt_gpt_real_decoder(build_google_tests, multi_gpu_model, lora_setup, - build_dir): - - if platform.system() != "Windows": - run_trt_gpt_model_real_decoder_multi_gpu_tests( - build_dir=build_dir, timeout=_cpp.default_test_timeout) - - -@pytest.mark.parametrize("build_google_tests", ["80", "86", "89", "90"], - indirect=True) -class TestDisagg: - - @pytest.mark.parametrize( - "kvcache_type", [KVCacheType.MPI, KVCacheType.UCX, KVCacheType.NIXL], - ids=["mpi_kvcache", "ucx_kvcache", "nixl_kvcache"]) - @pytest.mark.parametrize("nprocs", [2, 4, 8], - ids=["2proc", "4proc", "8proc"]) - @pytest.mark.parametrize("model", ["gpt", "llama"]) - def test_symmetric_executor(self, build_google_tests, model, nprocs, - kvcache_type, prepare_models_disagg, build_dir): - - if model == "gpt" and nprocs > 2: - pytest.skip( - "test_symmetric_executor only supports 2 processes for gpt") - - if platform.system() != "Windows": - prepare_models_disagg(model) - - run_disagg_symmetric_executor_tests(build_dir=build_dir, - model=model, - nprocs=nprocs, - kvcache_type=kvcache_type) - - @pytest.mark.parametrize( - "kvcache_type", [KVCacheType.MPI, KVCacheType.UCX, KVCacheType.NIXL], - ids=["mpi_kvcache", "ucx_kvcache", "nixl_kvcache"]) - @pytest.mark.parametrize("nprocs", [4, 6, 8], - ids=["4proc", "6proc", "8proc"]) - @pytest.mark.parametrize("model", ["llama"]) - def test_asymmetric_executor(self, build_google_tests, model, nprocs, - kvcache_type, prepare_models_disagg, - build_dir): - - if platform.system() != "Windows": - prepare_models_disagg(model_name=model) - - run_disagg_asymmetric_executor_tests(build_dir=build_dir, - model=model, - nprocs=nprocs, - kvcache_type=kvcache_type) - - @pytest.mark.parametrize( - "kvcache_type", [KVCacheType.MPI, KVCacheType.UCX, KVCacheType.NIXL], - ids=["mpi_kvcache", "ucx_kvcache", "nixl_kvcache"]) - @pytest.mark.parametrize("model", ["llama"]) - def test_orchestrator_params(self, build_google_tests, model, kvcache_type, - prepare_models_disagg, build_dir): - - if platform.system() != "Windows": - prepare_models_disagg(model) - - run_disagg_orchestrator_params_tests(build_dir=build_dir, - model=model, - kvcache_type=kvcache_type) - - @pytest.mark.parametrize("kvcache_type", - [KVCacheType.UCX, KVCacheType.NIXL], - ids=["ucx_kvcache", "nixl_kvcache"]) - @pytest.mark.parametrize("model", ["llama"]) - def test_spawn_orchestrator(self, build_google_tests, model, kvcache_type, - prepare_models_disagg, build_dir): - - if platform.system() != "Windows": - prepare_models_disagg(model) - - run_disagg_spawn_orchestrator_tests(build_dir=build_dir, - model=model, - kvcache_type=kvcache_type) diff --git a/tests/integration/defs/cpp/test_unit_tests.py b/tests/integration/defs/cpp/test_unit_tests.py index a55440bb1a48..730ebbf389ed 100644 --- a/tests/integration/defs/cpp/test_unit_tests.py +++ b/tests/integration/defs/cpp/test_unit_tests.py @@ -8,7 +8,7 @@ indirect=True) @pytest.mark.parametrize("test_group", [ "batch_manager", "common", "executor", "kernels", "layers", "runtime", - "thop", "utils" + "thop" ]) def test_unit_tests(build_google_tests, test_group, build_dir, lora_setup): diff --git a/tests/integration/defs/deterministic/mixtral_deterministic.py b/tests/integration/defs/deterministic/mixtral_deterministic.py deleted file mode 100644 index 53abff63d739..000000000000 --- a/tests/integration/defs/deterministic/mixtral_deterministic.py +++ /dev/null @@ -1,253 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2023-2025 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 argparse -import json -import os -import time -from concurrent.futures import ThreadPoolExecutor -from typing import Dict, List - -import jinja2 -from transformers import AutoTokenizer - -import tensorrt_llm.bindings.executor as trtllm - - -def get_args(): - parser = argparse.ArgumentParser() - parser.add_argument("--engine_dir", type=str, default="engine_outputs") - parser.add_argument("--tokenizer_dir", type=str, default="tokenizer_dir") - parser.add_argument("--payload", type=str, default="./payload.json") - parser.add_argument("--concurrency", type=int, default=50) - parser.add_argument("--check_deterministic_accuracy", - action="store_true", - default=False) - parser.add_argument("--deterministic_accuracy_threshold", - type=int, - default=1) - parser.add_argument("--batch", action="store_true", default=False) - parser.add_argument("--wait", type=float, default=0.0) - parser.add_argument("--output", type=str, default='out-strs') - return parser.parse_args() - - -def create_engine(engine): - parallel_config = trtllm.ParallelConfig( - communication_type=trtllm.CommunicationType.MPI, - communication_mode=trtllm.CommunicationMode.LEADER) - trt_scheduler_config = trtllm.SchedulerConfig( - trtllm.CapacitySchedulerPolicy.GUARANTEED_NO_EVICT) - kv_cache_config = trtllm.KvCacheConfig( - free_gpu_memory_fraction=0.9, - enable_block_reuse=True, - ) - extend_runtime_perf_knob_config = trtllm.ExtendedRuntimePerfKnobConfig() - extend_runtime_perf_knob_config.cuda_graph_mode = False - extend_runtime_perf_knob_config.multi_block_mode = False - executor_config = trtllm.ExecutorConfig( - 1, - iter_stats_max_iterations=100, - # nvbugs/4662826 - request_stats_max_iterations=0, - parallel_config=parallel_config, - # normalize_log_probs=False, - batching_type=trtllm.BatchingType.INFLIGHT, - # batching_type=trtllm.BatchingType.STATIC, - scheduler_config=trt_scheduler_config, - kv_cache_config=kv_cache_config, - enable_chunked_context=True, - extended_runtime_perf_knob_config=extend_runtime_perf_knob_config, - ) - - return trtllm.Executor(model_path=engine, - model_type=trtllm.ModelType.DECODER_ONLY, - executor_config=executor_config) - - -def create_request(payload_file, template_str, tokenizer): - json_data = { - 'model': - 'my-model', - 'messages': [ - { - 'role': 'user', - 'content': 'Hello there how are you?', - }, - { - 'role': 'assistant', - 'content': 'Good and you?', - }, - { - 'role': 'user', - 'content': 'Whats your name?', - }, - ], - 'max_tokens': - 1024, - 'temperature': - 0, - #'top_k':1, - #'nvext': {"top_k": 1}, - 'stream': - False - } - - json_data['messages'][2]['content'] = """ - Classify the sentiment expressed in the following text and provide the response in a single word Positive/Negative/Neutral. Explain your answer in 2 lines. - TEXT:: Today I will exaggerate, will be melodramatic (mostly the case when I am excited) and be naive (as always). Just came out from the screening of the Avengers Endgame ("Endgame")! The journey had started in the year 2008, when Tony Stark, during his capture in a cave in Afghanistan, had created a combat suit and came out of his captivity. - Then the combat suit made of iron was perfected and Tony Stark officially became the Iron Man!! The Marvel Cinematic Universe ("MCU") thus was initiated. The journey continued since then and in 2012 all the MCU heroes came together and formed the original "Avengers" (so much fun and good it was). - 21 Movies in the MCU and culminating into the Infinity War (2018) and finally into the Endgame! The big adventure for me started from Jurassic Park and then came Titanic, Lagaan, Dark Knight; and then came the Avengers in 2012. Saw my absolute favorite Sholay in the hall in 2014. In the above-mentioned genre, there are good movies, great movies and then there is the Endgame. - Today after a long long time, I came out of the hall with 100% happiness, satisfaction and over the top excitement/emotions. The movie is Epic, Marvel (in the real sense) and perfect culmination of the greatest cinematic saga of all time. It is amazing, humorous, emotional and has mind-blowing action! It is one of the finest Superhero Movie of all time. - Just pure Awesome! It's intelligent! - """ - with open(payload_file, 'r') as f: - msg_system = json.load(f) - msg_user = [] - msg_user.append({ - "role": - "user", - "content": - msg_system[0]["content"] + "\n\n" + msg_system[1]["content"] - }) - msg_user.extend(msg_system[2:]) - json_data['messages'] = msg_user - - environment = jinja2.Environment() - template = environment.from_string(template_str) - json_data['bos_token'] = '' - json_data['eos_token'] = '' - prompt = template.render(json_data) - - tokens = tokenizer.encode(prompt) - - sample_params = trtllm.SamplingConfig( - beam_width=1, # beam_width=1 for inflight batching - top_k=1, # SizeType topK - top_p=1.0, - top_p_min=None, - top_p_reset_ids=None, # SizeType topPResetIds - top_p_decay=None, # FloatType topPDecay - seed=1234, - temperature=1, - min_tokens=1, # SizeType minLength - beam_search_diversity_rate=None, # FloatType beamSearchDiversityRate - repetition_penalty=1, # FloatType repetitionPenalty - presence_penalty=0, # FloatType presencePenalty - frequency_penalty=0, # FloatType frequencyPenalty - length_penalty=1, # FloatType lengthPenalty - early_stopping= - None, # SizeType earlyStopping. Controls beam search, so irrelevant until we have beam_width > 1 - ) - #sample_params = trtllm.SamplingConfig(temperature=0, seed=1234) - - return trtllm.Request( - input_token_ids=tokens[1:], - max_tokens=1024, - sampling_config=sample_params, - streaming=False, - stop_words=None, - # stop_words=[[2]], # - ), prompt - - -def get_tokenizer(tokenizer_file): - return AutoTokenizer.from_pretrained(tokenizer_file) - - -def get_template(tokenizer_file): - with open(os.path.join(tokenizer_file, - "tokenizer_config.json")) as tok_config: - cfg = json.load(tok_config) - return cfg['chat_template'] - - -def enqueue_requests(pool, executor, request, concurrency=50, wait=0): - for _ in range(concurrency): - _ = pool.submit(executor.enqueue_request, request) - if wait > 0: - time.sleep(wait) - - -def main(): - args = get_args() - executor = create_engine(args.engine_dir) - if executor.can_enqueue_requests(): - template = get_template(args.tokenizer_dir) - tokenizer = get_tokenizer(args.tokenizer_dir) - concurrency = int(args.concurrency) - - request, prompt = create_request(args.payload, template, tokenizer) - os.makedirs(args.output, exist_ok=True) - with open(os.path.join(args.output, "prompt.txt"), 'w') as f: - f.write(prompt) - - try: - for _ in range(1): - outputs: Dict[str, List[trtllm.Result]] = {} - num_finished = 0 - - if not args.batch: - with ThreadPoolExecutor(max_workers=concurrency) as pool: - enqueue_requests(pool, - executor, - request, - concurrency=concurrency, - wait=args.wait) - else: - executor.enqueue_requests( - [request for _ in range(concurrency)]) - while num_finished < concurrency: - responses = executor.await_responses() - for response in responses: - if response.has_error(): - outputs[response.request_id] = response.error_msg - num_finished += 1 - else: - result = response.result - if result.is_final: - num_finished += 1 - if response.request_id not in outputs: - outputs[response.request_id] = [] - outputs[response.request_id].append(result) - output_strs = {} - for req_id, output in outputs.items(): - if isinstance(output, str): - raise RuntimeError(output) - elif isinstance(output, list): - if len(output) != 1: - raise RuntimeError("Expected list size of 1") - output_strs[req_id] = tokenizer.decode( - output[0].output_token_ids[0]) - with open(os.path.join(args.output, f"{req_id}.out"), - "w") as f: - f.write(output_strs[req_id]) - else: - raise RuntimeError("Unexpected output") - - output_set = set(output_strs.values()) - num_unique_responses = len(output_set) - if args.check_deterministic_accuracy: - assert num_unique_responses <= args.deterministic_accuracy_threshold, f"Expected num unique responses <= {args.deterministic_accuracy_threshold} while got {num_unique_responses} " - result_str = f"Num Unique responses in {len(outputs)}: {len(output_set)}" - print(result_str) - with open(os.path.join(args.output, "num_outputs"), 'w') as f: - f.write(result_str + '\n') - finally: - executor.shutdown() - - -if __name__ == "__main__": - main() diff --git a/tests/integration/defs/deterministic/payload.json b/tests/integration/defs/deterministic/payload.json deleted file mode 100644 index b3373b811fb9..000000000000 --- a/tests/integration/defs/deterministic/payload.json +++ /dev/null @@ -1,14 +0,0 @@ -[ - { - "role": "system", - "content": "DOCUMENTS:\n[\n {\n \"id\": \"doc_1\",\n \"title\": \"Employee Benefits - India\",\n \"content\": \"**Insurance Entity** **SPOC** **Escalation-1** **Final Escalation** IPD: Medi Assist - TPA Dedicated help \\& Support Mr. Hamza Zaidi Mr. Syed Sajid ++[servicenow@mediassist.in](mailto:thota.rajani@mediassist.in)++ ++[hamza.zaidi@mediassist.in](mailto:hamza.zaidi@mediassist.in)++ ++[syed.sajid@mediassist.in](mailto:syed.sajid@mediassist.in)++ 8951386235 7349122713 8105565671 OPD: Connect and Heal (CNH) Dedicated help \\& Support Mr.Rohit ++[servicenow.support@connectandheal.com](mailto:servicenow.support@connectandheal.com)++ ++[rohit.singh@connectandheal.com](mailto:rohit.singh@connectandheal.com)++ 91118 91118 8884726473 Marsh - Broker Mr. Sridhar Anneboina Ms. Anjali Agarwal Mr. K. Balajee Rao ++[Sridhar.anneboina@marsh.com](mailto:Sridhar.anneboina@marsh.com)++ ++[anjali.agarwal@marsh.com](mailto:anjali.agarwal@marsh.com)++ ++[balajee.rao@marsh.com](mailto:balajee.rao@marsh.com)++ 8976773547 9966200636 9160555531 Note: Any insurance escalations are to reach up to the last contact of Marsh and later if required can reach for internal support. ++[Internal Escalation: Please raise a request in the My ServiceNow Portal.](https://my.servicenow.com/esc?id=service_catalog&spa=1&sc_catalog=42c07a15db951700b2e2d34b5e9619f1&sc_category=undefined&sc_cat_item=36b09733db36df00014fffd51d9619a0)++ **Tele Consultations on CNH portal:**A Virtual Consultation with a medical/health practitioner online for 32 specialists. This service will help to reach and consult the specialists \\& doctors to discuss health concerns from home. No cost is implied for phone specialist consultations and is available for dependants as well. **Annual Health Screening** **on CNH portal** **:**A yearly free annual health check-up for employees and a discounted program for dependants. **For further information:** Read more about this program on the [benefits portal.](https://urldefense.com/v3/__https:/servicenow.okta.com/home/thomsonsonlinebenefits/0oa1g1v3i17DiSuob0h8/38923?fromHome=true__;!!N4vogdjhuJM!Vyvf3ceOERltigdu08bwjhFaS6SzovSGcXfj6H7vRqzpUz_C1AJ5rSTXDQh037j1frPr9vC0bKY$) **Need Help?** You can log in to the CNH portal via Quick Links for bookings and can find a user guide available in the [benefits portal.](https://urldefense.com/v3/__https:/servicenow.okta.com/home/thomsonsonlinebenefits/0oa1g1v3i17DiSuob0h8/38923?fromHome=true__;!!N4vogdjhuJM!Vyvf3ceOERltigdu08bwjhFaS6SzovSGcXfj6H7vRqzpUz_C1AJ5rSTXDQh037j1frPr9vC0bKY$) contact ++[servicenow.support@connectandheal.com](mailto:servicenow.support@connectandheal.com)++before booking your appointment to know more about the terms of this service. |\n| **Group Term Life Insurance** | | The current insurance provider is ICICI Prudential Life Insurance Co. Ltd. Our insurance is renewed every **31 March to 30 March.** **I**t is an employee-only benefit. Should in the event of death either natural or accident of the Insured Member whilst working for ServiceNow as a regular, full-time active employee, you or your beneficiary(ies) will receive a lump sum payment equal to the insured amount. This payment will be subject to income tax. The insured amount is **3 x Annual Base Salary** (coverage may be subject to medical underwriting and insurer's acceptance) and subject to a maximum of INR 4,00,00,000, per Insured Member * Age Group of Cover: 18 years to 80 years of age (age as of last birthday) **New to ServiceNow?** You are automatically covered under this policy from joining the day of ServiceNow. You need to nominate your dependents in the **[benefits portal.](https://urldefense.com/v3/__https:/servicenow.okta.com/home/thomsonsonlinebenefits/0oa1g1v3i17DiSuob0h8/38923?fromHome=true__;!!N4vogdjhuJM!Vyvf3ceOERltigdu08bwjhFaS6SzovSGcXfj6H7vRqzpUz_C1AJ5rSTXDQh037j1frPr9vC0bKY$)** **Policy Rules:** Please refer to the [**Insurance Manual**](sys_attachment.do?sys_id=5b6062bb47a3795030c37a24846d43eb)for more information. **What happens if I leave ServiceNow?** On your date of exit, the group term life plan coverages stop immediately. **Need Help?** Please contact our Broker SPOCs. **Internal support** - Please [raise a request](https://my.servicenow.com/esc?id=service_catalog&spa=1&sc_catalog=42c07a15db951700b2e2d34b5e9619f1&sc_category=undefined&sc_cat_item=36b09733db36df00014fffd51d9619a0) in the My ServiceNow Portal. **Insurance Entity** **SPOC** **Escalation-1** **Final Escalation** **Marsh - Broker** Mr. Sridhar Anneboina Ms. Anjali Agarwal Mr. K. Balajee Rao [Sridhar.anneboina@marsh.com](mailto:Indraneel.Singamsetty@marsh.com) [anjali.agarwal@marsh.com](mailto:anjali.agarwal@marsh.com) [balajee.rao@marsh.com](mailto:balajee.rao@marsh.com) 8976773547 9966200636 9160555531 |\n| **Group Personal Accident** | | The current insurance provider is Liberty General Insurance. Our insurance is renewed every **31 March to 30 March.** It is an employee-only benefit. This insurance covers you for both work-related \\& non-work-related accidents. Should in the event of an accident and Permanent Total or Partial Disability of the Insured Member whilst working for ServiceNow as a regular, full-time active employee, you or your beneficiary(ies) will receive a lump sum payment equal to the insured amount. This payment will be subject to income tax. The insured amount is **3x Annual Base Salary** (coverage may be subject to medical underwriting and insurer's acceptance) per Insured Member. This policy provides a lump sum benefit to the Insured Member or dependent(s), in the event the Insured Member sustains Bodily Injury as per policy terms. Age Group of Cover: 18 years to 65 years of age (age as of last birthday) **New to ServiceNow?**You are automatically covered under this policy from joining the day of ServiceNow. You need to nominate your dependents in the [benefits portal.](https://urldefense.com/v3/__https:/servicenow.okta.com/home/thomsonsonlinebenefits/0oa1g1v3i17DiSuob0h8/38923?fromHome=true__;!!N4vogdjhuJM!Vyvf3ceOERltigdu08bwjhFaS6SzovSGcXfj6H7vRqzpUz_C1AJ5rSTXDQh037j1frPr9vC0bKY$) **Policy Rules:** Please refer to the [**Insurance Manual**](sys_attachment.do?sys_id=5b6062bb47a3795030c37a24846d43eb) for more information. **What happens if I leave ServiceNow?** On your date of exit, the group personal accident plan coverages stop immediately. **Need Help?** Please contact our Broker SPOCs. **Internal support** - Please [raise a request](https://my.servicenow.com/esc?id=service_catalog&spa=1&sc_catalog=42c07a15db951700b2e2d34b5e9619f1&sc_category=undefined&sc_cat_item=36b09733db36df00014fffd51d9619a0) in the My ServiceNow Portal. **Insurance Entity** **SPOC** **Escalation-1** **Final Escalation** **Marsh - Broker** Mr. Sridhar Anneboina Ms. Anjali Agarwal Mr. K. Balajee Rao[\\[SN10\\]](#_msocom_10){#_anchor_10} [Sridhar.anneboina@marsh.com](mailto:Indraneel.Singamsetty@marsh.com) [anjali.agarwal@marsh.com](mailto:anjali.agarwal@marsh.com) [balajee.rao@marsh.com](mailto:balajee.rao@marsh.com) 8976773547 9966200636 9160555531 *** ** * ** *** |\n| #### **Retirement \\& Pension Benefits** | | ServiceNow contributes to the statutory norms of the Country and Provides PF and gratuity. **Provident Fund (PF)** - Both employer \\& employee stipulated funds A provident fund is a part of your salary that is deducted every month and deposited with the Employee Provident Fund Organization (EPFO) under your name. Normally 12% of your basic monthly salary is deposited in the EPFO while a matching percentage is contributed by ServiceNow. While your contribution is fully invested, the employer's contribution is partly invested along with your funds and partly utilized for providing other benefits like Pensions and Employee Deposit Linked Insurance (EDLI) premiums. When an employee leaves his/her job, a part of this amount can be withdrawn as per the governing withdrawal rules. If the employee attains retirement age, then the Provident Fund can be partly withdrawn and partly used for purchasing a life annuity plan that provides a monthly pension. **Policy Rules:** Please refer to Benefits Information in the [benefits portal.](https://urldefense.com/v3/__https:/servicenow.okta.com/home/thomsonsonlinebenefits/0oa1g1v3i17DiSuob0h8/38923?fromHome=true__;!!N4vogdjhuJM!Vyvf3ceOERltigdu08bwjhFaS6SzovSGcXfj6H7vRqzpUz_C1AJ5rSTXDQh037j1frPr9vC0bKY$) **Need to avail?** Please [raise a request](https://my.servicenow.com/esc?id=service_catalog&spa=1&sc_catalog=0c0369b4db555300b2e2d34b5e9619fe&sc_category=undefined&sc_cat_item=0284fe44130dd240796ad1e2e144b0f9) in the My ServiceNow Portal with Payroll. **New to ServiceNow?** Soon after joining ServiceNow, you will be automatically enrolled in PF, and you need to provide the required details to Payroll for processing. **Need Help?** Please [rai](https://my.servicenow.com/esc?id=service_catalog&spa=1&sc_catalog=0c0369b4db555300b2e2d34b5e9619fe&sc_category=undefined&sc_cat_item=0284fe44130dd240796ad1e2e144b0f9)[se a request](https://my.servicenow.com/esc?id=service_catalog&spa=1&sc_catalog=0c0369b4db555300b2e2d34b5e9619fe&sc_category=undefined&sc_cat_item=0284fe44130dd240796ad1e2e144b0f9)in the My ServiceNow Portal with Payroll. **Gratuity:** An employee completing almost 5 years or more of uninterrupted service can claim this benefit. Only in case the employee suffers from a permanent total disability that this benefit be accelerated and can be paid before completion of 5 years of service as well. **Policy Rules:** Please refer to Benefits Information in the [benefits portal.](https://urldefense.com/v3/__https:/servicenow.okta.com/home/thomsonsonlinebenefits/0oa1g1v3i17DiSuob0h8/38923?fromHome=true__;!!N4vogdjhuJM!Vyvf3ceOERltigdu08bwjhFaS6SzovSGcXfj6H7vRqzpUz_C1AJ5rSTXDQh037j1frPr9vC0bKY$) **Need to avail?** Under the following two conditions, an employee can claim this benefit: 1. Retirement/ Resignation/ Superannuation/Death after 5 years or more of uninterrupted service. 2. Suffering from a permanent total disability due to accident, sickness, or disease. **New to ServiceNow?** Soon after joining ServiceNow, you will be automatically enrolled in Gratuity. **Need Help?** Please [raise a request](https://my.servicenow.com/esc?id=service_catalog&spa=1&sc_catalog=0c0369b4db555300b2e2d34b5e9619fe&sc_category=undefined&sc_cat_item=0284fe44130dd240796ad1e2e144b0f9) in the My ServiceNow Portal with the payroll **NPS- National Pension System** The National Pension System is a Tax effective low-cost pension plan that helps employees manage finances post-retirement. Please refer to the NPS section in the [benefits portal](https://urldefense.com/v3/__https:/servicenow.okta.com/home/thomsonsonlinebenefits/0oa1g1v3i17DiSuob0h8/38923?fromHome=true__;!!N4vogdjhuJM!Vyvf3ceOERltigdu08bwjhFaS6SzovSGcXfj6H7vRqzpUz_C1AJ5rSTXDQh037j1frPr9vC0bKY$) to know more about NPS. Corporate NPS enables to save additional tax up to 10% of Salary (Basic) on employer's contribution in NPS u/s 80CCD(2) of Income Tax Act, over and above the limit of 1.5 lakhs u/s 80C and exclusive additional tax deduction benefit of Rs. 50,000/- u/s 80CCD(1B). It is a voluntary plan and employees once have a retirement account raise a request with an [HR ticket](https://my.servicenow.com/esc?id=service_catalog&spa=1&sc_catalog=42c07a15db951700b2e2d34b5e9619f1&sc_category=0ca38e41dba9a300e7e5f9c31d961955&sc_cat_item=e9ef54aadb2636809f5258a8dc9619c2) for PRAN account mapping with corporate and after HR confirmation can start contributing monthly from their salary via employer with a [payroll ticket](https://my.servicenow.com/esc?id=service_catalog&spa=1&sc_catalog=0c0369b4db555300b2e2d34b5e9619fe&sc_category=undefined&sc_cat_item=ce6fbf171be1f0101e41caae6e4bcb67). These funds will be deposited to the employee NPS account and invested. **Need to avail?** Under the following two conditions, an employee/nominee can claim this benefit as per the pension plan guidelines of NSDL. 1. Retirement at the age of 60 years 2. Death **Need Help?** Please reach the [NSDL Portal](https://npscra.nsdl.co.in/all-citizens-faq.php). **Internal support** - Please [raise a request](https://my.servicenow.com/esc?id=service_catalog&spa=1&sc_catalog=42c07a15db951700b2e2d34b5e9619f1&sc_category=undefined&sc_cat_item=e9ef54aadb2636809f5258a8dc9619c2) in the My ServiceNow Portal. #### **What happens if I leave ServiceNow?** **On your date of exit, the Retirement \\& Pension Benefits stop immediately.** |\n\n**If you need additional assistance or have questions about any of the information provided in this article,** please [raise a request](https://my.servicenow.com/esc?id=service_catalog&spa=1&sc_catalog=42c07a15db951700b2e2d34b5e9619f1&sc_category=undefined&sc_cat_item=36b09733db36df00014fffd51d9619a0) in the My ServiceNow Portal.\n\n**Additional Resources**\n\n[**Employee Benefits Summary**](sys_attachment.do?sys_id=1f6062bb47a3795030c37a24846d43e5)\n\nTime Away and [Holiday](https://my.servicenow.com/esc?id=events) -- [KB0056866](https://my.servicenow.com/esc?sys_id=d1080c2e87a42d102217437e0ebb3580&id=surf_kb_article&q=Time%20Away%20-%20India)\n\nDay Care Facility -- [KB0059432](https://my.servicenow.com/esc?sys_id=af448b25973e81100a63b84bf253afd3&id=surf_kb_article&q=Day%20Care%20Facility%20%E2%80%93%20India)\n\nTuition Reimbursement - [KB0007401](https://my.servicenow.com/esc?id=surf_kb_article&sys_id=290bf1d7db54a8108ae42b691396199c)\n\nStock Resource Center - [KB0034611](https://my.servicenow.com/esc?id=surf_kb_article&number=KB0034611)\n\nBusiness Travel - [KB0056930](https://my.servicenow.com/esc?id=surf_kb_article&number=KB0056930)\n\nEmployee Assistance Program - [KB0045470](https://my.servicenow.com/esc?id=surf_kb_article&number=KB0045470 \"KB0045470\")\n\nWhil/Rethink Benefits**-** [KB0051803](https://my.servicenow.com/esc?sys_id=fc0e5326dbf064505ad34f7813961974&id=surf_kb_article&q=Rethink%20benefits)\n\nGlobal Wellbeing Programs - [KB0058670](https://my.servicenow.com/esc?id=surf_kb_article&number=KB0058670)\n\nCleo - Parenting support -- [KB0058667](https://my.servicenow.com/esc?id=surf_kb_article&number=KB0058667)\n\nFertilityIQ - Family building resources - [KB0058660](https://my.servicenow.com/esc?id=surf_kb_article&number=KB0058660)\n\nGrayce - Caregiver support - [KB0058668](https://my.servicenow.com/esc?id=surf_kb_article&number=KB0058668)\n\nVirgin Pulse - Wellbeing platform - [KB0058669](https://my.servicenow.com/esc?id=surf_kb_article&number=KB0058669)\n\nFond - Discount Program - [KB0008923](https://my.servicenow.com/esc?id=surf_kb_article&number=KB0008923 \"KB0008923\")\n\n Employee Benefits - India\n\n|——————————-|---|—————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————---|\n| **Group Medical Insurance** | | The current medical insurance provider is Aditya Birla Health Insurance and the Third-Party Administrator (TPA) for claims management is Medi Assist Healthcare Services Private Limited for IPD and Connect and Heal (CNH) for OPD. Our insurance is renewed every **31 March to 30 March.** This policy provides Cashless/reimbursement of recognized costs of services or treatment which are recommended by a medical practitioner, and are medically necessary for the care and treatment of an injury or sickness, up to the limits shown in the list of benefits in the [++Insurance Manual++](sys_attachment.do?sys_id=5b6062bb47a3795030c37a24846d43eb) The insured amount is up to INR 5,00,000 and INR 15,000 per year of insurance which is inclusive of hospitalization coverages, Day Care procedures, maternity, and more for IPD and OPD expenses like dental and vision, generalist consultations, specialist consultations, and AYUSH. Employee-only medical insurance coverage for Critical illness \\& COVID of INR 5,00,000 and INR 50,000 This insurance coverage may be subject to sub-limits as shown in the list of benefits in the [Insurance manual](sys_attachment.do?sys_id=5b6062bb47a3795030c37a24846d43eb) and provides cover within Pan India. **Eligibility:** All regular, full-time active employees, employed by ServiceNow India, and their Eligible Dependents\\* **\\*Eligible Dependents:** * One legal spouse or legal LGBTQ partner * Children (4) under the age of 25 * Either Parents or Parents_In_Laws(Combination not allowed) **Policy rules:** Please refer to the [**Insurance manual**](sys_attachment.do?sys_id=5b6062bb47a3795030c37a24846d43eb) for more information. **Need to make a claim?**You may follow the claim process mentioned in the benefits portal under Group Medical Insurance benefit -- FAQ. (Cashless / Reimbursement) **New to ServiceNow?** Soon after joining ServiceNow, you will be automatically enrolled in this insurance plan, and you need to enroll your dependants in the [benefits portal](https://urldefense.com/v3/__https:/servicenow.okta.com/home/thomsonsonlinebenefits/0oa1g1v3i17DiSuob0h8/38923?fromHome=true__;!!N4vogdjhuJM!Vyvf3ceOERltigdu08bwjhFaS6SzovSGcXfj6H7vRqzpUz_C1AJ5rSTXDQh037j1frPr9vC0bKY$) and mark them for this insurance coverage within 13 days. You will be provided access to this benefits portal on your day 2 and also informed via email. **Life event?** If you are recently married or have a newborn baby, log the new dependent details on the benefits portal and mark them for this insurance cover within 30 days from the date of the event. For any further assistance, please raise a [ticket](https://my.servicenow.com/esc?id=service_catalog&spa=1&sc_catalog=42c07a15db951700b2e2d34b5e9619f1) in the My ServiceNow portal and our Global People Care team will be in touch. Once you have completed the online process, the medical cover will be set up and ready to go Via the TPA's portal, and mobile apps and membership cards will be available in the next 30 business days. **++Group Medical Insurance Top-up Coverage++** **++:++** The Mediclaim plan is an extension of company-provided medical insurance cover and you may choose this on a yearly basis for yourself and your family as an extra cover for medical expenses. **++Top-Up Insurance (Voluntary)--++** You also have the option to purchase an additional Mediclaim expense limit over and above the company-sponsored INR 5,00,000 and INR 15,000 limit via a \"Top-up Medical Insurance Plan\". Under this plan, you can choose to buy any one of the additional Sum Insured limits for IPD from INR 3,00,000, INR 5,00,000, INR 7,00,000, INR 10,00,000 and INR 15,00,000 for OPD from INR 7,000, INR 10,000, INR 15,000 by paying an additional premium from your salary. This Top-up plan is an extension of base coverages as per policy terms and will be automatically initiated when default sum insurance is exhausted. **++Note -++** The premium amount details can be found in the Group Medical Insurance Top-up sections in the Benefits portal and the respective premiums will be deducted in three equal installments from your salary based on the top-up coverage Selections. As with all medical schemes, there are a number of policy exclusions that may affect your ability to claim. For this reason, you should read the [++Insurance Manual++](sys_attachment.do?sys_id=5b6062bb47a3795030c37a24846d43eb) which details the terms and conditions of coverage. **Need to make a claim?**The claims will automatically trigger once the medical policy cover provided by the company is exhausted. **New to ServiceNow?** Soon after joining ServiceNow, you will be provided with options to choose from this insurance plan in the [benefits portal.](https://urldefense.com/v3/__https:/servicenow.okta.com/home/thomsonsonlinebenefits/0oa1g1v3i17DiSuob0h8/38923?fromHome=true__;!!N4vogdjhuJM!Vyvf3ceOERltigdu08bwjhFaS6SzovSGcXfj6H7vRqzpUz_C1AJ5rSTXDQh037j1frPr9vC0bKY$)Your dependants have automatically been considered for this cover from the company-provided base policy and the premium is pro-rated as per your Date of joining. **What happens if I leave ServiceNow?** On your date of exit, the group medical plan coverages stop immediately. **Need Help?** Please contact our TPA or Broker SPOCs.\"\n },\n {\n \"id\": \"doc_2\",\n \"title\": \"India Darwin Benefits Enrolment Portal Guide\",\n \"content\": \"India Darwin Benefits Enrolment Portal Guide\n\n**Q) How do I select my Benefits?**\n\n1. Sign into the **[My ServiceNow Benefits](https://servicenow.okta.com/home/bookmark/0oa1ttc0ujznXItCh0h8/2557)** or via [Okta](https://servicenow.okta.com/app/UserHome?session_hint=AUTHENTICATED). And select the 'Review your benefits Coverage (Go to the Darwin benefits enrolment portal)' tile in the left side corner.\n2. Select the notification bell in the top right-hand corner\n3. If you are a new joiner, you will see this option 'New Joiner Enrolment' notification\n4. If you click view or select benefits, the relevant insurance benefits will be available to select.\n5. Benefits are broken down into the following:\n\n* Financial\n* Physical\n\nIf you would like to know more about these, please select the 'more info' link found underneath each benefit which will take you to a handy FAQ document.\n\n6. Make your Insurance selections.\n\na. Click 'Edit' on the benefit you would like to select.\n\nb. Complete the sequence of steps\n\nc. Add selection to Cart\n\nd. Make sure to check out!\n\n7. Shortly after checking out, you will receive an email from DoNotReply@ServiceNowBenefits.com confirming your selections.\n\na. All selections made will show as green in your benefits portal.\n\n**Q) When can I add or change my dependents?**\n\n* You can add or remove dependents to your cover during the annual enrolment window, which occurs in March each year.\n* You can also apply to make a change to your cover if your personal circumstances change. This is called a lifestyle event and broadly includes things like births, marriage. You can apply for a lifestyle event by navigating to the person icon in the top right corner of the darwin benefits enrollment site and selecting \"Life Event\"\n* Please note;\n * When you add dependents to your benefits portal, they do not automatically add to the policies, the dependants must be added manually by yourself.\n * Kindly enroll new dependents (spouse/newborn) within one month of the occurrence of the life event.\n\n**Q) When can I select my benefits and do my previous selections carry over?**\n\nYou can add or remove dependents, change/add the Medical TOP UP selections to your cover during the annual enrolment window, which occurs in March each year.\n\n1. Sign into the **[My ServiceNow Benefits](https://servicenow.okta.com/home/bookmark/0oa1ttc0ujznXItCh0h8/2557)**or via [Okta](https://servicenow.okta.com/app/UserHome?session_hint=AUTHENTICATED). And select the 'Review your benefits Coverage (Go to the Darwin benefits enrolment portal)' tile in the left side corner.\n2. Select the notification bell in the top right-hand corner\n3. If you are a current employee, you will see this option 'Annual Enrolment Window' notification\n4. If you click view or select benefits, the relevant insurance benefits will be available to select.\n5. Benefits are broken down into the following:\n\n* Financial\n* Physical\n\nIf you would like to know more about these, please select the 'more info' link found underneath each benefit which will take you to a handy FAQ document.\n\n6. Make your Insurance selections.\n\na. Click 'Edit' on the each of the insurance benefit you would like to select(to navigate from one benefit to another click 'Go back' on top left corner of each benefits page)\n\nb. Complete the sequence of steps\n\nc. Add selection to Cart\n\nd. Make sure to check out!\n\nPlease note; Note: If there are no changes to any of the benefit page, the system notifies (Refer below the pop up notification) that there is no change to consider, no action to take. The selections will be carried forward by default.\n\nIf this article does not resolve your issue, please reach out to ServiceNow Benefit Helpdesk team by navigating to the 'question mark' icon in the top right corner of the [darwin enrollment site](https://servicenow.okta.com/home/thomsonsonlinebenefits/0oa1g1v3i17DiSuob0h8/38923?fromHome=true) and click on it, they will assist you.\"\n },\n {\n \"id\": \"doc_3\",\n \"title\": \"ServiceNow's Our World of Work – Employee FAQs\",\n \"content\": \"ServiceNow's Our World of Work – Employee FAQs\n\n**{#Top1}Our world of work -- Employee Frequently Asked Questions (FAQs)**\n\nThis document provides supporting details for our World of Work model. Our World of Work model is based on these guiding principles:\n\n* Keep employees healthy and safe\n* Live company values\n* Reinforce our brand and product opportunities\n* Drive employee engagement and experience\n* Commit to equity and inclusion in our culture\n* Further our Global Impact goals\n\nThe information contained in this FAQ is subject to local legal requirements, and may change at any time. Employees should leverage these key resources for additional information on our policies and practices:\n\n* [Our world of work KB hub](https://my.servicenow.com/esc?id=surf_kb_article&number=KB0057030)\n* [30-day work from anywhere benefit](https://my.servicenow.com/esc?id=surf_kb_article&number=KB0057039)\n* [Flexible work persona](https://my.servicenow.com/esc?id=surf_kb_article&number=KB0057029)\n* [Remote work persona](https://my.servicenow.com/esc?id=surf_kb_article&number=KB0057031)\n* [Required in office work persona](https://my.servicenow.com/esc?id=surf_kb_article&number=KB0057402)\n* [Work persona change request](https://my.servicenow.com/esc?id=surf_kb_article&number=KB0057042)\n\nClick on the questions below to be redirected to the answer:\n\n**General questions**\n\n[What are the benefits of a distributed, flexible work model?](#G1)\n\n[How is Our World of Work defined?](#G2)\n\n[Employees have been successfully working from home. Why doesn't this just continue for everyone in all organizations globally?](#G3)\n\n[What else is ServiceNow doing to drive improvements to collaboration, which was an opportunity area identified in the Employee Voice Survey (EVS)?](#GQ4)\n\n**Work personas**\n\n[How do I know what my work persona is?](#WP1)\n\n[What if I want to change my work persona?](#WP2)\n\n[How are work personas assigned?](#WP3)\n\n[If I submit a change to my work persona/work location, when will it be processed in our systems?](#WP4)\n\n[If my work persona is remote, can I relocate to live anywhere I like?](#WPN2)\n\n[If I change my work persona to remote, is there a compensation change? What about if I change my status to remote but still live in the same area?](#WP5)\n\n[Why have we clarified that our expectation for the flexible work persona is two or more days a week in the office?](#WPN3)\n\n[What happens if I don't come into the office two or more days a week?](#WPN4)\n\n[What if I have a personal circumstance that may make it difficult for me to meet the two or more days a week in-office expectation?](#WPN5)\n\n[What if I am in the flexible work persona, but none of my direct team is co-located in my assigned work location?](#WPNN2)\n\n[What is the messaging for contingent workers?](#WPNN1)\n\n**Policies and benefits**\n\n[Which policies support Our World of Work?](#P1)\n\n[Will I be reimbursed for my home office expenses?](#P2)\n\n[Can I use my Work from Anywhere (WFA) 30-day allotment toward the two or more days per week expectation?](#PB3)\n\n**Workplaces and workspaces**\n\n[How will workplaces evolve in our flexible, distributed model?](#WPL1)\n\n[Am I allowed to work from a ServiceNow office more than two days a week?](#WPL3)\n\n[What if there are multiple ServiceNow offices within a commutable distance for me? Which office am I expected to be aligned to?](#WWN1)\n\n[Is it possible to work from a location other than the contracted location for which I was hired?](#WPL4)\n\n[Will I have an assigned desk?](#WPL5)\n\n[What kind of workstation setup can I expect when I go into the office?](#WWN3 \"What kind of workstation setup can I expect when I go into the office?\")\n\n[What is the difference between a workplace and workspace?](#WPL6)\n\n[If we want to sit near our partners and internal customers, are we allowed access to those neighborhoods?](#WPL7)\n\n[How will the new work personas impact travel to ServiceNow workplace locations?](#WPL8)\n\n[Is my workplace attendance data being shared with anyone? If so, who is it being shared with?](#WWN4)\n\n[Why are some organizations more closely monitoring workplace attendance data than others?](#WWN5)\n\n[I'm in the flexible work persona, but work during \"off hours\" (evening or nighttime work start times). How are we accounting for physical workplace safety in meeting the two or more days per week expectation?](#WPNN1)\n\n[If I live in India, which states am I allowed to move to?](#WPL9)\n\n{#G1}\n\n**Q. What are the benefits of a distributed, flexible work model?**\n\n**A.** We have adopted a distributed, flexible work model that is right for Our World of Work. We reshaped our connections with each other, adapted the intent and purpose of the workplace, and used the power of the Now Platform to do our best work together. What about if I change my status to remote but still live in the same area?**\n\n**A**. Work personas by themselves do not affect compensation, but the market tier you work in will (only in the United States). If your work persona is flexible or required-in-office, your compensation is based on your office location. If you are remote, then your compensation is based on your remote work location. Refer to the [market tiers for the United States](https://my.servicenow.com/esc?id=surf_kb_article&number=KB0005540).\n\n[**Go to top**](#Top1 \"Go to top\")\n\n{#WPN3}\n\n**Q.** **Why have we clarified that our expectation for the flexible work persona is two or more days a week in the office?**\n\n**A.**By getting more specific on expectation of number of days spent in the office, we will bring more clarity for and ultimately drive consistency for employees in the flexible work persona across the company. We feel that two or more days a week will help facilitate the benefits and advantages of 'shoulder-to-shoulder' work. We know how valuable remote work is as well, and that's why we continue to lead Our World of Work with your flexibility and choice in mind.\n\n[**Go to top**](#Top1 \"Go to top\")\n\n{#WPN4}\n\n**Q. What happens if I don't come into the office two or more days a week?**\n\n**A.**Your manager may check in with you during regular discussions on whether you're adhering to the expectation of two or more days each week. If you're not regularly coming to the office, your manager may partner with you to create a plan that supports you meeting the expectation for being in the office. For now, there are no formal consequences for not meeting this expectation.\n\n[**Go to top**](#Top1 \"Go to top\")\n\n{#WPN5}\n\n**Q. What if I have a personal circumstance that may make it difficult for me to meet the two or more days a week in-office expectation?**\n\n**A.**Bill said it best: There's something incalculable in ServiceNow's culture about working \"shoulder-to-shoulder.\" We come together to do meaningful work while getting to know our colleagues. We believe that working with your team (and across teams) in person fosters even greater collaboration and innovation, builds social wellbeing via the relationships we develop, and actively cultivates ServiceNow's culture. Thus, we feel that those who live close to ServiceNow offices are a natural fit for the flexible persona. It is reasonable to ask you to be present in our offices without burdensome commutes that unduly impact your personal time.\n\nIf you have a personal circumstance that prevents you from meeting the expectations of the flexible work persona and you would like to request an accommodation, please use these resources:\n\n* Employees in the U.S. can reference [this workplace accommod...\"\n }\n]\n\n\nCATALOG_ITEMS:\n[\n {\n \"id\": \"cat_1\",\n \"title\": \"Benefits - General\",\n \"description\": \"\"\n },\n {\n \"id\": \"cat_2\",\n \"title\": \"Benefits Coverage\",\n \"description\": \"For questions related to your benefits including your coverage, enrolling in or changing your benefits, portal login issues, business traveler benefits, or benefit changes related to leaves of absence (LOA) please use this request\"\n },\n {\n \"id\": \"cat_3\",\n \"title\": \"Benefits - 401(k)\",\n \"description\": \"\"\n },\n {\n \"id\": \"cat_4\",\n \"title\": \"IN - Benefits & Insurance details\",\n \"description\": \"**Please add your dependent details for medical insurance policy.**\"\n },\n {\n \"id\": \"cat_5\",\n \"title\": \"Benefits - Leave / LOA\",\n \"description\": \"\"\n }\n]" - }, - { - "role": "user", - "content": "USER QUERY:\ni would like to know benefits\n\nTASK:\n- Your role as a question and answering bot is to respond to the given USER QUERY based on the information provided in the DOCUMENTS and CATALOG ITEMS\n- If the USER QUERY is not answerable based on the DOCUMENTS, say 'No answer found' and nothing else.\n- If there are any helpful CATALOG ITEMS to the USER QUERY, you should recommend them.\n\nIMPORTANT RULES:\n- Format your response in markdown format\n- Provide citations like (doc_1), (doc_2), (doc_3) or (cat_1), (cat_2), (cat_3) throughout your response.\n- Do not say \"Based on ...\" or \"Sure\"\n- Do not use any first person language such as \"I\".\n- Do not produce any HTML tags\n- Keep your answer concise.\n- Do not miss any helpful CATALOG ITEMS.\n\nRESPONSE FORMAT:\n\nuser_interested_to_read_docs: \nprovide_enough_details_for_user_to_understand_without_reading: \ndo_not_provide_task_instructions_to_user: \ndo_not_respond_to_unsafe_requests: \nlist_of_max_5_based_on_user_query_history:[max 5 key word topics...]\nlist_of_max_5_topics_requested_by_user: [max of 5 key word topics...]\nlist_of_max_5_topics_in_docs: [max of 5 key word topics...]\nlist_of_matching_catalog_items: [cat_1...]\nsay_no_answer_found: \n\n....\n" - }, - { - "role": "assistant", - "content": "\nuser_interested_to_read_docs: \nprovide_enough_details_for_user_to_understand_without_reading: \ndo_not_provide_task_instructions_to_user: \ndo_not_respond_to_unsafe_requests: " - } -] diff --git a/tests/integration/defs/deterministic/test_mixtral_deterministic.py b/tests/integration/defs/deterministic/test_mixtral_deterministic.py deleted file mode 100644 index 1796e9827d31..000000000000 --- a/tests/integration/defs/deterministic/test_mixtral_deterministic.py +++ /dev/null @@ -1,81 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 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 os - -import pytest -from defs.common import (convert_weights, generate_deterministic_cmd, - venv_mpi_check_call) -from defs.conftest import skip_pre_hopper -from defs.trt_test_alternative import check_call - - -@skip_pre_hopper -@pytest.mark.skip_less_device(4) -@pytest.mark.skip_less_device_memory(80000) -@pytest.mark.parametrize("data_type", ['float16', 'bfloat16']) -@pytest.mark.parametrize("llm_mixtral_model_root", - ['Mixtral-8x7B-Instruct-v0.1'], - indirect=True) -def test_llm_mixtral_4gpus_deterministic(llama_example_root, - llm_mixtral_model_root, - deterministic_test_root, llm_venv, - cmodel_dir, engine_dir, data_type): - tp_size, pp_size = 4, 1 - world_size = tp_size * pp_size - moe_tp_size = tp_size - - os.environ['FORCE_DETERMINISTIC'] = "1" - - print("Convert checkpoint...") - ckpt_dir = convert_weights(llm_venv=llm_venv, - example_root=llama_example_root, - cmodel_dir=cmodel_dir, - model="mixtral-instruct", - model_path=llm_mixtral_model_root, - tp_size=tp_size, - moe_tp_size=moe_tp_size, - pp_size=pp_size, - data_type=data_type, - workers=world_size) - - print("Build engines...") - build_cmd = [ - "trtllm-build", - f"--checkpoint_dir={ckpt_dir}", - f"--output_dir={engine_dir}", - f"--workers={world_size}", - "--use_paged_context_fmha=enable", - "--max_batch_size=256", - "--max_num_tokens=33280", - ] - - check_call(" ".join(build_cmd), shell=True, env=llm_venv._new_env) - - print("Run deterministic test...") - deterministic_accuracy_threshold = 1 - payload = os.path.join(deterministic_test_root, "payload.json") - deterministic_cmd = generate_deterministic_cmd( - deterministic_test_root, - engine_dir=engine_dir, - tokenizer_dir=llm_mixtral_model_root, - payload=payload, - deterministic_accuracy_threshold=deterministic_accuracy_threshold) - - venv_mpi_check_call( - llm_venv, ["mpirun", "-n", f"{world_size}", "--allow-run-as-root"], - deterministic_cmd) - - os.environ.pop('FORCE_DETERMINISTIC', None) diff --git a/tests/integration/defs/perf/test_perf.py b/tests/integration/defs/perf/test_perf.py index 6d3fb5093813..3cfea70c4f8e 100644 --- a/tests/integration/defs/perf/test_perf.py +++ b/tests/integration/defs/perf/test_perf.py @@ -1066,7 +1066,7 @@ def get_trtllm_bench_build_command(self, engine_dir) -> list: def get_prepare_data_command(self, engine_dir, input_len, output_len) -> list: data_cmd = [] - prepare_data_script = os.path.join(self._llm_root, "benchmarks", "cpp", + prepare_data_script = os.path.join(self._llm_root, "benchmarks", "prepare_dataset.py") if self._config.model_name in MODEL_PATH_DICT.keys(): diff --git a/tests/integration/defs/test_mlpf_results.py b/tests/integration/defs/test_mlpf_results.py deleted file mode 100644 index 71dbb12430dd..000000000000 --- a/tests/integration/defs/test_mlpf_results.py +++ /dev/null @@ -1,260 +0,0 @@ -""" -MLPerf target tests -""" -import os -import subprocess -from argparse import Namespace -from copy import deepcopy - -from defs.common import get_cpp_benchmark, get_trt_llm_lib_dir, venv_check_call -from defs.trt_test_alternative import check_call - -### End of utility functions -""" -Test: Runs the gptManagerBenchmark on LLama TRTLLM engine and checks accuracy of predictions -Steps: - 1. Quantize the model: step_quantize - 2. Build the engine: step_engine_build - 3. Run engine and get outputs: step_run_llm - 4. Check prediction accuracy: step_check_accuracy -""" - - -# Test step 1: Quantize the model -# MLPerf step: python examples/quantization/quantize.py --dtype=float16 --output_dir=<> --model_dir=<> --qformat=fp8 --kv_cache_dtype=fp8 --tp_size 2 -def step_quantize(tp_size, llm_venv, llm_root, model_root, model, - calib_dataset): - quantized_model_path = "{}/test_mlperf_quantized_models/{}-tp{}-pp1/".format( - llm_venv.get_working_directory(), model, tp_size) - tekit_example_dir = os.path.join(llm_root, "examples/") - - # Set MLPerf params explicitly - quantize_cmd = [ - f"{tekit_example_dir}/quantization/quantize.py", "--dtype=float16", - "--qformat=fp8", "--kv_cache_dtype=fp8", f"--tp_size={tp_size}", - f"--output_dir={quantized_model_path}", f"--model_dir={model_root}", - "--calib_size=1024", f"--calib_dataset={calib_dataset}" - ] - - venv_check_call(llm_venv, quantize_cmd) - - return quantized_model_path - - -# Test step 2: Build the TRTLLM engine -# MLPerf step: -# python3 -m tensorrt_llm.commands.build --gpt_attention_plugin=float16 --max_batch_size=896 --max_input_len=1024 --max_seq_len=2048 --max_beam_width=1 \ -# --max_num_tokens=4096 --output_dir=<> --checkpoint_dir=<> --context_fmha=enable --remove_input_padding=enable \ -# --paged_kv_cache=enable --workers=2 - - -def step_engine_build(quantized_model_path, system_config, engine_dir, - llm_venv): - - batch_size = system_config.batch_size - beam_width = system_config.beam_width - max_input_len = system_config.max_input_len - max_seq_len = system_config.max_seq_len - max_num_tokens = system_config.max_num_tokens - num_workers = system_config.num_workers - use_fp8_context_fmha = "enable" if system_config.fp8_fmha else "disable" - - build_cmd = [ - "trtllm-build", - "--gpt_attention_plugin=float16", - f"--max_batch_size={batch_size}", - f"--max_input_len={max_input_len}", - f"--max_seq_len={max_seq_len}", - f"--max_beam_width={beam_width}", - f"--max_num_tokens={max_num_tokens}", - f"--output_dir={engine_dir}", - f"--checkpoint_dir={quantized_model_path}", - "--context_fmha=enable", - f"--use_fp8_context_fmha={use_fp8_context_fmha}", - "--remove_input_padding=enable", - "--paged_kv_cache=enable", - f"--workers={num_workers}", - ] - - check_call(" ".join(build_cmd), shell=True, env=llm_venv._new_env) - return engine_dir - - -DEFAULT_RPARAMS = Namespace( - engine_dir=None, - api="executor", - # type="IFB", - dataset=None, - output_csv="gptmanager_bench_results.csv", - max_num_samples=24576, - beam_width=1, - warm_up=2, - eos_id=-1, - pad_id=-1, - max_tokens_in_paged_kvcache=None, - kv_cache_free_gpu_mem_fraction=None, - streaming=False, - enable_kv_cache_reuse=False, - enable_chunked_context=False, - return_context_logits=False, - return_generation_logits=False, - scheduler_policy="guaranteed_no_evict", - static_emulated_batch_size=None, - log_level="verbose", - log_iteration_data=False, - wait_sleep="25", - lora_dir=None, - lora_host_cache_bytes=None, - lora_num_device_mod_layers=None, - responses_json=None) -""" -./benchmarks/gptManagerBenchmark \ - --engine_dir <> \ - --dataset <> \ - --max_num_samples 24576 \ - --beam_width 1 \ - --eos_id 2 \ - --pad_id 2 \ - --kv_cache_free_gpu_mem_fraction 0.95 \ - --scheduler_policy max_utilization \ - --output_csv <> -""" - - -# Test step 3: Run the gptManagerBenchmark and get outputs -def step_run_llm(system_config, - engine_path, - dataset_path, - llm_venv, - llm_root, - kv_cache_free_gpu_mem_fraction=0.95): - tp, pp = system_config.tp_size, system_config.pp_size - eos_id, pad_id = system_config.eos_id, system_config.pad_id - max_num_samples = system_config.num_samples - beam_width = system_config.beam_width - - benchmark_exe = get_cpp_benchmark('gptManagerBenchmark', llm_root) - workspace_path = llm_venv.get_working_directory() - run_params = deepcopy(DEFAULT_RPARAMS) - run_params.beam_width = beam_width - run_params.engine_dir = engine_path - run_params.dataset = dataset_path - run_params.max_num_samples = max_num_samples - run_params.eos_id = eos_id - run_params.pad_id = pad_id - run_params.kv_cache_free_gpu_mem_fraction = kv_cache_free_gpu_mem_fraction - run_params.scheduler_policy = "max_utilization" - run_params.responses_json = os.path.join( - workspace_path, f"responses_test_mlperf_tp{tp}_pp{pp}.json") - run_params.output_csv = os.path.join( - workspace_path, f"perf_stats_test_mlperf_tp{tp}_pp{pp}.csv") - - run_params_dict = vars(run_params) - run_params_dict['type'] = "IFB" - - bench_cmd = [benchmark_exe] - for key, val in run_params_dict.items(): - if val is None or val is False: - continue - if val is True: - val = "" - bench_cmd.append("--" + str(key)) - bench_cmd.append(str(val)) - - envs = deepcopy(os.environ) - _ = envs.pop("CUDA_VISIBLE_DEVICES", "") - envs[ - "LD_LIBRARY_PATH"] = f'{get_trt_llm_lib_dir(llm_venv)}:{os.path.dirname(benchmark_exe)}:{envs.get("LD_LIBRARY_PATH", "")}' - - print( - f'CUDA_VISIBLE_DEVICES: {os.environ.get("CUDA_VISIBLE_DEVICES", None)}') - - num_ranks = tp * pp - if num_ranks > 1: - mpi_cmd = ["mpirun", "-n", f"{num_ranks}", "--allow-run-as-root"] - bench_cmd = mpi_cmd + bench_cmd - - print(f"Running gptManagerBenchmark using cmd: {' '.join(bench_cmd)}") - subprocess.check_output(bench_cmd, env=envs) - return run_params.responses_json - - -def step_check_accuracy(responses_file, dataset_path, model_root, llm_venv, - llm_root): - """ - python3 /code/tensorrt_llm/benchmarks/python/check_accuracy_mlperf.py - --dataset <> - --responses <> - --base_model <> - """ - accuracy_script = os.path.join( - llm_root, "benchmarks/python/check_accuracy_mlperf.py") - accuracy_check_cmd = [ - f"{accuracy_script}", "--dataset", f"{dataset_path}", "--responses", - f"{responses_file}", "--base_model", f"{model_root}" - ] - venv_check_call(llm_venv, accuracy_check_cmd) - - -LlamaBaseSystem = Namespace(tp_size=None, - pp_size=1, - batch_size=None, - max_input_len=1024, - max_seq_len=2048, - max_num_tokens=4096, - beam_width=1, - num_workers=None, - num_samples=24576, - eos_id=2, - pad_id=2, - fp8_fmha=False) - -GptjBaseSystem = Namespace(tp_size=1, - pp_size=1, - batch_size=None, - max_input_len=1919, - max_seq_len=2047, - max_num_tokens=4096, - beam_width=4, - num_workers=1, - num_samples=13368, - eos_id=50256, - pad_id=50256, - fp8_fmha=False) - - -def get_mlperf_system_config(model: str, system: str, fp8_fmha: bool): - if model == "llama_v2_70b_chat": - return get_mlperf_llama_system_config(system) - elif model == "gpt_j": - return get_mlperf_gptj_system_config(system, fp8_fmha) - raise RuntimeError(f"Unexpected model: {system}") - - -def get_mlperf_llama_system_config(system: str): - system_config = deepcopy(LlamaBaseSystem) - if system == "H100x2": - system_config.tp_size = 2 - system_config.batch_size = 896 - system_config.num_workers = 2 - elif system == "H200x1": - system_config.tp_size = 1 - system_config.batch_size = 806 - system_config.num_workers = 1 - else: - raise RuntimeError(f"No Llama config found for system: {system}") - - return system_config - - -def get_mlperf_gptj_system_config(system: str, fp8_fmha: bool): - system_config = deepcopy(GptjBaseSystem) - system_config.fp8_fmha = fp8_fmha - if system == "H100x1": - system_config.batch_size = 192 - elif system == "H200x1": - system_config.batch_size = 396 - else: - raise RuntimeError(f"No GPT-J config found for system: {system}") - - return system_config diff --git a/tests/integration/test_lists/test-db/l0_a10.yml b/tests/integration/test_lists/test-db/l0_a10.yml index a720fa3e9912..a9e4d7e6af2a 100644 --- a/tests/integration/test_lists/test-db/l0_a10.yml +++ b/tests/integration/test_lists/test-db/l0_a10.yml @@ -178,22 +178,6 @@ l0_a10: - thirdparty/test_git_modules.py::test_gitmodules # helper-script unit tests (CPU-only, ~3s) - unittest/scripts -- condition: - ranges: - system_gpu_count: - gte: 1 - lte: 1 - wildcards: - gpu: - - '*a10*' - linux_distribution_name: ubuntu* - terms: - stage: pre_merge - backend: cpp - tests: - # ------------- CPP tests --------------- - - cpp/test_e2e.py::test_model[-redrafter-86] - - cpp/test_e2e.py::test_model[-mamba-86] - condition: ranges: system_gpu_count: diff --git a/tests/integration/test_lists/test-db/l0_a30.yml b/tests/integration/test_lists/test-db/l0_a30.yml index 125fb8191b1b..d406ca8829a1 100644 --- a/tests/integration/test_lists/test-db/l0_a30.yml +++ b/tests/integration/test_lists/test-db/l0_a30.yml @@ -64,9 +64,6 @@ l0_a30: - cpp/test_unit_tests.py::test_unit_tests[layers-80] - cpp/test_unit_tests.py::test_unit_tests[runtime-80] - cpp/test_unit_tests.py::test_unit_tests[thop-80] - - cpp/test_unit_tests.py::test_unit_tests[utils-80] - - cpp/test_e2e.py::test_model[-gpt-80] - - cpp/test_e2e.py::test_benchmarks[gpt-80] TIMEOUT (90) - condition: ranges: system_gpu_count: diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 464c3e553f91..2ce7bb9ee7b1 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -294,7 +294,6 @@ l0_b200: tests: # ------------- TRT tests --------------- - unittest/llmapi/test_llm_quant.py # 3.5 mins on B200 - - unittest/disaggregated/test_router.py - unittest/disaggregated/test_openai_server_info.py - condition: ranges: diff --git a/tests/integration/test_lists/test-db/l0_dgx_h100.yml b/tests/integration/test_lists/test-db/l0_dgx_h100.yml index b11214e16256..2ed7c363fe3a 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_h100.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_h100.yml @@ -283,9 +283,6 @@ l0_dgx_h100: - cpp/test_multi_gpu.py::test_cache_transceiver[8proc-ucx_kvcache-90] ISOLATION - cpp/test_multi_gpu.py::test_cache_transceiver[8proc-mooncake_kvcache-90] ISOLATION - cpp/test_multi_gpu.py::test_user_buffer[2proc-90] - - cpp/test_multi_gpu.py::TestDisagg::test_symmetric_executor[gpt-2proc-mpi_kvcache-90] - - cpp/test_multi_gpu.py::TestDisagg::test_symmetric_executor[gpt-2proc-ucx_kvcache-90] - - cpp/test_multi_gpu.py::TestDisagg::test_symmetric_executor[gpt-2proc-nixl_kvcache-90] - condition: ranges: system_gpu_count: diff --git a/tests/integration/test_lists/test-db/l0_h100.yml b/tests/integration/test_lists/test-db/l0_h100.yml index 5602f66498ea..dcd7b2c9f021 100644 --- a/tests/integration/test_lists/test-db/l0_h100.yml +++ b/tests/integration/test_lists/test-db/l0_h100.yml @@ -299,8 +299,6 @@ l0_h100: - cpp/test_unit_tests.py::test_unit_tests[kernels-90] - cpp/test_unit_tests.py::test_unit_tests[layers-90] - cpp/test_unit_tests.py::test_unit_tests[thop-90] - - cpp/test_unit_tests.py::test_unit_tests[utils-90] - - cpp/test_e2e.py::test_model[-encoder-90] - condition: ranges: system_gpu_count: @@ -418,22 +416,6 @@ l0_h100: - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_single_gpu -k "CUTLASS and None" # Documentation URL validation (CPU-only, no GPU needed) - test_doc.py::test_url_validity -- condition: - ranges: - system_gpu_count: - gte: 1 - lte: 1 - wildcards: - gpu: - - '*h100*' - linux_distribution_name: ubuntu* - terms: - stage: post_merge - backend: cpp - tests: - # ------------- CPP tests --------------- - - cpp/test_e2e.py::test_model[-bart-90] - - cpp/test_e2e.py::test_benchmarks[bart-90] - condition: ranges: system_gpu_count: diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index bd566aae915c..5a3a29252c40 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -132,16 +132,6 @@ accuracy/test_llm_api_pytorch.py::TestStep3_7::test_nvfp4[tp_size=4-ep_size=4-mt accuracy/test_llm_api_pytorch_multimodal.py::TestMistralLarge3_675B::test_nvfp4_4gpus[latency_moe_trtllm] SKIP (https://nvbugs/6248827) accuracy/test_llm_api_pytorch_multimodal.py::TestStep3_7::test_nvfp4[mtp_nextn=3] SKIP (https://nvbugs/6367805) accuracy/test_llm_api_pytorch_ray.py::TestLlama3_1_8BInstruct::test_pp2_ray SKIP (https://nvbugs/6427411) -cpp/test_e2e.py::test_benchmarks[bart-90] SKIP (https://nvbugs/5550689) -cpp/test_e2e.py::test_benchmarks[gpt-80] SKIP (https://nvbugs/5550689) -cpp/test_e2e.py::test_model[-bart-90] SKIP (https://nvbugs/6162804) -cpp/test_e2e.py::test_model[-encoder-90] SKIP (waive Encoder-only test because it doesn't take batched input) -cpp/test_e2e.py::test_model[-gpt-80] SKIP (https://nvbugs/5983283) -cpp/test_e2e.py::test_model[-mamba-86] SKIP (https://nvbugs/5781665) -cpp/test_e2e.py::test_model[-redrafter-86] SKIP (https://nvbugs/5761642) -cpp/test_multi_gpu.py::TestDisagg::test_symmetric_executor[gpt-2proc-mpi_kvcache-90] SKIP (https://nvbugs/5755941) -cpp/test_multi_gpu.py::TestDisagg::test_symmetric_executor[gpt-2proc-nixl_kvcache-90] SKIP (https://nvbugs/6093820) -cpp/test_multi_gpu.py::TestDisagg::test_symmetric_executor[gpt-2proc-ucx_kvcache-90] SKIP (https://nvbugs/6093820) cpp/test_multi_gpu.py::test_cache_transceiver[8proc-mooncake_kvcache-90] SKIP (https://nvbugs/5838199) disaggregated/test_disaggregated.py::test_disaggregated_cancel_large_context_requests[DeepSeek-V3-Lite-bf16] SKIP (https://nvbugs/6105768) disaggregated/test_disaggregated.py::test_disaggregated_ctxpp2_genpp2[TinyLlama-1.1B-Chat-v1.0] SKIP (https://nvbugs/6427411) @@ -437,7 +427,6 @@ unittest/_torch/visual_gen/test_attention_integration.py::test_sage_attention_se unittest/_torch/visual_gen/test_attention_integration.py::test_sage_attention_self_attention[int8-2-1560] SKIP (https://nvbugs/6198760) unittest/_torch/visual_gen/test_wan21_i2v_pipeline.py::TestWanI2VBatchGeneration::test_batch_prompt_shape SKIP (https://nvbugs/6418822) unittest/auto_deploy/multigpu/custom_ops SKIP (https://nvbugs/6403920) -unittest/bindings/test_executor_bindings.py SKIP (TRTLLM-13781: legacy TensorRT examples removed; tests to be removed in follow-up PR3) unittest/disaggregated/test_kv_transfer.py::test_transfer_worker_v2[tp4_pp1_to_tp2_pp2] SKIP (https://nvbugs/6426834) unittest/executor/test_rpc.py::TestRpcCorrectness::test_incremental_task_async SKIP (https://nvbugs/5741476) unittest/executor/test_rpc_proxy.py SKIP (https://nvbugs/5605741) 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 48b650c3ba1f..6da6cd6e9c0a 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,8 +146,8 @@ 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", "cpp", "prepare_dataset.py") - script_dir = Path(root_dir, "benchmarks", "cpp") + 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 = [ 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 e227aae954a9..f766576cd454 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,7 @@ 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", "cpp") + script_dir = Path(root_dir, "benchmarks") # Generate a small dataset to run a test - matching workload configuration command = [ diff --git a/tests/unittest/bindings/binding_test_utils.py b/tests/unittest/bindings/binding_test_utils.py deleted file mode 100644 index e34316005e57..000000000000 --- a/tests/unittest/bindings/binding_test_utils.py +++ /dev/null @@ -1,58 +0,0 @@ -import logging as _log -import os as _os -import pathlib as _pl -import subprocess as _sp -import sys as _sys -import typing as _tp - -import numpy as _np - - -def run_command(command: _tp.Sequence[str], - cwd: _pl.Path, - *, - shell=False, - env=None) -> None: - _log.info("Running: cd %s && %s", str(cwd), " ".join(command)) - _sp.check_call(command, cwd=cwd, shell=shell, env=env) - - -def prepare_model_tests( - llm_root: _pl.Path, - resource_path: _pl.Path, - model_name: str, - model_cache_arg=[], -): - scripts_dir = resource_path / "scripts" - python_exe = _sys.executable - model_env = {**_os.environ, "PYTHONPATH": f"examples/{model_name}"} - build_engines = [ - python_exe, - str(scripts_dir / f"build_{model_name}_engines.py") - ] + model_cache_arg - run_command(build_engines, cwd=llm_root, env=model_env) - - model_env["PYTHONPATH"] = "examples" - generate_expected_output = [ - python_exe, - str(scripts_dir / f"generate_expected_{model_name}_output.py") - ] - run_command(generate_expected_output, cwd=llm_root, env=model_env) - - -def prepare_lora_configs(llm_root: _pl.Path, resource_path: _pl.Path, - lora_config_path: _pl.Path): - python_exe = _sys.executable - generate_lora_data_args_tp1 = [ - python_exe, - str(resource_path / "scripts" / "generate_test_lora_weights.py"), - f"--out-dir={str(lora_config_path)}", "--tp-size=1", - "--hidden-size=768", "--num-layers=12", "--config-ids-filter=0", - "--no-generate-cache-pages" - ] - run_command(generate_lora_data_args_tp1, cwd=llm_root) - - -def sequence_lengths(sequences: _np.ndarray, pad_id: int) -> _np.ndarray: - return _np.apply_along_axis(lambda x: _np.searchsorted(x, True), 1, - sequences == pad_id).astype("int32") diff --git a/tests/unittest/bindings/test_datatype_parity.py b/tests/unittest/bindings/test_datatype_parity.py new file mode 100644 index 000000000000..d3ca39f7d962 --- /dev/null +++ b/tests/unittest/bindings/test_datatype_parity.py @@ -0,0 +1,60 @@ +# 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. +"""Serialization/format-compatibility guard for the internal data-type enum. + +``tensorrt_llm.bindings.DataType`` is backed by ``tensorrt_llm::DataType`` +(``common/tllmDataType.h``), currently an alias of ``nvinfer1::DataType`` and +slated to become a standalone enum when the TensorRT-engine execution path is +removed. Its enumerator *integer values* must continue to match the legacy +``nvinfer1::DataType`` values so that previously-serialized executor configs +and KV-cache metadata remain byte-compatible, and so the public Python member +set is unchanged. +""" + +import pytest + +bindings = pytest.importorskip("tensorrt_llm.bindings") + +# Legacy nvinfer1::DataType integer values. These MUST remain stable. +LEGACY_DATATYPE_VALUES = { + "FLOAT": 0, + "HALF": 1, + "INT8": 2, + "INT32": 3, + "BOOL": 4, + "UINT8": 5, + "FP8": 6, + "BF16": 7, + "INT64": 8, + "NVFP4": 10, +} + + +def test_bindings_datatype_values_match_legacy(): + """Each exposed DataType enumerator keeps its legacy numeric value.""" + dt = bindings.DataType + for name, value in LEGACY_DATATYPE_VALUES.items(): + assert hasattr(dt, name), f"tensorrt_llm.bindings.DataType is missing '{name}'" + assert getattr(dt, name).value == value, ( + f"DataType.{name} numeric value changed to {getattr(dt, name).value}; " + f"expected {value} (breaks serialization compatibility)" + ) + + +def test_bindings_datatype_member_set_unchanged(): + """The public DataType member set still contains all legacy members.""" + members = {m for m in dir(bindings.DataType) if not m.startswith("_")} + missing = set(LEGACY_DATATYPE_VALUES) - members + assert not missing, f"tensorrt_llm.bindings.DataType lost members: {sorted(missing)}" diff --git a/tests/unittest/bindings/test_executor_bindings.py b/tests/unittest/bindings/test_executor_bindings.py index 438101bba3da..324126d02b84 100644 --- a/tests/unittest/bindings/test_executor_bindings.py +++ b/tests/unittest/bindings/test_executor_bindings.py @@ -1,854 +1,19 @@ import datetime +import inspect import json import os as _os import pickle -import random -import sys as _sys -import time import typing as tp -from pathlib import Path import numpy as np import pytest import torch -from binding_test_utils import * from pydantic import BaseModel import tensorrt_llm.bindings.executor as trtllm import tensorrt_llm.version as trtllm_version -from tensorrt_llm._utils import torch_to_numpy from tensorrt_llm.models.modeling_utils import PretrainedConfig -_sys.path.append(_os.path.join(_os.path.dirname(__file__), '..')) -import inspect - -from utils.cpp_paths import * -from utils.llm_data import llm_models_root -from utils.util import skip_pre_hopper - - -@pytest.fixture -def model_files(llm_root: Path, resource_path: Path, results_data_path: Path): - # Model engines and expected outputs need to be generated. - if not results_data_path.exists(): - model_cache = llm_models_root() - model_cache_arg = ["--model_cache", str(model_cache) - ] if model_cache is not None else [] - prepare_model_tests(llm_root, resource_path, "gpt", model_cache_arg) - - -@pytest.fixture -def lora_config_paths(llm_root: Path, resource_path: Path, - lora_config_path: Path): - if not lora_config_path.exists(): - prepare_lora_configs(llm_root, resource_path, lora_config_path) - return (lora_config_path / "source.npy", lora_config_path / "config.npy") - - -def get_expected_num_tokens(prompt_len, max_tokens, streaming, - exclude_input_from_output): - if not streaming and not exclude_input_from_output: - return prompt_len + max_tokens - return max_tokens - - -def test_executor_valid_ctor(model_files, model_path): - executor_config = trtllm.ExecutorConfig( - 1, kv_cache_config=trtllm.KvCacheConfig(free_gpu_memory_fraction=0.5)) - executor = trtllm.Executor(model_path, trtllm.ModelType.DECODER_ONLY, - executor_config) - - -def test_executor_from_memory(model_files, model_path): - executor_config = trtllm.ExecutorConfig( - 1, kv_cache_config=trtllm.KvCacheConfig(free_gpu_memory_fraction=0.5)) - engine_buffer = open(model_path / "rank0.engine", mode="rb").read() - json_config_str = open(model_path / "config.json", 'r').read() - executor = trtllm.Executor(engine_buffer, json_config_str, - trtllm.ModelType.DECODER_ONLY, executor_config) - - -def test_executor_with_managed_weights(model_files, model_path): - """Test executor constructor with standard dtypes in managed weights.""" - - executor_config = trtllm.ExecutorConfig( - 1, kv_cache_config=trtllm.KvCacheConfig(free_gpu_memory_fraction=0.5)) - engine_buffer = open(model_path / "rank0.engine", mode="rb").read() - json_config_str = open(model_path / "config.json", 'r').read() - - managed_weights = { - "weight_float32": - np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32), - "weight_int32": - np.array([[1, 2], [3, 4]], dtype=np.int32), - "weight_int64": - np.array([[1, 2], [3, 4]], dtype=np.int64), - "weight_int8": - np.array([[1, 2], [3, 4]], dtype=np.int8), - "weight_fp16": - np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float16), - "weight_bf16": - torch_to_numpy( - torch.tensor([[1.0, 2.0], [3.0, 4.0]], dtype=torch.bfloat16)), - "weight_fp8": - torch_to_numpy( - torch.tensor([[1.0, 2.0], [3.0, 4.0]], dtype=torch.float8_e4m3fn)), - } - - executor = trtllm.Executor(engine_buffer, json_config_str, - trtllm.ModelType.DECODER_ONLY, executor_config, - managed_weights) - - assert executor.can_enqueue_requests() == True - - -def test_executor_invalid_ctor(): - executor_config = trtllm.ExecutorConfig( - 1, kv_cache_config=trtllm.KvCacheConfig(free_gpu_memory_fraction=0.5)) - invalid_path = "Bla" - try: - executor = trtllm.Executor(invalid_path, trtllm.ModelType.DECODER_ONLY, - executor_config) - assert False, "Expected an error" - except Exception as e: - assert "File does not exist" in str(e) - - -def test_shutdown(model_files, model_path): - - beam_width = 1 - executor_config = trtllm.ExecutorConfig( - beam_width, - kv_cache_config=trtllm.KvCacheConfig(free_gpu_memory_fraction=0.5)) - executor = trtllm.Executor(model_path, trtllm.ModelType.DECODER_ONLY, - executor_config) - - # Create the request - max_tokens = 5 - input_tokens = [1, 2, 3, 4] - request = trtllm.Request(input_tokens, - max_tokens=max_tokens, - streaming=False, - sampling_config=trtllm.SamplingConfig()) - - # Enqueue the request - assert executor.can_enqueue_requests() == True - req_id = executor.enqueue_request(request) - - executor.shutdown() - assert executor.can_enqueue_requests() == False - - with pytest.raises(Exception): - executor.enqueue_request(request) - with pytest.raises(Exception): - executor.await_responses() - with pytest.raises(Exception): - executor.get_latest_iteration_stats() - with pytest.raises(Exception): - executor.get_latest_request_stats() - with pytest.raises(Exception): - executor.get_latest_debug_tensors() - with pytest.raises(Exception): - executor.cancel_request(req_id) - with pytest.raises(Exception): - executor.get_num_responses_ready(req_id) - - -def test_embedding_bias(model_files, model_path): - streaming = False - exclude_input_from_output = False - output_config = trtllm.OutputConfig() - output_config.exclude_input_from_output = exclude_input_from_output - - # Create executor - beam_width = 1 - executor_config = trtllm.ExecutorConfig( - beam_width, - kv_cache_config=trtllm.KvCacheConfig(free_gpu_memory_fraction=0.5)) - executor = trtllm.Executor(model_path, trtllm.ModelType.DECODER_ONLY, - executor_config) - - # Create the request - max_tokens = 5 - input_tokens = [1, 2, 3, 4] - # Set embedding bias so "biased_output" is always picked - biased_output = 10 - vocab_size_padded = 50257 - embedding_bias = torch.zeros(vocab_size_padded) - embedding_bias[biased_output] = torch.finfo(torch.float32).max - request = trtllm.Request(input_tokens, - max_tokens=max_tokens, - streaming=streaming, - sampling_config=trtllm.SamplingConfig(), - output_config=output_config, - embedding_bias=embedding_bias) - - # Enqueue the request - request_id = executor.enqueue_request(request) - - # Get the new tokens - tokens = [] - done = False - i = 0 - max_wait_ms = 10000 - while not done and i < max_wait_ms: - wait_time = datetime.timedelta(milliseconds=1) - responses = executor.await_responses(request_id, wait_time) - for response in responses: - assert not response.has_error( - ), f"Request id {request_id} failed with err {response.error_msg}" - result = response.result - done = result.is_final - new_tokens = result.output_token_ids[beam_width - 1] - tokens.extend(new_tokens) - i += 1 - assert i < max_wait_ms - assert len(tokens) == get_expected_num_tokens( - len(input_tokens), max_tokens, streaming, - exclude_input_from_output), f"{request_id}" - # All generated tokens should equal biased_output - assert tokens[-max_tokens:] == [biased_output] * max_tokens - - -@pytest.mark.parametrize("streaming", [False, True]) -@pytest.mark.parametrize("exclude_input_from_output", [False]) -def test_single_request(streaming: bool, exclude_input_from_output: bool, - model_files, model_path): - output_config = trtllm.OutputConfig() - output_config.exclude_input_from_output = exclude_input_from_output - - # Create executor - beam_width = 1 - executor_config = trtllm.ExecutorConfig( - beam_width, - kv_cache_config=trtllm.KvCacheConfig(free_gpu_memory_fraction=0.5)) - executor = trtllm.Executor(model_path, trtllm.ModelType.DECODER_ONLY, - executor_config) - - # Create the request - max_tokens = 5 - input_tokens = [1, 2, 3, 4] - request = trtllm.Request(input_tokens, - max_tokens=max_tokens, - streaming=streaming, - sampling_config=trtllm.SamplingConfig(), - output_config=output_config) - - # Enqueue the request - request_id = executor.enqueue_request(request) - - # Get the new tokens - tokens = [] - done = False - i = 0 - max_wait_ms = 10000 - while not done and i < max_wait_ms: - wait_time = datetime.timedelta(milliseconds=1) - responses = executor.await_responses(request_id, wait_time) - for response in responses: - assert not response.has_error( - ), f"Request id {request_id} failed with err {response.error_msg}" - result = response.result - done = result.is_final - new_tokens = result.output_token_ids[beam_width - 1] - tokens.extend(new_tokens) - i += 1 - assert i < max_wait_ms - assert len(tokens) == get_expected_num_tokens( - len(input_tokens), max_tokens, streaming, - exclude_input_from_output), f"{request_id}" - - executor.get_latest_iteration_stats() - executor.get_latest_request_stats() - executor.get_latest_debug_tensors() - - -def test_single_request_lora(model_files, model_path_lora, lora_config_paths): - streaming = False - exclude_input_from_output = False - output_config = trtllm.OutputConfig() - output_config.exclude_input_from_output = exclude_input_from_output - - # Create executor - beam_width = 1 - - peft_cache_config = trtllm.PeftCacheConfig(num_put_workers=4, - num_ensure_workers=4) - executor_config = trtllm.ExecutorConfig( - 1, - peft_cache_config=peft_cache_config, - kv_cache_config=trtllm.KvCacheConfig(free_gpu_memory_fraction=0.5)) - executor = trtllm.Executor(model_path_lora, trtllm.ModelType.DECODER_ONLY, - executor_config) - - # Create the request - max_tokens = 5 - input_tokens = [1, 2, 3, 4] - lora_weights = torch.tensor(np.load(lora_config_paths[0])).half() - lora_config = torch.tensor(np.load(lora_config_paths[1])) - request = trtllm.Request(input_tokens, - max_tokens=max_tokens, - streaming=streaming, - sampling_config=trtllm.SamplingConfig(), - output_config=output_config, - lora_config=trtllm.LoraConfig( - 0, lora_weights, lora_config)) - - # Enqueue the request - request_id = executor.enqueue_request(request) - - # Get the new tokens - tokens = [] - done = False - i = 0 - max_wait_ms = 10000 - while not done and i < max_wait_ms: - wait_time = datetime.timedelta(milliseconds=1) - responses = executor.await_responses(request_id, wait_time) - for response in responses: - assert not response.has_error( - ), f"Request id {request_id} failed with err {response.error_msg}" - result = response.result - done = result.is_final - new_tokens = result.output_token_ids[beam_width - 1] - tokens.extend(new_tokens) - i += 1 - assert i < max_wait_ms - assert len(tokens) == get_expected_num_tokens( - len(input_tokens), max_tokens, streaming, - exclude_input_from_output), f"{request_id}" - - -@pytest.mark.parametrize("streaming", [False, True]) -@pytest.mark.parametrize("exclude_input_from_output", [False]) -def test_multi_request(streaming: bool, exclude_input_from_output: bool, - model_files, model_path): - output_config = trtllm.OutputConfig() - output_config.exclude_input_from_output = exclude_input_from_output - - # Create executor - beam_width = 1 - executor_config = trtllm.ExecutorConfig( - beam_width, - kv_cache_config=trtllm.KvCacheConfig(free_gpu_memory_fraction=0.5)) - executor = trtllm.Executor(model_path, trtllm.ModelType.DECODER_ONLY, - executor_config) - - num_requests = 20 - max_prompt_len = 20 - max_max_tokens = 20 - end_id = -1 - - # Enqueue the requests - tokens = {} - expected_num_tokens = {} - for i in range(num_requests): - prompt_len = random.randint(1, max_prompt_len) - max_tokens = random.randint(1, max_max_tokens) - input_tokens = [1] * prompt_len - - # Some requests has num_return_sequences > 1. - num_return_sequences = 2 if i % 5 == 1 else 1 - - request = trtllm.Request(input_tokens, - max_tokens=max_tokens, - streaming=streaming, - sampling_config=trtllm.SamplingConfig( - num_return_sequences=num_return_sequences), - output_config=output_config, - end_id=end_id) - request_id = executor.enqueue_request(request) - tokens[request_id] = [ - [] for _ in range(request.sampling_config.num_return_sequences) - ] - expected_num_tokens[request_id] = get_expected_num_tokens( - prompt_len, max_tokens, streaming, exclude_input_from_output) - - # Get the new tokens for each request - num_finished = 0 - i = 0 - num_responses = 0 - max_wait_ms = 10000 - while num_finished < num_requests and i < max_wait_ms: - wait_time = datetime.timedelta(milliseconds=1) - responses = executor.await_responses(wait_time) - for response in responses: - num_responses += 1 - assert not response.has_error( - ), f"Request id {response.request_id} failed with err {response.error_msg}" - result = response.result - num_finished += result.is_final - new_tokens = result.output_token_ids[beam_width - 1] - tokens[response.request_id][result.sequence_index].extend( - new_tokens) - i += 1 - assert i < max_wait_ms - - for request_id in expected_num_tokens: - for actual_tokens in tokens[request_id]: - assert len(actual_tokens) == expected_num_tokens[request_id] - - -@pytest.mark.parametrize("streaming", [False, True]) -@pytest.mark.parametrize("exclude_input_from_output", [False]) -def test_multi_request_with_ids(streaming: bool, - exclude_input_from_output: bool, model_files, - model_path): - output_config = trtllm.OutputConfig() - output_config.exclude_input_from_output = exclude_input_from_output - - # Create executor - beam_width = 1 - executor_config = trtllm.ExecutorConfig( - beam_width, - kv_cache_config=trtllm.KvCacheConfig(free_gpu_memory_fraction=0.5)) - executor = trtllm.Executor(model_path, trtllm.ModelType.DECODER_ONLY, - executor_config) - - num_requests = 20 - max_prompt_len = 20 - max_max_tokens = 20 - end_id = -1 - - # Enqueue the requests - tokens = {} - expected_num_tokens = {} - for i in range(num_requests): - prompt_len = random.randint(1, max_prompt_len) - max_tokens = random.randint(1, max_max_tokens) - input_tokens = [1] * prompt_len - num_return_sequences = 2 if i % 5 == 1 else 1 - - request = trtllm.Request(input_tokens, - max_tokens=max_tokens, - streaming=streaming, - sampling_config=trtllm.SamplingConfig( - num_return_sequences=num_return_sequences), - output_config=output_config, - end_id=end_id) - request_id = executor.enqueue_request(request) - tokens[request_id] = [ - [] for _ in range(request.sampling_config.num_return_sequences) - ] - expected_num_tokens[request_id] = get_expected_num_tokens( - prompt_len, max_tokens, streaming, exclude_input_from_output) - - # Get the new tokens for each request - num_finished = 0 - i = 0 - num_responses = 0 - max_wait_ms = 10000 - while num_finished < num_requests and i < max_wait_ms: - wait_time = datetime.timedelta(milliseconds=1) - id_responses = executor.await_responses(list(tokens.keys()), wait_time) - for responses in id_responses: - for response in responses: - num_responses += 1 - # Allow response with error only if await_response processed a terminated request id - if response.has_error(): - terminated_request_error = "ReqId " + str( - response.request_id - ) + " has already been processed and was terminated." - assert response.error_msg == terminated_request_error, ( - f"Request id {response.request_id} failed with err " - f"{response.error_msg}") - else: - result = response.result - num_finished += result.is_final - new_tokens = result.output_token_ids[beam_width - 1] - tokens[response.request_id][result.sequence_index].extend( - new_tokens) - i += 1 - assert i < max_wait_ms - - for request_id in expected_num_tokens: - for seq_idx, actual_tokens in enumerate(tokens[request_id]): - assert len(actual_tokens) == expected_num_tokens[request_id] - - -@pytest.mark.parametrize("streaming", [False, True]) -@pytest.mark.parametrize("exclude_input_from_output", [False]) -def test_get_num_responses_ready(streaming: bool, - exclude_input_from_output: bool, model_files, - model_path): - output_config = trtllm.OutputConfig() - output_config.exclude_input_from_output = exclude_input_from_output - - # Create executor - executor_config = trtllm.ExecutorConfig( - 1, kv_cache_config=trtllm.KvCacheConfig(free_gpu_memory_fraction=0.5)) - executor = trtllm.Executor(model_path, trtllm.ModelType.DECODER_ONLY, - executor_config) - - max_prompt_len = 20 - max_max_tokens = 20 - - # Enqueue the requests - num_requests = random.randint(1, 50) - num_expected_responses = 0 - req_num_expected_responses = {} - for i in range(num_requests): - prompt_len = random.randint(1, max_prompt_len) - max_tokens = random.randint(1, max_max_tokens) - num_return_sequences = 2 if i % 5 == 1 else 1 - - request = trtllm.Request([1] * prompt_len, - max_tokens=max_tokens, - streaming=streaming, - sampling_config=trtllm.SamplingConfig( - num_return_sequences=num_return_sequences), - output_config=output_config) - request_id = executor.enqueue_request(request) - req_num_expected_responses[request_id] = ( - (max_tokens if streaming else 1) * num_return_sequences) - num_expected_responses += req_num_expected_responses[request_id] - - i = 0 - num_ready = 0 - max_wait_ms = 10000 - while num_ready < num_expected_responses and i < max_wait_ms: - num_ready = 0 - for request_id in req_num_expected_responses: - num_ready += executor.get_num_responses_ready(request_id) - time.sleep(0.001) - i += 1 - assert i < max_wait_ms - - for request_id in req_num_expected_responses: - num_ready = executor.get_num_responses_ready(request_id) - assert num_ready == req_num_expected_responses[request_id] - assert executor.get_num_responses_ready() == num_expected_responses - - -@pytest.mark.parametrize("batching_type", [trtllm.BatchingType.INFLIGHT]) -@pytest.mark.parametrize("streaming", [False, True]) -@pytest.mark.parametrize("beam_width", [1]) -@pytest.mark.parametrize("compute_log_probs", [False, True]) -@pytest.mark.parametrize("exclude_input_from_output", [False]) -@pytest.mark.parametrize("return_context_logits", [False, True]) -@pytest.mark.parametrize("return_generation_logits", [False, True]) -def test_token_comparison(batching_type: trtllm.BatchingType, streaming: bool, - beam_width: int, compute_log_probs: bool, - exclude_input_from_output: bool, - return_context_logits: bool, - return_generation_logits: bool, model_files, - model_path, model_path_return_logits, input_data_path, - results_data_path, results_data_path_beam_width_2): - if streaming and beam_width > 1: - pytest.skip("Test does not support streaming with beam search") - - vocab_size_padded = 50257 - pad_id = 50256 - remove_input = not exclude_input_from_output and not streaming - - def load_test_data(input_path, results_path): - # Inputs - assert input_path.is_file() - given_input = np.load(input_path).astype("int32") - input_shape = given_input.shape - assert len(input_shape) == 2 - max_input_length = input_shape[1] - given_input_lengths = sequence_lengths(given_input, pad_id) - assert np.all(given_input_lengths <= max_input_length) - # Expected results - assert results_path.is_file() - expected_outputs = np.load(results_path).astype("int32") - output_shape = expected_outputs.shape - assert len(output_shape) == 2 - assert input_shape[0] * beam_width == output_shape[0] - max_seq_length = output_shape[1] - max_tokens = max_seq_length - max_input_length - - end_ids = [pad_id for _ in range(len(given_input_lengths))] - expected_lengths = [] - for i in range(len(given_input_lengths)): - expected_lengths.append([ - given_input_lengths[i] + max_tokens for _ in range(beam_width) - ]) - - test_data = { - "expected_output_ids": expected_outputs, - "expected_output_lengths": expected_lengths, - "max_seq_length": max_seq_length, - "end_ids": end_ids - } - return given_input, given_input_lengths, max_input_length, test_data - - def validate_results_shapes(result, input_length, max_output_len, - beam_tokens): - if compute_log_probs: - assert result.cum_log_probs is not None - assert result.log_probs is not None - assert len(result.cum_log_probs) == beam_width - assert len(result.log_probs) == beam_width - for beam in range(beam_width): - expected_len = len( - beam_tokens[beam]) - (input_length if remove_input else 0) - assert len(result.log_probs[beam]) == expected_len - else: - assert result.cum_log_probs is None - assert result.log_probs is None - if return_context_logits: - assert result.context_logits is not None - assert len(result.context_logits.shape) == 2 - assert list(result.context_logits.shape) == [ - input_length, vocab_size_padded - ] - else: - assert result.context_logits is None - if return_generation_logits: - assert len(result.generation_logits.shape) == 3 - if streaming: - assert list(result.generation_logits.shape) == [ - max_output_len, beam_width, vocab_size_padded - ] or list(result.generation_logits.shape) == [ - 1, beam_width, vocab_size_padded - ] - else: - assert list(result.generation_logits.shape) == [ - beam_width, max_output_len, vocab_size_padded - ] - - def verify_output(beam_tokens, test_data, given_input_lengths): - - for batch_id, seq_tokens in beam_tokens.items(): - input_length = given_input_lengths[batch_id] - end_id = test_data["end_ids"][batch_id] - for tokens in seq_tokens: - for beam in range(beam_width): - - predicted_tokens = tokens[beam] - if remove_input: - predicted_tokens = predicted_tokens[input_length:] - expected_length = test_data["expected_output_lengths"][ - batch_id][beam] - input_length - assert len(predicted_tokens) == expected_length - - expected_tokens = test_data["expected_output_ids"][ - batch_id * beam_width + beam][input_length:] - - # From experiments find out when set return_context_logits - # or return_generation_logits, the predicted_tokens cannot match with expected_tokens - # Fixed by comparing partial output tokens like in c++ test - compare_length = 2 if ( - return_context_logits - or return_generation_logits) else len(predicted_tokens) - - for i in range(compare_length): - if expected_tokens[i] == end_id: - break - # Predicted: [21221, 290, 373, 257, 2888, 286, 262, 4141] - # Expected: [21221, 290, 257, 4255, 379, 262, 1957, 7072] - # generation logits are almost same at token ids 257 and 373, - # which causes unstable generation results. - assert predicted_tokens[i] == expected_tokens[i], \ - f"Predicted: {predicted_tokens} vs Expected: {expected_tokens}" - - output_config = trtllm.OutputConfig() - output_config.exclude_input_from_output = exclude_input_from_output - output_config.return_log_probs = compute_log_probs - output_config.return_generation_logits = return_generation_logits - output_config.return_context_logits = return_context_logits - # Change free_gpu_memory_fraction to solve OOM error - kv_cache_config = trtllm.KvCacheConfig(False, free_gpu_memory_fraction=0.3) - executor_config = trtllm.ExecutorConfig(beam_width) - executor_config.batching_type = batching_type - executor_config.kv_cache_config = kv_cache_config - if return_generation_logits: - executor_config.gather_generation_logits = True - - if return_context_logits or return_generation_logits: - model_path = model_path_return_logits - executor = trtllm.Executor(model_path, trtllm.ModelType.DECODER_ONLY, - executor_config) - - # Load test data - results_path = results_data_path if beam_width == 1 else results_data_path_beam_width_2 - given_input, given_input_lengths, max_input_length, test_data = load_test_data( - input_data_path, results_path) - - # Create requests from input data - num_requests = len(given_input_lengths) - requests = [] - req_max_tokens = [] - - for i in range(num_requests): - input_len = given_input_lengths[i] - max_tokens = test_data["max_seq_length"] - max_input_length - req_max_tokens.append(max_tokens) - req_tokens = given_input[i][:input_len] - num_return_sequences = 2 if i % 5 == 1 else 1 - request = trtllm.Request(req_tokens, - max_tokens=max_tokens, - streaming=streaming, - sampling_config=trtllm.SamplingConfig( - beam_width, - num_return_sequences=num_return_sequences), - output_config=output_config, - end_id=-1) - requests.append(request) - - req_ids = executor.enqueue_requests(requests) - - req_to_batch_id = {req_ids[i]: i for i in range(len(requests))} - tokens = { - i: [[[] for _ in range(beam_width)] - for _ in range(req.sampling_config.num_return_sequences)] - for i, req in enumerate(requests) - } - - num_finished = 0 - i = 0 - num_responses = 0 - max_wait_ms = 10000 - while num_finished < num_requests and i < max_wait_ms: - wait_time = datetime.timedelta(milliseconds=1) - responses = executor.await_responses(wait_time) - for response in responses: - num_responses += 1 - assert not response.has_error( - ), f"Request id {response.request_id} failed with err {response.error_msg}" - result = response.result - num_finished += result.is_final - - batch_id = req_to_batch_id[response.request_id] - for beam in range(beam_width): - new_tokens = result.output_token_ids[beam] - tokens[batch_id][result.sequence_index][beam] += new_tokens - - validate_results_shapes(result, given_input_lengths[batch_id], - req_max_tokens[batch_id], - tokens[batch_id][result.sequence_index]) - i += 1 - assert i < max_wait_ms - verify_output(tokens, test_data, given_input_lengths) - - -@pytest.mark.parametrize("streaming", [False, True]) -@pytest.mark.parametrize("beam_width", [1]) -def test_finish_reason(streaming: bool, beam_width: int, model_files, - model_path): - if streaming and beam_width > 1: - pytest.skip("Test does not support streaming with beam search") - executor = trtllm.Executor( - model_path, trtllm.ModelType.DECODER_ONLY, - trtllm.ExecutorConfig( - beam_width, - kv_cache_config=trtllm.KvCacheConfig(free_gpu_memory_fraction=0.5))) - requests = [ - # Finish due to length. - trtllm.Request([1, 2, 3, 4], - max_tokens=5, - streaming=streaming, - sampling_config=trtllm.SamplingConfig(beam_width)), - # Finish due to end id. - trtllm.Request([1, 2, 3, 4], - max_tokens=5, - streaming=streaming, - sampling_config=trtllm.SamplingConfig(beam_width), - end_id=4), - # Finish due to stop word. - trtllm.Request([1, 2, 3, 4], - max_tokens=5, - streaming=streaming, - sampling_config=trtllm.SamplingConfig(beam_width), - stop_words=[[4, 2]]), - ] - req_ids = executor.enqueue_requests(requests) - req_to_batch_id = {req_ids[i]: i for i in range(len(requests))} - - num_finished = 0 - i = 0 - num_responses = 0 - max_wait_ms = 10000 - while num_finished < len(requests) and i < max_wait_ms: - wait_time = datetime.timedelta(milliseconds=1) - responses = executor.await_responses(wait_time) - for response in responses: - num_responses += 1 - assert not response.has_error( - ), f"Request id {response.request_id} failed with err {response.error_msg}" - result = response.result - num_finished += result.is_final - batch_id = req_to_batch_id[response.request_id] - - # Non final results should have "NOT_FINISHED". Revise this when streaming + beam_width > 1 is enabled. - if not result.is_final: - assert all([ - r == trtllm.FinishReason.NOT_FINISHED - for r in result.finish_reasons - ]) - # Check if finish reason is correct. - elif batch_id == 0: - assert all([ - r == trtllm.FinishReason.LENGTH - for r in result.finish_reasons - ]) - elif batch_id == 1: - assert all([ - r == trtllm.FinishReason.END_ID - for r in result.finish_reasons - ]) - elif batch_id == 2: - assert all([ - r == trtllm.FinishReason.STOP_WORDS - for r in result.finish_reasons - ]) - i += 1 - assert i < max_wait_ms - - -def test_gpt_executor_timed_out(model_files, model_path): - beam_width = 1 - executor_config = trtllm.ExecutorConfig( - beam_width, - kv_cache_config=trtllm.KvCacheConfig(free_gpu_memory_fraction=0.5)) - executor = trtllm.Executor(model_path, trtllm.ModelType.DECODER_ONLY, - executor_config) - - # No requests enqueued, expect no responses - num_responses_ready = executor.get_num_responses_ready() - assert num_responses_ready == 0 - - wait_time = datetime.timedelta(milliseconds=10) - responses = executor.await_responses(wait_time) - assert len(responses) == 0 - - -def test_single_request_invalid_inputs(model_files, model_path): - streaming = True - beam_width = 1 - executor_config = trtllm.ExecutorConfig( - beam_width, - kv_cache_config=trtllm.KvCacheConfig(free_gpu_memory_fraction=0.5)) - executor = trtllm.Executor(model_path, trtllm.ModelType.DECODER_ONLY, - executor_config) - - max_tokens = 5 - input_tokens = [1, 2, 3, 4] - request = trtllm.Request(input_tokens, - max_tokens=max_tokens, - streaming=streaming) - # Invalid embedding bias shape - embedding_bias = torch.ones(1) - request.embedding_bias = embedding_bias - expected_error_msg = "embedding bias shape is not as expected" - - request_id = executor.enqueue_request(request) - - done = False - i = 0 - max_wait_ms = 10000 - while not done and i < max_wait_ms: - wait_time = datetime.timedelta(milliseconds=1) - responses = executor.await_responses(request_id, wait_time) - for response in responses: - assert response.has_error(), "Expected an error" - assert expected_error_msg in response.error_msg - done = True - i += 1 - assert done - def test_sampling_config(): beam_width = 1 @@ -1149,36 +314,6 @@ def test_lora_config(): assert (lora_config.config == config).all() -def test_wakeup(model_files, model_path): - import threading - - def resp_thread(stop_signal: threading.Event, executor: trtllm.Executor): - while not stop_signal.is_set(): - timeout = None - responses = executor.await_responses(timeout=timeout) - if stop_signal.is_set(): - return - for response in responses: - response.result.output_token_ids - - executor = trtllm.Executor( - model_path, trtllm.ModelType.DECODER_ONLY, - trtllm.ExecutorConfig(kv_cache_config=trtllm.KvCacheConfig( - free_gpu_memory_fraction=0.5))) - stop_signal = threading.Event() - thread = threading.Thread(target=resp_thread, args=(stop_signal, executor)) - thread.start() - request = trtllm.Request(input_token_ids=[1, 2, 3, 4], - max_tokens=5, - streaming=True) - executor.enqueue_request(request) - time.sleep(2) - stop_signal.set() - executor.shutdown() - thread.join() - assert not thread.is_alive() - - def test_guided_decoding_params(): guided_decoding_params = trtllm.GuidedDecodingParams( trtllm.GuidedDecodingParams.GuideType.JSON) @@ -1698,7 +833,7 @@ def test_decoding_mode(): assert mode.isEagle() -def test_speculative_decoding_config(): +def test_decoding_config(): config = trtllm.DecodingConfig() assert config.decoding_mode is None assert config.lookahead_decoding_config is None @@ -1975,427 +1110,6 @@ def test_peft_cache_config(): assert peft_cache_config.lora_prefetch_dir == lora_prefetch_dir -def test_logits_post_processor(model_files, model_path): - - # Define the logits post-processor callback - def logits_post_processor(req_id: int, logits: torch.Tensor, - ids: tp.List[tp.List[int]], stream_ptr: int, - client_id: tp.Optional[int]): - assert client_id == 123 - with torch.cuda.stream(torch.cuda.ExternalStream(stream_ptr)): - logits[:] = float("-inf") - logits[..., 42] = 0 - - # Create executor - beam_width = 1 - executor_config = trtllm.ExecutorConfig( - beam_width, - kv_cache_config=trtllm.KvCacheConfig(free_gpu_memory_fraction=0.5)) - executor_config.logits_post_processor_config = trtllm.LogitsPostProcessorConfig( - {"my_logits_pp": logits_post_processor}) - executor = trtllm.Executor(model_path, trtllm.ModelType.DECODER_ONLY, - executor_config) - - # Create the request - max_tokens = 5 - input_tokens = [1, 2, 3, 4] - request = trtllm.Request(input_tokens, - max_tokens=max_tokens, - streaming=False, - client_id=123) - request.logits_post_processor_name = "my_logits_pp" - - # Enqueue the request - request_id = executor.enqueue_request(request) - - # Get the new tokens - tokens = [] - done = False - i = 0 - max_wait_ms = 10000 - while not done and i < max_wait_ms: - wait_time = datetime.timedelta(milliseconds=1) - responses = executor.await_responses(request_id, wait_time) - for response in responses: - assert not response.has_error( - ), f"Request id {request_id} failed with err {response.error_msg}" - result = response.result - done = result.is_final - new_tokens = result.output_token_ids[beam_width - 1] - tokens.extend(new_tokens) - i += 1 - assert i < max_wait_ms - assert len(tokens) == get_expected_num_tokens(len(input_tokens), max_tokens, - False, False), f"{request_id}" - - # check that all output tokens are 42 - print(tokens) - assert tokens[-max_tokens:] == [42] * max_tokens - - -def test_logits_post_processor_batched(model_files, model_path): - - # Define the logits post-processor callback - def logits_post_processor_batched( - req_id_batch: tp.List[int], logits_batch: tp.List[torch.Tensor], - ids_batch: tp.List[tp.List[tp.List[int]]], stream_ptr: int, - client_id_batch: tp.List[tp.Optional[int]]): - for client_id in client_id_batch: - assert client_id == 123 - with torch.cuda.stream(torch.cuda.ExternalStream(stream_ptr)): - for logits in logits_batch: - logits[:] = float("-inf") - logits[..., 42] = 0 - - # Create executor - beam_width = 1 - executor_config = trtllm.ExecutorConfig( - beam_width, - kv_cache_config=trtllm.KvCacheConfig(free_gpu_memory_fraction=0.5)) - executor_config.logits_post_processor_config = trtllm.LogitsPostProcessorConfig( - None, logits_post_processor_batched) - executor = trtllm.Executor(model_path, trtllm.ModelType.DECODER_ONLY, - executor_config) - - # Create the request - max_tokens = 5 - input_tokens = [1, 2, 3, 4] - request = trtllm.Request(input_tokens, - max_tokens=max_tokens, - streaming=False, - client_id=123) - request.logits_post_processor_name = request.BATCHED_POST_PROCESSOR_NAME - - batch_size = 4 - # Enqueue the requests - request_ids = [] - for _ in range(batch_size): - request_id = executor.enqueue_request(request) - request_ids.append(request_id) - - # Get the new tokens - tokens = {req_id: [] for req_id in request_ids} - num_finished = 0 - i = 0 - max_wait_ms = 10000 - while num_finished < len(request_ids) and i < max_wait_ms: - responses = executor.await_responses(datetime.timedelta(milliseconds=1)) - for response in responses: - req_id = response.request_id - assert not response.has_error( - ), f"Request id {req_id} failed with err {response.error_msg}" - result = response.result - num_finished += 1 if result.is_final else 0 - new_tokens = result.output_token_ids[beam_width - 1] - tokens[req_id].extend(new_tokens) - assert i < max_wait_ms - - expected_num_tokens = get_expected_num_tokens(len(input_tokens), max_tokens, - False, False) - for req_id in request_ids: - assert len(tokens[req_id]) == expected_num_tokens, f"{req_id}" - - -@pytest.mark.skip("https://nvbugs/5082576") -def test_kv_event_stream(model_path): - - beam_width = 1 - executor_config = trtllm.ExecutorConfig( - beam_width, - kv_cache_config=trtllm.KvCacheConfig(True, - 4 * 64, - event_buffer_max_size=1024, - host_cache_size=3000000, - free_gpu_memory_fraction=0.5)) - - executor = trtllm.Executor(model_path, trtllm.ModelType.DECODER_ONLY, - executor_config) - - cache_manager = executor.get_kv_cache_event_manager() - - events = cache_manager.get_latest_events() - - assert len(events) == 1 - assert isinstance(events[0], trtllm.kv_cache.KVCacheEvent) - assert events[0].event_id == 0 - assert isinstance(events[0].data, trtllm.kv_cache.KVCacheCreatedData) - - for req in range(2): - input_tokens = list(range(req, req + 127)) - request = trtllm.Request(input_tokens, - max_tokens=5, - streaming=False, - sampling_config=trtllm.SamplingConfig()) - - id = executor.enqueue_request(request) - - responses = executor.await_responses(id) - - for response in responses: - assert not response.has_error() - if response.result.is_final: - time.sleep(0.1) - events = cache_manager.get_latest_events( - datetime.timedelta(milliseconds=100)) - - if req == 0: - assert events[0].event_id == 1 - assert isinstance(events[0].data, - trtllm.kv_cache.KVCacheStoredData) - assert events[0].data.parent_hash is None - assert len(events[0].data.blocks) == 1 - - assert events[1].data.parent_hash == events[0].data.blocks[ - 0].block_hash - assert len(events[1].data.blocks) == 2 - else: - # Swap a block to secondary - assert isinstance(events[0].data, - trtllm.kv_cache.KVCacheUpdatedData) - assert events[0].data.cache_level.old_value == 0 - assert events[0].data.cache_level.new_value == 1 - # Store the filled context block - assert isinstance(events[1].data, - trtllm.kv_cache.KVCacheStoredData) - assert len(events[1].data.blocks) == 1 - assert events[1].data.parent_hash is None - # Swap another block to secondary - assert isinstance(events[2].data, - trtllm.kv_cache.KVCacheUpdatedData) - assert events[2].data.cache_level.old_value == 0 - assert events[2].data.cache_level.new_value == 1 - assert isinstance(events[2].data.cache_level, - trtllm.kv_cache.KVCacheEventDiffInt) - # Remove the first block in secondary - assert isinstance(events[3].data, - trtllm.kv_cache.KVCacheRemovedData) - assert len(events[3].data.block_hashes) == 1 - assert events[3].data.block_hashes[0] == events[ - 0].data.block_hash - # Store the second context block and the decode block - assert isinstance(events[4].data, - trtllm.kv_cache.KVCacheStoredData) - assert len(events[4].data.blocks) == 2 - assert events[4].data.parent_hash == events[1].data.blocks[ - 0].block_hash - - -@pytest.mark.parametrize("streaming", [False, True]) -def test_request_perf_metrics(streaming: bool, model_path): - - # Create executor - beam_width = 1 - executor_config = trtllm.ExecutorConfig( - beam_width, - kv_cache_config=trtllm.KvCacheConfig(free_gpu_memory_fraction=0.5)) - executor = trtllm.Executor(model_path, trtllm.ModelType.DECODER_ONLY, - executor_config) - - # Create request - max_tokens = 5 - input_tokens = [1, 2, 3, 4] - output_config = trtllm.OutputConfig(return_perf_metrics=True) - request = trtllm.Request(input_tokens, - max_tokens=max_tokens, - streaming=streaming, - output_config=output_config) - - # Enqueue the request - request_id = executor.enqueue_request(request) - - def check_perf_metrics(perf_metrics, done, response_id): - assert perf_metrics is not None - - timing_metrics = perf_metrics.timing_metrics - assert timing_metrics.arrival_time < timing_metrics.first_scheduled_time - assert timing_metrics.first_scheduled_time < timing_metrics.first_token_time - if done: - assert timing_metrics.first_token_time < timing_metrics.last_token_time - else: - assert timing_metrics.last_token_time == datetime.timedelta(0) - - kv_cache_metrics = perf_metrics.kv_cache_metrics - assert kv_cache_metrics.num_total_allocated_blocks == 1 - assert kv_cache_metrics.num_new_allocated_blocks == 1 - assert kv_cache_metrics.num_reused_blocks == 0 - assert kv_cache_metrics.num_missed_blocks == 1 - assert kv_cache_metrics.kv_cache_hit_rate == 0 - - assert perf_metrics.first_iter == 0 - if done: - assert perf_metrics.iter == (max_tokens - 1) - assert perf_metrics.last_iter == max_tokens - 1 - else: - assert perf_metrics.iter == response_id - assert perf_metrics.last_iter is None - - # Get the new tokens - tokens = [] - done = False - i = 0 - max_wait_ms = 10000 - response_id = 0 - while not done and i < max_wait_ms: - wait_time = datetime.timedelta(milliseconds=1) - responses = executor.await_responses(request_id, wait_time) - for response in responses: - assert not response.has_error( - ), f"Request id {request_id} failed with err {response.error_msg}" - result = response.result - done = result.is_final - check_perf_metrics(result.request_perf_metrics, done, response_id) - new_tokens = result.output_token_ids[beam_width - 1] - tokens.extend(new_tokens) - response_id += 1 - i += 1 - assert i < max_wait_ms - assert len(tokens) == get_expected_num_tokens( - len(input_tokens), - max_tokens, - streaming=streaming, - exclude_input_from_output=False), f"{request_id}" - - -def test_request_perf_metrics_kv_cache(model_path): - - # Create executor - beam_width = 1 - executor_config = trtllm.ExecutorConfig(beam_width) - executor = trtllm.Executor(model_path, trtllm.ModelType.DECODER_ONLY, - executor_config) - - # Create request: model uses 32 tokens per block, so it will fill a full block - max_tokens = 32 - input_tokens = [1, 2, 3, 4] - request = trtllm.Request(input_tokens, max_tokens=max_tokens) - - # Enqueue the request - request_id = executor.enqueue_request(request) - # Store two blocks with a total of 35 reusable tokens (4 input + 32 output, but last token is not stored) - - # Get the response - responses = executor.await_responses(request_id) - assert not responses[0].has_error() - result = responses[0].result - assert result.is_final - - # Prepare second request using the first output - input_tokens = result.output_token_ids[beam_width - 1] + [1, 2, 3, 4] - output_config = trtllm.OutputConfig(return_perf_metrics=True) - request = trtllm.Request(input_tokens, - max_tokens=max_tokens, - output_config=output_config) - - # Enqueue the request - # New query has 36 tokens that match input and output of first request plus 4 additional tokens. - # First block reused completely (32 tokens), from second block we'll partially reuse 3 tokens. - request_id = executor.enqueue_request(request) - - # Get the response - responses = executor.await_responses(request_id) - assert not responses[0].has_error() - result = responses[0].result - assert result.is_final - - # Check KV cache metric: Two blocks will be reused, although second block is partially reused. - # Cache hit rate is 100% since granularity is blocks, not tokens. - kv_cache_metrics = result.request_perf_metrics.kv_cache_metrics - assert kv_cache_metrics.num_total_allocated_blocks == 0 - assert kv_cache_metrics.num_new_allocated_blocks == 0 - assert kv_cache_metrics.num_reused_blocks == 2 - assert kv_cache_metrics.num_missed_blocks == 0 - assert kv_cache_metrics.kv_cache_hit_rate == 1.0 - - -# Skip test for pre-Hopper: https://nvbugs/5404000 -@skip_pre_hopper -@pytest.mark.parametrize("exclude_input_from_output", [False, True]) -def test_request_perf_metrics_draft(model_path_draft_tokens_external, - exclude_input_from_output: bool): - - # Create executor - beam_width = 1 - executor_config = trtllm.ExecutorConfig( - beam_width, - kv_cache_config=trtllm.KvCacheConfig(free_gpu_memory_fraction=0.5)) - executor = trtllm.Executor(model_path_draft_tokens_external, - trtllm.ModelType.DECODER_ONLY, executor_config) - - # Create request - max_tokens = 5 - input_tokens = [1, 2, 3, 4] - - # Only first two tokens will be accepted -> 50% acceptance rate - draft_config = trtllm.ExternalDraftTokensConfig([2, 4, 9, 10]) - output_config = trtllm.OutputConfig( - exclude_input_from_output=exclude_input_from_output, - return_perf_metrics=True) - request = trtllm.Request(input_tokens, - max_tokens=max_tokens, - output_config=output_config, - external_draft_tokens_config=draft_config) - - # Enqueue the request - request_id = executor.enqueue_request(request) - - # Get the response - responses = executor.await_responses(request_id) - assert not responses[0].has_error() - result = responses[0].result - assert result.is_final - - # check the new tokens - new_tokens = result.output_token_ids[beam_width - 1] - if exclude_input_from_output: - assert new_tokens == [2, 4, 2] - else: - assert new_tokens == [1, 2, 3, 4, 2, 4, 2] - - # Check the perf metrics - perf_metrics = result.request_perf_metrics - assert perf_metrics is not None - - timing_metrics = perf_metrics.timing_metrics - assert timing_metrics.arrival_time < timing_metrics.first_scheduled_time - assert timing_metrics.first_scheduled_time < timing_metrics.first_token_time - assert timing_metrics.first_token_time <= timing_metrics.last_token_time - - assert perf_metrics.first_iter == 0 - assert perf_metrics.iter == 0 - assert perf_metrics.last_iter == 0 - - spec_dec_metrics = perf_metrics.speculative_decoding - assert spec_dec_metrics.acceptance_rate == 0.5 - assert spec_dec_metrics.total_accepted_draft_tokens == 2 - assert spec_dec_metrics.total_draft_tokens == 4 - - -def test_kv_event_stream_timeout(model_path): - - beam_width = 1 - executor_config = trtllm.ExecutorConfig( - beam_width, - kv_cache_config=trtllm.KvCacheConfig(True, - 4 * 64, - event_buffer_max_size=1024, - free_gpu_memory_fraction=0.5)) - - executor = trtllm.Executor(model_path, trtllm.ModelType.DECODER_ONLY, - executor_config) - - cache_manager = executor.get_kv_cache_event_manager() - - events = cache_manager.get_latest_events() - assert len(events) == 1 - - start = datetime.datetime.now() - events = cache_manager.get_latest_events(1000) - end = datetime.datetime.now() - # Make sure that it actually waited - assert abs(end - start) > datetime.timedelta(milliseconds=900) - assert len(events) == 0 - - def test_request_perf_metrics_pickle(): metrics = trtllm.RequestPerfMetrics() random_delta = datetime.timedelta(seconds=42, milliseconds=123) diff --git a/tests/unittest/others/test_leak.py b/tests/unittest/others/test_leak.py deleted file mode 100644 index 819744fb74b9..000000000000 --- a/tests/unittest/others/test_leak.py +++ /dev/null @@ -1,88 +0,0 @@ -import unittest - -import tensorrt_llm as tllm -import tensorrt_llm.profiler as profiler - -import psutil # isort:skip - - -def create_model(): - ''' Lots of parameters are created here, and thus memory increases - ''' - profiler.print_memory_usage('Before creating Module') - - config = { - 'architecture': "LlamaForCausalLM", - 'dtype': 'float32', - 'num_hidden_layers': 2, - 'num_attention_heads': 80, - 'hidden_size': 12800, - 'num_key_value_heads': 80, - 'vocab_size': 50000, - 'position_embedding_type': 'rope_gpt_neox', - 'max_position_embeddings': 2048, - 'hidden_act': 'silu' - } - config = tllm.models.PretrainedConfig.from_dict(config) - - # About 24GiB model size, big enough to detect leak and avoid noise and false positive - # and small enough to make sure CI single-gpu machine can run it. - model = tllm.models.LLaMAForCausalLM.from_config(config) - profiler.print_memory_usage('After creating Module') - return model - - -def create_optimize_network(): - builder = tllm.Builder() - model = create_model() - network = builder.create_network() - network.plugin_config.gpt_attention_plugin = 'float16' - network.plugin_config.paged_kv_cache = True - network.plugin_config.use_fp8_context_fmha = False - profiler.print_memory_usage('Before creating Network') - with tllm.net_guard(network): - # Forward - inputs = model.prepare_inputs(max_batch_size=1, - max_input_len=1024, - max_seq_len=1024 + 32, - max_num_tokens=1024, - use_cache=True, - max_beam_width=1) - model(**inputs) - profiler.print_memory_usage('After creating Network') - - # When the Network has gpt attention plugin layer, graph rewriting pattern matching is triggered, - # thus the Network._get_graph_impl will be called, and a lru_cache will be created to cache this Network object - # and thus these registered ndarrays inside the Network, these objects are destroyed only when the cache is full or the - # program ends - tllm.graph_rewriting.optimize(network) - - -def run(): - # Create a TRT builder to warm up the memory, and avoid the noise of leak detection. - # Builder creation will create global objects like kernels. - _ = tllm.Builder() - - used, _, _ = profiler.host_memory_info() - - for i in range(5): - # Ideally the memory used inside create_optimize_network will all be released after the function returns - profiler.print_memory_usage(f'create_optimize_network {i} started') - create_optimize_network() - profiler.print_memory_usage(f'create_optimize_network {i} returned') - - used_after, _, _ = profiler.host_memory_info() - mem_increase_in_gb = (used_after - used) / (1024**3) - # The model has more than 10GB, so if there is leak, it will be absolutely bigger than 1GB - assert mem_increase_in_gb < 1, f"Memory increased {mem_increase_in_gb} GB" - - -class TestHostMemLeak(unittest.TestCase): - - def test_host_mem_leak(self): - tllm.logger.set_level('info') - run() - - -if __name__ == '__main__': - unittest.main() diff --git a/tests/unittest/tools/test_prepare_dataset.py b/tests/unittest/tools/test_prepare_dataset.py index 948cde1e0928..08bcb0ab1e38 100644 --- a/tests/unittest/tools/test_prepare_dataset.py +++ b/tests/unittest/tools/test_prepare_dataset.py @@ -16,7 +16,7 @@ _DEFAULT_OUTPUT_STDEV = 10 _TEST_TASK_IDS = [0, 1, 2] _TOKENIZER_SUBPATH = "llama-models-v2/tinyllama-tarot-v1/" -_PREPARE_DATASET_SCRIPT_PATH = "benchmarks/cpp/prepare_dataset.py" +_PREPARE_DATASET_SCRIPT_PATH = "benchmarks/prepare_dataset.py" class TestPrepareDatasetLora: diff --git a/tests/unittest/tools/test_test_to_stage_mapping.py b/tests/unittest/tools/test_test_to_stage_mapping.py index 65b27e755b44..0af0d468547e 100644 --- a/tests/unittest/tools/test_test_to_stage_mapping.py +++ b/tests/unittest/tools/test_test_to_stage_mapping.py @@ -182,17 +182,22 @@ def test_search_functionality(stage_query, sample_test_cases): @pytest.mark.parametrize('file_format', ['txt', 'yml']) -def test_cli_functionality(tmp_path, sample_test_cases, file_format): +def test_cli_functionality(tmp_path, stage_query, sample_test_cases, + file_format): """Test CLI functionality with sample data.""" - if not sample_test_cases: - pytest.skip("No test cases available") + # Use the first sample that maps to at least one stage (some test-db + # files, e.g. multi-node perf-sanity lists, have no L0 stage). + test_case = next( + (t for t in sample_test_cases if stage_query.tests_to_stages([t])), + None) + if test_case is None: + pytest.skip("No sampled test maps to any stage") - # Use only first sample for CLI test test_file = tmp_path / f'sample_tests.{file_format}' if file_format == 'txt': - test_file.write_text(f'{sample_test_cases[0]}\n') + test_file.write_text(f'{test_case}\n') else: # yml - test_file.write_text(f'- {sample_test_cases[0]}\n') + test_file.write_text(f'- {test_case}\n') script = os.path.join(SCRIPTS_DIR, 'test_to_stage_mapping.py') cmd = [sys.executable, script, '--test-list', str(test_file)] @@ -200,7 +205,7 @@ def test_cli_functionality(tmp_path, sample_test_cases, file_format): lines = output.decode().strip().splitlines() # Should return at least one stage - assert lines, f"No stages returned for test '{sample_test_cases[0]}'" + assert lines, f"No stages returned for test '{test_case}'" def test_backend_filtering_consistency(stage_query): diff --git a/tests/unittest/utils/cpp_paths.py b/tests/unittest/utils/cpp_paths.py index b147b5271138..feac0f698cab 100644 --- a/tests/unittest/utils/cpp_paths.py +++ b/tests/unittest/utils/cpp_paths.py @@ -4,9 +4,6 @@ import pytest -import tensorrt_llm.bindings as _tb -from tensorrt_llm.bindings.internal.testing import ModelSpec - _sys.path.append(_os.path.join(_os.path.dirname(__file__), '..', '..', '..')) @@ -15,70 +12,3 @@ def llm_root() -> _pl.Path: environ_root = _os.environ.get("LLM_ROOT", None) return _pl.Path(environ_root) if environ_root is not None else _pl.Path( __file__).resolve().parent.parent.parent.parent - - -@pytest.fixture(scope="module") -def resource_path(llm_root: _pl.Path) -> _pl.Path: - return llm_root / "cpp" / "tests" / "resources" - - -@pytest.fixture(scope="module") -def data_path(resource_path: _pl.Path) -> _pl.Path: - return resource_path / "data" - - -@pytest.fixture(scope="module") -def input_data_path(data_path): - return data_path / "input_tokens.npy" - - -@pytest.fixture(scope="module") -def engine_path(resource_path: _pl.Path) -> _pl.Path: - return resource_path / "models" / "rt_engine" - - -def get_base_model_spec() -> ModelSpec: - model_spec_obj = ModelSpec('input_tokens.npy', _tb.DataType.HALF) - model_spec_obj.use_gpt_plugin().set_kv_cache_type( - _tb.KVCacheType.PAGED).use_packed_input() - return model_spec_obj - - -@pytest.fixture(scope="module") -def model_path(engine_path): - return engine_path / f"gpt2/{get_base_model_spec().get_model_path()}/tp1-pp1-cp1-gpu" - - -@pytest.fixture(scope="module") -def model_path_return_logits(engine_path): - return engine_path / f"gpt2/{get_base_model_spec().gather_logits().get_model_path()}/tp1-pp1-cp1-gpu" - - -@pytest.fixture -def model_path_lora(engine_path: _pl.Path) -> _pl.Path: - return engine_path / f"gpt2/{get_base_model_spec().use_lora_plugin().get_model_path()}/tp1-pp1-cp1-gpu" - - -@pytest.fixture -def model_path_draft_tokens_external(engine_path: _pl.Path) -> _pl.Path: - return engine_path / f"gpt2/{get_base_model_spec().use_draft_tokens_external_decoding().get_model_path()}/tp1-pp1-cp1-gpu" - - -@pytest.fixture -def lora_config_path(data_path: _pl.Path) -> _pl.Path: - return data_path / "lora-test-weights-gpt2-tp1" - - -@pytest.fixture(scope="module") -def results_data_path(data_path: _pl.Path) -> _pl.Path: - return data_path / f"gpt2/sampling/{get_base_model_spec().get_results_file()}" - - -@pytest.fixture(scope="module") -def results_data_path_beam_width_2(data_path: _pl.Path) -> _pl.Path: - return data_path / f"gpt2/beam_search_2/{get_base_model_spec().get_results_file()}" - - -@pytest.fixture(scope="module") -def results_data_path_fmhafp32acc(data_path: _pl.Path) -> _pl.Path: - return data_path / f"gpt2/sampling/{get_base_model_spec().enable_context_fmha_fp32_acc().get_results_file()}"