diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 4043efb35b63..d67e9577de86 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..3415985c118e 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 | @@ -1350,14 +1350,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 | diff --git a/benchmarks/README.md b/benchmarks/README.md index 5d89f412ac5c..392e1a2f251c 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -1,10 +1,36 @@ -# TensorRT-LLM Benchmarks +# TensorRT-LLM Benchmarking -## Overview +For benchmarking TensorRT-LLM, use +[`trtllm-bench`](../docs/source/developer-guide/perf-benchmarking.md) and +`trtllm-serve`. The legacy C++ runtime benchmarks (`gptManagerBenchmark`, +`bertBenchmark`, `disaggServerBenchmark`) were removed together with the +TensorRT backend. -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 100% rename from benchmarks/cpp/prepare_dataset.py rename to benchmarks/prepare_dataset.py 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 a3b8cd348ce4..3885daac6ae8 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" @@ -233,7 +234,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, @@ -241,7 +242,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 index 9a577b61ad51..71e19c7d3d6c 100644 --- a/cpp/include/tensorrt_llm/batch_manager/guidedDecoder.h +++ b/cpp/include/tensorrt_llm/batch_manager/guidedDecoder.h @@ -17,6 +17,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/bufferManager.h" #include "tensorrt_llm/runtime/iTensor.h" @@ -39,7 +40,8 @@ class GuidedDecoder using BitmaskT = uint32_t; GuidedDecoder(executor::GuidedDecodingConfig const& guidedDecodingConfig, SizeType32 maxNumSequences, - SizeType32 vocabSizePadded, nvinfer1::DataType logitsDtype, runtime::BufferManager const& runtimeBufferManager); + SizeType32 vocabSizePadded, tensorrt_llm::DataType logitsDtype, + runtime::BufferManager const& runtimeBufferManager); void build(ScheduledRequests const& scheduledRequests); void execute(DecoderInputBuffers const& decoderInputBuffers, runtime::BufferManager const& runtimeBufferManager); @@ -51,7 +53,7 @@ class GuidedDecoder SizeType32 mMaxNumSequences; SizeType32 mVocabSizePadded; SizeType32 mBitmaskSize; // CeilDiv(vocabSizePadded, 32) - nvinfer1::DataType mLogitsDtype; + tensorrt_llm::DataType mLogitsDtype; TensorPtr mLogitsBitmask; // [mMaxNumRequests, mBitmaskSize] TensorPtr mLogitsBitmaskHost; // [mMaxNumRequests, mBitmaskSize] 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 299bf1b570f4..09186d507cc9 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, @@ -2664,7 +2664,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..2c0f5492967e --- /dev/null +++ b/cpp/include/tensorrt_llm/common/tllmDataType.h @@ -0,0 +1,98 @@ +/* + * 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. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" + +#include + +//! \file tllmDataType.h +//! +//! Standalone, TensorRT-free runtime types that replace the few \c nvinfer1 +//! types the shared C++ core historically used as common currency: +//! \c tensorrt_llm::DataType, \c tensorrt_llm::Dims and \c tensorrt_llm::ILogger. +//! These are defined here so the retained tree (runtime, batch manager, executor, +//! kernels and the nanobind bridge) compiles and links without the TensorRT +//! library. +//! +//! 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 + +//! \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]; +}; + +//! \brief Severity levels for the internal logger, mirroring the legacy +//! \c nvinfer1::ILogger::Severity values. +enum class ILoggerSeverity : int32_t +{ + kINTERNAL_ERROR = 0, + kERROR = 1, + kWARNING = 2, + kINFO = 3, + kVERBOSE = 4, +}; + +//! \brief Minimal, TensorRT-free logger interface replacing +//! \c nvinfer1::ILogger for the retained C++ tree. +class ILogger +{ +public: + using Severity = ILoggerSeverity; + + virtual void log(Severity severity, char const* msg) noexcept = 0; + + virtual ~ILogger() = default; +}; + +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/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..6290acaccde9 100644 --- a/cpp/tensorrt_llm/batch_manager/CMakeLists.txt +++ b/cpp/tensorrt_llm/batch_manager/CMakeLists.txt @@ -32,34 +32,21 @@ 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 fe47f6c531ad..719a6c7416f8 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 @@ -51,7 +52,7 @@ void BufferIndexHolder::release() 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 2cbf9f514bd7..9ce096b7ed03 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" @@ -197,7 +198,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 { @@ -224,7 +225,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 248ac55d7c90..2920d2121521 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" @@ -290,7 +291,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) @@ -385,20 +386,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 2ec86df47e99..ca9ef67983fa 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" @@ -761,8 +761,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()); @@ -1279,8 +1279,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()); @@ -1358,7 +1358,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); } @@ -1397,7 +1398,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 3362574da902..4cbbc561c52b 100644 --- a/cpp/tensorrt_llm/batch_manager/dataTransceiver.h +++ b/cpp/tensorrt_llm/batch_manager/dataTransceiver.h @@ -28,6 +28,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" @@ -290,8 +291,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(); @@ -341,8 +342,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 index cb2264ec8003..e2586bf10e6a 100644 --- a/cpp/tensorrt_llm/batch_manager/guidedDecoder.cpp +++ b/cpp/tensorrt_llm/batch_manager/guidedDecoder.cpp @@ -19,6 +19,7 @@ #include "tensorrt_llm/batch_manager/decoderBuffers.h" #include "tensorrt_llm/batch_manager/llmRequest.h" #include "tensorrt_llm/common/envUtils.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/kernels/logitsBitmask.h" #include @@ -30,7 +31,7 @@ namespace tensorrt_llm::batch_manager { GuidedDecoder::GuidedDecoder(executor::GuidedDecodingConfig const& guidedDecodingConfig, SizeType32 maxNumSequences, - SizeType32 vocabSizePadded, nvinfer1::DataType logitsDtype, BufferManager const& runtimeBufferManager) + SizeType32 vocabSizePadded, tensorrt_llm::DataType logitsDtype, BufferManager const& runtimeBufferManager) : mGuidedDecodingBackend{guidedDecodingConfig.getBackend()} , mMaxNumSequences{maxNumSequences} , mVocabSizePadded{vocabSizePadded} @@ -201,13 +202,13 @@ void GuidedDecoder::execute(DecoderInputBuffers const& decoderInputBuffers, Buff *ITensor::slice(mLogitsBitmaskPtrVec, 0, batchIdx)); auto logitsBitmaskPtrVec = bufferCast(*mLogitsBitmaskPtrVec); - if (mLogitsDtype == nvinfer1::DataType::kFLOAT) + if (mLogitsDtype == tensorrt_llm::DataType::kFLOAT) { auto logitsPtrVec = bufferCast(*mLogitsPtrVec); tensorrt_llm::kernels::invokeLogitsBitmask( logitsPtrVec, logitsBitmaskPtrVec, batchIdx, mVocabSizePadded, stream.get()); } - else if (mLogitsDtype == nvinfer1::DataType::kHALF) + else if (mLogitsDtype == tensorrt_llm::DataType::kHALF) { auto logitsPtrVec = bufferCast(*mLogitsPtrVec); tensorrt_llm::kernels::invokeLogitsBitmask( 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 c67a4711cc03..3fa4d18490de 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 @@ -4261,7 +4262,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..2a0e14c29989 100644 --- a/cpp/tensorrt_llm/batch_manager/utils/debugUtils.h +++ b/cpp/tensorrt_llm/batch_manager/utils/debugUtils.h @@ -25,7 +25,6 @@ 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..bef39826847c 100644 --- a/cpp/tensorrt_llm/executor/CMakeLists.txt +++ b/cpp/tensorrt_llm/executor/CMakeLists.txt @@ -19,6 +19,7 @@ set(TARGET_DIR ${CMAKE_CURRENT_SOURCE_DIR}) # keep this list sorted alphabetically set(SRCS + kvCacheEvent.cpp cache_transmission/mpi_utils/connection.cpp cache_transmission/agent_utils/connection.cpp cache_transmission/transferAgent.cpp @@ -27,9 +28,7 @@ set(SRCS contextPhaseParams.cpp debugConfig.cpp decodingConfig.cpp - executor.cpp executorConfig.cpp - executorImpl.cpp executorKVCacheEventManager.cpp extendedRuntimePerfKnobConfig.cpp guidedDecodingConfig.cpp @@ -60,7 +59,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/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/kvCacheEvent.cpp b/cpp/tensorrt_llm/executor/kvCacheEvent.cpp new file mode 100644 index 000000000000..5158eb93c983 --- /dev/null +++ b/cpp/tensorrt_llm/executor/kvCacheEvent.cpp @@ -0,0 +1,44 @@ +/* + * 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. + */ + +// Definition of the (backend-agnostic) KVCacheEvent constructor. It previously +// lived in executor.cpp alongside the legacy TensorRT-engine Executor class +// implementation; that file was removed with the TensorRT backend, so the +// constructor is relocated here so the retained KV-cache event path +// (kvCacheEventManager, serialization, nanobind) continues to link. + +#include +#include + +namespace tensorrt_llm::executor +{ + +char const* version() noexcept +{ + return kTensorRtLlmVersion; +} + +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/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/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 8a7120022b09..1f3d4ce4696b 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 @@ -119,7 +120,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 1496ebbb883d..867be6a3ccfa 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/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 369bf721b099..3c95892c4acf 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..d8118ed82fd9 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,5 @@ 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 index f5deb048a180..443d9f8688de 100644 --- a/cpp/tests/e2e_tests/CMakeLists.txt +++ b/cpp/tests/e2e_tests/CMakeLists.txt @@ -14,4 +14,3 @@ # 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 index 875e12eb975f..f866f13f1fb9 100644 --- a/cpp/tests/e2e_tests/batch_manager/CMakeLists.txt +++ b/cpp/tests/e2e_tests/batch_manager/CMakeLists.txt @@ -16,7 +16,3 @@ # 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 index 7b262cacb27d..39678cd6ba0f 100644 --- a/cpp/tests/e2e_tests/batch_manager/guidedDecoderTest.cpp +++ b/cpp/tests/e2e_tests/batch_manager/guidedDecoderTest.cpp @@ -27,6 +27,7 @@ #include "tensorrt_llm/batch_manager/decoderBuffers.h" #include "tensorrt_llm/batch_manager/guidedDecoder.h" #include "tensorrt_llm/batch_manager/llmRequest.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/executor/executor.h" using namespace tensorrt_llm::runtime; @@ -60,7 +61,7 @@ class GuidedDecoderTest : public ::testing::Test void initData(std::filesystem::path tokenizerInfoPath, SizeType32 vocabSizePadded, VecTokens outputIds, std::vector expectedNumRejected) { - mLogitsDtype = nvinfer1::DataType::kFLOAT; + mLogitsDtype = tensorrt_llm::DataType::kFLOAT; mMaxNumRequests = 16; mVocabSizePadded = vocabSizePadded; @@ -191,7 +192,7 @@ class GuidedDecoderTest : public ::testing::Test private: SizeType32 mMaxNumRequests; SizeType32 mVocabSizePadded; - nvinfer1::DataType mLogitsDtype; + tensorrt_llm::DataType mLogitsDtype; std::vector mLogits; // [mBatchSize, mVocabSizePadded] std::vector mLogitsHost; // [mBatchSize, mVocabSizePadded] 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 daa4f284b95d..1cde7662176c 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 @@ -91,7 +91,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 1d21b94a45ad..76ef3a732ecd 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); @@ -821,7 +823,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); @@ -1156,7 +1159,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); @@ -1391,7 +1395,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); @@ -1595,7 +1600,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); @@ -1888,7 +1894,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); @@ -2157,7 +2164,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); @@ -2389,7 +2397,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); @@ -2446,7 +2454,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); @@ -2584,7 +2593,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); @@ -2691,7 +2700,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); @@ -2763,7 +2772,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); { @@ -2856,7 +2865,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); @@ -2953,7 +2962,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); @@ -2997,7 +3006,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); @@ -3039,7 +3048,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); @@ -3124,7 +3133,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); @@ -3228,7 +3237,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; @@ -3251,11 +3260,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)); @@ -3326,10 +3336,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); @@ -3496,11 +3506,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); @@ -3604,10 +3615,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); @@ -3728,7 +3739,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); @@ -3822,7 +3833,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; @@ -4006,7 +4017,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); @@ -4133,7 +4144,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; @@ -4216,7 +4227,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; @@ -4342,7 +4353,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; @@ -4402,7 +4413,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; @@ -4479,7 +4490,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}; @@ -4545,7 +4556,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}; @@ -4671,7 +4682,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}; @@ -4787,7 +4798,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); @@ -4839,7 +4850,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); @@ -4911,7 +4922,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; @@ -4931,7 +4942,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)); @@ -4966,7 +4977,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; @@ -5025,10 +5036,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++) @@ -5065,11 +5077,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)); @@ -5156,10 +5168,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); @@ -5318,10 +5330,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); @@ -5458,12 +5470,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); @@ -5666,7 +5678,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) @@ -7162,7 +7174,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}; @@ -7240,7 +7252,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}; @@ -7334,7 +7346,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}; @@ -7447,7 +7459,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); @@ -7592,7 +7605,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); @@ -7816,7 +7830,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, @@ -7928,7 +7942,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); @@ -8254,7 +8268,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); @@ -8273,7 +8287,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); @@ -8306,7 +8320,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, @@ -8635,7 +8649,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, @@ -8792,7 +8806,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; @@ -9126,7 +9140,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; @@ -9145,7 +9159,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; @@ -9868,7 +9882,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(); @@ -10004,7 +10018,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, @@ -10198,7 +10212,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{ @@ -10246,7 +10260,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{}; @@ -10300,7 +10314,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; @@ -10390,7 +10404,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; @@ -10443,7 +10457,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(); @@ -10549,11 +10563,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..34b2ae2a6edf 100644 --- a/cpp/tests/unit_tests/executor/CMakeLists.txt +++ b/cpp/tests/unit_tests/executor/CMakeLists.txt @@ -19,14 +19,6 @@ 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) 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/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/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/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 77f8ea03fb83..3cc8ea0fdf93 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); @@ -1319,7 +1320,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); @@ -1434,7 +1435,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); @@ -1571,7 +1572,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); @@ -1718,87 +1719,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, @@ -1812,7 +1813,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), @@ -1834,7 +1835,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), @@ -1856,7 +1857,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), @@ -1875,7 +1876,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), @@ -1894,7 +1895,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), @@ -1916,7 +1917,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), @@ -1938,7 +1939,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), @@ -1960,7 +1961,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), @@ -1982,7 +1983,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), @@ -2004,7 +2005,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), @@ -2023,7 +2024,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), @@ -2033,97 +2034,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 @@ -2134,7 +2135,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; @@ -2424,7 +2425,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; @@ -2640,7 +2641,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..530b8a79568c 100644 --- a/cpp/tests/unit_tests/multi_gpu/kernels/allReduce/gemmAllReduceTest.cu +++ b/cpp/tests/unit_tests/multi_gpu/kernels/allReduce/gemmAllReduceTest.cu @@ -27,10 +27,10 @@ #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" @@ -54,7 +54,6 @@ #include "cutlass/util/reference/host/tensor_fill.h" using namespace cutlass; -using namespace nvinfer1; using namespace tensorrt_llm::mpi; using namespace tensorrt_llm::runtime; using namespace tensorrt_llm::common; @@ -238,7 +237,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 +245,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 +253,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 65282c2ae406..1a7df55c8521 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} @@ -136,7 +136,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 && \ @@ -146,24 +145,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/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/apps/chat.py b/examples/apps/chat.py index 0e3f2d928158..661949a3a607 100755 --- a/examples/apps/chat.py +++ b/examples/apps/chat.py @@ -5,8 +5,8 @@ import colorama from transformers import AutoTokenizer, PreTrainedTokenizer -from tensorrt_llm._tensorrt_engine import LLM -from tensorrt_llm.llmapi import BuildConfig, KvCacheConfig, SamplingParams +from tensorrt_llm import LLM +from tensorrt_llm.llmapi import KvCacheConfig, SamplingParams class LlmConsole(code.InteractiveConsole): @@ -72,10 +72,6 @@ def main(model: str, tokenizer: str, tp_size: int): free_gpu_memory_fraction=0.8) kv_cache_config.enable_block_reuse = True - build_config = BuildConfig(max_batch_size=1, - max_input_len=6000, - max_num_tokens=10240) - sampling_params = SamplingParams(max_tokens=100, temperature=0.5, top_p=0.95, @@ -83,7 +79,9 @@ def main(model: str, tokenizer: str, tp_size: int): llm = LLM(model, tokenizer, - build_config=build_config, + max_batch_size=1, + max_input_len=6000, + max_num_tokens=10240, kv_cache_config=kv_cache_config, tensor_parallel_size=tp_size) diff --git a/examples/apps/fastapi_server.py b/examples/apps/fastapi_server.py index 510b281a701b..e1d90c37934e 100755 --- a/examples/apps/fastapi_server.py +++ b/examples/apps/fastapi_server.py @@ -18,9 +18,9 @@ from fastapi import FastAPI, Request from fastapi.responses import JSONResponse, Response, StreamingResponse -from tensorrt_llm._tensorrt_engine import LLM +from tensorrt_llm import LLM from tensorrt_llm.executor import CppExecutorError, RequestError -from tensorrt_llm.llmapi import BuildConfig, KvCacheConfig, SamplingParams +from tensorrt_llm.llmapi import KvCacheConfig, SamplingParams TIMEOUT_KEEP_ALIVE = 5 # seconds. @@ -120,8 +120,6 @@ def entrypoint(model_dir: str, port = port or 8000 logging.info(f"Starting server at {host}:{port}") - build_config = BuildConfig(max_batch_size=10, max_beam_width=max_beam_width) - kv_cache_config = KvCacheConfig( free_gpu_memory_fraction=kv_cache_free_gpu_memory_fraction) @@ -129,7 +127,8 @@ def entrypoint(model_dir: str, tokenizer, tensor_parallel_size=tp_size, pipeline_parallel_size=pp_size, - build_config=build_config, + max_batch_size=10, + max_beam_width=max_beam_width, kv_cache_config=kv_cache_config) server = LlmServer(llm=llm) diff --git a/examples/auto_deploy/llmc/create_standalone_package.py b/examples/auto_deploy/llmc/create_standalone_package.py index 3368cd373bd2..e896ed55a46c 100644 --- a/examples/auto_deploy/llmc/create_standalone_package.py +++ b/examples/auto_deploy/llmc/create_standalone_package.py @@ -370,7 +370,7 @@ 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(temp_dir)\n", 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/layer_wise_benchmarks/sample_performance_alignment.sh b/examples/layer_wise_benchmarks/sample_performance_alignment.sh index 46c5964541e6..00a96cc00468 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 a5d3966bb5c8..d12f6402ecb2 100644 --- a/jenkins/Build.groovy +++ b/jenkins/Build.groovy @@ -421,23 +421,11 @@ 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 sh "mkdir -p TensorRT-LLM/attribution" sh "cp ${LLM_ROOT}/cpp/build/attribution/missing_files.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 255d9f6e442d..519704f9c175 100644 --- a/jenkins/L0_MergeRequest.groovy +++ b/jenkins/L0_MergeRequest.groovy @@ -929,7 +929,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", @@ -944,13 +943,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", 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..a08494074d1e 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 diff --git a/pyproject.toml b/pyproject.toml index 99b65e954d00..7564f1a418ad 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", diff --git a/requirements.txt b/requirements.txt index f387028c0429..fc0a14c9c640 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-04.html#rel-26-04 uses 2.12.0a0. torch>=2.11.0,<=2.12.0a0 torchvision diff --git a/ruff-legacy-baseline.json b/ruff-legacy-baseline.json index ab0274cacc59..b57a3dd23ec4 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": 5338, - "total_files": 492 + "total_violations": 2307, + "total_files": 270 }, ".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, @@ -65,7 +65,7 @@ }, "cpp/tensorrt_llm/kernels/cutlass_kernels/python/generate_kernels.py": { "F403": 1, - "F405": 184 + "F405": 22 }, "docs/source/helper.py": { "D205": 2, @@ -144,22 +144,11 @@ "D212": 1, "F821": 1 }, - "examples/medusa/convert_checkpoint.py": { - "F821": 1 - }, "examples/mmlu.py": { "D205": 1, "D415": 1, "E741": 1 }, - "examples/models/contrib/bloom/convert_checkpoint.py": { - "D200": 2, - "D202": 2, - "D210": 3, - "D212": 2, - "D415": 2, - "E741": 1 - }, "examples/models/contrib/chatglm-6b/tokenization_chatglm.py": { "D205": 4, "D210": 4, @@ -179,55 +168,6 @@ "D212": 3, "D415": 3 }, - "examples/models/contrib/cogvlm/convert_checkpoint.py": { - "D200": 1, - "D300": 1, - "D403": 1, - "D415": 1 - }, - "examples/models/contrib/dbrx/convert_checkpoint.py": { - "D200": 1, - "D212": 1, - "D415": 1, - "E741": 1 - }, - "examples/models/contrib/dit/sample.py": { - "D200": 1, - "D205": 1, - "D212": 1, - "D415": 1 - }, - "examples/models/contrib/dit/utils_modelopt.py": { - "D200": 1, - "D212": 1 - }, - "examples/models/contrib/gptj/convert_checkpoint.py": { - "F821": 1 - }, - "examples/models/contrib/gptneox/convert_checkpoint.py": { - "D202": 1, - "D210": 2, - "D415": 1, - "E741": 2 - }, - "examples/models/contrib/grok/convert_checkpoint.py": { - "D200": 1, - "D300": 1, - "D403": 1, - "D415": 1 - }, - "examples/models/contrib/mmdit/sample.py": { - "D200": 1 - }, - "examples/models/contrib/mpt/convert_checkpoint.py": { - "D200": 1, - "D212": 1, - "D415": 1, - "E741": 2 - }, - "examples/models/contrib/opt/convert_checkpoint.py": { - "E741": 1 - }, "examples/models/contrib/sdxl/pipeline_stable_diffusion_xl.py": { "D202": 1, "D205": 6, @@ -235,122 +175,9 @@ "D414": 1, "D415": 2 }, - "examples/models/contrib/stdit/pipeline_tllm.py": { - "D200": 1, - "D205": 1, - "D212": 1, - "D415": 1 - }, - "examples/models/contrib/stdit/text_encoder.py": { - "D200": 1, - "D212": 1 - }, - "examples/models/contrib/stdit/utils.py": { - "D200": 2, - "D205": 2, - "D212": 4, - "D415": 2 - }, - "examples/models/contrib/stdit/vae.py": { - "D415": 1 - }, - "examples/models/contrib/stdit/video_transforms.py": { - "D200": 2, - "D205": 17, - "D212": 19, - "D410": 8, - "D411": 9, - "D415": 18 - }, - "examples/models/core/bert/convert_checkpoint.py": { - "D200": 1, - "D300": 1, - "D403": 1, - "D415": 1 - }, - "examples/models/core/commandr/convert_checkpoint.py": { - "D200": 1, - "D300": 1, - "D403": 1, - "D415": 1 - }, - "examples/models/core/enc_dec/helper.py": { - "D205": 1, - "D212": 1, - "D300": 1, - "D403": 1, - "D415": 1 - }, - "examples/models/core/enc_dec/run.py": { - "D202": 1, - "D208": 18, - "D212": 1, - "D300": 1, - "D403": 1, - "D415": 1 - }, - "examples/models/core/glm-4-9b/convert_checkpoint.py": { - "D200": 1, - "D300": 1, - "D403": 1, - "D415": 1 - }, - "examples/models/core/glm-4-9b/tokenization_chatglm.py": { - "D200": 1, - "D205": 1, - "D210": 2, - "D212": 4, - "D415": 3, - "F821": 1 - }, - "examples/models/core/gpt/convert_checkpoint.py": { - "D200": 1, - "D300": 1, - "D403": 1, - "D415": 1 - }, - "examples/models/core/internlm2/convert_checkpoint.py": { - "D210": 1, - "D415": 1, - "E741": 1 - }, - "examples/models/core/llama/convert_checkpoint.py": { - "D200": 2, - "D300": 2, - "D403": 2, - "D415": 2, - "E722": 1 - }, - "examples/models/core/mamba/convert_checkpoint.py": { - "D200": 1, - "D300": 1, - "D403": 1, - "D415": 1, - "F821": 6 - }, - "examples/models/core/mllama/convert_checkpoint.py": { - "D200": 1, - "D300": 1, - "D403": 1, - "D415": 1, - "E722": 1 - }, "examples/models/core/multimodal/run.py": { "E731": 1 }, - "examples/models/core/phi/convert_checkpoint.py": { - "D200": 1, - "D300": 1, - "D403": 1, - "D415": 1 - }, - "examples/models/core/qwen/convert_checkpoint.py": { - "D200": 1, - "D300": 1, - "D403": 1, - "D415": 1, - "E722": 1 - }, "examples/models/core/qwen2audio/run.py": { "D200": 1, "D212": 1, @@ -366,22 +193,6 @@ "examples/models/core/qwenvl/run_chat.py": { "E722": 1 }, - "examples/models/core/whisper/convert_checkpoint.py": { - "D415": 1 - }, - "examples/models/core/whisper/run.py": { - "E712": 1 - }, - "examples/models/core/whisper/whisper_utils.py": { - "D200": 1, - "D202": 1, - "D205": 3, - "D212": 4, - "D403": 1, - "D411": 2, - "D415": 4, - "D416": 2 - }, "examples/ngram/run_dtm_ngram.py": { "D200": 1, "D205": 1, @@ -423,12 +234,6 @@ "E741": 1, "F601": 2 }, - "examples/redrafter/convert_checkpoint.py": { - "D205": 1, - "D212": 1, - "D415": 1, - "E722": 1 - }, "examples/run.py": { "E711": 4 }, @@ -452,9 +257,6 @@ "D212": 1, "E741": 1 }, - "scripts/check_test_list.py": { - "D212": 1 - }, "scripts/dco_check.py": { "D212": 1 }, @@ -485,42 +287,23 @@ "tensorrt_llm/__init__.py": { "D202": 1, "D212": 1, - "E402": 27 - }, - "tensorrt_llm/_torch/attention_backend/sparse/dsa.py": { - "F821": 2 + "E402": 20 }, "tensorrt_llm/_torch/attention_backend/sparse/kernel.py": { "E731": 3 }, "tensorrt_llm/_torch/attention_backend/sparse/rocket.py": { - "E712": 2, - "F821": 4 - }, - "tensorrt_llm/_torch/attention_backend/sparse/utils.py": { - "F821": 4 + "E712": 2 }, "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 }, @@ -528,10 +311,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 }, @@ -562,12 +341,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 }, @@ -575,9 +348,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 }, @@ -588,7 +358,7 @@ "F821": 1 }, "tensorrt_llm/_torch/modules/fused_moe/interface.py": { - "E402": 4 + "E402": 3 }, "tensorrt_llm/_torch/modules/fused_moe/moe_load_balancer.py": { "E712": 2 @@ -611,29 +381,12 @@ "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, - "F821": 13 + "F811": 1 }, "tensorrt_llm/_torch/pyexecutor/py_executor.py": { "E712": 1, - "E721": 1, - "F821": 2 - }, - "tensorrt_llm/_torch/pyexecutor/resource_manager.py": { - "F821": 2 + "E721": 1 }, "tensorrt_llm/_torch/pyexecutor/seq_slot_manager.py": { "F821": 1 @@ -641,15 +394,9 @@ "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 @@ -663,16 +410,12 @@ "D415": 1 }, "tensorrt_llm/bench/build/build.py": { - "D202": 1, - "D205": 2, - "D208": 1, - "D210": 3, + "D205": 1, + "D210": 2, "D411": 1 }, "tensorrt_llm/bench/build/dataclasses.py": { - "D205": 1, - "D208": 1, - "D210": 5 + "D210": 3 }, "tensorrt_llm/bench/build/tuning.py": { "D205": 2, @@ -686,26 +429,6 @@ "D200": 6, "D212": 6 }, - "tensorrt_llm/builder.py": { - "D205": 8, - "D208": 23, - "D210": 1, - "D212": 4, - "D300": 7, - "D403": 1, - "D415": 4, - "E712": 2 - }, - "tensorrt_llm/commands/prune.py": { - "D200": 1, - "D212": 1, - "D300": 1 - }, - "tensorrt_llm/commands/refit.py": { - "D200": 1, - "D212": 1, - "D300": 1 - }, "tensorrt_llm/commands/serve.py": { "D202": 2, "D212": 2, @@ -725,20 +448,17 @@ }, "tensorrt_llm/executor/base_worker.py": { "D200": 1, - "D202": 3, - "D205": 1, - "D208": 2, - "D210": 8, + "D202": 2, + "D210": 7, "D212": 1, - "D300": 5, + "D300": 4, "D415": 2 }, "tensorrt_llm/executor/executor.py": { - "D205": 7, + "D205": 3, "D210": 1, - "D212": 4, "D410": 3, - "D411": 7 + "D411": 3 }, "tensorrt_llm/executor/ipc.py": { "D202": 2, @@ -760,9 +480,7 @@ "D202": 1, "D205": 1, "D208": 3, - "D210": 1, "D212": 1, - "D300": 1, "E722": 1 }, "tensorrt_llm/executor/ray_executor.py": { @@ -778,14 +496,8 @@ }, "tensorrt_llm/executor/result.py": { "D200": 3, - "D202": 1, "D205": 1, - "D210": 3, - "D212": 5, - "D300": 3, - "F402": 1, - "F811": 1, - "F821": 1 + "F402": 1 }, "tensorrt_llm/executor/rpc/rpc_client.py": { "D200": 1, @@ -809,7 +521,6 @@ }, "tensorrt_llm/executor/rpc_proxy.py": { "D205": 1, - "D212": 1, "D415": 1 }, "tensorrt_llm/executor/rpc_worker.py": { @@ -820,588 +531,97 @@ "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": { - "D202": 1, "D415": 1 }, "tensorrt_llm/inputs/registry.py": { - "D200": 3, - "D205": 6, - "D212": 17, + "D200": 2, + "D205": 5, + "D212": 13, "D301": 1, "D405": 1, - "D411": 1, - "D415": 3, + "D415": 2, "E722": 2 }, "tensorrt_llm/inputs/utils.py": { - "D202": 3, - "D208": 5, - "D212": 1, - "D214": 1, - "E731": 2 - }, - "tensorrt_llm/layers/attention.py": { - "E712": 1, - "F821": 1 - }, - "tensorrt_llm/layers/embedding.py": { - "D200": 1, - "D205": 3, - "D212": 6, - "D415": 1, - "D416": 1 - }, - "tensorrt_llm/layers/language_adapter.py": { - "D205": 1, - "D212": 1 - }, - "tensorrt_llm/layers/linear.py": { - "D212": 1, - "D300": 1, - "D403": 1, - "D415": 1, - "E711": 2 - }, - "tensorrt_llm/layers/moe.py": { - "D200": 1, - "D212": 1, - "D300": 1, - "D415": 1 - }, - "tensorrt_llm/layers/recurrent.py": { - "D205": 1, - "D212": 1, - "D300": 1, - "D415": 1 - }, - "tensorrt_llm/layers/ssm.py": { - "D205": 3, - "D212": 3, - "D300": 3, - "D415": 3 - }, - "tensorrt_llm/llmapi/build_cache.py": { - "D200": 9, - "D210": 2, - "D212": 13, - "D300": 12, - "D415": 8, - "E722": 1 - }, - "tensorrt_llm/llmapi/disagg_utils.py": { - "D205": 2, - "D210": 1, - "D212": 2, - "D403": 2, - "D415": 2, - "F822": 2 - }, - "tensorrt_llm/llmapi/llm.py": { - "D200": 2, - "D202": 1, - "D205": 4, - "D207": 2, - "D212": 2, - "D300": 4, - "D410": 2, - "D411": 2 - }, - "tensorrt_llm/llmapi/llm_args.py": { - "D200": 23, - "D202": 1, - "D205": 11, - "D208": 2, - "D210": 7, - "D212": 37, - "D300": 8, - "D415": 2 - }, - "tensorrt_llm/llmapi/llm_utils.py": { - "D200": 2, - "D205": 4, - "D209": 1, - "D210": 9, - "D212": 3, - "D300": 12 - }, - "tensorrt_llm/llmapi/mgmn_leader_node.py": { - "D205": 2, - "D212": 2, - "D300": 2 - }, - "tensorrt_llm/llmapi/mm_encoder.py": { - "D200": 1, - "D205": 3, - "D207": 1, - "D410": 1, - "D411": 1 - }, - "tensorrt_llm/llmapi/mpi_session.py": { - "D200": 1, - "D205": 2, - "D210": 7, - "D212": 3, - "D300": 10, - "D411": 1, - "D415": 1 - }, - "tensorrt_llm/llmapi/reasoning_parser.py": { - "D205": 4, - "D209": 4, - "D212": 1, - "D301": 1 - }, - "tensorrt_llm/llmapi/tracer.py": { - "D210": 3, - "D300": 3 - }, - "tensorrt_llm/llmapi/utils.py": { - "D200": 2, - "D202": 1, - "D205": 5, - "D209": 1, - "D210": 14, - "D212": 5, - "D300": 8, - "E722": 1 - }, - "tensorrt_llm/mapping.py": { - "D205": 1, - "D212": 1, - "D300": 1, - "D415": 4 - }, - "tensorrt_llm/models/baichuan/config.py": { - "F821": 2 - }, - "tensorrt_llm/models/baichuan/convert.py": { - "D202": 1, - "D205": 1, - "D208": 17, - "D212": 1, - "D415": 1, - "E731": 1, - "E741": 3 - }, - "tensorrt_llm/models/baichuan/model.py": { - "D200": 1, - "D210": 1, - "D300": 1, - "D415": 1, - "F821": 1 - }, - "tensorrt_llm/models/bert/convert.py": { - "D200": 4, - "D208": 1, - "D212": 5, - "D300": 1, - "D403": 4, - "D415": 5 - }, - "tensorrt_llm/models/bert/model.py": { - "D200": 2, - "D202": 1, - "D205": 1, - "D212": 3, - "D300": 1, - "D415": 2 - }, - "tensorrt_llm/models/chatglm/config.py": { - "F821": 1 - }, - "tensorrt_llm/models/chatglm/convert.py": { - "D200": 2, - "D208": 1, - "D212": 1, - "D300": 1, - "D415": 1, - "E741": 1 - }, - "tensorrt_llm/models/chatglm/model.py": { - "D200": 2, - "D210": 1, - "D300": 1, - "D415": 1, - "F821": 1 - }, - "tensorrt_llm/models/cogvlm/convert.py": { - "E741": 1 - }, - "tensorrt_llm/models/commandr/model.py": { - "D200": 1, - "D202": 1, - "D210": 1, - "D300": 1, - "D415": 1 - }, - "tensorrt_llm/models/convert_utils.py": { - "D200": 2, - "D202": 1, - "D205": 1, - "D208": 17, - "D210": 2, - "D212": 3, - "D410": 2, - "D411": 4, - "D415": 3, - "E731": 1, - "F821": 1 - }, - "tensorrt_llm/models/deepseek_v1/config.py": { - "F821": 1 - }, - "tensorrt_llm/models/deepseek_v1/convert.py": { - "D200": 1, - "D212": 1, - "D415": 1, - "E741": 2 - }, - "tensorrt_llm/models/deepseek_v1/model.py": { - "F821": 1 - }, - "tensorrt_llm/models/deepseek_v2/config.py": { - "F821": 1 - }, - "tensorrt_llm/models/deepseek_v2/convert.py": { - "E741": 4 - }, - "tensorrt_llm/models/deepseek_v2/model.py": { - "F821": 1 - }, - "tensorrt_llm/models/dit/model.py": { - "D200": 2, - "D205": 2, - "D208": 2, - "D212": 3, - "D300": 1, - "D415": 1 - }, - "tensorrt_llm/models/eagle/config.py": { - "F821": 1 - }, - "tensorrt_llm/models/eagle/model.py": { - "D202": 3, - "D205": 4, - "D212": 5, - "D300": 4, - "D415": 1, - "F821": 1 - }, - "tensorrt_llm/models/enc_dec/model.py": { - "D202": 2, - "D205": 2, - "D208": 4, - "D300": 2, - "E712": 1 - }, - "tensorrt_llm/models/falcon/config.py": { - "F821": 1 - }, - "tensorrt_llm/models/falcon/convert.py": { - "D202": 1, - "D210": 2, - "D415": 1, - "E741": 2 - }, - "tensorrt_llm/models/falcon/model.py": { - "D200": 1, - "D210": 1, - "D300": 1, - "D415": 1, - "F821": 1 - }, - "tensorrt_llm/models/gemma/config.py": { - "D202": 1, - "D415": 1 - }, - "tensorrt_llm/models/gemma/convert.py": { - "D205": 1, - "D212": 1, - "D415": 1 - }, - "tensorrt_llm/models/gemma/smoothquant.py": { - "D200": 1, - "D212": 1, - "D415": 1, - "E741": 1 - }, - "tensorrt_llm/models/gemma/utils/modules.py": { - "D200": 1 - }, - "tensorrt_llm/models/gemma/utils/params.py": { - "D207": 6 - }, - "tensorrt_llm/models/gemma/utils/positional_embeddings.py": { - "D200": 1 - }, - "tensorrt_llm/models/gemma/weight.py": { - "D200": 1, - "D212": 1, - "E741": 3 - }, - "tensorrt_llm/models/generation_mixin.py": { - "E712": 2 - }, - "tensorrt_llm/models/gpt/config.py": { - "F821": 1 - }, - "tensorrt_llm/models/gpt/convert.py": { - "D202": 1, - "D205": 1, - "D212": 1, - "E741": 2 - }, - "tensorrt_llm/models/gpt/model.py": { - "D200": 1, - "D210": 1, - "D300": 1, - "D415": 1, - "F821": 1 - }, - "tensorrt_llm/models/gptj/config.py": { - "D200": 1, - "D212": 1, - "F811": 1, - "F821": 1 - }, - "tensorrt_llm/models/gptj/convert.py": { - "E741": 1 - }, - "tensorrt_llm/models/gptj/model.py": { - "F821": 1 - }, - "tensorrt_llm/models/grok/convert.py": { - "D200": 2, - "D208": 1, - "D210": 1, - "D212": 1, - "D300": 2, - "D415": 2, - "E741": 2 - }, - "tensorrt_llm/models/llama/config.py": { - "F821": 1 - }, - "tensorrt_llm/models/llama/convert.py": { - "D200": 1, - "D205": 1, - "D208": 4, - "D210": 1, - "D212": 2, - "D300": 3, - "D415": 1, - "E741": 5 - }, - "tensorrt_llm/models/llama/model.py": { - "D200": 1, - "D210": 1, - "D300": 1, - "D415": 1 - }, - "tensorrt_llm/models/mamba/convert.py": { - "D210": 1, - "E741": 1 - }, - "tensorrt_llm/models/mamba/model.py": { - "D205": 1, - "D208": 2, - "D300": 1, - "E712": 1, - "F821": 1 - }, - "tensorrt_llm/models/medusa/config.py": { - "F821": 1 - }, - "tensorrt_llm/models/medusa/model.py": { - "E712": 1, - "F821": 1 - }, - "tensorrt_llm/models/medusa/weight.py": { - "E741": 2 - }, - "tensorrt_llm/models/mllama/config.py": { - "F821": 1 - }, - "tensorrt_llm/models/mllama/model.py": { - "D200": 1, - "D202": 1, - "D205": 1, - "D208": 2, - "D210": 1, - "D300": 2, - "D415": 1, - "E712": 1, - "F821": 1 - }, - "tensorrt_llm/models/mmdit_sd3/model.py": { - "D200": 1 - }, - "tensorrt_llm/models/model_weights_loader.py": { - "D415": 2, - "E722": 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/models/multimodal_encoders/config.py": { - "F821": 1 - }, - "tensorrt_llm/models/multimodal_encoders/model.py": { - "D200": 1, - "D202": 1, - "D205": 1, - "D208": 2, - "D210": 1, - "D300": 2, - "D415": 1 - }, - "tensorrt_llm/models/nemotron_nas/config.py": { - "F821": 1 - }, - "tensorrt_llm/models/nemotron_nas/convert.py": { - "D200": 1, - "D415": 1, - "E741": 1, - "F821": 7 - }, - "tensorrt_llm/models/nemotron_nas/model.py": { - "D200": 2, - "D202": 1, - "D205": 6, - "D212": 8, - "D415": 6, - "F821": 1 + "D202": 3, + "D208": 5, + "D212": 1, + "D214": 1, + "E731": 2 }, - "tensorrt_llm/models/phi/config.py": { - "F821": 1 + "tensorrt_llm/llmapi/disagg_utils.py": { + "D205": 2, + "D210": 1, + "D212": 2, + "D403": 2, + "D415": 2, + "F822": 2 }, - "tensorrt_llm/models/phi/model.py": { - "F821": 1 + "tensorrt_llm/llmapi/llm.py": { + "D200": 2, + "D205": 4 }, - "tensorrt_llm/models/phi3/config.py": { - "F821": 1 + "tensorrt_llm/llmapi/llm_args.py": { + "D205": 11, + "D212": 1 }, - "tensorrt_llm/models/phi3/model.py": { - "F821": 1 + "tensorrt_llm/llmapi/llm_utils.py": { + "D200": 1 }, - "tensorrt_llm/models/qwen/config.py": { - "F821": 1 + "tensorrt_llm/llmapi/mgmn_leader_node.py": { + "D205": 2, + "D212": 2, + "D300": 2 }, - "tensorrt_llm/models/qwen/convert.py": { - "D200": 1, - "D208": 1, - "D212": 1, - "D300": 1, - "D415": 1, - "E741": 2 + "tensorrt_llm/llmapi/mm_encoder.py": { + "D205": 2, + "D410": 1, + "D411": 1 }, - "tensorrt_llm/models/qwen/model.py": { + "tensorrt_llm/llmapi/mpi_session.py": { "D200": 1, - "D210": 1, - "D300": 1, - "D415": 1, - "F821": 1 - }, - "tensorrt_llm/models/recurrentgemma/model.py": { "D205": 1, - "D208": 2, - "D300": 1, - "E712": 1 + "D415": 1 }, - "tensorrt_llm/models/redrafter/model.py": { - "D200": 1, - "D202": 2, - "D205": 2, - "D212": 3, - "D300": 1, - "D415": 3 + "tensorrt_llm/llmapi/reasoning_parser.py": { + "D205": 4, + "D209": 4, + "D212": 1, + "D301": 1 }, - "tensorrt_llm/models/redrafter/redrafter_helper.py": { - "D200": 5, - "D202": 2, - "D205": 12, - "D208": 7, - "D212": 17, - "D300": 9, - "D403": 1, - "D410": 2, - "D411": 5, - "D415": 14 + "tensorrt_llm/llmapi/tracer.py": { + "D210": 3, + "D300": 3 }, - "tensorrt_llm/models/stdit/model.py": { - "D200": 1, + "tensorrt_llm/llmapi/utils.py": { + "D200": 2, "D202": 1, - "D205": 1, - "D208": 2, - "D300": 1 - }, - "tensorrt_llm/models/unet/embeddings.py": { - "D205": 1, - "D212": 1 + "D205": 5, + "D209": 1, + "D210": 13, + "D212": 5, + "D300": 7, + "E722": 1 }, - "tensorrt_llm/network.py": { - "D200": 5, + "tensorrt_llm/mapping.py": { "D205": 1, - "D208": 4, - "D210": 1, - "D212": 13, - "D300": 14, - "D415": 2, - "E731": 1, - "F811": 1, - "F821": 10 - }, - "tensorrt_llm/plugin/plugin.py": { - "D205": 2, - "D208": 10, "D212": 1, - "D415": 1, - "E712": 1, - "E721": 1, - "F821": 5 - }, - "tensorrt_llm/quantization/functional.py": { - "D205": 4, - "D212": 4, - "D300": 4, - "D411": 2, - "D415": 2 + "D300": 1, + "D415": 4 }, - "tensorrt_llm/quantization/layers.py": { + "tensorrt_llm/models/convert_utils.py": { "D200": 2, + "D202": 1, "D205": 1, - "D208": 8, + "D208": 17, + "D210": 2, "D212": 3, - "D415": 1, - "E712": 2 + "D410": 2, + "D411": 4, + "D415": 3, + "E731": 1, + "F821": 1 }, "tensorrt_llm/quantization/quantize_by_modelopt.py": { "D200": 1, @@ -1427,81 +647,10 @@ "D205": 2, "D212": 3, "D403": 2, - "D415": 2 - }, - "tensorrt_llm/runtime/__init__.py": { - "E402": 7 - }, - "tensorrt_llm/runtime/enc_dec_model_runner.py": { - "E731": 1 - }, - "tensorrt_llm/runtime/generation.py": { - "D200": 4, - "D202": 1, - "D205": 3, - "D208": 2, - "D212": 6, - "D300": 3, - "D403": 2, - "D415": 3, - "E711": 5, - "F403": 1, - "F405": 9 - }, - "tensorrt_llm/runtime/kv_cache_manager.py": { - "D200": 8, - "D202": 1, - "D205": 6, - "D212": 14, - "D415": 10, - "E712": 1, - "E741": 1 - }, - "tensorrt_llm/runtime/medusa_utils.py": { - "D200": 1, - "D212": 1 - }, - "tensorrt_llm/runtime/model_runner.py": { - "D200": 1, - "D205": 3, - "D212": 7, - "D410": 2, - "D411": 2, - "E741": 1 - }, - "tensorrt_llm/runtime/model_runner_cpp.py": { - "D200": 1, - "D205": 1, - "D212": 3, - "D410": 2, - "D411": 2, - "E711": 1 - }, - "tensorrt_llm/runtime/multimodal_model_runner.py": { - "D202": 1, - "D208": 27, - "D212": 6, - "D214": 9, - "D410": 2, - "D411": 3, - "E741": 2 - }, - "tensorrt_llm/runtime/processor_wrapper/mllama_processor_wrapper.py": { - "D200": 1, - "D212": 1, - "D300": 1, "D415": 1 }, - "tensorrt_llm/runtime/session.py": { - "D200": 4, - "D202": 1, - "D205": 6, - "D210": 1, - "D212": 6, - "D300": 11, - "D403": 2, - "D415": 8, - "E741": 3 + "tensorrt_llm/runtime/__init__.py": { + "E402": 1 }, "tensorrt_llm/scaffolding/contrib/DeepConf/deep_conf_controller.py": { "E712": 1 @@ -1547,10 +696,6 @@ "PLE0302": 1 }, "tensorrt_llm/scaffolding/task.py": { - "F811": 1, - "F821": 1 - }, - "tensorrt_llm/scaffolding/task_collection.py": { "F821": 1 }, "tensorrt_llm/scaffolding/worker.py": { @@ -1571,11 +716,6 @@ "D210": 1, "F811": 1 }, - "tensorrt_llm/serve/openai_server.py": { - "D205": 1, - "D212": 2, - "F821": 2 - }, "tensorrt_llm/serve/responses_utils.py": { "D200": 2, "D205": 4, @@ -1633,35 +773,10 @@ }, "tensorrt_llm/tokenizer/tokenizer.py": { "D202": 1, - "D205": 1, "D209": 1, "D210": 3, "D300": 5 }, - "tensorrt_llm/tools/plugin_gen/core.py": { - "D200": 4, - "D210": 6, - "D212": 8, - "D300": 14, - "D415": 2, - "F821": 1 - }, - "tensorrt_llm/tools/plugin_gen/plugin_gen.py": { - "D200": 10, - "D202": 1, - "D212": 10, - "D300": 10, - "D403": 2, - "D415": 3 - }, - "tensorrt_llm/tools/plugin_gen/shape_infer.py": { - "D200": 1, - "D212": 2, - "D300": 2, - "D403": 1, - "D415": 1, - "E731": 1 - }, "tensorrt_llm/tools/ppl.py": { "D200": 1, "D212": 1 @@ -1674,29 +789,14 @@ "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.py": { - "D300": 1, - "D415": 2 - }, - "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, - "D415": 3, - "E402": 7 + "D415": 3 }, "tests/integration/defs/common.py": { "D200": 1, @@ -1710,17 +810,16 @@ }, "tests/integration/defs/conftest.py": { "D200": 3, - "D202": 10, - "D205": 4, + "D202": 9, + "D205": 3, "D209": 1, "D212": 5, "D214": 1, - "D300": 96, - "D403": 52, - "D415": 96, + "D300": 92, + "D403": 51, + "D415": 92, "E722": 2, - "E741": 1, - "F811": 1 + "E741": 1 }, "tests/integration/defs/cpp/conftest.py": { "E402": 2 @@ -1730,7 +829,6 @@ "E741": 2 }, "tests/integration/defs/disaggregated/test_disaggregated.py": { - "D202": 2, "D205": 3, "D209": 1, "D212": 5, @@ -1766,12 +864,6 @@ "tests/integration/defs/examples/test_commandr.py": { "D300": 1 }, - "tests/integration/defs/examples/test_gemma.py": { - "D202": 1, - "D300": 6, - "D403": 5, - "D415": 5 - }, "tests/integration/defs/examples/test_gpt.py": { "D202": 3, "D300": 5, @@ -1788,10 +880,9 @@ }, "tests/integration/defs/examples/test_llama.py": { "D202": 2, - "D300": 6, - "D403": 4, - "D415": 4, - "E712": 1 + "D300": 4, + "D403": 2, + "D415": 2 }, "tests/integration/defs/examples/test_mamba.py": { "D300": 2, @@ -1816,9 +907,8 @@ }, "tests/integration/defs/examples/test_phi.py": { "D202": 1, - "D300": 3, - "D403": 1, - "D415": 3 + "D300": 2, + "D415": 2 }, "tests/integration/defs/examples/test_qwen.py": { "D202": 1, @@ -1833,11 +923,6 @@ "D300": 1, "D415": 1 }, - "tests/integration/defs/examples/test_recurrentgemma.py": { - "D202": 1, - "D300": 2, - "D415": 2 - }, "tests/integration/defs/llmapi/test_llm_api_qa.py": { "D200": 1, "D212": 1, @@ -1874,7 +959,7 @@ "D212": 3 }, "tests/integration/defs/perf/open_search_db_utils.py": { - "D200": 2, + "D200": 1, "D212": 4, "E402": 1 }, @@ -1941,37 +1026,21 @@ }, "tests/integration/defs/test_e2e.py": { "D200": 5, - "D202": 3, - "D205": 3, - "D210": 1, - "D212": 8, - "D300": 9, - "D403": 4, - "D415": 10, - "F811": 2 + "D415": 6, + "F811": 1 }, "tests/integration/defs/test_list_parser.py": { "D202": 5, "D205": 5, "D208": 1, "D209": 2, - "D212": 3, - "D214": 3, - "D410": 1, - "D411": 1, - "E731": 2, - "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 + "D212": 3, + "D214": 3, + "D410": 1, + "D411": 1, + "E731": 2, + "E741": 2, + "F811": 1 }, "tests/integration/defs/triton_server/conftest.py": { "D200": 3, @@ -1986,19 +1055,6 @@ "E741": 3, "F821": 1 }, - "tests/integration/defs/triton_server/local_venv.py": { - "D205": 4, - "D212": 4 - }, - "tests/integration/defs/triton_server/rcca/bug_4323566/inflight_batcher_llm_client_with_end_id.py": { - "E721": 2 - }, - "tests/integration/defs/triton_server/runner_interface.py": { - "D205": 3, - "D208": 4, - "D212": 3, - "D415": 2 - }, "tests/integration/defs/triton_server/test_list_parser.py": { "D202": 5, "D205": 5, @@ -2013,22 +1069,6 @@ "F811": 1, "F821": 2 }, - "tests/integration/defs/triton_server/test_triton_llm.py": { - "F403": 2, - "F405": 174 - }, - "tests/integration/defs/triton_server/test_triton_memleak.py": { - "F403": 2, - "F405": 7 - }, - "tests/integration/defs/triton_server/test_triton_multi_node.py": { - "F403": 2, - "F405": 3 - }, - "tests/integration/defs/triton_server/test_triton_rcca.py": { - "F403": 2, - "F405": 29 - }, "tests/integration/defs/triton_server/trt_test_alternative.py": { "D200": 1, "D212": 1, @@ -2048,16 +1088,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 @@ -2068,12 +1098,6 @@ "tests/unittest/_torch/misc/test_autotuner.py": { "E731": 1 }, - "tests/unittest/_torch/modeling/test_modeling_exaone4.py": { - "E402": 12 - }, - "tests/unittest/_torch/thop/parallel/test_mamba2_chunk_ss_update.py": { - "E731": 4 - }, "tests/unittest/api_stability/api_stability_core.py": { "D202": 1, "D205": 1, @@ -2083,15 +1107,8 @@ "tests/unittest/bindings/test_bindings_moe.py": { "D415": 6 }, - "tests/unittest/bindings/test_executor_bindings.py": { - "D202": 1, - "E712": 56, - "F403": 2, - "F405": 3, - "F811": 1 - }, "tests/unittest/conftest.py": { - "D200": 4, + "D200": 3, "D202": 2, "D212": 5 }, @@ -2150,30 +1167,17 @@ "D209": 2 }, "tests/unittest/llmapi/test_llm.py": { - "D200": 1, "D205": 1, "D209": 1, "D210": 1, - "D212": 1, - "D300": 2, - "E712": 1, - "E722": 1 + "D300": 1, + "E712": 1 }, "tests/unittest/llmapi/test_llm_args.py": { "D200": 1, - "D202": 2, - "D212": 2, - "E712": 20, - "F811": 1 - }, - "tests/unittest/llmapi/test_llm_multi_gpu.py": { - "D210": 2, - "D300": 2 + "E712": 19 }, "tests/unittest/llmapi/test_llm_pytorch.py": { - "D202": 1, - "D210": 2, - "D212": 1, "F811": 1, "F821": 3 }, @@ -2181,10 +1185,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, @@ -2194,35 +1194,8 @@ "tests/unittest/llmapi/test_serialization.py": { "E721": 1 }, - "tests/unittest/others/test_graph_rewriter.py": { - "D205": 1, - "D212": 1, - "D300": 1, - "D415": 1 - }, "tests/unittest/others/test_kv_cache_transceiver.py": { - "D205": 1, - "D212": 2 - }, - "tests/unittest/others/test_layer.py": { - "D205": 1, - "D212": 1, - "E711": 1 - }, - "tests/unittest/others/test_leak.py": { - "D200": 1, - "D210": 1, - "D300": 1, - "D415": 1 - }, - "tests/unittest/others/test_model_dtype.py": { - "D200": 1, - "D210": 1, - "D300": 1, - "D415": 1 - }, - "tests/unittest/others/test_module.py": { - "E741": 1 + "D212": 1 }, "tests/unittest/others/test_time_breakdown.py": { "D212": 1, @@ -2234,27 +1207,6 @@ "tests/unittest/scaffolding/test_scaffolding.py": { "F811": 10 }, - "tests/unittest/test_model_runner_cpp.py": { - "F403": 2, - "F405": 2 - }, - "tests/unittest/tools/plugin_gen/kernel_config.py": { - "F403": 1, - "F405": 22 - }, - "tests/unittest/tools/plugin_gen/test_core.py": { - "D200": 1, - "D212": 1, - "D300": 1, - "D403": 1, - "D415": 1, - "F403": 1, - "F405": 7 - }, - "tests/unittest/tools/plugin_gen/test_shape_infer.py": { - "F403": 1, - "F405": 4 - }, "tests/unittest/tools/test_prepare_dataset.py": { "D205": 2, "D212": 7 @@ -2263,76 +1215,6 @@ "D202": 1, "E402": 1 }, - "tests/unittest/trt/attention/test_gpt_attention_IFB.py": { - "E711": 1, - "E712": 1, - "E731": 2 - }, - "tests/unittest/trt/functional/test_moe.py": { - "D202": 1, - "D210": 4, - "D415": 5 - }, - "tests/unittest/trt/functional/test_unbind.py": { - "F811": 1 - }, - "tests/unittest/trt/model/test_gpt_e2e.py": { - "E402": 1 - }, - "tests/unittest/trt/model/test_llama.py": { - "E741": 1 - }, - "tests/unittest/trt/model/test_mamba.py": { - "E741": 1 - }, - "tests/unittest/trt/model/test_mistral.py": { - "E741": 1 - }, - "tests/unittest/trt/model/test_nemotron_nas.py": { - "D415": 1 - }, - "tests/unittest/trt/model/test_phi.py": { - "F601": 1 - }, - "tests/unittest/trt/model/test_unet.py": { - "F403": 1, - "F405": 13 - }, - "tests/unittest/trt/model_api/test_model_api_multi_gpu.py": { - "D200": 1, - "D300": 1, - "D415": 1 - }, - "tests/unittest/trt/model_api/test_model_level_api.py": { - "D200": 1, - "D205": 1, - "D208": 3, - "D300": 2, - "D403": 1, - "D415": 1 - }, - "tests/unittest/trt/quantization/test_moe_weight_only_quant_matmul.py": { - "E741": 1 - }, - "tests/unittest/trt/quantization/test_qserve_gemm.py": { - "E402": 3 - }, - "tests/unittest/trt/quantization/test_quant_layer.py": { - "D200": 2, - "D205": 1, - "D208": 1, - "D210": 1, - "D212": 2, - "D415": 3 - }, - "tests/unittest/utils/test_medusa_utils.py": { - "D202": 2, - "D205": 2, - "D212": 3, - "D405": 1, - "D411": 1, - "D415": 2 - }, "tests/unittest/utils/test_prebuilt_whl_cpp_extensions.py": { "D200": 1, "D212": 1 @@ -2348,92 +1230,14 @@ }, "tests/unittest/utils/util.py": { "D200": 1, - "D202": 1, - "D205": 2, "D210": 2, - "D212": 3, + "D212": 1, "D300": 3, "D403": 3, - "D405": 1, - "D410": 1, - "D411": 3, - "D415": 4, - "E731": 1 - }, - "triton_backend/all_models/disaggregated_serving/disaggregated_serving_bls/1/model.py": { - "D205": 3, - "D416": 1, - "E722": 1 - }, - "triton_backend/all_models/gpt/postprocessing/1/model.py": { - "D202": 1, - "D205": 4, - "D411": 1, - "D415": 2, - "D416": 1 - }, - "triton_backend/all_models/gpt/preprocessing/1/model.py": { - "D200": 1, - "D202": 1, - "D205": 5, - "D208": 1, - "D212": 2, - "D300": 1, - "D403": 1, - "D411": 1, - "D415": 3, - "D416": 1, - "E711": 1 - }, - "triton_backend/all_models/gpt/tensorrt_llm/1/model.py": { - "D205": 4, - "D416": 1 - }, - "triton_backend/all_models/inflight_batcher_llm/postprocessing/1/model.py": { - "D202": 1, - "D205": 4, - "D411": 1, - "D415": 2, - "D416": 1 - }, - "triton_backend/all_models/inflight_batcher_llm/preprocessing/1/model.py": { - "D200": 1, - "D202": 2, - "D205": 6, - "D208": 1, - "D210": 1, - "D212": 5, - "D300": 1, - "D403": 1, - "D411": 1, "D415": 3, - "D416": 1, - "E711": 3, - "F821": 1 - }, - "triton_backend/all_models/inflight_batcher_llm/tensorrt_llm/1/model.py": { - "D205": 5, - "D212": 1, - "D415": 1, - "D416": 1, - "E402": 1, - "E712": 2, - "E722": 1 - }, - "triton_backend/all_models/inflight_batcher_llm/tensorrt_llm_bls/1/lib/decode.py": { - "D415": 1 - }, - "triton_backend/all_models/inflight_batcher_llm/tensorrt_llm_bls/1/lib/triton_decoder.py": { - "F403": 1, - "F405": 43 - }, - "triton_backend/all_models/inflight_batcher_llm/tensorrt_llm_bls/1/model.py": { - "D205": 1 + "E731": 1 }, "triton_backend/all_models/llmapi/tensorrt_llm/1/helpers.py": { - "D205": 3, - "D212": 3, - "D415": 3, "E722": 1 }, "triton_backend/all_models/llmapi/tensorrt_llm/1/model.py": { @@ -2444,75 +1248,12 @@ "D411": 1, "D415": 2 }, - "triton_backend/all_models/multimodal/multimodal_encoders/1/model.py": { - "D202": 3, - "D205": 6, - "D212": 2, - "D411": 1, - "D415": 2, - "D416": 1, - "E711": 2 - }, - "triton_backend/all_models/multimodal/multimodal_encoders/1/multimodal_utils.py": { - "D212": 1, - "D214": 2, - "D410": 1, - "D411": 2 - }, - "triton_backend/all_models/tests/test_decode.py": { - "F403": 1, - "F405": 34 - }, "triton_backend/all_models/tests/test_llmapi_python_backend.py": { "E402": 2, "E712": 3, "F403": 1, "F405": 2 }, - "triton_backend/all_models/tests/test_multi_image_preprocess.py": { - "E402": 1 - }, - "triton_backend/all_models/tests/test_multimodal_encoders.py": { - "E402": 1 - }, - "triton_backend/all_models/tests/test_python_backend.py": { - "E402": 2, - "E712": 40, - "F403": 1, - "F405": 72 - }, - "triton_backend/all_models/tests/test_triton_decoder.py": { - "E402": 4, - "F601": 1 - }, - "triton_backend/all_models/whisper/whisper_bls/1/fbank.py": { - "D205": 2, - "D212": 2, - "D403": 1, - "D415": 2, - "D416": 1 - }, - "triton_backend/all_models/whisper/whisper_bls/1/model.py": { - "D205": 4, - "D212": 2, - "D415": 1, - "D416": 2 - }, - "triton_backend/ci/L0_backend_trtllm/base_metrics_verification_tests.py": { - "D205": 1, - "D212": 3, - "E402": 3 - }, - "triton_backend/inflight_batcher_llm/client/e2e_grpc_speculative_decoding_client.py": { - "E721": 1 - }, - "triton_backend/inflight_batcher_llm/client/end_to_end_grpc_client.py": { - "E721": 1, - "E722": 1 - }, - "triton_backend/inflight_batcher_llm/client/inflight_batcher_llm_client.py": { - "E721": 2 - }, "triton_backend/tools/inflight_batcher_llm/benchmark_core_model.py": { "E722": 1 }, @@ -2529,10 +1270,5 @@ "D205": 1, "D212": 1, "F811": 3 - }, - "triton_backend/tools/whisper/client.py": { - "D205": 1, - "D212": 1, - "E721": 3 } } diff --git a/ruff-legacy.toml b/ruff-legacy.toml index 2f9d1bd00745..9a0e4a65d217 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", 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..ca798ca6f953 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) @@ -1313,9 +1275,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 be0f840caaa5..07c0d95038af 100644 --- a/setup.py +++ b/setup.py @@ -116,15 +116,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', @@ -438,9 +436,6 @@ def extract_from_precompiled(precompiled_location: str, package_data: list[str], license_files=get_license(), entry_points={ 'console_scripts': [ - 'trtllm-build=tensorrt_llm.commands.build:main', - 'trtllm-prune=tensorrt_llm.commands.prune:main', - 'trtllm-refit=tensorrt_llm.commands.refit:main', 'trtllm-bench=tensorrt_llm.commands.bench:main', 'trtllm-serve=tensorrt_llm.commands.serve:main', 'trtllm-eval=tensorrt_llm.commands.eval:main' diff --git a/tensorrt_llm/__init__.py b/tensorrt_llm/__init__.py index 0b269405274f..d707ca0e0694 100644 --- a/tensorrt_llm/__init__.py +++ b/tensorrt_llm/__init__.py @@ -105,31 +105,23 @@ def _setup_vendored_triton_kernels(): import torch # noqa import tensorrt_llm._torch.models as torch_models -import tensorrt_llm.functional as functional import tensorrt_llm.math_utils as math_utils import tensorrt_llm.models as models import tensorrt_llm.quantization as quantization import tensorrt_llm.runtime as runtime import tensorrt_llm.tools as tools -from ._common import _init, default_net, default_trtnet, precision +from ._common import _init from ._mnnvl_utils import MnnvlMemory, MnnvlMoe, MoEAlltoallInfo from ._utils import (default_gpus_per_node, local_mpi_rank, local_mpi_size, mpi_barrier, mpi_comm, mpi_rank, mpi_world_size, - set_mpi_comm, str_dtype_to_torch, str_dtype_to_trt, - torch_dtype_to_trt) -from .builder import BuildConfig, Builder, BuilderConfig, build + set_mpi_comm, str_dtype_to_torch) from .disaggregated_params import DisaggregatedParams -from .functional import Tensor, constant from .llmapi import LLM, AsyncLLM, MultimodalEncoder -from .llmapi.llm_args import LlmArgs, TorchLlmArgs, TrtLlmArgs +from .llmapi.llm_args import LlmArgs, TorchLlmArgs from .logger import logger from .mapping import Mapping from .models.automodel import AutoConfig, AutoModelForCausalLM -from .module import Module -from .network import Network, net_guard -from .parameter import Parameter -from .python_plugin import PluginBase from .sampling_params import SamplingParams from .version import __version__ from .visual_gen import (ExtraParamSchema, VisualGen, VisualGenArgs, @@ -140,8 +132,6 @@ def _setup_vendored_triton_kernels(): 'AutoConfig', 'AutoModelForCausalLM', 'logger', - 'str_dtype_to_trt', - 'torch_dtype_to_trt', 'str_dtype_to_torch', 'default_gpus_per_node', 'local_mpi_rank', @@ -151,27 +141,12 @@ def _setup_vendored_triton_kernels(): 'mpi_rank', 'set_mpi_comm', 'mpi_world_size', - 'constant', - 'default_net', - 'default_trtnet', - 'precision', - 'net_guard', 'torch_models', - 'Network', 'Mapping', 'MnnvlMemory', 'MnnvlMoe', 'MoEAlltoallInfo', - 'PluginBase', - 'Builder', - 'BuilderConfig', - 'build', - 'BuildConfig', - 'Tensor', - 'Parameter', 'runtime', - 'Module', - 'functional', 'models', 'quantization', 'tools', @@ -180,7 +155,6 @@ def _setup_vendored_triton_kernels(): 'MultimodalEncoder', 'LlmArgs', 'TorchLlmArgs', - 'TrtLlmArgs', 'SamplingParams', 'VisualGenArgs', 'ExtraParamSchema', diff --git a/tensorrt_llm/_common.py b/tensorrt_llm/_common.py index 871120aabf13..d7feeededb9f 100644 --- a/tensorrt_llm/_common.py +++ b/tensorrt_llm/_common.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# 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"); @@ -12,35 +12,19 @@ # WITHOUT 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 contextlib + import ctypes import os import platform import threading import time -from functools import wraps from pathlib import Path -from typing import TYPE_CHECKING - -import numpy as np -# isort: off import torch -import tensorrt as trt - -# isort: on - -if TYPE_CHECKING: - from .network import Network -else: - Network = None -from ._utils import print_all_stacks, str_dtype_to_trt +from ._utils import print_all_stacks from .bindings import MpiComm from .logger import logger -from .plugin import _load_plugin_lib - -net = None _inited = False @@ -50,7 +34,6 @@ def _init(log_level: object = None) -> None: if _inited: return _inited = True - # Move to __init__ if log_level is not None: logger.set_level(log_level) @@ -60,11 +43,19 @@ def _init(log_level: object = None) -> None: logger.info("Starting TensorRT LLM init.") - # load plugin lib - _load_plugin_lib() - - # load FT decoder layer and torch custom ops project_dir = str(Path(__file__).parent.absolute()) + + # Promote libtensorrt_llm.so symbols to the process's global scope. The KV + # cache transfer agent wrappers (libtensorrt_llm_nixl_wrapper.so, + # libtensorrt_llm_ucx_wrapper.so) are dlopen'ed at runtime without linking + # against libtensorrt_llm.so and resolve its symbols from the global symbol + # table, while Python extension modules and their dependencies load with + # RTLD_LOCAL. This promotion was previously a side effect of loading the + # TensorRT plugin library with RTLD_GLOBAL. + if platform.system() != "Windows": + ctypes.CDLL(project_dir + "/libs/libtensorrt_llm.so", mode=ctypes.RTLD_GLOBAL) + + # Load FT decoder layer and torch custom ops. if platform.system() == "Windows": ft_decoder_lib = project_dir + "/libs/th_common.dll" else: @@ -97,210 +88,3 @@ def _print_stacks(): print_stacks_thread.start() logger.info("TensorRT LLM inited.") - - -def default_net() -> Network: - assert net, ( - "Use builder to create network first, and use `set_network` or `net_guard` to set it to default" - ) - return net - - -def default_trtnet(): - return default_net().trt_network - - -def set_network(network): - global net - net = network - - -def switch_net_dtype(cur_dtype): - prev_dtype = default_net().dtype - default_net().dtype = cur_dtype - return prev_dtype - - -@contextlib.contextmanager -def precision(dtype): - if isinstance(dtype, str): - dtype = str_dtype_to_trt(dtype) - prev_dtype = switch_net_dtype(dtype) - yield - switch_net_dtype(prev_dtype) - - -def serialize_engine(engine, path): - logger.info(f"Serializing engine to {path}...") - tik = time.time() - if isinstance(engine, trt.ICudaEngine): - engine = engine.serialize() - with open(path, "wb") as f: - f.write(engine) - tok = time.time() - t = time.strftime("%H:%M:%S", time.gmtime(tok - tik)) - logger.info(f"Engine serialized. Total time: {t}") - - -def deserialize_engine(path): - runtime = trt.Runtime(logger.trt_logger) - with open(path, "rb") as f: - logger.info(f"Loading engine from {path}...") - tik = time.time() - - engine = runtime.deserialize_cuda_engine(f.read()) - assert engine is not None - - tok = time.time() - t = time.strftime("%H:%M:%S", time.gmtime(tok - tik)) - logger.info(f"Engine loaded. Total time: {t}") - return engine - - -_field_dtype_to_np_dtype_dict = { - trt.PluginFieldType.FLOAT16: np.float16, - trt.PluginFieldType.FLOAT32: np.float32, - trt.PluginFieldType.FLOAT64: np.float64, - trt.PluginFieldType.INT8: np.int8, - trt.PluginFieldType.INT16: np.int16, - trt.PluginFieldType.INT32: np.int32, -} - - -def field_dtype_to_np_dtype(dtype): - ret = _field_dtype_to_np_dtype_dict.get(dtype) - assert ret is not None, f"Unsupported dtype: {dtype}" - return ret - - -def convert_capsule_to_void_p(capsule): - ctypes.pythonapi.PyCapsule_GetPointer.restype = ctypes.c_void_p - ctypes.pythonapi.PyCapsule_GetPointer.argtypes = [ctypes.py_object, ctypes.c_char_p] - return ctypes.pythonapi.PyCapsule_GetPointer(capsule, None) - - -def get_nparray_from_void_p(void_pointer, elem_size, field_dtype): - ctypes.pythonapi.PyMemoryView_FromMemory.restype = ctypes.py_object - ctypes.pythonapi.PyMemoryView_FromMemory.argtypes = [ - ctypes.c_char_p, - ctypes.c_ssize_t, - ctypes.c_int, - ] - logger.info(f"get_nparray: pointer = {void_pointer}, elem_size = {elem_size}") - char_pointer = ctypes.cast(void_pointer, ctypes.POINTER(ctypes.c_char)) - np_dtype = field_dtype_to_np_dtype(field_dtype) - buf_bytes = elem_size * np.dtype(np_dtype).itemsize - logger.info(f"get_nparray: buf_bytes = {buf_bytes}") - mem_view = ctypes.pythonapi.PyMemoryView_FromMemory( - char_pointer, buf_bytes, 0 - ) # number 0 represents PyBUF_READ - logger.info(f"get_nparray: mem_view = {mem_view}, field_dtype = {field_dtype}") - buf = np.frombuffer(mem_view, np_dtype) - return buf - - -def get_scalar_from_field(field): - void_p = convert_capsule_to_void_p(field.data) - np_array = get_nparray_from_void_p(void_p, 1, field.type) - return np_array[0] - - -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 - - -def check_max_num_tokens( - max_num_tokens, - opt_num_tokens, - max_batch_size, - max_input_len, - max_seq_len, - max_beam_width, - remove_input_padding, - enable_context_fmha, - tokens_per_block, - multiple_profiles, -): - if not remove_input_padding: - if max_num_tokens is not None or opt_num_tokens is not None: - max_num_tokens = max_batch_size * max_seq_len - logger.warning( - "remove_input_padding is not enabled, the specified " - "max_num_tokens/opt_num_tokens will be ignored." - ) - return max_num_tokens, opt_num_tokens - else: - if max_num_tokens is None: - max_num_tokens = max_seq_len * max_batch_size - logger.warning( - "remove_input_padding is enabled, while max_num_tokens " - "is not set, setting to max_batch_size*max_seq_len. \n" - "It may not be optimal to set max_num_tokens=max_batch_size*max_seq_len " - "when remove_input_padding is enabled, because the number " - "of packed input tokens are very likely to be smaller, " - "we strongly recommend to set max_num_tokens according " - "to your workloads." - ) - if opt_num_tokens is None and not multiple_profiles: - opt_num_tokens = min(max_batch_size * max_beam_width, max_num_tokens) - logger.warning( - "remove_input_padding is enabled, while opt_num_tokens " - "is not set, setting to max_batch_size*max_beam_width. \n" - ) - if max_num_tokens > 16384: - logger.warning( - "Specifying a `max_num_tokens` larger than 16384 is usually " - "not recommended, we do not expect perf gain with that and too " - "large `max_num_tokens` could possibly exceed the TensorRT " - "tensor volume, causing runtime errors. " - f"Got `max_num_tokens` = {max_num_tokens}" - ) - if max_num_tokens > max_seq_len * max_batch_size: - logger.warning( - f"max_num_tokens ({max_num_tokens}) shouldn't be greater than " - f"max_seq_len * max_batch_size ({max_seq_len * max_batch_size}), " - f"specifying to max_seq_len * max_batch_size ({max_seq_len * max_batch_size})." - ) - max_num_tokens = max_seq_len * max_batch_size - if max_num_tokens < max_input_len and not enable_context_fmha: - logger.warning( - f"When enable_context_fmha is not turned on, max_num_tokens ({max_num_tokens}) " - f"should be at least max_input_len ({max_input_len}), specifying to " - f"max_input_len ({max_input_len})." - ) - max_num_tokens = max_input_len - elif max_num_tokens < tokens_per_block and enable_context_fmha: - logger.warning( - f"When enable_context_fmha is turned on, max_num_tokens ({max_num_tokens}) " - f"should be at least tokens_per_block ({tokens_per_block}), specifying to " - f"tokens_per_block ({tokens_per_block}). At this time, you also need to enable " - f"context chunking at runtime, otherwise you may encounter errors." - ) - max_num_tokens = tokens_per_block - - if opt_num_tokens is not None and opt_num_tokens > max_num_tokens: - logger.warning( - f"opt_num_tokens ({opt_num_tokens}) shouldn't be greater than " - f"max_num_tokens ({max_num_tokens}), " - f"specifying to max_num_tokens ({max_num_tokens})." - ) - opt_num_tokens = max_num_tokens - - return max_num_tokens, opt_num_tokens diff --git a/tensorrt_llm/_mnnvl_utils.py b/tensorrt_llm/_mnnvl_utils.py index af4157786c06..6a2a2ad49e2e 100644 --- a/tensorrt_llm/_mnnvl_utils.py +++ b/tensorrt_llm/_mnnvl_utils.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# 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"); @@ -355,6 +355,12 @@ def close_mnnvl_memory(cls, ptr: int): @staticmethod @functools.cache def support_nvlink(dev_id: int, need_all_up: bool = True): + # ensure nvml is initialized; do not rely on other modules having + # initialized it as an import side effect. + try: + pynvml.nvmlDeviceGetCount() + except pynvml.NVMLError_Uninitialized: + pynvml.nvmlInit() handle = pynvml.nvmlDeviceGetHandleByIndex(dev_id) link_count = pynvml.NVML_NVLINK_MAX_LINKS active_links = 0 diff --git a/tensorrt_llm/_tensorrt_engine/__init__.py b/tensorrt_llm/_tensorrt_engine/__init__.py deleted file mode 100644 index 39669a168fd3..000000000000 --- a/tensorrt_llm/_tensorrt_engine/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -from tensorrt_llm.llmapi.llm import _TrtLLM as LLM - -__all__ = ['LLM'] diff --git a/tensorrt_llm/_torch/auto_deploy/llm_args.py b/tensorrt_llm/_torch/auto_deploy/llm_args.py index e1129641e5d6..f0c2af09cf6b 100644 --- a/tensorrt_llm/_torch/auto_deploy/llm_args.py +++ b/tensorrt_llm/_torch/auto_deploy/llm_args.py @@ -20,7 +20,6 @@ from pydantic_settings import BaseSettings, SettingsConfigDict from tensorrt_llm.llmapi.llm_args import ( - BuildConfig, EagleDecodingConfig, MTPDecodingConfig, TorchLlmArgs, @@ -75,13 +74,6 @@ class LlmArgs(DynamicYamlMixInForSettings, TorchLlmArgs, BaseSettings): model_config = _get_config_dict() - build_config: Optional[BuildConfig] = Field( - default_factory=BuildConfig, - description="!!! DO NOT USE !!! Internal only; needed for BaseLlmArgs compatibility.", - exclude_from_json=True, - frozen=True, - repr=False, - ) backend: Literal["_autodeploy"] = Field( default="_autodeploy", description="The backend to use for this LLM instance.", @@ -101,12 +93,6 @@ def ensure_no_beam_search(cls, value: Any) -> Any: raise ValueError("AutoDeploy does not support beam search (max_beam_width > 1).") return value - @field_validator("build_config", mode="before") - @classmethod - def ensure_no_build_config(cls, value: Any, info: ValidationInfo) -> Any: - msg = "build_config is not in use by AutoDeploy's LlmArgs" - return _check_for_default_value_only(cls, value, info, msg) - @field_validator( "tensor_parallel_size", "pipeline_parallel_size", diff --git a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py index 3d2695075986..1d24db4fcc64 100644 --- a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py @@ -25,10 +25,11 @@ import tensorrt_llm.quantization.utils.fp4_utils as fp4_utils from tensorrt_llm import deep_gemm +from tensorrt_llm._torch.distributed.allreduce_helper import \ + CustomAllReduceHelper from tensorrt_llm._utils import get_sm_version from tensorrt_llm.functional import AllReduceFusionOp, AllReduceStrategy from tensorrt_llm.logger import logger -from tensorrt_llm.plugin.plugin import CustomAllReduceHelper from tensorrt_llm.quantization.utils import fp8_quantize from ..autotuner import (AutoTuner, ConstraintSpec, DistributedTuningStrategy, diff --git a/tensorrt_llm/_torch/disaggregation/native/bounce/gather_scatter.py b/tensorrt_llm/_torch/disaggregation/native/bounce/gather_scatter.py index 9f9707808f69..21b81b250ff5 100644 --- a/tensorrt_llm/_torch/disaggregation/native/bounce/gather_scatter.py +++ b/tensorrt_llm/_torch/disaggregation/native/bounce/gather_scatter.py @@ -26,8 +26,7 @@ except ImportError: from cuda import cudart -from tensorrt_llm._utils import prefer_pinned -from tensorrt_llm.runtime.generation import CUASSERT +from tensorrt_llm._utils import CUASSERT, prefer_pinned try: import torch diff --git a/tensorrt_llm/_torch/disaggregation/native/bounce/impl.py b/tensorrt_llm/_torch/disaggregation/native/bounce/impl.py index d285a548c37a..8fefc2c761e4 100644 --- a/tensorrt_llm/_torch/disaggregation/native/bounce/impl.py +++ b/tensorrt_llm/_torch/disaggregation/native/bounce/impl.py @@ -34,7 +34,7 @@ TransferOp, TransferRequest, ) -from tensorrt_llm.runtime.generation import CUASSERT +from tensorrt_llm._utils import CUASSERT from .buffer import SlotAllocator from .config import SizingContext, fit_within_free diff --git a/tensorrt_llm/_torch/disaggregation/native/transfer.py b/tensorrt_llm/_torch/disaggregation/native/transfer.py index b808e6c360b3..fe5105bc8652 100644 --- a/tensorrt_llm/_torch/disaggregation/native/transfer.py +++ b/tensorrt_llm/_torch/disaggregation/native/transfer.py @@ -51,9 +51,8 @@ from tensorrt_llm._torch.disaggregation.resource.utils import get_unique_pool_memory_descs from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager -from tensorrt_llm._utils import nvtx_range +from tensorrt_llm._utils import CUASSERT, nvtx_range from tensorrt_llm.disaggregated_params import DisaggregatedParams, DisaggScheduleStyle -from tensorrt_llm.runtime.generation import CUASSERT if TYPE_CHECKING: from .bounce import Config diff --git a/tensorrt_llm/_torch/distributed/allreduce_helper.py b/tensorrt_llm/_torch/distributed/allreduce_helper.py new file mode 100644 index 000000000000..ef555be727f6 --- /dev/null +++ b/tensorrt_llm/_torch/distributed/allreduce_helper.py @@ -0,0 +1,155 @@ +# 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. + +import os +from typing import List, Tuple + +import torch + +from tensorrt_llm._ipc_utils import IpcMemory, can_access_peer +from tensorrt_llm.bindings.internal.runtime import ( + lamport_initialize, + max_workspace_size_lowprecision, +) +from tensorrt_llm.logger import logger +from tensorrt_llm.mapping import Mapping + + +def force_all_reduce_deterministic(): + return ( + os.getenv("FORCE_DETERMINISTIC", "0") == "1" + or os.getenv("FORCE_ALL_REDUCE_DETERMINISTIC", "0") == "1" + ) + + +class CustomAllReduceHelper: + """ + Helper for custom all-reduce workspace/IPC buffer allocation and sizing, + used by the PyTorch backend. + """ + + POINTERS_PER_RANK = 7 + POINTERS_OF_COUNTER = 3 + + @staticmethod + def max_workspace_size_auto(tp_size: int, support_deterministic=True) -> int: + """Calculate workspace size for allreduce fusion kernel. + + The workspace is used for lamport buffers in the fusion kernel. + Required size calculation: + - Each GPU needs 3 sub-buffers (for triple buffering) + - Each sub-buffer stores: max_num_tokens * hidden_size * dtype_size (bf16=2) + - The lamport allocation multiplies by tp_size, so: + lamport_size = 3 * size * tp_size (per GPU) + + Example: Llama 8B (hidden=4096), max_tokens=8192, bf16, TP=4 + - Data per sub-buffer: 8192 * 4096 * 2 = 64 MiB + - Total lamport: 3 * 64MB * 4 = 768 MiB per GPU + - Required 'size' parameter: 64 MiB (gets multiplied by tp_size in allocation) + + Default (67,108,864 = 64 MiB) supports: + - Models up to hidden_size=4096 with max_num_tokens=8192 + - Or hidden_size=8192 with max_num_tokens=4096 + + Override with TRTLLM_ALLREDUCE_FUSION_WORKSPACE_SIZE env var if needed for larger models. + """ + if force_all_reduce_deterministic() and support_deterministic: + workspace_size = os.getenv("FORCE_ALLREDUCE_KERNEL_WORKSPACE_SIZE", "1000000000") + return int(workspace_size) + + # Allow override via environment variable for edge cases + workspace_size_env = os.getenv("TRTLLM_ALLREDUCE_FUSION_WORKSPACE_SIZE") + if workspace_size_env: + size = int(workspace_size_env) + logger.info( + f"Using custom allreduce fusion workspace size: {size} bytes ({size / (1024**2):.1f} MiB)" + ) + return size + + # Default: 64 MiB - supports most common model configurations + # Increase via env var if you see CUDA illegal memory access errors with large models + default_size = 67_108_864 # Exactly 64 MiB + return default_size + + @staticmethod + def max_workspace_size_lowprecision(tp_size: int) -> int: + return max_workspace_size_lowprecision(tp_size) + + @staticmethod + def initialize_lowprecision_buffers(workspace: "torch.tensor", tp_size: int) -> None: + return torch.ops.trtllm.initialize_static_lowprecision_buffers(workspace, tp_size) + + @staticmethod + def allocate_lowprecision_workspace( + mapping: Mapping, size: int + ) -> Tuple[List[IpcMemory], "torch.tensor"]: + # Force pull mode and disable lamport when force deterministic is enabled, for reducing device memory usage. + is_p2p_supported = can_access_peer(mapping) + ipc_buffers_size = size + ipc_buffers_ping = IpcMemory(mapping, ipc_buffers_size, is_p2p_supported) + ipc_buffers_pong = IpcMemory(mapping, ipc_buffers_size, is_p2p_supported) + ipc_barriers_in = IpcMemory( + mapping, IpcMemory.IPC_BARRIERS_SIZE_PER_GPU * mapping.tp_size * 2, is_p2p_supported + ) + ipc_barriers_out = IpcMemory( + mapping, IpcMemory.IPC_BARRIERS_SIZE_PER_GPU * mapping.tp_size * 2, is_p2p_supported + ) + buffers = [ipc_buffers_ping, ipc_buffers_pong, ipc_barriers_in, ipc_barriers_out] + + return buffers, torch.tensor( + ipc_buffers_ping.serialize() + + ipc_buffers_pong.serialize() + + ipc_barriers_in.serialize() + + ipc_barriers_out.serialize() + + [0] + + [0], + dtype=torch.int64, + device="cpu", + ) + + @staticmethod + def allocate_allreduce_fusion_workspace( + mapping: Mapping, size: int + ) -> Tuple[List[IpcMemory], "torch.tensor"]: + is_p2p_supported = can_access_peer(mapping) + ipc_buffers_size = size * mapping.tp_size + ipc_buffers = IpcMemory(mapping, ipc_buffers_size, is_p2p_supported) + ipc_barriers = IpcMemory(mapping, 256 * mapping.tp_size, is_p2p_supported) + lamport_buffers_size = size * mapping.tp_size + lamport_buffers = IpcMemory(mapping, 3 * lamport_buffers_size, is_p2p_supported) + if is_p2p_supported: + lamport_initialize( + lamport_buffers.local_ptr, + 3 * lamport_buffers_size, + ) + # flag_buffer[0], atomic flag read counter + # flag_buffer[1], non-lamport flag + # flag_buffer[2], lamport flag + flag_buffer = torch.tensor([0, 0, 0], dtype=torch.int, device="cuda") + # layout_buffer[0], clear size for next lamport kernel + # layout_buffer[1], triple buffer offset for lamport kernel + layout_buffer = torch.tensor([0, lamport_buffers_size], dtype=torch.int64, device="cuda") + + buffers = [ipc_buffers, ipc_barriers, lamport_buffers, flag_buffer, layout_buffer] + + return buffers, torch.tensor( + ipc_buffers.serialize() + + ipc_barriers.serialize() + + lamport_buffers.serialize() + + [flag_buffer.data_ptr()] + + [layout_buffer.data_ptr()], + dtype=torch.int64, + device="cuda", + ) diff --git a/tensorrt_llm/_torch/distributed/ops.py b/tensorrt_llm/_torch/distributed/ops.py index 63bf9d76c2c0..95570ccdce4d 100644 --- a/tensorrt_llm/_torch/distributed/ops.py +++ b/tensorrt_llm/_torch/distributed/ops.py @@ -8,6 +8,8 @@ from torch import nn from tensorrt_llm._mnnvl_utils import HelixCpMnnvlMemory, MnnvlMemory +from tensorrt_llm._torch.distributed.allreduce_helper import \ + CustomAllReduceHelper from tensorrt_llm._torch.distributed.symm_mem_allreduce import \ SymmetricMemoryAllReduce from tensorrt_llm._torch.utils import get_model_extra_attrs @@ -19,7 +21,6 @@ AllReduceStrategy, MoEAllReduceParams) from tensorrt_llm.logger import logger from tensorrt_llm.mapping import Mapping -from tensorrt_llm.plugin.plugin import CustomAllReduceHelper # Feature flag: GEMM→NCCL-window zero-copy (writes GEMM output directly into # the window buffer so the allreduce needs no extra copy). On by default diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/qwen3_moe_weight_mapper.py b/tensorrt_llm/_torch/models/checkpoints/hf/qwen3_moe_weight_mapper.py index 2ad24e129be7..57efbbec82bc 100644 --- a/tensorrt_llm/_torch/models/checkpoints/hf/qwen3_moe_weight_mapper.py +++ b/tensorrt_llm/_torch/models/checkpoints/hf/qwen3_moe_weight_mapper.py @@ -1,20 +1,15 @@ -from typing import Union - from torch import nn from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.models.checkpoints.hf.qwen2_moe_weight_mapper import \ Qwen2MoeHfWeightMapper from tensorrt_llm._torch.models.modeling_utils import register_mapper -from tensorrt_llm.models.modeling_utils import DecoderModelForCausalLM @register_mapper("HF", "Qwen3MoeForCausalLM") class Qwen3MoeHfWeightMapper(Qwen2MoeHfWeightMapper): - def init_model_and_config(self, model: Union[nn.Module, - DecoderModelForCausalLM], - config: ModelConfig): + def init_model_and_config(self, model: nn.Module, config: ModelConfig): super().init_model_and_config(model, config) def should_skip_module(self, module_name: str) -> bool: diff --git a/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_two_sided_flashinfer.py b/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_two_sided_flashinfer.py index 5f7e5e194bac..879436f58b2b 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_two_sided_flashinfer.py +++ b/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_two_sided_flashinfer.py @@ -24,6 +24,7 @@ import os from typing import List, Optional, Tuple +import pynvml import torch from flashinfer.comm.mnnvl import MnnvlMemory as flashinfer_MnnvlMemory from flashinfer.comm.trtllm_alltoall import MnnvlMoe as flashinfer_MnnvlMoe @@ -89,6 +90,13 @@ def is_platform_supported() -> bool: """ Check if NVLINK two-sided comm is supported on current hardware. """ + # flashinfer's MnnvlMemory.supports_mnnvl() queries NVML without + # initializing it (only its initialize() guards), so satisfy that + # precondition on our side of the boundary. + try: + pynvml.nvmlDeviceGetCount() + except pynvml.NVMLError_Uninitialized: + pynvml.nvmlInit() return flashinfer_MnnvlMemory.supports_mnnvl() def supports_post_quant_dispatch(self) -> bool: diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 1a8effc634a8..bd17f1095b04 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -22,9 +22,10 @@ except ImportError: from cuda import cudart -from tensorrt_llm._utils import (customized_gc_thresholds, is_trace_enabled, - mpi_comm, mpi_disabled, nvtx_range, - set_thread_local_mpi_comm, trace_func) +from tensorrt_llm._utils import (CUASSERT, customized_gc_thresholds, + is_trace_enabled, mpi_comm, mpi_disabled, + nvtx_range, set_thread_local_mpi_comm, + trace_func) from tensorrt_llm.bindings.executor import (DisServingRequestStats, FinishReason, InflightBatchingStats, IterationStats, KvCacheStats, @@ -38,7 +39,6 @@ from tensorrt_llm.llmapi.llm_args import PeftCacheConfig, WaitingQueuePolicy from tensorrt_llm.logger import logger from tensorrt_llm.mapping import CpType -from tensorrt_llm.runtime.generation import CUASSERT from tensorrt_llm.runtime.kv_cache_manager_v2._exceptions import OutOfPagesError from tensorrt_llm.tools.layer_wise_benchmarks import get_calibrator from tensorrt_llm.tools.profiler.host_profile_tools.host_profiler import ( diff --git a/tensorrt_llm/_utils.py b/tensorrt_llm/_utils.py index cf7916446321..67f1c9d3e950 100644 --- a/tensorrt_llm/_utils.py +++ b/tensorrt_llm/_utils.py @@ -38,12 +38,15 @@ import nvtx from mpi4py import MPI from mpi4py.util import pkl5 -from packaging import version from typing_extensions import ParamSpec # isort: off import torch -import tensorrt as trt + +try: + from cuda.bindings import runtime as cudart +except ImportError: + from cuda import cudart try: from pynvml import ( @@ -107,6 +110,17 @@ def numpy_to_torch(x): return torch.from_numpy(x) +def CUASSERT(cuda_ret): + err = cuda_ret[0] + if err != cudart.cudaError_t.cudaSuccess: + raise RuntimeError( + f"CUDA ERROR: {err}, error code reference: https://nvidia.github.io/cuda-python/module/cudart.html#cuda.cudart.cudaError_t" + ) + if len(cuda_ret) > 1: + return cuda_ret[1:] + return None + + def numpy_to_dtype(x, dtype: str): if str_dtype_to_np(dtype) == x.dtype: return x @@ -124,28 +138,12 @@ def numpy_to_dtype(x, dtype: str): bool_array = partial(np.array, dtype=np.bool_) -def dims_array(x): - is_int64_dims = True - try: - trt.Dims([np.iinfo(np.int64).max]) - except TypeError: - is_int64_dims = False - return int64_array(x) if is_int64_dims else int32_array(x) - - def bf16_array(x): x = torch.tensor(x, dtype=torch.bfloat16) x = torch_to_numpy(x) return x -def numpy_array(data, trt_dtype): - # convenient wrapper due to numpy not support bf16 yet - if trt_dtype == trt.bfloat16: - return bf16_array(data) - return np.array(data, trt_dtype_to_np(trt_dtype)) - - def copy_torch_to_numpy(x: torch.Tensor, ndarray: np.array): if x.dtype == torch.bfloat16: torch.from_numpy(ndarray.view(np.int16)).copy_(x.view(torch.int16)) @@ -156,18 +154,6 @@ def copy_torch_to_numpy(x: torch.Tensor, ndarray: np.array): return ndarray -def trt_version(): - return trt.__version__ - - -def trt_gte(major: int, minor: int = 0): - """ - Check if TRT version is greater than or equal to major.minor - """ - trt_ver = version.parse(trt_version()) - return trt_ver.major >= major and trt_ver.minor >= minor - - def torch_version(): return torch.__version__ @@ -276,82 +262,6 @@ def torch_dtype_to_str(dtype): return _torch_dtype_to_str_dict[dtype] -_str_to_trt_dtype_dict = dict(float16=trt.float16, - float32=trt.float32, - int64=trt.int64, - int32=trt.int32, - int8=trt.int8, - bool=trt.bool, - bfloat16=trt.bfloat16, - fp8=trt.fp8, - nvfp4=trt.fp4) - - -def str_dtype_to_trt(dtype): - if dtype == "fp4": - # Special handling for FP4 since CI's trt version is not recent enough. - if not hasattr(trt, 'fp4'): - raise ValueError( - "fp4 unsupported, trt version needs to be upgraded.") - return trt.fp4 - - ret = _str_to_trt_dtype_dict.get(dtype) - assert ret is not None, f'Unsupported dtype: {dtype}' - return ret - - -_trt_to_str_dtype_dict = {v: k for k, v in _str_to_trt_dtype_dict.items()} - - -def trt_dtype_to_str(dtype: trt.DataType) -> str: - assert isinstance(dtype, trt.DataType) - return _trt_to_str_dtype_dict[dtype] - - -_np_to_trt_dtype_dict = { - np.int8: trt.int8, - np.int32: trt.int32, - np.int64: trt.int64, - np.float16: trt.float16, - np.float32: trt.float32, - np.bool_: trt.bool, - - # hash of np.dtype('int32') != np.int32 - np.dtype('int8'): trt.int8, - np.dtype('int32'): trt.int32, - np.dtype('int64'): trt.int64, - np.dtype('float16'): trt.float16, - np.dtype('float32'): trt.float32, - np.dtype('bool'): trt.bool, - np_bfloat16: trt.bfloat16, - np_float8: trt.fp8, -} - - -def np_dtype_to_trt(dtype): - ret = _np_to_trt_dtype_dict.get(dtype) - assert ret is not None, f'Unsupported dtype: {dtype}' - return ret - - -_trt_to_np_dtype_dict = { - trt.int8: np.int8, - trt.int32: np.int32, - trt.int64: np.int64, - trt.float16: np.float16, - trt.float32: np.float32, - trt.bool: np.bool_, - trt.bfloat16: np_bfloat16, - trt.fp8: np_float8, -} - - -def trt_dtype_to_np(dtype): - ret = _trt_to_np_dtype_dict.get(dtype) - assert ret is not None, f'Unsupported dtype: {dtype}' - return ret - - _torch_to_np_dtype_dict = { torch.bool: np.bool_, torch.uint8: np.uint8, @@ -398,54 +308,6 @@ def np_dtype_to_torch(dtype): return ret -_trt_to_torch_dtype_dict = { - trt.float16: torch.float16, - trt.float32: torch.float32, - trt.int64: torch.int64, - trt.int32: torch.int32, - trt.int8: torch.int8, - trt.bool: torch.bool, - trt.bfloat16: torch.bfloat16, - trt.fp8: torch.float8_e4m3fn, -} - - -def trt_dtype_to_torch(dtype): - ret = _trt_to_torch_dtype_dict.get(dtype) - assert ret is not None, f'Unsupported dtype: {dtype}' - return ret - - -def is_same_dtype(type_a: Union[str, trt.DataType], - type_b: Union[str, trt.DataType]) -> bool: - if isinstance(type_a, str): - type_a = str_dtype_to_trt(type_a) - - if isinstance(type_b, str): - type_b = str_dtype_to_trt(type_b) - - return type_a == type_b - - -_torch_to_trt_dtype_dict = { - torch.float16: trt.float16, - torch.float32: trt.float32, - torch.int64: trt.int64, - torch.int32: trt.int32, - torch.int8: trt.int8, - torch.float8_e4m3fn: trt.fp8, - torch.qint8: trt.int8, - torch.bool: trt.bool, - torch.bfloat16: trt.bfloat16 -} - - -def torch_dtype_to_trt(dtype): - ret = _torch_to_trt_dtype_dict.get(dtype) - assert ret is not None, f'Unsupported dtype: {dtype}' - return ret - - _torch_to_binding_dtype_dict = { torch.float16: DataType.HALF, torch.float32: DataType.FLOAT, @@ -1068,7 +930,7 @@ class TensorWrapper: def __init__( self, data_ptr: int, - dtype: Union[torch.dtype, str, np.dtype, trt.DataType, DataType], + dtype: Union[torch.dtype, str, np.dtype, DataType], shape: Sequence[int], strides: Optional[Sequence[int]] = None, ): @@ -1090,16 +952,13 @@ def shape(self): return getattr(self, "_shape", None) @dtype.setter - def dtype(self, dtype: Union[torch.dtype, str, np.dtype, trt.DataType, - DataType]): + def dtype(self, dtype: Union[torch.dtype, str, np.dtype, DataType]): if isinstance(dtype, torch.dtype): self._dtype = dtype elif isinstance(dtype, str): self._dtype = str_dtype_to_torch(dtype) elif isinstance(dtype, np.dtype): self._dtype = np_dtype_to_torch(dtype) - elif isinstance(dtype, trt.DataType): - self._dtype = trt_dtype_to_torch(dtype) elif isinstance(dtype, DataType): self._dtype = binding_to_torch_dtype(dtype) else: @@ -1128,10 +987,6 @@ def __cuda_array_interface__(self): 3, } - @staticmethod - def from_trt_desc(desc: trt.PluginTensorDesc, pointer: int): - return TensorWrapper(pointer, trt_dtype_to_torch(desc.type), desc.dims) - def convert_to_torch_tensor( tensor: Union[TensorWrapper, torch.Tensor]) -> torch.Tensor: diff --git a/tensorrt_llm/bench/benchmark/__init__.py b/tensorrt_llm/bench/benchmark/__init__.py index 53fa8b337d66..8386e115f505 100644 --- a/tensorrt_llm/bench/benchmark/__init__.py +++ b/tensorrt_llm/bench/benchmark/__init__.py @@ -5,7 +5,6 @@ from pydantic import AliasChoices, BaseModel, Field from tensorrt_llm import LLM as PyTorchLLM -from tensorrt_llm._tensorrt_engine import LLM from tensorrt_llm.bench.benchmark.utils.processes import IterationWriter from tensorrt_llm.bench.build.build import get_model_config from tensorrt_llm.bench.dataclasses.configuration import RuntimeConfig @@ -128,7 +127,7 @@ def get_llm(runtime_config: RuntimeConfig, kwargs: dict): Returns: An instance of the appropriate LLM class for the specified backend. """ - llm_cls = LLM + llm_cls = PyTorchLLM if runtime_config.backend != None: ignore_trt_only_args(kwargs, runtime_config.backend) diff --git a/tensorrt_llm/bench/benchmark/utils/asynchronous.py b/tensorrt_llm/bench/benchmark/utils/asynchronous.py index 18782061bccf..6fbf8f1f724b 100644 --- a/tensorrt_llm/bench/benchmark/utils/asynchronous.py +++ b/tensorrt_llm/bench/benchmark/utils/asynchronous.py @@ -12,8 +12,7 @@ from zmq import PUSH from zmq.asyncio import Context -from tensorrt_llm import SamplingParams -from tensorrt_llm._tensorrt_engine import LLM +from tensorrt_llm import LLM, SamplingParams from tensorrt_llm._utils import EnergyMonitor from tensorrt_llm.bench.dataclasses.general import InferenceRequest from tensorrt_llm.bench.dataclasses.reporting import PerfItemTuple, StatsKeeper diff --git a/tensorrt_llm/bench/benchmark/utils/general.py b/tensorrt_llm/bench/benchmark/utils/general.py index e35708df1bd2..e28121b18aba 100755 --- a/tensorrt_llm/bench/benchmark/utils/general.py +++ b/tensorrt_llm/bench/benchmark/utils/general.py @@ -23,7 +23,7 @@ QuantAlgo.NVFP4.value: "fp8", } -ALL_SUPPORTED_BACKENDS = ["pytorch", "_autodeploy", "tensorrt"] +ALL_SUPPORTED_BACKENDS = ["pytorch", "_autodeploy"] def get_settings_from_engine( diff --git a/tensorrt_llm/bench/build/build.py b/tensorrt_llm/bench/build/build.py index 3290828107b4..16bf6d23d9aa 100644 --- a/tensorrt_llm/bench/build/build.py +++ b/tensorrt_llm/bench/build/build.py @@ -1,28 +1,28 @@ from __future__ import annotations from pathlib import Path -from typing import Tuple, get_args -import click -from click_option_group import AllOptionGroup, optgroup - -from tensorrt_llm._torch.pyexecutor.config_utils import is_nemotron_hybrid, is_qwen3_hybrid, load_pretrained_config -from tensorrt_llm.bench.dataclasses.general import BenchmarkEnvironment -from tensorrt_llm.bench.utils.data import create_dataset_from_stream, initialize_tokenizer -from tensorrt_llm.bench.utils import VALID_QUANT_ALGOS -from tensorrt_llm.builder import BuildConfig -from tensorrt_llm._tensorrt_engine import LLM +from typing import Tuple + +from tensorrt_llm._torch.pyexecutor.config_utils import (is_nemotron_hybrid, + is_qwen3_hybrid, + load_pretrained_config) +from tensorrt_llm.bench.build.dataclasses import (ModelConfig, + NemotronHybridConfig, + Qwen3HybridConfig) +from tensorrt_llm.bench.build.tuning import calc_engine_setting +from tensorrt_llm.llmapi.llm_args import TorchLlmArgs from tensorrt_llm.llmapi.llm_utils import QuantConfig from tensorrt_llm.logger import logger from tensorrt_llm.quantization.mode import QuantAlgo -from tensorrt_llm.bench.build.dataclasses import ModelConfig, NemotronHybridConfig, Qwen3HybridConfig -from tensorrt_llm.bench.build.tuning import calc_engine_setting TUNED_QUANTS = { QuantAlgo.NVFP4, QuantAlgo.FP8, QuantAlgo.FP8_BLOCK_SCALES, QuantAlgo.NO_QUANT, None } -DEFAULT_MAX_BATCH_SIZE = BuildConfig.model_fields["max_batch_size"].default -DEFAULT_MAX_NUM_TOKENS = BuildConfig.model_fields["max_num_tokens"].default +# Sourced from TorchLlmArgs so the bench defaults track the args-class field +# defaults and can't drift. +DEFAULT_MAX_BATCH_SIZE = TorchLlmArgs.model_fields["max_batch_size"].default +DEFAULT_MAX_NUM_TOKENS = TorchLlmArgs.model_fields["max_num_tokens"].default def get_benchmark_engine_settings( @@ -98,273 +98,3 @@ def get_model_config(model_name: str, model_path: Path = None) -> ModelConfig: if is_qwen3_hybrid(pretrained_config): return Qwen3HybridConfig.from_hf(model_name, model_path) return ModelConfig.from_hf(model_name, model_path) - - -def apply_build_mode_settings(params): - """ Validate engine build options and update the necessary values for engine - build settings. - """ - dataset_path = params.get("dataset") - max_batch_size = params.get("max_batch_size") - target_input_len = params.get("target_input_len") - max_seq_len = params.get("max_seq_len") - tp_size = params.get("tp_size") - pp_size = params.get("pp_size") - - # Check of engine build method. User must choose one engine build option. - build_options = [dataset_path, max_batch_size, target_input_len] - # If no engine build option is provided, fall back to build engine with - # TRT-LLM's default max_batch_size and max_num_tokens. - if sum([bool(opt) for opt in build_options]) == 0: - logger.warning( - "No engine build option is selected, use TRT-LLM default " - "max_batch_size and max_num_tokens to build the engine.") - params['max_batch_size'] = DEFAULT_MAX_BATCH_SIZE - params['max_num_tokens'] = DEFAULT_MAX_NUM_TOKENS - elif sum([bool(opt) for opt in build_options]) > 1: - raise ValueError("Multiple engine build options detected, please " - "choose only one engine build option. Exiting.") - - # Check for supported parallelism mappings: only world size <= 8 for now. - if tp_size * pp_size > 8: - raise ValueError( - f"Parallelism mapping of TP{tp_size}-PP{pp_size} is " - "currently unsupported. Please try with a mapping with <=8 GPUs.") - - # If dataset is not specified, max_seq_len must be provided. - if not dataset_path and not max_seq_len: - raise ValueError("Unspecified max_seq_len for engine build. Exiting.") - - -@click.command(name="build") -@optgroup.group("Engine Configuration", - help="Configuration of the TensorRT LLM engine.") -@optgroup.option( - "--tp_size", - "-tp", - type=int, - default=1, - required=False, - help="Number of tensor parallel shards to run the benchmark with.", -) -@optgroup.option( - "--pp_size", - "-pp", - type=int, - default=1, - required=False, - help="Number of pipeline parallel shards to run the benchmark with.", -) -@optgroup.option( - "--quantization", - "-q", - type=click.Choice(tuple(get_args(VALID_QUANT_ALGOS))), - default=None, - help= - ("The quantization algorithm to be used when benchmarking. See the " - "documentations for more information.\n" - " - https://nvidia.github.io/TensorRT-LLM/precision.html" - " - https://github.com/NVIDIA/TensorRT-LLM/blob/main/docs/source/blogs/quantization-in-TRT-LLM.md" - ), -) -@optgroup.option( - "--max_seq_len", - default=None, - type=click.IntRange(min=1), - help="Maximum total length of one request, including prompt and outputs.", -) -@optgroup.option( - "--no_weights_loading", - type=bool, - default=False, - help= - "Do not load the weights from the checkpoint. Use dummy weights instead.") -@optgroup.option( - "--trust_remote_code", - type=bool, - default=False, - help= - "Trust remote code for the HF models that are not natively implemented in the transformers library. " - "This is needed when using LLM API when loading the HF config to build the engine." -) -@optgroup.group( - "Build Engine with Dataset Information", - cls=AllOptionGroup, - help="Optimize engine build parameters with user-specified dataset " - "statistics, e.g., average input/output length, max sequence length.", -) -@optgroup.option( - "--dataset", - type=click.Path(exists=True, - readable=True, - path_type=Path, - resolve_path=True), - default=None, - help="Dataset file to extract the sequence statistics for engine build.", -) -@optgroup.group( - "Build Engine with IFB Scheduler Limits", - cls=AllOptionGroup, - help="Optimize engine build parameters with user-specified inflight " - "batching scheduler settings.", -) -@optgroup.option( - "--max_batch_size", - default=None, - type=click.IntRange(min=1), - help="Maximum number of requests that the engine can schedule.", -) -@optgroup.option( - "--max_num_tokens", - default=None, - type=click.IntRange(min=1), - help="Maximum number of batched tokens the engine can schedule.", -) -@optgroup.group( - "[Experimental Feature] Build Engine with Tuning Heuristics Hints", - cls=AllOptionGroup, - help="Optimize engine build parameters with user-specified target " - "sequence length information.", -) -@optgroup.option( - "--target_input_len", - default=None, - type=click.IntRange(min=1), - help="Target (average) input length for tuning heuristics.", -) -@optgroup.option( - "--target_output_len", - default=None, - type=click.IntRange(min=1), - help="Target (average) sequence length for tuning heuristics.", -) -@click.pass_obj -def build_command( - bench_env: BenchmarkEnvironment, - **params, -) -> None: - """Build engines for benchmarking.""" - - apply_build_mode_settings(params) - # Collect configuration parameters from CLI parameters. - tp_size = params.get("tp_size") - pp_size = params.get("pp_size") - quantization = params.get("quantization") - max_seq_len: int = params.get("max_seq_len") - # Dataset options - dataset_path: Path = params.get("dataset") - # IFB scheduler options - max_batch_size = params.get("max_batch_size") - max_num_tokens = params.get("max_num_tokens") - # Tuning heuristics options - target_input_len: int = params.get("target_input_len") - target_output_len: int = params.get("target_output_len") - - load_format = "dummy" if params.get("no_weights_loading") else "auto" - trust_remote_code: bool = params.get("trust_remote_code") - model_name = bench_env.model - checkpoint_path = bench_env.checkpoint_path or model_name - model_config = get_model_config(model_name, bench_env.checkpoint_path) - engine_dir = Path(bench_env.workspace, model_name, - f"tp_{tp_size}_pp_{pp_size}") - - # Set the compute quantization. - quant_algo = QuantAlgo(quantization) if quantization is not None else None - quant_config = QuantConfig(quant_algo=quant_algo) - # If the quantization is NVFP4 or FP8, force the KV cache dtype to FP8. - if quant_algo in [QuantAlgo.NVFP4, QuantAlgo.FP8]: - quant_config.kv_cache_quant_algo = QuantAlgo.FP8 - - # Initialize the HF tokenizer for the specified model. - tokenizer = initialize_tokenizer(checkpoint_path) - # If we receive dataset from a path or stdin, parse and gather metadata. - if dataset_path: - logger.info("Found dataset.") - # Dataset Loading and Preparation - with open(dataset_path, "r") as dataset: - metadata, _ = create_dataset_from_stream( - tokenizer, - dataset, - ) - max_seq_len = metadata.max_sequence_length - target_input_len = metadata.avg_isl - target_output_len = metadata.avg_osl - logger.info(metadata.get_summary_for_print()) - - # Use user-specified engine settings if provided. - if max_batch_size and max_num_tokens: - logger.info("Use user-provided max batch size and max num tokens for " - "engine build and benchmark.") - # If not provided, use the engine setting provided by trtllm-bench. - else: - logger.info( - "Max batch size and max num tokens are not provided, " - "use tuning heuristics or pre-defined setting from trtllm-bench.") - max_batch_size, max_num_tokens = get_benchmark_engine_settings( - model_config, - quant_config, - tp_size, - pp_size, - target_input_len, - target_output_len, - ) - - # Construct a TRT-LLM build config. - build_config = BuildConfig(max_batch_size=max_batch_size, - max_seq_len=max_seq_len, - max_num_tokens=max_num_tokens) - - build_config.plugin_config.dtype = model_config.dtype - # Enable multiple profiles and paged context FMHA. - build_config.plugin_config.multiple_profiles = True - # build_config.plugin_config._reduce_fusion = True - - # Enable FHMA, and FP8 FMHA if NVFP4 or FP8 quantization is enabled. - # TODO: Revisit, there is an issue with enabling FHMA. If only - # paged FMHA is enabled with NVFP4 or FP8 quantization, the Builder - # will not enable the FP8 FMHA. - build_config.plugin_config.use_paged_context_fmha = True - if quant_algo in [QuantAlgo.NVFP4, QuantAlgo.FP8]: - build_config.plugin_config.use_fp8_context_fmha = True - # Enable nvfp4 gemm_plugin explicitly for Blackwell - if quant_algo == QuantAlgo.NVFP4: - build_config.plugin_config.gemm_plugin = "nvfp4" - - # Build the LLM engine with the LLMAPI. - llm = LLM(checkpoint_path, - tokenizer, - dtype=model_config.dtype, - tensor_parallel_size=tp_size, - pipeline_parallel_size=pp_size, - build_config=build_config, - quant_config=quant_config, - workspace=str(bench_env.workspace), - load_format=load_format, - trust_remote_code=trust_remote_code, - telemetry_config=bench_env.telemetry_config) - # Save the engine. - llm.save(engine_dir) - llm.shutdown() - - logger.info( - "\n===========================================================\n" - "= ENGINE BUILD INFO\n" - "===========================================================\n" - f"Model Name:\t\t{bench_env.model}\n" - f"Model Path:\t\t{bench_env.checkpoint_path}\n" - f"Workspace Directory:\t{bench_env.workspace}\n" - f"Engine Directory:\t{engine_dir}\n\n" - "===========================================================\n" - "= ENGINE CONFIGURATION DETAILS\n" - "===========================================================\n" - f"Max Sequence Length:\t\t{max_seq_len}\n" - f"Max Batch Size:\t\t\t{max_batch_size}\n" - f"Max Num Tokens:\t\t\t{max_num_tokens}\n" - f"Quantization:\t\t\t{quant_config.quant_algo}\n" - f"KV Cache Dtype:\t\t\t{quant_config.kv_cache_quant_algo}\n" - "===========================================================\n") - - logger.info( - "\n\n===========================================================\n" - f"ENGINE SAVED: {engine_dir}\n" - "===========================================================\n") diff --git a/tensorrt_llm/builder.py b/tensorrt_llm/builder.py deleted file mode 100644 index 764c83db3668..000000000000 --- a/tensorrt_llm/builder.py +++ /dev/null @@ -1,1292 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -import json -import math -import os -import shutil -import time -from pathlib import Path -from typing import Dict, Optional, Union - -import numpy as np -import tensorrt as trt -from pydantic import Field - -from ._common import _is_building, check_max_num_tokens, serialize_engine -from ._deprecation import emit_engine_arch_deprecation -from ._utils import (get_sm_version, np_bfloat16, np_float8, str_dtype_to_trt, - to_json_file, trt_gte) -from .functional import PositionEmbeddingType -from .graph_rewriting import optimize -from .llmapi.kv_cache_type import KVCacheType -from .llmapi.utils import StrictBaseModel -from .logger import logger -from .lora_helper import LoraConfig -from .models import PretrainedConfig, PretrainedModel -from .models.modeling_utils import SpeculativeDecodingMode, optimize_model -from .network import Network, net_guard -from .plugin import PluginConfig -from .quantization import QuantAlgo, QuantMode -from .version import __version__ - - -class ConfigEncoder(json.JSONEncoder): - - def default(self, obj): - if hasattr(obj, 'model_dump'): - # Handle Pydantic models (including DecodingBaseConfig and subclasses) - return obj.model_dump(mode='json') - else: - return super().default(obj) - - -class BuilderConfig(object): - - def __init__(self, **kwargs): - # intentionally use **kwargs, user should never call this ctor directly, - # use Builder.create_builder_config() instead - pass - - def _init(self, trt_builder_config, **kwargs): - self._trt_builder_config = trt_builder_config - for key, value in kwargs.items(): - setattr(self, key, value) - return self - - @property - def trt_builder_config(self) -> trt.IBuilderConfig: - return self._trt_builder_config - - def to_dict(self) -> Dict: - '''return a dict with keys - { - "builder_config": { - # all key values set by the _init function - }, - "plugin_config": { - # the network plugin_config (if any) attached to this BuilderConfig object - # inside the Builder.build_engine - } - } - ''' - config = {'builder_config': {}} - for k in self.__dict__.keys(): - if k not in ['_trt_builder_config', 'plugin_config']: - config['builder_config'][k] = self.__getattribute__(k) - if hasattr(self, 'plugin_config'): - assert isinstance(self.plugin_config, PluginConfig), \ - f"Found unexpected plugin_config object with type: {type(self.plugin_config)}" - config['plugin_config'] = self.plugin_config.model_dump(mode="json") - return config - - -class Builder(): - - _ALLOWED_PRECISIONS = [ - 'float32', 'float16', 'bfloat16', trt.DataType.HALF, trt.DataType.FLOAT, - trt.DataType.BF16 - ] - - def __init__(self): - super().__init__() - self._trt_builder = trt.Builder(logger.trt_logger) - self.strongly_typed = True - - @property - def trt_builder(self) -> trt.Builder: - return self._trt_builder - - def create_network(self) -> Network: - explicit_batch_flag = 0 - # Explicit batch flag will be deprecated in TRT 10 - if "EXPLICIT_BATCH" in trt.NetworkDefinitionCreationFlag.__members__.keys( - ): - explicit_batch_flag = 1 << int( - trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH) - - if self.strongly_typed: - return Network()._init( - self.trt_builder.create_network( - explicit_batch_flag - | (1 << int( - trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED)))) - else: - return Network()._init( - self.trt_builder.create_network(explicit_batch_flag)) - - def create_builder_config(self, - precision: Union[str, trt.DataType], - timing_cache: Union[str, Path, - trt.ITimingCache] = None, - tensor_parallel: int = 1, - use_refit: bool = False, - int8: bool = False, - strongly_typed: bool = True, - force_num_profiles: Optional[int] = None, - profiling_verbosity: str = "layer_names_only", - use_strip_plan: bool = False, - weight_streaming: bool = False, - precision_constraints: Optional[str] = "obey", - **kwargs) -> BuilderConfig: - ''' @brief Create a builder config with given precisions and timing cache - @param precision: one of allowed precisions, defined in Builder._ALLOWED_PRECISIONS - @param timing_cache: a timing cache object or a path to a timing cache file - @param tensor_parallel: number of GPUs used for tensor parallel - @param kwargs: any other arguments users would like to attach to the config object as attributes - @param refit: set to accelerate multi-gpu building, build engine for 1 gpu and refit for the others - @param int8: whether to build with int8 enabled or not. Can't be used together with refit option - @return: A BuilderConfig object, return None if failed - ''' - self.strongly_typed = self.strongly_typed and strongly_typed - - quant_mode = kwargs.get("quant_mode", QuantMode(0)) - if not strongly_typed and precision not in self._ALLOWED_PRECISIONS: - logger.error( - f"precision should be one of {self._ALLOWED_PRECISIONS}") - - config = self.trt_builder.create_builder_config() - if weight_streaming: - config.set_flag(trt.BuilderFlag.WEIGHT_STREAMING) - if not self.strongly_typed: - fp8 = quant_mode.has_fp8_qdq() or quant_mode.has_fp8_kv_cache() - if precision == 'float16' or precision == trt.DataType.HALF: - config.set_flag(trt.BuilderFlag.FP16) - if precision_constraints == 'obey': - config.set_flag(trt.BuilderFlag.OBEY_PRECISION_CONSTRAINTS) - elif precision == 'bfloat16' or precision == trt.DataType.BF16: - config.set_flag(trt.BuilderFlag.BF16) - if precision_constraints == 'obey': - config.set_flag(trt.BuilderFlag.OBEY_PRECISION_CONSTRAINTS) - if int8: - config.set_flag(trt.BuilderFlag.INT8) - if fp8: - config.set_flag(trt.BuilderFlag.FP8) - if precision_constraints == 'obey': - config.set_flag(trt.BuilderFlag.OBEY_PRECISION_CONSTRAINTS) - - if use_refit: - config.set_flag(trt.BuilderFlag.REFIT) - - # Use fine-grained refit when strip plan is enabled in TRT10.2+. - if use_strip_plan: - config.set_flag(trt.BuilderFlag.REFIT_INDIVIDUAL) - - if use_strip_plan: - config.set_flag(trt.BuilderFlag.STRIP_PLAN) - - # Set TRT Engine profiling verbosity - if profiling_verbosity == "detailed": - config.profiling_verbosity = trt.ProfilingVerbosity.DETAILED - elif profiling_verbosity == "none": - config.profiling_verbosity = trt.ProfilingVerbosity.NONE - else: - config.profiling_verbosity = trt.ProfilingVerbosity.LAYER_NAMES_ONLY - - # set timing cache - cache = None - if timing_cache is not None: - # use given cache - if isinstance(timing_cache, trt.ITimingCache): - cache = timing_cache - # read cache from file - elif isinstance(timing_cache, - (str, Path)) and os.path.exists(timing_cache): - with open(timing_cache, "rb") as f: - cache = config.create_timing_cache(f.read()) - else: - logger.warning( - "Invalid timing cache, using freshly created one") - if cache is None: - cache = config.create_timing_cache(b"") - # When user does not given any existing cache, internally always created one - # so the cache should never None here - assert cache is not None and isinstance(cache, trt.ITimingCache) - config.set_timing_cache(cache, ignore_mismatch=False) - - # set weight sparsity - weight_sparsity = kwargs.get("weight_sparsity", False) - if weight_sparsity: - config.set_flag(trt.BuilderFlag.SPARSE_WEIGHTS) - - # TODO: remove this constraint after trt 10.6 is integrated - if trt_gte(10, 6): - # set monitor memory - monitor_memory = kwargs.get("monitor_memory", False) - if monitor_memory: - config.set_flag(trt.BuilderFlag.MONITOR_MEMORY) - - return BuilderConfig()._init(config, - precision=precision, - tensor_parallel=tensor_parallel, - use_refit=use_refit, - int8=int8, - force_num_profiles=force_num_profiles, - strongly_typed=self.strongly_typed, - use_strip_plan=use_strip_plan, - **kwargs) - - def _add_optimization_profile(self, network: Network, - builder_config: BuilderConfig): - assert isinstance(builder_config, BuilderConfig) - assert isinstance(network, Network) - input_tensors = network._inputs - if len(input_tensors) == 0: - logger.warning("There are no inputs in the network!") - return - num_profiles = len(list(input_tensors.values())[0].profiles) - force_num_profiles = getattr(builder_config, "force_num_profiles", None) - for i in range(num_profiles): - logger.debug(f'Adding optimization profile {i+1}/{num_profiles}') - profile = self.trt_builder.create_optimization_profile() - for input_name in input_tensors.keys(): - if len(input_tensors[input_name].profiles) == 0: - continue - shape_profile = input_tensors[input_name].profiles[i] - min_shape = [*shape_profile.min] - opt_shape = [*shape_profile.opt] - max_shape = [*shape_profile.max] - profile.set_shape(input_name, min_shape, opt_shape, max_shape) - logger.debug( - f'{input_name}, min: {min_shape}, opt: {opt_shape}, max: {max_shape}, dimension names: {shape_profile.dimension_names}' - ) - ret = builder_config.trt_builder_config.add_optimization_profile( - profile) - logger.debug(f"Added optimization profile: #{ret}") - if force_num_profiles is not None and ( - i + 1 - ) == force_num_profiles and force_num_profiles < num_profiles: - logger.warning( - f"Only adding {force_num_profiles} profiles instead of {num_profiles}." - ) - break - assert self._validate_named_dimensions( - network, builder_config - ), "Validation of the tensor dimension ranges failed, please check the dimension ranges, find the offensive tensor and dimension name in above the error log" - - def _validate_named_dimensions(self, network: Network, - builder_config) -> bool: - ''' - For each profile, validate that the named dimensions of different input tensors in this profile all have same range. - TRT will validate the same condition, validate it earlier to make sure the modeling in TensorRT LLM are correct and - makes the error msg more user friendly. - ''' - valid = True - for profile_idx in range( - builder_config.trt_builder_config.num_optimization_profiles): - dimension_to_range = {} - for input_name, input_tensor in network._inputs.items(): - # it's legal that a Tensor does not have dim_range? - if len(input_tensor.profiles) != 0: - profile = input_tensor.profiles[profile_idx] - for dim_idx, dim_name in enumerate(profile.dimension_names): - if dim_name not in dimension_to_range: - dimension_to_range[dim_name] = [] - min, opt, max = profile.min[dim_idx], profile.opt[ - dim_idx], profile.max[dim_idx] - dimension_to_range[dim_name].append( - (input_name, (min, opt, max))) - for dim, ranges in dimension_to_range.items(): - unique_ranges = set([r[1] for r in ranges]) - logger.debug( - f"Validating dimension:{dim}, ranges for this dim are:{unique_ranges}" - ) - if len(unique_ranges) != 1: - logger.error( - f"Found illegal dimension setting for profile {profile_idx}, dimension name is: {dim}" - ) - logger.error( - "Offensive tensors which have this dimension are:\n" + - "\n".join([f"{r[1]} {dim} {r[0]}" for r in ranges])) - valid = False - return valid - - @_is_building - def refit_engine(self, network: Network, engine_buffer) -> trt.IHostMemory: - ''' - @brief: Refit one TensorRT engine using weights from the network, - user should guarantee that the engine is built with REFIT flag, and the network has the same structure with the engine. - @param engine_buffer: A serialized TensorRT engine. - @param network: Network object. - @return: A serialized TRT engine if refit successfully, None otherwise - ''' - assert isinstance(network, Network) - logger.info('Refit TRT engine') - runtime = trt.Runtime(logger.trt_logger) - engine = runtime.deserialize_cuda_engine(engine_buffer) - - tik = time.time() - - # Refit engine - refitter = trt.Refitter(engine, logger.trt_logger) - if network.named_parameters is not None: - for name, param in network.named_parameters: - if param._get_weights( - ) is None or not refitter.set_named_weights( - name, param._get_weights()): - logger.error(f'Failed to refit weight: {name}') - return None - else: - logger.error( - 'Please set named parameters before building multiple engines.') - return None - - if not refitter.refit_cuda_engine(): - logger.error('Failed to refit engine.') - return None - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - logger.info(f'Total time of refitting {engine.name}: {t}') - serialized_engine = engine.serialize() - return serialized_engine - - @_is_building - def build_engine(self, - network: Network, - builder_config: BuilderConfig, - managed_weights: dict = None) -> trt.IHostMemory: - ''' - @brief: Build one TensorRT engine from the network. - @param network: Network object. - @param builder_config: BuilderConfig object. - @return: A serialized TRT engine. - ''' - assert isinstance(network, Network) - builder_config.plugin_config = network.plugin_config - if builder_config.trt_builder_config.num_optimization_profiles == 0: - self._add_optimization_profile(network, builder_config) - logger.info( - f"Total optimization profiles added: {builder_config.trt_builder_config.num_optimization_profiles}" - ) - engine = None - - tik = time.time() - # Rename weights - if network.named_parameters is not None: - managed_parameters = [] - for name, param in network.named_parameters: - if param.is_managed(network): - assert managed_weights is not None, "managed_weights should be provided when enabled" - managed_parameters.append(param) - param.set_name(name, network) - continue - if param._get_weights(network) is None: - if not param.is_buffer: - logger.debug( - f"Parameter {name} {param.raw_value.shape} {param.raw_value.dtype} was created" - " but unused in forward method, so not materialized to TRT network" - ) - continue - if not param.set_name(name, network): - raise RuntimeError(f'Failed to set weight: {name}') - # This mark_weights_refittable has no side effect when refit_individual is not enabled. - network.trt_network.mark_weights_refittable(name) - - network._fill_weights() - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - logger.info( - f'Total time to initialize the weights in network {network.trt_network.name}: {t}' - ) - - # Build engine - logger.info(f'Build TensorRT engine {network.trt_network.name}') - tik = time.time() - engine = self.trt_builder.build_serialized_network( - network.trt_network, builder_config.trt_builder_config) - assert engine is not None, 'Engine building failed, please check the error log.' - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - logger.info(f'Total time of building {network.trt_network.name}: {t}') - - if managed_weights is not None and network.named_parameters is not None: - for param in managed_parameters: - name = param.name - value: np.ndarray = param._value - if value is None: - logger.error(f'Failed to get weight: {name}') - continue - if param.need_transpose: - # MOE has ndim=3 and uses plugin, no need to transpose - value = value.transpose(1, 0) # WAR for bug 4641821 - managed_weights[name] = value - - return engine - - @staticmethod - def save_timing_cache(builder_config: BuilderConfig, out_path: str) -> bool: - '''Serialize timing cache of given builder config to file specified by out_path - return True if the cache is successfully serialized, False otherwise - ''' - cache = builder_config.trt_builder_config.get_timing_cache() - if cache is None: - logger.warning( - 'No timing cache found in the given builder config, skip saving.' - ) - return False - with cache.serialize() as buffer: - with open(out_path, "wb") as f: - f.write(buffer) - f.flush() - os.fsync(f) - logger.info(f'Timing cache serialized to {out_path}') - return True - - @staticmethod - def save_config(builder_config: BuilderConfig, config_path: str): - config = builder_config.to_dict() - to_json_file(config, config_path) - logger.info(f'Config saved to {config_path}.') - - -class BuildConfig(StrictBaseModel): - """Configuration class for TensorRT LLM engine building parameters. - - This class contains all the configuration parameters needed to build a TensorRT LLM engine, - including sequence length limits, batch sizes, optimization settings, and various features. - """ - max_input_len: int = Field(default=1024, - description="Maximum length of input sequences.") - max_seq_len: Optional[int] = Field( - default=None, - description= - "The maximum possible sequence length for a single request, including both input and generated " - "output tokens.") - opt_batch_size: int = Field( - default=8, description="Optimal batch size for engine optimization.") - max_batch_size: int = Field( - default=2048, description="Maximum batch size the engine can handle.") - max_beam_width: int = Field( - default=1, description="Maximum beam width for beam search decoding.") - max_num_tokens: int = Field( - default=8192, - description="Maximum number of batched input tokens after padding is " - "removed in each batch.") - opt_num_tokens: Optional[int] = Field( - default=None, - description= - "Optimal number of batched input tokens for engine optimization.") - max_prompt_embedding_table_size: int = Field( - default=0, - description="Maximum size of prompt embedding table for prompt tuning.") - kv_cache_type: Optional[KVCacheType] = Field( - default=None, - description= - "Type of KV cache to use (CONTINUOUS or PAGED). If None, defaults to PAGED." - ) - gather_context_logits: bool = Field( - default=False, - description="Whether to gather logits during context phase.") - gather_generation_logits: bool = Field( - default=False, - description="Whether to gather logits during generation phase.") - strongly_typed: bool = Field(default=True, - description="Whether to use strongly_typed.") - force_num_profiles: Optional[int] = Field( - default=None, - description= - "Force a specific number of optimization profiles. If None, auto-determined." - ) - profiling_verbosity: str = Field( - default='layer_names_only', - description= - "Verbosity level for TensorRT profiling ('layer_names_only', 'detailed', 'none')." - ) - enable_debug_output: bool = Field( - default=False, - description="Whether to enable debug output during building.") - max_draft_len: int = Field( - default=0, - description="Maximum length of draft tokens for speculative decoding.") - speculative_decoding_mode: SpeculativeDecodingMode = Field( - default=SpeculativeDecodingMode.NONE, - description="Mode for speculative decoding (NONE, MEDUSA, EAGLE, etc.)." - ) - use_refit: bool = Field( - default=False, - description="Whether to enable engine refitting capabilities.") - input_timing_cache: Optional[str] = Field( - default=None, - description= - "Path to input timing cache file. If None, no input cache used.") - output_timing_cache: str = Field( - default='model.cache', description="Path to output timing cache file.") - lora_config: LoraConfig = Field( - default_factory=LoraConfig, - description="Configuration for LoRA (Low-Rank Adaptation) fine-tuning.") - weight_sparsity: bool = Field( - default=False, - description="Whether to enable weight sparsity optimization.") - weight_streaming: bool = Field( - default=False, - description="Whether to enable weight streaming for large models.") - plugin_config: PluginConfig = Field( - default_factory=PluginConfig, - description="Configuration for TensorRT LLM plugins.") - use_strip_plan: bool = Field( - default=False, - description="Whether to use stripped plan for engine building.") - max_encoder_input_len: int = Field( - default=1024, - description="Maximum encoder input length for encoder-decoder models.") - dry_run: bool = Field( - default=False, - description= - "Whether to perform a dry run without actually building the engine.") - visualize_network: Optional[str] = Field( - default=None, - description= - "Path to save network visualization. If None, no visualization generated." - ) - monitor_memory: bool = Field( - default=False, - description="Whether to monitor memory usage during building.") - use_mrope: bool = Field( - default=False, - description= - "Whether to use Multi-RoPE (Rotary Position Embedding) optimization.") - - # Since we have some overlapping between kv_cache_type, paged_kv_cache, and paged_state (later two will be deprecated in the future), - # we need to handle it given model architecture. - def update_kv_cache_type(self, model_architecture: str): - paged_kv_cache_attr = 'paged_state' if model_architecture in [ - 'MambaForCausalLM', 'RecurrentGemmaForCausalLM' - ] else 'paged_kv_cache' - assert self.plugin_config is not None - paged_kv_cache_val = getattr(self.plugin_config, paged_kv_cache_attr) - - if self.kv_cache_type is not None: - if paged_kv_cache_val is not None: - assert (paged_kv_cache_val == True - and self.kv_cache_type == KVCacheType.PAGED) or ( - paged_kv_cache_val == False - and self.kv_cache_type != KVCacheType.PAGED) - else: - setattr(self.plugin_config, paged_kv_cache_attr, - self.kv_cache_type == KVCacheType.PAGED) - else: - if paged_kv_cache_val is not None: - self.kv_cache_type = KVCacheType.PAGED if paged_kv_cache_val else KVCacheType.CONTINUOUS - else: - self.kv_cache_type = KVCacheType.PAGED - setattr(self.plugin_config, paged_kv_cache_attr, - self.kv_cache_type == KVCacheType.PAGED) - - assert self.kv_cache_type is not None and getattr( - self.plugin_config, paged_kv_cache_attr) is not None - - def override_attri(attr_name, value): - val = getattr(self.plugin_config, attr_name) - if val is not None and val != value: - logger.warning(f'Overriding {attr_name} to {value}') - setattr(self.plugin_config, attr_name, value) - - # Init other paged kvcache attri to false. For RecurrentGemma, we only support paged_state and paged_kv_cache have - # the same values. All other models should only consume either of the value and set other to False. - is_recurrent_gemma = model_architecture == 'RecurrentGemmaForCausalLM' - - if paged_kv_cache_attr == 'paged_state': - override_attri( - 'paged_kv_cache', - getattr(self.plugin_config, paged_kv_cache_attr) - if is_recurrent_gemma else False) - else: - override_attri('paged_state', False) - - @classmethod - def from_json_file(cls, config_file): - with open(config_file) as f: - config = json.load(f) - return BuildConfig(**config) - - -class EngineConfig: - - def __init__(self, pretrained_config: 'PretrainedConfig', - build_config: 'BuildConfig', version: str): - self.pretrained_config = pretrained_config - self.build_config = build_config - self.version = version - - @classmethod - def from_json_file(cls, config_file): - with open(config_file) as f: - return cls.from_json_str(f.read()) - - @classmethod - def from_json_str(cls, config_str): - config = json.loads(config_str) - return cls(PretrainedConfig.from_dict(config['pretrained_config']), - BuildConfig(**config['build_config']), config['version']) - - def to_dict(self): - build_config = self.build_config.model_dump(mode="json") - build_config.pop('dry_run', None) # Not an Engine Characteristic - build_config.pop('visualize_network', - None) # Not an Engine Characteristic - return { - 'version': self.version, - 'pretrained_config': self.pretrained_config.to_dict(), - 'build_config': build_config, - } - - -class Engine: - - def __init__( - self, - config: EngineConfig, - engine: Union[trt.IHostMemory, None], - managed_weights: dict[str, np.ndarray] = {}, - ): - self.config = config - self.engine = engine - self.managed_weights = managed_weights - if self.managed_weights is None: - self.managed_weights = {} - for name, value in self.managed_weights.items(): - if not value.flags['C_CONTIGUOUS']: - self.managed_weights[name] = np.ascontiguousarray(value) - - def save(self, engine_dir: str): - os.makedirs(engine_dir, exist_ok=True) - lora_config = self.config.build_config.lora_config - lora_dirs = lora_config.lora_dir - root_lora_dir = os.path.join(engine_dir, 'lora') - if len(lora_dirs) > 0: - os.makedirs(root_lora_dir, exist_ok=True) - for index, lora_dir in enumerate(lora_dirs): - if lora_config.lora_ckpt_source == 'hf': - target_lora_dir = f"{root_lora_dir}/{index}" - os.makedirs(target_lora_dir, exist_ok=True) - shutil.copy2(os.path.join(lora_dir, 'adapter_config.json'), - target_lora_dir) - weight_file = os.path.join(lora_dir, 'adapter_model.bin') - if os.path.exists(weight_file): - shutil.copy2(weight_file, target_lora_dir) - weight_file = os.path.join(lora_dir, - 'adapter_model.safetensors') - if os.path.exists(weight_file): - shutil.copy2(weight_file, target_lora_dir) - lora_config.lora_dir[index] = f"lora/{index}" - elif lora_config.lora_ckpt_source == 'nemo': - target_lora_file = f"{root_lora_dir}/{index}.nemo" - shutil.copyfile(lora_dir, target_lora_file) - lora_config.lora_dir[index] = f"lora/{index}.nemo" - else: - if os.path.exists(root_lora_dir) and os.path.isdir(root_lora_dir): - shutil.rmtree(root_lora_dir) - if self.config.pretrained_config.mapping.rank == 0: - config_dict = self.config.to_dict() - if self.config.pretrained_config.quant_algo == QuantAlgo.MIXED_PRECISION: - quant_dict = { - 'version': self.config.version, - } - quant_dict.update( - config_dict['pretrained_config']['quantization']) - config_dict['pretrained_config']['quantization'].pop( - 'quantized_layers', None) - with open(os.path.join(engine_dir, 'quant_cfg.json'), - "w", - encoding="utf-8") as f: - json.dump(quant_dict, f, indent=4, cls=ConfigEncoder) - - with open(os.path.join(engine_dir, 'config.json'), - "w", - encoding="utf-8") as f: - json.dump(config_dict, f, indent=4, cls=ConfigEncoder) - if self.engine is not None: - serialize_engine( - self.engine, - os.path.join( - engine_dir, - f'rank{self.config.pretrained_config.mapping.rank}.engine')) - if self.managed_weights is not None and len(self.managed_weights) > 0: - fn = os.path.join( - engine_dir, - f'rank{self.config.pretrained_config.mapping.rank}_managed_weights.safetensors' - ) - serialize_managed_weights(self.managed_weights, fn) - - @classmethod - def from_dir(cls, engine_dir: str, rank: int = 0): - with open(os.path.join(engine_dir, f'rank{rank}.engine'), 'rb') as f: - engine_buffer = f.read() - - mw_path = os.path.join(engine_dir, - f'rank{rank}_managed_weights.safetensors') - managed_weights = deserialize_managed_weights( - mw_path) if os.path.exists(mw_path) else None - - config = EngineConfig.from_json_file( - os.path.join(engine_dir, 'config.json')) - config.pretrained_config.set_rank(rank) - - return cls(config, engine_buffer, managed_weights) - - @classmethod - def from_buffer(cls, - engine_buffer: Union[trt.IHostMemory, bytes], - json_config_str: str, - rank: int = 0): - config = EngineConfig.from_json_str(json_config_str) - config.pretrained_config.set_rank(rank) - return cls(config, engine_buffer) - - -def get_engine_version(engine_dir: str) -> Union[None, str]: - engine_dir = Path(engine_dir) - config_path = engine_dir / "config.json" - with open(config_path, 'r') as f: - config = json.load(f) - - if 'version' not in config: - return None - - return config['version'] - - -def optimize_model_with_config(model: PretrainedModel, - build_config: BuildConfig): - gemm_swiglu_plugin = build_config.plugin_config.gemm_swiglu_plugin - low_latency_gemm_swiglu_plugin = build_config.plugin_config.low_latency_gemm_swiglu_plugin - if gemm_swiglu_plugin or low_latency_gemm_swiglu_plugin: - if not build_config.plugin_config.use_fused_mlp: - raise RuntimeError( - "GemmSwiGLU plugin requires --use_fused_mlp flag") - if gemm_swiglu_plugin not in [ - "fp8" - ] and low_latency_gemm_swiglu_plugin not in ["fp8"]: - raise RuntimeError( - f"GemmSwiGLU plugin currently has limited support: fp8 only, " - f"got: {gemm_swiglu_plugin}" - f"got: {low_latency_gemm_swiglu_plugin}") - - if build_config.plugin_config.lora_plugin is not None: - model.use_lora(build_config.lora_config) - - is_enc_dec = model.config.architecture in ["EncoderModel", "DecoderModel"] - # FusedMLP does not support RecurrentGemma FP8 currently. - is_recurrent_gemma = model.config.architecture in [ - "RecurrentGemmaForCausalLM" - ] - is_fp8 = model.config.quantization.quant_algo == QuantAlgo.FP8 - model = optimize_model( - model, - share_embedding_table=True, - use_ootb_moe=build_config.plugin_config.moe_plugin is None, - use_fused_mlp=(build_config.plugin_config.use_fused_mlp - and not is_enc_dec - and not (is_recurrent_gemma and is_fp8)), - gemm_swiglu_plugin_dtype=gemm_swiglu_plugin, - low_latency_gemm_swiglu_plugin_dtype=low_latency_gemm_swiglu_plugin, - use_fused_rg_lru=is_recurrent_gemma, - use_unfused_qkv_gemm=False, - use_prompt_tuning=(build_config.max_prompt_embedding_table_size > 0), - use_lora=build_config.plugin_config.lora_plugin is not None, - max_lora_rank=build_config.lora_config.max_lora_rank, - use_fp8_context_fmha=(model.config.quantization.quant_algo in [ - QuantAlgo.FP8, QuantAlgo.W4A8_AWQ, QuantAlgo.NVFP4 - ] and build_config.plugin_config.use_fp8_context_fmha), - fuse_fp4_quant=build_config.plugin_config.fuse_fp4_quant, - use_optimize_cross_qkv=True, - use_dora=build_config.plugin_config.dora_plugin) - - if is_enc_dec: - model.precompute_relative_attention_bias(build_config) - return model - - -def _init_max_seq_len(model_config, build_config): - """ - If max_seq_len is not specified, set it to max_position_embeddings * rotary_factor - Additional checks to ensure max_seq_len, max_input_len, and max_num_tokens have valid values. - """ - # Extract rotary scaling which will be used for checks and default value of max_seq_len - rotary_scaling = getattr(model_config, "rotary_scaling", None) - if rotary_scaling is not None: - rotary_type = rotary_scaling.get('type', - rotary_scaling.get('rope_type')) - rotary_factor = rotary_scaling.get( - 'factor', 1.0) if rotary_type not in ("su", "longrope", - "llama3") else 1 - else: - rotary_factor = 1 - - if model_config.architecture == "EncoderModel": - if build_config.max_seq_len is None: - build_config.max_seq_len = build_config.max_input_len - logger.info( - f'max_seq_len is not specified for EncoderModel, using --max_input_len.' - ) - assert build_config.max_input_len == build_config.max_seq_len, f"EncoderModel should have same --max_input_len ({build_config.max_input_len}) and --max_seq_len ({build_config.max_seq_len})." - - if build_config.max_seq_len is None: - # Step 1: Find the upper bound of max_seq_len - deduced_max_seq_len = 2048 - if model_config.max_position_embeddings is not None: - deduced_max_seq_len = model_config.max_position_embeddings - - # Step 2: Scale max_seq_len with rotary scaling - if rotary_factor != 1: - deduced_max_seq_len = math.ceil(deduced_max_seq_len * rotary_factor) - logger.warning( - f'max_seq_len is scaled to {deduced_max_seq_len} by rotary scaling {rotary_factor}' - ) - - # Step 3: Assign the new max_seq_len - build_config.max_seq_len = int(deduced_max_seq_len) - logger.info( - f'max_seq_len is not specified, using deduced value {deduced_max_seq_len}' - ) - else: - if not build_config.plugin_config.streamingllm and model_config.max_position_embeddings is not None \ - and model_config.position_embedding_type != PositionEmbeddingType.relative: - if build_config.max_seq_len > model_config.max_position_embeddings * rotary_factor: - logger.warning( - f'max_seq_len {build_config.max_seq_len} is larger than max_position_embeddings {model_config.max_position_embeddings} * rotary scaling {rotary_factor}, ' - 'the model accuracy might be affected') - - if build_config.max_input_len > build_config.max_seq_len: - logger.warning( - f'max_input_len is {build_config.max_input_len} is larger than max_seq_len {build_config.max_seq_len}, clipping it to max_seq_len' - ) - build_config.max_input_len = build_config.max_seq_len - - # Check and may modify max_num_tokens and opt_num_tokens (need to happen after max_seq_len is deduced) - max_num_tokens, opt_num_tokens = check_max_num_tokens( - max_num_tokens=build_config.max_num_tokens, - opt_num_tokens=build_config.opt_num_tokens, - max_batch_size=build_config.max_batch_size, - max_input_len=build_config.max_input_len, - max_seq_len=build_config.max_seq_len, - max_beam_width=build_config.max_beam_width, - remove_input_padding=build_config.plugin_config.remove_input_padding, - enable_context_fmha=build_config.plugin_config.context_fmha, - tokens_per_block=build_config.plugin_config.tokens_per_block, - multiple_profiles=build_config.plugin_config.multiple_profiles, - ) - build_config.max_num_tokens, build_config.opt_num_tokens = max_num_tokens, opt_num_tokens - - if build_config.plugin_config.remove_input_padding and build_config.plugin_config.context_fmha: - if build_config.max_input_len: - logger.warning( - 'padding removal and fMHA are both enabled, max_input_len is not required and will be ignored' - ) - else: - assert build_config.max_input_len is not None, 'padding removal and fMHA aren\'t both enabled, max_input_len is required' - if build_config.max_seq_len: - assert build_config.max_input_len <= build_config.max_seq_len, 'max_input_len should not be larger than max_seq_len' - - -def serialize_managed_weights(managed_weights: dict[str, np.ndarray], - path: str | Path, - metadata=None) -> None: - header = {} - if metadata is not None: - header["__metadata__"] = metadata - begin = 0 - for name, value in managed_weights.items(): - size = value.size * value.itemsize - if value.dtype == np.float32: - dtype = "F32" - elif value.dtype == np.float16: - dtype = "F16" - elif value.dtype == np_bfloat16: - dtype = "BF16" - elif value.dtype == np_float8: - dtype = "F8_E4M3" - elif value.dtype == np.int64: - dtype = "I64" - elif value.dtype == np.int32: - dtype = "I32" - elif value.dtype == np.int8: - dtype = "I8" - else: - raise RuntimeError(f"Unsupported dtype: {value.dtype}") - header[name] = { - "dtype": dtype, - "shape": value.shape, - "data_offsets": [begin, begin + size], - } - begin += size - - header_json = json.dumps(header) - header_json_len = len(header_json) - with open(path, "wb") as f: - logger.info( - f"Serializing {len(managed_weights)} managed weights to {path}...") - f.write(header_json_len.to_bytes(8, byteorder="little")) - f.write(header_json.encode()) - for name, value in managed_weights.items(): - logger.debug(f"Serializing managed weight: {name}") - buf = value.data - f.write(buf) - - -def deserialize_managed_weights(path: str | Path) -> dict[str, np.ndarray]: - with open(path, "rb") as f: - header_json_len = int.from_bytes(f.read(8), byteorder="little") - header_json = f.read(header_json_len).decode() - header = json.loads(header_json) - - managed_weights = {} - for name, info in header.items(): - dtype = info["dtype"] - shape = info["shape"] - data_offsets = info["data_offsets"] - if dtype == "F32": - dtype = np.float32 - elif dtype == "F16": - dtype = np.float16 - elif dtype == "BF16": - dtype = np_bfloat16 - elif dtype == "F8_E4M3": - dtype = np_float8 - elif dtype == "I64": - dtype = np.int64 - elif dtype == "I32": - dtype = np.int32 - else: - raise RuntimeError(f"Unsupported dtype: {dtype}") - - f.seek(data_offsets[0] + header_json_len + 8) - buf = f.read(data_offsets[1] - data_offsets[0]) - value = np.frombuffer(buf, dtype=dtype).reshape(shape) - managed_weights[name] = value - - return managed_weights - - -def build(model: PretrainedModel, build_config: BuildConfig) -> Engine: - '''Build engine from given model and optimization options specified in the build_config - WARNING: this function may change the given model object state in some optimization passes - to avoid cloning a model since normally the LLM models consumes large memory. - Create a new fresh model object if you need to build with different options. - ''' - emit_engine_arch_deprecation("builder.build()") - tic = time.time() - # avoid changing the input config - build_config = build_config.model_copy(deep=True) - build_config.plugin_config.dtype = model.config.dtype - build_config.update_kv_cache_type(model.config.architecture) - - _init_max_seq_len(model.config, build_config) - - if build_config.plugin_config.streamingllm: - build_config.plugin_config.use_paged_context_fmha = False - logger.warning( - "Paged Context FMHA is disabled because StreamingLLM is not supported when enabling paged KV context FMHA." - ) - if build_config.plugin_config.reduce_fusion and ( - model.config.mapping.tp_size == 1 or - (model.config.architecture != "LlamaForCausalLM" - and model.config.architecture != "Gemma2ForCausalLM" - and model.config.architecture != "MedusaForCausalLM")): - logger.warning('Overriding reduce_fusion to False') - build_config.plugin_config.reduce_fusion = False - if build_config.plugin_config.user_buffer and not build_config.plugin_config.reduce_fusion: - logger.warning('Overriding user_buffer to False') - build_config.plugin_config.user_buffer = False - if build_config.plugin_config.norm_quant_fusion and ( - build_config.plugin_config.reduce_fusion - or model.config.architecture != "LlamaForCausalLM" - or model.config.quantization.quant_algo != QuantAlgo.NVFP4): - logger.warning('Overriding norm_quant_fusion to False') - build_config.plugin_config.norm_quant_fusion = False - - if model.config.quantization.quant_algo == QuantAlgo.FP8 or \ - model.config.quantization.kv_cache_quant_algo == QuantAlgo.FP8: - build_config.strongly_typed = True - - if hasattr(model.config, 'max_draft_len'): - # If model.config has 'max_draft_len' but build_config not specified, - # use the value of model.config.max_draft_len to set the value of build_config.max_draft_len - if build_config.max_draft_len == 0: - build_config.max_draft_len = model.config.max_draft_len - - if hasattr(model.config, 'redrafter_num_beams') and hasattr( - model.config, 'redrafter_draft_len_per_beam'): - build_config.max_draft_len = model.config.redrafter_num_beams * model.config.redrafter_draft_len_per_beam - if build_config.speculative_decoding_mode != SpeculativeDecodingMode.EXPLICIT_DRAFT_TOKENS: - logger.warning( - 'speculative_decoding_mode is not EXPLICIT_DRAFT_TOKENS for ReDrafter model. Overwriting speculative_decoding_mode' - ) - build_config.speculative_decoding_mode = SpeculativeDecodingMode.EXPLICIT_DRAFT_TOKENS - - if build_config.speculative_decoding_mode != SpeculativeDecodingMode.NONE: - logger.info( - f'Increasing max_seq_len ({build_config.max_seq_len}) ' - f'by max_draft_len ({build_config.max_draft_len}) ' - 'to account for speculative decoding implementation specifics. ' - 'Maximum number of generated tokens remains the same. ' - f'New max_seq_len is set to {build_config.max_seq_len + build_config.max_draft_len}' - ) - build_config.max_seq_len += build_config.max_draft_len - - if build_config.speculative_decoding_mode == SpeculativeDecodingMode.EAGLE: - assert hasattr(model.config, 'num_eagle_layers') - num_eagle_layers = model.config.num_eagle_layers - logger.info( - f'Increasing max_seq_len ({build_config.max_seq_len}) ' - f'by num_eagle_layers ({num_eagle_layers}) ' - 'to account for EAGLE implementation specifics. ' - 'Maximum number of generated tokens remains the same. ' - f'New max_seq_len is set to {build_config.max_seq_len + num_eagle_layers}' - ) - build_config.max_seq_len += num_eagle_layers - - if build_config.speculative_decoding_mode != SpeculativeDecodingMode.NONE: - num_tokens = build_config.max_batch_size * (build_config.max_draft_len + - 1) - if build_config.max_num_tokens < num_tokens: - logger.info( - f'max_num_tokens ({build_config.max_num_tokens}) is smaller than ' - 'max_batch_size * (max_draft_len + 1) = ' - f'({build_config.max_batch_size} * ({build_config.max_draft_len} + 1)). ' - f'New max_num_tokens is set to {num_tokens}.') - build_config.max_num_tokens = num_tokens - - # Logics to control paged_context_fmha and fp8_context_fmha - if not build_config.plugin_config.context_fmha: - build_config.plugin_config.use_fp8_context_fmha = False - build_config.plugin_config.use_paged_context_fmha = False - logger.warning( - "Context FMHA is disabled, FP8 Context FMHA and Paged Context FMHA are disabled." - ) - elif model.config.quantization.quant_algo not in [ - QuantAlgo.FP8, QuantAlgo.W4A8_AWQ, QuantAlgo.NVFP4 - ]: - if build_config.plugin_config.use_fp8_context_fmha: - build_config.plugin_config.use_fp8_context_fmha = False - logger.warning( - "FP8 Context FMHA is disabled because it must be used together with the fp8 quantization workflow." - ) - if build_config.plugin_config.use_paged_context_fmha and model.config.quant_mode.has_fp8_kv_cache( - ): - build_config.plugin_config.use_paged_context_fmha = False - logger.warning( - "FP8 Paged Context FMHA is disabled because FP8 context FMHA is disabled." - ) - elif get_sm_version() < 89: - build_config.plugin_config.use_fp8_context_fmha = False - logger.warning( - "FP8 context FMHA is disabled because it is only supported on Ada and Hopper Arch." - ) - if build_config.plugin_config.use_paged_context_fmha and model.config.quant_mode.has_fp8_kv_cache( - ): - build_config.plugin_config.use_paged_context_fmha = False - logger.warning( - "FP8 Paged Context FMHA is disabled because FP8 context FMHA is disabled." - ) - elif build_config.plugin_config.use_paged_context_fmha: - if not model.config.quant_mode.has_fp8_kv_cache( - ) and build_config.plugin_config.use_fp8_context_fmha: - build_config.plugin_config.use_fp8_context_fmha = False - logger.warning( - "FP8 Paged Context FMHA is disabled because it must be used together with fp8 KV Cache." - ) - elif model.config.quant_mode.has_fp8_kv_cache( - ) and not build_config.plugin_config.use_fp8_context_fmha: - build_config.plugin_config.use_fp8_context_fmha = True - logger.warning( - "FP8 Context FMHA is enabled to support FP8 Paged Context FMHA." - ) - - if build_config.plugin_config.use_paged_context_fmha and model.config.quant_mode.has_int8_kv_cache( - ): - build_config.plugin_config.use_paged_context_fmha = False - logger.warning( - "Paged Context FMHA is disabled because it doesn't work with int8 kv cache currently." - ) - - if get_sm_version() >= 100 and get_sm_version() < 120: - if model.config.quant_mode.is_int8_weight_only( - ) or model.config.quant_mode.is_int4_weight_only( - ) or model.config.quant_mode.has_int8_kv_cache(): - raise RuntimeError( - "INT8/INT4 quantization is not supported on SM>=100.") - if model.config.quant_mode.has_act_and_weight_quant(): - raise RuntimeError("SmoothQuant is not supported on SM>=100.") - if model.config.quant_mode.has_per_channel_scaling( - ) or model.config.quant_mode.has_per_token_dynamic_scaling(): - raise RuntimeError( - "Per-channel or per-token scaling is not supported on SM>=100.") - - model = optimize_model_with_config(model, build_config) - - builder = Builder() - builder_config = builder.create_builder_config( - precision=model.config.dtype, - use_refit=build_config.use_refit, - timing_cache=build_config.input_timing_cache, - int8=(model.config.quant_mode.has_act_or_weight_quant() - and not model.config.quant_mode.has_per_group_scaling()) - or model.config.quant_mode.has_int8_kv_cache(), - strongly_typed=build_config.strongly_typed, - force_num_profiles=build_config.force_num_profiles, - profiling_verbosity=build_config.profiling_verbosity, - quant_mode=model.config.quant_mode, - use_strip_plan=build_config.use_strip_plan, - weight_sparsity=build_config.weight_sparsity, - weight_streaming=build_config.weight_streaming, - monitor_memory=build_config.monitor_memory, - ) - - network = builder.create_network() - network.plugin_config = build_config.plugin_config - - use_weight_only = model.config.quant_mode.is_weight_only() - per_group = model.config.quant_mode.has_per_group_scaling() - use_smooth_quant = model.config.quant_mode.has_act_and_weight_quant() - use_qserve = model.config.quant_mode.is_qserve_w4a8() - use_fp8_rowwise = model.config.quant_mode.has_fp8_rowwise() - disable_weight_only_quant_plugin = model.config.disable_weight_only_quant_plugin if hasattr( - model.config, 'disable_weight_only_quant_plugin') else False - use_fp8_rowwise = model.config.quant_mode.has_fp8_rowwise() - use_fp4_gemm = model.config.quant_mode.has_nvfp4() - if use_fp4_gemm and network.plugin_config._explicitly_disable_gemm_plugin is False: - logger.info( - 'NVFP4 quantization detected, by default enabling NVFP4 GEMM plugin. To use OOTB GEMM, please explicitly set gemm_plugin to "disable"' - ) - network.plugin_config.gemm_plugin = "nvfp4" - - if build_config.plugin_config.manage_weights: - if use_weight_only and disable_weight_only_quant_plugin: - raise RuntimeError( - "Manage weights of weight only quant works only with plugin currently." - ) - - if use_weight_only and not disable_weight_only_quant_plugin: - if per_group: - network.plugin_config.weight_only_groupwise_quant_matmul_plugin = model.config.dtype - else: - network.plugin_config.weight_only_quant_matmul_plugin = model.config.dtype - if use_smooth_quant and model.config.quantization._use_plugin_sq and build_config.plugin_config.smooth_quant_plugins: - network.plugin_config.set_smooth_quant_plugins(model.config.dtype) - if use_qserve: - network.plugin_config.set_qserve_plugins(model.config.dtype) - if use_fp8_rowwise: - network.plugin_config.set_fp8_rowwise_quant_plugins(model.config.dtype) - nccl_plugin = model.config.dtype if model.config.mapping.world_size > 1 else None - network.plugin_config.set_nccl_plugin(nccl_plugin) - - with net_guard(network): - # Prepare - network.set_named_parameters(model.named_parameters()) - - # Forward - prepare_input_args = { - "max_batch_size": - build_config.max_batch_size, - "max_input_len": - build_config.max_input_len, - "max_seq_len": - build_config.max_seq_len, - "use_cache": - build_config.kv_cache_type != KVCacheType.DISABLED, - "max_beam_width": - build_config.max_beam_width, - "max_num_tokens": - build_config.max_num_tokens, - "opt_num_tokens": - build_config.opt_num_tokens, - "prompt_embedding_table_size": - build_config.max_prompt_embedding_table_size, - "max_draft_len": - build_config.max_draft_len, - "speculative_decoding_draft_tokens_external": - build_config.speculative_decoding_mode == - SpeculativeDecodingMode.DRAFT_TOKENS_EXTERNAL, - "gather_context_logits": - build_config.gather_context_logits, - "lora_target_modules": - build_config.lora_config.lora_target_modules - } - - if model.config.architecture == "DecoderModel" or "mllama" in model.config.architecture.lower( - ): - prepare_input_args["max_seq_len"] = build_config.max_seq_len - prepare_input_args[ - "max_decoder_input_len"] = build_config.max_input_len - prepare_input_args[ - "max_encoder_input_len"] = build_config.max_encoder_input_len - - if model.config.architecture == "WhisperEncoder": - - prepare_input_args = { - "max_batch_size": build_config.max_batch_size, - } - - if build_config.speculative_decoding_mode == SpeculativeDecodingMode.EAGLE: - prepare_input_args[ - "spec_decoding_is_generation_length_variable"] = True - assert build_config.max_batch_size <= 512, "Max batch size > 512 is not supported for EAGLE" - assert build_config.max_draft_len <= 256, "Max draft len > 256 is not supported for EAGLE" - - if build_config.speculative_decoding_mode == SpeculativeDecodingMode.LOOKAHEAD_DECODING: - prepare_input_args[ - "spec_decoding_is_generation_length_variable"] = True - if model.config.architecture == "Qwen2VLForConditionalGeneration" or model.config.architecture == "Qwen2VLModel": - prepare_input_args[ - 'mrope_rotary_cos_sin_size'] = model.config.max_position_embeddings * model.config.rotary_embedding_dim - if build_config.speculative_decoding_mode == SpeculativeDecodingMode.EAGLE and not build_config.plugin_config.use_paged_context_fmha: - logger.warning( - "Paged Context FMHA is required for EAGLE. Turning it on") - build_config.plugin_config.use_paged_context_fmha = True - - inputs = model.prepare_inputs(**prepare_input_args) - model(**inputs) - - if build_config.enable_debug_output: - for k, v in model.named_network_outputs(): - network._mark_output(v, k, str_dtype_to_trt(model.config.dtype)) - - if model.config.architecture != "DecoderModel": - optimize(network) - - if build_config.visualize_network is not None: - with net_guard(network): - network.to_onnx(build_config.visualize_network) - - # Network -> Engine - logger.info( - f"Total time of constructing network from module object {time.time()-tic} seconds" - ) - managed_weights = {} if network.plugin_config.manage_weights else None - engine = None if build_config.dry_run else builder.build_engine( - network, builder_config, managed_weights) - engine_config = EngineConfig(model.config, build_config, __version__) - - if build_config.output_timing_cache is not None and model.config.mapping.rank == 0: - ok = builder.save_timing_cache(builder_config, - build_config.output_timing_cache) - assert ok, "Failed to save timing cache." - - import psutil - - # Get the current process - current_process = psutil.Process() - # Get resource usage for the current process (self) - rusage_s = current_process.memory_info() - # Get resource usage for all child processes - children = current_process.children(recursive=True) - rusage_c = [child.memory_info() for child in children] - logger.info( - f"Build phase peak memory: {rusage_s.rss / 1024 / 1024:.2f} MB, children: {sum([ru.rss for ru in rusage_c]) / 1024 / 1024:.2f} MB" - ) - - return Engine(engine_config, engine, managed_weights) diff --git a/tensorrt_llm/commands/bench.py b/tensorrt_llm/commands/bench.py index 269d4d1d3f2e..63873b5701a1 100644 --- a/tensorrt_llm/commands/bench.py +++ b/tensorrt_llm/commands/bench.py @@ -6,7 +6,6 @@ from tensorrt_llm.bench.benchmark.low_latency import latency_command from tensorrt_llm.bench.benchmark.throughput import throughput_command from tensorrt_llm.bench.benchmark.visual_gen import visual_gen_command -from tensorrt_llm.bench.build.build import build_command from tensorrt_llm.bench.dataclasses.general import BenchmarkEnvironment from tensorrt_llm.bench.dataset.prepare_dataset import prepare_dataset from tensorrt_llm.logger import logger, severity_map @@ -88,7 +87,6 @@ def main( ctx.obj.workspace.mkdir(parents=True, exist_ok=True) -main.add_command(build_command) main.add_command(throughput_command) main.add_command(latency_command) main.add_command(prepare_dataset) diff --git a/tensorrt_llm/commands/build.py b/tensorrt_llm/commands/build.py deleted file mode 100644 index 0f68382607bb..000000000000 --- a/tensorrt_llm/commands/build.py +++ /dev/null @@ -1,553 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -import argparse -import copy -import os -import time -import traceback -from concurrent.futures import ProcessPoolExecutor, as_completed -from importlib.machinery import SourceFileLoader -from multiprocessing import get_context -from typing import Optional, Union - -import torch - -from tensorrt_llm._deprecation import emit_engine_arch_deprecation -from tensorrt_llm._utils import (local_mpi_rank, local_mpi_size, mpi_barrier, - mpi_comm, mpi_rank, mpi_world_size) -from tensorrt_llm.builder import BuildConfig, Engine, build -from tensorrt_llm.llmapi.kv_cache_type import KVCacheType -from tensorrt_llm.logger import logger, severity_map -from tensorrt_llm.lora_helper import LoraConfig -from tensorrt_llm.lora_manager import LoraManager -from tensorrt_llm.models import MODEL_MAP, PretrainedConfig -from tensorrt_llm.models.modeling_utils import SpeculativeDecodingMode -from tensorrt_llm.plugin import PluginConfig, add_plugin_argument -from tensorrt_llm.quantization.mode import QuantAlgo - - -def parse_arguments(): - parser = argparse.ArgumentParser( - formatter_class=argparse.ArgumentDefaultsHelpFormatter) - parser.add_argument( - '--checkpoint_dir', - type=str, - default=None, - help="The directory path that contains TensorRT LLM checkpoint.") - parser.add_argument( - '--model_config', - type=str, - default=None, - help="The file path that saves TensorRT LLM checkpoint config.") - parser.add_argument( - '--build_config', - type=str, - default=None, - help="The file path that saves TensorRT LLM build config.") - parser.add_argument( - '--model_cls_file', - type=str, - default=None, - help="The file path that defines customized TensorRT LLM model.") - parser.add_argument('--model_cls_name', - type=str, - default=None, - help="The customized TensorRT LLM model class name.") - parser.add_argument( - '--output_dir', - type=str, - default='engine_outputs', - help= - "The directory path to save the serialized engine files and engine config file." - ) - - parser.add_argument( - '--max_batch_size', - type=int, - default=BuildConfig.model_fields["max_batch_size"].default, - help="Maximum number of requests that the engine can schedule.") - parser.add_argument( - '--max_input_len', - type=int, - default=BuildConfig.model_fields["max_input_len"].default, - help="Maximum input length of one request.") - parser.add_argument( - '--max_seq_len', - '--max_decoder_seq_len', - dest='max_seq_len', - type=int, - default=BuildConfig.model_fields["max_seq_len"].default, - help="Maximum total length of one request, including prompt and outputs. " - "If unspecified, the value is deduced from the model config.") - parser.add_argument( - '--max_beam_width', - type=int, - default=BuildConfig.model_fields["max_beam_width"].default, - help="Maximum number of beams for beam search decoding.") - parser.add_argument( - '--max_num_tokens', - type=int, - default=BuildConfig.model_fields["max_num_tokens"].default, - help= - "Maximum number of batched input tokens after padding is removed in each batch. " - "Currently, the input padding is removed by default; " - "you may explicitly disable it by specifying ``--remove_input_padding disable``." - ) - parser.add_argument( - '--opt_num_tokens', - type=int, - default=BuildConfig.model_fields["opt_num_tokens"].default, - help= - "Optimal number of batched input tokens after padding is removed in each batch " - "It equals to ``max_batch_size * max_beam_width`` by default, set this " - "value as close as possible to the actual number of tokens on your workload. " - "Note that this argument might be removed in the future.") - parser.add_argument( - '--max_encoder_input_len', - type=int, - default=BuildConfig.model_fields["max_encoder_input_len"].default, - help="Maximum encoder input length for enc-dec models. " - "Set ``max_input_len`` to 1 to start generation from decoder_start_token_id of length 1." - ) - parser.add_argument( - '--max_prompt_embedding_table_size', - '--max_multimodal_len', - type=int, - default=BuildConfig.model_fields["max_prompt_embedding_table_size"]. - default, - help= - "Maximum prompt embedding table size for prompt tuning, or maximum multimodal input size for multimodal models. " - "Setting a value > 0 enables prompt tuning or multimodal input.") - parser.add_argument( - '--kv_cache_type', - default=argparse.SUPPRESS, - type=KVCacheType, - help= - "Set KV cache type (continuous, paged, or disabled). For disabled case, KV cache is disabled and only context phase is allowed." - ) - parser.add_argument( - '--paged_kv_cache', - type=str, - default=argparse.SUPPRESS, - help= - "Deprecated. Enabling this option is equivalent to ``--kv_cache_type paged`` for transformer based models." - ) - - parser.add_argument( - '--input_timing_cache', - type=str, - default=BuildConfig.model_fields["input_timing_cache"].default, - help= - "The file path to read the timing cache. This option is ignored if the file does not exist." - ) - parser.add_argument( - '--output_timing_cache', - type=str, - default=BuildConfig.model_fields["output_timing_cache"].default, - help="The file path to write the timing cache.") - parser.add_argument( - '--profiling_verbosity', - type=str, - default=BuildConfig.model_fields["profiling_verbosity"].default, - choices=['layer_names_only', 'detailed', 'none'], - help= - "The profiling verbosity for the generated TensorRT engine. Setting to detailed allows inspecting tactic choices and kernel parameters." - ) - parser.add_argument( - '--strip_plan', - default=BuildConfig.model_fields["use_strip_plan"].default, - action='store_true', - help= - "Enable stripping weights from the final TensorRT engine under the assumption that the refit weights are identical to those provided at build time." - ) - parser.add_argument( - '--weight_sparsity', - default=BuildConfig.model_fields["weight_sparsity"].default, - action='store_true', - help="Enable weight sparsity.") - parser.add_argument( - '--weight_streaming', - default=BuildConfig.model_fields["weight_streaming"].default, - action='store_true', - help= - "Enable offloading weights to CPU and streaming loading at runtime.", - ) - parser.add_argument( - '--fast_build', - default=False, - action='store_true', - help= - "Enable features for faster engine building. This may cause some performance degradation and is currently incompatible with int8/int4 quantization without plugin.", - ) - - parser.add_argument('--workers', - type=int, - default=1, - help="The number of workers for building in parallel.") - parser.add_argument('--log_level', - type=str, - default='info', - choices=severity_map.keys(), - help="The logging level.") - parser.add_argument( - '--enable_debug_output', - default=BuildConfig.model_fields["enable_debug_output"].default, - action='store_true', - help="Enable debug output.") - parser.add_argument( - '--visualize_network', - type=str, - default=None, - help= - "The directory path to export TensorRT Network as ONNX prior to Engine build for debugging." - ) - parser.add_argument( - '--dry_run', - default=BuildConfig.model_fields["dry_run"].default, - action='store_true', - help= - "Run through the build process except the actual Engine build for debugging." - ) - parser.add_argument('--monitor_memory', - default=False, - action='store_true', - help="Enable memory monitor during Engine build.") - - logits_parser = parser.add_argument_group("Logits arguments") - logits_parser.add_argument('--logits_dtype', - type=str, - default=None, - choices=['float16', 'float32'], - help="The data type of logits.") - logits_parser.add_argument('--gather_context_logits', - action='store_true', - default=False, - help="Enable gathering context logits.") - logits_parser.add_argument('--gather_generation_logits', - action='store_true', - default=False, - help="Enable gathering generation logits.") - logits_parser.add_argument( - '--gather_all_token_logits', - action='store_true', - default=False, - help= - "Enable both ``gather_context_logits`` and ``gather_generation_logits``." - ) - - lora_parser = parser.add_argument_group("LoRA arguments") - lora_parser.add_argument( - '--lora_dir', - type=str, - default=None, - nargs="+", - help="The directory of LoRA weights. " - "If multiple directories are provided, the first one is used for configuration." - ) - lora_parser.add_argument('--lora_ckpt_source', - type=str, - default="hf", - choices=["hf", "nemo"], - help="The source type of LoRA checkpoint.") - lora_parser.add_argument( - '--lora_target_modules', - nargs='+', - default=None, - choices=LoraManager.LORA_MODULE_IDS.keys(), - help= - "The target module names that LoRA is applied. Only effective when ``lora_plugin`` is enabled." - ) - lora_parser.add_argument( - '--max_lora_rank', - type=int, - default=64, - help="Maximum LoRA rank for different LoRA modules. " - "It is used to compute the workspace size of LoRA plugin.") - - spec_parser = parser.add_argument_group("Speculative decoding arguments") - spec_parser.add_argument('--speculative_decoding_mode', - default=None, - choices=[ - "draft_tokens_external", "lookahead_decoding", - "medusa", "explicit_draft_tokens", "eagle" - ], - help="Mode of speculative decoding.") - spec_parser.add_argument( - '--max_draft_len', - type=int, - default=0, - help= - "Maximum lengths of draft tokens for speculative decoding target model." - ) - - plugin_config_parser = parser.add_argument_group("Plugin config arguments") - add_plugin_argument(plugin_config_parser) - return parser - - -def build_model( - build_config: BuildConfig, - rank: int = 0, - ckpt_dir: str = None, - model_config: Union[str, PretrainedConfig] = None, - model_cls=None, - dry_run: - bool = False, # return the modified BuildConfig without actually building the engine - **kwargs -) -> Union[Engine, BuildConfig]: - model_config = copy.deepcopy(model_config) - - logits_dtype = kwargs.get('logits_dtype') - if logits_dtype is not None: - model_config.logits_dtype = logits_dtype - - architecture = model_config.architecture - assert not build_config.plugin_config.streamingllm, \ - "StreamingLLM is no longer supported because attention sink cannot work with the non-cyclic kv cache kernel & runtime changes." - assert not build_config.plugin_config.pp_reduce_scatter or architecture == "MixtralForCausalLM", \ - "PP reduce scatter is only supported in the mixtral model." - - assert rank < model_config.mapping.world_size - - rank_config = copy.deepcopy(model_config) - rank_config.set_rank(rank) - - if model_cls is None: - assert architecture in MODEL_MAP, \ - f"Unsupported model architecture: {architecture}" - model_cls = MODEL_MAP[architecture] - if ckpt_dir is None: - model = model_cls(rank_config) - else: - model = model_cls.from_checkpoint(ckpt_dir, config=rank_config) - is_checkpoint_pruned = getattr(rank_config, 'is_pruned', False) - - if build_config.plugin_config.lora_plugin is not None: - lora_config = LoraConfig(lora_dir=kwargs['lora_dir'] or [], - lora_ckpt_source=kwargs['lora_ckpt_source'], - max_lora_rank=kwargs['max_lora_rank']) - if kwargs['lora_target_modules'] is not None: - # command line options is preferred over the modules in the lora dir - lora_config.lora_target_modules = kwargs['lora_target_modules'] - build_config.lora_config = lora_config - - if is_checkpoint_pruned or kwargs.pop('strip_plan', False): - build_config.use_strip_plan = True - build_config.use_refit = kwargs.get('refit', False) - - return build(model, build_config) - - -def build_and_save(rank, gpu_id, ckpt_dir, build_config, output_dir, log_level, - model_config, model_cls, **kwargs): - torch.cuda.set_device(gpu_id) - logger.set_level(log_level) - engine = build_model(build_config, - rank, - ckpt_dir, - model_config, - model_cls=model_cls, - **kwargs) - assert engine is not None - engine.save(output_dir) - return True - - -def parallel_build(model_config: PretrainedConfig, - ckpt_dir: Optional[str], - build_config: BuildConfig, - output_dir: str, - workers: int = 1, - log_level: str = 'info', - model_cls=None, - **kwargs): - - world_size = model_config.mapping.world_size - use_mpi = mpi_world_size() > 1 - - if not use_mpi and workers == 1: - for rank in range(world_size): - passed = build_and_save(rank, rank % workers, ckpt_dir, - build_config, output_dir, log_level, - model_config, model_cls, **kwargs) - assert passed, "Engine building failed, please check error log." - elif not use_mpi: - with ProcessPoolExecutor(mp_context=get_context('spawn'), - max_workers=workers) as p: - futures = [ - p.submit(build_and_save, rank, rank % workers, ckpt_dir, - build_config, output_dir, log_level, model_config, - model_cls, **kwargs) for rank in range(world_size) - ] - exceptions = [] - for future in as_completed(futures): - try: - future.result() - except Exception as e: - traceback.print_exc() - exceptions.append(e) - assert len(exceptions - ) == 0, "Engine building failed, please check error log." - else: - mpi_local_rank = local_mpi_rank() - node_gpu_count = local_mpi_size() - exceptions = [] - for engine_rank in range(world_size): - if engine_rank % mpi_world_size() != mpi_rank(): - continue - try: - build_and_save(engine_rank, mpi_local_rank % node_gpu_count, - ckpt_dir, build_config, output_dir, log_level, - model_config, model_cls, **kwargs) - except Exception as e: - traceback.print_exc() - exceptions.append(e) - mpi_barrier() - if len(exceptions) != 0: - print("Engine building failed, please check error log.", flush=True) - mpi_comm().Abort() - - -def main(): - emit_engine_arch_deprecation("trtllm-build") - - parser = parse_arguments() - args = parser.parse_args() - - if hasattr(args, 'gather_generation_logits'): - logger.warning( - 'Option --gather_generation_logits is deprecated, a build flag is not required anymore. Use --output_generation_logits at runtime instead.' - ) - - if args.gather_all_token_logits: - args.gather_context_logits = True - args.gather_generation_logits = True - if args.gather_context_logits and args.max_draft_len > 0: - raise RuntimeError( - "Gather context logits is not support with draft len > 0. " - "If want to get the accepted tokens' logits from target model, please just enable gather_generation_logits" - ) - - logger.set_level(args.log_level) - tik = time.time() - - if not os.path.exists(args.output_dir): - os.makedirs(args.output_dir, exist_ok=True) - - model_cls = None - if args.model_cls_file is not None: - assert args.model_cls_name is not None - loader = SourceFileLoader('models', args.model_cls_file) - mod = loader.load_module() - model_cls = getattr(mod, args.model_cls_name) - - workers = min(torch.cuda.device_count(), args.workers) - - if hasattr(args, 'paged_kv_cache'): - logger.warning( - 'Option --paged_kv_cache is deprecated, use --kv_cache_type=paged/disabled instead.' - ) - - plugin_config = PluginConfig.from_arguments(args) - if args.fast_build: - plugin_config.manage_weights = True - - kwargs = { - 'logits_dtype': args.logits_dtype, - 'use_fused_mlp': args.use_fused_mlp, - 'lora_dir': args.lora_dir, - 'lora_ckpt_source': args.lora_ckpt_source, - 'max_lora_rank': args.max_lora_rank, - 'lora_target_modules': args.lora_target_modules, - 'strip_plan': args.strip_plan, - 'refit': False, - } - speculative_decoding_mode = SpeculativeDecodingMode.from_arguments(args) - - ckpt_dir_or_model_config = args.checkpoint_dir if args.checkpoint_dir is not None else args.model_config - if ckpt_dir_or_model_config.lower().endswith('.json'): - config_path = ckpt_dir_or_model_config - ckpt_dir = None - else: - config_path = os.path.join(ckpt_dir_or_model_config, 'config.json') - ckpt_dir = ckpt_dir_or_model_config - model_config = PretrainedConfig.from_json_file(config_path) - - # avoid ValueError if not supported quantization is chosen with use_fused_mlp - quant_algo = model_config.quantization.quant_algo - if quant_algo and quant_algo not in (QuantAlgo.FP8, - QuantAlgo.MIXED_PRECISION): - kwargs['use_fused_mlp'] = False - - if args.build_config is None: - if args.multiple_profiles == "enable" and args.opt_num_tokens is not None: - raise RuntimeError( - "multiple_profiles is enabled, while opt_num_tokens is set. " - "They are not supposed to be working in the same time for now.") - - # This should only be used for debugging. - # The env var BUILDER_FORCE_NUM_PROFILES should override the number of - # optimization profiles during TRT build. - # BUILDER_FORCE_NUM_PROFILES must be less than or equal to the number of - # optimization profiles set by model's prepare_inputs(). - force_num_profiles_from_env = os.environ.get( - "BUILDER_FORCE_NUM_PROFILES", None) - if force_num_profiles_from_env is not None: - logger.warning( - f"Overriding # of builder profiles <= {force_num_profiles_from_env}." - ) - - build_config = BuildConfig( - max_input_len=args.max_input_len, - max_seq_len=args.max_seq_len, - max_batch_size=args.max_batch_size, - max_beam_width=args.max_beam_width, - max_num_tokens=args.max_num_tokens, - opt_num_tokens=args.opt_num_tokens, - max_prompt_embedding_table_size=args. - max_prompt_embedding_table_size, - kv_cache_type=getattr(args, "kv_cache_type", None), - gather_context_logits=args.gather_context_logits, - gather_generation_logits=args.gather_generation_logits, - strongly_typed=True, - force_num_profiles=force_num_profiles_from_env, - weight_sparsity=args.weight_sparsity, - profiling_verbosity=args.profiling_verbosity, - enable_debug_output=args.enable_debug_output, - max_draft_len=args.max_draft_len, - speculative_decoding_mode=speculative_decoding_mode, - input_timing_cache=args.input_timing_cache, - output_timing_cache=args.output_timing_cache, - dry_run=args.dry_run, - visualize_network=args.visualize_network, - max_encoder_input_len=args.max_encoder_input_len, - weight_streaming=args.weight_streaming, - monitor_memory=args.monitor_memory, - use_mrope=getattr(model_config, "qwen_type", None) == "qwen2_vl", - plugin_config=plugin_config) - else: - build_config = BuildConfig.from_json_file(args.build_config) - build_config.plugin_config = plugin_config - - parallel_build(model_config, ckpt_dir, build_config, args.output_dir, - workers, args.log_level, model_cls, **kwargs) - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - logger.info(f'Total time of building all engines: {t}') - - -if __name__ == '__main__': - main() diff --git a/tensorrt_llm/commands/eval.py b/tensorrt_llm/commands/eval.py index 2f0e4b9cd6b0..0e25bb80736f 100644 --- a/tensorrt_llm/commands/eval.py +++ b/tensorrt_llm/commands/eval.py @@ -19,18 +19,22 @@ import tensorrt_llm.profiler as profiler from .. import LLM as PyTorchLLM -from .._tensorrt_engine import LLM from ..evaluate import (AALCR, AIME2025, AIME2026, GSM8K, HLE, MMLU, MMMU, ArenaHard, CnnDailymail, CoVoST2, GPQADiamond, GPQAExtended, GPQAMain, GPQANemoSkills, IFBench, JsonModeEval, LongBenchV1, LongBenchV2, MMMUPro, SciCode) -from ..llmapi import BuildConfig, KvCacheConfig +from ..llmapi import KvCacheConfig +from ..llmapi.llm_args import TorchLlmArgs from ..llmapi.llm_utils import update_llm_args_with_extra_options from ..logger import logger, severity_map from ..usage import config as _telemetry_config from .utils import collect_explicit_cli_keys +# CLI defaults are sourced from the TorchLlmArgs field defaults so they stay in +# lock-step with the args class and can't drift. +_LLM_ARGS_FIELDS = TorchLlmArgs.model_fields + # Map Click parameter names to the LlmArgs field name (or merge-function CLI # scalar name) used by `update_llm_args_with_extra_options`. _CLICK_TO_LLM_ARG = { @@ -64,7 +68,7 @@ ) @click.option( "--backend", - type=click.Choice(["pytorch", "tensorrt"]), + type=click.Choice(["pytorch"]), default="pytorch", help="The backend to use for evaluation. Default is pytorch backend.") @click.option('--log_level', @@ -73,23 +77,23 @@ help="The logging level.") @click.option("--max_beam_width", type=int, - default=BuildConfig.model_fields["max_beam_width"].default, + default=_LLM_ARGS_FIELDS["max_beam_width"].default, help="Maximum number of beams for beam search decoding.") @click.option("--max_batch_size", type=int, - default=BuildConfig.model_fields["max_batch_size"].default, + default=_LLM_ARGS_FIELDS["max_batch_size"].default, help="Maximum number of requests that the engine can schedule.") @click.option( "--max_num_tokens", type=int, - default=BuildConfig.model_fields["max_num_tokens"].default, + default=_LLM_ARGS_FIELDS["max_num_tokens"].default, help= "Maximum number of batched input tokens after padding is removed in each batch." ) @click.option( "--max_seq_len", type=int, - default=BuildConfig.model_fields["max_seq_len"].default, + default=None, help="Maximum total length of one request, including prompt and outputs. " "If unspecified, the value is deduced from the model config.") @click.option("--tp_size", type=int, default=1, help='Tensor parallelism size.') @@ -187,13 +191,6 @@ def main(ctx, model: str, tokenizer: Optional[str], max_num_tokens=max_num_tokens, max_beam_width=max_beam_width, max_seq_len=max_seq_len) - elif backend == 'tensorrt': - llm_cls = LLM - build_config = BuildConfig(max_batch_size=max_batch_size, - max_num_tokens=max_num_tokens, - max_beam_width=max_beam_width, - max_seq_len=max_seq_len) - llm_args.update(build_config=build_config) else: raise click.BadParameter( f"{backend} is not a known backend, check help for available options.", diff --git a/tensorrt_llm/commands/prune.py b/tensorrt_llm/commands/prune.py deleted file mode 100644 index 5e7994e9b4da..000000000000 --- a/tensorrt_llm/commands/prune.py +++ /dev/null @@ -1,122 +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. -''' -Script that prunes TRT-LLM checkpoints. -''' -import argparse -import json -import os -from pathlib import Path -from typing import Dict - -import safetensors -import torch -from safetensors.torch import save_file - -from tensorrt_llm.logger import logger -from tensorrt_llm.models import MODEL_MAP, PretrainedConfig - -SUPPORTED_MODELS = list(MODEL_MAP.keys()) -PRUNABLE_WEIGHTS = [ - 'attention.qkv.weight', - 'attention.proj.weight', - 'mlp.fc.weight', - 'mlp.proj.weight', - 'mlp.gate.weight', -] - - -def can_prune(key: str) -> bool: - for w in PRUNABLE_WEIGHTS: - if w in key: - return True - return False - - -def load_config(config_path: Path) -> Dict[str, any]: - if not config_path.exists(): - return {} - - with open(str(config_path), 'r') as f: - return json.load(f) - - -def prune_and_save(ckpt_dir: str, out_dir: str, prune_all: bool): - logger.info(f'Checkpoint Dir: {ckpt_dir}, Out Dir: {out_dir}') - model_config = PretrainedConfig.from_json_file( - os.path.join(ckpt_dir, 'config.json')) - - architecture = model_config.architecture - if architecture not in MODEL_MAP: - raise RuntimeError(f'Unsupported model architecture: {architecture}') - - if not os.path.exists(out_dir): - os.makedirs(out_dir) - - for rank in range(model_config.mapping.world_size): - pruned_weights = {} - with safetensors.safe_open(os.path.join(ckpt_dir, - f'rank{rank}.safetensors'), - framework='pt', - device='cpu') as f: - for key in f.keys(): - tensor = f.get_tensor(key) - if prune_all or can_prune(key): - pruned_weights[key] = torch.tensor([], dtype=tensor.dtype) - else: - pruned_weights[key] = tensor - - save_file(pruned_weights, - os.path.join(out_dir, f'rank{rank}.safetensors')) - - config_path = Path(ckpt_dir, 'config.json') - with open(str(Path(out_dir, 'config.json')), 'w') as f: - config = load_config(config_path) - config['is_pruned'] = True - json.dump(config, f) - - -def main(): - parser = argparse.ArgumentParser() - parser.add_argument('--checkpoint_dir', type=str, default=None) - parser.add_argument('--prune_all', - default=False, - action='store_true', - help='Remove all weights in the checkpoint') - parser.add_argument( - '--out_dir', - type=str, - default=None, - help= - 'Path to write pruned checkpoint. Defaults to the same directory append with `.pruned`' - ) - args = parser.parse_args() - - if args.checkpoint_dir is None: - raise RuntimeError( - "No `--checkpoint_dir` supplied to checkpoint pruner.") - - if args.out_dir is None: - ckpt_path = Path(args.checkpoint_dir) - ckpt_name = ckpt_path.name - args.out_dir = str( - Path(args.checkpoint_dir).with_name(ckpt_name + '.pruned')) - - prune_and_save(os.path.abspath(args.checkpoint_dir), - os.path.abspath(args.out_dir), args.prune_all) - - -if __name__ == '__main__': - main() diff --git a/tensorrt_llm/commands/refit.py b/tensorrt_llm/commands/refit.py deleted file mode 100644 index b77296758f61..000000000000 --- a/tensorrt_llm/commands/refit.py +++ /dev/null @@ -1,183 +0,0 @@ -''' -Script that refits TRT-LLM engine(s) with weights in a TRT-LLM checkpoint. -''' -import argparse -import json -import os -import re -import shutil -import time -from pathlib import Path - -import tensorrt as trt - -from tensorrt_llm._common import _is_building -from tensorrt_llm._deprecation import emit_engine_arch_deprecation -from tensorrt_llm._utils import np_dtype_to_trt -from tensorrt_llm.builder import EngineConfig, optimize_model_with_config -from tensorrt_llm.models import MODEL_MAP, PretrainedConfig - -from ..logger import logger - -ENGINE_RE = re.compile(r'rank(\d+).engine') - - -@_is_building -def refit_engine(engine_path: str, refit_engine_dir: str, checkpoint_dir: str, - engine_config: EngineConfig, fixed_weights_names: list): - # This function loops through all weights in the model and does a textual match between - # checkpoint weight names and engine weight names. - rank = int(ENGINE_RE.fullmatch(os.path.basename(engine_path)).group(1)) - tik = time.time() - with open(engine_path, - "rb") as f, trt.Runtime(logger.trt_logger) as runtime: - engine = runtime.deserialize_cuda_engine(f.read()) - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - logger.info(f'Load TRT engine time: {t}') - - refitter = trt.Refitter(engine, logger.trt_logger) - refittable_weights = set(refitter.get_all_weights()) - - # Load model. - tik = time.time() - rank_config = PretrainedConfig.from_dict( - engine_config.pretrained_config.to_dict()) - rank_config.set_rank(rank) - - architecture = rank_config.architecture - assert architecture in MODEL_MAP, \ - f"Unsupported model architecture: {architecture}" - model_cls = MODEL_MAP[architecture] - model = model_cls.from_checkpoint(checkpoint_dir, config=rank_config) - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - logger.info(f'Load checkpoint(s) time: {t}') - - # There are weights preprocess during optimize model. - tik = time.time() - build_config = engine_config.build_config.model_copy(deep=True) - optimize_model_with_config(model, build_config) - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - logger.info(f'Preprocess weights time: {t}') - - # Refit engine. - tik = time.time() - refitted_weights = [] - for name, buf in model.named_parameters(): - if name in refittable_weights: - assert buf.is_inited, f"Failed because weight `{name}` is not initialized in model." - weight = buf._value - weights_value = trt.Weights(np_dtype_to_trt(weight.dtype), - weight.ctypes.data, weight.size) - assert refitter.set_named_weights( - name, weights_value), f'Failed to refit weight: `{name}`' - refitted_weights.append(name) - else: - if name not in fixed_weights_names: - logger.warning( - f"model weights `{name}` (shape: {buf._value.shape}) is not refittable, this means that we might not be able to update the engine using fine-tuned checkpoint!" - ) - - # Validate all refittable weights are provided. - if len(refitted_weights) != len(refittable_weights): - raise RuntimeError( - f'Missing refittable weights {refittable_weights.difference(refitted_weights)} from {checkpoint_dir}' - ) - - assert refitter.refit_cuda_engine(), f'Failed to refit engine.' - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - logger.info(f'Execute GPU refit graph time: {t}') - - tik = time.time() - refit_engine_path = os.path.join(refit_engine_dir, - os.path.basename(engine_path)) - with open(refit_engine_path, 'wb') as f: - logger.info(f'\nWriting refitted engine to `{refit_engine_path}`') - s_config = engine.create_serialization_config() - s_config.flags &= ~(1 << int(trt.SerializationFlag.EXCLUDE_WEIGHTS)) - f.write(engine.serialize_with_config(s_config)) - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - logger.info(f'Write TRT engine to disk time: {t}') - - del refitter - - -def refit(engine_dir: str, checkpoint_dir: str, engine_config: EngineConfig, - output_dir: str, fixed_weights_names: list): - refit_engine_dir = output_dir - os.makedirs(refit_engine_dir, exist_ok=True) - shutil.copyfile(os.path.join(engine_dir, 'config.json'), - os.path.join(refit_engine_dir, 'config.json')) - engine_paths = list(Path(engine_dir).glob('*.engine')) - for path in engine_paths: - refit_engine(path, refit_engine_dir, checkpoint_dir, engine_config, - fixed_weights_names) - - -def main(): - emit_engine_arch_deprecation("trtllm-refit") - - parser = argparse.ArgumentParser() - parser.add_argument( - '--engine_dir', - type=str, - default=None, - help= - 'Path to trt-llm engines. These engines must have been built from a pruned checkpoint, or otherwise be refittable.' - ) - parser.add_argument('--checkpoint_dir', - type=str, - default=None, - help='Path to checkpoint containing desired weights') - parser.add_argument('--output_dir', - type=str, - default=None, - help="Output path of the refit model") - parser.add_argument('--log_level', type=str, default='info') - args = parser.parse_args() - - logger.set_level(args.log_level) - if args.engine_dir is None or not Path(args.engine_dir).exists(): - raise RuntimeError( - f'Please supply a valid --engine_dir (found `{args.engine_dir}`)') - if args.checkpoint_dir is None or not Path(args.checkpoint_dir).exists(): - raise RuntimeError( - f'Please supply a valid --checkpoint_dir (found `{args.checkpoint_dir}`)' - ) - - engine_config = EngineConfig.from_json_file( - os.path.join(args.engine_dir, 'config.json')) - - with open(os.path.join(args.checkpoint_dir, 'config.json'), 'r') as f: - checkpoint_config = json.load(f) - - engine_arch = engine_config.pretrained_config.architecture - checkpoint_arch = checkpoint_config['architecture'] - if engine_arch != checkpoint_arch: - raise RuntimeError( - f'Engine Architecture and Checkpoint Architecture do not match. ' + - f'Engine Architecture: `{engine_arch}`, Checkpoint Architecture: `{checkpoint_arch}`' - ) - - # The fixed weights are not read from checkpoint, they are hardcoded buffer from the model itself. These values remain constant across different fine-tuned checkpoints. - fixed_wts_in_model = [] - model_cls = MODEL_MAP[engine_arch] - model = model_cls.from_config(engine_config.pretrained_config) - for name, param in model.named_parameters(): - if param.is_inited(): - fixed_wts_in_model.append(name) - - refit(engine_dir=os.path.normpath(args.engine_dir), - checkpoint_dir=os.path.normpath(args.checkpoint_dir), - engine_config=engine_config, - output_dir=os.path.normpath(args.output_dir), - fixed_weights_names=fixed_wts_in_model) - - -if __name__ == '__main__': - main() diff --git a/tensorrt_llm/commands/serve.py b/tensorrt_llm/commands/serve.py index bc949b3db948..2f39353c521c 100644 --- a/tensorrt_llm/commands/serve.py +++ b/tensorrt_llm/commands/serve.py @@ -11,7 +11,7 @@ import sys import uuid from pathlib import Path -from typing import Any, Dict, Mapping, Optional, Sequence, Set +from typing import Any, Dict, Optional, Sequence, Set import click import torch @@ -22,24 +22,20 @@ from tensorrt_llm import LLM as PyTorchLLM from tensorrt_llm import MultimodalEncoder -from tensorrt_llm._tensorrt_engine import LLM from tensorrt_llm._utils import mpi_rank from tensorrt_llm.commands._serve_stability import stability_option from tensorrt_llm.commands.utils import (collect_explicit_cli_keys, get_is_diffusion_only_model) from tensorrt_llm.executor.utils import LlmLauncherEnvs from tensorrt_llm.inputs.multimodal import MultimodalServerConfig -from tensorrt_llm.llmapi import (BuildConfig, CapacitySchedulerPolicy, - DynamicBatchConfig, KvCacheConfig, - SchedulerConfig) +from tensorrt_llm.llmapi import KvCacheConfig from tensorrt_llm.llmapi.disagg_utils import (DisaggClusterConfig, MetadataServerConfig, ServerRole, extract_disagg_cluster_config, parse_disagg_config_file, parse_metadata_server_config_file, validate_config_bool) -from tensorrt_llm.llmapi.llm_args import (MultimodalConfig, TorchLlmArgs, - TrtLlmArgs) +from tensorrt_llm.llmapi.llm_args import MultimodalConfig, TorchLlmArgs from tensorrt_llm.llmapi.llm_utils import update_llm_args_with_extra_dict from tensorrt_llm.llmapi.mpi_session import find_free_ipc_addr from tensorrt_llm.llmapi.reasoning_parser import (ReasoningParserFactory, @@ -157,9 +153,7 @@ def is_non_default_or_required(param_name, value, backend, explicit_cli_keys): for s in cli_derived_fields.get(param_name, ())): return True - if backend == "tensorrt": - llm_args_class = TrtLlmArgs - elif backend == "_autodeploy": + if backend == "_autodeploy": from tensorrt_llm._torch.auto_deploy.llm_args import \ LlmArgs as AutoDeployLlmArgs llm_args_class = AutoDeployLlmArgs @@ -179,19 +173,21 @@ def is_non_default_or_required(param_name, value, backend, explicit_cli_keys): return value != default +# CLI/API defaults are sourced from the TorchLlmArgs field defaults so they stay +# in lock-step with the args class and can't drift. +_LLM_ARGS_FIELDS = TorchLlmArgs.model_fields + + def get_llm_args( model: str, tokenizer: Optional[str] = None, custom_tokenizer: Optional[str] = None, post_processor_hook: Optional[str] = None, backend: str = "pytorch", - max_beam_width: int = BuildConfig.model_fields["max_beam_width"]. - default, - max_batch_size: int = BuildConfig.model_fields["max_batch_size"]. - default, - max_num_tokens: int = BuildConfig.model_fields["max_num_tokens"]. - default, - max_seq_len: int = BuildConfig.model_fields["max_seq_len"].default, + max_beam_width: int = _LLM_ARGS_FIELDS["max_beam_width"].default, + max_batch_size: int = _LLM_ARGS_FIELDS["max_batch_size"].default, + max_num_tokens: int = _LLM_ARGS_FIELDS["max_num_tokens"].default, + max_seq_len: int = None, tensor_parallel_size: int = 1, pipeline_parallel_size: int = 1, context_parallel_size: int = 1, @@ -250,19 +246,8 @@ def get_llm_args( dtype=kv_cache_dtype), "cp_config": cp_config, - "build_config": - BuildConfig(max_batch_size=max_batch_size, - max_num_tokens=max_num_tokens, - max_beam_width=max_beam_width, - max_seq_len=max_seq_len) if backend == "tensorrt" else None, "scheduler_config": - SchedulerConfig(capacity_scheduler_policy=CapacitySchedulerPolicy. - GUARANTEED_NO_EVICT, - dynamic_batch_config=DynamicBatchConfig( - enable_batch_size_tuning=True, - enable_max_num_tokens_tuning=False, - dynamic_batch_moving_average_window=128)) - if backend == "tensorrt" else None, + None, "max_batch_size": max_batch_size, "max_beam_width": @@ -419,9 +404,6 @@ def launch_server( # AutoDeploy does not support build_config llm_args.pop("build_config", None) llm = AutoDeployLLM(**llm_args) - elif backend == 'tensorrt' or backend == 'trt': - llm_args.pop("backend") - llm = LLM(**llm_args) else: raise click.BadParameter( f"{backend} is not a known backend, check help for available options.", @@ -489,9 +471,6 @@ async def serve_grpc_async(): from tensorrt_llm._torch.auto_deploy import LLM as AutoDeployLLM llm_args.pop("build_config", None) llm = AutoDeployLLM(**llm_args) - elif backend == "tensorrt" or backend == "trt": - llm_args.pop("backend") - llm = LLM(**llm_args) else: raise click.BadParameter( f"{backend} is not a known backend, check help for available options.", @@ -764,27 +743,6 @@ def launch_visual_gen_server( uvloop.run(server(host, port, sockets=[s])) -class ChoiceWithAlias(click.Choice): - - def __init__(self, - choices: Sequence[str], - aliases: Mapping[str, str], - case_sensitive: bool = True) -> None: - super().__init__(choices, case_sensitive) - self.aliases = aliases - - def to_info_dict(self) -> Dict[str, Any]: - info_dict = super().to_info_dict() - info_dict["aliases"] = self.aliases - return info_dict - - def convert(self, value: Any, param: Optional["click.Parameter"], - ctx: Optional["click.Context"]) -> Any: - if value in self.aliases: - value = self.aliases[value] - return super().convert(value, param, ctx) - - @click.command("serve") @click.argument("model", type=str) @stability_option( @@ -826,8 +784,7 @@ def convert(self, value: Any, param: Optional["click.Parameter"], status="beta") @stability_option( "--backend", - type=ChoiceWithAlias(["pytorch", "tensorrt", "_autodeploy"], - {"trt": "tensorrt"}), + type=click.Choice(["pytorch", "_autodeploy"]), default="pytorch", help="The backend to use to serve the model. Default is pytorch backend.", status="beta") @@ -847,26 +804,26 @@ def convert(self, value: Any, param: Optional["click.Parameter"], status="beta") @stability_option("--max_beam_width", type=int, - default=BuildConfig.model_fields["max_beam_width"].default, + default=_LLM_ARGS_FIELDS["max_beam_width"].default, help="Maximum number of beams for beam search decoding.", status="beta") @stability_option( "--max_batch_size", type=int, - default=BuildConfig.model_fields["max_batch_size"].default, + default=_LLM_ARGS_FIELDS["max_batch_size"].default, help="Maximum number of requests that the engine can schedule.", status="beta") @stability_option( "--max_num_tokens", type=int, - default=BuildConfig.model_fields["max_num_tokens"].default, + default=_LLM_ARGS_FIELDS["max_num_tokens"].default, help= "Maximum number of batched input tokens after padding is removed in each batch.", status="beta") @stability_option( "--max_seq_len", type=int, - default=BuildConfig.model_fields["max_seq_len"].default, + default=None, help="Maximum total length of one request, including prompt and outputs. " "If unspecified, the value is deduced from the model config.", status="beta") @@ -1363,7 +1320,7 @@ def _serve_visual_gen(): @stability_option( "--max_batch_size", type=int, - default=BuildConfig.model_fields["max_batch_size"].default, + default=_LLM_ARGS_FIELDS["max_batch_size"].default, help="Maximum number of requests that the engine can schedule.", status="beta") @stability_option( @@ -1503,7 +1460,7 @@ def serve_encoder(model: str, host: str, port: int, log_level: str, help="The logging level.") @click.option("--max_batch_size", type=int, - default=BuildConfig.model_fields["max_batch_size"].default, + default=_LLM_ARGS_FIELDS["max_batch_size"].default, help="Maximum batch size coalesced into a single encode() call.") @click.option( "--max_num_tokens", diff --git a/tensorrt_llm/disaggregated_params.py b/tensorrt_llm/disaggregated_params.py index 93b9a38f2215..5c91dae3f5de 100644 --- a/tensorrt_llm/disaggregated_params.py +++ b/tensorrt_llm/disaggregated_params.py @@ -4,11 +4,6 @@ import numpy as np -# isort: off -# needed before trying to import bindings to load tensorrt_libs -import tensorrt as trt # noqa -# isort: on - from tensorrt_llm.bindings import executor as tllme diff --git a/tensorrt_llm/evaluate/cnn_dailymail.py b/tensorrt_llm/evaluate/cnn_dailymail.py index f46cf65cc4c3..a676c0f5f69a 100644 --- a/tensorrt_llm/evaluate/cnn_dailymail.py +++ b/tensorrt_llm/evaluate/cnn_dailymail.py @@ -12,14 +12,13 @@ # WITHOUT 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 typing import Iterable, List, Optional, Union +from typing import Iterable, List, Optional import click import datasets import evaluate from .. import LLM as PyTorchLLM -from .._tensorrt_engine import LLM from ..llmapi import RequestOutput from ..logger import logger from ..sampling_params import SamplingParams @@ -124,7 +123,7 @@ def command(ctx, dataset_path: Optional[str], num_samples: int, apply_chat_template: bool, system_prompt: Optional[str], max_input_length: int, max_output_length: int, output_dir: Optional[str]) -> None: - llm: Union[LLM, PyTorchLLM] = ctx.obj + llm: PyTorchLLM = ctx.obj sampling_params = SamplingParams( max_tokens=max_output_length, truncate_prompt_tokens=max_input_length) diff --git a/tensorrt_llm/evaluate/covost2.py b/tensorrt_llm/evaluate/covost2.py index 74e5dea4cd22..0e2051b9464a 100644 --- a/tensorrt_llm/evaluate/covost2.py +++ b/tensorrt_llm/evaluate/covost2.py @@ -22,14 +22,13 @@ task is added here for apples-to-apples comparison. """ -from typing import Iterable, List, Optional, Tuple, Union +from typing import Iterable, List, Optional, Tuple import click import datasets import numpy as np from .. import LLM as PyTorchLLM -from .._tensorrt_engine import LLM from ..llmapi import RequestOutput from ..logger import logger from ..sampling_params import SamplingParams @@ -350,7 +349,7 @@ def command( output_dir: Optional[str], dump_samples_path: Optional[str], ) -> None: - llm: Union[LLM, PyTorchLLM] = ctx.obj + llm: PyTorchLLM = ctx.obj sampling_params = SamplingParams( max_tokens=max_output_length, truncate_prompt_tokens=max_input_length, diff --git a/tensorrt_llm/evaluate/json_mode_eval.py b/tensorrt_llm/evaluate/json_mode_eval.py index 8e6b0b814e50..ce35d474b333 100644 --- a/tensorrt_llm/evaluate/json_mode_eval.py +++ b/tensorrt_llm/evaluate/json_mode_eval.py @@ -14,7 +14,7 @@ # limitations under the License. import json import os -from typing import Iterable, List, Optional, Union +from typing import Iterable, List, Optional import click import datasets @@ -22,7 +22,6 @@ import numpy as np from .. import LLM as PyTorchLLM -from .._tensorrt_engine import LLM from ..llmapi import RequestOutput from ..logger import logger from ..sampling_params import GuidedDecodingParams, SamplingParams @@ -132,7 +131,7 @@ def command(ctx, dataset_path: Optional[str], num_samples: int, random_seed: int, system_prompt: Optional[str], max_input_length: int, max_output_length: int, output_dir: Optional[str]) -> None: - llm: Union[LLM, PyTorchLLM] = ctx.obj + llm: PyTorchLLM = ctx.obj sampling_params = SamplingParams( max_tokens=max_output_length, truncate_prompt_tokens=max_input_length) diff --git a/tensorrt_llm/evaluate/lm_eval.py b/tensorrt_llm/evaluate/lm_eval.py index d41ede2e4180..b077747b6c03 100644 --- a/tensorrt_llm/evaluate/lm_eval.py +++ b/tensorrt_llm/evaluate/lm_eval.py @@ -17,7 +17,7 @@ import os from contextlib import contextmanager from pathlib import Path -from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple, Union +from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple import click import numpy as np @@ -33,7 +33,6 @@ TemplateLM = object from .. import LLM as PyTorchLLM -from .._tensorrt_engine import LLM from ..inputs import (ConversationMessage, MultimodalDataTracker, add_multimodal_placeholders, convert_image_mode) from ..inputs.content_format import ContentFormat @@ -55,7 +54,7 @@ class LmEvalWrapper(TemplateLM): def __init__(self, - llm: Union[LLM, PyTorchLLM], + llm: PyTorchLLM, sampling_params: Optional[SamplingParams] = None, streaming: bool = False, chat_template_kwargs: Optional[dict[str, Any]] = None, @@ -199,7 +198,7 @@ class MultimodalLmEvalWrapper(LmEvalWrapper): """ def __init__(self, - llm: Union[LLM, PyTorchLLM], + llm: PyTorchLLM, sampling_params: Optional[SamplingParams] = None, streaming: bool = False, max_images: int = 999, @@ -263,7 +262,7 @@ def __init__(self, self.interleave = MULTIMODAL_PLACEHOLDER_REGISTRY.get_interleave_placeholders( self.model_type) - def _get_model_type(self, llm: Union[LLM, PyTorchLLM]) -> str: + def _get_model_type(self, llm: PyTorchLLM) -> str: """Extract model type from the model configuration.""" config_path = os.path.join(llm._hf_model_dir, 'config.json') @@ -602,7 +601,7 @@ def save_results(self, results: dict) -> None: logger.info(f"Results saved to {result_path}") def evaluate(self, - llm: Union[LLM, PyTorchLLM], + llm: PyTorchLLM, sampling_params: Optional[SamplingParams] = None, streaming: bool = False, scores_filter: str = None, @@ -663,7 +662,7 @@ def evaluate(self, @classmethod def command_harness(cls, ctx, **kwargs): - llm: Union[LLM, PyTorchLLM] = ctx.obj + llm: PyTorchLLM = ctx.obj # Resolve the post-processor: accept a callable (already-bound) or the # string key "strip_thinking_mmmu" coming from CLI flags. @@ -1336,7 +1335,7 @@ def _get_group_score(metrics: Dict[str, Any], return None def evaluate(self, - llm: Union[LLM, PyTorchLLM], + llm: PyTorchLLM, sampling_params: Optional[SamplingParams] = None, streaming: bool = False) -> float: import lm_eval @@ -1445,7 +1444,7 @@ def evaluate(self, @click.pass_context @staticmethod def command(ctx, **kwargs) -> None: - llm: Union[LLM, PyTorchLLM] = ctx.obj + llm: PyTorchLLM = ctx.obj evaluator = LongBenchV1( dataset_path=kwargs.pop("dataset_path", None), diff --git a/tensorrt_llm/evaluate/longbench_v2.py b/tensorrt_llm/evaluate/longbench_v2.py index 1b2ccaf221ba..40a7d236efa3 100644 --- a/tensorrt_llm/evaluate/longbench_v2.py +++ b/tensorrt_llm/evaluate/longbench_v2.py @@ -24,13 +24,12 @@ import os import re from datetime import datetime -from typing import Any, Dict, Iterable, List, Optional, Union +from typing import Any, Dict, Iterable, List, Optional import click from datasets import load_dataset from .. import LLM as PyTorchLLM -from .._tensorrt_engine import LLM from ..llmapi import RequestOutput from ..logger import logger from ..sampling_params import SamplingParams @@ -779,7 +778,7 @@ def command(ctx, dataset_path: str, prompts_dir: Optional[str], max_len: int, max_input_length: int, max_output_length: int, chat_template_kwargs: Optional[dict[str, Any]], temperature: float, top_p: float) -> None: - llm: Union[LLM, PyTorchLLM] = ctx.obj + llm: PyTorchLLM = ctx.obj sampling_params = SamplingParams(max_tokens=max_output_length, temperature=temperature, diff --git a/tensorrt_llm/evaluate/mmlu.py b/tensorrt_llm/evaluate/mmlu.py index 0b26e52a5af9..8463a639bd94 100644 --- a/tensorrt_llm/evaluate/mmlu.py +++ b/tensorrt_llm/evaluate/mmlu.py @@ -17,14 +17,13 @@ import json import math -from typing import Any, Iterable, List, Optional, Union +from typing import Any, Iterable, List, Optional import click import numpy as np import pandas as pd from .. import LLM as PyTorchLLM -from .._tensorrt_engine import LLM from ..llmapi import RequestOutput from ..logger import logger from ..sampling_params import SamplingParams @@ -312,7 +311,7 @@ def command(ctx, dataset_path: Optional[str], num_samples: int, system_prompt: Optional[str], max_input_length: int, max_output_length: int, check_accuracy: bool, accuracy_threshold: float, output_dir: Optional[str]) -> None: - llm: Union[LLM, PyTorchLLM] = ctx.obj + llm: PyTorchLLM = ctx.obj sampling_params = SamplingParams( max_tokens=max_output_length, truncate_prompt_tokens=max_input_length) diff --git a/tensorrt_llm/evaluate/nemo_skills_eval.py b/tensorrt_llm/evaluate/nemo_skills_eval.py index c1dd60285933..7a3c207047ce 100644 --- a/tensorrt_llm/evaluate/nemo_skills_eval.py +++ b/tensorrt_llm/evaluate/nemo_skills_eval.py @@ -46,7 +46,7 @@ import os import re import tempfile -from typing import Any, Iterable, List, Optional, Union +from typing import Any, Iterable, List, Optional import click @@ -749,9 +749,8 @@ def _level(field: str) -> tuple[float, float]: @classmethod def command_harness(cls, ctx, **kwargs) -> None: from .. import LLM as PyTorchLLM - from .._tensorrt_engine import LLM - llm: Union[LLM, PyTorchLLM] = ctx.obj + llm: PyTorchLLM = ctx.obj evaluator = cls( dataset_path=kwargs.pop("dataset_path", None), split=kwargs.pop("split", None), diff --git a/tensorrt_llm/executor/base_worker.py b/tensorrt_llm/executor/base_worker.py index c4cb84d6bbf8..5896b488f5a8 100644 --- a/tensorrt_llm/executor/base_worker.py +++ b/tensorrt_llm/executor/base_worker.py @@ -33,7 +33,6 @@ from .._utils import (global_mpi_rank, global_mpi_size, mpi_comm, mpi_rank, nvtx_range_debug) from ..bindings import executor as tllm -from ..builder import ConfigEncoder, Engine, EngineConfig from ..llmapi.llm_args import BaseLlmArgs, ExecutorMemoryType, PybindMirror from ..llmapi.tokenizer import TokenizerBase from ..llmapi.tracer import global_tracer @@ -41,7 +40,6 @@ from ..lora_manager import LoraManager from ..prompt_adapter_manager import PromptAdapterManager from ..runtime import ModelConfig -from ..runtime.model_runner import _engine_config_to_model_config from ..sampling_params import BatchedLogitsProcessor, SamplingParams from .executor import GenerationExecutor, IterationResultQueue from .ipc import FusedIpcQueue, IpcQueue @@ -91,7 +89,7 @@ class WorkerExit(GeneratorExit): def __init__( self, - engine: Union[Path, Engine], + engine: Path, executor_config: Optional[tllm.ExecutorConfig] = None, batched_logits_processor: Optional[BatchedLogitsProcessor] = None, postproc_worker_config: Optional[PostprocWorkerConfig] = None, @@ -209,56 +207,12 @@ 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) - - if isinstance(engine, Engine): - return tllm.Executor(engine.engine, - json.dumps(engine.config.to_dict(), - cls=ConfigEncoder), - tllm.ModelType.DECODER_ONLY, - executor_config=executor_config, - managed_weights=engine.managed_weights) - - 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 self._runtime_model_config: Optional[ModelConfig] = None - if self.rank == 0 and isinstance(self.engine, tllm.Executor): - if isinstance(self.engine, Engine): - engine_config = self.engine.config - else: - engine_config = EngineConfig.from_json_file( - f"{self._engine}/config.json") - self._runtime_model_config = _engine_config_to_model_config( - engine_config) - if engine_config.build_config.plugin_config.lora_plugin: - # TODO(azuker): Passing peft cache manager to LoraManager is used for LoRA optimization - # (see LoraManager constructor docstring). Getting the peft cache manager from this - # point in the TRT flow is currently not supported (it's at the CPP - # Executor->ExecutorImpl->TrtGptModel->mPeftCacheManager) therefore for now this LoRA - # optimization is not available in TRT-python flow. - self._lora_manager = LoraManager( - mapping=engine_config.pretrained_config.mapping, - model_config=self._runtime_model_config, - cpp_peft_cache_manager=None) - if engine_config.build_config.max_prompt_embedding_table_size > 0: - self._prompt_adapter_manager = PromptAdapterManager() if self._backend == "pytorch" and self._lora_config is not None: from tensorrt_llm._torch.pyexecutor.resource_manager import \ @@ -275,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 @@ -294,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/executor.py b/tensorrt_llm/executor/executor.py index ea698342a7ae..ba5b843ff347 100644 --- a/tensorrt_llm/executor/executor.py +++ b/tensorrt_llm/executor/executor.py @@ -20,7 +20,6 @@ from .._utils import mpi_world_size from ..bindings import executor as tllm -from ..builder import Engine from ..conversation_params import ConversationParams from ..disaggregated_params import DisaggregatedParams from ..llmapi.llm_args import BaseLlmArgs, TorchLlmArgs @@ -529,7 +528,7 @@ def _create_ipc_executor( @staticmethod def create( - engine: Union[Path, Engine], + engine: Path, executor_config: Optional[tllm.ExecutorConfig] = None, batched_logits_processor: Optional[BatchedLogitsProcessor] = None, model_world_size: int = 1, diff --git a/tensorrt_llm/executor/ray_gpu_worker.py b/tensorrt_llm/executor/ray_gpu_worker.py index 0b86d5255c8a..7a312aa003f4 100644 --- a/tensorrt_llm/executor/ray_gpu_worker.py +++ b/tensorrt_llm/executor/ray_gpu_worker.py @@ -5,7 +5,7 @@ from functools import wraps from pathlib import Path from queue import Queue -from typing import Any, List, Optional, Type, Union +from typing import Any, List, Optional, Type import ray import torch @@ -17,7 +17,6 @@ from .. import TorchLlmArgs from ..bindings import executor as tllm -from ..builder import Engine from ..llmapi.llm_args import BaseLlmArgs, ExecutorMemoryType from ..llmapi.tokenizer import TokenizerBase from ..llmapi.utils import configure_cpu_affinity @@ -214,7 +213,7 @@ class RayGPUWorker(RpcWorkerMixin, BaseWorker): def __init__( self, device_id: int, - engine: Union[Path, Engine], + engine: Path, executor_config: Optional[tllm.ExecutorConfig] = None, batched_logits_processor: Optional[BatchedLogitsProcessor] = None, postproc_worker_config: Optional[PostprocWorkerConfig] = None, diff --git a/tensorrt_llm/executor/rpc_worker.py b/tensorrt_llm/executor/rpc_worker.py index 5433cbd37b54..667fa985921f 100644 --- a/tensorrt_llm/executor/rpc_worker.py +++ b/tensorrt_llm/executor/rpc_worker.py @@ -1,7 +1,7 @@ from pathlib import Path from queue import Queue from threading import Event -from typing import Optional, Union +from typing import Optional import nvtx @@ -10,7 +10,6 @@ from .._utils import mpi_rank from ..bindings import executor as tllm -from ..builder import Engine from ..llmapi.llm_args import BaseLlmArgs from ..llmapi.tokenizer import TokenizerBase from ..logger import set_level @@ -48,7 +47,7 @@ class RpcWorker(RpcWorkerMixin, BaseWorker): def __init__( self, - engine: Union[Path, Engine], + engine: Path, executor_config: Optional[tllm.ExecutorConfig] = None, is_llm_executor: Optional[bool] = None, batched_logits_processor: Optional[BatchedLogitsProcessor] = None, @@ -111,7 +110,7 @@ def start(self): @staticmethod def main_task( - engine: Union[Path, Engine], + engine: Path, rpc_addr: str, *, executor_config: Optional[tllm.ExecutorConfig] = None, diff --git a/tensorrt_llm/executor/worker.py b/tensorrt_llm/executor/worker.py index c05e44f02cd3..293cc88a4d9b 100644 --- a/tensorrt_llm/executor/worker.py +++ b/tensorrt_llm/executor/worker.py @@ -5,7 +5,7 @@ import traceback from concurrent.futures import ProcessPoolExecutor from pathlib import Path -from typing import List, Optional, Union +from typing import List, Optional import zmq @@ -13,7 +13,6 @@ from .._utils import mpi_comm, mpi_rank, print_all_stacks from ..bindings import executor as tllm -from ..builder import Engine from ..llmapi.llm_args import BaseLlmArgs from ..llmapi.mpi_session import set_mpi_session_cpp from ..llmapi.tokenizer import TokenizerBase @@ -38,7 +37,7 @@ class GenerationExecutorWorker(RpcWorkerMixin, BaseWorker): def __init__( self, - engine: Union[Path, Engine], + engine: Path, executor_config: Optional[tllm.ExecutorConfig] = None, batched_logits_processor: Optional[BatchedLogitsProcessor] = None, postproc_worker_config: Optional[PostprocWorkerConfig] = None, @@ -153,11 +152,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() @@ -165,7 +159,7 @@ def block_subordinates(self): @print_traceback_on_error def worker_main( - engine: Path | Engine, + engine: Path, worker_queues: WorkerCommIpcAddrs, log_level: str, executor_config: Optional[tllm.ExecutorConfig] = None, diff --git a/tensorrt_llm/functional.py b/tensorrt_llm/functional.py old mode 100755 new mode 100644 index df80459359aa..f4b97746f615 --- a/tensorrt_llm/functional.py +++ b/tensorrt_llm/functional.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# 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"); @@ -12,661 +12,14 @@ # WITHOUT 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 math -import weakref -from collections import OrderedDict from enum import IntEnum -from functools import partial -from typing import List, Optional, Sequence, Tuple, Union +from typing import List, Optional import numpy as np import torch - -# isort: off -import tensorrt as trt -# isort: on - -from . import graph_rewriting as gw -from ._common import default_net, default_trtnet, precision -from ._utils import (QuantModeWrapper, bf16_array, bool_array, - dim_resolve_negative, dim_to_trt_axes, dims_array, - fp16_array, fp32_array, get_sm_version, int32_array, - int64_array, np_dtype_to_trt, str_dtype_to_trt, - trt_dtype_to_np, trt_dtype_to_str) -from .network import PluginInfo, get_np_weight, set_np_weight, set_plugin_info -from .plugin import TRT_LLM_PLUGIN_NAMESPACE, current_all_reduce_helper -from .quantization import QuantMode - - -class DimRange(object): - ''' - One DimRange object stores the ranges of all the dimensions of one tensor in one optimization profile. - For example, tensor has 2 dimensions. Then the data members are: - self.min = [dim 0 min, dim 1 min] - self.opt = [dim 0 opt, dim 1 opt] - self.max = [dim 0 max, dim 1 max] - For static dimension, it has min==opt==max, thus the shape param in the ctor can be an integer - ''' - - def __init__(self, shape: List[Union[int, List[int], Tuple[int, int, int]]], - names: List[str]): - ''' - Parameters: - shape: a list with length N, each element is an integer or a 3-elements tuple/list of int, - where N is the number of dimensions for a tensor. - When one element is an integer, it means that dimension is static. - Otherwise, when one element is a tuple/list, it means the dimension is dynamic. - The 3 elements in one tuple/list is ordered by (min, opt, max), and this function asserts - 0 <= min <= opt <= max. - - Example, for a 3 rank tensor, with 1st dimension being static and has value 16, and second dimension being dynamic with - min/opt/max values being 1/8/32, and 3rd dimension being static and has value 8. - The shape parameter could be: - [16, (1, 8, 32), 8] - It has same semantics of - [(16, 16, 16), (1, 8, 32), (8, 8, 8)] - ''' - self.min = [] - self.opt = [] - self.max = [] - self.dimension_names = names - assert len(names) == len( - shape - ), "Expecting shape list and name list must have same length, got {shape=}, {name=}" - for dim in shape: - if isinstance(dim, (list, tuple)): - assert len(dim) == 3 and 0 <= dim[0] <= dim[1] <= dim[2], \ - "Each dimension must specify a 3-elements tuple or list in the order of (min,opt,max), got {dim=}" - self.min.append(dim[0]) - self.opt.append(dim[1]) - self.max.append(dim[2]) - elif isinstance(dim, int): - self.min.append(dim) - self.opt.append(dim) - self.max.append(dim) - else: - raise AttributeError( - f'Dimension should be [min, opt, max] (dynamic shape) or int (specific value). Got {type(dim)}' - ) - - def __eq__(self, __value: object) -> bool: - return isinstance(__value, DimRange) and \ - self.dimension_names == __value.dimension_names and \ - self.min == __value.min and self.opt == __value.opt and self.max == __value.max - - def __repr__(self) -> str: - return str(self) - - def __str__(self) -> str: - return f"{self.dimension_names=} {self.min=}, {self.opt=}, {self.max=})" - - def __hash__(self) -> int: - return hash(str(self)) - - -class Tensor(object): - ''' - The class to represent dense tensors. - - A dense tensor is named, has a shape and contains typed elements. Each - dimension of a tensor can either be static or dynamic. Static dimensions - are known at engine compilation by TensorRT. Dynamic dimensions can take - values determined at runtime. The tensor can be located on the host (CPU) - or the device (GPU). - ''' - - def __init__(self, - name=None, - dtype=None, - shape=None, - dim_range=None, - is_network_input=True, - location=trt.TensorLocation.DEVICE, - network=None, - trt_tensor=None): - ''' - Parameters: - name : str - The name of the tensor. - - dtype : tensorrt.DataType - The type of the elements of the tensor. See the TensorRT - documentation for list of supported data types. - - shape : tensorrt.Dims - The dimensions of the tensor. In TensorRT-LLM, tensors can have - static or dynamic dimensions (it is possible to mix static and - dynamic dimensions). A static dimension is known when the - TensorRT engine is built. A dynamic dimension can be set when - the engine is executed. Use -1 for dynamic dimensions. - - dim_range : OrderedDict - An ordered dictionary (the positions of the elements matter) - that associates a name and a range of values to the dimensions. - For a static dimension, the range must be limited to a single - value. For a dynamic dimension, the range is defined by three - values [min, opt, max] where min and max are, respectively, the - smallest and largest possible values of that dimension. The - opt value is used by TensorRT to optimize the engine for the - most common case. - - Assume there is N optimization profiles, each item dim_range dict is ordered by: - (dynamic dimension name : [profile 0 (min, opt, max), profile 1 (min, opt, max), ... profile N(min, opt, max)]) - or it's following when the dimension is static (can think as min==opt==max): - (static dimension name : [profile 0 value, profile 1 value, ... profile N value]) - For static dimension the profile 0-N value must be same, (TODO: can it be simplified to be only 1 value?) - And number of keys is equal to number of optimization profiles. - - is_network_input : bool - A boolean indicating if that tensor is an input of the network. - Inputs must be provided by the user to run the engine. - - location : tensorrt.TensorLocation - A flag to indicate where the tensor will be located. It can be - on the host (CPU) or the device (GPU). - - network: Network - A parent Network instance, that helps to fine the users of this tensor. - - trt_tensor: trt.ITensor - Construct with the ITensor instance directly, and no shape profiles are required. - ''' - - # Layout of self.profiles - # Opt profile 0: dim 0 (min, opt, max), dim 1 (min, opt, max) ... dim M - # Opt profile 1: dim 0 (min, opt, max), dim 1 (min, opt, max) ... dim M - # ... - # Opt profile N: dim 0 ... dim M - - # So from the dim_range arg to self.profiles conversion, there is a layout transpose - # dim_range arg is: {M dimension x N profiles}, while self.profiles layout is {N profiles x M dimensions} - if isinstance(dtype, str): - dtype = str_dtype_to_trt(dtype) - - self.profiles = [] - - self.is_tensor_wrapper = False # specially for the graph rewriter - - # work as a wrapper for a trt.ITensor, this is used specially in the graph rewriter - if trt_tensor is not None: - self.is_tensor_wrapper = True - assert network is not None - self.trt_tensor = trt_tensor - self._network = weakref.ref(network) - assert not is_network_input, "is_network_input should be False when trt_tensor is not None" - return - - # be cautious here, the weakref is critical to avoid circular referencing before Network and Tensor - # using strong reference will likely cause significant peak memory increase, since Network objects - # holds the weights data. - self._network = weakref.ref(default_net()) - self.is_network_input = is_network_input - if is_network_input: - if dim_range is not None: - assert isinstance(dim_range, OrderedDict) - assert len( - dim_range - ) >= 1, f"Each input tensor shall have at least one dimension, tensor '{name}' found {dim_range=}" - - found_profiles = [ - len(ranges) for _, ranges in dim_range.items() - ] - assert all( - [x == found_profiles[0] for x in found_profiles] - ), f"Expecting all the dimensions in the dim_range has same number of profiles, tensor '{name}' got {dim_range=}" - - num_opt_profile = len(list(dim_range.items())[0][1]) - assert num_opt_profile >= 1 - for i in range(num_opt_profile): - range_shape = [] - dimension_names = [] - for dim, ranges in dim_range.items(): - assert isinstance(ranges, (list, tuple)) - range_shape.append(ranges[i]) - dimension_names.append(dim) - self.profiles.append(DimRange(range_shape, dimension_names)) - - default_net()._add_input(self, name, dtype, shape, dim_range) - self.name = name - self.dtype = dtype - self.shape = shape - self.location = location - - @property - def network(self): - return self._network() - - @property - def name(self): - ''' - The name of the tensor. - ''' - return self.trt_tensor.name - - @name.setter - def name(self, name): - ''' - Set the name of the tensor. - ''' - if name is not None: - self.trt_tensor.name = name - - @property - def dtype(self): - ''' - The type of the elements in the tensor. - ''' - return self.trt_tensor.dtype - - @dtype.setter - def dtype(self, dtype): - ''' - Set the type of the elements in the tensor. - ''' - if dtype is not None: - self.trt_tensor.dtype = dtype - - @property - def shape(self): - ''' - The shape of the tensor. - ''' - return self.size() - - @shape.setter - def shape(self, shape): - ''' - Set the shape of the tensor. See __init__. - ''' - if shape is not None: - self.trt_tensor.shape = shape - - @property - def location(self): - ''' - The physical location of the tensor (on the host or the device). - ''' - return self.trt_tensor.location - - @location.setter - def location(self, location): - ''' - Set the physical location of the tensor (on the host or the device). See __init__. - ''' - if location is not None: - self.trt_tensor.location = location - - def mark_output(self, - name: Optional[str] = None, - dtype: Optional[Union[str, trt.DataType]] = None): - ''' - Mark a tensor as a network output. - - When a tensor is marked as an output, its content can be obtained after - the execution of the TensorRT engine. The user is responsible for - allocating buffers to store the output tensors when preparing the - execution of the TensorRT engine. - ''' - if name is None: - name = self.name - - if isinstance(dtype, str): - dtype = str_dtype_to_trt(dtype) - - assert dtype is None or isinstance(dtype, trt.DataType) - default_net()._mark_output(self, name, dtype) - - def __add__(self, b): - ''' - See functional.add. - ''' - return add(self, b) - - def __radd__(self, b): - ''' - See functional.add. - ''' - return add(b, self) - - def __sub__(self, b): - ''' - See functional.sub. - ''' - return sub(self, b) - - def __rsub__(self, b): - ''' - See functional.sub. - ''' - return sub(b, self) - - def __mul__(self, b): - ''' - See functional.mul. - ''' - return mul(self, b) - - def __rmul__(self, b): - ''' - See functional.mul. - ''' - return mul(b, self) - - def __truediv__(self, b): - ''' - See functional.div. - ''' - return div(self, b) - - def __floordiv__(self, b): - ''' - See functional.floordiv. - ''' - return floordiv(self, b) - - def __mod__(self, b): - ''' - See functional.floordiv. - ''' - return modulo(self, b) - - def __lt__(self, b): - ''' - See functional.lt. - ''' - return lt(self, b) - - def __gt__(self, b): - ''' - See functional.gt. - ''' - return gt(self, b) - - def __eq__(self, b): - ''' - See functional.eq. - ''' - if self.is_tensor_wrapper: - # for graph rewriter - return hash(self) == hash(b) - else: - # for creating the network - return eq(self, b) - - def __ge__(self, b): - ''' - Maps to functional.gt or functional.eq. - ''' - return op_or(self.__gt__(b), self.__eq__(b)) - - def __le__(self, b): - ''' - Maps to functional.lt or functional.eq. - ''' - return op_or(self.__lt__(b), self.__eq__(b)) - - def view(self, shape, zero_is_placeholder=True): - ''' - See functional.view. - ''' - return view(self, shape, zero_is_placeholder) - - def flatten(self, start_dim=0, end_dim=-1): - ''' - See functional.flatten. - ''' - return flatten(self, start_dim, end_dim) - - def permute(self, dims): - ''' - See functional.permute. - ''' - return permute(self, dims) - - def transpose(self, dim0, dim1): - ''' - See functional.transpose. - ''' - return transpose(self, dim0, dim1) - - def mean(self, dim, keepdim=False): - ''' - See functional.mean. - ''' - return mean(self, dim, keepdim) - - def max(self, dim, keepdim=False): - ''' - See functional.max. - ''' - return max(self, dim, keepdim) - - def abs(self): - ''' - See functional.abs. - ''' - return abs(self) - - def sqrt(self): - ''' - See functional.sqrt. - ''' - return sqrt(self) - - def squeeze(self, dim, zero_is_placeholder): - ''' - See functional.squeeze. - ''' - return squeeze(self, dim, zero_is_placeholder) - - def unsqueeze(self, dim): - ''' - See functional.squeeze. - ''' - return unsqueeze(self, dim) - - def log(self): - ''' - See functional.log. - ''' - return log(self) - - def cast(self, dtype): - ''' - See functional.cast. - ''' - return cast(self, dtype) - - def size(self, dim=None): - ''' - Returns the shape of the tensor if the dim parameter is None. - Otherwise, returns a size of the dimension indicated by dim. The - behavior is undefined if dim is negative or exceeds the rank of the - tensor. - ''' - if dim is None: - return self.trt_tensor.shape - - return self.trt_tensor.shape[dim] - - def rank(self): - ''' - Returns the rank (i.e. the number of dimensions) of the tensor. - ''' - return len(self.trt_tensor.shape) - - def ndim(self): - ''' - Returns the rank (i.e. the number of dimensions) of the tensor. - ''' - return self.rank() - - def split(self, split_size_or_sections, dim=0): - ''' - See functional.split. - ''' - return split(self, split_size_or_sections, dim) - - def select(self, dim, index): - ''' - See functional.select. - ''' - return select(self, dim, index) - - def unbind(self, dim=0): - ''' - See functional.unbind. - ''' - return unbind(self, dim) - - def repeat(self, sizes): - ''' - See functional.repeat - ''' - return repeat(self, sizes) - - def is_dynamic(self, dim=None): - ''' - If the argument 'dim' is None, that function returns a boolean that - indicates if the tensor contains a dynamic dimension (True) or not - (False). In that case, the first dimension is excluded (as it usually - corresponds to the batch size). If the argument is an integer, that - functions returns a boolean that indicates if the dimension 'dim' is - dynamic (True) or not (False). - ''' - if dim is not None: - return self.trt_tensor.shape[dim] == -1 - - for i, s in enumerate(self.trt_tensor.shape): - if i != 0 and s == -1: - return True - - return False - - # graph writer related functions - - def get_parent(self): - ''' Get the layer that produces this tensor. ''' - return self.network.get_tensor_parent(self) - - def get_users(self): - ''' Get the layers that use this tensor as an input. ''' - return self.network.get_tensor_users(self) - - def replace_all_uses_with(self, new_tensor): - ''' - Replace all uses of this tensor as an input to consumer layers - ''' - - self.network.is_graph_altered = True - users = self.get_users() - for user in users: - inputs_changed = 0 - for i in range(user.num_inputs): - if user.get_inputs(i)[0].trt_tensor is self.trt_tensor: - inputs_changed += 1 - user.set_input(i, new_tensor.trt_tensor) - assert inputs_changed >= 1, "Tensor not found in layer inputs" - - # update the FLayerMetadata as well - flayer = gw.FLayerInfoMemo.instance().get(user.name) - flayer and flayer.replace_input_with(self, new_tensor) - - def is_trt_wrapper(self): - ''' - Check if there is a trt.ITensor member inside, which is required for - graph rewriter. In order to differentiate usages, it may be necessary - to have an inheritance hierarchy. - ''' - if hasattr(self, 'trt_tensor'): - return True - else: - return False - - def __hash__(self): - if self.is_trt_wrapper(): - return id(self.trt_tensor) - else: - return id(None) - - def __repr__(self): - return f"TensorRT LLM Tensor: {self.name=} {self.dtype=} {self.shape=}" - - def __xor__(self, b): - ''' - Maps to functional.gt or functional.eq. - ''' - print(f"self.shape: {self.shape}, b.shape: {b.shape}") - a, b = broadcast_helper(self, b) - print(f"a.shape: {a.shape}, b.shape: {b.shape}") - return op_xor(a, b) - - -def _create_tensor(trt_tensor: trt.ITensor, producer: trt.ILayer) -> Tensor: - ''' - A helper function to create a TensorRT LLM Tensor object that encapsulates - the connection between the TensorRT tensor (trt.ITensor) and the layer - (trt.ILayer) that produces it. - - That function is expected to be used as: - - # Insert a new layer in the network using the TensorRT API: - layer = default_trtnet().add_(...) - # Extract the first output of that layer and connect it to the layer. - return _create_tensor(layer.get_output(0), layer) - - That function also sets the precision of the layer/producer to the default - precision of the network. - - Parameters: - trt_tensor : trt.ITensor - The TensorRT tensor to connect to its producer (the layer). - - producer : trt.ILayer - The producer. - - Returns: - The TensorRT LLM tensor (functional.Tensor) that encapsulates the - TensorRT tensor and the layer that produces it. The former is - accessible through the attribute 'trt_tensor' and the latter using the - attribute 'producer'. - ''' - assert trt_tensor is not None - assert producer is not None - - # Set the layer name since this is the only - # centralized location to pass the name from - # module space to the TRT IR - default_net()._set_layer_name(producer) - - assert trt_tensor.shape.__len__( - ) >= 0, f"tensor {trt_tensor.name} has an invalid shape" - tensor = Tensor(name=trt_tensor.name, - dtype=trt_tensor.dtype, - shape=trt_tensor.shape, - is_network_input=False) - tensor.trt_tensor = trt_tensor - tensor.producer = producer - - # tb.print_stack(limit=10) # FOR DEBUGGING: filter producer.name if needed - if default_net().dtype is not None and not default_net().strongly_typed: - if producer.type not in [ - trt.LayerType.SHAPE, trt.LayerType.CONSTANT, - trt.LayerType.GATHER, trt.LayerType.CONCATENATION - ]: - producer.precision = default_net().dtype - assert tensor is not None - - if gw.FLayerInfoMemo.instance().cur_flayer is not None: - gw.FLayerInfoMemo.instance().cur_flayer.layer_name = producer.name - - return tensor - - -def _add_plugin_info(layer, plugin_creator: trt.IPluginCreator, - plugin_name: str, pfc: trt.PluginFieldCollection) -> None: - plugin_info = PluginInfo(plugin_creator, plugin_name, pfc) - set_plugin_info(default_net().trt_network, layer.name, plugin_info) +from torch import Tensor class RotaryScalingType(IntEnum): @@ -691,7 +44,7 @@ def from_string(s): try: return RotaryScalingType[key] except KeyError: - raise ValueError(f'Unsupported rotary scaling type: {s}') + raise ValueError(f"Unsupported rotary scaling type: {s}") class PositionEmbeddingType(IntEnum): @@ -705,7 +58,9 @@ class PositionEmbeddingType(IntEnum): chatglm = 7 yarn = 8 mrope = 9 - deferred = 10 # Apply customized positional embedding by using an external embedder. K will be cached before embedding. + # Apply customized positional embedding by using an external embedder. + # K will be cached before embedding. + deferred = 10 def is_rope(self) -> bool: return self in [ @@ -736,7 +91,7 @@ def from_string(s): try: return PositionEmbeddingType[s] except KeyError: - raise ValueError(f'Unsupported position embedding type: {s}') + raise ValueError(f"Unsupported position embedding type: {s}") class AttentionMaskType(IntEnum): @@ -749,6943 +104,418 @@ class AttentionMaskType(IntEnum): custom_mask = 6 -class LayerNormType(IntEnum): - LayerNorm = 0 - RmsNorm = 1 - GroupNorm = 2 - - -class LayerNormPositionType(IntEnum): - pre_layernorm = 0 - post_layernorm = 1 - - -class MLPType(IntEnum): - MLP = 0 - GatedMLP = 1 - FusedGatedMLP = 2 - - -class SliceInputType(IntEnum): - data = 0 - start = 1 - size = 2 - stride = 3 - fill_value = 4 - axes = 5 - - -def activation(input: Tensor, act_type: trt.ActivationType) -> Tensor: - ''' - Add an activation function. - - Parameters: - input : Tensor - The input tensor on which the activation function is applied. - - act_type : trt.ActivationType - The type of the activation (RELU, TANH, SIGMOID, ...). - - The following closures are defined in functional.*: - - relu for op=trt.ActivationType.RELU - tanh for op=trt.ActivationType.TANH - sigmoid for op=trt.ActivationType.SIGMOID - - Returns: - The tensor produced by the activation layer. - ''' - layer = default_trtnet().add_activation(input.trt_tensor, act_type) - return _create_tensor(layer.get_output(0), layer) - - -def int_clip(input: Tensor, lower: int, upper: int) -> Tensor: - assert lower <= upper, f"Lower bound must be less than or equal to upper bound i.e. {lower} <= {upper}" - res = minimum(input, upper) - res = maximum(res, lower) - return res - - -def clip(input: Tensor, alpha: float, beta: float) -> Tensor: - ''' - Add a CLIP operation that sets the range to [alpha, beta]. +class AllReduceStrategy(IntEnum): + NCCL = 0 + MIN_LATENCY = 1 + UB = 2 + AUTO = 3 + ONESHOT = 4 + TWOSHOT = 5 + LOWPRECISION = 6 + MNNVL = 7 + NCCL_SYMMETRIC = 8 + SYMM_MEM = 9 # PyTorch symmetric memory with MULTIMEM - Parameters: - input : Tensor - The input tensor on which the activation function is applied. - alpha : float - The lower bound of the CLIP function. +class AllReduceFusionOp(IntEnum): + NONE = 0 + RESIDUAL_RMS_NORM = 1 + LAST_PROCESS_FOR_UB = 2 + RESIDUAL_RMS_PREPOST_NORM = 3 + RESIDUAL_RMS_NORM_QUANT_FP8 = 4 + RESIDUAL_RMS_NORM_QUANT_NVFP4 = 5 + RESIDUAL_RMS_NORM_OUT_QUANT_FP8 = 6 + RESIDUAL_RMS_NORM_OUT_QUANT_NVFP4 = 7 + MOE_FINALIZE_ALLREDUCE_RESIDUAL_RMS_NORM = 8 + RMS_NORM = 9 - beta : float - The upper bound of the CLIP function. - Returns: - The tensor produced by the activation layer. - ''' - layer = default_trtnet().add_activation(input.trt_tensor, - trt.ActivationType.CLIP) - layer.alpha = alpha - layer.beta = beta - return _create_tensor(layer.get_output(0), layer) +class AllReduceParams: + + def __init__( + self, + strategy: AllReduceStrategy = AllReduceStrategy.AUTO, + fusion_op: AllReduceFusionOp = AllReduceFusionOp.NONE, + bias: Optional[Tensor] = None, + residual: Optional[Tensor] = None, + norm_weight: Optional[Tensor] = None, + scale: Optional[Tensor] = None, + norm_pre_residual_weight: Optional[Tensor] = None, + eps: float = 1e-06, + enable_allreduce: bool = True, + trigger_completion_at_end: bool = True, + ): + self.strategy = strategy + self.fusion_op = fusion_op + self.bias = bias + self.residual = residual + self.norm_weight = norm_weight + self.scale = scale + self.norm_pre_residual_weight = norm_pre_residual_weight + self.eps = eps + # For torch path only, has no effect on TRT path + self.enable_allreduce = enable_allreduce + self.trigger_completion_at_end = trigger_completion_at_end + assert fusion_op in (AllReduceFusionOp.NONE.value, + AllReduceFusionOp.RMS_NORM.value) or (residual + is not None) + def has_affine(self): + return 1 if self.norm_weight is not None else 0 -relu = partial(activation, act_type=trt.ActivationType.RELU) -tanh = partial(activation, act_type=trt.ActivationType.TANH) -sigmoid = partial(activation, act_type=trt.ActivationType.SIGMOID) + def has_bias(self): + return 1 if self.bias is not None else 0 + def has_scale(self): + return 1 if self.scale is not None else 0 -def silu(input: Tensor) -> Tensor: - ''' - Add a SiLU (`x * sigmoid(x)`) operation. - Parameters: - input : Tensor - The input tensor on which the activation function is applied. +class MoEAllReduceParams(AllReduceParams): - Returns: - The tensor produced by the activation layer. - ''' - return input * sigmoid(input) + def __init__( + self, + device_num_experts: Optional[Tensor] = None, + expert_scale_factor: Optional[Tensor] = None, + expanded_idx_to_permuted_idx: Optional[Tensor] = None, + shared_expert_output: Optional[Tensor] = None, + bias: Optional[Tensor] = None, + residual: Optional[Tensor] = None, + norm_weight: Optional[Tensor] = None, + scale: Optional[Tensor] = None, + norm_pre_residual_weight: Optional[Tensor] = None, + eps: float = 1e-06, + enable_allreduce: bool = True, + is_cutlass_min_latency: bool = False, + ): + super().__init__( + bias=bias, + residual=residual, + norm_weight=norm_weight, + scale=scale, + norm_pre_residual_weight=norm_pre_residual_weight, + eps=eps, + enable_allreduce=enable_allreduce, + ) + self.device_num_experts = device_num_experts + self.expert_scale_factor = expert_scale_factor + self.expanded_idx_to_permuted_idx = expanded_idx_to_permuted_idx + self.shared_expert_output = shared_expert_output + self.is_cutlass_min_latency = is_cutlass_min_latency + def is_valid(self): + if self.is_cutlass_min_latency: + return (self.device_num_experts is not None + and self.expert_scale_factor is not None + and self.shared_expert_output is not None) + else: + return self.expanded_idx_to_permuted_idx is not None -def swiglu(input: Tensor) -> Tensor: - ''' - Add a SwiGLU (`x * SiLU(gate)`) operation. - That function takes a tensor, splits it into two halves along the last - dimension, applies SiLU to the second half and multiply the results. The - behavior is undefined if the last dimension is not even. +class RopeEmbeddingUtils: - Parameters: - input : Tensor - The input tensor on which the activation function is applied. + @staticmethod + # ref: https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_rope_utils.py#L298 + def apply_llama3_scaling(inv_freqs: np.ndarray, rope_scaling_config: dict): + scale_factor = rope_scaling_config.get("factor", 8.0) + low_freq_factor = rope_scaling_config.get("low_freq_factor", 1.0) + high_freq_factor = rope_scaling_config.get("high_freq_factor", 4.0) + old_context_len = rope_scaling_config.get( + "original_max_position_embeddings", 8192) - Returns: - The tensor produced by the activation layer. - ''' - x, gate = chunk(input, 2, dim=-1) - return silu(gate) * x + low_freq_wavelen = old_context_len / low_freq_factor + high_freq_wavelen = old_context_len / high_freq_factor + new_inv_freqs = [] + for inv_freq in inv_freqs: + wavelen = 2 * math.pi / inv_freq + if wavelen < high_freq_wavelen: + new_inv_freqs.append(inv_freq) + elif wavelen > low_freq_wavelen: + new_inv_freqs.append(inv_freq / scale_factor) + else: + assert low_freq_wavelen != high_freq_wavelen + smooth = (old_context_len / wavelen - low_freq_factor) / ( + high_freq_factor - low_freq_factor) + new_inv_freqs.append((1 - smooth) * inv_freq / scale_factor + + smooth * inv_freq) + return np.array(new_inv_freqs, dtype=inv_freqs.dtype) + @staticmethod + def create_sinusoidal_positions(num_pos: int, + dim: int, + theta: float = 10000.0, + dtype=np.float32): + inv_freq = 1.0 / (theta**(np.arange(0, dim, 2) / dim)).astype(dtype) + sinusoid_inp = np.einsum("i , j -> i j", + np.arange(num_pos, dtype=dtype), + inv_freq, + dtype=dtype) + concat = np.concatenate((np.sin(sinusoid_inp), np.cos(sinusoid_inp)), + axis=1) + return np.expand_dims(concat, axis=0).astype(dtype) -def squared_relu(x: Tensor) -> Tensor: - ''' - Add a Squared ReLU operation. + @staticmethod + def create_sinusoidal_positions_for_attention_plugin( + num_pos: int, + dim: int, + theta: float = 10000.0, + scale: float = 1.0, + scale_type: RotaryScalingType = RotaryScalingType.none, + # Other scaling configs that only used by certain scaling types. + rope_scaling_config: dict = None, + duplicate_data: bool = False, + dtype=np.float32, + ): + if scale_type == RotaryScalingType.linear: + scale = 1.0 / scale + if scale_type == RotaryScalingType.llama3: + assert rope_scaling_config is not None, "rotary_scaling config must be provided." + inv_freq = 1.0 / (theta**(np.arange(0, dim, 2) / dim)).astype(dtype) + inv_freq = RopeEmbeddingUtils.apply_llama3_scaling( + inv_freq, rope_scaling_config) + elif scale_type == RotaryScalingType.dynamic: + # Make sure scaling_alpha exists in rope_scaling + # Ref: https://huggingface.co/tencent/Hunyuan-A13B-Instruct-FP8/blob/main/modeling_hunyuan.py#L346 + assert rope_scaling_config["alpha"] is not None, ( + "rope_scaling_config.alpha must be provided.") + scaling_alpha = rope_scaling_config["alpha"] + adjusted_base = theta * (scaling_alpha**(dim / (dim - 2))) + inv_freq = 1.0 / (adjusted_base**( + np.arange(0, dim, 2, dtype=dtype) / dim)).astype(dtype) + else: + inv_freq = scale / (theta + **(np.arange(0, dim, 2) / dim)).astype(dtype) + sinusoid_inp = np.expand_dims( + np.einsum("i , j -> i j", + np.arange(num_pos, dtype=dtype), + inv_freq, + dtype=dtype), + axis=-1, + ) + if duplicate_data: + sinusoid_inp = np.concatenate((sinusoid_inp, sinusoid_inp), axis=-2) + # fuse cos/sin into float2 (cos, sin). + concat = np.concatenate( + (np.cos(sinusoid_inp), np.sin(sinusoid_inp)), + axis=-1) # np.cos(sinusoid_inp).shape = (32768, 64, 1) - This function applies ReLU and squares the output. + return inv_freq, concat.reshape(1, -1).astype(dtype) - Parameters: - input : Tensor - The input tensor on which the activation function is applied. - - Returns: - The tensor produced by the activation layer. - ''' - return pow(relu(x), 2.0) - - -def cast(input: Tensor, dtype: Union[str, trt.DataType]): - ''' - Add a cast operation. - - For an input tensor of type INT8, this function sets the dynamic range of - the input to [-127, 127] for automatic dequantization. For a cast into - INT8, that function sets the dynamic range of the output to [-127, 127] for - automatic quantization. - - Parameters: - input : Tensor - The input tensor on which the cast is applied. - - dtype : str or trt.DataType - The data type of the output tensor after the cast. When 'dtype' is - provided as a string, it must be a name amongst the valid names. - See _str_to_trt_dtype_dict in _utils.py for a list of supported - types and type names. - - Returns: - The tensor produced by the inserted layer. - ''' - if isinstance(dtype, str): - cvt_dtype = str_dtype_to_trt(dtype) - elif isinstance(dtype, trt.DataType): - cvt_dtype = dtype - else: - raise TypeError("%s is not supported" % type(dtype)) - - if input.dtype == cvt_dtype: - # If input type and cast dtype are the same, do nothing - return input - - layer = default_trtnet().add_cast(input.trt_tensor, cvt_dtype) - if not default_net().strongly_typed: - layer.set_output_type(0, cvt_dtype) - output = _create_tensor(layer.get_output(0), layer) - if not default_net().strongly_typed: - if input.dtype == str_dtype_to_trt('int8'): - layer.get_input(0).set_dynamic_range(-127, 127) - if cvt_dtype == str_dtype_to_trt('int8'): - layer.get_output(0).set_dynamic_range(-127, 127) - - return output - - -def flip(input: Tensor, dims: Sequence[int]) -> Tensor: - ''' - Reverses the order of an n-D tensor along given axis in dims. - - That flip operation maps to a TensorRT ISliceLayer. For the dimensions - listed in dims it copies the elements from the last one to the first one - (from (N-1) down to 0 with a step of -1). For the dimensions not in 'dims', - it copies the elements from the first one to the last one (from 0 to N-1 - with a step of 1). - - Parameters: - input : Tensor - The input tensor on which the cast is applied. - - dims : list or tuple - The axes to flip. Negative indices are supported. - - Returns: - The tensor produced by the inserted layer. - ''' - assert not input.is_dynamic() - - ndim = input.ndim() - - for index, value in enumerate(dims): - assert -ndim <= value < ndim - if -ndim <= value < 0: - dims[index] += ndim - - assert len(dims) == len(set(dims)) - - start_values = [ - input.size()[i] - 1 if i in dims else 0 for i in range(ndim) - ] - stride_values = [-1 if i in dims else 1 for i in range(ndim)] - - layer = default_trtnet().add_slice(input.trt_tensor, - start=start_values, - shape=input.size(), - stride=stride_values) - - return _create_tensor(layer.get_output(0), layer) - - -def interpolate(input: Tensor, - size: Union[int, List[int]] = None, - scale_factor: Union[float, List[float]] = None, - mode: str = 'nearest', - align_corners: bool = False, - recompute_scale_factor: bool = False, - antialias: bool = False) -> Tensor: - ## - ## TODO: Document that function! - ## - - assert not input.is_dynamic() - - input_ndim = input.ndim() - - assert 2 < input_ndim < 6, "Only 3D, 4D and 5D input Tensors supported" - assert (size is not None) ^ ( - scale_factor - is not None), "Only one of out_shape or scales should be defined" - - assert mode in ('nearest', 'linear', 'bilinear', 'bicubic', 'trilinear', - 'nearest-exact') - - if mode == 'trilinear' and input_ndim != 5: - raise ValueError("trilinear only supports 5D tensor") - - if mode == "bilinear" and input_ndim != 4: - raise ValueError("bilinear only supports 4D tensor") - - if mode == "linear" and input_ndim != 3: - raise ValueError("linear only supports 3D tensor") - - layer = default_trtnet().add_resize(input.trt_tensor) - - input_shape = input.size() - - updated_shape = [] - if scale_factor: - scale_len = 1 if isinstance(scale_factor, - (float, int)) else len(scale_factor) - if scale_len == 1 and isinstance(scale_factor, (float, int)): - updated_scale = [scale_factor for _ in range(input_ndim - 2)] - - else: - updated_scale = scale_factor - updated_shape = [ - int(math.floor(updated_scale[i - 2] * - input_shape[i])) if i > 1 else input_shape[i] - for i in range(input_ndim) - ] - - else: - size_len = 1 if isinstance(size, int) else len(size) - assert size_len == input_ndim - 2 - if size_len == 1 and isinstance(size, int): - updated_size = [size for _ in range(input_ndim - 2)] - else: - updated_size = size - - updated_shape = [ - input_shape[i] if i < 2 else updated_size[i - 2] - for i in range(input_ndim) - ] - layer.shape = updated_shape - - if mode in ['nearest', 'nearest-exact'] or mode is None: - layer.resize_mode = trt.InterpolationMode.NEAREST - layer.coordinate_transformation = trt.ResizeCoordinateTransformation.ASYMMETRIC - elif mode in ['linear', 'bilinear', 'trilinear']: - layer.resize_mode = trt.InterpolationMode.LINEAR - if align_corners: - layer.coordinate_transformation = trt.ResizeCoordinateTransformation.ALIGN_CORNERS - else: - layer.coordinate_transformation = trt.ResizeCoordinateTransformation.HALF_PIXEL - # TODO, need to confirm the align_corners effect on bilinear mode. - if mode == 'bilinear': - layer.coordinate_transformation = trt.ResizeCoordinateTransformation.HALF_PIXEL - - elif mode in ['bicubic']: - layer.resize_mode = trt.InterpolationMode.CUBIC - - layer.coordinate_transformation = trt.ResizeCoordinateTransformation.HALF_PIXEL - - else: - layer.resize_mode = trt.InterpolationMode.NEAREST - layer.coordinate_transformation = trt.ResizeCoordinateTransformation.ASYMMETRIC - - return _create_tensor(layer.get_output(0), layer) - - -def matmul(input: Tensor, - mat2: Tensor, - transa: bool = False, - transb: bool = False, - use_fp32_acc: bool = True) -> Tensor: - ''' - Add a matrix multiplication. - - That operation maps to a tensorrt.IMatrixMultiplyLayer layer. As explained - in the TensorRT documentation, it computes the inner product between the - two inputs after applying an optional transposition on the inputs. - - Parameters: - input : Tensor - The first tensor (often called A). - - mat2 : Tensor - The second tensor (often called B). - - transa : bool - Is the first input transposed? Set to 'True' if you want the first - input to be transposed, 'False' otherwise. - - transb : bool - Is the second input transposed? Set to 'True' if you want the - second input to be transposed, 'False' otherwise. - - use_fp32_acc: bool - Set to 'True' if for accuracy reason, this fp16 matmul needs to use - fp32 accumulation. This can be a per model and per matmul decision. - Returns: - The tensor produced by the inserted layer. - ''' - # This option is only supported for fp16, but not bf16 or any other precisions. - use_fp32_acc = use_fp32_acc and input.dtype == trt.DataType.HALF and mat2.dtype == trt.DataType.HALF - - if use_fp32_acc: - input = cast(input, 'float32') - mat2 = cast(mat2, 'float32') - - input, mat2 = broadcast_helper(input, mat2) - op0 = trt.MatrixOperation.TRANSPOSE if transa \ - else trt.MatrixOperation.NONE - op1 = trt.MatrixOperation.TRANSPOSE if transb \ - else trt.MatrixOperation.NONE - layer = default_trtnet().add_matrix_multiply(input.trt_tensor, op0, - mat2.trt_tensor, op1) - output = _create_tensor(layer.get_output(0), layer) - if use_fp32_acc: - output = cast(output, "float16") - - return output - - -def gemm_swiglu(input: Tensor, - weight: Tensor, - bias: Optional[Tensor] = None, - scale_d0: float = 1.0, - scale_d1: float = 1.0, - scale_output: float = 1.0) -> Tensor: - ''' - Add a matrix multiplication, followed by SwiGLU (`x * SiLU(gate)`) operation. - - The second SwiGLU operation takes the preceding tensor, splits it into two halves - along the last dimension, applies SiLU to the second half and multiply the results. The - behaviour is undefined if the last dimension is not even. - - Parameters: - input : Tensor - The first tensor (often called A). - - weight : Tensor - The second tensor (often called B). - - bias : Optional[Tensor] - The per-channel bias. The plugin with fp8 dtype does not support bias yet. - - scale_d0 : float - The scale for dequantizing x, used for fp8 - - scale_d1 : float - The scale for dequantizing gate, used for fp8 - - scale_output : float - The scale for quantizing output, used for fp8 - - Returns: - The tensor produced by the inserted layer. - ''' - plg_creator = trt.get_plugin_registry().get_plugin_creator( - 'GemmSwiglu', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert plg_creator is not None - - p_dtype = default_net().plugin_config.gemm_swiglu_plugin - if p_dtype == "fp8": - assert bias is None, "fp8 gemm_swiglu does not support bias yet" - - pf_type = trt.PluginField( - "type_id", np.array([int(str_dtype_to_trt(p_dtype))], np.int32), - trt.PluginFieldType.INT32) - pf_has_bias = trt.PluginField( - "has_bias", np.array(np.int8(0 if bias is None else 1), np.int8), - trt.PluginFieldType.INT8) - pf_scale_d0 = trt.PluginField("scale_d0", - np.array(scale_d0, dtype=np.float32), - trt.PluginFieldType.FLOAT32) - pf_scale_d1 = trt.PluginField("scale_d1", - np.array(scale_d1, dtype=np.float32), - trt.PluginFieldType.FLOAT32) - pf_scale_output = trt.PluginField("scale_output", - np.array(scale_output, dtype=np.float32), - trt.PluginFieldType.FLOAT32) - - pfc = trt.PluginFieldCollection( - [pf_type, pf_has_bias, pf_scale_d0, pf_scale_d1, pf_scale_output]) - gemm_swiglu_plug = plg_creator.create_plugin("gemm_swiglu", pfc) - - # TODO(anchengc) pass nullptr when no bias - if bias is None: - bias = constant( - np.zeros([weight.shape[0]], dtype=trt_dtype_to_np(input.dtype))) - plug_inputs = [input.trt_tensor, weight.trt_tensor, bias.trt_tensor] - - layer = default_trtnet().add_plugin_v2(plug_inputs, gemm_swiglu_plug) - - return _create_tensor(layer.get_output(0), layer) - - -def constant(ndarray: np.ndarray, - as_dtype: trt.DataType | None = None, - as_shape=None) -> Tensor: - ''' - Add a constant layer. - - TensorRT graphs encapsulate constant values in the form of constant layers - (tensorrt.IConstantLayer). That function creates such a layer from a Numpy - array of values. After compilation of the network by TensorRT, those - weights are stored in the serialized TensorRT engine. - - Parameters: - ndarray : numpy.ndarray - The array of values (weights) encapsulated by this constant layer. - - Returns: - The tensor produced by the inserted layer. - ''' - trt_dtype = np_dtype_to_trt(ndarray.dtype) if as_dtype is None else as_dtype - trt_shape = trt.Dims( - ndarray.shape) if as_shape is None else trt.Dims(as_shape) - trt_count = 1 - for i in range(len(trt_shape)): - trt_count *= trt_shape[i] - weights = trt.Weights(trt_dtype, ndarray.ctypes.data, trt_count) - # Prevent underlying numpy array from going out of scope - default_net().register_ndarray(ndarray) - layer = default_trtnet().add_constant(trt_shape, weights) - if not default_net().strongly_typed: - layer.set_output_type(0, trt_dtype) - tensor = _create_tensor(layer.get_output(0), layer) - # TODO: remove this WAR after https://nvbugs/4359151 fixed. - set_np_weight(default_trtnet(), layer.name, ndarray) - return tensor - - -# TODO: TensorRT uses sizes of the output dimensions. -# DL framework uses ends usually. Will change it to ends. -def slice(input: Tensor, - starts: Union[Tensor, Sequence[int]], - sizes: Union[Tensor, Sequence[int]], - strides: Union[Tensor, Sequence[int]] = None, - mode: trt.SampleMode = None, - fill_value: Union[float, Tensor] = None) -> Tensor: - ''' - Add an operation to extract a slice from a tensor. - - As described in the TensorRT documentation of the ISliceLayer, the slice - layer has two variants: Static and dynamic. - - For static slicing, this function takes the starts and sizes values in the - different dimensions to slice at layer creation time via a sequence of - integers. For dynamic slicing, it accepts starts and sizes as - tensorrt.ITensor`s. - - The slice layer selects for each dimension a start location from within the - input tensor, and copies elements to the output tensor using a stride of 1 - across the input tensor. Start and size tensors must be 1-D int32 shape - tensors if not specified as a sequence of integers. - - As an example, on input = [[0, 2, 4], [1, 3, 5]], the call to - - slice(input, start=[1, 0], size=[1, 2]) - - will produce the tensor [[1, 3]] as output. The slice operator when - executed by TensorRT will copy one row (because size[0] == 1) starting from - the 2nd row (because start[0] == 1) and two columns (size[1] == 2) starting - from the 1st column (because start[1] == 0). - - In pseudo-code the behavior of that operation can be described as follows - for a 2D tensor (and easily be extended to more dimensions): - - output = Tensor(shape=sizes) - for ii in range(sizes[0]): - for jj in range(sizes[1]): - output[ii][jj] = input[starts[0]+ii][starts[1]+jj] - - Note that it is common in deep-learning frameworks to use ranges - [start:end] for similar operations. It can be emulated by setting the sizes - argument such that in each dimension [start:start+size] == [start:end] i.e. - size = end-start. - - TensorRT supports different slice modes but that function restricts that - choice to `mode == tensorrt.SampleMode.STRICT_BOUNDS`. - - Parameters: - input : Tensor - The input tensor on which the slicing is performed. - - starts : Union[Tensor, Sequence[int]] - The starting points, in the input tensor, and each dimension. - - sizes : Union[Tensor, Sequence[int]] - The number of elements in each dimension of the sliced tensor (output). - - strides : Union[Tensor, Sequence[int]] - The step be taken from start, in input tensor. - - mode : trt.SampleMode - The mode that controls how the slice operation handles out of bounds coordinates. - - Returns: - The tensor produced by the slice layer. - ''' - input_ndim = input.ndim() - - trt_starts = starts - if isinstance(starts, Tensor): - trt_starts = [0 for _ in range(input_ndim)] # unused dummy value - - trt_sizes = sizes - if isinstance(sizes, Tensor): - trt_sizes = [1 for _ in range(input_ndim)] # unused dummy value - - trt_strides = strides - if isinstance(strides, Tensor) or strides is None: - trt_strides = [1 for _ in range(input_ndim)] - - if fill_value is not None and isinstance(fill_value, float): - fill_value = constant(fp32_array(fill_value)) - - layer = default_trtnet().add_slice(input.trt_tensor, - start=trt_starts, - shape=trt_sizes, - stride=trt_strides) - if mode is not None: - layer.mode = mode - - if isinstance(starts, Tensor): - layer.set_input(1, starts.trt_tensor) - - if isinstance(sizes, Tensor): - layer.set_input(2, sizes.trt_tensor) - - if isinstance(strides, Tensor): - layer.set_input(3, strides.trt_tensor) - - if mode is trt.SampleMode.FILL and isinstance(fill_value, Tensor): - layer.set_input(4, fill_value.trt_tensor) - - return _create_tensor(layer.get_output(0), layer) - - -def pad(input: Tensor, - pad: Union[Sequence[int], Tensor], - mode: str = 'constant', - value: Optional[float] = None) -> Tensor: - ''' - Add a pad layer. - - The padding layer adds zero-padding at the start and end of the input tensor. And the - padding size by which to pad some dimensions of input are described starting from the - last dimension and moving forward. - - `[len(pad) / 2]` dimensions of input will be padded. For example, to pad only the last - dimension of the input tensor, then pad has the form [padding_left, padding_right]; to - pad the last 2 dimensions of the input tensor, then use [padding_left, padding_right, - padding_top, padding_bottom]; to pad the last 3 dimensions, use [padding_left, - padding_right, padding_top, padding_bottom, padding_front, padding_back]. - - Parameters: - input : Tensor - The input tensor on which the padding_2d is performed. - pad : sequence of int - An m-elements tuple for padding, where its length m meets the requirement that - m <= 2*input dimensions, and m is even. - mode : str - Only \'constant\' is supported. - value : float - Fill value for 'constant' padding. Default: 0. - - Returns: - The tensor produced by the inserted layer. - ''' - assert mode == "constant", "Only `'constant'` is supported now." - - if isinstance(pad, list) or isinstance(pad, tuple): - assert ( - len(pad) % 2 == 0 and len(pad) <= 2 * input.ndim() - ), "The length of `pad` should be even and less than 2*input.ndim" - pad = constant(np.array(pad).astype(np.int32)).view([-1, 2]) - elif isinstance(pad, Tensor): - pad = pad.flatten() - assert ( - pad.size(0) % 2 == 0 and pad.size(0) <= 2 * input.ndim() - ), "The length of `pad` should be even and less than 2*input.ndim" - pad = pad.cast("int32").view([-1, 2]) - else: - raise NotImplementedError(f"pad type {type(pad)} not supported") - if value is None: - value = 0 - - pad = concat([constant(np.zeros((1, 2), dtype=np.int32)), - pad]) # pre-padding the indices - padding_index = [0] * input.ndim() - padding_index[-(pad.size(0) - 1):] = list(range(pad.size(0) - 1, 0, - -1)) # reverse the indices - pad = index_select(pad, - dim=0, - index=constant(np.array(padding_index, dtype=np.int32))) - pre_padding, post_padding = chunk(pad, chunks=2, dim=1) - start = (pre_padding.flatten() * (-1)).cast('int32') - extend_size = (pre_padding + post_padding).flatten() - size = (extend_size + shape(input)).cast('int32') - layer = default_trtnet().add_slice(input.trt_tensor, - start=[0] * input.ndim(), - shape=[0] * input.ndim(), - stride=[1] * input.ndim()) - layer.mode = trt.SampleMode.FILL - layer.set_input(SliceInputType.start, start.trt_tensor) - layer.set_input(SliceInputType.size, size.trt_tensor) - layer.set_input(SliceInputType.fill_value, - constant_to_tensor_(value, dtype=input.dtype).trt_tensor) - output = _create_tensor(layer.get_output(0), layer) - return output - - -def rand(shape: Tensor, - low: float = 0, - high: float = 1, - dtype: Union[str, trt.DataType] = 'float32') -> Tensor: - ''' - This operation adds a fill layer that generates a random (uniform) tensor with the specified shape and data type. - - Parameters: - shape: Tensor - The shape of the tensor needed to be generated. - low: float - The minimum value (inclusive) of the range used for random. - high: float - The maximum value (inclusive) of the range used for random. - dtype: Union[str, trt.DataType] - The desired data type for the output tensor. - Returns: - The generated random tensor produced by the fill layer. - ''' - # NOTE: DISABLED FOR NOW UNTIL THE FILL LAYER (RANDOM_UNIFORM) in TRT IS FIXED - assert False, "The rand() op is temporarily disabled." - low = constant(fp32_array(low)) - high = constant(fp32_array(high)) - trt_dtype = dtype if isinstance(dtype, - trt.DataType) else str_dtype_to_trt(dtype) - - layer = default_trtnet().add_fill([0], trt.FillOperation.RANDOM_UNIFORM, - trt_dtype) - - layer.set_input(0, shape.trt_tensor) - layer.set_input(1, low.trt_tensor) - layer.set_input(2, high.trt_tensor) - return _create_tensor(layer.get_output(0), layer) - - -def categorical_sample(probs: Tensor, rand_data: Tensor = None) -> Tensor: - ''' - This is a sampling operation and an equivalent of torch.distributions.Categorical.sample() - i.e. given a probability distribution tensor, it samples an index of that tensor. - See: https://pytorch.org/docs/stable/distributions.html#torch.distributions.categorical.Categorical.sample - NOTE: This assumes that the given probabilities are **not** normalized. - - Parameters: - probs: Tensor - A 1-D floating point tensor representing the probability distributions. - rand_data: Tensor (optional) - A random tensor of same shape as `probs` tensor. - If not provided, this function will add a rand() op to generate it and use for sampling. - Returns: - A tensor containing a single index of the `probs` tensor representing the sample. - ''' - probs = probs / sum(probs, dim=-1, keepdim=True) - rand_shape = [] - assert probs.ndim() > 0 - for i in range(probs.ndim() - 1): - rand_shape.append(shape(probs, i)) - rand_shape = concat(rand_shape) - if rand_data is None: - rand_data = rand(rand_shape, low=0, high=1, dtype=probs.dtype) - assert rand_shape == shape(rand_data) - rand_data = expand(unsqueeze(rand_data, -1), shape(probs)) - cum_probs = cumsum(probs, dim=-1) - cmp = cast(cum_probs >= rand_data, probs.dtype) - samples = argmax(cmp, dim=-1) - return samples - - -class Conditional: - ''' - Add an operation to conditionally execute two code paths/subgraphs. - - Usage: - 1. conditional = Conditional(condition) - 2. input_1_ = conditional.add_input(input_1) - ... - input_n_ = conditional.add_input(input_n) - 3. Construct the graph to get true_output_value and false_output_value using input_1_, ..., input_n_ - 4. output = conditional.add_output(true_output_value, false_output_value) - ''' - - def __init__(self, condition: Tensor): - self.layer = default_trtnet().add_if_conditional() - if condition.ndim() > 0: - condition = view(condition, []) - self.layer.set_condition(condition.trt_tensor) - - def add_input(self, input: Tensor) -> Tensor: - in_node = self.layer.add_input(input.trt_tensor) - return _create_tensor(in_node.get_output(0), in_node) - - def add_output(self, true_value: Tensor, false_value: Tensor) -> Tensor: - out_node = self.layer.add_output(true_value.trt_tensor, - false_value.trt_tensor) - return _create_tensor(out_node.get_output(0), out_node) - - -# TODO: support step. -def arange(start: Union[Tensor, int], end: Union[Tensor, int], - dtype: str) -> Tensor: - ''' - Add an operation to fill a 1D tensor. - - The tensor is filled with the values between start and end with a step of 1 - between the different elements. In pseudo-code, it corresponds to a tensor - populated with the values: - - output = Tensor([dtype(ii) for ii in range(start, end, 1)]) - - For example, a call to arange(3, 6, 'int32') will add an operation to the - TensorRT graph that will produce [3, 4, 5] when executed. The call to - arange(2, 5, 'float32') will add a layer to generate [2.0, 3.0, 4.0]. - - This operation is implemented using a tensorrt.IFillLayer in - trt.FillOperation.LINSPACE mode. - - Parameters: - start : Union[Tensor, int] - The starting point of the range. - - end : Union[Tensor, int] - The end point of the range. - - dtype : str - The type of the elements. See _str_to_trt_dtype_dict in _utils.py - for a list of supported types and type names. - - Returns: - The tensor produced by the fill layer. It is a 1D tensor containing - `end-start` elements of type `dtype`. - ''' - res_dtype = str_dtype_to_trt(dtype) - if isinstance(start, int): - assert isinstance(end, int) - array_func = int32_array if res_dtype == trt.int32 else int64_array - start = constant(array_func(start)) - end = constant(array_func(end)) - elif isinstance(start, Tensor): - assert isinstance(end, Tensor) - assert start.dtype == trt.int32 or start.dtype == trt.int64 - assert end.dtype == trt.int32 or end.dtype == trt.int64 - if start.dtype != end.dtype: - if start.dtype == trt.int32: # end == trt.int64 - if res_dtype == trt.int32: - end = cast(end, "int32") - else: - start = cast(start, "int64") - else: # start == trt.int64 and end == trt.int32 - if res_dtype == trt.int32: - start = cast(start, "int32") - else: - end = cast(end, "int64") - else: - raise TypeError("%s is not supported" % type(start)) - - assert start.dtype == end.dtype, f"start type ({start.dtype}) != end type ({end.dtype})" - step = constant_to_tensor_(1, dtype=start.dtype, to_array=True) - - num = end - start - num = num.view([1]).cast(trt.int64) - - layer = default_trtnet().add_fill([0], trt.FillOperation.LINSPACE, - start.dtype) - layer.set_input(0, num.trt_tensor) # rank = 1 - layer.set_input(1, start.trt_tensor) # rank = 0 - layer.set_input(2, step.trt_tensor) # rank = 1 - tensor = _create_tensor(layer.get_output(0), layer) - if tensor.dtype != res_dtype: - tensor = tensor.cast(dtype) - return tensor - - -def expand(input: Tensor, expand_shape: Tensor) -> Tensor: - ''' - Add an operation to expand a tensor. - - The operation expands the input tensor in the singleton dimensions to the - size indicated by the corresponding dimension in the `expand_shape` tensor. - In other words, given an input tensor with dimensions of size 1, those - dimensions will be expanded to the size in `expand_shape`. - - For example, a tensor of shape [4, 3, 1, 3] will be expanded to a tensor of - shape [4, 3, 2, 3] by the layer created using expand(input, [4, 3, 2, 3]). - - The expansion may either replicate the values or be mapped to a view with a - stride of 0 in the expanded dimensions. For example, for a tensor [[3, 2]] of - shape [1, 2], - - expand([[3, 2]], [2, 2]) - - can be used to expand the input to [[3, 2], [3, 2]]. - - This operation is implemented using a tensorrt.ISliceLayer. The current - implementation does not verify that non singleton dimensions are not - shrunk. In other words, for an input of shape [4, 1, 2], - - expand(input, [3, 2, 2]) - - will produce a tensor of shape [3, 2, 2]. That behavior is subject to - change in the future. - - Parameters: - input : Tensor - The input tensor. - - expand_shape : Tensor - The new shape of the expanded tensor. - - Returns: - The tensor produced by the expand layer. - ''' - ndim = input.rank() - layer = default_trtnet().add_slice( - input.trt_tensor, - start=[0 for _ in range(ndim)], - shape=[1 for _ in range(ndim)], # unused dummy value - stride=[1 for _ in range(ndim)] # unused dummy value - ) - - # The stride is either: - # 0 for dimensions of size 1 (i.e. shape(input, i) - 1 == 1 - 1 == 0) or, - # 1 for dimensions of size > 1 since minimum(value >= 1, 1) == 1. - stride_tensor = concat( - [minimum((shape(input, i) - 1), 1) for i in range(ndim)]) - - layer.set_input(2, expand_shape.trt_tensor) - layer.set_input(3, stride_tensor.trt_tensor) - return _create_tensor(layer.get_output(0), layer) - - -def einsum(einsum_eq: str, inputs: Sequence[Tensor]) -> Tensor: - ''' - Add an Einsum operation. - - That operation maps to tensorrt.IEinsumLayer. As explained in the TensorRT - documentation, this layer implements a summation over the elements of the - inputs along dimensions specified by the equation parameter, based on the - Einstein summation convention. The layer can have one or more inputs of - rank >= 0. All the inputs must be of same data type. This layer supports - all TensorRT data types except bool. There is one output tensor of the same - type as the input tensors. The shape of output tensor is determined by the - equation. - - The equation specifies ASCII lower-case letters for each dimension in the - inputs in the same order as the dimensions, separated by comma for each - input. The dimensions labeled with the same subscript must match or be - able to be broadcasted. Repeated subscript labels in one input take the diagonal. - Repeating a label across multiple inputs means that those axes will be - multiplied. Omitting a label from the output means values along those axes - will be summed. In implicit mode, the indices which appear once in the - expression will be part of the output in increasing alphabetical order. In - explicit mode, the output can be controlled by specifying output subscript - labels by adding an arrow (‘->’) followed by subscripts for the output. For - example, “ij,jk->ik” is equivalent to “ij,jk”. Ellipsis (‘…’) can be used - in place of subscripts to broadcast the dimensions. See the TensorRT - Developer Guide for more details on equation syntax. - - Many common operations can be expressed using the Einsum equation. For - example: - Matrix Transpose: ij->ji - Sum: ij-> Matrix-Matrix - Multiplication: ik,kj->ij - Dot Product: i,i-> - Matrix-Vector Multiplication: ik,k->i - Batch Matrix Multiplication: ijk,ikl->ijl - Batch Diagonal: …ii->…i - - Note that TensorRT does not support ellipsis or diagonal operations so, - neither, does TensorRT-LLM. - - Parameters: - einsum_eq : str - The Einsum equation. - - inputs: Sequence[Tensor] - The sequence of inputs consumed by the Einsum operation. - - Returns: - The tensor produced by the Einsum operation. - ''' - layer = default_trtnet().add_einsum([i.trt_tensor for i in inputs], - einsum_eq) - return _create_tensor(layer.get_output(0), layer) - - -def permute(input: Tensor, dims: Sequence[int]) -> Tensor: - ''' - Add an operation to permute the dimensions of a tensor. - - The dimensions of the input tensor are permuted according to the sequence - of dimensions in 'dims'. That operation maps to tensorrt.IShuffleLayer where - the second transposition is described by the indices in 'dims'. - - Given a tensor of rank N, the result of the permutation is a tensor of rank - N in which the i-th input dimension maps to the dims[i]-th dimension. - - For example, permute(input, [1, 0]) will transpose a 2D tensor by permuting - the rows and columns. - - Parameters: - input : Tensor - The input tensor to permute. - - dims : Sequence[int] - The description of the permutation. - - Returns: - The tensor produced by the permutation layer. - ''' - dims = dim_resolve_negative(tuple(dims), input.ndim()) - layer = default_trtnet().add_shuffle(input.trt_tensor) - layer.second_transpose = dims - return _create_tensor(layer.get_output(0), layer) - - -def transpose(input: Tensor, dim0: int, dim1: int) -> Tensor: - ''' - Add an operation to transpose two dimensions of a tensor. - - That operation produces a tensor in which the dimensions 'dim0' and 'dim1' - are permuted. The other dimensions, if the rank of the tensor is greater - than 2, remain untouched. - - That function is a helper built on the 'functional.permute' function. - - Parameters: - input : Tensor - The input tensor to transpose. - - dim0 : int - The first dimension to transpose. - - dim1 : int - The second dimension to transpose. - - Returns: - The tensor produced by the permutation layer. - ''' - permutation = list(range(input.ndim())) - permutation[dim0] = dim1 - permutation[dim1] = dim0 - - return permute(input, permutation) - - -def view(input: Tensor, - shape: Union[Tensor, Sequence[int]], - zero_is_placeholder: bool = True) -> Tensor: - ''' - Add an operation to create a view of a tensor. - - That operation adds a tensorrt.IShuffleLayer to the network. If the 'shape' - parameter is a Tensor, that view is dynamic. Otherwise, it is a static - view. - - Note that TensorRT limits the number of inferred dimensions to 1. It means - that the shape sequence or tensor cannot contain more than one -1. This - function enforces that constraint and will assert if it is not respected. - - Parameters: - input : Tensor - The input tensor to transpose. - - shape : Union[Tensor, Sequence[int]] - The shape of the new tensor. - - zero_is_placeholder : bool - When that parameter is True, the 0s in 'shape' are replaced by the - sizes of the corresponding dimensions from the 'input'. Otherwise, - the dimensions corresponding to 0s are shrunk. - - Returns: - The tensor produced by the view/shuffle layer. - ''' - - # TensorRT demands that at most one dimension is permitted to be specified as -1 - def assert_no_more_than_one_inferred_dim(list): - inferred_dim_list = [i for i in list if i == -1] - assert len(inferred_dim_list) <= 1 - - layer = default_trtnet().add_shuffle(input.trt_tensor) - layer.zero_is_placeholder = zero_is_placeholder - if isinstance(shape, Tensor): - assert_no_more_than_one_inferred_dim(shape.shape) - layer.set_input(1, shape.trt_tensor) - elif isinstance(shape, (list, tuple)): - assert_no_more_than_one_inferred_dim(shape) - layer.reshape_dims = tuple(shape) - else: - raise TypeError("%s is not supported" % type(shape)) - return _create_tensor(layer.get_output(0), layer) - - -def flatten(input: Tensor, start_dim: int = 0, end_dim: int = -1): - ''' - Flattens input by reshaping it into a one-dimensional tensor. - - If start_dim or end_dim are passed, only dimensions starting with start_dim and - ending with end_dim are flattened. The order of elements in input is unchanged. - - Parameters: - input : Tensor - The input tensor to flatten. - - start_dim : int - The first dim to flatten. - - end_dim : int - The last dim to flatten. - - Returns: - The tensor produced by the flatten layer. - - ''' - shape = input.shape - ndim = input.ndim() - if start_dim < 0: start_dim += ndim - if end_dim < 0: end_dim += ndim - new_shape = list() - for i in range(start_dim): - new_shape.append(shape[i]) - if end_dim - start_dim >= 0: - flat_dim = 1 - for i in range(start_dim, end_dim + 1): - flat_dim *= shape[i] - new_shape.append(flat_dim) - for i in range(end_dim + 1, ndim): - new_shape.append(shape[i]) - return view(input, new_shape) - - -def expand_dims(input: Tensor, - dim: Union[int, Sequence[int]], - shape_cast_dtype=None) -> Tensor: - ''' - Add an operation to expand the tensor shape with singleton dimensions. - - That function adds a tensorrt.IShuffleLayer to the network. Given an 'input' - of rank N and a sequence of M dimensions, the output tensor produced by - this operation (when executed by TensorRT) will have a rank of N+M. Singleton - dimensions will be inserted at the different positions in 'dim'. - - The pseudo-code for that operation is: - - new_shape, ii = [], 0 - for jj in range(input.rank() + len(dim)): - new_shape.append(1 if jj in dims else input.shape[ii++]) - - For example, for a tensor of shape [3, 4, 1, 5] - - expand_dims(input, [0, 2]) - - will produce a tensor of shape [1, 3, 1, 4, 1, 5]. - - Parameters: - input : Tensor - The input tensor to expand. - - dim : Union[int, Sequence[int]] - The positions in the output tensor where to insert singleton - dimensions. - - Returns: - The tensor produced by the shuffle layer. - ''' - if isinstance(dim, int): - dim = (dim, ) - - out_ndim = len(dim) + input.ndim() - - input_shape = shape(input, cast_to_dtype=shape_cast_dtype) - out_shapes = [] - j = 0 - for i in range(out_ndim): - if i in dim: - out_shapes.append(1) - else: - out_shapes.append(gather(input_shape, 0, j)) - j = j + 1 - - out_shape = concat(out_shapes) - - return view(input, out_shape, zero_is_placeholder=False) - - -# NOTE: Jointly added with Apple -def squeeze(input: Tensor, - dim: Optional[Union[int, Sequence[int]]] = None, - zero_is_placeholder: bool = False): - ''' - Add an operation to remove singleton dimensions of a tensor. - - This functions creates an operation that removes singleton dimension - (dimension of size 1) at positions 'dim' in the input tensor. It works with - negative values for the 'dim'. - - For example, for a tensor 'input' of shape [1, 4, 1, 4]: - - squeeze(input, 0) will produce an output of shape [4, 1, 4], - squeeze(input, 2) will produce an output of shape [1, 4, 4], - squeeze(input, [0, 2]) will produce an output of shape [4, 4], - squeeze(input, [-2]) will produce an output of shape [1, 4, 4], - - Parameters: - input : Tensor - The input tensor for which the singleton dimensions will be removed. - - dim : Union[int, Sequence[int]] - The index of the singleton dimensions in the input tensor. - - Returns: - The tensor produced by the layer. - ''' - if dim is None: - dim = list(range(input.ndim())) - if isinstance(dim, int): - dim = (dim, ) - dim = dim_resolve_negative(dim, input.ndim()) - - new_shape = [] - for i, s in enumerate(input.shape): - if s == 1 and i in dim: - continue - new_shape.append(shape(input, i)) - - new_shape = concat(new_shape) if len(new_shape) > 0 else [] - input = input.view(new_shape, zero_is_placeholder=zero_is_placeholder) - return input - - -def unsqueeze(input: Tensor, axis: int): - ''' - Add an operation to insert a singleton dimension to a tensor. - - That functions creates an operation that insert a singleton dimension - (dimension of size 1) at position 'axis' in the output tensor. It works with - negative values for the 'axis'. - - For example, for a tensor 'input' of shape [4, 4]: - - unsqueeze(input, 0) will produce an output of shape [1, 4, 4], - unsqueeze(input, 1) will produce an output of shape [4, 1, 4], - unsqueeze(input, -1) will produce an output of shape [4, 4, 1], - unsqueeze(input, -2) will produce an output of shape [4, 1, 4], - - Parameters: - input : Tensor - The input tensor to expand with a singleton dimension. - - axis : int - The index of the singleton dimension in the output tensor. - - Returns: - The tensor produced by the layer. - ''' - if axis < 0: - axis = axis + input.ndim() + 1 - - return expand_dims(input, axis) - - -def stack(inputs: Sequence[Tensor], dim: int = 0) -> Tensor: - ''' - Add an operation to contact input tensors along a new dimension. - - The function creates an operation that creates a new dim for all the - input tensors and then concatenates them along that new dim. -. - - All the tensors in 'inputs' must have the same shape. - - for ii in range(inputs[0].rank()): - assert all(inp.shape[ii] == inputs[0].shape[ii] for inp in inputs) - - The shape of the output tensor is defined as: - - output.rank() = inputs[0].rank() + 1 - - output.shape[dim] = len(inputs) - - for ii in range(inputs[0].rank()): - if ii < dim: - output.shape[ii] = inputs[0].shape[ii] - else: - output.shape[ii+1] = inputs[0].shape[ii] - - For example, given a sequence of two 2D tensors [[0, 1], [2, 3]] and - [[4, 5], [6, 7]] both of shape [2, 2], - - stack(inputs, 0) - - will produce [[[0, 1], [2, 3]], [[4, 5], [6, 7]]] of shape [2, 2, 2] and - - stack(inputs, 1) - - will produce [[[0, 1], [4, 5]], [[2, 3], [6, 7]]] of shape [2, 2, 2]. - - Parameters: - inputs : Sequence[Tensor] - The sequence of tensors to stack. - - dim : int - The dimension in which the stack is performed. - - Returns: - A tensor that contains the input tensors stacked along a new dimension. - ''' - return concat([unsqueeze(inp, axis=dim) for inp in inputs], dim=dim) - - -def expand_dims_like(left: Union[Tensor, int, float], right: Tensor) -> Tensor: - ''' - Add an operation to expand the first tensor to the same rank as the second - tensor. - - That function takes a first tensor. It also accepts an integer or a float, - in which case it creates a constant tensor from it. In both cases, the rank - of that first tensor is compared to the rank of the second tensor. If they - are of the same rank, the first tensor is returned. Otherwise, the first - tensor is expanded on the left to match the rank of the second tensor. - - Note that the shapes do not have to match, only the rank is considered in - that function. - - For example, for a pair of tensors of shapes [3, 4] and [4, 3, 2], the - first tensor will be expanded to a tensor of rank 3 and shape [1, 3, 4]. - - Parameters: - left : Union[Tensor, int, float] - The first tensor to expand. When a scalar value is provided as a - parameter, that function first creates a tensor before expanding it - (if needed). - - right : Tensor - The reference tensor to match. - - Returns: - The tensor produced by the shuffle layer. - ''' - if isinstance(left, int): - left = constant(dims_array([left])) - elif isinstance(left, float): - if isinstance(right, Tensor) and right.dtype == trt.DataType.HALF: - left = constant(fp16_array([left])) - else: - left = constant(fp32_array([left])) - left_ndim = left.ndim() - right_ndim = right.ndim() - if right_ndim > left_ndim: - new_ndim = list(range(right_ndim - left_ndim)) - return expand_dims(left, new_ndim) - return left - - -# If dim is None, return a 1-D TensorRT LLM tensor of the size -# If dim is not None, return a 0-D TensorRT LLM tensor of the dimension size -def shape(input: Tensor, - dim: Optional[int] = None, - cast_to_dtype: Optional[Union[str, trt.DataType]] = None, - clip_before_cast: Sequence[int] = None) -> Tensor: - ''' - Add an operation to create a shape tensor. - - The shape tensor can either be the shape of the input tensor when the - parameter dim is None or a scalar (tensor of rank 0) that corresponds to - the size of dim-th dimension. - - Parameters: - input : Tensor - The input tensor from which we want to extract the shape or the - size in one dimension. - - dim : Optional[int] - The dimension from which to extract the size. If it is None, the - entire shape of the input tensor is returned. - - Returns: - A tensor that contains the shape of the input tensor (if 'dim' is None) - or the size in the dimension 'dim' of the input tensor. If 'dim' is - 'None', that tensor has the same rank as the input tensor, otherwise - its rank is 0. - ''' - layer = default_trtnet().add_shape(input.trt_tensor) - res = _create_tensor(layer.get_output(0), layer) - if cast_to_dtype is not None: - if clip_before_cast is not None and (cast_to_dtype == 'int32' - or cast_to_dtype == trt.int32): - assert len( - clip_before_cast - ) == 2, f"This parameter only expects a tuple of 2 integers (lower, upper) but got {clip_before_cast}" - res = int_clip(res, clip_before_cast[0], clip_before_cast[1]) - res = cast(res, cast_to_dtype) - - if dim is None: - return res - - return gather(res, dim=0, indices=dim).view([]) - - -def gather(input: Tensor, dim: int, indices: Union[Tensor, int]) -> Tensor: - ''' - Add an operation to gather elements from a tensor. - - That function implements the GatherElements operator from the ONNX - specification as described in - - https://github.com/onnx/onnx/blob/main/docs/Operators.md#GatherElements - - The input and indices arguments must have the same rank >= 1. The operation - will produce a tensor with the same shape as the indices tensor. The axis - is the dimension to gather on. - - As shown in the ONNX description, for a 3D tensor, the output is: - - out[i][j][k] = input[indices[i][j][k]][j][k] if axis = 0, - out[i][j][k] = input[i][indices[i][j][k]][k] if axis = 1, - out[i][j][k] = input[i][j][indices[i][j][k]] if axis = 2. - - For example, - - gather([[4, 2], [5, 3]], 0, [[1, 0], [0, 1]]) - - will produce [[5, 2], [4, 3]]. - - gather([[1, 2, 3], [4, 5, 6], 1, [[1], [0]]) - - will produce [[2], [4]]. See the ONNX documentation for more examples. - - That operation maps to the TensorRT IGatherLayer. - - Parameters: - input : Tensor - The input tensor to gather elements from. - - dim : int - The dimension to gather on. - - indices : Union[Tensor, int] - The positions in the 'dim' dimension to gather from. - - Returns: - The tensor containing the gathered elements. It has the same shape as - the indices tensor. - ''' - if isinstance(indices, int): - indices = constant(int32_array([indices])) - - # The input and indices tensors must have the same rank. - assert input.rank() == indices.rank() - - layer = default_trtnet().add_gather_v2(input.trt_tensor, - indices.trt_tensor, - mode=trt.GatherMode.ELEMENT) - - if dim < 0: - dim = input.ndim() + dim - layer.axis = dim - return _create_tensor(layer.get_output(0), layer) - - -def select(input: Tensor, dim: int, index: Union[Tensor, int]) -> Tensor: - ''' - Add an operation to select a slice of elements from a tensor. - - Given an input tensor, that function creates an operation that selects the - index-th slice of elements in the dimension 'dim' to create a new tensor. - The output tensor has a shape in which the input dimension 'dim' is - removed. - - The 'index' can either be an integer or a 1D tensor containing a single - element. - - For example, on input=[[4, 2, 5], [2, 1, 2], [4, 7, 1]], which has a shape - [3, 3], - - select(input, 0, 1) - - will create a tensor of shape [3] that contains the [2, 1, 2]. - - Regarding the shape of the output tensor, the dimension 'dim' is removed. - It means that for a tensor of shape [4, 2, 6, 3], - - select(input, 2, 4) - - will select the 5th slice (index == 4) from the 3rd dimension (dim == 2) - and return a tensor of shape [4, 2, 3] (i.e. the 3rd dimension is removed). - - That operation maps to the TensorRT IGatherLayer. - - Parameters: - input : Tensor - The input tensor to select from. - - dim : int - The dimension to select from. - - index : Union[Tensor, int] - The index of the slice in the 'dim' dimension to select. - - Returns: - The tensor containing the selected slice. - ''' - if isinstance(index, int): - index = constant(int32_array([index])) - assert index.rank() == 1 and index.size( - 0) == 1, f"index should have rank 1, got {index.rank()}" - - new_shape = [] - for i in range(input.rank()): - if i != dim: - new_shape.append(shape(input, i)) - - layer = default_trtnet().add_gather(input.trt_tensor, index.trt_tensor, dim) - return _create_tensor(layer.get_output(0), layer).view(concat(new_shape)) - - -def index_select(input: Tensor, dim: int, index: Tensor) -> Tensor: - ''' - Add an operation to select slices of elements from a tensor. - - Given an input tensor, that function creates an operation that selects the - slices of elements in the dimension 'dim' at the indices listed in 'index' - to create a new tensor. The output tensor has the same rank as the input - tensor. - - The 'index' is a tensor of rank 1. - - For example, on input=[[4, 2, 5], [2, 1, 2], [4, 7, 1]], which has a shape - [3, 3], - - index_select(input, 0, [0, 1]) - - will create a tensor of shape [2, 3] that contains the [[4, 2, 5], [2, 1, 2]]. - - Regarding the shape of the output tensor, the dimension 'dim' has the same - size as the 'index' tensor. It means that for a input tensor of shape [4, 2, 6, 3], - - index_select(input, 2, [1, 4]) - - will select the 2nd and 5th slices (index == 1 or 4) from the 3rd dimension - (dim == 2) and return a tensor of shape [4, 2, 2, 3] (i.e. the 3rd - dimension is shrunk to 2). - - Note that this operation can also be used to expand a tensor in the 'dim' - dimension, for example, on input [[0, 1], [2, 3]], - - index_select(input, 1, [0, 0, 0]) - - will produce a tensor of shape [2, 3] containing [[0, 0, 0], [2, 2, 2]]. - - That operation maps to the TensorRT IGatherLayer. - - Parameters: - input : Tensor - The input tensor to select from. - - dim : int - The dimension to select from. - - index : Tensor - The indices of the slices in the 'dim' dimension to select. - - Returns: - The tensor containing the selected slices. - ''' - assert index.rank() == 1, f"index should have rank 1, got {index.rank()}" - - new_shape = [] - for i in range(input.rank()): - if i != dim: - new_shape.append(shape(input, i)) - else: - new_shape.append(shape(index, 0)) - - layer = default_trtnet().add_gather(input.trt_tensor, index.trt_tensor, dim) - return _create_tensor(layer.get_output(0), layer).view(concat(new_shape)) - - -# NOTE: Jointly added with Apple -def scatter(input: Tensor, dim: int, indices: Tensor, - updates: Tensor) -> Tensor: - ''' - This operation adds a layer that creates an output tensor by element-wise - copying values from the input tensor and then updating values by the given - `indices` and `updates` tensors. - For a 2D input tensor, it first copies the input to output, - then updates the output tensor like the following for each entry in `updates`: - output[indices[i][j]][j] = updates[i][j] if dim=0 - output[i][indices[i][j]] = updates[i][j] if dim=1 - If the `input` tensor is [[1, 2, 3], [4, 5, 6]], - the indices tensor is [[1, 2], [0, 1]], - the updates tensor is [[-1, -2], [-3, -4]], and dim=1 - the output tensor will be [[1, -1, -2], [-3, -4, 6]]. - Parameters: - input: Tensor - The input data that needs to be updated. - dim: int - The axis on which the scatter is to be performed. - indices: Tensor - An integer tensor of the same rank as input that indicates the positions to be updated. - updates: Tensor - A data tensor of same shape as the `indices` tensor that contains the update values. - Returns: - A tensor created by the element-wise scatter layer. - ''' - layer = default_trtnet().add_scatter(input.trt_tensor, - indices.trt_tensor, - updates.trt_tensor, - mode=trt.ScatterMode.ELEMENT) - layer.axis = dim - return _create_tensor(layer.get_output(0), layer) - - -def gather_nd(input: Tensor, indices: Tensor, batch_dims: int = 1) -> Tensor: - ''' - Adds a layer that performs a gather with some element-wise dimensions. - See: https://onnx.ai/onnx/operators/onnx__GatherND.html - The gather is performed on dim=batch_dims. - - Parameters: - input: Tensor - The tensor on which the gather operation is performed. - indices: Tensor - The tensor that indicates which entries to be gathered. - batch_dims: int - The number of first dimensions that should be skipped before gather starts. - Returns: - A tensor created by the gather layer with GatherMode.ND. - ''' - gather_layer = default_trtnet().add_gather_v2(input.trt_tensor, - indices.trt_tensor, - mode=trt.GatherMode.ND) - gather_layer.num_elementwise_dims = batch_dims - return _create_tensor(gather_layer.get_output(0), gather_layer) - - -def nonzero(input: Tensor) -> Tensor: - ''' - Adds a layer that finds the indices of non-zero values of the input tensor. - - Parameters: - input: Tensor - The input tensor for which we need to find the indices of non-zero values. - Returns: - A tensor of shape [D, C] where D is the number of dimensions of `input` and - C is the number of non-zero values in it. - Each column of this 2D tensor represents the index tuple for each non-zero value. - ''' - non_zero_layer = default_trtnet().add_non_zero(input.trt_tensor) - return _create_tensor(non_zero_layer.get_output(0), non_zero_layer) - - -def masked_select(input: Tensor, mask: Tensor) -> Tensor: - ''' - Add an operation to select elements from a tensor according to a boolean - mask tensor. - - Given an input tensor, that function creates an operation that selects - elements at the indices indicated by the boolean mask tensor to create - a new tensor. The output tensor is a 1-D tensor. - - The input tensor must have rank >= 1. The shapes of the input tensor and - the mask tensor don’t need to match, but they must be able to be broadcasted. - - For example, on input=[[4, 2, 5], [2, 1, 2], [4, 7, 1]], which has a shape - [3, 3], - - masked_select(input, [[True, False, True], [False, True, False], [True, False, True]]) - - will create a tensor of shape [5] that contains the [4, 5, 1, 4, 1]. - - masked_select(input, [[True], [False], [True]]) - - will create a tensor of shape [6] that contains the [4, 2, 5, 4, 7, 1]. - - masked_select(input, [[False, False, True]]) - - will create a tensor of shape [3] that contains the [5, 2, 1]. - - masked_select(input, [False]) - - will create a tensor of shape [0] which is empty. - - That operation is implemented by NonZero, Shuffle and GatherV2 layers - in TensorRT. - - Parameters: - input : Tensor - The input tensor to select from. - - mask : Tensor - The boolean mask tensor that indicates elements to select. - - Returns: - The 1-D tensor containing the selected elements. - ''' - assert input.rank() >= 1, "input should have rank >= 1" - input, mask = broadcast_helper(input, mask) - expanded_mask = expand(mask, shape(input)) - - non_zero_layer = default_trtnet().add_non_zero(expanded_mask.trt_tensor) - - shuffle_layer = default_trtnet().add_shuffle(non_zero_layer.get_output(0)) - shuffle_layer.second_transpose = (1, 0) - - gather_layer = default_trtnet().add_gather_v2(input.trt_tensor, - shuffle_layer.get_output(0), - mode=trt.GatherMode.ND) - return _create_tensor(gather_layer.get_output(0), gather_layer) - - -def cumsum(input: Tensor, dim: int, prefer_plugin: bool = True) -> Tensor: - ''' - Add an operation to calculate inclusive cumulative sum of elements of - a tensor in a given dimension. - - Given an input tensor, that function creates an operation that calculates - inclusive cumulative sum of elements in the dimension 'dim' to create - a new tensor. The output tensor has the same shape as the input tensor. - - The input tensor must have rank >= 1. The 'dim' must be valid, and negative - value is supported. - - For example, on input=[[4, 2, 5], [2, 1, 2], [4, 7, 1]], which has a shape - [3, 3], - - cumsum(input, 0) - - will produce [[4, 2, 5], [6, 3, 7], [10, 10, 8]]. - - cumsum(input, 1) - - will produce [[4, 6, 11], [2, 3, 5], [4, 11, 12]]. - - That operation is implemented by TensorRT ILoopLayer. - - Parameters: - input : Tensor - The input tensor to calculate the inclusive cumulative sum. - - dim : int - The dimension to calculate the inclusive cumulative sum. Negative - value is supported. - - prefer_plugin : bool - Whether to use the cumsumLastDim plugin if dim is last dim. - - Returns: - The tensor containing the inclusive cumulative sum of input. - ''' - assert input.rank() >= 1, "input should have rank >= 1" - assert dim < input.rank() and dim >= -input.rank( - ), f"dim should be in [{-input.rank()}, {input.rank()}) when input have rank {input.rank()}" - - dim = dim_resolve_negative(dim, input.ndim())[0] - - if dim == input.ndim() - 1: - if prefer_plugin: - last_dim = input.size(-1) - if last_dim == -1: # dynamic? - last_dim = shape(input, -1) - old_shape = shape(input) - if input.ndim() == 1: - input_2d = unsqueeze( - input, 0) # special handling of rank-1 dynamic tensor - elif input.ndim() != 2: - input_2d = input.view(concat([-1, last_dim]), - zero_is_placeholder=False) - else: - input_2d = input - cumsum_last_dim_plg_creator = trt.get_plugin_registry( - ).get_plugin_creator('CumsumLastDim', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert cumsum_last_dim_plg_creator is not None - input_length = trt.PluginField( - "input_length", np.array(input_2d.size(-1), dtype=np.int32), - trt.PluginFieldType.INT32) - pf_type = trt.PluginField("type_id", - np.array([int(input_2d.dtype)], np.int32), - trt.PluginFieldType.INT32) - pfc = trt.PluginFieldCollection([input_length, pf_type]) - cumsum_last_dim_plug = cumsum_last_dim_plg_creator.create_plugin( - "cumsum_last_dim", pfc) - plug_inputs = [input_2d] - plug_inputs = [i.trt_tensor for i in plug_inputs] - layer = default_trtnet().add_plugin_v2(plug_inputs, - cumsum_last_dim_plug) - _add_plugin_info(layer, cumsum_last_dim_plg_creator, - "cumsum_last_dim", pfc) - output = _create_tensor(layer.get_output(0), layer) - output = output.view(old_shape, zero_is_placeholder=False) - return output - else: - # credit to Apple - reduction_length = shape(input, -1) - reduction_range = arange(constant_to_tensor_(0, - dtype='int64', - to_array=False), - reduction_length, - dtype='int64') - lower_triangle = cast(unsqueeze(reduction_range, 0) - <= unsqueeze(reduction_range, 1), - dtype=input.dtype) - output = sum(unsqueeze(input, -2) * lower_triangle, dim=-1) - return output - else: - slice_shape = [] - for i in range(input.ndim()): - if i != dim: - slice_shape.append(shape(input, i)) - - zero_tensor = constant_to_tensor_(0, input.dtype, False) - if len(slice_shape) > 0: - zero_tensor = expand_dims(zero_tensor, - [i for i in range(len(slice_shape))]) - slice_shape = concat(slice_shape) - zero_tensor = expand(zero_tensor, slice_shape) - - loop_layer = default_trtnet().add_loop() - trip_limit = shape(input, dim).trt_tensor - loop_layer.add_trip_limit(trip_limit, trt.TripLimit.COUNT) - - iterator_layer = loop_layer.add_iterator(input.trt_tensor, dim) - cur_slice = iterator_layer.get_output(0) - - running_sum_layer = loop_layer.add_recurrence(zero_tensor.trt_tensor) - running_sum = running_sum_layer.get_output(0) - - cur_sum_layer = default_trtnet().add_elementwise( - cur_slice, running_sum, trt.ElementWiseOperation.SUM) - cur_sum = cur_sum_layer.get_output(0) - running_sum_layer.set_input(1, cur_sum) - - loop_output_layer = loop_layer.add_loop_output( - cur_sum, trt.LoopOutput.CONCATENATE, dim) - loop_output_layer.set_input(1, trip_limit) - return _create_tensor(loop_output_layer.get_output(0), - loop_output_layer) - - -def masked_scatter(input: Tensor, mask: Tensor, source: Tensor) -> Tensor: - ''' - Add the masked_scatter base on PyTorch definition. - - See https://pytorch.org/docs/stable/generated/torch.Tensor.masked_scatter_.html#torch-tensor-masked-scatter for a - description of that function. - - Parameters: - input : Tensor - The input tensor. - - mask : Tensor - The boolean mask tensor that indicates elements to select. - - source: Tensor - The tensor to copy from - Returns: - The tensor containing the source tensor selected by mask. - - ''' - assert input.rank() >= 1, "input should have rank >= 1" - input, mask = broadcast_helper(input, mask) - expanded_mask = expand(mask, shape(input)) - - non_zero_layer = default_trtnet().add_non_zero(expanded_mask.trt_tensor) - - shuffle_layer = default_trtnet().add_shuffle(non_zero_layer.get_output(0)) - shuffle_layer.second_transpose = (1, 0) - source = source.view([-1]) - - scatter_layer = default_trtnet().add_scatter(input.trt_tensor, - shuffle_layer.get_output(0), - source.trt_tensor, - mode=trt.ScatterMode.ND) - - return _create_tensor(scatter_layer.get_output(0), scatter_layer) - - -def concat(inputs: Sequence[Union[Tensor, int]], dim: int = 0) -> Tensor: - ''' - Add an operation to concatenate tensors. - - The function creates an operation that concatenates the tensors from the - sequence 'inputs'. The concatenation is done along the dimension 'dim'. - - All the tensors in 'inputs' must have the same shape expect for the - dimension 'dim'. - - for ii in range(inputs[0].rank()): - assert (ii == dim) or all(inp.shape[ii] == inputs[0].shape[ii] for inp in inputs) - - The shape of the output tensor is defined as: - - for ii in range(inputs[0].rank()): - # Same size as all the inputs in dimension ii != dim. - output.shape[ii] = inputs[0].shape[ii] - - # Sum of the sizes in the different inputs in dimension 'dim'. - if ii == dim: - for jj in range(1, len(inputs)): - output.shape[ii] += inputs[jj].shape[ii] - - For example, given a sequence of two 2D tensors [[0, 1], [2, 3]] and - [[4, 5], [6, 7]] both of shape [2, 2], - - concat(inputs, 0) - - will produce [[0, 1], [2, 3], [4, 5], [6, 7]] of shape [4, 2] and - - concat(inputs, 1) - - will produce [[0, 1, 4, 5], [2, 3, 6, 7]] of shape [2, 4]. - - Parameters: - inputs : Sequence[Union[Tensor, int]] - The sequence of tensors to concatenate. For integers, that function - creates constant tensors. - - dim : int - The dimension in which the concatenation is performed. - - Returns: - A tensor that contains the concatenation of the tensors. - ''' - assert len( - inputs - ) > 0, f"Number of inputs ({len(inputs)}) to the concatenation layer must be > 0." - tmp = [] - inputs = constants_to_tensors_(*inputs) - for i in inputs: - if i.rank() == 0: - tmp.append(i.view([1])) - else: - tmp.append(i) - - layer = default_trtnet().add_concatenation([i.trt_tensor for i in tmp]) - layer.axis = dim_resolve_negative(dim, tmp[0].ndim())[0] - return _create_tensor(layer.get_output(0), layer) - - -def softmax(input: Tensor, dim: Optional[int] = None) -> Tensor: - ''' - Add an operation to compute softmax on a tensor. - - That operation computes the softmax on the input tensor in the dimension - 'dim' if specified. Otherwise, it is applied on the last dimension. - - It inserts a ISoftmaxLayer to the TensorRT graph. - - Parameters: - input : Tensor - The input tensor on which to apply softmax. - - dim : Optional[int] - The dimension used to apply softmax. - - Returns: - The output tensor of the softmax layer. - ''' - if dim is None: - dim = input.ndim() - 1 - if dim < 0: - dim = input.ndim() + dim - axes = dim_to_trt_axes(dim) - - layer = default_trtnet().add_softmax(input.trt_tensor) - layer.axes = axes - - return _create_tensor(layer.get_output(0), layer) - - -def _lookup_plugin(input: Tensor, weight: Tensor, rank: int, - per_token_scale: Tensor) -> Tensor: - ''' - Add an operation to perform lookup in a tensor. - - That operation performs the lookup needed by embedding layers. Given a - 'weight' tensor of shape [rows, cols], it produces a tensor of shape - [inputs.size(0), cols] where the ith row corresponds to the input[i] row in - the weight tensor. - - It inserts a IPluginV2Layer. - - Parameters: - input : Tensor - The input tensor contains the indices to perform the lookup. - - weight : Tensor - The table to gather from. - - rank : int - The mpi rank. - - Returns: - The output tensor of the lookup layer. - ''' - plg_creator = trt.get_plugin_registry().get_plugin_creator( - 'Lookup', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert plg_creator is not None - - p_dtype = per_token_scale.dtype - pf_type = trt.PluginField("type_id", np.array([int(p_dtype)], np.int32), - trt.PluginFieldType.INT32) - - rank = trt.PluginField("rank", np.array([int(rank)], np.int32), - trt.PluginFieldType.INT32) - - pfc = trt.PluginFieldCollection([pf_type, rank]) - lookup_plug = plg_creator.create_plugin("lookup", pfc) - plug_inputs = [input.trt_tensor, weight.trt_tensor] - if per_token_scale is not None: - plug_inputs.append(per_token_scale.trt_tensor) - weight.trt_tensor.set_dynamic_range(-127, 127) - layer = default_trtnet().add_plugin_v2(plug_inputs, lookup_plug) - _add_plugin_info(layer, plg_creator, "lookup", pfc) - return _create_tensor(layer.get_output(0), layer) - - -def embedding(input: Tensor, - weight: Tensor, - tp_size=1, - tp_group=None, - sharding_dim=0, - tp_rank=None, - per_token_scale=None, - padding=None) -> Tensor: - ''' - Add an operation to perform embedding lookup. - - That operation performs the embedding lookup. The 'input' tensor contains - the identifiers of the rows of 'weight' to gather. - - 1. Distribute the embedding lookup table over multiple GPU - When 'tp_size' is greater than 1 and the 'tp_group' is defined, this - embedding lookup is distributed among multiple GPUs. - - When 'sharding_dim==0', each GPU stores a subset of the rows of the embedding - table rows(that number of rows per GPU is given by weights.shape[0] and the offset to - the 1st row stored on the GPU is given by rank * weights.shape[0]). Each - parallel rank will query all the indices and set 0s for the weights that - are not stored on the associated GPU. To compute the final result, a - parallel all-reduce operation is added to the TensorRT graph. That lookup - can be performed using either the plugin or the operators TensorRT support. - - When'sharding_dim==1', each GPU stores a subset of the embedding table's columns. - Each rank can obtain a portion of the embedding results. - Then the embedding is collected using the all-gather operation. - Related transposition operations are also used to obtain the final results. - - 2. Store embedding lookup table as a whole - When 'tp_size' is not greater than 1, the embedding lookup table will not - be divided. In this case, when the default_net().plugin_config.lookup_plugin is set, - the operation is implemented using a plugin (without the all-reduce operation). - Otherwise, this operation is implemented using the standard IGatherLayer in TensorRT. - - Parameters: - input : Tensor - The input tensor the contains the indices to perform the lookup. - - weight : Tensor - The table to gather from. - - tp_size : int - The number of GPUs collaborating to perform that embedding. - - tg_group : Optional[List[int]] - The group of world ranks participating in the all-reduce when - tp_size > 1. - - sharding_dim : int - sharding_dim = 0 means that we shard the embedding table in vocab dim; - sharding_dim = 1 means that we shard the embedding table in embedding dim. - - tp_rank : int - The tensor parallelism rank. Used to calculate offset in TP on vocab dim. - - padding: Tensor - Additional padding added to the end of the embedding table before feeding into gather op. - - Returns: - The tensor produced by the embedding lookup layer. - ''' - - # Per token scale is only supported by lookup plugin so if per_token_scale is not None, we must use lookup plugin - # Otherwise, we prefer to use ootb - use_lookup_plugin = per_token_scale is not None - - if padding is not None: - padded_weight = concat([weight, padding], dim=0) - else: - padded_weight = weight - - # Distribute embedding lookup table across multiple GPU - if tp_size > 1 and tp_group is not None: - if sharding_dim == 0: # TP on vocab_size dimension - if tp_rank is None: - raise ValueError( - "Rank cannot be none for tensor parallelism on vocab dim") - - if use_lookup_plugin: - x = _lookup_plugin(input, weight, tp_rank, per_token_scale) - x = allreduce(x, tp_group) - else: - shape_weight = shape(weight) - vocab_size = slice(shape_weight, starts=[0], sizes=[1]) - tmp_input = input - vocab_size * tp_rank - - # Identify the valid indices - is_qualified = op_and(tmp_input >= 0, tmp_input < vocab_size) - is_qualified_expand = expand_dims(is_qualified, - [is_qualified.ndim()]) - - # Replace the invalid ones to zero - placeholder_input = where(is_qualified, tmp_input, 0) - - # Get the temporal results - layer = default_trtnet().add_gather( - padded_weight.trt_tensor, placeholder_input.trt_tensor, 0) - tmp_output = _create_tensor(layer.get_output(0), layer) - - # Set zero for invalid results - placeholder_tmp = cast(is_qualified_expand, tmp_output.dtype) - placeholder = placeholder_tmp - placeholder_tmp - x = where(is_qualified_expand, tmp_output, placeholder) - - # Use all reduce to collect the results - x = allreduce(x, tp_group) - - elif sharding_dim == 1: # TP on hidden dimension - layer = default_trtnet().add_gather(padded_weight.trt_tensor, - input.trt_tensor, 0) - x = _create_tensor(layer.get_output(0), layer) - - # [dim0, local_dim] -> [dim0 * tp_size, local_dim] --> [dim0, local_dim * tp_size] - x = allgather(x, tp_group, gather_dim=-1) - - else: - raise ValueError( - 'Tensor Parallelism only support splitting Embedding lookup along hidden (sharding_dim==1) and vocab (sharding_dim==0) dimensionis' - ) - - # Store embedding lookup table as a whole - else: - if use_lookup_plugin: - x = _lookup_plugin(input, - padded_weight, - rank=0, - per_token_scale=per_token_scale) - else: - layer = default_trtnet().add_gather(padded_weight.trt_tensor, - input.trt_tensor, 0) - x = _create_tensor(layer.get_output(0), layer) - return x - - -def constant_to_tensor_(input: Union[Tensor, int, float, bool], - dtype: Union[trt.DataType, str] = None, - to_array=True) -> Tensor: - if dtype is None: - # deduce the type from the given value - # NOTE: bool is a subtype of int, so bool needs to be checked first - if isinstance(input, bool): - dtype = trt.bool - elif isinstance(input, int): - dtype = trt.int32 - else: - dtype = trt.float32 - - if not isinstance(input, Tensor): - if isinstance(dtype, str): - dtype = str_dtype_to_trt(dtype) - array_fn_dict = { - trt.int64: int64_array, - trt.int32: int32_array, - trt.float32: fp32_array, - trt.float16: fp16_array, - trt.bfloat16: bf16_array, - trt.bool: bool_array, - } - assert dtype in array_fn_dict - return constant(array_fn_dict[dtype]([input] if to_array else input)) - - return input - - -def constants_to_tensors_( - *inputs: Union[Tensor, int, float]) -> Tuple[Tensor, ...]: - ''' - Helper function to create tensors from multiple inputs. - - For each inputs, that function first creates a constant tensor if the input - is an integer or a float. Then, if any input is int64, it upcasts other - integer inputs to int64. - - Parameters: - inputs : Tuple[Union[Tensor, int, float], ...] - The inputs to create tensors from. - - Returns: - A tuple of tensors. - ''' - has_int64: bool = False - for i in inputs: - if isinstance(i, int) and (i >= 2**31 or i < -2**31)\ - or isinstance(i, Tensor) and i.dtype == trt.int64: - has_int64 = True - break - - if not has_int64: - return tuple(constant_to_tensor_(i) for i in inputs) - - result = [] - for i in inputs: - if isinstance(i, int) or isinstance(i, Tensor) and i.dtype == trt.int32: - result.append( - constant_to_tensor_(i, trt.int64 if has_int64 else trt.int32)) - else: - result.append(constant_to_tensor_(i)) - return tuple(result) - - -def broadcast_helper(left: Union[Tensor, int, float], - right: Union[Tensor, int, float]) -> Tuple[Tensor, Tensor]: - ''' - Helper function to perform a broadcast. - - For each input, that function first creates a constant tensor if the input - is an integer or a float. Then, if needed, it expands the smaller tensor to - make sure its rank is the same as the larger one. - - Parameters: - left : Union[Tensor, int, float] - The first input. If that input is an integer or a float, the - function creates a constant tensor. - - right : Union[Tensor, int, float] - The second input. If that input is an integer or a float, the - function creates a constant tensor. - - Returns: - A pair of tensors of same rank. - ''' - if not default_net().strongly_typed: - left = constant_to_tensor_(left) - right = constant_to_tensor_(right) - else: - left = constant_to_tensor_( - left, right.dtype if isinstance(right, Tensor) else None) - right = constant_to_tensor_(right, left.dtype) - - if left.rank() == right.rank(): - return (left, right) - - if left.rank() < right.rank(): - left = expand_dims_like(left, right) - return (left, right) - - if left.rank() > right.rank(): - right = expand_dims_like(right, left) - return (left, right) - - -def elementwise_binary(left: Union[Tensor, int, - float], right: Union[Tensor, int, float], - op: trt.ElementWiseOperation) -> Tensor: - ''' - Add an elementwise operation with two inputs. - - For each input, that function first creates a constant tensor if the input - is an integer or a float. Then, if needed, it expands the smaller tensor to - make sure its rank is the same as the larger one. Then, it performs the - elementwise operation 'op'. - - The following closures are defined in functional.*: - - add for op=trt.ElementWiseOperation.SUM - sub for op=trt.ElementWiseOperation.SUB - mul for op=trt.ElementWiseOperation.PROD - div for op=trt.ElementWiseOperation.DIV - floordiv for op=trt.ElementWiseOperation.FLOOR_DIV - gt for op=trt.ElementWiseOperation.GREATER - lt for op=trt.ElementWiseOperation.LESS - op_and for op=trt.ElementWiseOperation.AND - op_or for op=trt.ElementWiseOperation.OR - eq for op=trt.ElementWiseOperation.EQUAL - minimum for op=trt.ElementWiseOperation.MIN - maximum for op=trt.ElementWiseOperation.MAX - pow for op=trt.ElementWiseOperation.POW - - It is implemented using the IElementWiseLayer from TensorRT. - - Parameters: - left : Union[Tensor, int, float] - The first input. If that input is an integer or a float, the - function creates a constant tensor. - - right : Union[Tensor, int, float] - The second input. If that input is an integer or a float, the - function creates a constant tensor. - - op : trt.ElementWiseOperation - The binary operation to perform. - - Returns: - The tensor produced by this elementwise operation. - ''' - left, right = broadcast_helper(left, right) - if left.dtype == trt.int32 and right.dtype == trt.int64: - left = cast(left, trt.int64) - if left.dtype == trt.int64 and right.dtype == trt.int32: - right = cast(right, trt.int64) - layer = default_trtnet().add_elementwise(left.trt_tensor, right.trt_tensor, - op) - return _create_tensor(layer.get_output(0), layer) - - -add = partial(elementwise_binary, op=trt.ElementWiseOperation.SUM) -sub = partial(elementwise_binary, op=trt.ElementWiseOperation.SUB) -mul = partial(elementwise_binary, op=trt.ElementWiseOperation.PROD) -div = partial(elementwise_binary, op=trt.ElementWiseOperation.DIV) -floordiv = partial(elementwise_binary, op=trt.ElementWiseOperation.FLOOR_DIV) -gt = partial(elementwise_binary, op=trt.ElementWiseOperation.GREATER) -lt = partial(elementwise_binary, op=trt.ElementWiseOperation.LESS) -op_and = partial(elementwise_binary, op=trt.ElementWiseOperation.AND) -op_or = partial(elementwise_binary, op=trt.ElementWiseOperation.OR) -eq = partial(elementwise_binary, op=trt.ElementWiseOperation.EQUAL) -minimum = partial(elementwise_binary, op=trt.ElementWiseOperation.MIN) -maximum = partial(elementwise_binary, op=trt.ElementWiseOperation.MAX) -pow = partial(elementwise_binary, op=trt.ElementWiseOperation.POW) -op_xor = partial(elementwise_binary, op=trt.ElementWiseOperation.XOR) - - -def modulo(x: Tensor, y: Union[Tensor, int]) -> Tensor: - ''' - This function adds an element-wise modulo (x % y) operation for a given tensor. - Since there is no TensorRT layer that can directly perform this, - this function implements it using some of the basic operations. - - Returns: - A tensor that represents (x % y) modulo operation. - ''' - return x - (x // y) * y - - -def where(condition: Union[Tensor, bool], left: Union[Tensor, int, float], - right: Union[Tensor, int, float]) -> Tensor: - ''' - Add a where (aka select or if-then-else) operation. - - Assuming the three input parameters have the same shape, that function creates - the operation to compute a tensor of the same shape such that: - - for ii in range(mul(condition.shape)): - output[ii] = left[ii] if condition[ii] else right[ii] - - For each input, that function first creates a constant tensor if the - condition is boolean or the left/right input is an integer or a float. - Then, if needed, it expands the smaller tensor to make sure its - rank is the same as the larger one. Then, it performs the selection. - - It is implemented using the ISelectLayer from TensorRT. - - Parameters: - condition : Union[Tensor, bool] - The condition. If that input is a boolean, the function - creates a constant tensor. - - left : Union[Tensor, int, float] - The first input. If that input is an integer or a float, the - function creates a constant tensor. - - right : Union[Tensor, int, float] - The second input. If that input is an integer or a float, the - function creates a constant tensor. - - Returns: - The tensor produced by this where operation. - ''' - # Convert to tensors. - condition = constant_to_tensor_(condition) - left, right = constants_to_tensors_(left, right) - - # Find the tensor with the largest rank of the three. - largest = condition - if largest.rank() < left.rank(): - largest = left - if largest.rank() < right.rank(): - largest = right - - # Expand the tensors to match the largest one. - if condition is not largest: - condition = expand_dims_like(condition, largest) - if left is not largest: - left = expand_dims_like(left, largest) - if right is not largest: - right = expand_dims_like(right, largest) - - # Insert the operation. - layer = default_trtnet().add_select(condition.trt_tensor, left.trt_tensor, - right.trt_tensor) - return _create_tensor(layer.get_output(0), layer) - - -def unary(input: Tensor, op: trt.UnaryOperation) -> Tensor: - ''' - Add an elementwise operation on a single input. - - The following closures are defined in functional.*: - - round for op=trt.UnaryOperation.ROUND - sqrt for op=trt.UnaryOperation.SQRT - exp for op=trt.UnaryOperation.EXP - sin for op=trt.UnaryOperation.SIN - cos for op=trt.UnaryOperation.COS - abs for op=trt.UnaryOperation.ABS - log for op=trt.UnaryOperation.LOG - - It is implemented using the IUnaryLayer from TensorRT. - - Parameters: - input : Tensor - The input tensor. - - op : trt.UnaryOperation - The unary operation to perform. - - Returns: - The tensor produced by this elementwise operation. - ''' - layer = default_trtnet().add_unary(input.trt_tensor, op) - return _create_tensor(layer.get_output(0), layer) - - -round = partial(unary, op=trt.UnaryOperation.ROUND) -sqrt = partial(unary, op=trt.UnaryOperation.SQRT) -exp = partial(unary, op=trt.UnaryOperation.EXP) -sin = partial(unary, op=trt.UnaryOperation.SIN) -cos = partial(unary, op=trt.UnaryOperation.COS) -abs = partial(unary, op=trt.UnaryOperation.ABS) -log = partial(unary, op=trt.UnaryOperation.LOG) -not_op = partial(unary, op=trt.UnaryOperation.NOT) - - -def log_softmax(input: Tensor, dim: int) -> Tensor: - ''' - This function is equivalent of torch.nn.functional.log_softmax() i.e. - it performs log(softmax(input)) in a safer and faster way. - - Parameters: - input: Tensor - The data tensor on which log_softmax to be computed. - dim: int - The dimension of the input tensor along which log_softmax will be computed. - Returns: - A tensor of same shape as input with log_softmax computed on the specified dim. - ''' - x_max = max(input, dim=dim, keepdim=True) - x = input - x_max - return x - log(sum(exp(x), dim=dim, keepdim=True)) - - -def reduce(input: Tensor, - op: trt.ReduceOperation, - dim: Union[int, Tuple[int]], - keepdim: bool = False) -> Tensor: - ''' - Add an reduction operation to do along a dimension. - - It is implemented using the IReduceLayer from TensorRT. - - Parameters: - input : Tensor - The input tensor. - - op : trt.ReduceOperation - The reduction operation to perform. - Options: SUM, PROD, MAX, MIN, AVG - - dim : int - The dimension along which the reduction is performed. - - keepdim : bool - Is the dimension kept in the reduced tensor? When True the - dimension is kept, it is removed from the shape otherwise. - - Returns: - The tensor produced by this reduction operation. - ''' - dim = dim_resolve_negative(dim, input.ndim()) - axes = dim_to_trt_axes(dim) - - layer = default_trtnet().add_reduce(input.trt_tensor, - op, - axes, - keep_dims=keepdim) - return _create_tensor(layer.get_output(0), layer) - - -prod = partial(reduce, op=trt.ReduceOperation.PROD) -min = partial(reduce, op=trt.ReduceOperation.MIN) - - -def mean(input: Tensor, - dim: Union[int, Tuple[int]], - keepdim: bool = False) -> Tensor: - ''' - Add an operation to compute the mean along a dimension. - - Computes the mean along the dimension 'dim' of the input tensor. - - It is implemented using the IReduceLayer from TensorRT. - - Parameters: - input : Tensor - The input tensor. - - dim : int - The dimension along which the mean is computed. - - keepdim : bool - Is the dimension kept in the reduced tensor? When True the - dimension is kept, it is removed from the shape otherwise. - - Returns: - The tensor produced by this reduction operation. - ''' - return reduce(input, op=trt.ReduceOperation.AVG, dim=dim, keepdim=keepdim) - - -def max(input: Tensor, dim: int, keepdim: bool = False) -> Tensor: - ''' - Add an operation to compute the max along a dimension. - - Computes the max along the dimension 'dim' of the input tensor. - - It is implemented using the IReduceLayer from TensorRT. - - Parameters: - input : Tensor - The input tensor. - - dim : int - The dimension along which the mean is computed. - - keepdim : bool - Is the dimension kept in the reduced tensor? When True the - dimension is kept, it is removed from the shape otherwise. - - Returns: - The tensor produced by this reduction operation. - ''' - return reduce(input, op=trt.ReduceOperation.MAX, dim=dim, keepdim=keepdim) - - -def sum(input: Tensor, dim: int, keepdim: bool = False) -> Tensor: - ''' - Add an operation to compute the sum along a dimension. - - Computes the sum along the dimension 'dim' of the input tensor. - - It is implemented using the IReduceLayer from TensorRT. - - Parameters: - input : Tensor - The input tensor. - - dim : int - The dimension along which the mean is computed. - - keepdim : bool - Is the dimension kept in the reduced tensor? When True the - dimension is kept, it is removed from the shape otherwise. - - Returns: - The tensor produced by this reduction operation. - ''' - return reduce(input, op=trt.ReduceOperation.SUM, dim=dim, keepdim=keepdim) - - -def identity(input: Tensor) -> Tensor: - ''' - Add an identity operation. - - Parameters: - input : Tensor - The input tensor. - - Returns: - The tensor produced by this identity operation. - ''' - if not default_net().plugin_config.identity_plugin: - layer = default_trtnet().add_identity(input.trt_tensor) - else: - plg_creator = trt.get_plugin_registry().get_plugin_creator( - 'Identity', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert plg_creator is not None - pfc = trt.PluginFieldCollection() - id_plug = plg_creator.create_plugin("identity", pfc) - plug_inputs = [input.trt_tensor] - layer = default_trtnet().add_plugin_v2(plug_inputs, id_plug) - _add_plugin_info(layer, plg_creator, "identity", pfc) - return _create_tensor(layer.get_output(0), layer) - - -def argmax(input: Tensor, dim: int, keepdim: bool = False) -> Tensor: - ''' - Add an argmax operation. - - As explained in the ONNX documentation, - - https://github.com/onnx/onnx/blob/main/docs/Operators.md#argmax - - that function creates a layer computing the indices of the max elements of - the input tensor's element along the provided dim. The resulting tensor - has the same rank as the input if keepdims is True. If keepdims is False, - then the resulting tensor has the reduced dimension pruned. - - Parameters: - input : Tensor - The input tensor. - - dim : int - The dimension in which to compute the argmax indices. - - keepdim : bool - Do we keep the dimension along which the reduction is performed? - Yes, if set to True, no otherwise. - - Returns: - The tensor produced by this argmax operation. - ''' - dim = dim_resolve_negative(dim, input.ndim()) - axes = dim_to_trt_axes(dim) - - layer = default_trtnet().add_topk(input.trt_tensor, trt.TopKOperation.MAX, - 1, axes) - output = layer.get_output(1) - - if keepdim: - return _create_tensor(output, layer) - - output = _create_tensor(output, layer) - a = list(range(input.ndim())) - for d in dim: - a.pop(d) - indices = constant(int32_array(a)) - output_shape = shape(output) - new_shape = gather(output_shape, 0, indices) - return view(output, new_shape) - - -def gelu(x: Tensor) -> Tensor: - ''' - Add a GELU operation. - - Parameters: - input : Tensor - The input tensor on which the activation function is applied. - - Returns: - The tensor produced by the activation layer. - ''' - return 0.5 * x * ( - tanh(math.sqrt(2.0 / math.pi) * (x + 0.044715 * pow(x, 3.0))) + 1.0) - - -def geglu(x: Tensor) -> Tensor: - ''' - Add a Gated-GELU operation. - - That function takes a tensor, splits it into two halves along the last - dimension, applies GELU to the second half and multiply the results. The - behavior is undefined if the last dimension is not even. - - Parameters: - input : Tensor - The input tensor on which the activation function is applied. - - Returns: - The tensor produced by the activation layer. - ''' - a, b = chunk(x, 2, dim=-1) - return a * gelu(b) - - -def quick_gelu(x: Tensor) -> Tensor: - return x * sigmoid(1.702 * x) - - -def gegelu(x: Tensor, limit: Optional[float] = None) -> Tensor: - # a, b = x[..., ::2], x[..., 1::2] - ndim = x.ndim() - a_starts = [0 for i in range(ndim)] - b_starts = [1 if i == (ndim - 1) else 0 for i in range(ndim)] - shapes = concat([ - shape(x, i) / 2 if i == (ndim - 1) else shape(x, i) for i in range(ndim) - ]) - strides = [2 if i == (ndim - 1) else 1 for i in range(ndim)] - - a = slice(x, a_starts, shapes, strides) - b = slice(x, b_starts, shapes, strides) - - if limit is not None: - a = clip(a, alpha=float(-1e20), beta=limit) - b = clip(b, alpha=-limit, beta=limit) - - # C = B + 1 - const1 = arange(constant(int32_array(1)), constant(int32_array(2)), - trt_dtype_to_str(b.dtype)) - for _ in range(ndim - 1): - const1 = expand_dims(const1, 0) - - b_shape = concat([shape(b, i) for i in range(ndim)]) - const1_arr = expand(const1, b_shape) - - return quick_gelu(a) * (b + const1_arr) - - -def group_norm(input: Tensor, - num_groups: int, - weight: Optional[Tensor] = None, - bias: Optional[Tensor] = None, - eps: float = 1e-05): - - ## - ## TODO: Document that function! - ## - - assert not input.is_dynamic(1) - num_channels = input.size()[1] - - ndim = input.ndim() - old_shape = shape(input) - new_shape = concat([ - input.size(0), - num_groups, - num_channels // num_groups, - ] + [input.size(i) for i in range(2, ndim)]) - x = input.view(new_shape) - - # instance norm - w_shape = [1, num_groups] + [1 for i in range(ndim - 1)] - instance_weight = constant(np.ones(w_shape, dtype=trt_dtype_to_np(x.dtype))) - instance_bias = constant(np.zeros(w_shape, dtype=trt_dtype_to_np(x.dtype))) - axes_mask = 0 - for i in range(2, x.ndim()): - axes_mask |= 1 << i - layer = default_trtnet().add_normalization(x.trt_tensor, - instance_weight.trt_tensor, - instance_bias.trt_tensor, - axes_mask) - layer.epsilon = eps - y = _create_tensor(layer.get_output(0), layer) - y = y.view(old_shape) - - new_shape = concat([num_channels] + [1 for _ in range(2, ndim)]) - if weight is not None: - y = y * weight.view(new_shape) - if bias is not None: - y = y + bias.view(new_shape) - - return y - - -def softplus(input: Tensor, beta: float, threshold: float) -> Tensor: - ''' - Add the softplus activation base on PyTorch definition. - - See https://pytorch.org/docs/stable/generated/torch.nn.functional.softplus.html#torch-nn-functional-softplus for a - description of that function. - - Parameters: - input : Tensor - Input TensorRT LLM Tensor. - beta : float - The parameter for softplus computation. - threshold : float - The threshold for reverting to the linear function when input * beta > threshold - - Returns: - The output tensor created by that layer. - ''' - sf_layer = default_trtnet().add_activation(input.trt_tensor, - trt.ActivationType.SOFTPLUS) - sf_layer.alpha = 1 / beta - sf_layer.beta = beta - - prod_tensor = input * beta - result = prod_tensor > threshold - - return where(result, input, _create_tensor(sf_layer.get_output(0), - sf_layer)) - - -def outer(input: Tensor, vec2: Tensor) -> Tensor: - ''' - Add an operation to compute the outer product between two tensors. - - That operation creates an Einsum node. - - Parameters: - input : Tensor - The first input tensor. - - vec2 : Tensor - The second input tensor. - - Returns: - The output tensor produced by this layer. - ''' - return einsum('i,j->ij', [input, vec2]) - - -def avg_pool2d(input: Tensor, - kernel_size: Tuple[int], - stride: Optional[Tuple[int]] = None, - padding: Optional[Tuple[int]] = (0, 0), - ceil_mode: bool = False, - count_include_pad: bool = True) -> Tensor: - - ## - ## TODO: Document that function! - ## - - assert not input.is_dynamic() - ndim = input.ndim() - if ndim == 3: - input = expand_dims(input, 0) - - layer = default_trtnet().add_pooling_nd(input.trt_tensor, - trt.PoolingType.AVERAGE, - kernel_size) - if stride is None: - stride = kernel_size - layer.stride_nd = stride - - output = _create_tensor(layer.get_output(0), layer) - - if ndim == 3: - return output.view( - concat([output.size(1), - output.size(2), - output.size(3)])) - - return output - - -def _get_trt_weight(weight: Tensor) -> Tuple[trt.Weights, bool]: - is_weight_constant = (weight.producer is not None - and weight.producer.type == trt.LayerType.CONSTANT) - if is_weight_constant: - ndarray = get_np_weight(default_trtnet(), weight.producer.name) - if ndarray is not None: - trt_weight = trt.Weights(np_dtype_to_trt(ndarray.dtype), - ndarray.ctypes.data, - int(np.prod(ndarray.shape))) - else: - weight.producer.__class__ = trt.IConstantLayer - trt_weight = weight.producer.weights - else: - trt_weight = trt.Weights() - - return trt_weight, is_weight_constant - - -def conv1d(input: Tensor, - weight: Tensor, - bias: Optional[Tensor] = None, - stride: int = 1, - padding: int = 0, - dilation: int = 1, - groups: int = 1) -> Tensor: - - noutput = weight.size()[0] - kernel_size = weight.size()[-2] - kernel_shape = trt.Dims([kernel_size, 1]) - - trt_weight, is_weight_constant = _get_trt_weight(weight) - weight_tensor = weight - - if bias is not None: - bias_tensor = bias - trt_bias, is_bias_constant = _get_trt_weight(bias) - else: - bias_tensor = None - trt_bias = None - - input_shuffled = stack([input], dim=input.ndim()) - - layer = default_trtnet().add_convolution_nd(input_shuffled.trt_tensor, - noutput, kernel_shape, - trt_weight, trt_bias) - layer.stride_nd = (stride, 2) - layer.padding_nd = (padding, 0) - layer.dilation_nd = (dilation, 2) - layer.num_groups = groups - - if not is_weight_constant: - layer.set_input(1, weight_tensor.trt_tensor) - if bias_tensor is not None and not is_bias_constant: - layer.set_input(2, bias_tensor.trt_tensor) - - output_2d = _create_tensor(layer.get_output(0), layer) - output_1d = squeeze(output_2d, dim=-1) - return output_1d - - -def conv2d(input: Tensor, - weight: Tensor, - bias: Optional[Tensor] = None, - stride: Tuple[int, int] = (1, 1), - padding: Tuple[int, int] = (0, 0), - dilation: Tuple[int, int] = (1, 1), - groups: int = 1, - pre_padding: Optional[Tuple[int, int]] = None, - post_padding: Optional[Tuple[int, int]] = None) -> Tensor: - ## - ## TODO: Document that function! - ## - - ndim = input.ndim() - if ndim == 3: - input = expand_dims(input, 0) - - noutput = weight.size()[0] - kernel_size = (weight.size()[-2], weight.size()[-1]) - kernel_shape = trt.Dims(list(kernel_size)) - - trt_weight, is_weight_constant = _get_trt_weight(weight) - weight_tensor = weight - - if bias is not None: - bias_tensor = bias - trt_bias, is_bias_constant = _get_trt_weight(bias) - else: - bias_tensor = None - trt_bias = None - - layer = default_trtnet().add_convolution_nd(input.trt_tensor, noutput, - kernel_shape, trt_weight, - trt_bias) - layer.stride_nd = stride - layer.padding_nd = padding - layer.dilation_nd = dilation - layer.num_groups = groups - layer.dilation_nd = dilation - if pre_padding: - layer.pre_padding = pre_padding - if post_padding: - layer.post_padding = post_padding - - if not is_weight_constant: - layer.set_input(1, weight_tensor.trt_tensor) - if bias_tensor is not None and not is_bias_constant: - layer.set_input(2, bias_tensor.trt_tensor) - - output = _create_tensor(layer.get_output(0), layer) - - if ndim == 3: - return output.view( - concat([output.size(1), - output.size(2), - output.size(3)])) - - return output - - -def conv3d(input: Tensor, - weight: Tensor, - bias: Optional[Tensor] = None, - stride: Union[int, Tuple[int, int]] = (1, 1, 1), - padding: Union[int, Tuple[int, int]] = (0, 0, 0), - dilation: Union[int, Tuple[int, int]] = (1, 1, 1), - groups: int = 1) -> Tensor: - ## - ## TODO: Document this function! - ## - - ndim = input.ndim() - # TRT requires the input of Conv3D layer to be 5-dimentional tensor. - if ndim == 4: - input = expand_dims(input, 0) - assert input.ndim() == 5 - - if isinstance(stride, int): - stride = tuple([stride] * 3) - if isinstance(padding, int): - padding = tuple([padding] * 3) - if isinstance(dilation, int): - dilation = tuple([dilation] * 3) - - noutput = weight.size()[0] - kernel_size = (weight.size()[-3], weight.size()[-2], weight.size()[-1]) - kernel_shape = trt.Dims(list(kernel_size)) - - trt_weight, is_weight_constant = _get_trt_weight(weight) - weight_tensor = weight - - if bias is not None: - bias_tensor = bias - trt_bias, is_bias_constant = _get_trt_weight(bias) - else: - bias_tensor = None - trt_bias = None - - layer = default_trtnet().add_convolution_nd(input.trt_tensor, noutput, - kernel_shape, trt_weight, - trt_bias) - layer.stride_nd = stride - layer.padding_nd = padding - layer.dilation_nd = dilation - layer.num_groups = groups - layer.dilation_nd = dilation - - if not is_weight_constant: - layer.set_input(1, weight_tensor.trt_tensor) - if bias_tensor is not None and not is_bias_constant: - layer.set_input(2, bias_tensor.trt_tensor) - - output = _create_tensor(layer.get_output(0), layer) - return output - - -def conv_transpose2d(input: Tensor, - weight: Tensor, - bias: Optional[Tensor] = None, - stride: Tuple[int, int] = (1, 1), - padding: Tuple[int, int] = (0, 0), - output_padding: Tuple[int, int] = (0, 0), - dilation: Tuple[int, int] = (1, 1), - groups: int = 1) -> Tensor: - ## - ## TODO: Document that function! - ## - - assert not input.is_dynamic() - - ndim = input.ndim() - if ndim == 3: - input = expand_dims(input, 0) - - noutput = weight.size()[1] - kernel_size = (weight.size()[-2], weight.size()[-1]) - kernel_shape = trt.Dims(list(kernel_size)) - - trt_weight, is_weight_constant = _get_trt_weight(weight) - weight_tensor = weight - - if bias is not None: - bias_tensor = bias - trt_bias, is_bias_constant = _get_trt_weight(bias) - else: - bias_tensor = None - trt_bias = None - - layer = default_trtnet().add_deconvolution_nd(input.trt_tensor, noutput, - kernel_shape, trt_weight, - trt_bias) - layer.stride_nd = stride - layer.padding_nd = padding - layer.num_groups = groups - - if not is_weight_constant: - layer.set_input(1, weight_tensor.trt_tensor) - if bias_tensor is not None and not is_bias_constant: - layer.set_input(2, bias_tensor.trt_tensor) - - output = _create_tensor(layer.get_output(0), layer) - - if ndim == 3: - return output.view( - concat([output.size(1), - output.size(2), - output.size(3)])) - - return output - - -def split(tensor: Tensor, - split_size_or_sections: Union[int, Sequence[int]], - dim: int = 0) -> Sequence[Tensor]: - ''' - Add an operation that splits a tensor into sub-tensors. - - This operation creates a list of tensors that are obtained from the input - tensor by slicing it along the dimension 'dim'. If 'split_size_or_sections' - is an integer, the tensor is split into 'input.shape[dim] / - split_size_or_sections' slices. If 'split_size_or_sections' is a list of - sizes, the tensor is split into 'len(split_size_or_sections)' slices and - the size of the ith slice is given by 'split_size_or_sections[i]'. - - There are several constraints with the current implementation: - - - The input tensor must be static (no dynamic dimension), - - If 'split_size_or_sections' is an integer, the number of elements in - the 'dim' dimension of the input must be a multiple of - 'split_size_or_sections': 'input.shape[dim] % split_size_or_sections == 0'. - - If 'split_size_or_sections' is a sequence, the sum of the elements in - 'split_size_or_sections' must be equal to the size in the dimension - 'dim': 'input.shape[dim] == sum(ii for ii in split_size_or_sections)'. - - That operation is implemented using a 'slice' operation for each output - slice. - - Parameters: - tensor : Tensor - The input tensor to slice. - - split_size_or_sections : Union[int, Sequence[int]] - If it is an integer, it encodes the size of each slice. Otherwise, - if it is a sequence, it is the size of each slice. - - dim : int - The dimension of the tensor to slice. - - Returns: - The list of tensors produced by the different operations. - ''' - assert not tensor.is_dynamic(dim) - - ndim = tensor.ndim() - if dim < 0: - dim += ndim - dim_value = tensor.size()[dim] - starts = [constant(dims_array([0])) for _ in range(ndim)] - sizes = [shape(tensor, i) for i in range(ndim)] - - if isinstance(split_size_or_sections, int): - # TODO: support non-divisible cases - assert dim_value % split_size_or_sections == 0 - num_sections = dim_value // split_size_or_sections - sizes[dim] = constant(dims_array([split_size_or_sections])) - - outputs = [] - for i in range(num_sections): - starts[dim] = constant(dims_array([split_size_or_sections * i])) - outputs.append(slice(tensor, concat(starts), concat(sizes))) - return outputs - else: - total_size = 0 - for i in split_size_or_sections: - total_size += i - assert dim_value == total_size - num_sections = len(split_size_or_sections) - - outputs = [] - for i in range(num_sections): - if i > 0: - starts[dim] = starts[dim] + sizes[dim] - sizes[dim] = constant(dims_array([split_size_or_sections[i]])) - outputs.append(slice(tensor, concat(starts), concat(sizes))) - return outputs - - -def chunk(tensor: Tensor, chunks: int, dim: int = 0) -> Tensor: - ''' - Add an operation that splits a tensor into sub-tensors. - - This operation creates a list of tensors that are obtained from the input - tensor by chunking it along the dimension 'dim'. It produces 'chunks' - sub-tensors. - - That operation is only defined for static tensors (no dynamic dimension) - and the size of the tensor in the dimension 'dim' must be a multiple of - 'chunks': 'input.shape[dim] % chunks == 0'. - - It maps to 'split' with 'split_size = input.shape[dim] / chunks'. - - Parameters: - tensor : Tensor - The input tensor to slice. - - chunks : int - The number of slices to split the input tensor into. - - dim : int - The dimension of the tensor to slice. - - Returns: - The list of tensors produced by the different operations. - ''' - assert not tensor.is_dynamic(dim) - - ndim = tensor.ndim() - if dim < 0: - dim += ndim - dim_value = tensor.size()[dim] - assert dim_value % chunks == 0 - - return split(tensor, dim_value // chunks, dim) - - -def unbind(input: Tensor, dim: int = 0): - ''' - Removes a tensor dimension. - - Returns a tuple of all slices along a given dimension, already without it. - ''' - ndim = input.ndim() - outputs = split(input, 1, dim) - output_shape = [input.shape[i] for i in range(ndim) if i != dim] - return [output.view(output_shape) for output in outputs] - - -class AllReduceStrategy(IntEnum): - NCCL = 0 - MIN_LATENCY = 1 - UB = 2 - AUTO = 3 - ONESHOT = 4 - TWOSHOT = 5 - LOWPRECISION = 6 - MNNVL = 7 - NCCL_SYMMETRIC = 8 - SYMM_MEM = 9 # PyTorch symmetric memory with MULTIMEM - - -class AllReduceFusionOp(IntEnum): - NONE = 0 - RESIDUAL_RMS_NORM = 1 - LAST_PROCESS_FOR_UB = 2 - RESIDUAL_RMS_PREPOST_NORM = 3 - RESIDUAL_RMS_NORM_QUANT_FP8 = 4 - RESIDUAL_RMS_NORM_QUANT_NVFP4 = 5 - RESIDUAL_RMS_NORM_OUT_QUANT_FP8 = 6 - RESIDUAL_RMS_NORM_OUT_QUANT_NVFP4 = 7 - MOE_FINALIZE_ALLREDUCE_RESIDUAL_RMS_NORM = 8 - RMS_NORM = 9 - - -class AllReduceParams(): - - def __init__(self, - strategy: AllReduceStrategy = AllReduceStrategy.AUTO, - fusion_op: AllReduceFusionOp = AllReduceFusionOp.NONE, - bias: Optional[Tensor] = None, - residual: Optional[Tensor] = None, - norm_weight: Optional[Tensor] = None, - scale: Optional[Tensor] = None, - norm_pre_residual_weight: Optional[Tensor] = None, - eps: float = 1e-06, - enable_allreduce: bool = True, - trigger_completion_at_end: bool = True): - self.strategy = strategy - self.fusion_op = fusion_op - self.bias = bias - self.residual = residual - self.norm_weight = norm_weight - self.scale = scale - self.norm_pre_residual_weight = norm_pre_residual_weight - self.eps = eps - # For torch path only, has no effect on TRT path - self.enable_allreduce = enable_allreduce - self.trigger_completion_at_end = trigger_completion_at_end - assert fusion_op in (AllReduceFusionOp.NONE.value, - AllReduceFusionOp.RMS_NORM.value) or (residual - is not None) - - def has_affine(self): - return 1 if self.norm_weight is not None else 0 - - def has_bias(self): - return 1 if self.bias is not None else 0 - - def has_scale(self): - return 1 if self.scale is not None else 0 - - def update_strategy(self): - if self.strategy == AllReduceStrategy.AUTO and default_net( - ).plugin_config.user_buffer: - self.strategy = AllReduceStrategy.UB - - -class MoEAllReduceParams(AllReduceParams): - - def __init__(self, - device_num_experts: Optional[Tensor] = None, - expert_scale_factor: Optional[Tensor] = None, - expanded_idx_to_permuted_idx: Optional[Tensor] = None, - shared_expert_output: Optional[Tensor] = None, - bias: Optional[Tensor] = None, - residual: Optional[Tensor] = None, - norm_weight: Optional[Tensor] = None, - scale: Optional[Tensor] = None, - norm_pre_residual_weight: Optional[Tensor] = None, - eps: float = 1e-06, - enable_allreduce: bool = True, - is_cutlass_min_latency: bool = False): - super().__init__( - bias=bias, - residual=residual, - norm_weight=norm_weight, - scale=scale, - norm_pre_residual_weight=norm_pre_residual_weight, - eps=eps, - enable_allreduce=enable_allreduce, - ) - self.device_num_experts = device_num_experts - self.expert_scale_factor = expert_scale_factor - self.expanded_idx_to_permuted_idx = expanded_idx_to_permuted_idx - self.shared_expert_output = shared_expert_output - self.is_cutlass_min_latency = is_cutlass_min_latency - - def is_valid(self): - if self.is_cutlass_min_latency: - return (self.device_num_experts is not None - and self.expert_scale_factor is not None - and self.shared_expert_output is not None) - else: - return (self.expanded_idx_to_permuted_idx is not None) - - -def create_allreduce_plugin( - network: trt.INetworkDefinition, - tensor: trt.ITensor, - workspace: Optional[trt.ITensor], - group: np.array, - dtype: trt.DataType, - all_reduce_params: AllReduceParams, -): - allreduce_plg_creator = trt.get_plugin_registry().get_plugin_creator( - 'AllReduce', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert allreduce_plg_creator is not None - - pf_group = trt.PluginField("group", group, trt.PluginFieldType.INT32) - pf_dtype = trt.PluginField("type_id", np.array([int(dtype)], np.int32), - trt.PluginFieldType.INT32) - pfc = [pf_group, pf_dtype] - p_strategy = trt.PluginField( - "strategy", np.array([int(all_reduce_params.strategy)], np.int8), - trt.PluginFieldType.INT8) - pfc.append(p_strategy) - p_fusion_op = trt.PluginField( - "fusion_op", np.array([int(all_reduce_params.fusion_op)], np.int8), - trt.PluginFieldType.INT8) - pfc.append(p_fusion_op) - p_eps = trt.PluginField( - "eps", np.array([float(all_reduce_params.eps)], np.float32), - trt.PluginFieldType.FLOAT32) - pfc.append(p_eps) - p_affine = trt.PluginField( - "affine", np.array([int(all_reduce_params.has_affine())], np.int8), - trt.PluginFieldType.INT8) - pfc.append(p_affine) - p_bias = trt.PluginField( - "bias", np.array([int(all_reduce_params.has_bias())], np.int8), - trt.PluginFieldType.INT8) - pfc.append(p_bias) - p_scale = trt.PluginField( - "scale", np.array([int(all_reduce_params.has_scale())], np.int8), - trt.PluginFieldType.INT8) - pfc.append(p_scale) - - pfc = trt.PluginFieldCollection(pfc) - ar_plug = allreduce_plg_creator.create_plugin("allreduce", pfc) - plug_inputs = [tensor] - if all_reduce_params.strategy not in { - AllReduceStrategy.NCCL, AllReduceStrategy.UB, - AllReduceStrategy.NCCL_SYMMETRIC - }: - plug_inputs.append(workspace) - if all_reduce_params.fusion_op != AllReduceFusionOp.NONE: - if all_reduce_params.has_bias() == 1: - plug_inputs.append(all_reduce_params.bias.trt_tensor) - if all_reduce_params.residual is not None: - plug_inputs.append(all_reduce_params.residual.trt_tensor) - if all_reduce_params.has_affine() == 1: - plug_inputs.append(all_reduce_params.norm_weight.trt_tensor) - if all_reduce_params.fusion_op == AllReduceFusionOp.RESIDUAL_RMS_PREPOST_NORM: - plug_inputs.append( - all_reduce_params.norm_pre_residual_weight.trt_tensor) - if all_reduce_params.has_scale() == 1: - plug_inputs.append(all_reduce_params.scale.trt_tensor) - - layer = network.add_plugin_v2(plug_inputs, ar_plug) - return layer, allreduce_plg_creator, pfc - - -allreduce_ub_counter = 0 - - -def allreduce( - tensor: Tensor, - group: List[int], - all_reduce_params: Optional[AllReduceParams] = AllReduceParams() -) -> Tensor: - ''' - Add an operation that performs a collective all-reduce. - - Let's define 'world_size' as the length of the 'group' list. That functions - creates a layer to compute the sum of 'world_size' tensors distributed - amongst the 'world_size' participating ranks (one GPU per rank). - - The list 'group' contains the identifiers of the ranks participating into - the collective operation. - - The tensors in the different ranks must be 1D tensors (or views) and the output - tensor will have that same shape. The output tensor will be replicated on - the 'world_size' ranks. - - That operation is implemented using a plugin that wraps the NCCL all-reduce - collective operation. See - https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/usage/collectives.html#allreduce - for details. - - Parameters: - tensor : Tensor - The input tensor. - - group : List[int] - The ranks participating into the all-reduce operation. - - strategy: AllReduceStrategy - NCCL delegates all-reduce to NCCL while ONESHOT and TWOSHOT are custom latency-optimal algorithms. - AUTO chooses amongst the three based on a message-size heuristic. - - Returns: - The tensor produced by that layer. - ''' - - global allreduce_ub_counter - allreduce_ub_counter += 1 - - if all_reduce_params is None: - all_reduce_params = AllReduceParams() - all_reduce_params.update_strategy() - - # TODO(TRTLLM-996): remove this WAR when custom allreduce is supported - # for encoder models in C++ runtime. - workspace = None - if all_reduce_params.strategy != AllReduceStrategy.NCCL and all_reduce_params.strategy != AllReduceStrategy.UB: - if current_all_reduce_helper().workspace is None: - all_reduce_params.strategy = AllReduceStrategy.NCCL_SYMMETRIC - else: - workspace = current_all_reduce_helper().workspace.trt_tensor - if all_reduce_params.strategy == AllReduceStrategy.UB: - tensor.mark_output("allreduce_ub_0_" + str(allreduce_ub_counter)) - dtype = default_net().plugin_config.nccl_plugin - layer, allreduce_plg_creator, pfc = create_allreduce_plugin( - network=default_trtnet(), - tensor=tensor.cast(dtype).trt_tensor, - workspace=workspace, - group=np.array(group, dtype=np.int32), - dtype=str_dtype_to_trt(dtype), - all_reduce_params=all_reduce_params, - ) - _add_plugin_info(layer, allreduce_plg_creator, "allreduce", pfc) - if all_reduce_params.fusion_op != AllReduceFusionOp.NONE: - inter_output = _create_tensor(layer.get_output(1), - layer).cast(tensor.dtype) - if all_reduce_params.strategy == AllReduceStrategy.UB and all_reduce_params.has_scale( - ) == 1: - final_output = _create_tensor(layer.get_output(0), layer) - if all_reduce_params.fusion_op == AllReduceFusionOp.RESIDUAL_RMS_NORM_QUANT_NVFP4: - scale_factor = _create_tensor(layer.get_output(2), layer) - else: - final_output = _create_tensor(layer.get_output(0), - layer).cast(tensor.dtype) - if all_reduce_params.strategy == AllReduceStrategy.UB: - if all_reduce_params.has_scale() == 1: - final_output.mark_output("allreduce_ub_1_" + - str(allreduce_ub_counter)) - if all_reduce_params.fusion_op == AllReduceFusionOp.RESIDUAL_RMS_NORM_QUANT_NVFP4: - scale_factor.mark_output("allreduce_ub_2_" + - str(allreduce_ub_counter)) - return (final_output, scale_factor), inter_output - else: - assert all_reduce_params.fusion_op == AllReduceFusionOp.LAST_PROCESS_FOR_UB - inter_output.mark_output("allreduce_ub_1_" + - str(allreduce_ub_counter)) - return final_output, inter_output - else: - final_output = _create_tensor(layer.get_output(0), - layer).cast(tensor.dtype) - return final_output - - -def allgather(tensor: Tensor, group: List[int], gather_dim: int = 0) -> Tensor: - ''' - Add an operation that performs a collective all-gather. - - Let's define 'group_size' as the length of the 'group' list. That functions - creates a layer to gather 'group_size' tensors distributed - amongst the 'group_size' participating ranks (one GPU per rank). - - The list 'group' contains the identifiers of the ranks participating into - the collective operation. - - Note that 'group' here can be either TP group or PP group, because allgather communication is not limited to a specific split pattern. Therefore 'group_size' does not need to equal MPI 'world_size'. - - The tensors in the different ranks must be 1D tensors (or views) and the - output tensor will have that same shape. - - Given the 'section_size = input.shape[0] / group_size', each rank - contributes a section of its input tensor that correspond to - 'rank*section_size:(rank+1)*section_size'. - - That operation is implemented using a plugin that wraps the NCCL all-gather - collective operation. See - https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/usage/collectives.html#allgather - for details. - - Parameters: - tensor : Tensor - The input tensor. - - group : List[int] - The ranks participating into the all-gather operation. - - gather_dim: int = 0 - Gather along given dimension. By default 0, i.e. treated as 1D tensor. - - Returns: - The tensor produced by that layer. - ''' - allgather_plg_creator = trt.get_plugin_registry().get_plugin_creator( - 'AllGather', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert allgather_plg_creator is not None - - group_size = len(group) - group = trt.PluginField("group", np.array(group, dtype=np.int32), - trt.PluginFieldType.INT32) - - p_dtype = default_net().plugin_config.nccl_plugin - pf_type = trt.PluginField( - "type_id", np.array([int(str_dtype_to_trt(p_dtype))], np.int32), - trt.PluginFieldType.INT32) - - pfc = trt.PluginFieldCollection([group, pf_type]) - allgather = allgather_plg_creator.create_plugin("allgather", pfc) - plug_inputs = [tensor.cast(p_dtype).trt_tensor] - - layer = default_trtnet().add_plugin_v2(plug_inputs, allgather) - _add_plugin_info(layer, allgather_plg_creator, "allgather", pfc) - - x = _create_tensor(layer.get_output(0), layer).cast(tensor.dtype) - - # gather along a given dimension other than dim0 - if gather_dim != 0: - # also support -1 type of dim representation - if gather_dim < 0: - gather_dim = x.ndim() + gather_dim - - # plugin above gathers as 1D flattened tensor - # 1. [dim0, ...dimi, ...dimN] -> [group_size * dim0, ...dimi, ...dimN] - - # now we need to gather-by-dim via split-concat - # 2. [group_size * dim0, ...dimi, ...dimN] -> [dim0, ...group_size * dimi, ...dimN] - # 2.1 split - split_size = shape(x, dim=0) / group_size - ndim = x.ndim() - starts = [constant(dims_array([0])) for _ in range(ndim)] - sizes = [shape(x, dim=d) for d in range(ndim)] - sizes[0] = split_size - sections = [] - for i in range(group_size): - starts[0] = split_size * i - sections.append(slice(x, concat(starts), concat(sizes))) - # 2.2 concat - x = concat(sections, dim=gather_dim) - - return x - - -def reduce_scatter(tensor: Tensor, group: List[int]) -> Tensor: - - plg_creater = trt.get_plugin_registry().get_plugin_creator( - 'ReduceScatter', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert plg_creater is not None - - p_dtype = default_net().plugin_config.nccl_plugin - pf_type = trt.PluginField( - "type_id", np.array([int(str_dtype_to_trt(p_dtype))], np.int32), - trt.PluginFieldType.INT32) - group = trt.PluginField("group", np.array(group, dtype=np.int32), - trt.PluginFieldType.INT32) - pfc = trt.PluginFieldCollection([group, pf_type]) - - reduce_scatter_plug = plg_creater.create_plugin("reduce_scatter", pfc) - plug_inputs = [tensor.cast(p_dtype).trt_tensor] - - layer = default_trtnet().add_plugin_v2(plug_inputs, reduce_scatter_plug) - _add_plugin_info(layer, plg_creater, "reduce_scatter", pfc) - - return _create_tensor(layer.get_output(0), layer).cast(tensor.dtype) - - -def send(tensor: Tensor, tgt: int) -> Tensor: - ''' - Add an operation that performs a send from a rank to another. - - The send operation sends a tensor from one rank to another. If a rank 'i' - sends a tensor to a rank 'j', the rank 'j' must have a corresponding 'recv' - operation from rank 'i'. See 'recv'. - - That operation is implemented using a plugin that wraps the NCCL send - point-to-point operation. See - https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/api/p2p.html#ncclsend - for details. - - Parameters: - tensor : Tensor - The input tensor. - - tgt : int - The rank that receives the tensor. - - Returns: - The tensor produced by that layer. - ''' - send_plg_creator = trt.get_plugin_registry().get_plugin_creator( - 'Send', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert send_plg_creator is not None - - tgt = trt.PluginField("tgt_rank", np.array(tgt, dtype=np.int32), - trt.PluginFieldType.INT32) - - p_dtype = default_net().plugin_config.nccl_plugin - pf_type = trt.PluginField( - "type_id", np.array([int(str_dtype_to_trt(p_dtype))], np.int32), - trt.PluginFieldType.INT32) - - pfc = trt.PluginFieldCollection([tgt, pf_type]) - send_plug = send_plg_creator.create_plugin("send", pfc) - plug_inputs = [tensor.cast(p_dtype).trt_tensor] - - layer = default_trtnet().add_plugin_v2(plug_inputs, send_plug) - _add_plugin_info(layer, send_plg_creator, "send", pfc) - return _create_tensor(layer.get_output(0), layer).cast(tensor.dtype) - - -def recv(tensor: Tensor, src: int) -> Tensor: - ''' - Add an operation that performs a recv to a rank from another. - - The recv operation receives a tensor from on a rank from another. If a rank 'i' - receives a tensor from a rank 'j', the rank 'j' must have a corresponding 'send' - operation to rank 'j'. See 'send'. - - That operation is implemented using a plugin that wraps the NCCL recv - point-to-point operation. See - https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/api/p2p.html#ncclrecv - for details. - - Parameters: - tensor : Tensor - The input tensor. - - src : int - The rank that sends the tensor to. - - Returns: - The tensor produced by that layer. - ''' - recv_plg_creator = trt.get_plugin_registry().get_plugin_creator( - 'Recv', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert recv_plg_creator is not None - - src = trt.PluginField("src_rank", np.array(src, dtype=np.int32), - trt.PluginFieldType.INT32) - p_dtype = default_net().plugin_config.nccl_plugin - pf_type = trt.PluginField( - "type_id", np.array([int(str_dtype_to_trt(p_dtype))], np.int32), - trt.PluginFieldType.INT32) - - pfc = trt.PluginFieldCollection([src, pf_type]) - recv_plug = recv_plg_creator.create_plugin("recv", pfc) - plug_inputs = [tensor.cast(p_dtype).trt_tensor] - - layer = default_trtnet().add_plugin_v2(plug_inputs, recv_plug) - _add_plugin_info(layer, recv_plg_creator, "recv", pfc) - return _create_tensor(layer.get_output(0), layer).cast(tensor.dtype) - - -def gemm_allreduce(a: Tensor, - b: Tensor, - group: List[int], - transa: bool = False, - transb: bool = False, - alpha: Optional[Union[np.ndarray, Tensor]] = None, - output_dtype: Optional[trt.DataType] = None, - fp8_inputs_override: bool = False, - a_sf: Optional[Tensor] = None, - b_sf: Optional[Tensor] = None): - ''' - Add an operation that performs fused GEMM+AllReduce. - - Parameters: - a: Tensor - Input tensor A - b: Tensor - Input tensor B - a_sf: Optional[Tensor] - Input tensor for scaling input A - b_sf: Optional[Tensor] - Input tensor for scaling input B - group: List[int] - Ranks participating in collective - transa: bool - Whether or not input tensor A is transposed - transb: bool - Whether or not input tensor B is transposed - alpha: float - Alpha for GEMM -> beta * C + (alpha * acc) - output_dtype: trt.DataType - Output type for plugin. If it is None, we - will use type set in plugin_config. - fp8_inputs_override: bool - TRT graph does not detect FP8 inputs correctly. This - flag is used to override the derived input tensor - types so that our plugin knows to issue FP8 MMAs. - - Returns: - Returns GEMM output tensor which has been reduced across ranks. - ''' - - # Output tensor needs to be bound to externally managed - # memory so keep track of layer index so we can assign - # output tensor unique label. - if not hasattr(gemm_allreduce, 'layer_idx'): - gemm_allreduce.layer_idx = 0 - - # Check inputs - assert isinstance(a.dtype, trt.DataType) - assert isinstance(b.dtype, trt.DataType) - - if fp8_inputs_override: - assert ( - isinstance(alpha, np.ndarray) and alpha.dtype == np.float32 - and alpha.size == 1 - ), "`alpha` must be passed as a float32 ndarray if `fp8_inputs_override` is enabled for gemm_allreduce_plugin" - assert a.dtype == trt.fp8 - assert b.dtype == trt.fp8 - - if output_dtype is None: - output_dtype = str_dtype_to_trt( - default_net().plugin_config.gemm_allreduce_plugin) - assert output_dtype in [trt.float16, trt.bfloat16] - - alpha_is_tensor = isinstance(alpha, Tensor) - if alpha is None or alpha_is_tensor: - alpha_value = np.array(1.0, dtype=np.float32) - else: - alpha_value = alpha - - plugin_creator = trt.get_plugin_registry().get_plugin_creator( - 'GemmAllReduce', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert plugin_creator is not None - - trt_type_a = trt.fp8 if fp8_inputs_override else a.dtype - trt_type_b = trt.fp8 if fp8_inputs_override else b.dtype - - # create plugin fields - field_list = [] - field_list.append( - trt.PluginField('type_a', np.array([int(trt_type_a)], np.int32), - trt.PluginFieldType.INT32)) - field_list.append( - trt.PluginField('type_b', np.array([int(trt_type_b)], np.int32), - trt.PluginFieldType.INT32)) - field_list.append( - trt.PluginField('type_d', np.array([int(output_dtype)], np.int32), - trt.PluginFieldType.INT32)) - field_list.append( - trt.PluginField('transa', np.array(transa, dtype=np.int32), - trt.PluginFieldType.INT32)) - field_list.append( - trt.PluginField('transb', np.array(transb, dtype=np.int32), - trt.PluginFieldType.INT32)) - field_list.append( - trt.PluginField('group', np.array(group, dtype=np.int32), - trt.PluginFieldType.INT32)) - field_list.append( - trt.PluginField('has_sfa', np.array([int(a_sf is not None)], np.int8), - trt.PluginFieldType.INT8)) - field_list.append( - trt.PluginField('has_sfb', np.array([int(b_sf is not None)], np.int8), - trt.PluginFieldType.INT8)) - field_list.append( - trt.PluginField('alpha_is_ptr', np.array([int(alpha_is_tensor)], - np.int8), - trt.PluginFieldType.INT8)) - field_list.append( - trt.PluginField('alpha', alpha_value.flatten(), - trt.PluginFieldType.FLOAT32)) - - # create plugin - fields = trt.PluginFieldCollection(field_list) - plugin = plugin_creator.create_plugin("gemm_allreduce", fields) - # define symbolic input tensors. - # note this does NOT allocate memory. - inputs = [a.trt_tensor, b.trt_tensor] - if a_sf is not None: - inputs += [a_sf.trt_tensor] - if b_sf is not None: - inputs += [b_sf.trt_tensor] - if alpha_is_tensor: - inputs += [alpha.trt_tensor] - - layer = default_trtnet().add_plugin_v2(inputs, plugin) - _add_plugin_info(layer, plugin_creator, "gemm_allreduce", fields) - # define symbolic output tensors - # both output tensors point to same physical memory but - # one has unicast address and other has multicast address - uc_output = _create_tensor(layer.get_output(0), layer) - mc_output = _create_tensor(layer.get_output(1), layer) - ipc_output = _create_tensor(layer.get_output(2), layer) - assert uc_output is not None - assert mc_output is not None - assert ipc_output is not None - # mark outputs so that we can bind our own allocated memory in runtime - # (see generation.py) - uc_output.mark_output(f'gemm_allreduce_uc_out_{gemm_allreduce.layer_idx}') - mc_output.mark_output(f'gemm_allreduce_mc_out_{gemm_allreduce.layer_idx}') - ipc_output.mark_output(f'gemm_allreduce_ipc_out_{gemm_allreduce.layer_idx}') - gemm_allreduce.layer_idx += 1 - - return uc_output - - -def bert_attention(tensor: Tensor, - input_lengths: Tensor, - num_heads: int, - head_size: int, - q_scaling: float, - relative_attention: bool = False, - relative_attention_bias: Tensor = None, - max_distance: int = 0, - max_input_length: Tensor = None, - sage_attn: bool = False, - sage_attn_q_block_size: int = 0, - sage_attn_k_block_size: int = 0, - sage_attn_v_block_size: int = 0, - cp_group: list[int] = None, - cp_size: int = 1, - cp_rank: int = 0) -> Tuple[Tensor]: - ''' - Add an operation that performs the multi-head attention in BERT. - - The multi-head attention (MHA) is the sequence of a batched matmul, a - softmax and a batched matmul as described in - https://arxiv.org/abs/1706.03762. That function adds an operation that - performs those computations using a single GPU kernel. - - The input tensor contains the Q, K and V elements. It is a 2D tensor and - its shape is '[sum_of_tokens, 3*hidden_dim]' where the 'sum_of_tokens' is - the sum of the sequence lengths in the batch. - - In MHA, the output of the Q*K^T product is scaled by a constant value that - is computed as: - - 1.f / (q_scaling * sqrt(head_size)). - - That 'q_scaling' constant is the last argument of that function. - - That layer is implemented using a plugin (see bertAttentionPlugin). - - Parameters: - tensor : Tensor - The QKV input tensor. - - input_lengths : Tensor - The length of each sequence. It is a 1D tensor of size 'batch_size'. - - num_heads : int - The number of heads. - - head_size : int - The size of each head. - - q_scaling : float - The factor to compute the scaling factor to scale the output of the - 'Q*K^T' product. - - relative_attention: bool = False - If enable relative attention. - - relative_attention_bias: Tensor = None - The relative attention bias [num_heads, max_seq_len, max_seq_len], or The relative attention embedding table for implicit mode, [num_heads, num_buckets]. - - max_distance: int = 0 - The maximum distance of relative position in attention, for implicit mode. - Default value is 0, meaning to use the regular mode of relative attention bias. - Implicit mode is only enabled when passing in non-zero positive max_distance value. - See relative attention bias in docs/source/advanced/gpt-attention.md - - max_input_length: Tensor = None - The maximum input sequence length represented by Tensor shape. Requires for remove_input_padding to pre-define plugin workspace size. - - sage_attn: bool = False - SageAttention is a 8-bit implementation of attention kernel. It's input q, k, v and output datatypes are 16-bit. It performance dynamic quantization for q, k, v - tensor every time before attention. https://github.com/thu-ml/SageAttention - - sage_attn_q_quant_size: int = 0 - dynamic quant block size along sequence dimension of q tensor. Each quant block will share one scale. - - sage_attn_k_quant_size: int = 0 - dynamic quant block size along sequence dimension of k tensor. Each quant block will share one scale. - - sage_attn_v_quant_size: int = 0 - dynamic quant block size along sequence dimension of v tensor. Each quant block will share one scale. - - cp_group: list[int] = None - The communication group for context parallel - - cp_size: int = 1 - The communication size for context parallel - - cp_rank: int = 0 - The communication rank for context parallel - - Returns: - The tensor produced by that layer. - ''' - attn_plg_creator = trt.get_plugin_registry().get_plugin_creator( - 'BertAttention', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert attn_plg_creator is not None - - nheads = trt.PluginField("num_heads", np.array(num_heads, dtype=np.int32), - trt.PluginFieldType.INT32) - head_size = trt.PluginField("head_size", np.array(head_size, - dtype=np.int32), - trt.PluginFieldType.INT32) - q_scaling = trt.PluginField("q_scaling", - np.array(q_scaling, dtype=np.float32), - trt.PluginFieldType.FLOAT32) - context_fmha_type = trt.PluginField( - "context_fmha_type", - np.array(np.int8(default_net().plugin_config.context_fmha_type), - dtype=np.int8), trt.PluginFieldType.INT8) - p_dtype = default_net().plugin_config.bert_attention_plugin - pf_type = trt.PluginField( - "type_id", np.array([int(str_dtype_to_trt(p_dtype))], np.int32), - trt.PluginFieldType.INT32) - do_relative_attention = trt.PluginField( - "do_relative_attention", - np.array(np.int8(relative_attention), dtype=np.int8), - trt.PluginFieldType.INT8) - max_distance = trt.PluginField("max_distance", - np.array(max_distance, dtype=np.int32), - trt.PluginFieldType.INT32) - remove_padding = trt.PluginField( - "remove_padding", - np.array(np.int8(default_net().plugin_config.remove_input_padding), - dtype=np.int8), trt.PluginFieldType.INT8) - - sage_attn = trt.PluginField("sage_attn", - np.array(np.int8(sage_attn), dtype=np.int8), - trt.PluginFieldType.INT8) - - sage_attn_q_block_size = trt.PluginField( - "sage_attn_q_block_size", - np.array(sage_attn_q_block_size, dtype=np.int32), - trt.PluginFieldType.INT32) - - sage_attn_k_block_size = trt.PluginField( - "sage_attn_k_block_size", - np.array(sage_attn_k_block_size, dtype=np.int32), - trt.PluginFieldType.INT32) - - sage_attn_v_block_size = trt.PluginField( - "sage_attn_v_block_size", - np.array(sage_attn_v_block_size, dtype=np.int32), - trt.PluginFieldType.INT32) - - if cp_size > 1: - # transpose q,k,v inside qkv to make kv contiguous, which is required by ring attention - # (b, s, 3d) - query, key, value = chunk(tensor, 3, dim=-1) - bs = shape(query, 0) - seq_len = shape(query, 1) - # (b, s, d) -> (b, s, 2d) -> (2b, s, d) - kv = concat([key, value], - dim=-1).view(concat((2 * bs, seq_len, query.shape[-1]))) - tensor = concat((query, kv), - dim=0).view(concat((bs, seq_len, query.shape[-1] * 3))) - - cp_size = trt.PluginField("cp_size", np.array(cp_size, dtype=np.int32), - trt.PluginFieldType.INT32) - cp_rank = trt.PluginField("cp_rank", np.array(cp_rank, dtype=np.int32), - trt.PluginFieldType.INT32) - cp_group = cp_group or [0] - cp_group = np.array(cp_group, dtype=np.int32) - cp_group = trt.PluginField("cp_group", cp_group, trt.PluginFieldType.INT32) - - pfc = trt.PluginFieldCollection([ - nheads, head_size, q_scaling, context_fmha_type, pf_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 - ]) - - attn_plug = attn_plg_creator.create_plugin("padding_attn", pfc) - plug_inputs = [tensor, input_lengths] - if max_input_length is not None: - # for remove padding mode - plug_inputs += [max_input_length] - if relative_attention_bias is not None: - # for relative attention mode - plug_inputs += [relative_attention_bias] - - plug_inputs = [i.trt_tensor for i in plug_inputs] - - layer = default_trtnet().add_plugin_v2(plug_inputs, attn_plug) - _add_plugin_info(layer, attn_plg_creator, "padding_attn", pfc) - assert layer.num_outputs == 1, \ - f"Plugin outputs number mismatch with expected, got {layer.num_outputs}, expected 1" - output = _create_tensor(layer.get_output(0), layer) - assert output is not None - return output - - -class RopeEmbeddingUtils: - - @staticmethod - # ref: https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_rope_utils.py#L298 - def apply_llama3_scaling(inv_freqs: np.ndarray, rope_scaling_config: dict): - - scale_factor = rope_scaling_config.get("factor", 8.0) - low_freq_factor = rope_scaling_config.get("low_freq_factor", 1.0) - high_freq_factor = rope_scaling_config.get("high_freq_factor", 4.0) - old_context_len = rope_scaling_config.get( - "original_max_position_embeddings", 8192) - - low_freq_wavelen = old_context_len / low_freq_factor - high_freq_wavelen = old_context_len / high_freq_factor - new_inv_freqs = [] - for inv_freq in inv_freqs: - wavelen = 2 * math.pi / inv_freq - if wavelen < high_freq_wavelen: - new_inv_freqs.append(inv_freq) - elif wavelen > low_freq_wavelen: - new_inv_freqs.append(inv_freq / scale_factor) - else: - assert low_freq_wavelen != high_freq_wavelen - smooth = (old_context_len / wavelen - low_freq_factor) / ( - high_freq_factor - low_freq_factor) - new_inv_freqs.append((1 - smooth) * inv_freq / scale_factor + - smooth * inv_freq) - return np.array(new_inv_freqs, dtype=inv_freqs.dtype) - - @staticmethod - def create_sinusoidal_positions(num_pos: int, - dim: int, - theta: float = 10000.0, - dtype=np.float32): - inv_freq = 1.0 / (theta**(np.arange(0, dim, 2) / dim)).astype(dtype) - sinusoid_inp = np.einsum("i , j -> i j", - np.arange(num_pos, dtype=dtype), - inv_freq, - dtype=dtype) - concat = np.concatenate((np.sin(sinusoid_inp), np.cos(sinusoid_inp)), - axis=1) - return np.expand_dims(concat, axis=0).astype(dtype) - - @staticmethod - def create_sinusoidal_positions_for_attention_plugin( - num_pos: int, - dim: int, - theta: float = 10000.0, - scale: float = 1.0, - scale_type: RotaryScalingType = RotaryScalingType.none, - # Other scaling configs that only used by certain scaling types. - rope_scaling_config: dict = None, - duplicate_data: bool = False, - dtype=np.float32): - if scale_type == RotaryScalingType.linear: - scale = 1.0 / scale - if scale_type == RotaryScalingType.llama3: - assert rope_scaling_config is not None, "rotary_scaling config must be provided." - inv_freq = 1.0 / (theta**(np.arange(0, dim, 2) / dim)).astype(dtype) - inv_freq = RopeEmbeddingUtils.apply_llama3_scaling( - inv_freq, rope_scaling_config) - elif scale_type == RotaryScalingType.dynamic: - # Make sure scaling_alpha exists in rope_scaling - # Ref: https://huggingface.co/tencent/Hunyuan-A13B-Instruct-FP8/blob/main/modeling_hunyuan.py#L346 - assert rope_scaling_config[ - "alpha"] is not None, "rope_scaling_config.alpha must be provided." - scaling_alpha = rope_scaling_config["alpha"] - adjusted_base = theta * (scaling_alpha**(dim / (dim - 2))) - inv_freq = 1.0 / (adjusted_base**( - np.arange(0, dim, 2, dtype=dtype) / dim)).astype(dtype) - else: - inv_freq = scale / (theta - **(np.arange(0, dim, 2) / dim)).astype(dtype) - sinusoid_inp = np.expand_dims(np.einsum("i , j -> i j", - np.arange(num_pos, dtype=dtype), - inv_freq, - dtype=dtype), - axis=-1) - if duplicate_data: - sinusoid_inp = np.concatenate((sinusoid_inp, sinusoid_inp), axis=-2) - # fuse cos/sin into float2 (cos, sin). - concat = np.concatenate( - (np.cos(sinusoid_inp), np.sin(sinusoid_inp)), - axis=-1) #np.cos(sinusoid_inp).shape = (32768, 64, 1) - - return inv_freq, concat.reshape(1, -1).astype(dtype) - - @staticmethod - def create_sinusoidal_positions_for_cogvlm_attention_plugin( - num_pos: int, - dim: int, - theta: float = 10000.0, - scale: float = 1.0, - scale_type: RotaryScalingType = RotaryScalingType.none, - vision_start: int = 1, - vision_length: int = 1225, - dtype=np.float32): - if scale_type == RotaryScalingType.linear: - scale = 1.0 / scale - inv_freq = scale / (theta**(np.arange(0, dim, 2) / dim)).astype(dtype) - position_id = np.hstack([ - np.arange(0, vision_start + 1, dtype=dtype), - np.full(vision_length, vision_start + 1, dtype=dtype), - np.arange(vision_start + 2, - num_pos - (vision_length - 1), - dtype=dtype) - ]) - sinusoid_inp = np.expand_dims(np.einsum("i , j -> i j", - position_id, - inv_freq, - dtype=dtype), - axis=-1) - # fuse cos/sin into float2 (cos, sin). - concat = np.concatenate((np.cos(sinusoid_inp), np.sin(sinusoid_inp)), - axis=-1) - - return inv_freq, concat.reshape(1, -1).astype(dtype) - - def create_sinusoidal_positions_long_rope_for_attention_plugin( - num_pos: int, - num_orig_pos: int, - dim: int, - theta: float = 10000.0, - scaling_short_factors: Tensor = 1.0, - scaling_long_factors: Tensor = 1.0, - short_mscale=None, - long_mscale=None, - dtype=np.float32): - - def _calc_mscale(scale): - if scale <= 1.0: - return 1.0 - return math.sqrt(1 + math.log(scale) / math.log(num_orig_pos)) - - if short_mscale is None: - short_mscale = _calc_mscale(num_pos / num_orig_pos) - long_mscale = short_mscale - - def _compute_sinusoidal_positions(scale_factors, is_short, - for_attention_plugin): - inv_freq = 1 / (scale_factors * - (theta**(np.arange(0, dim, 2) / dim)).astype(dtype)) - sinusoid_inp = np.einsum("i , j -> i j", - np.arange(num_pos, dtype=dtype), - inv_freq, - dtype=dtype) - - if for_attention_plugin: - sinusoid_inp = np.expand_dims(sinusoid_inp, axis=-1) - concat = np.concatenate( - (np.cos(sinusoid_inp), np.sin(sinusoid_inp)), axis=-1) - else: - concat = np.concatenate( - (np.sin(sinusoid_inp), np.cos(sinusoid_inp)), axis=1) - concat = np.expand_dims(concat, axis=0) - - mscale = short_mscale if is_short else long_mscale - concat = concat.astype(dtype) * mscale - - # gpt attention plugins also need inv_freq. - if for_attention_plugin: - return inv_freq.reshape(1, -1), concat.reshape(1, -1) - else: - return concat - - return _compute_sinusoidal_positions( - scaling_short_factors, True, False), _compute_sinusoidal_positions( - scaling_long_factors, - False, False), _compute_sinusoidal_positions( - scaling_short_factors, True, - True), _compute_sinusoidal_positions( - scaling_long_factors, False, True), short_mscale - - @staticmethod - def create_sinusoidal_positions_long_rope(num_pos: int, - dim: int, - theta: float, - original_max_pos: int, - short_factor: List[float], - long_factor: List[float], - dtype=np.float32, - max_seq_len: Optional[int] = None, - duplicate_data: bool = False): - short_factor = np.array(short_factor, dtype=np.float32) - long_factor = np.array(long_factor, dtype=np.float32) - - inv_freq = 1.0 / (theta**(np.arange(0, dim, 2, dtype=np.float32) / dim)) - t_pos = np.arange(np.max([num_pos, original_max_pos]), dtype=np.float32) - - # Choose proper freqs based on max_seq_len. - factor = long_factor if max_seq_len is None or max_seq_len > original_max_pos else short_factor - inv_freq = inv_freq / factor - freqs = np.einsum("i,j->ij", t_pos, inv_freq) - sinusoid_inp = freqs.astype(np.float32)[..., np.newaxis] - - # Apply scaling - scale = num_pos / original_max_pos - if scale <= 1.0: - scaling_factor = 1.0 - else: - scaling_factor = np.sqrt(1.0 + - np.log(scale) / np.log(original_max_pos)) - - if duplicate_data: - sinusoid_inp = np.concatenate((sinusoid_inp, sinusoid_inp), axis=-2) - - # fuse cos/sin into float2 (cos, sin). - concat = np.concatenate( - (np.cos(sinusoid_inp) * scaling_factor, - np.sin(sinusoid_inp) * scaling_factor), - axis=-1, - ) - - return None, concat.reshape(1, -1).astype(dtype) - - @staticmethod - def create_fake_weight(dim: int, dtype=np.half): - return np.random.rand(dim).astype(dtype) - - # Note: When not using deepseek_yarn, make sure to set mscale_all_dim to 0.0. - @staticmethod - def create_sinusoidal_positions_yarn( - num_pos: int, - dim: int, - base: int = 10000, - scaling_factor: float = 1.0, - original_max_position_embeddings: int = 4096, - beta_fast: int = 32, - beta_slow: int = 1, - mscale: float = 1.0, - mscale_all_dim: float = 1.0, - duplicate_data: bool = True, - dtype=torch.float32): - - # Copy from https://huggingface.co/deepseek-ai/DeepSeek-V2/blob/main/modeling_deepseek.py - # Inverse dim formula to find dim based on number of rotations - def yarn_find_correction_dim(num_rotations, dim, base, - max_position_embeddings): - return (dim * math.log(max_position_embeddings / - (num_rotations * 2 * math.pi))) / ( - 2 * math.log(base)) - - # Find dim range bounds based on rotations - def yarn_find_correction_range(low_rot, high_rot, dim, base, - max_position_embeddings): - low = math.floor( - yarn_find_correction_dim(low_rot, dim, base, - max_position_embeddings)) - high = math.ceil( - yarn_find_correction_dim(high_rot, dim, base, - max_position_embeddings)) - if low < 0: - low = 0 - if high > dim - 1: - high = dim - 1 - return low, high # Clamp values just in case - - def yarn_get_mscale(scale, mscale): - if scale <= 1: - return 1.0 - return 0.1 * mscale * math.log(scale) + 1.0 - - def yarn_linear_ramp_mask(min, max, dim): - if min == max: - max += 0.001 # Prevent singularity - - linear_func = (torch.arange(dim, dtype=dtype) - min) / (max - min) - ramp_func = torch.clamp(linear_func, 0, 1) - return ramp_func - - pos_freqs = base**(torch.arange(0, dim, 2, dtype=dtype) / dim) - freq_extra = 1.0 / pos_freqs - freq_inter = 1.0 / (scaling_factor * pos_freqs) - - low, high = yarn_find_correction_range( - beta_fast, - beta_slow, - dim, - base, - original_max_position_embeddings, - ) - inv_freq_mask = (1 - yarn_linear_ramp_mask(low, high, dim // 2)) - inv_freq = freq_inter * (1 - inv_freq_mask) + freq_extra * inv_freq_mask - t = torch.arange(num_pos, dtype=dtype) - sinusoid_inp = torch.einsum("i,j -> ij", t, inv_freq).unsqueeze(-1) - - _mscale = float( - yarn_get_mscale(scaling_factor, mscale) / - yarn_get_mscale(scaling_factor, mscale_all_dim)) - - if duplicate_data: - emb = torch.cat((sinusoid_inp, sinusoid_inp), dim=-2) - else: - emb = sinusoid_inp - - concat = torch.cat((torch.cos(emb) * _mscale, torch.sin(emb) * _mscale), - dim=-1) - return inv_freq.numpy(), concat.reshape((1, -1)).to(dtype).numpy() - - @staticmethod - def rotate_every_two(tensor: Tensor) -> Tensor: - assert tensor.ndim() == 4 - - shape_tensor = concat([ - shape(tensor, i) / 2 if i == (tensor.ndim() - - 1) else shape(tensor, i) - for i in range(tensor.ndim()) - ]) - x1 = slice(tensor, [0, 0, 0, 0], shape_tensor, [1, 1, 1, 2]) - x2 = slice(tensor, [0, 0, 0, 1], shape_tensor, [1, 1, 1, 2]) - x1 = expand_dims(x1, 4) - x2 = expand_dims(x2, 4) - zero = constant( - np.ascontiguousarray( - np.zeros([1], dtype=trt_dtype_to_np(tensor.dtype)))) - x2 = zero - x2 - x = concat([x2, x1], 4) - return view( - x, concat([shape(x, 0), - shape(x, 1), - shape(x, 2), - shape(x, 3) * 2])) - - @staticmethod - def rotate_half(tensor: Tensor) -> Tensor: - # [bs, num_attention_kv_heads, seqlen, attention_head_size] - assert tensor.ndim() == 4 - shape_tensor = concat([ - shape(tensor, i) / 2 if i == (tensor.ndim() - - 1) else shape(tensor, i) - for i in range(tensor.ndim()) - ]) - last_dim = shape(tensor, tensor.ndim() - 1) / 2 - x1 = slice(tensor, [0, 0, 0, 0], shape_tensor, [1, 1, 1, 1]) - x2 = slice(tensor, concat([0, 0, 0, last_dim]), shape_tensor, - [1, 1, 1, 1]) - zero = constant( - np.ascontiguousarray( - np.zeros([1], dtype=trt_dtype_to_np(tensor.dtype)))) - x2 = zero - x2 - x = concat([x2, x1], 3) - return x - - @staticmethod - def apply_rotary_pos_emb( - tensor: Tensor, - position_embedding: List[Tensor] = None, - pos_emb_type: PositionEmbeddingType = PositionEmbeddingType.rope_gptj - ) -> Tensor: - - rotate_func = None - if pos_emb_type == PositionEmbeddingType.rope_gpt_neox or pos_emb_type == PositionEmbeddingType.long_rope: - assert len(position_embedding) == 2 - cos, sin = position_embedding - sin = expand_dims(sin, 2) - cos = expand_dims(cos, 2) - sin = concat([sin, sin], 3) - cos = concat([cos, cos], 3) - rotate_func = RopeEmbeddingUtils.rotate_half - elif pos_emb_type == PositionEmbeddingType.rope_gptj: - assert len(position_embedding) == 2 - cos, sin = position_embedding - sin = expand_dims(sin, 2) - cos = expand_dims(cos, 2) - sin = repeat_interleave(sin, 2, 3) - cos = repeat_interleave(cos, 2, 3) - rotate_func = RopeEmbeddingUtils.rotate_every_two - elif pos_emb_type == PositionEmbeddingType.chatglm: - assert len(position_embedding) == 4 - cos0, cos1, sin0, sin1 = position_embedding - shape_tensor = concat([ - shape(tensor, i) / 2 if i == (tensor.ndim() - - 1) else shape(tensor, i) - for i in range(tensor.ndim()) - ]) - last_dim = shape(tensor, tensor.ndim() - 1) / 2 - x_part0 = slice(tensor, [0, 0, 0, 0], shape_tensor, [1, 1, 1, 1]) - x_part1 = slice(tensor, concat([0, 0, 0, last_dim]), shape_tensor, - [1, 1, 1, 1]) - - y_part0 = (x_part0 * - cos0) + (RopeEmbeddingUtils.rotate_half(x_part0) * sin0) - y_part1 = (x_part1 * - cos1) + (RopeEmbeddingUtils.rotate_half(x_part1) * sin1) - - result = concat([y_part0, y_part1], dim=3) - return result.view(shape(tensor)) - - else: - raise ValueError('The PositionEmbeddingType is not RoPE') - return (tensor * cos) + (rotate_func(tensor) * sin) - - @staticmethod - def apply_rotary_pos_emb_chatglm(qkv, position_embedding, - num_attention_heads, attention_head_size, - max_position_embeddings, - rotary_embedding_scale, - remove_input_padding) -> Tensor: - - half_head_size = attention_head_size // 2 - input = qkv[0] if isinstance(qkv, list) else qkv - input_shape = shape(input) - batch_size = 1 if remove_input_padding else shape(input, 0) - seqlen = shape(input, 0 if remove_input_padding else 1) - if isinstance(qkv, list): - query, key, value = qkv - else: - qkv = qkv.view( - concat([ - batch_size, - seqlen, - num_attention_heads, - 3, - attention_head_size, - ])) - query, key, value = split(qkv, 1, dim=3) - q_shape = concat([ - batch_size, - seqlen, - num_attention_heads, - attention_head_size, - ]) - query = query.view(q_shape) - key = key.view(q_shape) - value = value.view(q_shape) - - embedding_weight = RopeEmbeddingUtils.create_sinusoidal_positions( - max_position_embeddings, half_head_size) - embedding_weight /= rotary_embedding_scale - embedding_weight = np.split(embedding_weight.squeeze(0), 2, axis=1) - embedding_weight = np.concatenate( - [ - embedding_weight[0], - embedding_weight[0], - embedding_weight[1], - embedding_weight[1], - ], - axis=1, - ) - - if remove_input_padding: - position_embedding = unsqueeze(position_embedding, 0) - - embedding_weight = embedding_weight.astype(trt_dtype_to_np(query.dtype)) - embedding_weight = constant(embedding_weight) - position_embedding = embedding(position_embedding, embedding_weight) - position_embedding, block_embedding = split( - position_embedding, - 1, - dim=1, - ) - sin0, cos0 = split(position_embedding, half_head_size, dim=3) - sin1, cos1 = split(block_embedding, half_head_size, dim=3) - - new_shape = concat([ - batch_size, - seqlen, - 1, - half_head_size, - ]) - position_embedding = [ - tensor.view(new_shape) for tensor in [cos0, cos1, sin0, sin1] - ] - - query = RopeEmbeddingUtils.apply_rotary_pos_emb( - tensor=query, - position_embedding=position_embedding, - pos_emb_type=PositionEmbeddingType.chatglm) - key = RopeEmbeddingUtils.apply_rotary_pos_emb( - tensor=key, - position_embedding=position_embedding, - pos_emb_type=PositionEmbeddingType.chatglm) - - if isinstance(qkv, list): - qkv = [ - query.view(input_shape), - key.view(input_shape), - value.view(input_shape), - ] - else: - qkv = concat([query, key, value], dim=2) - qkv = qkv.view(input_shape) - - return qkv - - @staticmethod - def apply_rotary_pos_emb_cogvlm(qkv, position_embedding, - num_attention_heads, attention_head_size, - max_position_embeddings, - rotary_embedding_scale, - remove_input_padding) -> Tensor: - input = qkv[0] if isinstance(qkv, list) else qkv - input_shape = shape(input) - batch_size = 1 if remove_input_padding else shape(input, 0) - seqlen = shape(input, 0 if remove_input_padding else 1) - if isinstance(qkv, list): - query, key, value = qkv - else: - qkv = qkv.view( - concat([ - batch_size, - seqlen, - 3, - num_attention_heads, - attention_head_size, - ])) - query, key, value = split(qkv, 1, dim=2) - q_shape = concat([ - batch_size, - seqlen, - num_attention_heads, - attention_head_size, - ]) - query = query.view(q_shape) - key = key.view(q_shape) - value = value.view(q_shape) - - embedding_weight = RopeEmbeddingUtils.create_sinusoidal_positions( - max_position_embeddings, attention_head_size).squeeze(0) - embedding_weight /= rotary_embedding_scale # [max_position_embeddings, attention_head_size] - - if remove_input_padding: - position_embedding = unsqueeze(position_embedding, 0) # [1, seqlen] - - embedding_weight = constant(embedding_weight) # float32 - position_embedding = embedding( - position_embedding, - embedding_weight) # [1, seqlen, attention_head_size] - sin, cos = split(position_embedding, attention_head_size // 2, - dim=-1) # [1, seqlen, attention_head_size//2] - - input_dtype = query.dtype - fp32_query = cast(query, "float32") - fp32_key = cast(key, "float32") - fp32_query = RopeEmbeddingUtils.apply_rotary_pos_emb( - tensor=fp32_query, - position_embedding=[cos, sin], - pos_emb_type=PositionEmbeddingType.rope_gpt_neox) - fp32_key = RopeEmbeddingUtils.apply_rotary_pos_emb( - tensor=fp32_key, - position_embedding=[cos, sin], - pos_emb_type=PositionEmbeddingType.rope_gpt_neox) - - query = cast(fp32_query, input_dtype) - key = cast(fp32_key, input_dtype) - - if isinstance(qkv, list): - qkv = [ - query.view(input_shape), - key.view(input_shape), - value.view(input_shape), - ] - else: - qkv = concat([query, key, value], dim=2) - qkv = qkv.view(input_shape) - - return qkv - - -@gw.record_signature -def gpt_attention( - *, - qkv: Tensor, - past_key_value: Tensor, - attention_mask: Optional[Tensor] = None, - attention_packed_mask: Optional[Tensor] = None, - sequence_length: Tensor, - host_past_key_value_lengths: Optional[Tensor], - host_max_attention_window_sizes: Tensor, - host_sink_token_length: Tensor, - context_lengths: Optional[Tensor], - cache_indirection: Optional[Tensor], - host_request_types: Tensor, - layer_idx: int, - num_heads: int, - num_kv_heads: int, - hidden_size_per_head: int, - q_scaling: float, - attn_logit_softcapping_scale: float = 0.0, - rotary_embedding_dim: int = 0, - rotary_embedding_base: float = 10000.0, - rotary_embedding_scale_type: RotaryScalingType = RotaryScalingType.none, - rotary_embedding_short_m_scale: float = 1.0, - rotary_embedding_long_m_scale: float = 1.0, - rotary_embedding_scale: float = 1.0, - rotary_embedding_max_positions: int = 1024, - rotary_embedding_original_max_positions: int = 1024, - position_embedding_type: PositionEmbeddingType = PositionEmbeddingType. - learned_absolute, - rotary_inv_freq: Optional[Tensor] = None, - rotary_cos_sin: Optional[Tensor] = None, - kv_orig_quant_scale: Optional[Tensor] = None, - kv_quant_orig_scale: Optional[Tensor] = None, - attention_output_orig_quant_scale: Optional[Tensor] = None, - attention_output_sf_scale: Optional[Tensor] = None, - kv_cache_quant_mode: Union[QuantModeWrapper, QuantMode] = QuantMode(0), - max_context_length: Optional[int] = None, - mask_type: AttentionMaskType = AttentionMaskType.causal, - block_sparse_block_size: int = 64, - block_sparse_homo_head_pattern: bool = False, - block_sparse_num_local_blocks: int = 16, - block_sparse_vertical_stride: int = 8, - alibi_slopes: Optional[Tensor] = None, - tp_size: int = 1, - tp_rank: int = 0, - vision_start: int = -1, - vision_length: int = -1, - kv_cache_block_offsets: Optional[Tensor] = None, - host_kv_cache_block_offsets: Tensor = None, - host_kv_cache_pool_pointers: Tensor = None, - host_kv_cache_pool_mapping: Tensor = None, - do_cross_attention: bool = False, - cross_kv: Optional[Tensor] = None, # for cross attention - cross_kv_length: Optional[Tensor] = None, # for cross attention - encoder_input_lengths: Optional[Tensor] = None, # for cross attention - relative_attention_bias: Optional[Tensor] = None, # for relative attention - logn_scaling: Optional[Tensor] = None, # for logn scaling - max_distance: int = 0, # for relative attention - host_context_lengths: Optional[Tensor] = None, # for pad-free input mode - qkv_bias: Optional[Tensor] = None, - use_cache: bool = True, - spec_decoding_is_generation_length_variable: bool = False, - spec_decoding_max_generation_length: int = 0, - spec_decoding_generation_lengths: Tensor = None, - spec_decoding_position_offsets: Tensor = None, - spec_decoding_packed_mask: Tensor = None, - spec_decoding_use: Tensor = None, - long_rope_rotary_inv_freq: Optional[Tensor] = None, - long_rope_rotary_cos_sin: Optional[Tensor] = None, - mrope_rotary_cos_sin: Tensor = None, - mrope_position_deltas: Tensor = None, - host_runtime_perf_knobs: Optional[Tensor] = None, - host_context_progress: Tensor = None, - is_mla_enabled_flag: bool = False, - q_lora_rank: int = 0, - kv_lora_rank: int = 0, - qk_nope_head_dim: int = 0, - qk_rope_head_dim: int = 0, - v_head_dim: int = 0, - q_b_proj: Optional[Tensor] = None, - kv_b_proj: Optional[Tensor] = None, - k_b_proj_trans: Optional[Tensor] = None, - skip_attn=None, - cp_group: List[int] = [0], - cp_size: int = 1, - cp_rank: int = 0, - num_kv_heads_origin: int = -1, -) -> Tuple[Tensor, Optional[Tensor]]: - ''' - Add an operation that performs the multi-head attention in GPT-like models. - - The signature of the function will change in the future release - we are in - the process of simplifying the API. The current version is still - work-in-progress! The following API is provided with hints regarding the - arguments that are likely to be removed or merged with others in the future - release. - - See docs/source/advanced/gpt-attention.md for the documentation of that function. - - Parameters: - qkv: Tensor (On GPU) - The input QKV tensor. Its shape is [batch_beam_size, max_seqlen, qkv_dim] in padded mode and [num_tokens, qkv_dim] in - packed mode. Where qkv_dim depends on using MQA, GQA, or MHA. See QKV Input in docs/source/advanced/gpt-attention.md, - - past_key_value: Tensor (On GPU) - The tensor that stores KV cache data. Its shape is - [max_batch_size * max_beam_width, 2, num_kv_heads, max_seqlen, hidden_dim_per_head] - in contiguous mode and - [max_blocks, 2, num_kv_heads, num_tokens_per_block, hidden_dim_per_head] - in paged mode. See KV Cache in docs/source/advanced/gpt-attention.md, - - attention_mask: Tensor (On GPU) - The tensor that stores the attention mask for unfused MHA or MMHA. - Its shape is [num_tokens, max_kv_seqlen]. - - attention_packed_mask: Tensor (On GPU) - The tensor that stores the packed custom mask for fmha. - Its shape is [num_tokens, max_kv_seqlen / 32], where each bit represents one mask position. - - sequence_lengths: Tensor (On GPU) - The tensor that stores the length of each sequence. Its shape is - [batch_size]. See QKV Input in docs/source/advanced/gpt-attention.md, - - host_past_key_value_lengths: Tensor (On CPU) - An INT32 tensor of shape [batch_size], - - host_max_attention_window_sizes: Tensor (On CPU) - An INT32 tensor of shape [1]. - by default, the max_attention_window_size is determined by the shape of cache_indir_table. - And we support independent max_attention_window_size for each layer. - This controls the sliding-window-attention kv-cache features. - - context_lengths: Tensor (On GPU) - The tensor that stores the context-phase sequence length of each request. Its shape - is [batch_size]. See QKV Input in doc/functional.py, - - cache_indirection: Tensor (On GPU) - The tensor to reconstruct the paths when using beam-search. Its - shape is [batch_size, beam_width, max_seqlen]. See Beam-Search in - docs/source/advanced/gpt-attention.md, - - host_request_types: 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/source/advanced/gpt-attention.md, - - layer_idx: int - The index of this attention layer, used to access kv_cache_block_offsets, - - num_heads: int - The number of heads, - - num_kv_heads: int - The number of KV heads, generic to handle MHA/MQA/GQA, - - hidden_size_per_head: int - The hidden size per head, - - q_scaling: float - The value used to compute the scaling factor applied to the output - of the Q*K^T product. See Scaling Factors in docs/source/advanced/gpt-attention.md, - - attn_logit_softcapping_scale: float - The scale * tanh(value / scale) used to compute the scaling factor applied to the output - of the Q*K^T product. - - rotary_embedding_dim: int - The dimension to compute RoPE. Use 0 when position_embedding_type is not RoPE. - - rotary_embedding_base: float - The theta value to use for RoPE. Ignored when position_embedding_type is not RoPE. - - rotary_embedding_scale_type: RotaryScalingType - The scaling type of RoPE. Ignored when position_embedding_type is not RoPE. - Possible rotary scaling type: - * RotaryScalingType.none - * RotaryScalingType.linear - * RotaryScalingType.dynamic - * RotaryScalingType.longrope - * RotaryScalingType.llama3 - - rotary_embedding_scale: float - The scale value to use for linear/dynamic scaling in RoPE. - Ignored when position_embedding_type is not RoPE. - Must be set to 1 (default) if rotary_embedding_scale_type is `none`. - - rotary_inv_freq: float Tensor - The rotary inv freq with shape [head_size / 2]. - - rotary_cos_sin: float2(cos/sin) Tensor - The rotary cos/sin cache, which will be reused among different requests. - It is taken as constant tensor. - - rotary_embedding_max_positions: int - Needed only for `dynamic` RoPE scaling. Ignored otherwise. - - position_embedding_type: PositionEmbeddingType - The position embedding type: - * PositionEmbeddingType.learned_absolute - * PositionEmbeddingType.relative - * PositionEmbeddingType.rope_gptj - * PositionEmbeddingType.rope_gpt_neox - * PositionEmbeddingType.alibi - * PositionEmbeddingType.alibi_with_scale - - kv_orig_quant_scale: Tensor - The tensor to store the scaling factor for quantization to INT8/FP8 - in the KV cache. Its shape is [1]. See INT8/FP8 KV Cache in - docs/source/advanced/gpt-attention.md, - - kv_quant_orig_scale: Tensor - The tensor to store the scaling factor for dequantization from - INT8/FP8 in the KV cache. Its shape is [1]. See INT8/FP8 KV Cache - in docs/source/advanced/gpt-attention.md, - - attention_output_orig_quant_scale: Tensor - The tensor to store the scaling factor for quantization to FP8 - in the KV cache. Its shape is [1]. - - kv_cache_quant_mode: QuantMode (int flags) - Do we enable the INT8 or FP8 KV cache? - - max_context_length: int32_t - The length of the longest input sequence. See QKV Input in - docs/source/advanced/gpt-attention.md, - - mask_type: int = 1 - The type of mask: - * tensorrt_llm.layers.AttentionMaskType.padding for BERT, - * tensorrt_llm.layers.AttentionMaskType.causal for GPT, - * tensorrt_llm.layers.AttentionMaskType.sliding_window_causal for GPT, - * tensorrt_llm.layers.AttentionMaskType.bidirectional for ChatGLM-6B, - * tensorrt_llm.layers.AttentionMaskType.bidirectionalglm for GLM-10B, - * tensorrt_llm.layers.AttentionMaskType.blocksparse for Phi-3-small, - * tensorrt_llm.layers.AttentionMaskType.custom_mask for any models. - - block_sparse_block_size: int - Block size in block sparse attention - - block_sparse_homo_head_pattern: bool - Do all attention heads share same vertical stride pattern? - - block_sparse_num_local_blocks: int - Number of active blocks near diagonal - - block_sparse_vertical_stride: int - Stride of active blocks in vertical dimension - - alibi_slopes: Tensor - The ALiBi slopes. The ALiBi bias is computed on-the-fly in the kernel - when possible, - - tp_size: int - The number of processes/GPUs when tensor parallelism is activated, - - tp_rank: int - The rank of that process (when running tensor parallelism), - - kv_cache_block_offsets: - The tensor of block offsets for the KV cache. Its shape is - [num_layers, max_batch_size, max_beam_width, 2, max_blocks_per_sequence * 2], - See KV cache section in docs/source/advanced/gpt-attention.md, on gpu, - - host_kv_cache_block_offsets: - The same as kv_cache_block_offsets, but on cpu, - - host_kv_cache_pool_pointers: - The tensor of pool pointers for the KV cache. Its shape is [num_layers, 2], - See KV cache section in docs/source/advanced/gpt-attention.md, on gpu, - - host_kv_cache_pool_mapping: - The tensor of pool mapping for the different memory pools. Its shape is [num_layers,2] - for each layer, the index of the pool, and the index of the layer within the pool, - - do_cross_attention: bool = False - Do we use this as cross attention instead of self attention, - - cross_kv: Tensor = None - The KV tensor of encoder output hidden states. Its shape is [batch_size, max_seqlen, 2 * kvHeadNum * headSize] in padded mode and [1, num_tokens, 2 * kvHeadNum * headSize] in - packed mode, - - cross_kv_length: Tensor = None - The length of the longest encoder output sequence, - - encoder_input_lengths: Tensor - The tensor that stores the length of each encoder input sequence. Its shape is [batch_size], - - logn_scaling: Tensor = None - The logn scaling tensor [max_position_embedding_len], which is applied to q in order to help extrapolation - - relative_attention_bias: Tensor = None - The relative attention bias [num_heads, max_seq_len, max_seq_len], or The relative attention embedding table for implicit mode, [num_heads, num_buckets]. - - max_distance: int = 0 - The maximum distance of relative position in attention, for implicit mode. - Default value is 0, meaning to use the regular mode of relative attention bias. - Implicit mode is only enabled when passing in non-zero positive max_distance value. - See relative attention bias in docs/source/advanced/gpt-attention.md - - host_context_lengths: Tensor = None (On CPU) - A host tensor that contains the lengths of the different inputs, - - qkv_bias: Tensor = None, - The qkv bias tensor. - - use_cache: bool = False - Do we need to store kv cache ? not needed if there is no generation phase. - - spec_decoding_is_generation_length_variable: bool = False, - Whether the generation lengths can be different for each sequence in a batch. - For Medusa, this should be set False. - For Redrafter, this should be set to True. - - spec_decoding_max_generation_length: int = 1, - The maximum number of tokens possible in the generation phase per sequence. - - spec_decoding_generation_lengths: Tensor = None, - The generation phase tokens' lengths for each sequence. - Shape: [batch_size] - - spec_decoding_position_offsets: Tensor = None, - The speculative decoding tokens's position offsets (shared by all sequences). - Shape: [batch_size, num_draft_tokens + 1]. - - spec_decoding_packed_mask: Tensor = None, - The speculative decoding tokens's attention mask (packed into uint32_t bits). - remove_input_padding is False: - Shape: [batch_size, num_draft_tokens + 1, divUp(num_draft_tokens + 1, 32)]. - remove_input_padding is True: - Shape: [sum(spec_decoding_generation_lengths), divUp(num_draft_tokens + 1, 32)]. - - long_rope_rotary_inv_freq: float Tensor - Additional rotary inv freq used for longer sequence lengths. Shape: [head_size / 2] - - long_rope_rotary_cos_sin: float2(cos/sin) Tensor - Additional rotary cos/sin cache used for longer sequence lengths. - - is_mla_enable: bool = False - Do we need to enable deepseekv2 mla? - - host_runtime_perf_knobs: Tensor = None, - The runtime perf knobs bit mask, controls whether to use certain perf knob in the runtime. - - host_context_progress: Tensor = None, - The structure used to track layer-wise progress in context phase. - - skip_attn: Tensor = None, - A bool tensor on CPU. If it is true, don't run attention plugin, returning directly. - - num_kv_heads_origin: int - The origin number of KV heads, without the process of TP - - Returns: - The tensor produced by that layer. - ''' - - assert host_request_types is not None - assert (alibi_slopes is not None) == (position_embedding_type.is_alibi()) - assert (mrope_rotary_cos_sin - is not None) == (position_embedding_type.is_mrope()) - attn_plg_creator = trt.get_plugin_registry().get_plugin_creator( - 'GPTAttention', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert attn_plg_creator is not None - assert host_context_lengths is not None or not default_net( - ).plugin_config.remove_input_padding - assert isinstance(max_context_length, int) - assert host_max_attention_window_sizes is not None - assert host_sink_token_length is not None - - paged_kv_cache_flag = default_net().plugin_config.paged_kv_cache - if isinstance(qkv, list): - is_unfuse_qkv_gemm = 1 - else: - is_unfuse_qkv_gemm = 0 - - default_net().plugin_config.context_fmha_type - if do_cross_attention and not paged_kv_cache_flag: - pass - if logn_scaling is not None: - use_logn_scaling = 1 - else: - use_logn_scaling = 0 - - if num_kv_heads_origin < 1: - num_kv_heads_origin = num_kv_heads - - unfuse_qkv_gemm = trt.PluginField( - "unfuse_qkv_gemm", np.array(np.int8(is_unfuse_qkv_gemm), dtype=np.int8), - trt.PluginFieldType.INT8) - - layer_idx = trt.PluginField("layer_idx", np.array(layer_idx, - dtype=np.int32), - trt.PluginFieldType.INT32) - nheads = trt.PluginField("num_heads", np.array(num_heads, dtype=np.int32), - trt.PluginFieldType.INT32) - vision_start = trt.PluginField("vision_start", - np.array(vision_start, dtype=np.int32), - trt.PluginFieldType.INT32) - vision_length = trt.PluginField("vision_length", - np.array(vision_length, dtype=np.int32), - trt.PluginFieldType.INT32) - num_kv_heads = trt.PluginField("num_kv_heads", - np.array(num_kv_heads, dtype=np.int32), - trt.PluginFieldType.INT32) - num_kv_heads_origin = trt.PluginField( - "num_kv_heads_origin", np.array(num_kv_heads_origin, dtype=np.int32), - trt.PluginFieldType.INT32) - head_size = trt.PluginField("head_size", - np.array(hidden_size_per_head, dtype=np.int32), - trt.PluginFieldType.INT32) - unidirectional = trt.PluginField("unidirectional", - np.array(1, dtype=np.int32), - trt.PluginFieldType.INT32) - q_scaling = trt.PluginField("q_scaling", - np.array(q_scaling, dtype=np.float32), - trt.PluginFieldType.FLOAT32) - attn_logit_softcapping_scale = trt.PluginField( - "attn_logit_softcapping_scale", - np.array(attn_logit_softcapping_scale, dtype=np.float32), - trt.PluginFieldType.FLOAT32) - rotary_embedding_dim = trt.PluginField( - "rotary_embedding_dim", np.array(rotary_embedding_dim, dtype=np.int32), - trt.PluginFieldType.INT32) - rotary_embedding_base = trt.PluginField( - "rotary_embedding_base", - np.array(rotary_embedding_base, dtype=np.float32), - trt.PluginFieldType.FLOAT32) - rotary_embedding_scale_type = trt.PluginField( - "rotary_embedding_scale_type", - np.array(rotary_embedding_scale_type, dtype=np.int8), - trt.PluginFieldType.INT8) - rotary_embedding_scale = trt.PluginField( - "rotary_embedding_scale", - np.array(rotary_embedding_scale, dtype=np.float32), - trt.PluginFieldType.FLOAT32) - rotary_embedding_short_m_scale = trt.PluginField( - "rotary_embedding_short_m_scale", - np.array(rotary_embedding_short_m_scale, dtype=np.float32), - trt.PluginFieldType.FLOAT32) - rotary_embedding_long_m_scale = trt.PluginField( - "rotary_embedding_long_m_scale", - np.array(rotary_embedding_long_m_scale, dtype=np.float32), - trt.PluginFieldType.FLOAT32) - rotary_embedding_max_positions = trt.PluginField( - "rotary_embedding_max_positions", - np.array(rotary_embedding_max_positions, dtype=np.int32), - trt.PluginFieldType.INT32) - rotary_embedding_original_max_positions = trt.PluginField( - "rotary_embedding_original_max_positions", - np.array(rotary_embedding_original_max_positions, dtype=np.int32), - trt.PluginFieldType.INT32) - position_embedding_type = trt.PluginField( - "position_embedding_type", - np.array(int(position_embedding_type), dtype=np.int8), - trt.PluginFieldType.INT8) - context_fmha_type = trt.PluginField( - "context_fmha_type", - np.array(np.int8(default_net().plugin_config.context_fmha_type), - dtype=np.int8), trt.PluginFieldType.INT8) - remove_input_padding = trt.PluginField( - "remove_input_padding", - np.array(np.int8(default_net().plugin_config.remove_input_padding), - dtype=np.int8), trt.PluginFieldType.INT8) - is_spec_decoding_enabled = trt.PluginField( - "is_spec_decoding_enabled", - np.array(np.int8(spec_decoding_packed_mask is not None), dtype=np.int8), - trt.PluginFieldType.INT8) - spec_decoding_is_generation_length_variable = trt.PluginField( - "spec_decoding_is_generation_length_variable", - np.array(np.int8(spec_decoding_is_generation_length_variable), - dtype=np.int8), trt.PluginFieldType.INT8) - spec_decoding_max_generation_length = trt.PluginField( - "spec_decoding_max_generation_length", - np.array(spec_decoding_max_generation_length, dtype=np.int32), - trt.PluginFieldType.INT32) - is_mla_enabled = trt.PluginField( - "is_mla_enabled", np.array(is_mla_enabled_flag, dtype=np.int8), - trt.PluginFieldType.INT8) - q_lora_rank = trt.PluginField("q_lora_rank", - np.array(q_lora_rank, dtype=np.int32), - trt.PluginFieldType.INT32) - kv_lora_rank = trt.PluginField("kv_lora_rank", - np.array(kv_lora_rank, dtype=np.int32), - trt.PluginFieldType.INT32) - qk_nope_head_dim = trt.PluginField( - "qk_nope_head_dim", np.array(qk_nope_head_dim, dtype=np.int32), - trt.PluginFieldType.INT32) - qk_rope_head_dim = trt.PluginField( - "qk_rope_head_dim", np.array(qk_rope_head_dim, dtype=np.int32), - trt.PluginFieldType.INT32) - v_head_dim = trt.PluginField("v_head_dim", - np.array(v_head_dim, dtype=np.int32), - trt.PluginFieldType.INT32) - p_dtype = default_net().plugin_config.gpt_attention_plugin - pf_type = trt.PluginField( - "type_id", np.array([int(str_dtype_to_trt(p_dtype))], np.int32), - trt.PluginFieldType.INT32) - # reset mask_type to custom_mask. - if (attention_mask is not None) or (attention_packed_mask is not None): - # context fmha needs packed mask. - assert attention_packed_mask is not None - if get_sm_version() < 100: - mask_type = AttentionMaskType.custom_mask - - mask_type_filed = trt.PluginField("mask_type", - np.array([int(mask_type)], np.int32), - trt.PluginFieldType.INT32) - block_sparse_block_size = trt.PluginField( - "block_sparse_block_size", np.array([block_sparse_block_size], - np.int32), - trt.PluginFieldType.INT32) - block_sparse_homo_head_pattern = trt.PluginField( - "block_sparse_homo_head_pattern", - np.array(np.int8(block_sparse_homo_head_pattern), np.int8), - trt.PluginFieldType.INT8) - block_sparse_num_local_blocks = trt.PluginField( - "block_sparse_num_local_blocks", - np.array([block_sparse_num_local_blocks], np.int32), - trt.PluginFieldType.INT32) - block_sparse_vertical_stride = trt.PluginField( - "block_sparse_vertical_stride", - np.array([block_sparse_vertical_stride], np.int32), - trt.PluginFieldType.INT32) - tp_size = trt.PluginField("tp_size", np.array(tp_size, dtype=np.int32), - trt.PluginFieldType.INT32) - tp_rank = trt.PluginField("tp_rank", np.array(tp_rank, dtype=np.int32), - trt.PluginFieldType.INT32) - if isinstance(kv_cache_quant_mode, QuantModeWrapper): - # Now in TRT-LLM only use global kv_cache, so it's enough to get the first quant mode from list - kv_cache_quant_mode = kv_cache_quant_mode[0] - kv_cache_quant_mode_field = trt.PluginField( - "kv_cache_quant_mode", np.array(kv_cache_quant_mode, dtype=np.int32), - trt.PluginFieldType.INT32) - paged_kv_cache = trt.PluginField( - "paged_kv_cache", np.array(paged_kv_cache_flag, dtype=np.int32), - trt.PluginFieldType.INT32) - tokens_per_block = trt.PluginField( - "tokens_per_block", - np.array(default_net().plugin_config.tokens_per_block, dtype=np.int32), - trt.PluginFieldType.INT32) - max_context_length = trt.PluginField("max_context_length", - np.array(max_context_length, np.int32), - trt.PluginFieldType.INT32) - pos_shift_enabled = trt.PluginField( - "pos_shift_enabled", - np.array(np.int8(default_net().plugin_config.streamingllm), - dtype=np.int8), trt.PluginFieldType.INT8) - dense_context_fmha = trt.PluginField( - "dense_context_fmha", - np.array(np.int8(default_net().plugin_config.streamingllm), - dtype=np.int8), trt.PluginFieldType.INT8) - if qkv_bias is None: - qkv_bias_enabled = trt.PluginField("qkv_bias_enabled", - np.array(0, dtype=np.int8), - trt.PluginFieldType.INT8) - else: - qkv_bias_enabled = trt.PluginField("qkv_bias_enabled", - np.array(1, dtype=np.int8), - trt.PluginFieldType.INT8) - do_cross_attention_field = trt.PluginField( - "do_cross_attention", - np.array(np.int8(do_cross_attention), dtype=np.int8), - trt.PluginFieldType.INT8) - max_distance = trt.PluginField("max_distance", - np.array(max_distance, dtype=np.int32), - trt.PluginFieldType.INT32) - use_paged_context_fmha_field = trt.PluginField( - "use_paged_context_fmha", - np.array(np.int8(default_net().plugin_config.use_paged_context_fmha), - dtype=np.int8), trt.PluginFieldType.INT8) - use_fp8_context_fmha_field = trt.PluginField( - "use_fp8_context_fmha", - np.array(np.int8(default_net().plugin_config.use_fp8_context_fmha), - dtype=np.int8), trt.PluginFieldType.INT8) - has_full_attention_mask_field = trt.PluginField( - "has_full_attention_mask", - np.array(np.int8(attention_mask is not None), dtype=np.int8), - trt.PluginFieldType.INT8) - use_cache_pf = trt.PluginField("use_cache", - np.array([use_cache], dtype=np.int32), - trt.PluginFieldType.INT32) - fuse_fp4_quant = default_net().plugin_config.fuse_fp4_quant - fuse_fp4_quant_pf = trt.PluginField( - "fuse_fp4_quant", np.array(np.int8(fuse_fp4_quant), dtype=np.int8), - trt.PluginFieldType.INT8) - skip_attn_pf = trt.PluginField( - "skip_attn", np.array([skip_attn is not None], dtype=np.int8), - trt.PluginFieldType.INT8) - cp_size = trt.PluginField("cp_size", np.array(cp_size, dtype=np.int32), - trt.PluginFieldType.INT32) - cp_rank = trt.PluginField("cp_rank", np.array(cp_rank, dtype=np.int32), - trt.PluginFieldType.INT32) - cp_group = np.array(cp_group, dtype=np.int32) - cp_group = trt.PluginField("cp_group", cp_group, trt.PluginFieldType.INT32) - use_logn_scaling = trt.PluginField( - "use_logn_scaling", np.array(np.int8(use_logn_scaling), dtype=np.int8), - trt.PluginFieldType.INT8) - - pfc = trt.PluginFieldCollection([ - layer_idx, nheads, 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, context_fmha_type, - kv_cache_quant_mode_field, remove_input_padding, mask_type_filed, - block_sparse_block_size, block_sparse_homo_head_pattern, - block_sparse_num_local_blocks, block_sparse_vertical_stride, - paged_kv_cache, tokens_per_block, pf_type, max_context_length, - qkv_bias_enabled, do_cross_attention_field, max_distance, - pos_shift_enabled, dense_context_fmha, use_paged_context_fmha_field, - use_fp8_context_fmha_field, has_full_attention_mask_field, use_cache_pf, - 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_pf, skip_attn_pf, cp_size, cp_rank, cp_group, - use_logn_scaling - ]) - - attn_plug = attn_plg_creator.create_plugin("causal_attn", pfc) - assert attn_plug - plug_inputs = [*qkv] if is_unfuse_qkv_gemm else [qkv] - if attention_mask is not None and mask_type == AttentionMaskType.custom_mask: - # useFullCustomMask - plug_inputs += [attention_mask] - if attention_packed_mask is not None and get_sm_version() < 100: - # usePackedCustomMask - plug_inputs += [attention_packed_mask] - if use_cache: - plug_inputs += [ - sequence_length, - host_past_key_value_lengths, - host_max_attention_window_sizes, - host_sink_token_length, - context_lengths, - cache_indirection, - host_request_types, - ] - else: - plug_inputs += [ - host_max_attention_window_sizes, - host_sink_token_length, - context_lengths, - host_request_types, - ] - if use_cache: - if paged_kv_cache_flag: - assert kv_cache_block_offsets is not None, "Paged kv cache is enabled, the kv_cache_block_offsets tensor shall not be None" - assert host_kv_cache_block_offsets is not None, "Paged kv cache is enabled, the host_kv_cache_block_offsets tensor shall not be None" - assert host_kv_cache_pool_pointers is not None, "Paged kv cache is enabled, the host_kv_cache_pool_pointers tensor shall not be None" - assert host_kv_cache_pool_mapping is not None, "Paged kv cache is enabled, the host_kv_cache_pool_mapping tensor shall not be None" - plug_inputs += [ - kv_cache_block_offsets, host_kv_cache_block_offsets, - host_kv_cache_pool_pointers, host_kv_cache_pool_mapping - ] - else: - plug_inputs += [past_key_value] - - if use_cache and kv_cache_quant_mode.has_kv_cache_quant(): - plug_inputs += [kv_orig_quant_scale, kv_quant_orig_scale] - - if attention_output_orig_quant_scale is not None: - assert default_net( - ).plugin_config.use_fp8_context_fmha, "FP8 Context FMHA needs to be enabled" - plug_inputs += [attention_output_orig_quant_scale] - - if fuse_fp4_quant: - assert attention_output_sf_scale is not None, "attention_output_sf_scale must be provided when fuse_fp4_quant is enabled." - plug_inputs += [attention_output_sf_scale] - - if rotary_inv_freq is not None: - plug_inputs += [rotary_inv_freq] - if rotary_cos_sin is not None: - plug_inputs += [rotary_cos_sin] - - if alibi_slopes is not None: - plug_inputs += [alibi_slopes] - - if relative_attention_bias is not None: - plug_inputs += [relative_attention_bias] - - if do_cross_attention: - plug_inputs += [cross_kv, cross_kv_length, encoder_input_lengths] - - if default_net().plugin_config.remove_input_padding: - plug_inputs += [host_context_lengths] - - if qkv_bias is not None: - plug_inputs += [qkv_bias] - - if spec_decoding_packed_mask is not None: - # add position_ids as well only if speculative decoding mode - assert spec_decoding_position_offsets is not None - assert spec_decoding_generation_lengths is not None - assert spec_decoding_use is not None - plug_inputs += [ - spec_decoding_generation_lengths, spec_decoding_packed_mask, - spec_decoding_position_offsets, spec_decoding_use - ] - - if long_rope_rotary_inv_freq is not None: - assert long_rope_rotary_cos_sin is not None - plug_inputs += [long_rope_rotary_inv_freq, long_rope_rotary_cos_sin] - - if mrope_rotary_cos_sin is not None: - assert mrope_position_deltas is not None - plug_inputs += [ - mrope_rotary_cos_sin, - mrope_position_deltas, - ] - if host_runtime_perf_knobs is not None: - plug_inputs += [host_runtime_perf_knobs] - - if host_context_progress is not None: - plug_inputs += [host_context_progress] - - if is_mla_enabled_flag: - assert q_b_proj is not None - assert kv_b_proj is not None - assert k_b_proj_trans is not None - plug_inputs += [q_b_proj, kv_b_proj, k_b_proj_trans] - - if skip_attn is not None: - plug_inputs += [skip_attn] - - if logn_scaling is not None: - plug_inputs += [logn_scaling] - - for idx, i in enumerate(plug_inputs): - assert i is not None, f"Found None input for {idx} th item in plugin inputs {plug_inputs}" - - plug_inputs = [i.trt_tensor for i in plug_inputs] - layer = default_trtnet().add_plugin_v2(plug_inputs, attn_plug) - _add_plugin_info(layer, attn_plg_creator, "causal_attn", pfc) - output = _create_tensor(layer.get_output(0), layer) - expected_outputs = 1 - - # The output scaling factor tensor. - output_sf = None - if fuse_fp4_quant: - output_sf = _create_tensor(layer.get_output(expected_outputs), layer) - expected_outputs += 1 - - present_key_value = None - if use_cache and not paged_kv_cache_flag: - present_key_value = _create_tensor(layer.get_output(expected_outputs), - layer) - assert present_key_value is not None - expected_outputs += 1 - - assert layer.num_outputs == expected_outputs, \ - f"Plugin outputs number mismatch with expected, got {layer.num_outputs}, expected {expected_outputs}" - - if kv_cache_quant_mode.has_int8_kv_cache( - ) and not default_net().strongly_typed: - if not paged_kv_cache_flag: - # past key value - layer.get_input(8).set_dynamic_range(-127, 127) - # present key value - layer.get_output(expected_outputs - 1).set_dynamic_range(-127, 127) - else: - layer.get_input(0).set_dynamic_range(-127, 127) - layer.get_input(1).set_dynamic_range(-127, 127) - layer.get_output(expected_outputs - 1).set_dynamic_range(-127, 127) - - assert output is not None - if fuse_fp4_quant: - assert output_sf is not None - return (output, output_sf), present_key_value - return output, present_key_value - - -def assertion(condition: Tensor, message: str = '') -> None: - default_trtnet().add_assertion(condition.trt_tensor, message) - - -def layer_norm(input: Tensor, - normalized_shape: Union[int, Tuple[int]], - weight: Optional[Tensor] = None, - bias: Optional[Tensor] = None, - eps: float = 1e-05, - use_diff_of_squares: bool = True) -> Tensor: - ''' - Add a layer-norm operation on a tensor. - - That operation applies the layer-normalization to its input tensor. In its - simplest form, for large language models, the 'normalized_shape' should be - set to the hidden dimension of the activation tensor. Otherwise, it is the - shape of the normalized fraction of the tensor (starting from the - right-most dimension). - - The 'weight' tensor corresponds to 'gamma' in the layer-norm formula and - 'bias' is 'beta'. The 'eps' value is added to the variance before computing - the squared-root. - - This implementation (when using the plugin) supports an additional flag to - enable/disable the use of a difference of squares ('Var = Mean(X^2) - - Mean(X)^2'). - - Parameters: - input : Tensor - The tensor to normalize. - - normalized_shape : Union[int, Tuple[int]] - The shape of the sub-tensor that is normalized. Use 'hidden_dim' to - normalize the inner-most dimension of an activation tensor in LLMs. - - weight : Optional[Tensor] = None - The 'gamma' term in layer-norm. Its shape must be - 'normalized_shape'. - - bias : Optional[Tensor] = None - The 'beta' term in layer-norm. Its shape must be - 'normalized_shape'. - - eps : float - The epsilon term to be added to the variance in the squared-root. - - use_diff_of_squares : bool - Does the plugin use the difference of squares to compute the - variance? - - Returns: - The output tensor of that operation. - ''' - input, weight = broadcast_helper(input, weight) - input, bias = broadcast_helper(input, bias) - if isinstance(normalized_shape, int): # FIXME: better way? - axis = input.ndim() - 1 - else: - axis = input.ndim() - len(normalized_shape) - axes_mask = 0 - for i in range(axis, input.ndim()): - axes_mask |= 1 << i - layer = default_trtnet().add_normalization(input.trt_tensor, - weight.trt_tensor, - bias.trt_tensor, axes_mask) - layer.epsilon = eps - return _create_tensor(layer.get_output(0), layer) - - -def rms_norm(input: Tensor, - normalized_shape: Union[int, Tuple[int]], - num_groups: int = 1, - weight: Optional[Tensor] = None, - eps: float = 1e-06) -> Tensor: - ''' - Add a RMS norm operation on a tensor. - - That operation applies the rms-normalization to its input tensor. In its - simplest form, for large language models, the 'normalized_shape' should be - set to the hidden dimension of the activation tensor. Otherwise, it is the - shape of the normalized fraction of the tensor (starting from the - right-most dimension). - - The 'weight' tensor corresponds to 'gamma' in the rms-norm formula. - The 'eps' value is added to the variance before computing the squared-root. - - Parameters: - input: Tensor - The tensor to normalize. - - normalized_shape : Union[int, Tuple[int]] - The shape of the sub-tensor that is normalized. Use 'hidden_dim' to - normalize the inner-most dimension of an activation tensor in LLMs. - - num_groups: int = 1 - The group size. - - weight : Optional[Tensor] = None - The 'gamma' term in layer-norm. Its shape must be - 'normalized_shape'. - - eps : float - The epsilon term to be added to the variance in the squared-root.weig - Returns: - The output tensor of that operation. - ''' - normalized_shape = [normalized_shape] if isinstance( - normalized_shape, int) else normalized_shape - - dim = tuple([-i - 1 for i in range(len(normalized_shape))]) - - if num_groups > 1: - assert len(normalized_shape) == 1 - num_channels = input.size()[-1] - ndim = input.ndim() - old_shape = shape(input) - new_shape = concat([input.size(i) for i in range(ndim - 1)] + - [num_groups, num_channels // num_groups]) - input = input.view(new_shape) - - with precision("float32"): - input_dtype = input.dtype - fp32_input = cast(input, "float32") - varx = pow(fp32_input, 2.0) - - varx = varx.mean(dim=dim, keepdim=True) - denom = varx + eps - denom = denom.sqrt() - fp32_y = fp32_input / denom - y = cast(fp32_y, input_dtype) - - if num_groups > 1: - y = y.view(old_shape) - - if weight is not None: - y = y * weight - - return y - - -def rearrange(inputs: Union[Tensor, Sequence[Tensor]], expression: str, - **kwargs) -> Tensor: - ''' - Add a rearrange operation on a tensor. - - This operation is a reader-friendly smart element reordering for multidimensional tensors, - including functionality of transpose (axes permutation), reshape (view), squeeze, unsqueeze, - stack, concatenate and other operations. Please see: https://einops.rocks/api/rearrange/ - - For example, if the shape of input tensor is [32, 30, 40, 3], and run: - `rearrange(x, 'b (h h1) (w w1) c -> b h w 1 (c h1 w1) 1', h1=2, w1=2)` - it would produce a tensor with shape as [32, 15, 20, 1, 12, 1]. - - Parameters: - input: Union[Tensor, Sequence[Tensor]] - If it is a tensor, it will directly operate on it. - Otherwise, if it is a sequence, it will concat it to a tensor and then - operates on it. - - expression : str - The expression about how to reorder the tensor in a reader-friendly way. - - kwargs: - Keyword arguments to set some identifiers with specific values. - - Returns: - The output tensor of this operation. - ''' - import re - - def _init_expression(expr): - expr_items = expr.split(" ") - tmp_name_index = 0 - for idx, item in enumerate(expr_items): - values = re.findall(r'\b\d+\b', item) - if len(values) > 0: - prefix = "(" if "(" in item else "" - subfix = ")" if ")" in item else "" - expr_items[ - idx] = f"{prefix}NumericId{tmp_name_index}Val{values[0]}{subfix}" - tmp_name_index += 1 - return " ".join(expr_items) - - def _get_all_identifier(expr): - return re.findall(r'\b[a-zA-Z_]+\d*\b', expr) - - def _get_all_symbols(expr): - return re.findall(r'\b\w+\b', expr) - - def _get_dim_expr(expr): - return [ - _get_all_symbols(match.group()) - for match in re.finditer(r'\b\w+\b|\(.*?\)', expr) - ] - - src_shape_expr, _, dst_shape_expr = expression.partition("->") - unknown_identifiers = re.findall(r'[^a-zA-Z0-9_\(\)]', - src_shape_expr + dst_shape_expr) - assert len( - unknown_identifiers) > 0, f"Unknown identifiers: {unknown_identifiers}" - src_identifiers = _get_all_identifier(src_shape_expr) - dst_identifiers = _get_all_identifier(dst_shape_expr) - assert (len(src_identifiers) == len(set(src_identifiers)) - and len(dst_identifiers) == len(set(dst_identifiers)) - ), "Indexing expression contains duplicate dimension." - assert (set(src_identifiers) == set(dst_identifiers) - ), "Identifiers only on one side of expression (should be on both)." - - new_expression = _init_expression(expression) - src_shape_expr, _, dst_shape_expr = new_expression.partition("->") - - # concat if inputs are sequence of tensors - if isinstance(inputs, Sequence): - inputs = concat([unsqueeze(t, 0) for t in inputs], dim=0) - assert ( - inputs.ndim() == len(_get_dim_expr(src_shape_expr)) - ), f"inputs.ndim() is {inputs.ndim()} while indexing expression has {len(_get_dim_expr(src_shape_expr))}" - - src_symbols = _get_all_symbols(src_shape_expr) - dst_symbols = _get_all_symbols(dst_shape_expr) - - # find all the symbols-values mapping and store them in symbol_map - symbol_map = { - symbol: { - "updated": False, - "value": None - } - for symbol in set(src_symbols + dst_symbols) - } - for symbol in symbol_map: - if "NumericId" in symbol: - symbol_map[symbol]["value"] = int(symbol.partition("Val")[-1]) - symbol_map[symbol]["updated"] = True - for symbol, value in kwargs.items(): - symbol_map[symbol]["value"] = value - symbol_map[symbol]["updated"] = True - - for idx, dim_expr in enumerate(_get_dim_expr(src_shape_expr)): - if len(dim_expr) == 1: - symbol = dim_expr[0] - if not symbol_map[symbol]["updated"]: - symbol_map[symbol]["value"] = shape(inputs, idx) - symbol_map[symbol]["updated"] = True - else: - divisors = [] - unknown_symbol = None - for symbol in dim_expr: - if not symbol_map[symbol]["updated"]: - unknown_symbol = symbol - else: - divisors.append(symbol_map[symbol]["value"]) - if unknown_symbol is not None: - assert len(divisors) > 0 - divisor = prod(cast(concat(divisors), "int64"), dim=-1) - symbol_map[unknown_symbol]["value"] = shape(inputs, - idx) / divisor - symbol_map[unknown_symbol]["updated"] = True - - for symbol, item in symbol_map.items(): - assert (item["updated"] - ), f"{symbol} cannot be inferred, please set it manually" - - dst_dims = [] - for dim_expr in _get_dim_expr(dst_shape_expr): - if len(dim_expr) == 1: - dst_dims.append(symbol_map[dim_expr[0]]["value"]) - else: - accumulator = prod(cast( - concat([symbol_map[symbol]["value"] for symbol in dim_expr]), - "int64"), - dim=-1) - dst_dims.append(accumulator) - dst_dims = cast(concat(dst_dims, dim=-1), "int64") - - src_indices = {symbol: idx for idx, symbol in enumerate(src_identifiers)} - permute_dims = [src_indices[symbol] for symbol in dst_identifiers] - - symbol_shape = cast( - concat([symbol_map[symbol]["value"] for symbol in src_identifiers], - dim=-1), "int64") - tensor = inputs.view(symbol_shape) - tensor = permute(tensor, permute_dims) - tensor = tensor.view(dst_dims) - return tensor - - -def repeat(input: Tensor, sizes: Sequence[int]) -> Tensor: - ''' - Repeats the tensor along the specified dimensions. - - Parameters: - input : Tensor - The tensor to be repeated. - sizes : Sequence[int] - The number of times to repeat the tensor along each dimension. - - Returns: - A tensor except for repeated input tensors along specified dim. - - ''' - assert input.ndim() <= len(sizes), \ - "Number of dimensions of repeat dims can not be smaller than number of dimensions of tensor" - repeated_tensor = input - for k in range(-1, -len(sizes) - 1, -1): - repeated_tensor = concat([repeated_tensor] * sizes[k], dim=k) - return repeated_tensor - - -def repeat_interleave(tensor: Tensor, repeats: int, dim: int) -> Tensor: - ''' - Repeats elements of a tensor along an axis. - - Parameters: - repeats : int - The number of repetitions along axis specified. - dim : int - The dimension along which repetitions are performed. - - Returns: - A tensor with the same shape as input except for repeated elements along specified dim. - - TODO: Allow repeats to be a list of integers and dim to be unspecified. - ''' - expanded_tensor = expand_dims(tensor, dim + 1) - tile_output_size = concat([ - repeats if i == (dim + 1) else shape(expanded_tensor, i) - for i in range(expanded_tensor.ndim()) - ]) - tile = expand(expanded_tensor, tile_output_size) - tile_reshape_size = [shape(tensor, i) for i in range(tensor.ndim())] - tile_reshape_size[dim] = tile_reshape_size[dim] * repeats - tensor = tile.view(concat(tile_reshape_size)) - return tensor - - -def meshgrid2d(x: Tensor, y: Tensor) -> Tuple[Tensor]: - ''' - Creates grids (2D) of coordinates specified by the 1D inputs (only supports `indexing=\'xy\'`). - - Parameters: - x : Tensor - The first input (1D) tensor. - y : Tensor - The second input (1D) tensor. - - Returns: - The tuple of two tensors produced. - - TODO: Add full support for torch.meshgrid. - See https://pytorch.org/docs/stable/generated/torch.meshgrid.html#torch-meshgrid - ''' - if x.ndim() == 1: - x = expand_dims(x, 0) - if y.ndim() == 1: - y = expand_dims(y, 0) - grid_x = repeat_interleave(x, shape(y, 1), - 1).view([x.shape[-1], y.shape[-1]]) - grid_y = repeat(y, (x.shape[-1], 1)) - return (grid_x, grid_y) - - -def generate_logn_scaling(seq_length: int = 8192, - max_position_embeddings: int = 32768) -> np.ndarray: - ''' - Compute the Log-N scaling vector for Qwen inference extrapolation - - Parameters: - seq_length : int - The max seq length in training (default to 8192 in Qwen-1) - max_position_embeddings : int - The max position embeddings. (default to 32768 in Qwen-1) - - Returns: - A constant np.ndarray that contains logn scaling vector - ''' - logn_list = [ - math.log(i, seq_length) if i > seq_length else 1 - for i in range(1, max_position_embeddings + 1) - ] - return np.asarray(logn_list, dtype=np.float32) - - -def generate_alibi_slopes(num_heads: int, - tp_size: int = 1, - tp_rank: int = 0, - alibi_scale: float = 1.0, - alibi_bias_max: int = 8) -> np.ndarray: - ''' - Compute the ALiBi slopes as described in https://arxiv.org/abs/2211.05100. - - Parameters: - num_heads : int - The number of heads. - dtype : trt.DataType - The data type of the returned slopes - tp_size : int - The tensor parallelism size - tp_rank : int - The tensor parallelism rank - - Returns: - A constant tensor that contains the ALiBi slopes. - ''' - start_head_id = 0 - end_head_id = num_heads - - if tp_size > 1: - rank_heads = num_heads // tp_size - start_head_id = rank_heads * tp_rank - end_head_id = start_head_id + rank_heads - - closest_power_of_2 = 2**np.floor(np.log2(num_heads)) - # FT's implementation - # https://github.com/NVIDIA/FasterTransformer/blob/main/src/fastertransformer/kernels/gen_relative_pos_bias.cu#L248 - slopes_ft = [] - for h_id in range(start_head_id, end_head_id): - if h_id < closest_power_of_2: - slopes_ft.append( - np.power( - 2**(-(2**-(np.log2(closest_power_of_2) - - np.log2(alibi_bias_max)))), h_id + 1)) - else: - slopes_ft.append( - np.power( - 2**(-(2**-(np.log2(closest_power_of_2 * 2) - - np.log2(alibi_bias_max)))), - (h_id - closest_power_of_2) * 2 + 1)) - slopes = np.asarray(slopes_ft, dtype=np.float32) - - slopes = alibi_scale * slopes - slopes = slopes.reshape(1, (end_head_id - start_head_id), 1, 1) - return slopes - - -def generate_alibi_biases(slopes: Tensor, key_length: Tensor) -> Tensor: - ''' - Compute the ALiBi biases as described in https://arxiv.org/abs/2211.05100. - - The ALiBi biases are added to the result of the Q*K^T product in the - multi-head attention block. - - Parameters: - slopes : Tensor - The slopes. - - key_length : Tensor - The size of the K vector per head. - - Returns: - A constant tensor that contains the ALiBi biases. - ''' - # We don't need to care about the batch size or query length since we can just broadcast - # across the batch and query dimensions - - trt_0 = constant(int32_array(0)) - arange_shape = concat([1, 1, 1, key_length]) - - arange_tensor = arange(trt_0, key_length, "float32").view(arange_shape) - return slopes * arange_tensor - - -def expand_mask(mask: Tensor, tgt_len: Optional[Tensor] = None) -> Tensor: - ''' - Expand an attention mask. - - That function adds the sequence of operations to expand from a tensor of - shape '[batch_size, src_seq_len]' to a tensor of shape - '[batch_size, 1, tgt_seq_len, src_seq_len]'. It can be used to create the - mask applied to the Q*K^T product before the softmax operation in the - multi-head attention block. - - Parameters: - mask : Tensor - The input mask - - tgt_len : Optional[Tensor] - The dimension of the 3rd dimension in the output tensor. If None, - the 2nd dimension of the input is used. - - Returns: - The tensor created by that sequence of operations. - ''' - bsz = shape(mask, 0) - src_len = shape(mask, 1) - tgt_len = tgt_len if tgt_len is not None else src_len - - mask = mask.view(concat([bsz, 1, 1, src_len])) - - mask = expand(mask, concat([bsz, 1, tgt_len, src_len])) - mask = where(mask == 0, float('-inf'), 0.0) - return mask - - -def gather_last_token_logits(hidden_states: Tensor, last_token_ids: Tensor, - remove_input_padding: bool) -> Tensor: - ''' - Extract the logits that correspond to the last token from the hidden states. - - That function adds the operations to extract the logits of the last tokens - in a batch of sequences. - - Depending on whether 'remove_input_padding' is 'True' or 'False', that - function assumes inputs of different shapes. - - When 'remove_input_padding' is 'True', the 'hidden_states' tensor is - assumed to be packed. It has a shape '[num_tokens, hidden_dim]' where - 'num_tokens' is the sum of the lengths of the sequences in the batch and - 'hidden_dim' is the hidden dimension. The 'last_tokens_ids' is a 1D tensor - that encodes the inclusive prefix-sums of the lengths of the sequences in - the batch. - - When 'remove_input_padding' is 'False', the 'hidden_states' tensor is - assumed to be padded. It has a shape '[batch_size, max_seqlen, hidden_dim]' - where 'max_seqlen' is the length of the longest sequence in the batch and - 'hidden_dim' is the hidden dimension. The 'last_token_ids' is a 1D tensor - that encodes the length of each sequence in the batch. - - In both cases, that function produces a tensor of shape '[batch_size, - hidden_size]' where the row at index 'i' corresponds to the logits of the - last token from the 'i'-th sequence. - - Parameters: - hidden_states : Tensor - The hidden states - - last_token_ids : Tensor - The inclusive prefix-sum of the lengths or the lengths of the - sequences in the batch. - - remove_input_padding : bool - Indicate if the hidden_states are packed ('True') or padded - ('False'). - - Returns: - The tensor created by that sequence of operations. - ''' - if last_token_ids is None: - return hidden_states - - if remove_input_padding: - hidden_states = index_select(hidden_states, 0, - last_token_ids - 1) # [seq_len, hidden] - - hidden_states = hidden_states.view( - concat([shape(last_token_ids, 0), - shape(hidden_states, 1)])) - else: - ndim = last_token_ids.ndim() - if ndim == 1: - # only calculate logits for the last token - # [batch_size, seqlen, hidden_size] -> [batch_size, hidden_size] - last_token_ids = last_token_ids.view( - concat([shape(last_token_ids, 0), 1, 1])) - last_token_ids = expand( - last_token_ids, - concat([shape(last_token_ids, 0), 1, - shape(hidden_states, 2)])) - last_token_ids = last_token_ids - 1 - hidden_states = gather( - hidden_states, dim=1, indices=last_token_ids).view( - concat([shape(hidden_states, 0), - shape(hidden_states, 2)])) - elif ndim == 2: # speculative decoding needs last few token's logits - # last_token_ids is of shape [batch_size, num_last_tokens] - # So [batch_size, seqlen, hidden_size] -> [batch_size, num_last_tokens, hidden_size] - last_token_ids = last_token_ids.view( - concat([shape(last_token_ids, 0), - shape(last_token_ids, 1), 1])) - last_token_ids = expand( - last_token_ids, - concat([ - shape(last_token_ids, 0), - shape(last_token_ids, 1), - shape(hidden_states, 2) - ])) - hidden_states = gather(hidden_states, dim=1, indices=last_token_ids) - return hidden_states - - -ACT2FN = { - 'relu': relu, - 'tanh': tanh, - 'gelu': gelu, - 'gelu_new': gelu, - 'gelu_fast': gelu, - 'gelu_pytorch_tanh': gelu, - 'openai-gelu': gelu, - 'geglu': geglu, - 'gegelu': gegelu, - 'identity': identity, - 'silu': silu, - 'softplus': softplus, - 'relu2': squared_relu, - 'squared-relu': squared_relu, - 'swiglu': swiglu, - 'fast-swiglu': swiglu, - 'sigmoid': sigmoid, - 'quick_gelu': quick_gelu, -} - -GATED_ACT_2_ACT = { - 'swiglu': 'silu', - 'fast-swiglu': 'silu', - 'geglu': 'gelu', -} - - -def is_gated_activation(activation): - ''' - Is a given activation function gated? - - Parameters: - activation : str - The name of the activation function. - - Returns: - True if the function is gated, False otherwise. - ''' - assert activation in ACT2FN - return activation in GATED_ACT_2_ACT - - -def non_gated_version(activation): - ''' - Given an activation function, get the non-gated version. - - If the activation function is non-gated, it returns the same activation - function name. - - For example, that function returns 'silu' for 'swiglu' and 'relu' for - 'relu'. - - Parameters: - activation : str - The name of the activation function. - - Returns: - The name of the non-gated activation function. - ''' - if is_gated_activation(activation): - return GATED_ACT_2_ACT[activation] - return activation - - -def lora_plugin( - input: Tensor = None, - in_hidden_size: int = 0, - out_hidden_sizes: List[int] = [0], - host_request_types: Tensor = None, - transa: bool = False, - transb: bool = False, - host_context_lengths: Tensor = None, # for pad-free input mode - max_low_rank: int = 0, - lora_ranks: List[Tensor] = None, - lora_weights_pointers: List[Tensor] = None, - weight_index: int = 0, -): - ''' - Parameters: - input : Tensor (On GPU) - The input tensor. Its shape is [batch_size, seq_len, dim] or [num_tokens, dim] for remove_input_padding - - in_hidden_size/out_hidden_size : int - the lora computation workflow is - [M, in_hidden_size] -> [M, low_rank] -> [M, out_hidden_size] - - host_request_types : Tensor = None - 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/source/advanced/gpt-attention.md, - - transa : bool - Is the first input transposed? Set to 'True' if you want the first - input to be transposed, 'False' otherwise. - - transb : bool - Is the second input transposed? Set to 'True' if you want the - second input to be transposed, 'False' otherwise. - - host_context_lengths: cpu Tensor = None - A host tensor that contains the lengths of the different inputs, - - max_low_rank : int - Maximum low_rank, used to determine the workspace size. - - lora_ranks : cpu Tensor with shape [batch_size] - The low_rank of each request - - lora_weights_pointers : cpu int64 Tensor with shape [batch_size, 3] - The weights pointers of each request. Consist of in_pointer, out_pointer and possibly a scales vector pointer. - - weight_index : int - The index of weight if the weight pointer pointing to multiple weights. - - Return: - The tensor produced by that layer. - - ''' - assert host_context_lengths is not None or not default_net( - ).plugin_config.remove_input_padding - - trt.get_plugin_registry().plugin_creator_list - in_hidden_size_field = trt.PluginField( - "in_hidden_size", np.array(in_hidden_size, dtype=np.int32), - trt.PluginFieldType.INT32) - out_hidden_size_field_list = [ - trt.PluginField(f"out_hidden_size_{i}", np.array(o, dtype=np.int32), - trt.PluginFieldType.INT32) - for i, o in enumerate(out_hidden_sizes) - ] - transa = 1 if transa else 0 - transa = trt.PluginField("transa", np.array(transa, dtype=np.int32), - trt.PluginFieldType.INT32) - transb = 1 if transb else 0 - transb = trt.PluginField("transb", np.array(transb, dtype=np.int32), - trt.PluginFieldType.INT32) - - plg_creator = trt.get_plugin_registry().get_plugin_creator( - 'Lora', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert plg_creator is not None - - p_dtype = default_net().plugin_config.lora_plugin - pf_type = trt.PluginField( - "type_id", np.array([int(str_dtype_to_trt(p_dtype))], np.int32), - trt.PluginFieldType.INT32) - remove_input_padding = trt.PluginField( - "remove_input_padding", - np.array(np.int8(default_net().plugin_config.remove_input_padding), - dtype=np.int8), trt.PluginFieldType.INT8) - max_low_rank_field = trt.PluginField("max_low_rank", - np.array(max_low_rank, dtype=np.int32), - trt.PluginFieldType.INT32) - weight_index_field = trt.PluginField("weight_index", - np.array(weight_index, dtype=np.int32), - trt.PluginFieldType.INT32) - num_lora_modules = len(out_hidden_sizes) - num_lora_modules_field = trt.PluginField( - "num_lora_modules", np.array(num_lora_modules, dtype=np.int32), - trt.PluginFieldType.INT32) - - pfc = trt.PluginFieldCollection([ - in_hidden_size_field, transa, transb, num_lora_modules_field, pf_type, - remove_input_padding, max_low_rank_field, weight_index_field - ] + out_hidden_size_field_list) - lora_plug = plg_creator.create_plugin("lora", pfc) - - plug_inputs = [input.cast(p_dtype), host_request_types - ] + lora_ranks + lora_weights_pointers - - if default_net().plugin_config.remove_input_padding: - plug_inputs += [host_context_lengths] - - plug_inputs = [i.trt_tensor for i in plug_inputs] - layer = default_trtnet().add_plugin_v2(plug_inputs, lora_plug) - - if num_lora_modules == 1: - return _create_tensor(layer.get_output(0), layer).cast(input.dtype) - else: - return [ - _create_tensor(layer.get_output(i), layer).cast(input.dtype) - for i in range(num_lora_modules) - ] - - -def dora_plugin(activations: Tensor, - out_hidden_sizes: list[int], - lora_weights_pointers: list[Tensor], - host_request_types: Tensor, - host_context_lengths: Tensor | None = None) -> Tensor: - ''' - The DoRA plugin applies column-wise scaling to the output of a LoRA layer. - - Parameters: - input : Tensor (On GPU) - The input tensor. Its shape is [batch_size, seq_len, dim] or [num_tokens, dim] for remove_input_padding - - out_hidden_sizes : list[int] - The output hidden size of each adapter in the related LoRA module. - For example, for a qkv projection out_hidden_sizes should be [q_dim, k_dim, v_dim]. - - host_request_types : Tensor = None - 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/source/advanced/gpt-attention.md, - - host_context_lengths: cpu Tensor = None - A host tensor that contains the lengths of the different inputs, - - Return: - The tensor produced by that layer. - - ''' - assert host_context_lengths is not None or not default_net( - ).plugin_config.remove_input_padding - - dora_plg_creator = trt.get_plugin_registry().get_creator( - 'Dora', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert dora_plg_creator is not None - - out_hidden_sizes = trt.PluginField( - f"out_hidden_sizes", np.array(out_hidden_sizes, dtype=np.int32), - trt.PluginFieldType.INT32) - - remove_input_padding = trt.PluginField( - "remove_input_padding", - np.array(np.int8(default_net().plugin_config.remove_input_padding), - dtype=np.int8), trt.PluginFieldType.INT8) - - lora_dtype = default_net().plugin_config.lora_plugin - type_id = trt.PluginField( - "type", np.array(int(str_dtype_to_trt(lora_dtype)), np.int32), - trt.PluginFieldType.INT32) - - pfc = trt.PluginFieldCollection( - [type_id, remove_input_padding, out_hidden_sizes]) - - dora_plug = dora_plg_creator.create_plugin("dora", pfc, - trt.TensorRTPhase.BUILD) - - plug_inputs = [activations.cast(lora_dtype), host_request_types - ] + lora_weights_pointers - - if default_net().plugin_config.remove_input_padding: - plug_inputs += [host_context_lengths] - - plug_inputs = [i.trt_tensor for i in plug_inputs] - - layer = default_trtnet().add_plugin_v3(plug_inputs, [], dora_plug) - _add_plugin_info(layer, dora_plg_creator, "dora", pfc) - output = _create_tensor(layer.get_output(0), layer).cast(activations.dtype) - return output - - -def mamba_conv1d(input: Tensor, - conv_state_or_ptr: Tensor, - conv_weight: Tensor, - conv_bias: Tensor, - host_request_types: Tensor, - last_token_ids: Tensor, - dim: int, - dconv: int, - dtype: str, - pre_stride: int = 0, - post_stride: int = 0, - host_context_lengths: Optional[Tensor] = None, - slot_mapping: Optional[Tensor] = None, - apply_silu: bool = True): - ''' - Parameters: - input : Tensor (On GPU) - The input tensor. Its shape is [batch_size, seq_len, dim] or [num_tokens, dim] for remove_input_padding - - conv_state_or_ptr : Tensor (On GPU or CPU) - The conv state tensor. Its shape is [batch_size, dconv - 1, dim] - Or the CPU tensor of shape [1] for the pointer of paged states. - - conv_weight : Tensor (On GPU) - The weight tensor. Its shape is [1, dconv, dim] - - conv_bias : Tensor (On GPU) - The bias tensor. Its shape is [dim] - - host_request_types : Tensor (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/source/advanced/gpt-attention.md, - - last_token_ids : Tensor (On GPU) - The inclusive prefix-sum of the lengths or the lengths of the - sequences in the batch. - - dim : int - The hidden dimension of conv1d - - dconv : int - The window size of conv1d - - dtype: str - data type - - pre_stride : int = 0 - The (pre) stride size of the input tensor. - The valid values of the input tensor are input[..., pre_stride: dim-post_stride] - - post_stride : int = 0 - The (post) stride size of the input tensor. - The valid values of the input tensor are input[..., pre_stride: dim-post_stride] - - host_context_lengths: Tensor (On CPU) (Optional) - A host tensor that contains the lengths of the different inputs, - - slot_mapping: Tensor (On GPU) (Optional) - Real page index in state. Its shape is [dim], used for paged state, each page shape is [dconv, dim] - - apply_silu: bool - Is there a SiLU operation after the conv1d? When True apply - SiLU activation function after the conv1d. - ''' - assert host_request_types is not None - mamba_conv1d_plg_creator = trt.get_plugin_registry().get_plugin_creator( - 'MambaConv1d', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert mamba_conv1d_plg_creator is not None - - dim = trt.PluginField("dim", np.array(dim, dtype=np.int32), - trt.PluginFieldType.INT32) - dconv = trt.PluginField("dconv", np.array(dconv, dtype=np.int32), - trt.PluginFieldType.INT32) - pre_stride = trt.PluginField("pre_stride", - np.array(pre_stride, dtype=np.int32), - trt.PluginFieldType.INT32) - post_stride = trt.PluginField("post_stride", - np.array(post_stride, dtype=np.int32), - trt.PluginFieldType.INT32) - pf_type = trt.PluginField( - "type_id", np.array([int(str_dtype_to_trt(dtype))], np.int32), - trt.PluginFieldType.INT32) - remove_input_padding = trt.PluginField( - "remove_input_padding", - np.array(np.int8(default_net().plugin_config.remove_input_padding), - dtype=np.int8), trt.PluginFieldType.INT8) - paged_state = trt.PluginField( - "paged_state", - np.array(np.int8(default_net().plugin_config.paged_state), - dtype=np.int8), trt.PluginFieldType.INT8) - apply_silu = trt.PluginField("apply_silu", - np.array(np.int8(apply_silu), dtype=np.int8), - trt.PluginFieldType.INT8) - - pfc = trt.PluginFieldCollection([ - dim, dconv, pre_stride, post_stride, pf_type, remove_input_padding, - paged_state, apply_silu - ]) - mamba_conv1d_plug = mamba_conv1d_plg_creator.create_plugin( - "mamba_conv1d", pfc) - plug_inputs = [ - input, conv_state_or_ptr, conv_weight, conv_bias, host_request_types, - last_token_ids - ] - if default_net().plugin_config.remove_input_padding: - plug_inputs += [host_context_lengths] - if default_net().plugin_config.paged_state: - plug_inputs += [slot_mapping] - plug_inputs = [i.trt_tensor for i in plug_inputs] - - layer = default_trtnet().add_plugin_v2(plug_inputs, mamba_conv1d_plug) - _add_plugin_info(layer, mamba_conv1d_plg_creator, "mamba_conv1d", pfc) - output = _create_tensor(layer.get_output(0), layer) - if default_net().plugin_config.paged_state: - return output, None - else: - present_state = _create_tensor(layer.get_output(1), layer) - return output, present_state - - -def selective_scan(input: Tensor, - state_or_ptr: Tensor, - delta: Tensor, - delta_bias: Tensor, - A: Tensor, - BC: Tensor, - D: Tensor, - host_request_types: Tensor, - last_token_ids: Tensor, - dim: int, - dstate: int, - dt_rank: int, - delta_softplus: bool, - dtype: str, - z: Optional[Tensor] = None, - host_context_lengths: Optional[Tensor] = None, - slot_mapping: Optional[Tensor] = None, - nheads: int = 1, - ngroups: int = 1, - chunk_size: int = 256, - mamba_version: str = 'Mamba1'): - ''' - Parameters: - input : Tensor (On GPU) - The input tensor. Its shape is [batch_size, seq_len, dim] - - state_or_ptr : Tensor (On GPU or CPU) - The ssm state tensor. Its shape is [batch_size, dstate, dim] - Or the CPU tensor of shape [1] for the pointer of paged states. - - delta : Tensor (On GPU) - The delta tensor. - mamba: Its shape is [batch_size, seq_len, dim] or [num_tokens, dim] for remove_input_padding - mamba2: Its shape is [batch_size, seq_len, nheads] or [num_tokens, nheads] for remove_input_padding - - delta_bias : Tensor (On GPU) - The delta bias tensor. - mamba: Its shape is [dim] - mamba2: Its shape is [nheads] - - A : Tensor (On GPU) - A matrix. - mamba: Its shape is [dstate, dim] - mamba2: Its shape is [nheads] - - BC : Tensor (On GPU) - B and C matrix. - mamba: Its shape is [batch_size, seq_len, dstate * 2] or [num_tokens, dstate * 2] for remove_input_padding - mamba2: Its shape is [batch_size, seq_len, ngroups * dstate * 2] or [num_tokens, ngroups * dstate * 2] for remove_input_padding - - D : Tensor (On GPU) - D matrix. - mamba: Its shape is [dim] - mamba2: Its shape is [nheads] - - host_request_types : Tensor (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/source/advanced/gpt-attention.md - - last_token_ids : Tensor (On GPU) - The inclusive prefix-sum of the lengths or the lengths of the - sequences in the batch. - - dim : int - The inner dimension of SSM block - - dstate : int - The state dimension of SSM block - - dt_rank: int - The rank dimension of dt_proj - - delta_softplus : bool - Do we apply softplus to the delta. - - dtype: str - data type - - z : Tensor (On GPU) (Optional) - The z tensor. Its shape is [batch_size, seq_len, dim] or [num_tokens, dim] for remove_input_padding - - host_context_lengths: Tensor (On CPU) (Optional) - A host tensor that contains the lengths of the different inputs, - - slot_mapping: Tensor (On GPU) (Optional) - Real page index in state. Its shape is [dim], used for paged state, each page shape is [dstate, dim] - - nheads: int (Optional) - The number of heads. - - ngroups: int (Optional) - The number of groups. - - chunk_size: int (Optional) - The chunk_size is used for the chunk_scan kernel. - - mamba_version: int (Optional) - Mamba version, support Mamba1 as default. - ''' - assert host_request_types is not None - selective_scan_plg_creator = trt.get_plugin_registry().get_plugin_creator( - 'SelectiveScan', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert selective_scan_plg_creator is not None - - dim = trt.PluginField("dim", np.array(dim, dtype=np.int32), - trt.PluginFieldType.INT32) - dstate = trt.PluginField("dstate", np.array(dstate, dtype=np.int32), - trt.PluginFieldType.INT32) - dt_rank = trt.PluginField("dt_rank", np.array(dt_rank, dtype=np.int32), - trt.PluginFieldType.INT32) - nheads = trt.PluginField("nheads", np.array(nheads, dtype=np.int32), - trt.PluginFieldType.INT32) - ngroups = trt.PluginField("ngroups", np.array(ngroups, dtype=np.int32), - trt.PluginFieldType.INT32) - chunk_size = trt.PluginField("chunk_size", - np.array(chunk_size, dtype=np.int32), - trt.PluginFieldType.INT32) - delta_softplus = trt.PluginField( - "delta_softplus", np.array(np.int8(delta_softplus), dtype=np.int8), - trt.PluginFieldType.INT8) - pf_type = trt.PluginField( - "type_id", np.array([int(str_dtype_to_trt(dtype))], np.int32), - trt.PluginFieldType.INT32) - remove_input_padding = trt.PluginField( - "remove_input_padding", - np.array(np.int8(default_net().plugin_config.remove_input_padding), - dtype=np.int8), trt.PluginFieldType.INT8) - paged_state = trt.PluginField( - "paged_state", - np.array(np.int8(default_net().plugin_config.paged_state), - dtype=np.int8), trt.PluginFieldType.INT8) - if z is None: - z_enabled = trt.PluginField("z_enabled", np.array(0, dtype=np.int8), - trt.PluginFieldType.INT8) - else: - z_enabled = trt.PluginField("z_enabled", np.array(1, dtype=np.int8), - trt.PluginFieldType.INT8) - is_mamba2 = trt.PluginField( - "is_mamba2", - np.array(1 if mamba_version == 'Mamba2' else 0, dtype=np.int8), - trt.PluginFieldType.INT8) - - pfc = trt.PluginFieldCollection([ - dim, dstate, dt_rank, nheads, ngroups, chunk_size, delta_softplus, - pf_type, remove_input_padding, paged_state, z_enabled, is_mamba2 - ]) - selective_scan_plug = selective_scan_plg_creator.create_plugin( - "selective_scan", pfc) - - plug_inputs = [ - input, state_or_ptr, delta, delta_bias, A, BC, D, host_request_types, - last_token_ids - ] - if default_net().plugin_config.remove_input_padding: - plug_inputs += [host_context_lengths] - if default_net().plugin_config.paged_state: - plug_inputs += [slot_mapping] - if z is not None: - plug_inputs += [z] - plug_inputs = [i.trt_tensor for i in plug_inputs] - - layer = default_trtnet().add_plugin_v2(plug_inputs, selective_scan_plug) - _add_plugin_info(layer, selective_scan_plg_creator, "selective_scan", pfc) - output = _create_tensor(layer.get_output(0), layer) - if default_net().plugin_config.paged_state: - return output, None - else: - present_state = _create_tensor(layer.get_output(1), layer) - return output, present_state - - -def rg_lru(input: Tensor, - A: Tensor, - state_or_ptr: Tensor, - host_request_types: Tensor, - last_token_ids: Tensor, - dim: int, - dtype: str, - block_size: int = 0, - y: Optional[Tensor] = None, - y_bias: Optional[Tensor] = None, - gate: Optional[Tensor] = None, - gate_bias: Optional[Tensor] = None, - gate_x: Optional[Tensor] = None, - gate_x_bias: Optional[Tensor] = None, - gate_a: Optional[Tensor] = None, - gate_a_bias: Optional[Tensor] = None, - slot_mapping: Optional[Tensor] = None): - ''' - Parameters: - input : Tensor (On GPU) - The input tensor. Its shape is [batch_size, seq_len, dim] - - A : Tensor (On GPU) - A matrix. Its shape is [dim] - - state_or_ptr : Tensor (On GPU or CPU) - The lru state tensor. Its shape is [batch_size, dstate, dim] - Or the CPU tensor of shape [1] for the pointer of paged states. - - host_request_types : Tensor (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/source/advanced/gpt-attention.md, - - last_token_ids : Tensor (On GPU) - The inclusive prefix-sum of the lengths or the lengths of the - sequences in the batch. - - dim : int - The inner dimension of RG_LRU block - - block_size : int - The block size of the block diagonal linear layer. It is used to - support the cases that enable fused gate. - - dtype: str - data type - - y : Tensor (On GPU) (Optional) - The y tensor. Its shape is [batch_size, seq_len, dim] - - y_bias : Tensor (On GPU) (Optional) - The y_bias tensor. Its shape is [dim]. If y_bias is not None, we - will fuse GELU(y + y_bias) in this function. - - gate : Tensor (On GPU) (Optional) - The gate tensor. Its shape is [batch_size, seq_len, 2 * dim]. - If gate is not None, we will fuse the gate_x and gate_a, otherwise - use those two tensors. - - gate_bias : Tensor (On GPU) (Optional) - The gate_bias tensor. Its shape is [2 * block_num, dim // block_num]. - If gate_bias is not None, we will fuse the bias add in this function. - - gate_x : Tensor (On GPU) (Optional) - The gate_x tensor. Its shape is [batch_size, seq_len, dim] - - gate_x_bias : Tensor (On GPU) (Optional) - The gate_x_bias tensor. Its shape is [block_num, dim // block_num]. - If gate_x_bias is not None, we will fuse the bias add in this function. - - gate_a : Tensor (On GPU) (Optional) - The gate_a tensor. Its shape is [batch_size, seq_len, dim] - - gate_a_bias : Tensor (On GPU) (Optional) - The gate_a_bias tensor. Its shape is [block_num, dim // block_num]. - If gate_a_bias is not None, we will fuse the bias add in this function. - - slot_mapping: Tensor (On GPU) (Optional) - Real page index in state. Its shape is [dim], used for paged state, each page shape is [dstate, dim] - ''' - assert host_request_types is not None - lru_plg_creator = trt.get_plugin_registry().get_plugin_creator( - 'LRU', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert lru_plg_creator is not None - assert (gate_x_bias is None) == (gate_a_bias is None) - enable_fuse_gate = gate is not None - has_gate_bias = (gate_bias is not None) or (gate_x_bias is not None) - if enable_fuse_gate: - assert gate is not None - assert block_size > 0 - if has_gate_bias: - assert gate_bias is not None - else: - assert gate_x is not None and gate_a is not None - if has_gate_bias: - assert gate_x_bias is not None and gate_a_bias is not None - - dim = trt.PluginField("dim", np.array(dim, dtype=np.int32), - trt.PluginFieldType.INT32) - block_size = trt.PluginField("block_size", - np.array(block_size, dtype=np.int32), - trt.PluginFieldType.INT32) - pf_type = trt.PluginField( - "type_id", np.array([int(str_dtype_to_trt(dtype))], np.int32), - trt.PluginFieldType.INT32) - remove_input_padding = trt.PluginField( - "remove_input_padding", - np.array(np.int8(default_net().plugin_config.remove_input_padding), - dtype=np.int8), trt.PluginFieldType.INT8) - paged_state = trt.PluginField( - "paged_state", - np.array(np.int8(default_net().plugin_config.paged_state), - dtype=np.int8), trt.PluginFieldType.INT8) - - if y is None: - y_enabled = trt.PluginField("y_enabled", np.array(0, dtype=np.int8), - trt.PluginFieldType.INT8) - else: - y_enabled = trt.PluginField("y_enabled", np.array(1, dtype=np.int8), - trt.PluginFieldType.INT8) - - if y_bias is None: - y_bias_enabled = trt.PluginField("y_bias_enabled", - np.array(0, dtype=np.int8), - trt.PluginFieldType.INT8) - else: - y_bias_enabled = trt.PluginField("y_bias_enabled", - np.array(1, dtype=np.int8), - trt.PluginFieldType.INT8) - - if enable_fuse_gate: - fuse_gate_enabled = trt.PluginField("fuse_gate_enabled", - np.array(1, dtype=np.int8), - trt.PluginFieldType.INT8) - else: - fuse_gate_enabled = trt.PluginField("fuse_gate_enabled", - np.array(0, dtype=np.int8), - trt.PluginFieldType.INT8) - - if has_gate_bias: - gate_bias_enabled = trt.PluginField("gate_bias_enabled", - np.array(1, dtype=np.int8), - trt.PluginFieldType.INT8) - else: - gate_bias_enabled = trt.PluginField("gate_bias_enabled", - np.array(0, dtype=np.int8), - trt.PluginFieldType.INT8) - - pfc = trt.PluginFieldCollection([ - dim, block_size, pf_type, remove_input_padding, paged_state, y_enabled, - y_bias_enabled, fuse_gate_enabled, gate_bias_enabled - ]) - lru_plug = lru_plg_creator.create_plugin("rg_lru", pfc) - - plug_inputs = [ - input, - A, - state_or_ptr, - host_request_types, - last_token_ids, - ] - if default_net().plugin_config.paged_state: - plug_inputs += [slot_mapping] - if y is not None: - plug_inputs += [y] - if y_bias is not None: - plug_inputs += [y_bias] - if enable_fuse_gate: - plug_inputs += [gate] - if has_gate_bias: - plug_inputs += [gate_bias] - else: - plug_inputs += [gate_x, gate_a] - if has_gate_bias: - plug_inputs += [gate_x_bias, gate_a_bias] - plug_inputs = [i.trt_tensor for i in plug_inputs] - - layer = default_trtnet().add_plugin_v2(plug_inputs, lru_plug) - _add_plugin_info(layer, lru_plg_creator, "rg_lru", pfc) - output = _create_tensor(layer.get_output(0), layer) - if default_net().plugin_config.paged_state: - return output, None - else: - present_state = _create_tensor(layer.get_output(1), layer) - return output, present_state - - -def topk(input: Tensor, - k: Union[Tensor, int], - dim: int, - largest: bool = True, - prefer_plugin: bool = True) -> Tuple[Tensor, Tensor]: - ''' - Add an topk operation. - - As explained in the ONNX documentation, - - https://github.com/onnx/onnx/blob/main/docs/Operators.md#topk - - NOTE: One distinction from the ONNX topk op, the output is always sorted - with TensorRT layer. - - Retrieve the top-K largest elements along a specified axis. - Given an input tensor of shape [a_1, a_2, ..., a_n, r] - and integer argument k, return two outputs: - Value tensor of shape [a_1, a_2, ..., a_{axis-1}, k, a_{axis+1}, ... a_n] which contains the values of the top k elements along the specified axis - Index tensor of shape [a_1, a_2, ..., a_{axis-1}, k, a_{axis+1}, ... a_n] which contains the indices of the top k elements (original indices from the input tensor). - - Parameters: - input : Tensor - The input tensor. + @staticmethod + def create_sinusoidal_positions_for_cogvlm_attention_plugin( + num_pos: int, + dim: int, + theta: float = 10000.0, + scale: float = 1.0, + scale_type: RotaryScalingType = RotaryScalingType.none, + vision_start: int = 1, + vision_length: int = 1225, + dtype=np.float32, + ): + if scale_type == RotaryScalingType.linear: + scale = 1.0 / scale + inv_freq = scale / (theta**(np.arange(0, dim, 2) / dim)).astype(dtype) + position_id = np.hstack([ + np.arange(0, vision_start + 1, dtype=dtype), + np.full(vision_length, vision_start + 1, dtype=dtype), + np.arange(vision_start + 2, + num_pos - (vision_length - 1), + dtype=dtype), + ]) + sinusoid_inp = np.expand_dims(np.einsum("i , j -> i j", + position_id, + inv_freq, + dtype=dtype), + axis=-1) + # fuse cos/sin into float2 (cos, sin). + concat = np.concatenate((np.cos(sinusoid_inp), np.sin(sinusoid_inp)), + axis=-1) - k : int - A single positive value corresponding to the number of top elements to retrieve + return inv_freq, concat.reshape(1, -1).astype(dtype) - dim: int - The dimension in which to compute the topk indices. + def create_sinusoidal_positions_long_rope_for_attention_plugin( + num_pos: int, + num_orig_pos: int, + dim: int, + theta: float = 10000.0, + scaling_short_factors: Tensor = 1.0, + scaling_long_factors: Tensor = 1.0, + short_mscale=None, + long_mscale=None, + dtype=np.float32, + ): - largest: bool - Controls whether to return largest or smallest elements + def _calc_mscale(scale): + if scale <= 1.0: + return 1.0 + return math.sqrt(1 + math.log(scale) / math.log(num_orig_pos)) - prefer_plugin : bool - Whether to use the topkLastDim plugin if dim is last dim and k is static. + if short_mscale is None: + short_mscale = _calc_mscale(num_pos / num_orig_pos) + long_mscale = short_mscale + def _compute_sinusoidal_positions(scale_factors, is_short, + for_attention_plugin): + inv_freq = 1 / (scale_factors * + (theta**(np.arange(0, dim, 2) / dim)).astype(dtype)) + sinusoid_inp = np.einsum("i , j -> i j", + np.arange(num_pos, dtype=dtype), + inv_freq, + dtype=dtype) - Returns: - The tensors (values, indices) produced by this topk operation. - ''' - dim = dim_resolve_negative(dim, input.ndim())[0] - if prefer_plugin and dim == input.ndim() - 1 and not isinstance(k, Tensor): - last_dim = input.size(-1) - if last_dim == -1: # dynamic? - last_dim = shape(input, -1) - # since we might need to flatten the input to 2d tensor, - # we need to prepare the output shape - out_shape = [] - for i in range(input.ndim() - 1): - out_shape.append(shape(input, i)) - out_shape = concat(out_shape + [k]) - if input.ndim() == 1: - input_2d = unsqueeze(input, - 0) # special handling of rank-1 dynamic tensor - elif input.ndim() != 2: - input_2d = input.view(concat([-1, last_dim]), - zero_is_placeholder=False) - else: - input_2d = input - plg_creator = trt.get_plugin_registry().get_plugin_creator( - "TopkLastDim", "1", TRT_LLM_PLUGIN_NAMESPACE) - assert plg_creator is not None - is_largest = trt.PluginField( - "is_largest", np.array(1 if largest else 0, dtype=np.int32), - trt.PluginFieldType.INT32) - k = trt.PluginField("k", np.array(k, dtype=np.int32), - trt.PluginFieldType.INT32) - pf_type = trt.PluginField("type_id", - np.array([int(input_2d.dtype)], np.int32), - trt.PluginFieldType.INT32) - pfc = trt.PluginFieldCollection([pf_type, k, is_largest]) - topk_last_dim_plug = plg_creator.create_plugin("topk_last_dim", pfc) - plug_inputs = [input_2d] - plug_inputs = [i.trt_tensor for i in plug_inputs] - layer = default_trtnet().add_plugin_v2(plug_inputs, topk_last_dim_plug) - _add_plugin_info(layer, plg_creator, "topk_last_dim", pfc) - values = _create_tensor(layer.get_output(0), layer) - indices = _create_tensor(layer.get_output(1), layer) - values = values.view(out_shape, zero_is_placeholder=False) - indices = indices.view(out_shape, zero_is_placeholder=False) - else: - # non-plugin path - axes = dim_to_trt_axes(dim) - layer = default_trtnet().add_topk( - input.trt_tensor, - trt.TopKOperation.MAX if largest else trt.TopKOperation.MIN, - k=k if not isinstance(k, Tensor) else 1, - axes=axes) - if isinstance(k, Tensor): - if k.ndim() == 1: - k = squeeze(k, 0) - layer.set_input(1, k.trt_tensor) - values = _create_tensor(layer.get_output(0), layer) - indices = _create_tensor(layer.get_output(1), layer) + if for_attention_plugin: + sinusoid_inp = np.expand_dims(sinusoid_inp, axis=-1) + concat = np.concatenate( + (np.cos(sinusoid_inp), np.sin(sinusoid_inp)), axis=-1) + else: + concat = np.concatenate( + (np.sin(sinusoid_inp), np.cos(sinusoid_inp)), axis=1) + concat = np.expand_dims(concat, axis=0) - return values, indices + mscale = short_mscale if is_short else long_mscale + concat = concat.astype(dtype) * mscale + # gpt attention plugins also need inv_freq. + if for_attention_plugin: + return inv_freq.reshape(1, -1), concat.reshape(1, -1) + else: + return concat -def scatter_nd(input: Tensor, mask: Tensor, source: Tensor) -> Tensor: - ''' - Scatter_nd is a tensor operation that writes or updates values in a tensor based on indices. + return ( + _compute_sinusoidal_positions(scaling_short_factors, True, False), + _compute_sinusoidal_positions(scaling_long_factors, False, False), + _compute_sinusoidal_positions(scaling_short_factors, True, True), + _compute_sinusoidal_positions(scaling_long_factors, False, True), + short_mscale, + ) - Parameters: - input: Tensor - The input tensor to be updated - mask: Tensor - A tensor of indices specifying the locations in data to be updated. - source: Tensor - A tensor of values to be written or scattered into data. - Returns: - New tensor with the same shape as the input tensor data, - where the values from the source tensor are scattered or written into the output tensor - at the locations specified by the mask tensor. - ''' - scatter_layer = default_trtnet().add_scatter(input.trt_tensor, - mask.trt_tensor, - source.trt_tensor, - mode=trt.ScatterMode.ND) - return _create_tensor(scatter_layer.get_output(0), scatter_layer) + @staticmethod + def create_sinusoidal_positions_long_rope( + num_pos: int, + dim: int, + theta: float, + original_max_pos: int, + short_factor: List[float], + long_factor: List[float], + dtype=np.float32, + max_seq_len: Optional[int] = None, + duplicate_data: bool = False, + ): + short_factor = np.array(short_factor, dtype=np.float32) + long_factor = np.array(long_factor, dtype=np.float32) + inv_freq = 1.0 / (theta**(np.arange(0, dim, 2, dtype=np.float32) / dim)) + t_pos = np.arange(np.max([num_pos, original_max_pos]), dtype=np.float32) -def low_latency_gemm(input: Tensor, - mat2: Tensor, - alpha: Optional[np.ndarray] = None, - strict_dtype: Optional[trt.DataType] = None) -> Tensor: - if not default_net().plugin_config.low_latency_gemm_plugin: - raise RuntimeError("Low Latency GEMM is only support with plugin") - elif default_net().plugin_config.low_latency_gemm_plugin != "fp8": - raise RuntimeError("Low Latency GEMM plugin only support fp8") - else: - plg_creator = trt.get_plugin_registry().get_plugin_creator( - "LowLatencyGemm", "1", TRT_LLM_PLUGIN_NAMESPACE) - assert plg_creator is not None - if ((input.dtype != trt.fp8) or ((mat2.dtype) != trt.fp8)): - raise TypeError("Low Latency GEMM only support fp8 input") - if (alpha): - assert (isinstance(alpha, np.ndarray) and alpha.dtype == np.float32 - and alpha.size - == 1), "`alpha` must be passed as a float32 ndarray" - alpha = alpha if alpha else np.array(1.0, dtype=np.float32) - alpha = trt.PluginField("alpha", alpha.flatten(), - trt.PluginFieldType.FLOAT32) + # Choose proper freqs based on max_seq_len. + factor = (long_factor if max_seq_len is None + or max_seq_len > original_max_pos else short_factor) + inv_freq = inv_freq / factor + freqs = np.einsum("i,j->ij", t_pos, inv_freq) + sinusoid_inp = freqs.astype(np.float32)[..., np.newaxis] - if strict_dtype is not None: - assert isinstance(strict_dtype, trt.DataType) - p_dtype = strict_dtype - if (p_dtype not in [trt.float32, trt.float16, trt.bfloat16]): - raise ValueError( - "strict_dtype must be float32, float16 or bfloat16 in low latency gemm plugin" - ) + # Apply scaling + scale = num_pos / original_max_pos + if scale <= 1.0: + scaling_factor = 1.0 else: - raise RuntimeError( - "need to use strict dtype in low latency gemm plugin fp8") - pf_type = trt.PluginField("type_id", np.array([int(p_dtype)], np.int32), - trt.PluginFieldType.INT32) - pfc = trt.PluginFieldCollection([alpha, pf_type]) - low_latency_gemm_plug = plg_creator.create_plugin( - "low_latency_gemm", pfc) - plug_inputs = [input.trt_tensor, mat2.trt_tensor] - layer = default_trtnet().add_plugin_v2(plug_inputs, - low_latency_gemm_plug) - _add_plugin_info(layer, plg_creator, "low_latency_gemm", pfc) - return _create_tensor(layer.get_output(0), layer) - - -class SideStreamIDType(IntEnum): - disable = 0 - moe = 1 - - -def low_latency_gemm_swiglu(input: Tensor, - weight: Tensor, - scale_d0: float = 1.0, - scale_d1: float = 1.0, - scale_output: float = 1.0) -> Tensor: - ''' - Add a matrix multiplication, followed by SwiGLU (`x * SiLU(gate)`) operation. - - The second SwiGLU operation takes the preceding tensor, splits it into two halves - along the last dimension, applies SiLU to the second half and multiply the results. The - behaviour is undefined if the last dimension is not even. - - Parameters: - input : Tensor - The first tensor (often called A). - - weight : Tensor - The second tensor (often called B). - - scale_d0 : float - The scale for dequantizing x, used for fp8 - - scale_d1 : float - The scale for dequantizing gate, used for fp8 - - scale_output : float - The scale for quantizing output, used for fp8 - - Returns: - The tensor produced by the inserted layer. - ''' - plg_creator = trt.get_plugin_registry().get_plugin_creator( - 'LowLatencyGemmSwiglu', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert plg_creator is not None - - p_dtype = default_net().plugin_config.low_latency_gemm_swiglu_plugin - pf_type = trt.PluginField( - "type_id", np.array([int(str_dtype_to_trt(p_dtype))], np.int32), - trt.PluginFieldType.INT32) - pf_scale_d0 = trt.PluginField("scale_d0", - np.array(scale_d0, dtype=np.float32), - trt.PluginFieldType.FLOAT32) - pf_scale_d1 = trt.PluginField("scale_d1", - np.array(scale_d1, dtype=np.float32), - trt.PluginFieldType.FLOAT32) - pf_scale_output = trt.PluginField("scale_output", - np.array(scale_output, dtype=np.float32), - trt.PluginFieldType.FLOAT32) - - pfc = trt.PluginFieldCollection( - [pf_type, pf_scale_output, pf_scale_d0, pf_scale_d1]) - low_latency_gemm_swiglu_plug = plg_creator.create_plugin( - "low_latency_gemm_swiglu", pfc) - - plug_inputs = [input.trt_tensor, weight.trt_tensor] - - layer = default_trtnet().add_plugin_v2(plug_inputs, - low_latency_gemm_swiglu_plug) - - return _create_tensor(layer.get_output(0), layer) - + scaling_factor = np.sqrt(1.0 + + np.log(scale) / np.log(original_max_pos)) -def cuda_stream_sync(input_list: List[Tensor], - side_stream_id: SideStreamIDType) -> Tensor: - ''' - Wait for the side stream on the main stream. - output = input_list[0] + if duplicate_data: + sinusoid_inp = np.concatenate((sinusoid_inp, sinusoid_inp), axis=-2) - Parameters: - input_list : List[Tensor] (On GPU) - The list of input tensors. - side_stream_id : int (On CPU) - The side stream ID. - ''' - plg_creator = trt.get_plugin_registry().get_plugin_creator( - "CudaStream", "1", TRT_LLM_PLUGIN_NAMESPACE) - assert plg_creator is not None + # fuse cos/sin into float2 (cos, sin). + concat = np.concatenate( + (np.cos(sinusoid_inp) * scaling_factor, + np.sin(sinusoid_inp) * scaling_factor), + axis=-1, + ) - p_side_stream_id = trt.PluginField("side_stream_id", - np.array(side_stream_id, dtype=np.int32), - trt.PluginFieldType.INT32) - p_num_inputs = trt.PluginField("num_inputs", - np.array(len(input_list), dtype=np.int32), - trt.PluginFieldType.INT32) - pf_type = trt.PluginField( - "type_id", np.array([int(input_list[0].dtype)], dtype=np.int32), - trt.PluginFieldType.INT32) - pfc = trt.PluginFieldCollection([p_side_stream_id, p_num_inputs, pf_type]) - plug = plg_creator.create_plugin("cuda_stream", pfc) - plug_inputs = [input.trt_tensor for input in input_list] + return None, concat.reshape(1, -1).astype(dtype) - layer = default_trtnet().add_plugin_v2(plug_inputs, plug) - _add_plugin_info(layer, plg_creator, "cuda_stream", pfc) - output = _create_tensor(layer.get_output(0), layer) - return output + @staticmethod + def create_fake_weight(dim: int, dtype=np.half): + return np.random.rand(dim).astype(dtype) + # Note: When not using deepseek_yarn, make sure to set mscale_all_dim to 0.0. + @staticmethod + def create_sinusoidal_positions_yarn( + num_pos: int, + dim: int, + base: int = 10000, + scaling_factor: float = 1.0, + original_max_position_embeddings: int = 4096, + beta_fast: int = 32, + beta_slow: int = 1, + mscale: float = 1.0, + mscale_all_dim: float = 1.0, + duplicate_data: bool = True, + dtype=None, + ): + if dtype is None: + dtype = torch.float32 -def cp_split_plugin( - input_ids: Tensor, - host_request_types: Tensor, - host_context_lengths: Tensor, # for pad-free input mode - cp_size: int = 1, - cp_rank: int = 0, -) -> Tensor: - ''' - Add an operation to perform splitting for context parallelism. + # Copy from https://huggingface.co/deepseek-ai/DeepSeek-V2/blob/main/modeling_deepseek.py + # Inverse dim formula to find dim based on number of rotations + def yarn_find_correction_dim(num_rotations, dim, base, + max_position_embeddings): + return (dim * math.log(max_position_embeddings / + (num_rotations * 2 * math.pi))) / ( + 2 * math.log(base)) - This operation split the input_ids into cp_size chunks, and return the cp_rank-th - chunk. - When the seqlen % cp_size != 0, the chunk sizes of each rank would be - [seqlen // cp_size, seqlen // cp_size, ..., seqlen - (seqlen // cp_size) * cp_size] + # Find dim range bounds based on rotations + def yarn_find_correction_range(low_rot, high_rot, dim, base, + max_position_embeddings): + low = math.floor( + yarn_find_correction_dim(low_rot, dim, base, + max_position_embeddings)) + high = math.ceil( + yarn_find_correction_dim(high_rot, dim, base, + max_position_embeddings)) + if low < 0: + low = 0 + if high > dim - 1: + high = dim - 1 + return low, high # Clamp values just in case - It inserts a IPluginV3Layer. + def yarn_get_mscale(scale, mscale): + if scale <= 1: + return 1.0 + return 0.1 * mscale * math.log(scale) + 1.0 - Parameters: - input : Tensor - The input tensor contains the indices to split. + def yarn_linear_ramp_mask(min, max, dim): + if min == max: + max += 0.001 # Prevent singularity - host_request_types: 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, + linear_func = (torch.arange(dim, dtype=dtype) - min) / (max - min) + ramp_func = torch.clamp(linear_func, 0, 1) + return ramp_func - host_context_lengths: Tensor = None (On CPU) - A host tensor that contains the lengths of the different inputs + pos_freqs = base**(torch.arange(0, dim, 2, dtype=dtype) / dim) + freq_extra = 1.0 / pos_freqs + freq_inter = 1.0 / (scaling_factor * pos_freqs) - Returns: - The output split tensor. - The length of the output split tensor. - The index for rebuilding the sequence - ''' - plg_creator = trt.get_plugin_registry().get_creator( - 'CpSplit', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert plg_creator is not None + low, high = yarn_find_correction_range( + beta_fast, + beta_slow, + dim, + base, + original_max_position_embeddings, + ) + inv_freq_mask = 1 - yarn_linear_ramp_mask(low, high, dim // 2) + inv_freq = freq_inter * (1 - inv_freq_mask) + freq_extra * inv_freq_mask + t = torch.arange(num_pos, dtype=dtype) + sinusoid_inp = torch.einsum("i,j -> ij", t, inv_freq).unsqueeze(-1) - cp_size = trt.PluginField("cp_size", np.array([int(cp_size)], np.int32), - trt.PluginFieldType.INT32) - cp_rank = trt.PluginField("cp_rank", np.array([int(cp_rank)], np.int32), - trt.PluginFieldType.INT32) + _mscale = float( + yarn_get_mscale(scaling_factor, mscale) / + yarn_get_mscale(scaling_factor, mscale_all_dim)) - pfc = trt.PluginFieldCollection([cp_size, cp_rank]) - cp_split_plug = plg_creator.create_plugin("cp_split", pfc, - trt.TensorRTPhase.BUILD) - plug_inputs = [ - input_ids.trt_tensor, host_request_types.trt_tensor, - host_context_lengths.trt_tensor - ] + if duplicate_data: + emb = torch.cat((sinusoid_inp, sinusoid_inp), dim=-2) + else: + emb = sinusoid_inp - layer = default_trtnet().add_plugin_v3(plug_inputs, [], cp_split_plug) - _add_plugin_info(layer, plg_creator, "cp_split", pfc) - return _create_tensor(layer.get_output(0), - layer), _create_tensor(layer.get_output(2), layer) + concat = torch.cat((torch.cos(emb) * _mscale, torch.sin(emb) * _mscale), + dim=-1) + return inv_freq.numpy(), concat.reshape((1, -1)).to(dtype).numpy() diff --git a/tensorrt_llm/graph_rewriting.py b/tensorrt_llm/graph_rewriting.py deleted file mode 100644 index 64aadd2da7e4..000000000000 --- a/tensorrt_llm/graph_rewriting.py +++ /dev/null @@ -1,645 +0,0 @@ -import inspect -import weakref -from copy import copy -from dataclasses import dataclass, field -from functools import wraps -from typing import Any, Callable, ClassVar, Dict, List, Optional, Set, Tuple, TypeVar - -import tensorrt as trt - -from ._utils import trt_gte -from .logger import logger -from .network import Network - - -class Layer: - """Layer is a wrapper for TensorRT's ILayer with several python-friendly helper functions.""" - - def __init__(self, network: Network, trt_layer: trt.ILayer): - self._network = weakref.ref(network) - self.trt_layer = trt_layer - - assert isinstance(self.network, Network) - assert isinstance(self.trt_layer, trt.ILayer) - - @property - def network(self): - return self._network() - - def get_inputs(self, *indices: int): - """Get the input tensors of the layer. - - Parameters: - idx: the indices of the input tensor, will return all inputs if left empty - - Returns: - List[Tensor] - """ - from .functional import Tensor - - indices = indices if indices else range(self.trt_layer.num_inputs) - - ret = [] - for i in indices: - assert i < self.trt_layer.num_inputs, ( - f"Invalid input index {i} for layer {self.trt_layer.name}" - ) - - tensor = self.trt_layer.get_input(i) - tensor = Tensor(trt_tensor=tensor, network=self.network, is_network_input=False) - ret.append(tensor) - return ret - - def get_outputs(self, *indices: int): - """Get the output tensor of the layer. - - Parameters: - idx: the index of the output tensor - - Returns: - List[Tensor] - """ - from .functional import Tensor - - indices = indices if indices else range(self.trt_layer.num_outputs) - - ret = [] - for i in indices: - assert i < self.trt_layer.num_outputs, ( - f"Invalid output index {i} for layer {self.trt_layer.name}" - ) - - tensor = self.trt_layer.get_output(i) - tensor = Tensor(trt_tensor=tensor, network=self.network, is_network_input=False) - ret.append(tensor) - return ret - - def is_removed(self): - return self.network.is_removed_layer(self) - - def mark_as_removed(self): - """Mark the layer as removed, this will remove the layer from the network.""" - # NOTE, since INetwork python API doesn't provide a way to remove a layer, we actually mark - # the layer as removed in the network. - self.network.mark_removed_layer(self) - - # remove the FLayerInfo if exists - FLayerInfoMemo.instance().remove(self.name) - - def __eq__(self, other: "Layer") -> bool: - if isinstance(other, Layer): - return self.trt_layer == other.trt_layer - if isinstance(other, trt.tensorrt.ILayer): - return self.trt_layer == other - return False - - def __getattr__(self, name: str) -> Any: - return getattr(self.trt_layer, name) - - # Refer to https://docs.nvidia.com/deeplearning/tensorrt/api/python_api/infer/Graph/Layers.html?highlight=elementwise#layers - # for a complete list of TRT layers. - TRT_LAYER_TYPE_TO_LAYER = { - trt.LayerType.CONVOLUTION: trt.IConvolutionLayer, - trt.LayerType.ACTIVATION: trt.IActivationLayer, - trt.LayerType.POOLING: trt.IPoolingLayer, - trt.LayerType.LRN: trt.ILRNLayer, - trt.LayerType.SCALE: trt.IScaleLayer, - trt.LayerType.SOFTMAX: trt.ISoftMaxLayer, - trt.LayerType.DECONVOLUTION: trt.IDeconvolutionLayer, - trt.LayerType.CONCATENATION: trt.IConcatenationLayer, - trt.LayerType.ELEMENTWISE: trt.IElementWiseLayer, - trt.LayerType.UNARY: trt.IUnaryLayer, - trt.LayerType.PADDING: trt.IPaddingLayer, - trt.LayerType.SHUFFLE: trt.IShuffleLayer, - trt.LayerType.REDUCE: trt.IReduceLayer, - trt.LayerType.TOPK: trt.ITopKLayer, - trt.LayerType.GATHER: trt.IGatherLayer, - trt.LayerType.MATRIX_MULTIPLY: trt.IMatrixMultiplyLayer, - trt.LayerType.RAGGED_SOFTMAX: trt.IRaggedSoftMaxLayer, - trt.LayerType.CONSTANT: trt.IConstantLayer, - trt.LayerType.IDENTITY: trt.IIdentityLayer, - trt.LayerType.PLUGIN_V2: trt.IPluginV2Layer, - trt.LayerType.PLUGIN_V3: trt.IPluginV3Layer, - trt.LayerType.SLICE: trt.ISliceLayer, - trt.LayerType.SHAPE: trt.IShapeLayer, - trt.LayerType.PARAMETRIC_RELU: trt.IParametricReLULayer, - trt.LayerType.RESIZE: trt.IResizeLayer, - trt.LayerType.TRIP_LIMIT: trt.ITripLimitLayer, - trt.LayerType.RECURRENCE: trt.IRecurrenceLayer, - trt.LayerType.ITERATOR: trt.IIteratorLayer, - trt.LayerType.LOOP_OUTPUT: trt.ILoopOutputLayer, - trt.LayerType.SELECT: trt.ISelectLayer, - trt.LayerType.FILL: trt.IFillLayer, - trt.LayerType.QUANTIZE: trt.IQuantizeLayer, - trt.LayerType.DEQUANTIZE: trt.IDequantizeLayer, - trt.LayerType.CONDITION: trt.IConditionLayer, - trt.LayerType.CONDITIONAL_INPUT: trt.IIfConditionalInputLayer, - trt.LayerType.CONDITIONAL_OUTPUT: trt.IIfConditionalOutputLayer, - trt.LayerType.ASSERTION: trt.IAssertionLayer, - trt.LayerType.SCATTER: trt.IScatterLayer, - trt.LayerType.EINSUM: trt.IEinsumLayer, - trt.LayerType.GRID_SAMPLE: trt.IGridSampleLayer, - trt.LayerType.ONE_HOT: trt.IOneHotLayer, - trt.LayerType.NON_ZERO: trt.INonZeroLayer, - trt.LayerType.NMS: trt.INMSLayer, - trt.LayerType.REVERSE_SEQUENCE: trt.IReverseSequenceLayer, - trt.LayerType.NORMALIZATION: trt.INormalizationLayer, - trt.LayerType.CAST: trt.ICastLayer, - } - - if trt_gte(10, 8): - TRT_LAYER_TYPE_TO_LAYER[trt.LayerType.DYNAMIC_QUANTIZE] = trt.IQuantizeLayer - - def as_layer(self) -> Any: - """Convert to a actual TensorRT layer object. - - This can be IPluginV2Layer or IConvolutionLayer, so that we can access the actual layer information. - """ - if self.type in self.TRT_LAYER_TYPE_TO_LAYER: - # bypass TRT's bug of retrieving a specific ILayer type in TensorRT - self.trt_layer.__class__ = self.TRT_LAYER_TYPE_TO_LAYER[self.type] - return self.trt_layer - raise NotImplementedError(f"Unknown layer type: {self.type}") - - def __hash__(self): - return id(self.trt_layer) - - -@dataclass -class _Pattern: - name: str - # args helps to pass in/out some information - args: Dict[str, Any] = field(default_factory=dict, init=False) - - def log_info(self, msg: str): - logger.info(f"Pattern {self.name}: {msg}") - - def log_error(self, msg: str): - logger.error(f"Pattern {self.name}: {msg}") - - def log_warn(self, msg: str): - logger.warning(f"Pattern {self.name}: {msg}") - - -class PatternRewriter(_Pattern): - """A pattern rewriter is a class that can match a pattern in the graph and rewrite the matched pattern. - - There are two ways to implement a pattern rewriter, either override match() and rewrite() separately, or - override match_and_rewrite(). - """ - - def __init__( - self, - name: str, - root_layer: Optional[Set[trt.LayerType]] = None, - separate_match_rewrite=False, - ): - """Constructor. - - Parameters: - name: the name of the rewrite pattern - root_layer: the root layer types to start the pattern matching, if not provided, the pattern - will traverse all the layers in the graph. - separate_match_rewrite: if set to True, the pattern should override match() and rewrite() - separately, otherwise, the pattern should override match_and_rewrite() - """ - super().__init__(name) - self.root_layer = root_layer - self._separate_match_rewrite = separate_match_rewrite - - def match(self, layer: Layer) -> bool: - raise NotImplementedError() - - def rewrite(self, layer: Layer) -> None: - raise NotImplementedError() - - def match_and_rewrite(self, layer: Layer) -> bool: - raise NotImplementedError() - - -class PatternAnalyzer(_Pattern): - def __init__(self, name: str, root_layer: Optional[Set[trt.LayerType]]) -> None: - super().__init__(name) - self.root_layer = root_layer - - def match(self, layer: Layer) -> bool: - raise NotImplementedError() - - def analyze(self, subgraph: List[Layer]) -> None: - raise NotImplementedError() - - -class _PatternManager: - PatternType = TypeVar("PatternType") - - def __init__(self): - # records of (benefit, pattern, id) - self.patterns: Dict[str, Tuple[int, _PatternManager.PatternType]] = {} - - def add(self, label: str, pattern: "_PatternManager.PatternType", benefit: int = 0): - assert label not in self.patterns, f"Pattern {label} already exists" - self.patterns[label] = (benefit, pattern) - - def get(self, label: str) -> "_PatternManager.PatternType": - return self.patterns[label][1] - - -class RewritePatternManager(_PatternManager): - def rewrite(self, net: Network, args=None): - modified = True - # TODO: we can optimize this by asking TRT to expose a graph iterator consistent even after - # the graph is modified. - while modified: - modified = False - # Since the graph iterator is hold by the underlying INetwork, we can only rebuild the - # graph cache and match the nodes again. - for layer in net.get_layers(): - if layer.is_removed(): - continue - for profit, pattern in sorted(self.patterns.values(), key=lambda x: x[0]): - pattern.args = args - - if pattern.root_layer is not None and layer.type not in pattern.root_layer: - continue - if pattern._separate_match_rewrite: - if pattern.match(layer): - pattern.rewrite(layer) - modified = True - else: - if pattern.match_and_rewrite(layer): - modified = True - - @staticmethod - def instance(): - return _global_rewrite_pattern_manager - - -class AnalysisPatternManager(_PatternManager): - def analyze(self, graph: Network, args=None): - for layer in graph.get_layers(): - if layer.name in graph.removed_layers: - continue - for benefit, pattern in sorted(self.patterns.values(), key=lambda x: x[0]): - pattern.args = args - - if pattern.root_layer is not None and layer.type not in pattern.root_layer: - continue - if pattern.match(layer): - subgraph = pattern.match(layer) - pattern.analyze(subgraph) - - @staticmethod - def instance(): - return _global_analysis_pattern_manager - - -@dataclass -class FLayerInfo: - """The FLayerInfo is used to track the functional layers in the INetwork and help graph rewriting. - - The lifetime of a FLayer is the same as the corresponding plugin instance in the INetwork. Once the - plugin instance is removed by the graph rewriting, the FLayer will be removed as well. - - WHY this is needed? - In the current implementation, for functional methods, once it is called in Python, it will lower - to a plugin instance in the INetwork. - However, the plugin interface is black box with customized logic, we cannot retrieve necessary - information from it. This is quite different from ILayer, which provides a set of APIs to retrieve - the information. - Therefore, we need to record the high level information in the FLayerInfo, and keep - it consistent during the graph rewriting. - """ - - layer_kind: str # the method name in the functional.py - # Record the raw inputs of the functional layer to be used in the graph rewrite - # NOTE: the raw inputs contains both the constants and Tensors, the Tensors will be also updated by - # graph rewriting APIs such as `replace_all_uses_with` - raw_inputs: Dict[str, Any] - - raw_outputs: List[Any] = field(default_factory=list, init=False) - - # the corresponding ILayer name - layer_name: str = field(init=False, default="") - - # the signature of the functional layer - signature: Any = field(init=False, default=None) - - def __post_init__(self): - from .functional import Tensor - - assert self.layer_kind - - def replace_with_symbols(arg) -> Any: - if arg is None: - return None - if isinstance(arg, Tensor): - return Tensor - if isinstance(arg, (list, tuple)): - return [replace_with_symbols(x) for x in arg] - if isinstance(arg, dict): - return {k: replace_with_symbols(v) for k, v in arg.items()} - - return arg - - def amend_tensor(arg) -> Any: - if arg is None: - return None - if isinstance(arg, Tensor): - arg.network = self.network - if isinstance(arg, (list, tuple)): - [replace_with_symbols(x) for x in arg] - if isinstance(arg, dict): - {k: replace_with_symbols(v) for k, v in arg.items()} - - return arg - - self.signature = ( - self.layer_kind, - {name: replace_with_symbols(value) for name, value in self.raw_inputs.items()}, - ) - - amend_tensor(self.raw_inputs) - - def set_outputs(self, outputs: List[Any]): - self.raw_outputs = outputs - - def get_input(self, name: str) -> Any: - return self.raw_inputs[name] - - def clone_inputs(self): - """Get a shallow copy of the inputs.""" - return copy(self.raw_inputs) - - def replace_input_with(self, src, dst): - """Replace the input `src` with the input `dst` in the raw_inputs. - - src: Tensor - dst: Tensor - """ - from .functional import Tensor - - def replace(arg: Any): - if isinstance(arg, Tensor): - if arg.trt_tensor is src.trt_tensor: - return dst - return arg - elif isinstance(arg, (list, tuple)): - return [replace(x) for x in arg] - elif isinstance(arg, dict): - return {k: replace(v) for k, v in arg.items()} - return arg - - replace(self.raw_inputs) - - def replace_outputs_uses_with(self, net: Network, new_outs: List[Any]): - """Replace the output users with the new outputs. - - new_outs: List[Tensor], the new outputs to replace with - """ - from .functional import Tensor - - assert len(self.raw_outputs) == len(new_outs) - for old_out, new_out in zip(self.raw_outputs, new_outs): - assert type(old_out) is type(new_out), ( - f"rewrite error, the output type {type(old_out)} is different from the new output " - f"type {type(new_out)} not match the original output type {type(old_out)}" - ) - - def _swap_tensor_info(new, deprecated): - name = deprecated.trt_tensor.name - deprecated.trt_tensor.name = name + "_deprecated" - from .functional import cast - - new = cast(new, deprecated.dtype) - new.trt_tensor.name = name - - def _reset_network_output_tensors(network, out, new_out): - net_outputs = list() - num_outputs = network._trt_network.num_outputs - need_to_mark = False - for i in range(num_outputs): - net_outputs.append(network._trt_network.get_output(i)) - if out.trt_tensor is net_outputs[i]: - need_to_mark = True - if need_to_mark is False: - return - for output in net_outputs: - network.trt_network.unmark_output(output) - for i in range(num_outputs): - if net_outputs[i] is out.trt_tensor: - network.trt_network.mark_output(new_out.trt_tensor) - new_out.trt_tensor.dtype = out.trt_tensor.dtype - else: - network.trt_network.mark_output(net_outputs[i]) - - def replace_all_uses_with(out, new_out): - if isinstance(out, Tensor): - assert isinstance(new_out, Tensor) - out.replace_all_uses_with(new_out) - _swap_tensor_info(new_out, out) - _reset_network_output_tensors(net, out, new_out) - elif isinstance(out, list): - assert isinstance(new_out, list) - for x, y in zip(out, new_out): - replace_all_uses_with(x, y) - elif isinstance(out, dict): - assert isinstance(new_out, dict) - for k, v in out.items(): - replace_all_uses_with(v, new_out[k]) - elif isinstance(out, tuple): - assert isinstance(new_out, tuple) - for x, y in zip(out, new_out): - replace_all_uses_with(x, y) - - replace_all_uses_with(self.raw_outputs, new_outs) - - def __hash__(self) -> int: - return hash(self.signature) - - def __repr__(self) -> str: - return "".format(self.signature) - - @staticmethod - def _get_spec(arg): - """Get the spec that could impact on the Module's topology in the `forward` method.""" - from .functional import Tensor - - # For scalars, we track their value since they are constant - if arg is None: - return None - elif isinstance(arg, (bool, int, str)): - return arg - # For tensors, currently we only track their type, since they are variables - elif isinstance(arg, Tensor): - return Tensor - elif isinstance(arg, (list, tuple)): - return [FLayerInfo._get_spec(x) for x in arg] - # NOTE Free to add more types here is broken, carefully note that, from the engine building angle, - # all the constants should be captured while for the network variables, their types as placeholders - # are enough. - else: - raise TypeError(f"unsupported input type detected: {type(arg)}") - - -@dataclass -class FLayerInfoMemo: - """FLayerInfoMemo holds the FLayer of all the necessary functional layers.""" - - data: Dict[str, FLayerInfo] = field(default_factory=dict, init=False) - - cur_flayer: ClassVar[Optional[FLayerInfo]] = None - - def add(self, layer_name: str, layer: FLayerInfo) -> None: - assert layer_name not in self.data, f"FLayer {layer_name} already exists in FLayerMemo" - self.data[layer_name] = layer - - def create(self, fn: Callable, *args, **kwargs) -> FLayerInfo: - """Add a FLayer to the memo.""" - return FLayerInfo(fn.__name__, self.get_function_arg_dict(fn, *args, **kwargs)) - - def get(self, layer_name: str) -> Optional[FLayerInfo]: - return self.data.get(layer_name, None) - - def remove(self, layer_name: str) -> None: - if layer_name in self.data: - del self.data[layer_name] - - @staticmethod - def instance() -> "FLayerInfoMemo": - """A singleton instance of FLayerMemo.""" - from ._common import default_net - - return default_net().flayer_memo - - @staticmethod - def get_function_arg_dict(f: Callable, *args, **kwargs): - """Get the input argument dict of a function.""" - sig = inspect.signature(f) - - bound_args = sig.bind(*args, **kwargs) - bound_args.apply_defaults() - - return {k: v for k, v in bound_args.arguments.items() if k != "self"} - - -class FLayerScope: - """FLayerScope is used to capture the plugin within a functional method.""" - - def __init__(self, fn, *args, **kwargs): - self.layer = FLayerInfoMemo.instance().create(fn, *args, **kwargs) - - def __enter__(self): - assert FLayerInfoMemo.cur_flayer is None, "FLayerMemo is not reentrant" - # There is no FLayer hierarchy, since the functional layers are not nested - FLayerInfoMemo.cur_flayer = self.layer - - def __exit__(self, exc_type, exc_val, exc_tb): - FLayerInfoMemo.cur_flayer = None - if exc_type is None: - assert self.layer.layer_name != "", ( - f"FLayer {self.layer.layer_kind} without a plugin name detected" - ) - FLayerInfoMemo.instance().add(self.layer.layer_name, self.layer) - - -def record_signature(f): - """Helps to decorate a functional method and record its metadata with a FLayerInfo.""" - - @wraps(f) - def wrapper(*args, **kwargs): - with FLayerScope(f, *args, **kwargs): - outs = f(*args, **kwargs) - FLayerInfoMemo.cur_flayer.set_outputs(outs) - return outs - - return wrapper - - -# singletons -_global_rewrite_pattern_manager = RewritePatternManager() -_global_analysis_pattern_manager = AnalysisPatternManager() - - -class FuseAttentionWithBiasPass(PatternRewriter): - def __init__(self): - super().__init__(name="fuse_attention_with_bias", separate_match_rewrite=False) - - @staticmethod - def is_attention_plugin(layer: Layer) -> bool: - if layer.as_layer().type != trt.LayerType.PLUGIN_V2: - return False - p = layer.as_layer().plugin - conds = [ - p.plugin_namespace == "tensorrt_llm", - p.plugin_type == "GPTAttention", - p.num_outputs == 2, - ] - return all(conds) - - @staticmethod - def is_elementwise_sum(layer: Layer) -> bool: - l = layer.as_layer() # noqa: E741 - if l.type != trt.LayerType.ELEMENTWISE: - return False - return l.op == trt.ElementWiseOperation.SUM - - @staticmethod - def get_eltwise_inputs(layer: Layer): - const_inputs = [] - mutable_inputs = [] - - from .functional import Tensor - - def const_foldable(tensor: Tensor, depth=0) -> bool: - max_depth = 10 - layer = tensor.get_parent() - if layer is None or depth > max_depth: - return False - if layer.type == trt.LayerType.CONSTANT and len(layer.get_inputs()) == 0: - return True - for _ in layer.get_inputs(): - if not const_foldable(_, depth + 1): - return False - return True - - for input in layer.get_inputs(): - if const_foldable(input): - const_inputs.append(input) - else: - mutable_inputs.append(input) - return const_inputs, mutable_inputs - - def match_and_rewrite(self, layer: Layer) -> bool: - from tensorrt_llm.network import net_guard - - with net_guard(layer.network): - if not self.is_attention_plugin(layer): - return False - plugin_flayer = FLayerInfoMemo.instance().get(layer.name) - input = plugin_flayer.raw_inputs["qkv"] - if input is None or isinstance(input, list) or len(list(input.get_users())) != 1: - return False - parent_layer = input.get_parent() - if not self.is_elementwise_sum(parent_layer): - return False - eltwise_const_inputs, eltwise_mutable_inputs = self.get_eltwise_inputs(parent_layer) - if len(eltwise_const_inputs) != 1 or len(eltwise_mutable_inputs) != 1: - return False - if plugin_flayer.raw_inputs["qkv_bias"] is not None: - return False - plugin_flayer.raw_inputs["qkv"] = eltwise_mutable_inputs[0] - plugin_flayer.raw_inputs["qkv_bias"] = eltwise_const_inputs[0] - from .functional import gpt_attention - - new_outputs = gpt_attention(**plugin_flayer.raw_inputs) - plugin_flayer.replace_outputs_uses_with(layer.network, new_outputs) - return True - - -def optimize(net): - patterns = RewritePatternManager() - patterns.add( - label="fuse_attention_with_bias", - pattern=FuseAttentionWithBiasPass(), - ) - patterns.rewrite(net) diff --git a/tensorrt_llm/layers/__init__.py b/tensorrt_llm/layers/__init__.py deleted file mode 100755 index cdfb8a6f897f..000000000000 --- a/tensorrt_llm/layers/__init__.py +++ /dev/null @@ -1,79 +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. -from .activation import Mish -from .attention import (Attention, AttentionMaskParams, AttentionMaskType, - AttentionParams, BertAttention, BlockSparseAttnParams, - CogVLMAttention, DeepseekV2Attention, - KeyValueCacheParams, MropeParams, PositionEmbeddingType, - SpecDecodingParams) -from .cast import Cast -from .conv import Conv1d, Conv2d, Conv3d, ConvTranspose2d -from .embedding import Embedding, PromptTuningEmbedding -from .language_adapter import LanguageAdapter, LanguageAdapterConfig -from .linear import ColumnLinear, Linear, RowLinear -from .lora import Lora, LoraParams, LoraRuntimeParams -from .mlp import MLP, FusedGatedMLP, GatedMLP -from .moe import MOE, MoeConfig, SharedMoE -from .normalization import GroupNorm, LayerNorm, RmsNorm -from .pooling import AvgPool2d -from .recurrent import FusedRgLru, GroupedLinear, Recurrent, RgLru -from .ssm import Mamba, Mamba2 - -__all__ = [ - 'LayerNorm', - 'RmsNorm', - 'ColumnLinear', - 'Linear', - 'RowLinear', - 'AttentionMaskType', - 'PositionEmbeddingType', - 'Attention', - 'BertAttention', - 'CogVLMAttention', - 'DeepseekV2Attention', - 'GroupNorm', - 'Embedding', - 'PromptTuningEmbedding', - 'Conv2d', - 'ConvTranspose2d', - 'Conv1d', - 'Conv3d', - 'AvgPool2d', - 'Mish', - 'MLP', - 'GatedMLP', - 'FusedGatedMLP', - 'Cast', - 'AttentionParams', - 'AttentionMaskParams', - 'SpecDecodingParams', - 'MropeParams', - 'KeyValueCacheParams', - 'BlockSparseAttnParams', - 'Lora', - 'LoraParams', - 'LoraRuntimeParams', - 'MOE', - 'MoeConfig', - 'SharedMoE', - 'Mamba', - 'Mamba2', - 'Recurrent', - 'GroupedLinear', - 'RgLru', - 'FusedRgLru', - 'LanguageAdapter', - 'LanguageAdapterConfig', -] diff --git a/tensorrt_llm/layers/activation.py b/tensorrt_llm/layers/activation.py deleted file mode 100644 index 94ea283a1cf3..000000000000 --- a/tensorrt_llm/layers/activation.py +++ /dev/null @@ -1,22 +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. -from ..functional import softplus, tanh -from ..module import Module - - -class Mish(Module): - - def forward(self, input): - return input * tanh(softplus(input, beta=1.0, threshold=20.0)) diff --git a/tensorrt_llm/layers/attention.py b/tensorrt_llm/layers/attention.py deleted file mode 100755 index 29b63a4258cb..000000000000 --- a/tensorrt_llm/layers/attention.py +++ /dev/null @@ -1,2808 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -import math -from typing import List, Optional - -import numpy as np -import tensorrt as trt -import torch - -from .._common import default_net, precision -from .._utils import (fp32_array, int32_array, is_same_dtype, set_obj_attrs, - trt_dtype_to_np, trt_dtype_to_str) - -# isort: off -from ..functional import ( - ACT2FN, AllReduceParams, AttentionMaskType, Conditional, LayerNormType, - PositionEmbeddingType, RopeEmbeddingUtils, RotaryScalingType, Tensor, - allgather, arange, bert_attention, cast, clip, concat, constant, embedding, - expand, expand_dims, expand_mask, generate_alibi_biases, identity, - generate_alibi_slopes, generate_logn_scaling, gpt_attention, matmul, - minimum, repeat_interleave, shape, slice, softmax, split, unsqueeze, where) -# isort: on -from ..mapping import Mapping -from ..module import Module, ModuleList -from ..parameter import Parameter -from ..quantization import QuantMode -from ..quantization.functional import dequantize, quantize -from .linear import ColumnLinear, RowLinear -from .lora import LoraRuntimeParams -from .normalization import GroupNorm, LayerNorm, RmsNorm - -layernorm_map = { - LayerNormType.LayerNorm: LayerNorm, - LayerNormType.RmsNorm: RmsNorm, - LayerNormType.GroupNorm: GroupNorm, -} - - -def make_causal_mask(bsz, tgt_len, past_key_values_length, dtype): - _range = arange(start=constant(int32_array(0)), - end=tgt_len, - dtype=trt_dtype_to_str(dtype)) - mask = repeat_interleave(_range, tgt_len, 0).view(concat([tgt_len, - tgt_len])) - mask = where(mask < mask.transpose(-1, -2), 1.0, 0.0) - - zero = constant(fp32_array(0)) - zero = expand_dims(zero, [0, 1]) - zero = expand(zero, concat([tgt_len, past_key_values_length])) - mask = concat([zero, mask], dim=1) - mask *= np.finfo(trt_dtype_to_np(dtype)).min.item() - mask = mask.view(concat([1, 1, tgt_len, tgt_len + past_key_values_length])) - mask = expand(mask, - concat([bsz, 1, tgt_len, tgt_len + past_key_values_length])) - return mask - - -def compute_relative_bias(query_length, - key_length, - num_buckets, - max_distance, - bidirectional, - rel_attn_table, - tp_size=1, - tp_group=None, - tp_rank=None): - - def make_relative_position_bucket(relative_position, bidirectional, - num_buckets, max_distance): - relative_buckets = 0 - if bidirectional: - num_buckets //= 2 - relative_buckets += where(relative_position > 0, num_buckets, 0) - relative_position = relative_position.abs() - else: - relative_position = 0 - minimum(relative_position, 0) - - max_exact = num_buckets // 2 - is_small = relative_position < max_exact - - max_exact_fp = constant(fp32_array(max_exact)) - tmp = cast(relative_position, "float32") / max_exact_fp - tmp = tmp.log() - const1 = math.log(max_distance / max_exact) - const2 = constant(fp32_array(num_buckets - max_exact)) - relative_position_if_large = tmp / const1 * const2 - relative_position_if_large = cast(relative_position_if_large, "int32") - relative_position_if_large = max_exact + relative_position_if_large - relative_position_if_large = minimum(relative_position_if_large, - num_buckets - 1) - - relative_buckets += where(is_small, relative_position, - relative_position_if_large) - return relative_buckets - - context_position = arange(start=constant(int32_array(0)), - end=query_length, - dtype=trt_dtype_to_str(trt.int32)) - context_position = unsqueeze(context_position, -1) - memory_position = arange(start=constant(int32_array(0)), - end=key_length, - dtype=trt_dtype_to_str(trt.int32)) - memory_position = unsqueeze(memory_position, 0) - relative_position = memory_position - context_position - relative_position_bucket = make_relative_position_bucket( - relative_position, # shape (query_length, key_length) - bidirectional, - num_buckets, - max_distance, - ) - # shape (query_length, key_length, num_heads) - values = embedding(relative_position_bucket, - rel_attn_table, - tp_size=tp_size, - tp_group=tp_group, - tp_rank=tp_rank) - # shape (1, num_heads, query_length, key_length) - values = unsqueeze(values.permute([2, 0, 1]), 0) - return values - - -class AttentionMaskParams(object): - - def __init__(self, - self_attention_mask: Tensor = None, - self_attention_packed_mask: Tensor = None, - cross_attention_mask: Tensor = None, - cross_attention_packed_mask: Tensor = None): - self.self_attention_mask = self_attention_mask - self.self_attention_packed_mask = self_attention_packed_mask - self.cross_attention_mask = cross_attention_mask - self.cross_attention_packed_mask = cross_attention_packed_mask - - -class AttentionParams(object): - - def __init__(self, - sequence_length: Tensor = None, - context_lengths: Tensor = None, - host_context_lengths: Tensor = None, - max_context_length: int = None, - host_request_types: Tensor = None, - encoder_input_lengths: Tensor = None, - encoder_max_input_length: Tensor = None, - host_runtime_perf_knobs: Tensor = None, - host_context_progress: Tensor = None): - self.sequence_length = sequence_length - self.context_lengths = context_lengths - self.host_context_lengths = host_context_lengths - # max allowed context length. Required to - # compute scratch memory size. - self.max_context_length = max_context_length - self.host_request_types = host_request_types - - self.encoder_input_lengths = encoder_input_lengths - self.encoder_max_input_length = encoder_max_input_length - - self.host_runtime_perf_knobs = host_runtime_perf_knobs - - self.host_context_progress = host_context_progress - - # const parameters that will be reused by all layers. - self.embed_positions = None - self.rotary_inv_freq = None - self.embed_positions_for_gpt_attention = None - - # auxiliary params to support models with non-homegeneous attn layers requiring - # a different set of rope params. e.g. Gemma3. - self.embed_positions_local = None - self.rotary_inv_freq_local = None - self.embed_positions_for_gpt_attention_local = None - - # long rope const parameters - self.long_rope_embed_positions = None - self.long_rope_rotary_inv_freq = None - self.long_rope_embed_positions_for_gpt_attention = None - self.short_mscale = 1.0 - self.long_mscale = 1.0 - - def fill_attention_const_params_for_rope( - self, - embed_positions: Tensor = None, - rotary_inv_freq: Tensor = None, - embed_positions_for_gpt_attention: Tensor = None, - embed_positions_local: Tensor = None, - rotary_inv_freq_local: Tensor = None, - embed_positions_for_gpt_attention_local: Tensor = None): - self.embed_positions = embed_positions - self.rotary_inv_freq = rotary_inv_freq - self.embed_positions_for_gpt_attention = embed_positions_for_gpt_attention - self.embed_positions_local = embed_positions_local - self.rotary_inv_freq_local = rotary_inv_freq_local - self.embed_positions_for_gpt_attention_local = embed_positions_for_gpt_attention_local - return self - - def fill_attention_const_params_for_long_rope( - self, embed_positions, long_rope_embed_positions, rotary_inv_freq, - long_rope_rotary_inv_freq, embed_positions_for_gpt_attention, - long_rope_embed_positions_for_gpt_attention, short_mscale, - long_mscale): - self.embed_positions = embed_positions - self.long_rope_embed_positions = long_rope_embed_positions - self.rotary_inv_freq = rotary_inv_freq - self.long_rope_rotary_inv_freq = long_rope_rotary_inv_freq - self.embed_positions_for_gpt_attention = embed_positions_for_gpt_attention - self.long_rope_embed_positions_for_gpt_attention = long_rope_embed_positions_for_gpt_attention - self.short_mscale = short_mscale - self.long_mscale = long_mscale - return self - - def is_valid_cross_attn(self, do_cross_attention): - if do_cross_attention: - if self.encoder_input_lengths is None: - return False - if self.encoder_max_input_length is None: - return False - return True - - def is_valid(self, gpt_attention_plugin, remove_input_padding, - use_kv_cache): - if gpt_attention_plugin: - if use_kv_cache and self.sequence_length is None: - return False - if self.context_lengths is None: - return False - if self.host_request_types is None: - return False - if self.max_context_length is None: - return False - if self.host_runtime_perf_knobs is None: - return False - if self.host_context_progress is None: - return False - - if remove_input_padding: - if self.host_context_lengths is None: - return False - if not gpt_attention_plugin: - return False - - return True - - -class SpecDecodingParams: - - def __init__(self, - spec_decoding_is_generation_length_variable: bool = False, - spec_decoding_max_generation_length: int = 1, - spec_decoding_generation_lengths: Tensor = None, - spec_decoding_position_offsets: Tensor = None, - spec_decoding_packed_mask: Tensor = None, - spec_decoding_use: Tensor = None): - - self.spec_decoding_is_generation_length_variable = spec_decoding_is_generation_length_variable - self.spec_decoding_max_generation_length = spec_decoding_max_generation_length - self.spec_decoding_generation_lengths = spec_decoding_generation_lengths - self.spec_decoding_position_offsets = spec_decoding_position_offsets - self.spec_decoding_packed_mask = spec_decoding_packed_mask - self.spec_decoding_use = spec_decoding_use - - -class MropeParams: - - def __init__( - self, - mrope_rotary_cos_sin: Tensor = None, - mrope_position_deltas: Tensor = None, - ): - self.mrope_rotary_cos_sin = mrope_rotary_cos_sin - self.mrope_position_deltas = mrope_position_deltas - - -class KeyValueCacheParams: - - def __init__(self, - past_key_value: List[Tensor] = None, - host_past_key_value_lengths: Tensor = None, - host_max_attention_window_sizes: Tensor = None, - host_sink_token_length: Tensor = None, - kv_cache_block_offsets: Tensor = None, - host_kv_cache_block_offsets: Tensor = None, - host_kv_cache_pool_pointers: Tensor = None, - host_kv_cache_pool_mapping: Tensor = None, - cache_indirection: Tensor = None, - past_key_value_length: Tensor = None, - cross_kv_cache_block_offsets: Tensor = None, - host_cross_kv_cache_block_offsets: Tensor = None, - host_cross_kv_cache_pool_pointers: Tensor = None, - host_cross_kv_cache_pool_mapping: Tensor = None): - self.past_key_value = past_key_value - self.host_past_key_value_lengths = host_past_key_value_lengths - self.host_max_attention_window_sizes = host_max_attention_window_sizes - self.host_sink_token_length = host_sink_token_length - self.kv_cache_block_offsets = kv_cache_block_offsets - self.host_kv_cache_block_offsets = host_kv_cache_block_offsets - self.host_kv_cache_pool_pointers = host_kv_cache_pool_pointers - self.host_kv_cache_pool_mapping = host_kv_cache_pool_mapping - self.cross_kv_cache_block_offsets = cross_kv_cache_block_offsets - self.host_cross_kv_cache_block_offsets = host_cross_kv_cache_block_offsets - self.host_cross_kv_cache_pool_pointers = host_cross_kv_cache_pool_pointers - self.host_cross_kv_cache_pool_mapping = host_cross_kv_cache_pool_mapping - self.cache_indirection = cache_indirection - # self.past_key_value_length = past_key_value_length - - def get_first_past_key_value(self): - if self.past_key_value is None: - return None - return self.past_key_value[0] - - def fill_none_tensor_list(self, list_size): - if self.past_key_value is None: - self.past_key_value = tuple([None] * list_size) - - def is_valid(self, gpt_attention_plugin): - if gpt_attention_plugin: - if self.host_past_key_value_lengths is None: - return False - if self.host_max_attention_window_sizes is None: - return False - if self.host_sink_token_length is None: - return False - if self.cache_indirection is None: - return False - - return True - - -class BlockSparseAttnParams: - - def __init__(self, - block_size: int = 64, - homo_head_pattern: bool = False, - num_local_blocks: int = 16, - vertical_stride: int = 8): - self.block_size = block_size - self.homo_head_pattern = homo_head_pattern - self.num_local_blocks = num_local_blocks - self.vertical_stride = vertical_stride - - -class Attention(Module): - - def __init__(self, - *, - local_layer_idx, - hidden_size, - num_attention_heads, - num_kv_heads=None, - max_position_embeddings=1024, - num_layers=1, - apply_query_key_layer_scaling=False, - attention_head_size=None, - qk_layernorm=False, - layernorm_type=LayerNormType.LayerNorm, - layernorm_share=True, - inner_layernorm=False, - eps=1e-05, - attention_mask_type=AttentionMaskType.padding, - bias=True, - dtype=None, - position_embedding_type=PositionEmbeddingType.learned_absolute, - rotary_embedding_base=10000.0, - rotary_embedding_base_local=1.0, - rotary_embedding_scaling=None, - rotary_embedding_percentage=1.0, - rope_scaling_short_factors=None, - rope_scaling_long_factors=None, - rope_scaling_short_mscale=None, - rope_scaling_long_mscale=None, - original_max_position_embeddings=1024, - tp_group=None, - tp_size=1, - tp_rank=0, - quant_mode: QuantMode = QuantMode(0), - q_scaling=1.0, - cross_attention=False, - relative_attention=False, - max_distance=0, - num_buckets=0, - dense_bias=None, - clip_qkv=None, - alibi_bias_max=8, - skip_cross_kv=False, - max_attn_value=0.0, - block_sparse_params=None, - use_implicit_relative_attention=False, - reorder=False, - enable_qkv=True, - cp_group=[0], - cp_size=1, - cp_rank=0, - max_seqlen_for_logn_scaling=8192, - use_logn_scaling=False, - is_local=False): - super().__init__() - - self.local_layer_idx = local_layer_idx - self.cross_attention = cross_attention - self.attention_mask_type = attention_mask_type - self.attention_head_size = hidden_size // num_attention_heads if attention_head_size is None else attention_head_size - assert num_attention_heads % tp_size == 0, \ - "num_attention_heads must be divisible by tp_size" - self.num_attention_heads = num_attention_heads // tp_size - self.num_attention_kv_heads = ( - num_kv_heads + tp_size - 1 - ) // tp_size if num_kv_heads is not None else self.num_attention_heads - self.num_kv_heads = num_kv_heads if num_kv_heads is not None else self.num_attention_heads - self.hidden_size = hidden_size - self.attention_hidden_size = self.attention_head_size * self.num_attention_heads - self.max_position_embeddings = max_position_embeddings - self.original_max_position_embeddings = original_max_position_embeddings - self.bias = bias - self.tp_group = tp_group - self.tp_size = tp_size - self.tp_rank = tp_rank - self.dtype = dtype - self.dense_bias = dense_bias - if dense_bias is None: - self.dense_bias = bias - self.cp_group = cp_group - self.cp_size = cp_size - self.cp_rank = cp_rank - self.is_local = is_local - - self.num_layers = num_layers - self.apply_query_key_layer_scaling = apply_query_key_layer_scaling - self.norm_factor = math.sqrt(self.attention_head_size) - self.q_scaling = q_scaling - if self.apply_query_key_layer_scaling: - self.norm_factor *= self.num_layers - self.q_scaling *= self.num_layers - # Whether to scale ALiBi bias. Mathematically, it's equivalent to - # normalizing QK after adding bias. - # - False, inv_sqrt_Dh * Q*K^T + alibi_bias - # - True, inv_sqrt_Dh * Q*K^T + inv_sqrt_Dh * alibi_bias - self.scale_alibi_bias = position_embedding_type == PositionEmbeddingType.alibi_with_scale - self.alibi_bias_max = alibi_bias_max - self.position_embedding_type = position_embedding_type - - self.relative_attention = relative_attention - self.max_distance = max_distance - self.num_buckets = num_buckets - self.rotary_embedding_base = rotary_embedding_base - self.rotary_embedding_base_local = rotary_embedding_base_local - self.rotary_embedding_scaling = rotary_embedding_scaling - self.rotary_embedding_scale_type = RotaryScalingType.none - self.rotary_embedding_scale = 1.0 - self.short_mscale = 1.0 - self.long_mscale = 1.0 - self.rotary_embedding_percentage = rotary_embedding_percentage - self.use_implicit_relative_attention = self.relative_attention and use_implicit_relative_attention - self.max_seqlen_for_logn_scaling = max_seqlen_for_logn_scaling - self.use_logn_scaling = use_logn_scaling - if rotary_embedding_scaling is not None: - rotary_scaling_type = rotary_embedding_scaling.get( - "type", rotary_embedding_scaling.get("rope_type")) - self.rotary_embedding_scale_type = RotaryScalingType.from_string( - rotary_scaling_type) - - self.rotary_embedding_scale = rotary_embedding_scaling.get( - "factor", 1.0) - - self.rotary_embedding_dim = 0 - if self.position_embedding_type.is_rope(): - self.rotary_embedding_dim = int(self.attention_head_size * - rotary_embedding_percentage) - elif self.position_embedding_type.is_alibi(): - alibi_scale = 1. / self.norm_factor if self.scale_alibi_bias else 1. - alibi_slopes = generate_alibi_slopes( - self.num_attention_heads * self.tp_size, - tp_size=self.tp_size, - tp_rank=self.tp_rank, - alibi_scale=alibi_scale, - alibi_bias_max=self.alibi_bias_max) - self.register_parameter( - 'alibi_slopes', - Parameter(alibi_slopes, dtype='float32', is_buffer=True)) - - if self.use_logn_scaling: - logn_scaling = generate_logn_scaling( - self.max_seqlen_for_logn_scaling, self.max_position_embeddings) - self.register_parameter( - 'logn_scaling', - Parameter(logn_scaling, dtype='float32', is_buffer=True)) - - self.quant_mode = quant_mode - self.max_attn_value = max_attn_value - self.register_parameter('kv_cache_scaling_factor', None) - self.register_parameter('attention_output_orig_quant_scale', None) - self.register_parameter('attention_output_sf_scale', None) - - self.block_sparse_params = block_sparse_params if block_sparse_params is not None else BlockSparseAttnParams( - ) - - # The output feature size is therefore (h/tp + 2*kvh/tp) * d, where h is num_heads, - # d is head_size, kvh is the num_kv_heads and tp is tensor_parallel_size. - # In ColumnLinear op, the output dim is calculated by (h + 2*kvh) * d / tp, - # which matches the desired output size (h/tp + 2*kvh/tp) * d after splitting - - # out dim is not necessarily hidden_size + kv specific size (in MQA/GQA), but num_heads * heads_size - # example: d_model != num_heads * head_size in Flan-T5/ByT5/Gemma - if enable_qkv: - self.qkv = ColumnLinear( - hidden_size, - tp_size * self.num_attention_heads * self.attention_head_size + - (2 * tp_size * self.num_attention_kv_heads * - self.attention_head_size), - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - gather_output=False, - is_qkv=True) - self.dense = RowLinear(tp_size * self.num_attention_heads * - self.attention_head_size, - hidden_size, - bias=self.dense_bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size) - - # see optimize_model's add_lora for LoRA initialization - self.qkv_lora = None - self.qkv_dora = None - - # per-layer relative attention table - if self.use_implicit_relative_attention: - self.rel_attn_table = Parameter(shape=(num_attention_heads // - tp_size, num_buckets), - dtype=dtype) - - # qk layernorm - self.qk_layernorm = qk_layernorm - self.layernorm_type = layernorm_type - self.layernorm_share = layernorm_share - ln_type = layernorm_map[layernorm_type] - if self.qk_layernorm: - # layernorm_share indicates whether all the QK head in one layer shares the same norm parameters or not - if layernorm_share: - self.q_layernorm = ln_type(self.attention_head_size, - eps=eps, - dtype=dtype) - self.k_layernorm = ln_type(self.attention_head_size, - eps=eps, - dtype=dtype) - else: - assert ln_type == LayerNorm - self.q_layernorm = ln_type( - (self.num_attention_heads, self.attention_head_size), - eps=eps, - dtype=dtype, - bias=False, - tp_size=tp_size, - tp_dim=0) - self.k_layernorm = ln_type( - (self.num_attention_kv_heads, self.attention_head_size), - eps=eps, - dtype=dtype, - bias=False, - tp_size=tp_size, - tp_dim=0) - - self.inner_layernorm = ln_type(self.hidden_size, dtype=dtype, - eps=eps) if inner_layernorm else None - if clip_qkv is not None: - self.clip_qkv = fp32_array([clip_qkv]) - else: - self.clip_qkv = None - - self.skip_cross_kv = skip_cross_kv - - @staticmethod - def create_attention_const_params(model_cls, config): - # get rotary parameters. - hidden_size = config.hidden_size - num_attention_heads = config.num_attention_heads - attention_head_size = config.head_size - max_position_embeddings = config.max_position_embeddings - position_embedding_type = config.position_embedding_type - rotary_embedding_base = getattr(config, 'rotary_base', 10000.0) - rotary_embedding_scaling = getattr(config, 'rotary_scaling', None) - rotary_embedding_percentage = getattr(config, 'rotary_pct', 1.0) - # only rope need the const parameters. - if not position_embedding_type.is_rope(): - return - # attention head size - attention_head_size = hidden_size // num_attention_heads if attention_head_size is None else attention_head_size - # rotary embedding dim. - rotary_embedding_dim = getattr( - config, 'rotary_dim', - int(attention_head_size * rotary_embedding_percentage)) - # rotary scaling. - rotary_embedding_scale_type = RotaryScalingType.none - rotary_embedding_scale = 1.0 - if rotary_embedding_scaling is not None: - rotary_scaling_type = rotary_embedding_scaling.get( - "type", rotary_embedding_scaling.get("rope_type")) - rotary_embedding_scale_type = RotaryScalingType.from_string( - rotary_scaling_type) - rotary_embedding_scale = rotary_embedding_scaling.get("factor", 1.0) - - if position_embedding_type == PositionEmbeddingType.long_rope: - rope_scaling_short_factors, rope_scaling_long_factors = None, None - rope_scaling_short_mscale, rope_scaling_long_mscale = None, None - original_max_position_embeddings = max_position_embeddings - - if hasattr(config, "longrope_scaling_short_factors"): - rope_scaling_short_factors = np.asarray( - config.longrope_scaling_short_factors).astype(np.float32) - rope_scaling_long_factors = np.asarray( - config.longrope_scaling_long_factors).astype(np.float32) - - original_max_position_embeddings = config.original_max_position_embeddings - - if config.architecture == "Phi3SmallForCausalLM" or config.architecture == "PhiMoEForCausalLM": - rope_scaling_short_mscale = config.longrope_short_mscale - rope_scaling_long_mscale = config.longrope_long_mscale - - embed_positions, long_rope_embed_positions, \ - (rotary_inv_freq, embed_positions_for_gpt_attention), \ - (long_rope_rotary_inv_freq, long_rope_embed_positions_for_gpt_attention), mscale \ - = RopeEmbeddingUtils.create_sinusoidal_positions_long_rope_for_attention_plugin( - max_position_embeddings, - original_max_position_embeddings, rotary_embedding_dim, - rotary_embedding_base, rope_scaling_short_factors, - rope_scaling_long_factors, rope_scaling_short_mscale, rope_scaling_long_mscale) - - if rope_scaling_short_mscale is not None: - assert rope_scaling_long_mscale is not None - short_mscale = rope_scaling_short_mscale - long_mscale = rope_scaling_long_mscale - else: - short_mscale = long_mscale = mscale - - model_cls.register_parameter( - 'embed_positions', - Parameter(embed_positions, dtype='float32', is_buffer=True)) - model_cls.register_parameter( - 'long_rope_embed_positions', - Parameter(long_rope_embed_positions, - dtype='float32', - is_buffer=True)) - model_cls.register_parameter( - 'rotary_inv_freq', - Parameter(rotary_inv_freq, dtype='float32', is_buffer=True)) - model_cls.register_parameter( - 'long_rope_rotary_inv_freq', - Parameter(long_rope_rotary_inv_freq, - dtype='float32', - is_buffer=True)) - model_cls.register_parameter( - 'embed_positions_for_gpt_attention', - Parameter(embed_positions_for_gpt_attention, - dtype='float32', - is_buffer=True)) - model_cls.register_parameter( - 'long_rope_embed_positions_for_gpt_attention', - Parameter(long_rope_embed_positions_for_gpt_attention, - dtype='float32', - is_buffer=True)) - model_cls.short_mscale = short_mscale - model_cls.long_mscale = long_mscale - elif rotary_embedding_scale_type == RotaryScalingType.yarn: - beta_fast = rotary_embedding_scaling.get("beta_fast", 32.0) - beta_slow = rotary_embedding_scaling.get("beta_slow", 1.0) - mscale = rotary_embedding_scaling.get("mscale", 1.0) - mscale_all_dim = rotary_embedding_scaling.get("mscale_all_dim", 0.0) - original_max_position_embeddings = rotary_embedding_scaling.get( - "original_max_position_embeddings", 4096) - rotary_inv_freq, embed_positions_for_gpt_attention = RopeEmbeddingUtils.create_sinusoidal_positions_yarn( - max_position_embeddings, rotary_embedding_dim, - rotary_embedding_base, rotary_embedding_scale, - original_max_position_embeddings, beta_fast, beta_slow, mscale, - mscale_all_dim, False) - - embed_positions = RopeEmbeddingUtils.create_sinusoidal_positions( - max_position_embeddings, - rotary_embedding_dim, - ) - model_cls.register_parameter( - 'embed_positions', - Parameter(embed_positions, dtype='float32', is_buffer=True)) - model_cls.register_parameter( - 'rotary_inv_freq', - Parameter(rotary_inv_freq, dtype='float32', is_buffer=True)) - model_cls.register_parameter( - 'embed_positions_for_gpt_attention', - Parameter(embed_positions_for_gpt_attention, - dtype='float32', - is_buffer=True)) - else: - - def register_rope_params(rotary_base, rotary_embedding_scale, - rotary_embedding_scale_type, - rotary_embedding_scaling, - names_to_register): - # Rotary const weights. - embed_positions = RopeEmbeddingUtils.create_sinusoidal_positions( - max_position_embeddings, - rotary_embedding_dim, - ) - - rotary_inv_freq, embed_positions_for_gpt_attention = RopeEmbeddingUtils.create_sinusoidal_positions_for_attention_plugin( - max_position_embeddings, rotary_embedding_dim, rotary_base, - rotary_embedding_scale, rotary_embedding_scale_type, - rotary_embedding_scaling) - model_cls.register_parameter( - names_to_register[0], - Parameter(embed_positions, dtype='float32', is_buffer=True)) - model_cls.register_parameter( - names_to_register[1], - Parameter(rotary_inv_freq, dtype='float32', is_buffer=True)) - model_cls.register_parameter( - names_to_register[2], - Parameter(embed_positions_for_gpt_attention, - dtype='float32', - is_buffer=True)) - - register_rope_params( - rotary_base=rotary_embedding_base, - rotary_embedding_scale=rotary_embedding_scale, - rotary_embedding_scale_type=rotary_embedding_scale_type, - rotary_embedding_scaling=rotary_embedding_scaling, - names_to_register=[ - 'embed_positions', 'rotary_inv_freq', - 'embed_positions_for_gpt_attention' - ]) - - # For models with non-homegeneous attention layers requiring a second set of rope params. e.g. Gemma3. - rotary_embedding_base_local = getattr(config, - 'rope_local_base_freq', None) - if rotary_embedding_base_local is not None: - register_rope_params( - rotary_base=rotary_embedding_base_local, - rotary_embedding_scale=1.0, - rotary_embedding_scale_type=RotaryScalingType.none, - rotary_embedding_scaling=None, - names_to_register=[ - 'embed_positions_local', 'rotary_inv_freq_local', - 'embed_positions_for_gpt_attention_local' - ]) - - @staticmethod - def fill_attention_params(model_cls, attention_params): - if model_cls.position_embedding_type.is_rope(): - if attention_params is None: - attention_params = AttentionParams() - if model_cls.position_embedding_type == PositionEmbeddingType.long_rope: - return attention_params.fill_attention_const_params_for_long_rope( - model_cls.embed_positions.value, - model_cls.long_rope_embed_positions.value, - model_cls.rotary_inv_freq.value, - model_cls.long_rope_rotary_inv_freq.value, - model_cls.embed_positions_for_gpt_attention.value, - model_cls.long_rope_embed_positions_for_gpt_attention.value, - model_cls.short_mscale, model_cls.long_mscale) - else: - return attention_params.fill_attention_const_params_for_rope( - model_cls.embed_positions.value, - model_cls.rotary_inv_freq.value, - model_cls.embed_positions_for_gpt_attention.value, - model_cls.embed_positions_local.value if hasattr( - model_cls, "embed_positions_local") else None, - model_cls.rotary_inv_freq_local.value if hasattr( - model_cls, "rotary_inv_freq_local") else None, - model_cls.embed_positions_for_gpt_attention_local.value - if hasattr( - model_cls, - "embed_positions_for_gpt_attention_local") else None) - # Fill nothing. - return attention_params - - def _get_output_orig_quant_scale(self): - attention_output_orig_quant_scale = self.attention_output_orig_quant_scale.value if self.attention_output_orig_quant_scale is not None else None - if attention_output_orig_quant_scale is not None and ( - default_net().plugin_config.gemm_plugin == 'nvfp4' - or self.quant_mode.has_nvfp4()): - # The scale was intended for nvfp4 quantization: max_value * scale = fp4_max * fp8_max - # So if we want to quantize the output to fp8, the scale should be divided by fp4_max - attention_output_orig_quant_scale = attention_output_orig_quant_scale / 6.0 - return attention_output_orig_quant_scale - - def forward( - self, - hidden_states: Tensor, - attention_mask=None, - attention_packed_mask=None, - use_cache=False, - spec_decoding_params=None, - mrope_params=None, - kv_cache_params=None, - attention_params=None, - encoder_output: Optional[Tensor] = None, - position_embedding=None, - norm_before_bmm1=False, - lora_layer_params=None, - cross_kv_cache_gen: Optional[Tensor] = None, - cross_kv_reuse: Optional[Tensor] = None, - all_reduce_params: Optional[AllReduceParams] = None, - skip_attn=None, - ): - attention_input = hidden_states - - assert isinstance(hidden_states, (Tensor, tuple)) - - spec_decoding_params = SpecDecodingParams( - ) if spec_decoding_params is None else spec_decoding_params - - mrope_params = MropeParams() if mrope_params is None else mrope_params - logn_scaling = None - if self.use_logn_scaling: - logn_scaling = self.logn_scaling.value - - alibi_slopes = None - if self.position_embedding_type.is_alibi(): - alibi_slopes = self.alibi_slopes.value - if default_net().plugin_config.gpt_attention_plugin: - alibi_slopes = cast(alibi_slopes, hidden_states.dtype) - - qkv_lora_params = None - if lora_layer_params is not None: - if not self.cross_attention: - qkv_lora_params = lora_layer_params.get_runtime_params( - 0, "attn_qkv") - else: - qkv_lora_params = lora_layer_params.get_runtime_params( - 0, "cross_attn_qkv") - - unfuse_qkv_gemm = self.qkv is None - if unfuse_qkv_gemm: - qkv_gemm = [self.q, self.k, self.v] - qkv = [gemm(hidden_states) for gemm in qkv_gemm] - if default_net( - ).plugin_config.lora_plugin and qkv_lora_params is not None: - lora = self.qkv.lora(hidden_states, qkv_lora_params) - kv_size = self.attention_head_size * self.num_attention_kv_heads - qkv_lora = split(lora, - [self.attention_hidden_size, kv_size, kv_size], - dim=1) - qkv = [tensor + lora for tensor, lora in zip(qkv, qkv_lora)] - else: - qkv = self.qkv(hidden_states, qkv_lora_params) - - if self.clip_qkv is not None: - qkv = clip(qkv, -self.clip_qkv, self.clip_qkv) - - if default_net().plugin_config.remove_input_padding: - if unfuse_qkv_gemm: - for tensor in qkv: - assert tensor.ndim() == 2 - else: - assert qkv.ndim() == 2 - - if default_net( - ).plugin_config.lora_plugin and qkv_lora_params is None and lora_layer_params is not None: - if not self.cross_attention: - q_lora_params = lora_layer_params.get_runtime_params( - 0, "attn_q") - k_lora_params = lora_layer_params.get_runtime_params( - 0, "attn_k") - v_lora_params = lora_layer_params.get_runtime_params( - 0, "attn_v") - else: - q_lora_params = lora_layer_params.get_runtime_params( - 0, "cross_attn_q") - k_lora_params = lora_layer_params.get_runtime_params( - 0, "cross_attn_k") - v_lora_params = lora_layer_params.get_runtime_params( - 0, "cross_attn_v") - - assert (q_lora_params is not None and k_lora_params is not None and v_lora_params is not None) or \ - (q_lora_params is None and k_lora_params is None and v_lora_params is None), "q_lora_params, k_lora_params and v_lora_params should be all enabled or all disabled at the same time." - - if q_lora_params is not None and k_lora_params is not None and v_lora_params is not None: - qkv_lora_runtime_params = LoraRuntimeParams( - lora_ranks=[ - q_lora_params.lora_ranks[0], - k_lora_params.lora_ranks[0], - v_lora_params.lora_ranks[0], - ], - lora_weights_pointers=[ - q_lora_params.lora_weights_pointers[0], - k_lora_params.lora_weights_pointers[0], - v_lora_params.lora_weights_pointers[0], - ], - host_request_types=q_lora_params.host_request_types, - host_context_lengths=q_lora_params.host_context_lengths, - max_encoder_context_length=q_lora_params. - max_encoder_context_length, - host_encoder_input_lengths=q_lora_params. - host_encoder_input_lengths, - partial_lora_mask=lora_layer_params.partial_lora_mask, - ) - - q_lora, k_lora, v_lora = self.qkv_lora(hidden_states, - qkv_lora_runtime_params) - qkv_lora = concat([q_lora, k_lora, v_lora], - dim=q_lora.rank() - 1) - qkv = qkv + qkv_lora - if self.qkv_dora is not None: - qkv = self.qkv_dora(qkv, qkv_lora_runtime_params) - if self.qk_layernorm: - base_shape = shape(qkv, 0) if qkv.ndim() == 2 else concat( - [shape(qkv, 0), shape(qkv, 1)]) - qkv_sections = [ - self.num_attention_heads, self.num_attention_kv_heads, - self.num_attention_kv_heads - ] - total_heads = sum(qkv_sections) - if self.num_attention_heads != self.num_attention_kv_heads: - qkv = qkv.view( - concat([base_shape, total_heads, self.attention_head_size])) - query, key, value = split(qkv, qkv_sections, dim=qkv.ndim() - 2) - else: - qkv = qkv.view( - concat([ - base_shape, self.num_attention_heads, 3, - self.attention_head_size - ])) - query, key, value = split(qkv, 1, dim=qkv.ndim() - 2) - q_shape = concat([ - base_shape, self.num_attention_heads, - self.attention_head_size - ]) - query = query.view(q_shape) - key = key.view(q_shape) - value = value.view(q_shape) - - normalized_shape = None - if not self.layernorm_share: - normalized_shape = self.attention_head_size - query = self.q_layernorm(query, normalized_shape=normalized_shape) - key = self.k_layernorm(key, normalized_shape=normalized_shape) - qkv = concat([query, key, value], dim=query.ndim() - 2) - qkv = qkv.view( - concat([base_shape, total_heads * self.attention_head_size])) - if self.position_embedding_type == PositionEmbeddingType.chatglm: - qkv = RopeEmbeddingUtils.apply_rotary_pos_emb_chatglm( - qkv, - position_embedding, - self.num_attention_heads, - self.attention_head_size, - self.max_position_embeddings, - self.rotary_embedding_scale, - default_net().plugin_config.remove_input_padding, - ) - self.rotary_embedding_scale_type = RotaryScalingType.none - self.rotary_embedding_scale = 1.0 - - paged_kv_cache = default_net().plugin_config.paged_kv_cache - - assert attention_params is None or attention_params.is_valid( - default_net().plugin_config.gpt_attention_plugin, - default_net().plugin_config.remove_input_padding, use_cache) - - if use_cache: - assert kv_cache_params is None or kv_cache_params.is_valid( - default_net().plugin_config.gpt_attention_plugin) - - past_key_value = None if kv_cache_params is None else kv_cache_params.get_first_past_key_value( - ) - - # if cross attention, cross QKV only needs to be calculated once in the - # 1st decoding step --> write to cross KV cache --> remains constant - # during the entire decoding steps. - # 1st and >1st steps are distinguished by a boolean tensor `cross_kv_cache_gen` passed at runtime - # also, cross KV cache max length is set from encoder output seqlen, - # this maps to the max context length concept in decoder-only models - cross_kv = None - if self.cross_attention and encoder_output: - assert isinstance(encoder_output, Tensor) - - def compute_cross_kv(encoder_output): - if hasattr(self, 'kv'): - # We optimize the graph by adding kv in the cross attention layer, preventing computing the - # query of encoder_output. - assert qkv_lora_params is None, "Not support LoRA when we only compute key/value in cross atteniton" - # see optimization_model's optimize_cross_qkv - cross_kv = self.kv(encoder_output, qkv_lora_params) - base_shape = shape( - cross_kv, 0) if cross_kv.ndim() == 2 else concat( - [shape(cross_kv, 0), - shape(cross_kv, 1)]) - if self.qk_layernorm: - cross_kv = cross_kv.view( - concat([ - base_shape, 2 * self.num_attention_kv_heads, - self.attention_head_size - ])) - - key, value = split(cross_kv, [ - self.num_attention_kv_heads, - self.num_attention_kv_heads - ], - dim=cross_kv.ndim() - 2) - - key = self.k_layernorm(key) - cross_kv = concat([key, value], dim=key.ndim() - 2) - else: - cross_qkv = self.qkv(encoder_output, qkv_lora_params) - base_shape = shape( - cross_qkv, 0) if cross_qkv.ndim() == 2 else concat( - [shape(cross_qkv, 0), - shape(cross_qkv, 1)]) - - cross_qkv = cross_qkv.view( - concat([ - base_shape, self.num_attention_heads + - 2 * self.num_attention_kv_heads, - self.attention_head_size - ])) - - if self.qk_layernorm: - _, key, value = split(cross_qkv, [ - self.num_attention_heads, - self.num_attention_kv_heads, - self.num_attention_kv_heads - ], - dim=cross_qkv.ndim() - 2) - - key = self.k_layernorm(key) - cross_kv = concat([key, value], dim=key.ndim() - 2) - else: - _, cross_kv = split(cross_qkv, [ - self.num_attention_heads, - self.num_attention_kv_heads * 2 - ], - dim=cross_qkv.ndim() - 2) - cross_kv = cross_kv.view( - concat([ - base_shape, 2 * self.num_attention_kv_heads * - self.attention_head_size - ])) - - if default_net( - ).plugin_config.lora_plugin and qkv_lora_params is None and lora_layer_params is not None: - _, cross_k_lora, cross_v_lora = self.qkv_lora( - encoder_output, - qkv_lora_runtime_params, - is_cross_attention=True) - cross_kv_lora = concat([cross_k_lora, cross_v_lora], - dim=cross_k_lora.rank() - 1) - cross_kv = cross_kv + cross_kv_lora - if self.qkv_dora is not None: - cross_kv = self.qkv_dora(cross_kv, - qkv_lora_runtime_params, - is_cross_attention=True) - - return cross_kv - - if self.skip_cross_kv: - conditional = Conditional(cross_kv_cache_gen) - cond_in1 = conditional.add_input(encoder_output) - cond_in2 = conditional.add_input(cross_kv_reuse) - - ## True branch: context phase, compute cross qkv - cross_kv_true = compute_cross_kv(cond_in1) - - ## False branch: generation phase, no compute but need to obey shape constraints - # because TRT's IfConditional requires the output shape of two subgraphs to be identical - # our 1st attempt was to stack encoder_output [B, S, H] or [N, H] --> cross qkv [B, S, 3*H] or [N, 3*H], - # but it still introduces unnecessary concat. A better solution is to create a dummy torch tensor `cross_kv_resue` - # with the correct shape and reuse it in every generation step - cross_kv_false = cond_in2 - cross_kv = conditional.add_output(cross_kv_true, cross_kv_false) - else: - cross_kv = compute_cross_kv(encoder_output) - - if default_net().plugin_config.gpt_attention_plugin: - if self.cross_attention and (past_key_value is not None): - past_key_value = kv_cache_params.past_key_value[1] - assert self.attention_mask_type in [ - AttentionMaskType.causal, AttentionMaskType.bidirectional, - AttentionMaskType.bidirectionalglm, - AttentionMaskType.blocksparse - ], 'Plugin only support masked MHA.' - - # KV cache scales. - if self.kv_cache_scaling_factor is not None: - kv_orig_quant_scale = self.kv_cache_rcp_scaling_factor.value - kv_quant_orig_scale = self.kv_cache_scaling_factor.value - else: - kv_orig_quant_scale = None - kv_quant_orig_scale = None - - # The output SF scale, needed when fuse_fp4_quant is enabled. - attention_output_sf_scale = self.attention_output_sf_scale.value if self.attention_output_sf_scale is not None else None - - # The rotary inv freq can be pre-computed. - rotary_inv_freq = getattr(attention_params, "rotary_inv_freq", None) - # Rotary cos/sin cache. - rotary_cos_sin = getattr(attention_params, - "embed_positions_for_gpt_attention", None) - rotary_inv_freq_local = getattr(attention_params, - "rotary_inv_freq_local", None) - rotary_cos_sin_local = getattr( - attention_params, "embed_positions_for_gpt_attention_local", - None) - - long_rope_rotary_inv_freq = getattr(attention_params, - "long_rope_rotary_inv_freq", - None) - long_rope_rotary_cos_sin = getattr( - attention_params, "long_rope_embed_positions_for_gpt_attention", - None) - - if self.position_embedding_type == PositionEmbeddingType.learned_absolute: - rotary_inv_freq = None - rotary_cos_sin = None - - # check if the cache is provided. - if self.position_embedding_type.is_rope(): - assert (rotary_inv_freq is not None) and ( - rotary_cos_sin is not None - ), "rotary_inv_freq and embed_positions_for_gpt_attention must be provided." - if self.position_embedding_type == PositionEmbeddingType.long_rope: - assert long_rope_rotary_inv_freq is not None - assert long_rope_rotary_cos_sin is not None - - context, past_key_value = gpt_attention( - qkv=qkv, - past_key_value=past_key_value, - attention_mask=attention_mask, - attention_packed_mask=attention_packed_mask, - sequence_length=attention_params.sequence_length, - host_past_key_value_lengths=kv_cache_params. - host_past_key_value_lengths, - host_max_attention_window_sizes=kv_cache_params. - host_max_attention_window_sizes, - host_sink_token_length=kv_cache_params.host_sink_token_length, - context_lengths=attention_params.context_lengths, - cache_indirection=kv_cache_params.cache_indirection, - host_request_types=attention_params.host_request_types, - layer_idx=self.local_layer_idx, - num_heads=self.num_attention_heads, - num_kv_heads=self.num_attention_kv_heads, - num_kv_heads_origin=self.num_kv_heads, - hidden_size_per_head=self.attention_head_size, - q_scaling=self.q_scaling, - rotary_embedding_dim=self.rotary_embedding_dim, - rotary_embedding_base=self.rotary_embedding_base - if not self.is_local else self.rotary_embedding_base_local, - rotary_embedding_scale_type=self.rotary_embedding_scale_type - if not self.is_local else RotaryScalingType.none, - rotary_embedding_short_m_scale=attention_params.short_mscale, - rotary_embedding_long_m_scale=attention_params.long_mscale, - rotary_embedding_scale=self.rotary_embedding_scale - if not self.is_local else 1.0, - rotary_embedding_max_positions=self.max_position_embeddings, - rotary_embedding_original_max_positions=self. - original_max_position_embeddings, - position_embedding_type=self.position_embedding_type, - rotary_inv_freq=rotary_inv_freq - if not self.is_local else rotary_inv_freq_local, - rotary_cos_sin=rotary_cos_sin - if not self.is_local else rotary_cos_sin_local, - kv_orig_quant_scale=kv_orig_quant_scale, - kv_quant_orig_scale=kv_quant_orig_scale, - attention_output_orig_quant_scale=self. - _get_output_orig_quant_scale(), - attention_output_sf_scale=attention_output_sf_scale, - kv_cache_quant_mode=self.quant_mode, - max_context_length=attention_params.max_context_length, - mask_type=self.attention_mask_type, - block_sparse_block_size=self.block_sparse_params.block_size, - block_sparse_homo_head_pattern=self.block_sparse_params. - homo_head_pattern, - block_sparse_num_local_blocks=self.block_sparse_params. - num_local_blocks, - block_sparse_vertical_stride=self.block_sparse_params. - vertical_stride, - alibi_slopes=alibi_slopes, - tp_size=self.tp_size, - tp_rank=self.tp_rank, - kv_cache_block_offsets=kv_cache_params.kv_cache_block_offsets - if not self.cross_attention else - kv_cache_params.cross_kv_cache_block_offsets, - host_kv_cache_block_offsets=kv_cache_params. - host_kv_cache_block_offsets if not self.cross_attention else - kv_cache_params.host_cross_kv_cache_block_offsets, - host_kv_cache_pool_pointers=kv_cache_params. - host_kv_cache_pool_pointers if not self.cross_attention else - kv_cache_params.host_cross_kv_cache_pool_pointers, - host_kv_cache_pool_mapping=kv_cache_params. - host_kv_cache_pool_mapping if not self.cross_attention else - kv_cache_params.host_cross_kv_cache_pool_mapping, - do_cross_attention=self.cross_attention, - cross_kv=cross_kv, - cross_kv_length=attention_params.encoder_max_input_length, - encoder_input_lengths=attention_params.encoder_input_lengths, - logn_scaling=logn_scaling, - relative_attention_bias=self.rel_attn_table.value - if self.relative_attention else None, - max_distance=self.max_distance, - host_context_lengths=attention_params.host_context_lengths, - use_cache=use_cache, - spec_decoding_is_generation_length_variable=spec_decoding_params - .spec_decoding_is_generation_length_variable, - spec_decoding_max_generation_length=spec_decoding_params. - spec_decoding_max_generation_length, - spec_decoding_generation_lengths=spec_decoding_params. - spec_decoding_generation_lengths, - spec_decoding_position_offsets=spec_decoding_params. - spec_decoding_position_offsets, - spec_decoding_packed_mask=spec_decoding_params. - spec_decoding_packed_mask, - spec_decoding_use=spec_decoding_params.spec_decoding_use, - long_rope_rotary_inv_freq=long_rope_rotary_inv_freq, - long_rope_rotary_cos_sin=long_rope_rotary_cos_sin, - mrope_rotary_cos_sin=mrope_params.mrope_rotary_cos_sin, - mrope_position_deltas=mrope_params.mrope_position_deltas, - attn_logit_softcapping_scale=self.max_attn_value, - host_runtime_perf_knobs=attention_params. - host_runtime_perf_knobs, - host_context_progress=attention_params.host_context_progress, - skip_attn=skip_attn, - cp_size=self.cp_size, - cp_rank=self.cp_rank, - cp_group=self.cp_group) - - else: - # plain TensorRT mode - assert paged_kv_cache == False - - assert logn_scaling is None, "plan TensorRT mode does not support logn scaling now" - - def transpose_for_scores(x, - rotary: bool = False, - is_kv: bool = False): - _num_attention_heads = self.num_attention_kv_heads if is_kv else self.num_attention_heads - new_x_shape = concat([ - shape(x, 0), - shape(x, 1), _num_attention_heads, self.attention_head_size - ]) - if rotary: - return x.view(new_x_shape) - else: - return x.view(new_x_shape).permute([0, 2, 1, 3]) - - # qkv after projection is of shape - # [bs, seqlen, (num_attention_heads + 2 * num_attention_kv_heads), attention_head_size]. - # The projected and split qkv after transpose_for_scores(): - # Q[bs, num_attention_heads, seqlen, attention_head_size] - # K[bs, num_attention_kv_heads, seqlen, attention_head_size] - # V[bs, num_attention_kv_heads, seqlen, attention_head_size] - kv_size = self.attention_head_size * self.num_attention_kv_heads - if unfuse_qkv_gemm: - query, key, value = qkv[0], qkv[1], qkv[2] - else: - query, key, value = split( - qkv, [self.attention_hidden_size, kv_size, kv_size], dim=2) - - # in cross attention mode, replace kv by encoder_output - if self.cross_attention and encoder_output is not None: - key, value = split(cross_kv, [kv_size, kv_size], dim=2) - - query = transpose_for_scores( - query, rotary=self.position_embedding_type.is_rope()) - key = transpose_for_scores( - key, is_kv=True, rotary=self.position_embedding_type.is_rope()) - value = transpose_for_scores(value, is_kv=True) - - if self.position_embedding_type.is_rope(): - if self.position_embedding_type == PositionEmbeddingType.long_rope: - sequence_length = shape(hidden_states, 1) - floor_seq_length = maximum( - sequence_length, self.original_max_position_embeddings) - - starts = concat([0, 0, 0]) - shapes = concat( - [1, floor_seq_length, self.rotary_embedding_dim]) - short = slice(attention_params.embed_positions, starts, - shapes) - long = slice(attention_params.long_rope_embed_positions, - starts, shapes) - - embed_positions = concat([short, long], dim=0) - select = where( - sequence_length - <= self.original_max_position_embeddings, 0, 1) - embed_positions = slice(embed_positions, - concat([select, 0, 0]), - sizes=shape(short)) - embed_positions = cast(embed_positions, self.dtype) - elif is_same_dtype(self.dtype, trt.bfloat16): - embed_positions = cast(attention_params.embed_positions, - trt.bfloat16) - else: - embed_positions = cast(attention_params.embed_positions, - query.dtype) - - if self.rotary_embedding_dim is not None: - # When shape(hidden_states, 1) > 1(Context phase), the embedding start from 0, - # otherwise (Generation phase) move start to position - if not use_cache: - # Only context phase is involved when kv cache is disabled. - start = 0 - else: - start = where( - shape(hidden_states, 1) > 1, 0, - shape(past_key_value, 3)) - size = where( - shape(hidden_states, 1) > 1, shape(hidden_states, 1), 1) - sincos = slice(embed_positions, concat([0, start, 0]), - concat([1, size, self.rotary_embedding_dim])) - sin, cos = split(sincos, - self.rotary_embedding_dim // 2, - dim=-1) - - key_rot_size = concat([ - shape(key, 0), - shape(key, 1), - shape(key, 2), self.rotary_embedding_dim - ]) - query_rot_size = concat([ - shape(query, 0), - shape(query, 1), - shape(query, 2), self.rotary_embedding_dim - ]) - remaining = shape(key, 3) - self.rotary_embedding_dim - key_pass_size = concat([ - shape(key, 0), - shape(key, 1), - shape(key, 2), remaining - ]) - query_pass_size = concat([ - shape(query, 0), - shape(query, 1), - shape(query, 2), remaining - ]) - k_rot = slice(key, [0, 0, 0, 0], key_rot_size) - k_pass = slice(key, [0, 0, 0, self.rotary_embedding_dim], - key_pass_size) - - q_rot = slice(query, [0, 0, 0, 0], query_rot_size) - q_pass = slice(query, [0, 0, 0, self.rotary_embedding_dim], - query_pass_size) - - k_rot = RopeEmbeddingUtils.apply_rotary_pos_emb( - k_rot, [cos, sin], self.position_embedding_type) - q_rot = RopeEmbeddingUtils.apply_rotary_pos_emb( - q_rot, [cos, sin], self.position_embedding_type) - - key = concat([k_rot, k_pass], dim=3) - query = concat([q_rot, q_pass], dim=3) - else: - key = RopeEmbeddingUtils.apply_rotary_pos_emb( - key, [cos, sin], self.position_embedding_type) - query = RopeEmbeddingUtils.apply_rotary_pos_emb( - query, [cos, sin], self.position_embedding_type) - - key = key.permute([0, 2, 1, 3]) - query = query.permute([0, 2, 1, 3]) - - if past_key_value is not None and not self.cross_attention: - if self.kv_cache_scaling_factor is not None: - past_key_value = dequantize( - past_key_value, - self.kv_cache_scaling_factor.value, - output_type=self.dtype) - - # past_key_value [bs, 2, num_heads, max_seq_len, head_dim] - past_key, past_value = split(past_key_value, 1, dim=1) - - key_shape = concat([ - shape(past_key, 0), - shape(past_key, 2), - shape(past_key, 3), - shape(past_key, 4) - ]) - past_key = past_key.view(key_shape, zero_is_placeholder=False) - past_value = past_value.view(key_shape, - zero_is_placeholder=False) - - key = concat([past_key, key], dim=2) - value = concat([past_value, value], dim=2) - - if use_cache: - key_inflated_shape = concat([ - shape(key, 0), 1, - shape(key, 1), - shape(key, 2), - shape(key, 3) - ]) - inflated_key = key.view(key_inflated_shape, - zero_is_placeholder=False) - inflated_value = value.view(key_inflated_shape, - zero_is_placeholder=False) - past_key_value = concat([inflated_key, inflated_value], dim=1) - - # TRT quantizes the tensor value by doing `cast(clip(fp_value / scale))` while - # the plugin quantizes it by doing `cast(clip(fp_value * scale))`. - if self.kv_cache_scaling_factor is not None: - past_key_value = quantize( - past_key_value, - self.kv_cache_scaling_factor.value, - dtype='fp8' - if self.quant_mode.has_fp8_kv_cache() else 'int8') - - # MQA broadcast - if self.num_attention_heads // self.num_attention_kv_heads > 1: - key = repeat_interleave( - key, - self.num_attention_heads // self.num_attention_kv_heads, 1) - value = repeat_interleave( - value, - self.num_attention_heads // self.num_attention_kv_heads, 1) - - key_length = shape(key, 2) - - # The following code creates a 2D tensor with 0s in the lower triangular (including the diagonal) and - # +INF in the upper triangular parts. This bias tensor will be added to the output of the Q*K^T matrix - # multiplication (BMM1). The +INF elements will be transformed to 0s by the Softmax operator that - # follows. The elements that corresponds to 0s in the bias are unaffected by the bias tensor. - # - # Note that when we added to another bias tensor B (for example, with AliBi), the values in the lower- - # triangular part of the B tensor are not affected and the upper-triangular ones are set to +INF. - if self.attention_mask_type == AttentionMaskType.causal and not self.cross_attention: - if self.position_embedding_type.is_alibi(): - query_length = shape(query, 2) - # bsz, tatget_length, past_key_value_length - buffer = make_causal_mask(shape(query, 0), query_length, - key_length - query_length, - trt.float32) - starts = concat([0, 0, 0, 0]) - sizes = concat([1, 1, query_length, key_length]) - generated_mask = slice(buffer, starts, sizes) - - else: - query_length = shape(query, 2) - starts = concat([0, 0, key_length - query_length, 0]) - sizes = concat([1, 1, query_length, key_length]) - if self.position_embedding_type == PositionEmbeddingType.long_rope: - buf_shape = (self.original_max_position_embeddings, - self.original_max_position_embeddings) - else: - buf_shape = (self.max_position_embeddings, - self.max_position_embeddings) - select_buf = np.expand_dims( - np.tril(np.ones(buf_shape)).astype(bool), (0, 1)) - - select_buf = np.logical_not(select_buf) - mask_buf = np.zeros_like(select_buf, np.float32) - mask_buf[select_buf] = float('-inf') - buffer = constant(mask_buf) - generated_mask = slice(buffer, starts, sizes) - - elif self.attention_mask_type == AttentionMaskType.bidirectional and not self.cross_attention: - query_length = shape(query, 2) - zero_buf = np.expand_dims( - np.zeros((self.max_position_embeddings, - self.max_position_embeddings), - dtype=np.float32), (0, 1)) - - zero_buf[:, :, :-1, -1] = 1 - zero_buf *= -10000 - - mask = constant(zero_buf) - - # context phase, query_length - mask_size = where(query_length > 1, query_length, 1) - mask_start = where(query_length > 1, - self.max_position_embeddings - mask_size, 1) - start = concat([0, 0, mask_start, mask_start]) - size = concat([1, 1, mask_size, mask_size]) - generated_mask = slice(mask, start, size) - - if attention_mask is not None: - if self.cross_attention: - batch_size = shape(attention_mask, 0) - query_len = shape(attention_mask, 1) - encoder_input_len = shape(attention_mask, 2) - attention_mask = attention_mask.view( - concat([batch_size, 1, query_len, encoder_input_len])) - attention_mask = where(attention_mask == 0, float('-inf'), - 0.0) - else: - attention_mask = expand_mask(attention_mask, - shape(query, 2)) - bias = attention_mask - if self.position_embedding_type.is_alibi(): - alibi_biases = generate_alibi_biases(alibi_slopes, key_length) - bias = alibi_biases if bias is None else bias + alibi_biases - - if self.relative_attention: - query_length = shape(query, 2) - if self.use_implicit_relative_attention: - relative_bias = compute_relative_bias( - query_length + key_length - 1, - key_length, - self.num_buckets, - self.max_distance, - False, # bidirectional - self.rel_attn_table.value.transpose(1, 0), - tp_size=self.tp_size, - tp_group=self.tp_group, - tp_rank=self.tp_rank) - else: - relative_bias = unsqueeze(self.rel_attn_table.value, 0) - start = concat([0, 0, query_length + key_length - 2, 0]) - size = concat([ - shape(relative_bias, 0), - shape(relative_bias, 1), 1, key_length - ]) - relative_bias = slice(relative_bias, start, size) - - key = key.permute([0, 1, 3, 2]) - model_type = query.dtype - with precision('float32'): - # FIXME the "with precision('float32') does not really work and lead to nan" - # in some cases - query = cast(query, 'float32') - key = cast(key, 'float32') - if norm_before_bmm1: - # Apply norm on query earlier to prevent matmul fp16 overflow. - query /= (self.q_scaling * self.norm_factor) - attention_scores = matmul(query, key) - if not norm_before_bmm1: - attention_scores = attention_scores / (self.q_scaling * - self.norm_factor) - if self.max_attn_value > 0: - attention_scores = self.max_attn_value * ACT2FN['tanh']( - attention_scores / self.max_attn_value) - - if self.attention_mask_type in [ - AttentionMaskType.causal, - AttentionMaskType.bidirectional - ] and not self.cross_attention: - - bias = generated_mask if bias is None else bias + generated_mask - - if bias is not None: - bias = cast(bias, attention_scores.dtype) - attention_scores = attention_scores + bias - - if self.relative_attention: - attention_scores = attention_scores + relative_bias - - attention_probs = softmax(attention_scores, dim=-1) - attention_probs = cast(attention_probs, model_type) - - # A dummy reshape WAR for mha fusion - attention_probs = attention_probs.view( - concat([ - shape(attention_probs, 0), - shape(attention_probs, 1), - shape(attention_probs, 2), - shape(value, 2) - ])) - context = matmul(attention_probs, value, - use_fp32_acc=False).permute([0, 2, 1, 3]) - context = context.view( - concat([ - shape(context, 0), - shape(context, 1), self.attention_hidden_size - ])) - - dense_lora_params = None - if lora_layer_params is not None: - dense_lora_params = lora_layer_params.get_runtime_params( - 0, "attn_dense") - - if skip_attn is not None and not default_net( - ).plugin_config.use_fp8_context_fmha: - # This case is used when we can skip this attention layer directly. - # The output would be undefined and not used if skip_attn is not None - # and set skip_attn as True during runtime - # But when use_fp8_context_fmha is enabled, the output data type of - # attention_plugin is fp8. Since TRT's conditional layer does not support - # FP8 data type yet, we cannot use it to skip the computation in such case. - - dense_conditional = Conditional(skip_attn) - skip_case = dense_conditional.add_input(attention_input) - context = dense_conditional.add_input(context) - - if self.inner_layernorm is not None: - context = self.inner_layernorm(context) - context = self.dense(context, - lora_runtime_params=dense_lora_params, - all_reduce_params=all_reduce_params) - - if skip_attn is not None and not default_net( - ).plugin_config.use_fp8_context_fmha: - context = dense_conditional.add_output(skip_case, context) - - if use_cache: - return (context, past_key_value) - else: - return context - - def set_rel_attn_table(self, max_seq_len, precomputed_relative_attention): - self.rel_attn_table = Parameter(shape=(self.num_attention_heads, - max_seq_len + 1, - max_seq_len + 1), - dtype=self.dtype) - self.rel_attn_table.value = precomputed_relative_attention - - def postprocess(self, tllm_key, weights, **kwargs): - - if tllm_key.endswith("kv_cache_scaling_factor"): - if weights is None: - return {tllm_key: torch.ones(1, ).float()} - elif isinstance(weights, torch.Tensor): - return {tllm_key: weights.float()} - elif None in weights: - return {tllm_key: torch.ones(1, ).float()} - else: - return {tllm_key: max(weights).float()} - elif tllm_key.endswith("kv_cache_rcp_scaling_factor"): - if weights is None: - return {tllm_key: torch.ones(1, ).float()} - elif isinstance(weights, torch.Tensor): - return {tllm_key: torch.reciprocal(weights.float())} - elif None in weights: - return {tllm_key: torch.ones(1, ).float()} - else: - return {tllm_key: torch.reciprocal(max(weights).float())} - else: - return {tllm_key: weights} - - -class BertAttention(Module): - - def __init__(self, - hidden_size, - num_attention_heads, - max_position_embeddings=1024, - num_layers=1, - attention_head_size=None, - num_kv_heads=None, - q_scaling=1.0, - apply_query_key_layer_scaling=False, - bias=True, - dtype=None, - tp_group=None, - tp_size=1, - tp_rank=0, - cp_group=None, - cp_size=1, - cp_rank=0, - relative_attention=False, - max_distance=0, - num_buckets=0, - quant_mode=QuantMode(0)): - super().__init__() - - self.attention_head_size = hidden_size // num_attention_heads if attention_head_size is None else attention_head_size - self.num_attention_heads = num_attention_heads // tp_size - self.num_attention_kv_heads = ( - num_kv_heads + tp_size - 1 - ) // tp_size if num_kv_heads is not None else self.num_attention_heads - self.hidden_size = hidden_size - self.attention_hidden_size = self.attention_head_size * self.num_attention_heads - self.max_position_embeddings = max_position_embeddings - self.norm_factor = math.sqrt(self.attention_head_size) - self.tp_group = tp_group - self.tp_size = tp_size - self.tp_rank = tp_rank - self.cp_group = cp_group - self.cp_size = cp_size - self.cp_rank = cp_rank - - self.num_layers = num_layers - self.apply_query_key_layer_scaling = apply_query_key_layer_scaling - self.norm_factor = math.sqrt(self.attention_head_size) - self.q_scaling = q_scaling - if self.apply_query_key_layer_scaling: - self.norm_factor *= self.num_layers - self.q_scaling *= self.num_layers - - self.dtype = dtype - # add quant mode to control quantization - self.quant_mode = quant_mode - - self.relative_attention = relative_attention - self.max_distance = max_distance - self.num_buckets = num_buckets - - # out dim is not necessarily hidden_size + kv specific size (in MQA/GQA), but num_heads * heads_size - # example: d_model != num_heads * head_size in Flan-T5 - self.qkv = ColumnLinear(hidden_size, - tp_size * self.attention_hidden_size + - (2 * tp_size * self.num_attention_kv_heads * - self.attention_head_size), - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - gather_output=False, - is_qkv=True) - self.dense = RowLinear(tp_size * self.num_attention_heads * - self.attention_head_size, - hidden_size, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size) - - # see optimize_model's add_lora for LoRA initialization - self.qkv_lora = None - - # per-layer relative attention table - if relative_attention: - self.rel_attn_table = Parameter(shape=(num_attention_heads // - tp_size, num_buckets), - dtype=dtype) - - def forward(self, - hidden_states: Tensor, - attention_mask=None, - input_lengths=None, - max_input_length=None, - lora_layer_params=None): - assert isinstance(hidden_states, Tensor) - - qkv_lora_params = None - if lora_layer_params is not None: - qkv_lora_params = lora_layer_params.get_runtime_params( - 0, "attn_qkv") - - qkv = self.qkv(hidden_states, qkv_lora_params) - - if default_net().plugin_config.remove_input_padding: - assert qkv.ndim() == 2 - - if default_net( - ).plugin_config.lora_plugin and qkv_lora_params is None and lora_layer_params is not None: - q_lora_params = lora_layer_params.get_runtime_params(0, "attn_q") - k_lora_params = lora_layer_params.get_runtime_params(0, "attn_k") - v_lora_params = lora_layer_params.get_runtime_params(0, "attn_v") - - assert (q_lora_params is not None and k_lora_params is not None and v_lora_params is not None) or \ - (q_lora_params is None and k_lora_params is None and v_lora_params is None), "q_lora_params, k_lora_params and v_lora_params should be all enabled or all disabled at the same time." - - if q_lora_params is not None and k_lora_params is not None and v_lora_params is not None: - qkv_lora_params = LoraRuntimeParams( - lora_ranks=[ - q_lora_params.lora_ranks[0], - k_lora_params.lora_ranks[0], - v_lora_params.lora_ranks[0], - ], - lora_weights_pointers=[ - q_lora_params.lora_weights_pointers[0], - k_lora_params.lora_weights_pointers[0], - v_lora_params.lora_weights_pointers[0], - ], - host_request_types=q_lora_params.host_request_types, - host_context_lengths=q_lora_params.host_context_lengths) - - q_lora, k_lora, v_lora = self.qkv_lora(hidden_states, - qkv_lora_params) - qkv_lora = concat([q_lora, k_lora, v_lora], - dim=q_lora.rank() - 1) - qkv = qkv + qkv_lora - - if default_net().plugin_config.bert_attention_plugin: - # TRT plugin mode - assert input_lengths is not None - context = bert_attention( - qkv, - input_lengths, - self.num_attention_heads, - self.attention_head_size, - q_scaling=self.q_scaling, - relative_attention=self.relative_attention, - max_distance=self.max_distance, - relative_attention_bias=self.rel_attn_table.value - if self.relative_attention else None, - max_input_length=max_input_length, - cp_group=self.cp_group, - cp_size=self.cp_size, - cp_rank=self.cp_rank) - else: - # plain TRT mode - def transpose_for_scores(x): - new_x_shape = concat([ - shape(x, 0), - shape(x, 1), self.num_attention_heads, - self.attention_head_size - ]) - return x.view(new_x_shape).permute([0, 2, 1, 3]) - - kv_size = self.attention_head_size * self.num_attention_kv_heads - query, key, value = split( - qkv, [self.attention_hidden_size, kv_size, kv_size], dim=2) - if self.cp_size > 1 and self.cp_group is not None: - key = allgather(key, self.cp_group, gather_dim=1) - value = allgather(value, self.cp_group, gather_dim=1) - query = transpose_for_scores(query) - key = transpose_for_scores(key) - value = transpose_for_scores(value) - - key = key.permute([0, 1, 3, 2]) - attention_scores = matmul(query, key, use_fp32_acc=False) - attention_scores = attention_scores / (self.q_scaling * - self.norm_factor) - - if self.relative_attention: - query_len = shape(attention_scores, 2) - key_len = shape(attention_scores, 3) - bias = compute_relative_bias( - query_len, - key_len, - self.num_buckets, - self.max_distance, - True, # bidirectional - self.rel_attn_table.value.transpose(1, 0), - tp_size=self.tp_size, - tp_group=self.tp_group, - tp_rank=self.tp_rank) - attention_scores = attention_scores + bias - - if attention_mask is not None: - attention_mask = expand_mask(attention_mask, shape(query, 2)) - attention_mask = cast(attention_mask, attention_scores.dtype) - attention_scores = attention_scores + attention_mask - - attention_probs = softmax(attention_scores, dim=-1) - - context = matmul(attention_probs, value, - use_fp32_acc=False).permute([0, 2, 1, 3]) - context = context.view( - concat([ - shape(context, 0), - shape(context, 1), self.attention_hidden_size - ])) - - dense_lora_params = None - if lora_layer_params is not None: - dense_lora_params = lora_layer_params.get_runtime_params( - 0, "attn_dense") - context = self.dense(context, lora_runtime_params=dense_lora_params) - - return context - - -class CogVLMAttention(Attention): - - def __init__( - self, - *, - local_layer_idx, - hidden_size, - num_attention_heads, - num_kv_heads=None, - max_position_embeddings=1024, - attention_mask_type=AttentionMaskType.causal, - bias=True, - dtype=None, - position_embedding_type=PositionEmbeddingType.learned_absolute, - rotary_embedding_base=10000.0, - rotary_embedding_scaling=None, - tp_group=None, - tp_size=1, - tp_rank=0, - quant_mode: QuantMode = QuantMode(0), - dense_bias=None, - ): - super().__init__(local_layer_idx=local_layer_idx, - hidden_size=hidden_size, - num_attention_heads=num_attention_heads, - num_kv_heads=num_kv_heads, - max_position_embeddings=max_position_embeddings, - dtype=dtype, - attention_mask_type=attention_mask_type, - bias=bias, - position_embedding_type=position_embedding_type, - rotary_embedding_base=rotary_embedding_base, - rotary_embedding_scaling=rotary_embedding_scaling, - tp_group=tp_group, - tp_size=tp_size, - tp_rank=tp_rank, - quant_mode=quant_mode) - - self.vis_qkv = ColumnLinear( - hidden_size, - tp_size * self.num_attention_heads * self.attention_head_size + - (2 * tp_size * self.num_attention_kv_heads * - self.attention_head_size), - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - gather_output=False, - is_qkv=True) - self.vis_dense = RowLinear(tp_size * self.num_attention_heads * - self.attention_head_size, - hidden_size, - bias=self.dense_bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size) - - def forward(self, - hidden_states: Tensor, - use_cache=False, - kv_cache_params=None, - attention_params=None, - vision_token_mask=None, - position_embedding=None): - - assert isinstance(hidden_states, Tensor) - assert (default_net().plugin_config.gpt_attention_plugin) - - vision_qkv = self.vis_qkv(hidden_states) - language_qkv = self.qkv(hidden_states) - qkv = where(vision_token_mask, vision_qkv, language_qkv) - - qkv = RopeEmbeddingUtils.apply_rotary_pos_emb_cogvlm( - qkv, position_embedding, self.num_attention_heads, - self.attention_head_size, self.max_position_embeddings, - self.rotary_embedding_scale, - default_net().plugin_config.remove_input_padding) - - assert attention_params is None or attention_params.is_valid( - default_net().plugin_config.gpt_attention_plugin, - default_net().plugin_config.remove_input_padding, use_cache) - assert kv_cache_params is None or kv_cache_params.is_valid( - default_net().plugin_config.gpt_attention_plugin) - - past_key_value = None if kv_cache_params is None else kv_cache_params.get_first_past_key_value( - ) - - if default_net().plugin_config.gpt_attention_plugin: - if self.cross_attention and (past_key_value is not None): - past_key_value = kv_cache_params.past_key_value[1] - assert self.attention_mask_type in [ - AttentionMaskType.causal, AttentionMaskType.bidirectional, - AttentionMaskType.bidirectionalglm - ], 'Plugin only support masked MHA.' - - # KV cache scales. - kv_orig_quant_scale = self.kv_cache_rcp_scaling_factor.value if self.quant_mode.has_kv_cache_quant( - ) else None - kv_quant_orig_scale = self.kv_cache_scaling_factor.value if self.quant_mode.has_kv_cache_quant( - ) else None - - context, past_key_value = gpt_attention( - qkv=qkv, - past_key_value=past_key_value, - sequence_length=attention_params.sequence_length, - host_past_key_value_lengths=kv_cache_params. - host_past_key_value_lengths, - host_max_attention_window_sizes=kv_cache_params. - host_max_attention_window_sizes, - host_sink_token_length=kv_cache_params.host_sink_token_length, - context_lengths=attention_params.context_lengths, - cache_indirection=kv_cache_params.cache_indirection, - host_request_types=attention_params.host_request_types, - layer_idx=self.local_layer_idx, - num_heads=self.num_attention_heads, - num_kv_heads=self.num_attention_kv_heads, - num_kv_heads_origin=self.num_kv_heads, - hidden_size_per_head=self.attention_head_size, - q_scaling=self.q_scaling, - position_embedding_type=self.position_embedding_type, - kv_orig_quant_scale=kv_orig_quant_scale, - kv_quant_orig_scale=kv_quant_orig_scale, - attention_output_orig_quant_scale=self. - _get_output_orig_quant_scale(), - kv_cache_quant_mode=self.quant_mode, - max_context_length=attention_params.max_context_length, - mask_type=self.attention_mask_type, - alibi_slopes=None, - tp_size=self.tp_size, - tp_rank=self.tp_rank, - kv_cache_block_offsets=kv_cache_params.kv_cache_block_offsets, - host_kv_cache_block_offsets=kv_cache_params. - host_kv_cache_block_offsets, - host_kv_cache_pool_pointers=kv_cache_params. - host_kv_cache_pool_pointers, - host_kv_cache_pool_mapping=kv_cache_params. - host_kv_cache_pool_mapping, - do_cross_attention=self.cross_attention, - cross_kv=None, - cross_kv_length=attention_params.encoder_max_input_length, - encoder_input_lengths=attention_params.encoder_input_lengths, - relative_attention_bias=self.rel_attn_table.value - if self.relative_attention else None, - max_distance=self.max_distance, - host_context_lengths=attention_params.host_context_lengths, - use_cache=use_cache, - spec_decoding_position_offsets=None, - spec_decoding_packed_mask=None, - mrope_rotary_cos_sin=None, - mrope_position_deltas=None, - host_runtime_perf_knobs=attention_params. - host_runtime_perf_knobs, - host_context_progress=attention_params.host_context_progress, - ) - - vision_dense = self.vis_dense(context) - language_dense = self.dense(context) - context = where(vision_token_mask, vision_dense, language_dense) - - if use_cache: - return (context, past_key_value) - else: - return context - - -class DeepseekV2Attention(Attention): - - def __init__( - self, - *, - local_layer_idx, - hidden_size, - num_attention_heads, - q_lora_rank, - kv_lora_rank, - qk_nope_head_dim=None, - qk_rope_head_dim=None, - v_head_dim=None, - eps=1e-06, - attention_mask_type=AttentionMaskType.causal, - dtype=None, - position_embedding_type=PositionEmbeddingType.learned_absolute, - max_position_embeddings=1024, - rotary_embedding_base=10000.0, - rotary_embedding_scaling=None, - rotary_embedding_beta_fast=32, - rotary_embedding_beta_slow=1, - rotary_embedding_mscale=1, - rotary_embedding_mscale_all_dim=0, - rotary_embedding_origin_max_position=4096, - rotary_scaling=None, - tp_group=None, - tp_size=1, - tp_rank=0, - quant_mode: QuantMode = QuantMode(0), - ): - super().__init__(local_layer_idx=local_layer_idx, - hidden_size=hidden_size, - num_attention_heads=num_attention_heads, - num_kv_heads=1, - max_position_embeddings=max_position_embeddings, - attention_head_size=kv_lora_rank + qk_rope_head_dim, - dtype=dtype, - attention_mask_type=attention_mask_type, - position_embedding_type=position_embedding_type, - rotary_embedding_base=rotary_embedding_base, - rotary_embedding_scaling=rotary_embedding_scaling, - tp_group=tp_group, - tp_size=tp_size, - tp_rank=tp_rank, - quant_mode=quant_mode, - bias=False, - dense_bias=False, - enable_qkv=False) - - self.tp_size = tp_size - - if q_lora_rank is None: - self.q_lora_rank = hidden_size - self.is_deepseek_v2_lite = True - else: - self.q_lora_rank = q_lora_rank - self.is_deepseek_v2_lite = False - - self.kv_lora_rank = kv_lora_rank - self.qk_nope_head_dim = qk_nope_head_dim - self.qk_rope_head_dim = qk_rope_head_dim - self.v_head_dim = v_head_dim - self.rotary_embedding_dim = 0 - self.rotary_scaling = rotary_scaling - self.shard_dim = 1 - - def yarn_get_mscale(scale=1, mscale=1): - if scale <= 1: - return 1.0 - return 0.1 * mscale * math.log(scale) + 1.0 - - assert self.rotary_scaling is not None - if self.rotary_scaling is not None: - mscale_all_dim = self.rotary_scaling.get("mscale_all_dim", 0) - scaling_factor = self.rotary_scaling["factor"] - if mscale_all_dim: - mscale = yarn_get_mscale(scaling_factor, mscale_all_dim) - self.q_scaling = 1.0 / (mscale * mscale) - - _, embed_positions_for_gpt_attention = RopeEmbeddingUtils.create_sinusoidal_positions_yarn( - self.max_position_embeddings, self.qk_rope_head_dim, - self.rotary_embedding_base, self.rotary_scaling["factor"], - rotary_embedding_origin_max_position, rotary_embedding_beta_fast, - rotary_embedding_beta_slow, rotary_embedding_mscale, - rotary_embedding_mscale_all_dim) - self.register_parameter( - 'embed_positions_for_gpt_attention', - Parameter(embed_positions_for_gpt_attention, dtype='float32')) - - self.rotary_embedding_scale_type = RotaryScalingType.none - self.rotary_embedding_scale = 1.0 - - if self.is_deepseek_v2_lite: - self.fused_a = ColumnLinear( - hidden_size, - kv_lora_rank + qk_rope_head_dim, - bias=self.dense_bias, - dtype=dtype, - ) - else: - self.fused_a = ColumnLinear( - hidden_size, - q_lora_rank + kv_lora_rank + qk_rope_head_dim, - bias=self.dense_bias, - dtype=dtype, - ) - self.q_a_layernorm = RmsNorm(q_lora_rank, dtype=dtype, eps=eps) - - self.kv_a_layernorm = RmsNorm(kv_lora_rank, dtype=dtype, eps=eps) - - self.kv_b_proj = Parameter( - shape=(self.num_attention_heads * self.qk_nope_head_dim * 2, - self.kv_lora_rank), - dtype=dtype) - self.k_b_proj_trans = Parameter( - shape=(self.num_attention_heads * self.kv_lora_rank, - self.qk_nope_head_dim), - dtype=dtype) - self.q_b_proj = Parameter( - shape=(self.num_attention_heads * - (self.qk_nope_head_dim + self.qk_rope_head_dim), - self.q_lora_rank), - dtype=dtype) - self.dense = RowLinear(tp_size * self.num_attention_heads * - self.v_head_dim, - hidden_size, - bias=self.dense_bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size) - set_obj_attrs(self.q_b_proj, { - "weight_loader": self.weight_loader, - }) - set_obj_attrs(self.kv_b_proj, { - "weight_loader": self.weight_loader, - }) - set_obj_attrs(self.k_b_proj_trans, { - "weight_loader": self.weight_loader, - }) - - def weight_loader(self, mapping: Mapping, param: Parameter, - loaded_weight: torch.Tensor): - # use_parallel_embedding - tp_rank = mapping.tp_rank - if self.tp_size > 1: - sharding_dim = self.sharding_dim - shard_size = param._shape[sharding_dim] - start_idx = tp_rank * shard_size - loaded_weight = loaded_weight.narrow(sharding_dim, start_idx, - shard_size) - param.value = loaded_weight - - def postprocess(self, tllm_key, weights, **kwargs): - - def split_matrix_tp(v, tp_size, idx, dim=0): - if tp_size == 1: - return v - if len(v.shape) == 1: - return torch.chunk(v, tp_size)[idx].contiguous() - else: - return torch.chunk(v, tp_size, dim=dim)[idx].contiguous() - - if tllm_key.find("q_b_proj") != -1: - q_b_proj_weight = weights.unflatten( - 0, - [ - self.num_attention_heads * self.tp_size, - self.qk_nope_head_dim + self.qk_rope_head_dim, - ], - ) - - q_b_proj_weight = split_matrix_tp( - q_b_proj_weight, - self.tp_size, - self.tp_rank, - dim=0, - ) - weights = q_b_proj_weight.reshape( - self.num_attention_heads * self.tp_size * - (self.qk_nope_head_dim + self.qk_rope_head_dim) // self.tp_size, - self.q_lora_rank) - - elif tllm_key.find("kv_b_proj") != -1: - kv_b_proj_weight = weights.unflatten( - 0, - [ - self.num_attention_heads * self.tp_size, - self.qk_nope_head_dim + self.v_head_dim, - ], - ) - kv_b_proj_weight = split_matrix_tp( - kv_b_proj_weight, - self.tp_size, - self.tp_rank, - dim=0, - ) - k_nope_weight, v_weight = kv_b_proj_weight.split( - [self.qk_nope_head_dim, self.v_head_dim], - dim=1, - ) - weights = torch.concat([ - k_nope_weight.reshape( - self.num_attention_heads * self.tp_size * - self.qk_nope_head_dim // self.tp_size, self.kv_lora_rank), - v_weight.reshape( - self.num_attention_heads * self.tp_size * self.v_head_dim // - self.tp_size, self.kv_lora_rank) - ], - dim=0) - - elif tllm_key.find("k_b_proj_trans") != -1: - kv_b_proj = weights.unflatten(0, [ - self.num_attention_heads * self.tp_size, - self.qk_nope_head_dim + self.v_head_dim - ]) - kv_b_proj = split(kv_b_proj, self.tp_size, self.tp_rank, dim=0) - k_nope_weight, v_weight = kv_b_proj.split( - [self.qk_nope_head_dim, self.v_head_dim], - dim=1, - ) - weights = k_nope_weight.transpose(2, 1).reshape( - self.num_attention_heads * self.kv_lora_rank, - self.qk_nope_head_dim) - - return {tllm_key: weights} - - def forward(self, - hidden_states: Tensor, - use_cache=False, - spec_decoding_params=None, - kv_cache_params=None, - attention_params=None): - assert default_net().plugin_config.remove_input_padding - - spec_decoding_params = SpecDecodingParams( - ) if spec_decoding_params is None else spec_decoding_params - - if default_net().plugin_config.remove_input_padding: - assert hidden_states.ndim() == 2 - - default_net().plugin_config.paged_kv_cache - - assert attention_params is None or attention_params.is_valid( - default_net().plugin_config.gpt_attention_plugin, - default_net().plugin_config.remove_input_padding, use_cache) - - if use_cache: - assert kv_cache_params is None or kv_cache_params.is_valid( - default_net().plugin_config.gpt_attention_plugin) - - past_key_value = None if kv_cache_params is None else kv_cache_params.get_first_past_key_value( - ) - - if self.is_deepseek_v2_lite: - compressed_kv, k_pe = self.fused_a(hidden_states).split( - [self.kv_lora_rank, self.qk_rope_head_dim], -1) - compressed_kv = self.kv_a_layernorm(compressed_kv) - input_qkv = concat([hidden_states, compressed_kv, k_pe], dim=-1) - else: - compressed_q, compressed_kv, k_pe = self.fused_a( - hidden_states).split([ - self.q_lora_rank, self.kv_lora_rank, self.qk_rope_head_dim - ], -1) - compressed_q = self.q_a_layernorm(compressed_q) - compressed_kv = self.kv_a_layernorm(compressed_kv) - input_qkv = concat([compressed_q, compressed_kv, k_pe], dim=-1) - - if default_net().plugin_config.gpt_attention_plugin: - if self.cross_attention and (past_key_value is not None): - past_key_value = kv_cache_params.past_key_value[1] - assert self.attention_mask_type in [ - AttentionMaskType.causal, - AttentionMaskType.bidirectional, - AttentionMaskType.bidirectionalglm, - ], 'Plugin only support masked MHA.' - - # KV cache scales. - if self.kv_cache_scaling_factor is not None: - kv_orig_quant_scale = self.kv_cache_rcp_scaling_factor.value - kv_quant_orig_scale = self.kv_cache_scaling_factor.value - else: - kv_orig_quant_scale = None - kv_quant_orig_scale = None - - rotary_cos_sin = self.embed_positions_for_gpt_attention.value - - context, past_key_value = gpt_attention( - qkv=input_qkv, - past_key_value=past_key_value, - sequence_length=attention_params.sequence_length, - host_past_key_value_lengths=kv_cache_params. - host_past_key_value_lengths, - host_max_attention_window_sizes=kv_cache_params. - host_max_attention_window_sizes, - host_sink_token_length=kv_cache_params.host_sink_token_length, - context_lengths=attention_params.context_lengths, - cache_indirection=kv_cache_params.cache_indirection, - host_request_types=attention_params.host_request_types, - layer_idx=self.local_layer_idx, - num_heads=self.num_attention_heads, - num_kv_heads=1, - num_kv_heads_origin=1, - hidden_size_per_head=self.kv_lora_rank + self.qk_rope_head_dim, - q_scaling=self.q_scaling, - position_embedding_type=self.position_embedding_type, - rotary_inv_freq=None, - rotary_cos_sin=rotary_cos_sin, - kv_orig_quant_scale=kv_orig_quant_scale, - kv_quant_orig_scale=kv_quant_orig_scale, - attention_output_orig_quant_scale=self. - _get_output_orig_quant_scale(), - kv_cache_quant_mode=self.quant_mode, - max_context_length=attention_params.max_context_length, - mask_type=self.attention_mask_type, - block_sparse_block_size=self.block_sparse_params.block_size, - block_sparse_homo_head_pattern=self.block_sparse_params. - homo_head_pattern, - block_sparse_num_local_blocks=self.block_sparse_params. - num_local_blocks, - block_sparse_vertical_stride=self.block_sparse_params. - vertical_stride, - alibi_slopes=None, - tp_size=self.tp_size, - tp_rank=self.tp_rank, - kv_cache_block_offsets=kv_cache_params.kv_cache_block_offsets - if not self.cross_attention else - kv_cache_params.cross_kv_cache_block_offsets, - host_kv_cache_block_offsets=kv_cache_params. - host_kv_cache_block_offsets if not self.cross_attention else - kv_cache_params.host_cross_kv_cache_block_offsets, - host_kv_cache_pool_pointers=kv_cache_params. - host_kv_cache_pool_pointers if not self.cross_attention else - kv_cache_params.host_cross_kv_cache_pool_pointers, - host_kv_cache_pool_mapping=kv_cache_params. - host_kv_cache_pool_mapping, - do_cross_attention=self.cross_attention, - cross_kv=None, - cross_kv_length=attention_params.encoder_max_input_length, - encoder_input_lengths=attention_params.encoder_input_lengths, - relative_attention_bias=self.rel_attn_table.value - if self.relative_attention else None, - max_distance=self.max_distance, - host_context_lengths=attention_params.host_context_lengths, - use_cache=use_cache, - spec_decoding_is_generation_length_variable=spec_decoding_params - .spec_decoding_is_generation_length_variable, - spec_decoding_max_generation_length=spec_decoding_params. - spec_decoding_max_generation_length, - spec_decoding_generation_lengths=spec_decoding_params. - spec_decoding_generation_lengths, - spec_decoding_position_offsets=spec_decoding_params. - spec_decoding_position_offsets, - spec_decoding_packed_mask=spec_decoding_params. - spec_decoding_packed_mask, - spec_decoding_use=spec_decoding_params.spec_decoding_use, - attn_logit_softcapping_scale=self.max_attn_value, - host_runtime_perf_knobs=attention_params. - host_runtime_perf_knobs, - host_context_progress=attention_params.host_context_progress, - is_mla_enabled_flag=True, - q_lora_rank=self.q_lora_rank, - kv_lora_rank=self.kv_lora_rank, - qk_nope_head_dim=self.qk_nope_head_dim, - qk_rope_head_dim=self.qk_rope_head_dim, - v_head_dim=self.v_head_dim, - fused_q_proj=self.fused_q_proj.value, - q_b_proj=self.q_b_proj.value, - kv_b_proj=self.kv_b_proj.value) - - context = self.dense(context) - - if use_cache: - return (context, past_key_value) - else: - return context - - -class DiffusersAttention(Module): - - def __init__(self, - *, - query_dim: int, - cross_attention_dim: Optional[int] = None, - heads: int = 8, - kv_heads: Optional[int] = None, - dim_head: int = 64, - dropout: float = 0.0, - bias: bool = False, - upcast_attention: bool = False, - upcast_softmax: bool = False, - cross_attention_norm: Optional[str] = None, - cross_attention_norm_num_groups: int = 32, - qk_norm: Optional[str] = None, - added_kv_proj_dim: Optional[int] = None, - added_proj_bias: Optional[bool] = True, - norm_num_groups: Optional[int] = None, - spatial_norm_dim: Optional[int] = None, - out_bias: bool = True, - scale_qk: bool = True, - only_cross_attention: bool = False, - eps: float = 1e-5, - rescale_output_factor: float = 1.0, - residual_connection: bool = False, - out_dim: int = None, - out_context_dim: int = None, - context_pre_only=None, - pre_only=False, - elementwise_affine: bool = True, - is_causal: bool = False, - attn_forward_funcname: str = 'joint_attn_forward', - mapping=Mapping(), - dtype=None): - super().__init__() - - self.cp_size = mapping.cp_size - self.cp_group = mapping.cp_group - self.tp_group = mapping.tp_group - self.tp_size = mapping.tp_size - self.tp_rank = mapping.tp_rank - self.dtype = dtype - self.attn_forward_func = getattr(self, attn_forward_funcname) - - self.inner_dim = out_dim if out_dim is not None else dim_head * heads - self.inner_kv_dim = self.inner_dim if kv_heads is None else dim_head * kv_heads - self.query_dim = query_dim - self.use_bias = bias - self.is_cross_attention = cross_attention_dim is not None - self.cross_attention_dim = cross_attention_dim if cross_attention_dim is not None else query_dim - - ## [TODO] Not supported yet. - # self.upcast_attention = upcast_attention - # self.upcast_softmax = upcast_softmax - # self.rescale_output_factor = rescale_output_factor - # self.residual_connection = residual_connection - # self.dropout = dropout - - self.fused_projections = False - self.out_dim = out_dim if out_dim is not None else query_dim - self.context_pre_only = context_pre_only - self.pre_only = pre_only - self.is_causal = is_causal - - self.scale_qk = scale_qk - self.scale = dim_head**-0.5 if self.scale_qk else 1.0 - - # Params for `Attention` Module - self.heads = out_dim // dim_head if out_dim is not None else heads - self.heads = self.heads // self.tp_size - self.dim_head = dim_head - # default attn settings - self.norm_factor = math.sqrt(dim_head) - self.q_scaling = 1.0 - self.max_distance = 0 - - self.added_kv_proj_dim = added_kv_proj_dim - self.only_cross_attention = only_cross_attention - if self.added_kv_proj_dim is None and self.only_cross_attention: - raise ValueError( - "`only_cross_attention` can only be set to True if `added_kv_proj_dim` is not None. Make sure to set either `only_cross_attention=False` or define `added_kv_proj_dim`." - ) - - if norm_num_groups is not None: - self.group_norm = GroupNorm(num_channels=query_dim, - num_groups=norm_num_groups, - eps=eps, - affine=True, - dtype=dtype) - else: - self.group_norm = None - - if spatial_norm_dim is not None: - raise NotImplementedError("SpatialNorm is not supported yet.") - else: - self.spatial_norm = None - - if qk_norm is None: - self.norm_q = None - self.norm_k = None - elif qk_norm == "layer_norm": - self.norm_q = LayerNorm(dim_head, - eps=eps, - elementwise_affine=elementwise_affine, - dtype=dtype) - self.norm_k = LayerNorm(dim_head, - eps=eps, - elementwise_affine=elementwise_affine, - dtype=dtype) - elif qk_norm == "fp32_layer_norm": - self.norm_q = LayerNorm(dim_head, - eps=eps, - elementwise_affine=False, - bias=False, - dtype=dtype) - self.norm_k = LayerNorm(dim_head, - eps=eps, - elementwise_affine=False, - bias=False, - dtype=dtype) - elif qk_norm == "rms_norm": - self.norm_q = RmsNorm(dim_head, eps=eps, dtype=dtype) - self.norm_k = RmsNorm(dim_head, eps=eps, dtype=dtype) - elif qk_norm == "rms_norm_across_heads": - # LTX applies qk norm across all heads - self.norm_q = RmsNorm(dim_head * heads, eps=eps, dtype=dtype) - self.norm_k = RmsNorm(dim_head * kv_heads, eps=eps, dtype=dtype) - elif qk_norm in ["layer_norm_across_heads", "l2"]: - raise NotImplementedError( - f"qk_norm {qk_norm} is not supported yet.") - else: - raise ValueError( - f"unknown qk_norm: {qk_norm}. Should be None,'layer_norm','fp32_layer_norm','rms_norm'" - ) - - if cross_attention_norm is None: - self.norm_cross = None - elif cross_attention_norm == "layer_norm": - self.norm_cross = LayerNorm(self.cross_attention_dim, dtype=dtype) - elif cross_attention_norm == "group_norm": - if self.added_kv_proj_dim is not None: - # The given `encoder_hidden_states` are initially of shape - # (batch_size, seq_len, added_kv_proj_dim) before being projected - # to (batch_size, seq_len, cross_attention_dim). The norm is applied - # before the projection, so we need to use `added_kv_proj_dim` as - # the number of channels for the group norm. - norm_cross_num_channels = added_kv_proj_dim - else: - norm_cross_num_channels = self.cross_attention_dim - self.norm_cross = GroupNorm( - num_channels=norm_cross_num_channels, - num_groups=cross_attention_norm_num_groups, - eps=1e-5, - affine=True, - dtype=dtype) - else: - raise ValueError( - f"unknown cross_attention_norm: {cross_attention_norm}. Should be None, 'layer_norm' or 'group_norm'" - ) - - # [TODO] check `gather_output` - self.to_q = ColumnLinear(query_dim, - self.inner_dim, - bias=bias, - tp_group=self.tp_group, - tp_size=self.tp_size, - gather_output=False, - dtype=dtype) - if not self.only_cross_attention: - self.to_k = ColumnLinear(self.cross_attention_dim, - self.inner_kv_dim, - bias=bias, - tp_group=self.tp_group, - tp_size=self.tp_size, - gather_output=False, - dtype=dtype) - self.to_v = ColumnLinear(self.cross_attention_dim, - self.inner_kv_dim, - bias=bias, - tp_group=self.tp_group, - tp_size=self.tp_size, - gather_output=False, - dtype=dtype) - else: - self.to_k = None - self.to_v = None - - self.added_proj_bias = added_proj_bias - if self.added_kv_proj_dim is not None: - self.add_k_proj = ColumnLinear(added_kv_proj_dim, - self.inner_kv_dim, - bias=added_proj_bias, - tp_group=self.tp_group, - tp_size=self.tp_size, - gather_output=False, - dtype=dtype) - self.add_v_proj = ColumnLinear(added_kv_proj_dim, - self.inner_kv_dim, - bias=added_proj_bias, - tp_group=self.tp_group, - tp_size=self.tp_size, - gather_output=False, - dtype=dtype) - if self.context_pre_only is not None: - self.add_q_proj = ColumnLinear(added_kv_proj_dim, - self.inner_dim, - bias=added_proj_bias, - tp_group=self.tp_group, - tp_size=self.tp_size, - gather_output=False, - dtype=dtype) - else: - self.add_q_proj = None - self.add_k_proj = None - self.add_v_proj = None - - if not self.pre_only: - self.to_out = ModuleList([ - RowLinear(self.inner_dim, - self.out_dim, - bias=out_bias, - tp_group=self.tp_group, - tp_size=self.tp_size, - dtype=dtype) - ]) - else: - self.to_out = None - - if self.context_pre_only is not None and not self.context_pre_only: - self.to_add_out = RowLinear(self.inner_dim, - self.out_dim, - bias=out_bias, - tp_group=self.tp_group, - tp_size=self.tp_size, - dtype=dtype) - else: - self.to_add_out = None - - if qk_norm is not None and added_kv_proj_dim is not None: - if qk_norm == "fp32_layer_norm": - self.norm_added_q = LayerNorm(dim_head, - elementwise_affine=False, - bias=False, - eps=eps, - dtype=dtype) - self.norm_added_k = LayerNorm(dim_head, - elementwise_affine=False, - bias=False, - eps=eps, - dtype=dtype) - elif qk_norm == "rms_norm": - self.norm_added_q = RmsNorm(dim_head, eps=eps, dtype=dtype) - self.norm_added_k = RmsNorm(dim_head, eps=eps, dtype=dtype) - else: - raise ValueError( - f"unknown qk_norm: {qk_norm}. Should be one of `None,'layer_norm','fp32_layer_norm','rms_norm'`" - ) - else: - self.norm_added_q = None - self.norm_added_k = None - - def joint_attn_forward(self, - hidden_states: Tensor, - encoder_hidden_states: Optional[Tensor] = None, - attention_mask: Optional[Tensor] = None, - max_input_length: Optional[Tensor] = None, - *args, - **kwargs): - if attention_mask is not None: - raise NotImplementedError() - residual = identity(hidden_states) - batch_size = shape(hidden_states, 0) - - # `sample` projections. - query = self.to_q(hidden_states) - key = self.to_k(hidden_states) - value = self.to_v(hidden_states) - - head_dim = self.dim_head - inner_dim = head_dim * self.heads - - query = query.view(concat([batch_size, -1, self.heads, - head_dim])).permute([0, 2, 1, 3]) - key = key.view(concat([batch_size, -1, self.heads, - head_dim])).permute([0, 2, 1, 3]) - value = value.view(concat([batch_size, -1, self.heads, - head_dim])).permute([0, 2, 1, 3]) - - if self.norm_q is not None: - query = self.norm_q(query) - if self.norm_k is not None: - key = self.norm_k(key) - - # `context` projections. - if encoder_hidden_states is not None: - encoder_hidden_states_query_proj = self.add_q_proj( - encoder_hidden_states) - encoder_hidden_states_key_proj = self.add_k_proj( - encoder_hidden_states) - encoder_hidden_states_value_proj = self.add_v_proj( - encoder_hidden_states) - - encoder_hidden_states_query_proj = encoder_hidden_states_query_proj.view( - concat([batch_size, -1, self.heads, - head_dim])).permute([0, 2, 1, 3]) - encoder_hidden_states_key_proj = encoder_hidden_states_key_proj.view( - concat([batch_size, -1, self.heads, - head_dim])).permute([0, 2, 1, 3]) - encoder_hidden_states_value_proj = encoder_hidden_states_value_proj.view( - concat([batch_size, -1, self.heads, - head_dim])).permute([0, 2, 1, 3]) - - if self.norm_added_q is not None: - encoder_hidden_states_query_proj = self.norm_added_q( - encoder_hidden_states_query_proj) - if self.norm_added_k is not None: - encoder_hidden_states_key_proj = self.norm_added_k( - encoder_hidden_states_key_proj) - - query = concat([query, encoder_hidden_states_query_proj], dim=2) - key = concat([key, encoder_hidden_states_key_proj], dim=2) - value = concat([value, encoder_hidden_states_value_proj], dim=2) - - # Transpose from [batch_size, num_heads, seq_len, head_dim] back to - # [batch_size, seq_len, num_heads * head_dim] for attention plugin. - query = query.permute([0, 2, 1, - 3]).view(concat([batch_size, -1, inner_dim])) - key = key.permute([0, 2, 1, 3]).view(concat([batch_size, -1, - inner_dim])) - value = value.permute([0, 2, 1, - 3]).view(concat([batch_size, -1, inner_dim])) - - if default_net().plugin_config.bert_attention_plugin: - # TRT plugin mode - assert self.cp_size == 1 - shape(query, 1) - qkv = concat([query, key, value], dim=-1) - input_lengths = expand( - shape(qkv, 1).unsqueeze(0), - shape(qkv, 0).unsqueeze(0)).cast("int32") - - hidden_states = bert_attention(qkv, - input_lengths, - self.heads, - head_dim, - q_scaling=self.q_scaling, - relative_attention=False, - max_distance=self.max_distance, - max_input_length=max_input_length) - else: - # plain TRT mode - def transpose_for_scores(x): - new_x_shape = concat( - [shape(x, 0), - shape(x, 1), self.heads, head_dim]) - return x.view(new_x_shape).permute([0, 2, 1, 3]) - - if self.cp_size > 1 and self.cp_group is not None: - key = allgather(key, self.cp_group, gather_dim=1) - value = allgather(value, self.cp_group, gather_dim=1) - query = transpose_for_scores(query) - key = transpose_for_scores(key) - value = transpose_for_scores(value) - - key = key.permute([0, 1, 3, 2]) - attention_scores = matmul(query, key, use_fp32_acc=True) - attention_scores = attention_scores / (self.q_scaling * - self.norm_factor) - - attention_probs = softmax(attention_scores, dim=-1) - - context = matmul(attention_probs, value, - use_fp32_acc=True).permute([0, 2, 1, 3]) - hidden_states = context.view( - concat([shape(context, 0), - shape(context, 1), inner_dim])) - - if encoder_hidden_states is not None: - # Split the attention outputs. - slice_seq_len = shape(residual, 1) - encoder_hidden_states = slice(hidden_states, - starts=concat([0, slice_seq_len, 0]), - sizes=concat([ - batch_size, - (shape(hidden_states, 1) - - slice_seq_len), inner_dim - ])) - hidden_states = slice(hidden_states, - starts=[0, 0, 0], - sizes=concat( - [batch_size, slice_seq_len, inner_dim])) - - if not self.context_pre_only: - encoder_hidden_states = self.to_add_out(encoder_hidden_states) - - # linear proj - hidden_states = self.to_out[0](hidden_states) - if encoder_hidden_states is not None: - return hidden_states, encoder_hidden_states - else: - return hidden_states - - def forward(self, - hidden_states: Tensor, - encoder_hidden_states: Optional[Tensor] = None, - attention_mask: Optional[Tensor] = None, - max_input_length: Optional[Tensor] = None, - *args, - **kwargs): - return self.attn_forward_func( - hidden_states=hidden_states, - encoder_hidden_states=encoder_hidden_states, - attention_mask=attention_mask, - max_input_length=max_input_length, - *args, - **kwargs) diff --git a/tensorrt_llm/layers/cast.py b/tensorrt_llm/layers/cast.py deleted file mode 100644 index 0dd647984aff..000000000000 --- a/tensorrt_llm/layers/cast.py +++ /dev/null @@ -1,29 +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. -from ..functional import cast -from ..module import Module - - -class Cast(Module): - - def __init__(self, output_dtype: str = 'float32') -> None: - super().__init__() - assert output_dtype in ('float32', 'float16', 'bfloat16', 'bool', - 'int32', 'int8'), TypeError( - "%s is not supported" % output_dtype) - self.output_dtype = output_dtype - - def forward(self, x): - return cast(x, self.output_dtype) diff --git a/tensorrt_llm/layers/conv.py b/tensorrt_llm/layers/conv.py deleted file mode 100644 index 49234c70962e..000000000000 --- a/tensorrt_llm/layers/conv.py +++ /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. -from typing import Tuple - -from ..functional import conv1d, conv2d, conv3d, conv_transpose2d -from ..module import Module -from ..parameter import Parameter - - -class Conv2d(Module): - - def __init__( - self, - in_channels: int, - out_channels: int, - kernel_size: Tuple[int, int], - stride: Tuple[int, int] = (1, 1), - padding: Tuple[int, int] = (0, 0), - dilation: Tuple[int, int] = (1, 1), - groups: int = 1, - bias: bool = True, - padding_mode: str = 'zeros', # TODO: refine this type - dtype=None) -> None: - super().__init__() - if groups <= 0: - raise ValueError('groups must be a positive integer') - if in_channels % groups != 0: - raise ValueError('in_channels must be divisible by groups') - if out_channels % groups != 0: - raise ValueError('out_channels must be divisible by groups') - - self.in_channels = in_channels - self.out_channels = out_channels - self.kernel_size = kernel_size - self.stride = stride - self.padding = padding - self.dilation = dilation - self.groups = groups - self.padding_mode = padding_mode - - self.weight = Parameter(shape=(out_channels, in_channels // groups, - *kernel_size), - dtype=dtype) - if bias: - self.bias = Parameter(shape=(out_channels, ), dtype=dtype) - else: - self.register_parameter('bias', None) - - def forward(self, input): - return conv2d(input, self.weight.value, - None if self.bias is None else self.bias.value, - self.stride, self.padding, self.dilation, self.groups) - - -class ConvTranspose2d(Module): - - def __init__( - self, - in_channels: int, - out_channels: int, - kernel_size: Tuple[int, int], - stride: Tuple[int, int] = (1, 1), - padding: Tuple[int, int] = (0, 0), - output_padding: Tuple[int, int] = (0, 0), - dilation: Tuple[int, int] = (1, 1), - groups: int = 1, - bias: bool = True, - padding_mode: str = 'zeros', # TODO: refine this type - dtype=None) -> None: - super().__init__() - if groups <= 0: - raise ValueError('groups must be a positive integer') - if in_channels % groups != 0: - raise ValueError('in_channels must be divisible by groups') - if out_channels % groups != 0: - raise ValueError('out_channels must be divisible by groups') - - self.in_channels = in_channels - self.out_channels = out_channels - self.kernel_size = kernel_size - self.stride = stride - self.padding = padding - self.output_padding = output_padding - self.dilation = dilation - self.groups = groups - self.padding_mode = padding_mode - - self.weight = Parameter(shape=(in_channels, out_channels // groups, - *kernel_size), - dtype=dtype) - - if bias: - self.bias = Parameter(shape=(out_channels, ), dtype=dtype) - else: - self.register_parameter('bias', None) - - def _output_padding(self, - input, - output_size, - stride, - padding, - kernel_size, - num_spatial_dims: int, - dilation=None): - if output_size is None: - ret = self.output_padding - else: - has_batch_dim = input.dim() == num_spatial_dims + 2 - num_non_spatial_dims = 2 if has_batch_dim else 1 - if len(output_size) == num_non_spatial_dims + num_spatial_dims: - output_size = output_size[num_non_spatial_dims:] - if len(output_size) != num_spatial_dims: - raise ValueError( - "ConvTranspose{}D: for {}D input, output_size must have {} or {} elements (got {})" - .format(num_spatial_dims, input.dim(), num_spatial_dims, - num_non_spatial_dims + num_spatial_dims, - len(output_size))) - - min_sizes = [] - max_sizes = [] - for d in range(num_spatial_dims): - dim_size = ( - (input.size(d + num_non_spatial_dims) - 1) * stride[d] - - 2 * padding[d] + - (dilation[d] if dilation is not None else 1) * - (kernel_size[d] - 1) + 1) - min_sizes.append(dim_size) - max_sizes.append(min_sizes[d] + stride[d] - 1) - - for i in range(len(output_size)): - size = output_size[i] - min_size = min_sizes[i] - max_size = max_sizes[i] - if size < min_size or size > max_size: - raise ValueError(( - "requested an output size of {}, but valid sizes range " - "from {} to {} (for an input of {})").format( - output_size, min_sizes, max_sizes, - input.size()[2:])) - - res = [] - for d in range(num_spatial_dims): - res.append(output_size[d] - min_sizes[d]) - - ret = res - return ret - - def forward(self, input, output_size=None): - num_spatial_dims = 2 - output_padding = self._output_padding(input, output_size, self.stride, - self.padding, self.kernel_size, - num_spatial_dims, self.dilation) - - return conv_transpose2d(input, self.weight.value, - None if self.bias is None else self.bias.value, - self.stride, self.padding, output_padding, - self.dilation, self.groups) - - -class Conv1d(Module): - - def __init__( - self, - in_channels: int, - out_channels: int, - kernel_size: int, - stride: int = 1, - padding: int = 0, - dilation: int = 1, - groups: int = 1, - bias: bool = True, - padding_mode: str = 'zeros', # TODO: refine this type - dtype=None) -> None: - super().__init__() - if groups <= 0: - raise ValueError('groups must be a positive integer') - if in_channels % groups != 0: - raise ValueError('in_channels must be divisible by groups') - if out_channels % groups != 0: - raise ValueError('out_channels must be divisible by groups') - - self.in_channels = in_channels - self.out_channels = out_channels - self.kernel_size = kernel_size - self.stride = stride - self.padding = padding - self.dilation = dilation - self.groups = groups - self.padding_mode = padding_mode - - self.weight = Parameter(shape=(out_channels, in_channels // groups, - kernel_size, 1), - dtype=dtype) - if bias: - self.bias = Parameter(shape=(out_channels, ), dtype=dtype) - else: - self.register_parameter('bias', None) - - def forward(self, input): - return conv1d(input, self.weight.value, - None if self.bias is None else self.bias.value, - self.stride, self.padding, self.dilation, self.groups) - - -class Conv3d(Module): - - def __init__( - self, - in_channels: int, - out_channels: int, - kernel_size: Tuple[int, int, int], - stride: Tuple[int, int, int] = (1, 1, 1), - padding: Tuple[int, int, int] = (0, 0, 0), - dilation: Tuple[int, int, int] = (1, 1, 1), - groups: int = 1, - bias: bool = True, - padding_mode: str = 'zeros', # TODO: refine this type - dtype=None) -> None: - super().__init__() - if groups <= 0: - raise ValueError('groups must be a positive integer') - if in_channels % groups != 0: - raise ValueError('in_channels must be divisible by groups') - if out_channels % groups != 0: - raise ValueError('out_channels must be divisible by groups') - - self.in_channels = in_channels - self.out_channels = out_channels - self.kernel_size = kernel_size - self.stride = stride - self.padding = padding - self.dilation = dilation - self.groups = groups - self.padding_mode = padding_mode - - self.weight = Parameter(shape=(out_channels, in_channels // groups, - kernel_size[0], kernel_size[1], - kernel_size[2]), - dtype=dtype) - if bias: - self.bias = Parameter(shape=(out_channels, ), dtype=dtype) - else: - self.register_parameter('bias', None) - - def forward(self, input): - return conv3d(input, self.weight.value, - None if self.bias is None else self.bias.value, - self.stride, self.padding, self.dilation, self.groups) diff --git a/tensorrt_llm/layers/embedding.py b/tensorrt_llm/layers/embedding.py deleted file mode 100644 index fd77aa9e643a..000000000000 --- a/tensorrt_llm/layers/embedding.py +++ /dev/null @@ -1,673 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -import math -from typing import Optional, Sequence, Union - -import numpy as np -import torch - -from .._utils import set_obj_attrs, str_dtype_to_torch, trt_dtype_to_np -from ..functional import (ACT2FN, Tensor, arange, concat, constant, cos, div, - embedding, exp, expand, identity, meshgrid2d, outer, - pad, shape, sin, slice, unsqueeze, where) -from ..mapping import Mapping -from ..module import Module -from ..parameter import Parameter -from .linear import ColumnLinear, Linear, RowLinear - - -class Embedding(Module): - """ - The embedding layer takes input indices (x) and the embedding lookup table (weight) as input. - And output the corresponding embeddings according to input indices. - The size of weight is [num_embeddings, embedding_dim] - - Four parameters (tp_size, tp_group, sharding_dim, tp_rank) are involved in tensor parallelism. - Only when "tp_size > 1 and tp_group is not None", tensor parallelism is enabled. - When "sharding_dim == 0", the weight is shared in the vocabulary dimension. - tp_rank must be set when sharding_dim == 0. - When "sharding_dim == 1", the weight is shard in the hidden dimension. - """ - - def __init__(self, - num_embeddings: int, - embedding_dim: int, - dtype: Optional[str] = None, - tp_size: int = 1, - tp_group: Optional[list] = None, - sharding_dim: int = 0, - tp_rank: Optional[int] = None): - super().__init__() - # num_embeddings records the total vocab size no matter using TP or not - self.num_embeddings = num_embeddings - self.embedding_dim = embedding_dim - self.tp_size = tp_size - self.tp_group = tp_group - self.sharding_dim = sharding_dim - self.tp_rank = tp_rank - self.dtype = dtype - self.tp_dim = sharding_dim - - if sharding_dim == 1: - shape = (self.num_embeddings, self.embedding_dim // self.tp_size) - elif sharding_dim == 0: - shape = (math.ceil(self.num_embeddings / self.tp_size), - self.embedding_dim) - - self.weight = Parameter(shape=shape, dtype=dtype) - - self.weight_padding_size = ((8 - shape[0] % 8) % 8, shape[1]) - - set_obj_attrs(self.weight, { - "weight_loader": self.weight_loader, - }) - - def forward(self, x): - # The embedding weight is padded to the multiple of 8. - # The reason is that when lm_head and vocab_embedding are using the same embedding weight, - # previously weights can't be depulicated in the engine because gemm will pad the weight to the multiple of 8. - # If we also pad the embedding weight to the multiple of 8, the weights can be successfully deduplicated. - # This will not affect the input and output of the gather op and perf impact is negligible. - if self.weight_padding_size[0] != 0: - padding_values = np.zeros(self.weight_padding_size, - dtype=trt_dtype_to_np( - self.weight.value.dtype)) - padding = constant(padding_values) - else: - padding = None - - return embedding(x, - self.weight.value, - tp_size=self.tp_size, - tp_group=self.tp_group, - sharding_dim=self.sharding_dim, - tp_rank=self.tp_rank, - padding=padding) - - def weight_loader(self, mapping: Mapping, param: Parameter, - loaded_weight: torch.Tensor): - # use_parallel_embedding - tp_rank = mapping.tp_rank - if self.tp_size > 1: - sharding_dim = self.sharding_dim - shard_size = param._shape[sharding_dim] - start_idx = tp_rank * shard_size - loaded_weight = loaded_weight.narrow(sharding_dim, start_idx, - shard_size) - param.value = loaded_weight - - def postprocess(self, tllm_key, weights, **kwargs): - if weights is None: - return {} - weights = weights.to(str_dtype_to_torch(self.dtype)) - return {tllm_key: weights} - - -class PromptTuningEmbedding(Embedding): - """ - PromptTuningEmbedding handles fine-tuned prompts with virtual tokens. At runtime, - a supplementary embedding dictionary is passed. Tokens whose ids are >= vocab_size are embedded - with that additional dictionary. - The prompt tuning dictionary holds multiple tasks, and each sequence is assigned a given task. - Prompt-tuned tokens from a given sequence use the adequate task dictionary, as defined by the `tasks` input. - """ - - def __init__(self, - num_embeddings, - embedding_dim, - vocab_size=None, - dtype=None, - tp_size=1, - tp_group=None, - sharding_dim=0, - tp_rank=0): - super().__init__(num_embeddings, embedding_dim, dtype, tp_size, - tp_group, sharding_dim, tp_rank) - if vocab_size is None: - vocab_size = num_embeddings - self.vocab_size = vocab_size - - def forward(self, tokens, prompt_embedding_table, tasks, task_vocab_size): - """ - Pass all tokens through both normal and prompt embedding tables. - Tokens are masked so that "normal" embedding only see "normal" tokens. Same logic for "prompt" embedding. - After those two embedding, combine results based on whether the token was "normal" or "prompt-tuned". - - Parameters: - tokens : Tensor - the ids to embed, size [batch_size, seq_len] - - prompt_embedding_table : Tensor - the additional embedding table for prompt-tuned tokens, size [num_tasks * num_tokens_per_task, hidden_size] - - tasks: Tensor - the task required by each token, size [batch_size, seq_len] - - task_vocab_size: Tensor - the number of tokens used for each task, should be equal to prompt_embedding_table's num_tokens_per_task, size [1] - - Returns: - Tokens' embedding - """ - # do not use ">=" because internally the layer works with floating points - prompt_tokens_mask = tokens > (self.vocab_size - 1) - - # clip tokens in the [0, vocab_size) range - normal_tokens = where(prompt_tokens_mask, self.vocab_size - 1, tokens) - normal_embeddings = embedding(normal_tokens, self.weight.value, - self.tp_size, self.tp_group, - self.sharding_dim, self.tp_rank) - - # put virtual tokens in the [0, max_prompt_vocab_size) range - prompt_tokens = where(prompt_tokens_mask, tokens - self.vocab_size, 0) - # add offsets to match the concatenated embedding tables - tasks = tasks * task_vocab_size - - # tasks: [batch_size, seq_len] - # prompt_tokens: [batch_size, seq_len] - # if speculative decoding is enabled the shape of prompt_tokens is [batch_size, seq_len + max_draft_len], - # so we need to expand tasks to [batch_size, seq_len + max_draft_len] - tasks = expand(tasks, shape(prompt_tokens)) - prompt_tokens = prompt_tokens + tasks - prompt_embeddings = embedding(prompt_tokens, prompt_embedding_table) - - # prompt_tokens_mask: [batch_size, seq_len] -> [batch_size, seq_len, 1] - # combine the correct sources of embedding: normal/prompt - return where(unsqueeze(prompt_tokens_mask, -1), prompt_embeddings, - normal_embeddings) - - -class LabelEmbedding(Module): - - def __init__(self, - num_classes: int, - hidden_size: int, - dropout_prob: float = 0.0, - mapping=Mapping(), - dtype=None): - super().__init__() - use_cfg_embedding = dropout_prob > 0 - self.embedding_table = Embedding(num_classes + use_cfg_embedding, - hidden_size, - tp_size=mapping.tp_size, - tp_group=mapping.tp_group, - dtype=dtype) - self.num_classes = num_classes - - def token_drop(self, labels: Tensor, force_drop_ids: Tensor): - labels = where(force_drop_ids == 1, self.num_classes, labels) - return labels - - def forward(self, labels: Tensor, force_drop_ids: Optional[Tensor] = None): - if force_drop_ids is not None: - labels = self.token_drop(labels, force_drop_ids) - embeddings = self.embedding_table(labels) - return embeddings - - -def get_1d_sincos_pos_embed_from_grid(embed_dim: int, pos: Tensor): - if embed_dim % 2 != 0: - raise ValueError("embed_dim must be divisible by 2") - - omega = torch.arange(embed_dim // 2, dtype=torch.float64) - omega /= embed_dim / 2.0 - omega = 1.0 / 10000**omega # (D/2,) - omega = constant(omega.numpy().astype(np.float32)) - - pos = pos.view([-1]) # (M,) - out = outer(pos, omega) # (M, D/2), outer product - - emb_sin = sin(out) # (M, D/2) - emb_cos = cos(out) # (M, D/2) - - emb = concat([emb_sin, emb_cos], dim=1) # (M, D) - return emb - - -def get_2d_sincos_pos_embed_from_grid(embed_dim: int, grid: Sequence[Tensor]): - if embed_dim % 2 != 0: - raise ValueError("embed_dim must be divisible by 2") - - # use half of dimensions to encode grid_h - emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, - grid[0]) # (H*W, D/2) - emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, - grid[1]) # (H*W, D/2) - - emb = concat([emb_h, emb_w], dim=1) # (H*W, D) - return emb - - -def get_2d_sincos_pos_embed( - embed_dim: int, - grid_size: Union[int, Sequence[int]], - cls_token: bool = False, - extra_tokens: int = 0, - interpolation_scale: float = 1.0, - base_size: int = 16, -): - if isinstance(grid_size, int): - grid_size = (grid_size, grid_size) - - grid_h = div( - div(arange(0, grid_size[0], 'float32'), - float(grid_size[0] / base_size)), interpolation_scale) - grid_w = div( - div(arange(0, grid_size[1], 'float32'), - float(grid_size[1] / base_size)), interpolation_scale) - grid_h, grid_w = meshgrid2d(grid_w, grid_h) # here w goes first - pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, [ - grid_h.view([1, grid_size[1], grid_size[0]]), - grid_w.view([1, grid_size[1], grid_size[0]]) - ]) - if cls_token and extra_tokens > 0: - pos_embed = concat([ - constant( - np.zeros(shape=(extra_tokens, embed_dim), - dtype=trt_dtype_to_np(pos_embed.dtype))), pos_embed - ], - dim=0) - return pos_embed - - -class SD3PatchEmbed(Module): - """ - 2D Image to Patch Embedding with support for SD3 cropping. - """ - - def __init__( - self, - height: int = 224, - width: int = 224, - patch_size: int = 16, - in_channels: int = 3, - embed_dim: int = 768, - layer_norm: bool = False, - flatten: bool = True, - bias: bool = True, - interpolation_scale: int = 1, - pos_embed_type: str = "sincos", - pos_embed_max_size: Optional[int] = None, # For SD3 cropping - dtype=None): - from diffusers.models.embeddings import \ - get_2d_sincos_pos_embed as get_2d_sincos_pos_embed_torch - - from .conv import Conv2d - from .normalization import LayerNorm - - super().__init__() - - num_patches = (height // patch_size) * (width // patch_size) - self.flatten = flatten - self.layer_norm = layer_norm - self.pos_embed_max_size = pos_embed_max_size - - self.proj = Conv2d(in_channels, - embed_dim, - kernel_size=(patch_size, patch_size), - stride=(patch_size, patch_size), - bias=bias, - dtype=dtype) - if layer_norm: - self.norm = LayerNorm(embed_dim, - elementwise_affine=False, - eps=1e-6, - dtype=dtype) - else: - self.norm = None - - self.patch_size = patch_size - self.height, self.width = height // patch_size, width // patch_size - self.base_size = height // patch_size - self.interpolation_scale = interpolation_scale - - # Calculate positional embeddings based on max size or default - if pos_embed_max_size: - grid_size = pos_embed_max_size - else: - grid_size = int(num_patches**0.5) - - if pos_embed_type is None: - self.pos_embed = None - elif pos_embed_type == "sincos": - pos_embed = get_2d_sincos_pos_embed_torch( - embed_dim, - grid_size, - base_size=self.base_size, - interpolation_scale=self.interpolation_scale, - output_type="pt", - ) - self.pos_embed = Parameter( - pos_embed.detach().cpu().float().unsqueeze(0), dtype=dtype) - else: - raise ValueError( - f"Unsupported pos_embed_type: {self.pos_embed_type}") - - def cropped_pos_embed(self, height, width): - """Crops positional embeddings for SD3 compatibility.""" - if self.pos_embed_max_size is None: - raise ValueError("`pos_embed_max_size` must be set for cropping.") - - height = height // self.patch_size - width = width // self.patch_size - if height > self.pos_embed_max_size: - raise ValueError( - f"Height ({height}) cannot be greater than `pos_embed_max_size`: {self.pos_embed_max_size}." - ) - if width > self.pos_embed_max_size: - raise ValueError( - f"Width ({width}) cannot be greater than `pos_embed_max_size`: {self.pos_embed_max_size}." - ) - - top = (self.pos_embed_max_size - height) // 2 - left = (self.pos_embed_max_size - width) // 2 - spatial_pos_embed = identity(self.pos_embed.value).view( - [1, self.pos_embed_max_size, self.pos_embed_max_size, -1]) - spatial_pos_embed = slice(spatial_pos_embed, - starts=[0, top, left, 0], - sizes=concat([ - shape(spatial_pos_embed, 0), height, - width, - shape(spatial_pos_embed, 3) - ])) - spatial_pos_embed = spatial_pos_embed.view( - concat( - [1, -1, - shape(spatial_pos_embed, - spatial_pos_embed.ndim() - 1)])) - return spatial_pos_embed - - def forward(self, latent): - # [TODO] to support height and width for runtime - if self.pos_embed_max_size is not None: - height, width = latent.shape[-2:] - else: - height, width = latent.shape[-2] // self.patch_size, latent.shape[ - -1] // self.patch_size - latent = self.proj(latent) - if self.flatten: - latent = latent.flatten(2).transpose(1, 2) # BCHW -> BNC - if self.layer_norm: - latent = self.norm(latent) - if self.pos_embed is None: - return latent.cast(latent.dtype) - # Interpolate or crop positional embeddings as needed - if self.pos_embed_max_size: - pos_embed = self.cropped_pos_embed(height, width) - else: - if self.height != height or self.width != width: - pos_embed = get_2d_sincos_pos_embed( - embed_dim=self.pos_embed.value.shape[-1], - grid_size=(height, width), - base_size=self.base_size, - interpolation_scale=self.interpolation_scale, - ) - pos_embed = unsqueeze(pos_embed.cast('float32'), axis=0) - else: - pos_embed = self.pos_embed.value - - pos_embed = pos_embed.cast(latent.dtype) - output = (latent + pos_embed).cast(latent.dtype) - return output - - -def get_timestep_embedding( - timesteps: Tensor, - embedding_dim: int, - flip_sin_to_cos: bool = False, - downscale_freq_shift: float = 1, - scale: float = 1, - max_period: int = 10000, -) -> Tensor: - """ - This matches the implementation in Denoising Diffusion Probabilistic Models: Create sinusoidal timestep embeddings. - - Args - timesteps (Tensor): - a 1-D Tensor of N indices, one per batch element. These may be fractional. - embedding_dim (int): - the dimension of the output. - flip_sin_to_cos (bool): - Whether the embedding order should be `cos, sin` (if True) or `sin, cos` (if False) - downscale_freq_shift (float): - Controls the delta between frequencies between dimensions - scale (float): - Scaling factor applied to the embeddings. - max_period (int): - Controls the maximum frequency of the embeddings - Returns - Tensor: an [N x dim] Tensor of positional embeddings. - """ - assert len(timesteps.shape) == 1, "Timesteps should be a 1d-array" - - half_dim = embedding_dim // 2 - exponent = -math.log(max_period) * np.arange( - start=0, stop=half_dim, dtype=np.float32) - exponent = exponent / (half_dim - downscale_freq_shift) - exponent = constant(exponent) - - emb = exp(exponent) - emb = unsqueeze(timesteps, -1).cast('float32') * unsqueeze(emb, 0) - - # scale embeddings - emb = scale * emb - - # flip sine and cosine embeddings - if flip_sin_to_cos: - emb = concat([cos(emb), sin(emb)], dim=-1) - else: - emb = concat([sin(emb), cos(emb)], dim=-1) - - # zero pad - if embedding_dim % 2 == 1: - emb = pad(emb, (0, 1, 0, 0)) - return emb - - -class TimestepEmbedding(Module): - - def __init__(self, - in_channels: int, - time_embed_dim: int, - act_fn: str = "silu", - out_dim: int = None, - post_act_fn: Optional[str] = None, - cond_proj_dim=None, - sample_proj_bias=True, - mapping=None, - dtype=None): - super().__init__() - tp_group = mapping.tp_group - tp_size = mapping.tp_size - self.linear_1 = ColumnLinear(in_channels, - time_embed_dim, - sample_proj_bias, - tp_group=tp_group, - tp_size=tp_size, - dtype=dtype, - gather_output=False) - - if cond_proj_dim is not None: - self.cond_proj = Linear(cond_proj_dim, - in_channels, - bias=False, - dtype=dtype) - else: - self.cond_proj = None - - self.act = ACT2FN[act_fn] - - if out_dim is not None: - time_embed_dim_out = out_dim - else: - time_embed_dim_out = time_embed_dim - self.linear_2 = RowLinear(time_embed_dim, - time_embed_dim_out, - sample_proj_bias, - tp_group=tp_group, - tp_size=tp_size, - dtype=dtype) - - if post_act_fn is None: - self.post_act = None - else: - self.post_act = ACT2FN[post_act_fn] - - def forward(self, sample, condition=None): - if condition is not None: - sample = sample + self.cond_proj(condition) - sample = self.linear_1(sample) - - if self.act is not None: - sample = self.act(sample) - - sample = self.linear_2(sample) - - if self.post_act is not None: - sample = self.post_act(sample) - return sample - - -class Timesteps(Module): - - def __init__(self, - num_channels: int, - flip_sin_to_cos: bool, - downscale_freq_shift: float, - scale: int = 1): - super().__init__() - self.num_channels = num_channels - self.flip_sin_to_cos = flip_sin_to_cos - self.downscale_freq_shift = downscale_freq_shift - self.scale = scale - - def forward(self, timesteps) -> Tensor: - t_emb = get_timestep_embedding( - timesteps, - self.num_channels, - flip_sin_to_cos=self.flip_sin_to_cos, - downscale_freq_shift=self.downscale_freq_shift, - scale=self.scale, - ) - return t_emb - - -class PixArtAlphaTextProjection(Module): - """ - Projects caption embeddings. Also handles dropout for classifier-free guidance. - - Adapted from https://github.com/PixArt-alpha/PixArt-alpha/blob/master/diffusion/model/nets/PixArt_blocks.py - """ - - def __init__(self, - in_features, - hidden_size, - out_features=None, - act_fn="gelu_tanh", - mapping=None, - dtype=None): - super().__init__() - if out_features is None: - out_features = hidden_size - tp_group = mapping.tp_group - tp_size = mapping.tp_size - self.linear_1 = ColumnLinear(in_features=in_features, - out_features=hidden_size, - bias=True, - tp_group=tp_group, - tp_size=tp_size, - dtype=dtype, - gather_output=False) - self.act_1 = ACT2FN[act_fn] - self.linear_2 = RowLinear(in_features=hidden_size, - out_features=out_features, - bias=True, - tp_group=tp_group, - tp_size=tp_size, - dtype=dtype) - - def forward(self, caption): - hidden_states = self.linear_1(caption) - hidden_states = self.act_1(hidden_states) - hidden_states = self.linear_2(hidden_states) - return hidden_states - - -class CombinedTimestepTextProjEmbeddings(Module): - - def __init__(self, - embedding_dim, - pooled_projection_dim, - mapping=Mapping(), - dtype=None): - super().__init__() - - self.time_proj = Timesteps(num_channels=256, - flip_sin_to_cos=True, - downscale_freq_shift=0) - self.timestep_embedder = TimestepEmbedding(in_channels=256, - time_embed_dim=embedding_dim, - mapping=mapping, - dtype=dtype) - self.text_embedder = PixArtAlphaTextProjection(pooled_projection_dim, - embedding_dim, - act_fn="silu", - mapping=mapping, - dtype=dtype) - - def forward(self, timestep: Tensor, pooled_projection: Tensor): - timesteps_proj = self.time_proj(timestep) - timesteps_emb = self.timestep_embedder( - timesteps_proj.cast(dtype=pooled_projection.dtype)) # (N, D) - - pooled_projections = self.text_embedder(pooled_projection) - - conditioning = timesteps_emb + pooled_projections - self.register_network_output('output', conditioning) - return conditioning - - -class CombinedTimestepLabelEmbeddings(Module): - - def __init__(self, - num_classes, - embedding_dim, - class_dropout_prob=0.0, - mapping=Mapping(), - dtype=None): - super().__init__() - self.time_proj = Timesteps(num_channels=256, - flip_sin_to_cos=True, - downscale_freq_shift=1) - self.timestep_embedder = TimestepEmbedding(in_channels=256, - time_embed_dim=embedding_dim, - mapping=mapping, - dtype=dtype) - self.class_embedder = LabelEmbedding(num_classes, - embedding_dim, - class_dropout_prob, - mapping=mapping, - dtype=dtype) - - def forward(self, - timestep: Tensor, - class_labels: Tensor, - hidden_dtype: Optional[str] = 'float32'): - timesteps_proj = self.time_proj(timestep) - timesteps_emb = self.timestep_embedder( - timesteps_proj.cast(dtype=hidden_dtype)) # (N, D) - class_labels = self.class_embedder(class_labels) # (N, D) - conditioning = timesteps_emb + class_labels # (N, D) - return conditioning diff --git a/tensorrt_llm/layers/language_adapter.py b/tensorrt_llm/layers/language_adapter.py deleted file mode 100755 index b9fa790b7d58..000000000000 --- a/tensorrt_llm/layers/language_adapter.py +++ /dev/null @@ -1,94 +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. -from dataclasses import asdict, dataclass, field - -from ..mapping import Mapping -from ..module import Module -from ..quantization import QuantMode -from .moe import MOE, MoeConfig - - -@dataclass -class LanguageAdapterConfig: - num_languages: int | None = None - top_k: int = 1 - normalization_mode: MoeConfig.ExpertScaleNormalizationMode = MoeConfig.ExpertScaleNormalizationMode.DEVICE_LIMITED - ffn_hidden_size: int = 1024 - language_list: list = field(default_factory=list) - - def validate(self) -> "LanguageAdapterConfig": - if (self.num_languages is None or self.num_languages == 0 - or self.top_k == 0): - raise ValueError( - "LanguageAdapterConfig's num_languages and top_k must not be set to 0" - ) - if (self.normalization_mode - != MoeConfig.ExpertScaleNormalizationMode.DEVICE_LIMITED): - raise ValueError( - "LanguageAdapterConfig's normalization_mode must be set to DEVICE_LIMITED since it skips both softmax and renormalization" - ) - if (len(self.language_list) == 0): - raise ValueError( - "LanguageAdapterConfig's language_list must not be empty") - return self - - @classmethod - def from_dict(cls, config: dict): - return cls(**config) - - def to_dict(self): - return asdict(self) - - def to_MOE_config(self): - return MoeConfig( - num_experts=self.num_languages, - top_k=self.top_k, - normalization_mode=self.normalization_mode, - ) - - -class LanguageAdapter(Module): - """ - Language Adapter module that uses MOE plugin with static expert selection passed in as a parameter in request. - A language MLP is selected by user for each request. - see https://arxiv.org/pdf/2005.00052 for more details. - """ - - def __init__( - self, - language_adapter_config: LanguageAdapterConfig, - hidden_size: int, - hidden_act: str, - mapping: Mapping = Mapping(), - has_mlp_bias: bool = True, - dtype=None, - quant_mode=QuantMode(0), - ): - super().__init__() - self.config = language_adapter_config - self.config.validate() - - self.layers = MOE( - hidden_size=hidden_size, - ffn_hidden_size=language_adapter_config.ffn_hidden_size, - hidden_act=hidden_act, - dtype=dtype, - bias=has_mlp_bias, - tp_group=mapping.tp_group, - tp_size=mapping.tp_size, - quant_mode=quant_mode, - static_routing=True, - moe_config=self.config.to_MOE_config(), - mapping=mapping) diff --git a/tensorrt_llm/layers/linear.py b/tensorrt_llm/layers/linear.py deleted file mode 100644 index b5539ded43cf..000000000000 --- a/tensorrt_llm/layers/linear.py +++ /dev/null @@ -1,551 +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. -from abc import ABCMeta, abstractmethod -from typing import Optional - -import numpy as np -import tensorrt as trt -import torch - -from .._common import default_net, default_trtnet -from .._utils import set_obj_attrs, str_dtype_to_torch, str_dtype_to_trt -from ..functional import (AllReduceFusionOp, AllReduceParams, Tensor, - _add_plugin_info, _create_tensor, allgather, - allreduce, cast, gemm_allreduce, low_latency_gemm, - matmul) -from ..mapping import Mapping -from ..module import Module -from ..parameter import Parameter -from ..plugin import TRT_LLM_PLUGIN_NAMESPACE -from .lora import LoraRuntimeParams - - -def _gemm_plugin(input: Tensor, - mat2: Tensor, - transa: bool = False, - transb: bool = False, - pad_lda: int = 0, - pad_ldb: int = 0, - pad_ldc: int = 0, - use_fp8: bool = False, - alpha: Optional[np.ndarray] = None, - strict_dtype: Optional[trt.DataType] = None) -> Tensor: - ''' - output = op(mat2)op(input) - - Parameters: - input : Tensor (On GPU) - The input tensor. - - mat2 : Tensor (On GPU) - The mat2 tensor. - - transa : bool - Is the input tensor transposed? Set to 'True' if you want the - input tensor to be transposed, 'False' otherwise. - - transb : bool - Is the mat2 tensor transposed? Set to 'True' if you want the - mat2 tensor to be transposed, 'False' otherwise. - - pad_lda: int - Padding to the lead dimension of input tensor. It is used to - support the strided GEMM that only uses the sub-tensor for - computation. The GEMM plugin computation is - [N, K] x [K, M+pad_lda] -> [N, M] if transa, - [N, K] x [K+pad_lda, M] -> [N, M] if not transa. - - pad_ldb: int - Padding to the lead dimension of mat2 tensor. It is used to - support the strided GEMM that only uses the sub-tensor for - computation. The GEMM plugin computation is - [N, K+pad_ldb] x [K, M] -> [N, M] if transb, - [N+pad_ldb, K] x [K, M] -> [N, M] if not transb. - - pad_ldc: int - Padding to the lead dimension of output tensor. It is used to - support the strided GEMM that only uses the sub-tensor for - computation. The GEMM plugin computation is - [N, K] x [K, M] -> [N+pad_ldc, M]. - - use_fp8: bool - Do we use fp8 GEMM. - - alpha: float - Alpha for fp8 GEMM. - - strict_dtype: trt.DataType - Set the data type for the GEMM plugin. If it is None, the data - type is the gemm_plugin type set in the plugin_config. - ''' - plg_creator = trt.get_plugin_registry().get_plugin_creator( - "Gemm", "1", TRT_LLM_PLUGIN_NAMESPACE) - assert plg_creator is not None - - if use_fp8: - assert ( - isinstance(alpha, np.ndarray) and alpha.dtype == np.float32 - and alpha.size == 1 - ), "`alpha` must be passed as a float32 ndarray if `use_fp8` is enabled for _gemm_plugin" - assert input.dtype == trt.fp8 - assert mat2.dtype == trt.fp8 - - transa = 1 if transa else 0 - transa = trt.PluginField("transa", np.array(transa, dtype=np.int32), - trt.PluginFieldType.INT32) - transb = 1 if transb else 0 - transb = trt.PluginField("transb", np.array(transb, dtype=np.int32), - trt.PluginFieldType.INT32) - pad_lda = trt.PluginField("pad_lda", np.array(pad_lda, dtype=np.int32), - trt.PluginFieldType.INT32) - pad_ldb = trt.PluginField("pad_ldb", np.array(pad_ldb, dtype=np.int32), - trt.PluginFieldType.INT32) - pad_ldc = trt.PluginField("pad_ldc", np.array(pad_ldc, dtype=np.int32), - trt.PluginFieldType.INT32) - use_fp8 = 1 if use_fp8 else 0 - use_fp8 = trt.PluginField("use_fp8", np.array(use_fp8, dtype=np.int32), - trt.PluginFieldType.INT32) - alpha = alpha if alpha else np.array(1.0, dtype=np.float32) - alpha = trt.PluginField("alpha", alpha.flatten(), - trt.PluginFieldType.FLOAT32) - - if strict_dtype is not None: - assert isinstance(strict_dtype, trt.DataType) - p_dtype = strict_dtype - else: - p_dtype = str_dtype_to_trt(default_net().plugin_config.gemm_plugin) - assert p_dtype != trt.fp8, "need to use strict dtype in gemm plugin fp8" - pf_type = trt.PluginField("type_id", np.array([int(p_dtype)], np.int32), - trt.PluginFieldType.INT32) - pfc = trt.PluginFieldCollection( - [transa, transb, pad_lda, pad_ldb, pad_ldc, pf_type, use_fp8, alpha]) - gemm_plug = plg_creator.create_plugin("gemm", pfc) - plug_inputs = [input.trt_tensor, mat2.trt_tensor] - - layer = default_trtnet().add_plugin_v2(plug_inputs, gemm_plug) - _add_plugin_info(layer, plg_creator, "gemm", pfc) - return _create_tensor(layer.get_output(0), layer) - - -class LinearBase(Module, metaclass=ABCMeta): - - def __init__( - self, - local_in_features, - local_out_features, - bias=True, - dtype=None, - tp_group=None, - tp_size=1, - share_weight=None, - strict_dtype=False, - pad_lda=0, - pad_ldc=0, - prefer_managed_weight=True, - ): - super().__init__() - self.in_features = local_in_features - self.out_features = local_out_features - self.dtype = dtype - self.pad_lda = pad_lda - self.pad_ldc = pad_ldc - self.prefer_managed_weight = prefer_managed_weight - - self.share_weight = share_weight - if not share_weight: - self.weight = Parameter( - shape=(self.out_features, self.in_features), - dtype=dtype, - prefer_managed=self.prefer_managed_weight, - ) - set_obj_attrs( - self.weight, - { - "weight_loader": self.weight_loader, - }, - ) - else: - self.weight = share_weight - - self.tp_size = tp_size - self.tp_group = tp_group - self.strict_dtype = str_dtype_to_trt( - self.dtype) if strict_dtype else None - - if bias: - self.bias = Parameter(shape=(self.out_features, ), dtype=dtype) - assert pad_ldc == 0, "not support pad_ldc with bias" - else: - self.register_parameter("bias", None) - - # see optimize_model's add_lora for LoRA initialization - self.lora = None - self.dora = None - - def weight_loader(self, mapping: Mapping, param: Parameter, - loaded_weight: torch.Tensor) -> None: - tp_rank = mapping.tp_rank - shard_size = param._shape[self.tp_split_dim()] - start_idx = tp_rank * shard_size - loaded_weight = loaded_weight.narrow(self.tp_split_dim(), start_idx, - shard_size) - - param.value = loaded_weight - - @classmethod - @abstractmethod - def tp_split_dim(cls) -> int: - pass - - def get_weight(self) -> Tensor: - if default_net( - ).plugin_config.manage_weights and self.prefer_managed_weight: - gemm_plugin = default_net().plugin_config.gemm_plugin - # nvfp4 plugin does not use this code path - use_gemm_plugin = gemm_plugin is not None and gemm_plugin != 'nvfp4' - use_low_latency_gemm_plugin = default_net( - ).plugin_config.low_latency_gemm_plugin == 'fp8' - use_gemm_allreduce_plugin = default_net( - ).plugin_config.gemm_allreduce_plugin is not None - return self.weight.get_managed_tensor(network=default_net()) - else: - return self.weight.get_constant_tensor(network=default_net()) - - def multiply_and_lora( - self, - x, - weight, - gemm_plugin: Optional[str] = None, - low_latency_gemm_plugin: Optional[str] = None, - use_fp8: bool = False, - alpha: Optional[np.ndarray] = None, - lora_runtime_params: Optional[LoraRuntimeParams] = None, - lora_hidden_state: Optional[Tensor] = None, - ): - hidden_state = x - if low_latency_gemm_plugin: - strict_dtype = str_dtype_to_trt(self.dtype) if isinstance( - self.dtype, str) else self.dtype - x = low_latency_gemm(x, weight, alpha, strict_dtype) - elif gemm_plugin and gemm_plugin != 'nvfp4': # nvfp4 gemm plugin has its own implementation - if gemm_plugin == 'fp8': - strict_dtype = str_dtype_to_trt(self.dtype) if isinstance( - self.dtype, str) else self.dtype - else: - strict_dtype = self.strict_dtype - x = _gemm_plugin(x, - weight, - transb=True, - pad_lda=self.pad_lda, - pad_ldc=self.pad_ldc, - use_fp8=use_fp8, - alpha=alpha, - strict_dtype=strict_dtype) - else: - x = matmul(x, weight, transb=True) - - if default_net( - ).plugin_config.lora_plugin and lora_runtime_params is not None: - x = x + self.lora( - hidden_state - if lora_hidden_state is None else lora_hidden_state, - lora_runtime_params=lora_runtime_params, - ) - if self.dora is not None: - x = self.dora(x, lora_runtime_params=lora_runtime_params) - - return x - - @abstractmethod - def collect_and_bias(self, x: Tensor) -> Tensor: - pass - - def multiply_collect( - self, - x, - weight, - gemm_plugin: Optional[str] = None, - low_latency_gemm_plugin: Optional[str] = None, - use_fp8: bool = False, - alpha: Optional[np.ndarray] = None, - lora_runtime_params: Optional[LoraRuntimeParams] = None, - lora_hidden_state: Optional[Tensor] = None, - **kwargs): - x = self.multiply_and_lora( - x, - weight, - gemm_plugin=gemm_plugin, - low_latency_gemm_plugin=low_latency_gemm_plugin, - use_fp8=use_fp8, - alpha=alpha, - lora_runtime_params=lora_runtime_params, - lora_hidden_state=lora_hidden_state, - ) - return self.collect_and_bias(x, **kwargs) - - def forward(self, - x, - lora_runtime_params: Optional[LoraRuntimeParams] = None, - lora_hidden_state: Optional[Tensor] = None, - **kwargs) -> Tensor: - return self.multiply_collect( - x, - self.get_weight(), - gemm_plugin=default_net().plugin_config.gemm_plugin, - use_fp8=False, - lora_runtime_params=lora_runtime_params, - lora_hidden_state=lora_hidden_state, - **kwargs) - - -class Linear(LinearBase): - - def __init__( - self, - in_features, - out_features, - bias=True, - dtype=None, - tp_group=None, - tp_size=1, - gather_output=True, - share_weight=None, - strict_dtype=False, - pad_lda=0, - pad_ldc=0, - prefer_managed_weight=True, - is_qkv=False, - ): - super().__init__( - local_in_features=in_features, - local_out_features=out_features // tp_size, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - share_weight=share_weight, - strict_dtype=strict_dtype, - pad_lda=pad_lda, - pad_ldc=pad_ldc, - prefer_managed_weight=prefer_managed_weight, - ) - self.gather_output = gather_output - self.is_qkv = is_qkv - self.tp_dim = 0 - if bias: - set_obj_attrs( - self.bias, - { - "weight_loader": self.weight_loader, - }, - ) - - @classmethod - def tp_split_dim(cls) -> int: - return 0 - - def collect_and_bias(self, x, **kwargs): - - if self.bias is not None: - bias = cast(self.bias.value, x.dtype) - x = x + bias - - if self.gather_output and self.tp_size > 1 and self.tp_group is not None: - # [dim0, local_dim] -> [dim0 * tp_size, local_dim] --> [dim0, local_dim * tp_size] - x = allgather(x, self.tp_group, gather_dim=-1) - - return x - - def postprocess(self, tllm_key, weights, **kwargs): - using_head_as_leading_dim = kwargs.get("using_head_as_leading_dim", - False) - config = kwargs.get("config", None) - if self.is_qkv: - if isinstance(weights, list): - head_size = config.hidden_size // config.num_attention_heads if config.head_size is None else config.head_size - if getattr(config, "remove_duplicated_kv_heads", False): - if config.remove_duplicated_kv_heads: - k, v = weights[1:] - k = k.reshape([ - k.shape[0] // head_size // 2, 2, head_size, - self.in_features - ]) - v = v.reshape([ - v.shape[0] // head_size // 2, 2, head_size, - self.in_features - ]) - assert (k[:, 0] == k[:, 1]).all() - assert (v[:, 0] == v[:, 1]).all() - k = k[:, 0].reshape([-1, self.in_features]) - v = v[:, 0].reshape([-1, self.in_features]) - weights[1] = k - weights[2] = v - # Duplicate kv heads in case of invalid TP size - tp_size = config.mapping.tp_size - num_kv_heads = config.num_key_value_heads - if num_kv_heads < tp_size: - for qkv_idx in range(3): - v = weights[qkv_idx] - if qkv_idx > 0: - assert tp_size % num_kv_heads == 0 - reps = tp_size // num_kv_heads - if tllm_key.endswith("bias"): - v = v.reshape(num_kv_heads, - head_size)[:, None, :].expand( - num_kv_heads, reps, head_size) - v = v.reshape(num_kv_heads * reps * head_size) - else: - v = v.reshape(num_kv_heads, head_size, - -1)[:, None, :, :].expand( - num_kv_heads, reps, head_size, - v.shape[1]) - v = v.reshape(num_kv_heads * reps * head_size, - -1) - weights[qkv_idx] = v.chunk( - tp_size, self.tp_dim)[config.mapping.tp_rank] - - weights = torch.cat(weights) - if using_head_as_leading_dim: - # Reorder [n_head, 3, head_dim, ...] into [3, n_head, head_dim, ...] - assert config.num_attention_heads == config.num_key_value_heads, "using_head_as_leading_dim require head_size to be multiple of 3." - num_heads = config.num_attention_heads - head_dim = self.out_features // (3 * num_heads) - w = weights.reshape(num_heads, 3, head_dim, -1) - w = w.transpose(0, 1) - if w.shape[-1] > 1: - weights = w.reshape(-1, self.in_features) # Weight - else: - weights = w.reshape(-1) # Bias - - if isinstance(weights, list): - weights = torch.cat(weights) - - weights = weights.to(str_dtype_to_torch(self.dtype)) - return {tllm_key: weights} - - -ColumnLinear = Linear - - -class RowLinear(LinearBase): - - def __init__( - self, - in_features, - out_features, - bias=True, - dtype=None, - tp_group=None, - tp_size=1, - strict_dtype: bool = False, - pad_lda=0, - prefer_managed_weight=True, - is_expert=False, - ): - super().__init__( - local_in_features=in_features // tp_size, - local_out_features=out_features, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - strict_dtype=strict_dtype, - pad_lda=pad_lda, - prefer_managed_weight=prefer_managed_weight, - ) - - self.tp_dim = 1 - self.tp_size = tp_size - self.is_expert = is_expert - - @classmethod - def tp_split_dim(cls) -> int: - return 1 - - def multiply_collect( - self, - x, - weight, - gemm_plugin: Optional[str] = None, - low_latency_gemm_plugin: Optional[str] = None, - use_fp8: bool = False, - alpha: Optional[np.ndarray] = None, - lora_runtime_params: Optional[LoraRuntimeParams] = None, - lora_hidden_state: Optional[Tensor] = None, - **kwargs): - - gemm_allreduce_plugin = default_net( - ).plugin_config.gemm_allreduce_plugin - if gemm_allreduce_plugin: - if lora_runtime_params != None or lora_hidden_state != None: - raise RuntimeError( - "gemm_allreduce_plugin not supported with lora.") - - output_dtype = self.dtype - if isinstance(output_dtype, str): - output_dtype = str_dtype_to_trt(output_dtype) - - x = gemm_allreduce( - a=x, - b=weight, - transa=False, # row-major - transb=True, # col-major - alpha=alpha, - group=self.tp_group, # ranks participating - output_dtype=output_dtype, - fp8_inputs_override=use_fp8) - - if self.bias is not None: - bias = cast(self.bias.value, x.dtype) - if self.is_expert: - x = x + bias / self.tp_size - else: - x = x + bias - return x - else: - return super().multiply_collect(x, weight, gemm_plugin, - low_latency_gemm_plugin, use_fp8, - alpha, lora_runtime_params, - lora_hidden_state, **kwargs) - - def collect_and_bias(self, x, **kwargs): - all_reduce_params: Optional[AllReduceParams] = kwargs.get( - "all_reduce_params", None) - if self.tp_size > 1 and self.tp_group is not None: - need_bias = self.bias is not None - fuse_bias_into_all_reduce = ( - need_bias and (all_reduce_params is not None) - and (all_reduce_params.fusion_op - == AllReduceFusionOp.RESIDUAL_RMS_NORM)) - if fuse_bias_into_all_reduce: - all_reduce_params.bias = self.bias.value - if not self.is_expert: - x = allreduce(x, - self.tp_group, - all_reduce_params=all_reduce_params) - if need_bias and not fuse_bias_into_all_reduce: - bias = cast(self.bias.value, x.dtype) - x = x + bias - else: - if need_bias and not fuse_bias_into_all_reduce: - bias = cast(self.bias.value, x.dtype) - x = x + bias / self.tp_size - return x - - if self.bias is not None: - bias = cast(self.bias.value, x.dtype) - x = x + bias - - return x diff --git a/tensorrt_llm/layers/lora.py b/tensorrt_llm/layers/lora.py deleted file mode 100644 index 26e93040dd5c..000000000000 --- a/tensorrt_llm/layers/lora.py +++ /dev/null @@ -1,181 +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. - -from typing import List - -import numpy as np - -from .._common import default_net -from ..functional import Tensor, constant, dora_plugin, lora_plugin, where -from ..module import Module - - -class LoraRuntimeParams(object): - - def __init__( - self, - lora_ranks: List[Tensor] = None, - lora_weights_pointers: List[Tensor] = None, - host_request_types: Tensor = None, - host_context_lengths: Tensor = None, - max_encoder_context_length: Tensor = None, - host_encoder_input_lengths: Tensor = None, - weight_index: int = 0, - partial_lora_mask: Tensor = None, - ): - - self.lora_ranks = lora_ranks - self.lora_weights_pointers = lora_weights_pointers - self.host_request_types = host_request_types - self.host_context_lengths = host_context_lengths - self.max_encoder_context_length = max_encoder_context_length - self.host_encoder_input_lengths = host_encoder_input_lengths - self.weight_index = weight_index - self.partial_lora_mask = partial_lora_mask # Partial LoRA for https://arxiv.org/abs/2401.16420 - - -class Lora(Module): - - def __init__(self, - in_hidden_size: int = 0, - out_hidden_sizes: List[int] = [0], - max_low_rank: int = 0) -> None: - super().__init__() - - self.in_hidden_size = in_hidden_size - self.out_hidden_sizes = out_hidden_sizes - self.max_low_rank = max_low_rank - - def forward(self, - x, - lora_runtime_params: LoraRuntimeParams = None, - is_cross_attention: bool = False): - if default_net().plugin_config.lora_plugin: - result = lora_plugin( - x, - in_hidden_size=self.in_hidden_size, - out_hidden_sizes=self.out_hidden_sizes, - host_request_types=lora_runtime_params.host_request_types, - transb=True, - # For cross attention, host_encoder_input_lengths should be used instead of host_context_lengths - host_context_lengths=lora_runtime_params.host_context_lengths - if not is_cross_attention else - lora_runtime_params.host_encoder_input_lengths, - max_low_rank=self.max_low_rank, - lora_ranks=lora_runtime_params.lora_ranks, - lora_weights_pointers=lora_runtime_params.lora_weights_pointers, - weight_index=lora_runtime_params.weight_index, - ) - if lora_runtime_params.partial_lora_mask is not None: - zero_tensor = constant(np.array([0.0], dtype=np.float16)) - if isinstance(result, List): - result = [ - where(lora_runtime_params.partial_lora_mask, r, - zero_tensor) for r in result - ] - elif isinstance(result, Tensor): - result = where(lora_runtime_params.partial_lora_mask, - result, zero_tensor) - else: - assert False - else: - assert False, "Not support lora without plugin" - - return result - - -class Dora(Module): - - def __init__(self, out_hidden_sizes: List[int] = [0]) -> None: - super().__init__() - self.out_hidden_sizes = out_hidden_sizes - - def forward(self, - x, - lora_runtime_params: LoraRuntimeParams = None, - is_cross_attention: bool = False): - assert lora_runtime_params.weight_index == 0, "DoRA does not support weight_index != 0" - if default_net().plugin_config.lora_plugin and default_net( - ).plugin_config.dora_plugin: - result = dora_plugin( - x, - out_hidden_sizes=self.out_hidden_sizes, - host_request_types=lora_runtime_params.host_request_types, - host_context_lengths=lora_runtime_params.host_context_lengths - if not is_cross_attention else - lora_runtime_params.host_encoder_input_lengths, - lora_weights_pointers=lora_runtime_params.lora_weights_pointers, - ) - else: - assert False, "Not support dora without plugin" - - return result - - -class LoraParams(object): - - def __init__( - self, - lora_ranks=None, # : List[dict[Tensor]] - lora_weights_pointers=None, # : List[dict[Tensor]] - host_context_lengths: Tensor = None, - max_encoder_context_length: Tensor = None, # For cross attention - host_request_types: Tensor = None, - host_encoder_input_lengths: Tensor = None, # For cross attention - weight_index: int = 0, - partial_lora_mask: Tensor = None, - ): - - self.lora_ranks = lora_ranks - self.lora_weights_pointers = lora_weights_pointers - - self.host_context_lengths = host_context_lengths - self.max_encoder_context_length = max_encoder_context_length - self.host_request_types = host_request_types - self.host_encoder_input_lengths = host_encoder_input_lengths - self.weight_index = weight_index - - self.partial_lora_mask = partial_lora_mask # Partial LoRA for https://arxiv.org/abs/2401.16420 - - def get_layer_params(self, layer_idx: int): - return LoraParams( - lora_ranks=[self.lora_ranks[layer_idx]], - lora_weights_pointers=[self.lora_weights_pointers[layer_idx]], - host_context_lengths=self.host_context_lengths, - max_encoder_context_length=self.max_encoder_context_length, - host_request_types=self.host_request_types, - host_encoder_input_lengths=self.host_encoder_input_lengths, - weight_index=self.weight_index, - partial_lora_mask=self.partial_lora_mask) - - def get_runtime_params(self, layer_idx: int, lora_module: str): - if f"{lora_module}_lora_ranks" in self.lora_ranks[layer_idx]: - return LoraRuntimeParams( - lora_ranks=[ - self.lora_ranks[layer_idx][f"{lora_module}_lora_ranks"] - ], - lora_weights_pointers=[ - self.lora_weights_pointers[layer_idx] - [f"{lora_module}_lora_weights_pointers"] - ], - host_context_lengths=self.host_context_lengths, - max_encoder_context_length=self.max_encoder_context_length, - host_request_types=self.host_request_types, - host_encoder_input_lengths=self.host_encoder_input_lengths, - weight_index=self.weight_index, - partial_lora_mask=self.partial_lora_mask, - ) - else: - return None diff --git a/tensorrt_llm/layers/mlp.py b/tensorrt_llm/layers/mlp.py deleted file mode 100644 index cbd2a764059a..000000000000 --- a/tensorrt_llm/layers/mlp.py +++ /dev/null @@ -1,565 +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. -from typing import Optional - -import tensorrt as trt - -from .._common import default_net -from ..functional import (ACT2FN, AllReduceParams, cast, chunk, concat, - gemm_swiglu, is_gated_activation, - low_latency_gemm_swiglu) -from ..mapping import Mapping -from ..module import Module -from ..quantization import QuantMode -from ..quantization.functional import quantize -from ..quantization.layers import FP8Linear, FP8RowLinear -from .linear import ColumnLinear, RowLinear -from .lora import LoraRuntimeParams -from .normalization import LayerNorm - - -def fc_gate_lora(hidden_states, lora, fused_gate_up_lora, lora_layer_params): - if lora_layer_params is not None: - mlp_fc_lora_params = lora_layer_params.get_runtime_params( - 0, "mlp_h_to_4h") - mlp_gate_lora_params = lora_layer_params.get_runtime_params( - 0, "mlp_gate") - mlp_gate_up_lora_params = lora_layer_params.get_runtime_params( - 0, "mlp_gate_up") - - if mlp_gate_up_lora_params is not None: - assert fused_gate_up_lora is not None - mlp_gate_up_lora = fused_gate_up_lora(hidden_states, - mlp_gate_up_lora_params) - return mlp_gate_up_lora - - elif mlp_fc_lora_params is not None and mlp_gate_lora_params is not None: - mlp_in_lora_params = LoraRuntimeParams( - lora_ranks=[ - mlp_fc_lora_params.lora_ranks[0], - mlp_gate_lora_params.lora_ranks[0] - ], - lora_weights_pointers=[ - mlp_fc_lora_params.lora_weights_pointers[0], - mlp_gate_lora_params.lora_weights_pointers[0] - ], - host_request_types=mlp_fc_lora_params.host_request_types, - host_context_lengths=mlp_fc_lora_params.host_context_lengths) - - mlp_fc_lora, mlp_gate_lora = lora(hidden_states, mlp_in_lora_params) - mlp_in_result = concat([mlp_gate_lora, mlp_fc_lora], - dim=mlp_fc_lora.rank() - 1) - return mlp_in_result - return None - - -def fc_gate_dora(hidden_states, dora, fused_gate_up_dora, lora_layer_params): - if lora_layer_params is not None: - mlp_fc_lora_params = lora_layer_params.get_runtime_params( - 0, "mlp_h_to_4h") - mlp_gate_lora_params = lora_layer_params.get_runtime_params( - 0, "mlp_gate") - mlp_gate_up_lora_params = lora_layer_params.get_runtime_params( - 0, "mlp_gate_up") - - if mlp_gate_up_lora_params is not None: - assert fused_gate_up_dora is not None - return fused_gate_up_dora(hidden_states, mlp_gate_up_lora_params) - - if mlp_fc_lora_params is not None and mlp_gate_lora_params is not None: - mlp_in_lora_params = LoraRuntimeParams( - lora_ranks=[ - mlp_fc_lora_params.lora_ranks[0], - mlp_gate_lora_params.lora_ranks[0] - ], - lora_weights_pointers=[ - mlp_fc_lora_params.lora_weights_pointers[0], - mlp_gate_lora_params.lora_weights_pointers[0] - ], - host_request_types=mlp_fc_lora_params.host_request_types, - host_context_lengths=mlp_fc_lora_params.host_context_lengths) - - return dora(hidden_states, mlp_in_lora_params) - return None - - -class MLP(Module): - - def __init__( - self, - hidden_size, - ffn_hidden_size, - hidden_act, - bias=True, - dtype=None, - tp_group=None, - tp_size=1, - quant_mode=QuantMode(0), - inner_layernorm=False, - eps=1e-05, - is_expert=False, - ): - super().__init__() - if hidden_act not in ACT2FN: - raise ValueError( - 'unsupported activation function: {}'.format(hidden_act)) - fc_output_size = 2 * ffn_hidden_size if hidden_act in [ - 'swiglu', 'gegelu' - ] else ffn_hidden_size - self.inner_layernorm = LayerNorm(ffn_hidden_size, dtype=dtype, - eps=eps) if inner_layernorm else None - - self.fc = ColumnLinear(hidden_size, - fc_output_size, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - gather_output=False) - self.proj = RowLinear(ffn_hidden_size, - hidden_size, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - is_expert=is_expert) - - self.hidden_size = hidden_size - self.ffn_hidden_size = ffn_hidden_size - self.hidden_act = hidden_act - self.dtype = dtype - self.bias = bias - self.tp_group = tp_group - self.tp_size = tp_size - self.quant_mode = quant_mode - self.eps = eps - self.is_expert = is_expert - # see optimize_model's add_lora for LoRA initialization - self.lora = None - self.dora = None - - def forward(self, hidden_states, lora_layer_params=None, gegelu_limit=None): - if lora_layer_params is not None: - assert lora_layer_params.get_runtime_params( - 0, "mlp_gate_up" - ) is None, f"LoRA module 'mlp_gate_up' is not supported in {self}" - if is_gated_activation(self.hidden_act): - inter = self.fc(hidden_states) - lora_result = fc_gate_lora(hidden_states, self.lora, None, - lora_layer_params) - if lora_result is not None: - inter = inter + lora_result - if self.dora is not None: - inter = fc_gate_dora(inter, self.dora, - self.fused_gate_up_dora, - lora_layer_params) - else: - mlp_fc_lora_params = None - if lora_layer_params is not None: - mlp_fc_lora_params = lora_layer_params.get_runtime_params( - 0, "mlp_h_to_4h") - inter = self.fc(hidden_states, mlp_fc_lora_params) - - mlp_proj_lora_params = None - if lora_layer_params is not None: - mlp_proj_lora_params = lora_layer_params.get_runtime_params( - 0, "mlp_4h_to_h") - - if self.hidden_act == 'gegelu': - inter = ACT2FN[self.hidden_act](inter, gegelu_limit) - else: - inter = ACT2FN[self.hidden_act](inter) - if self.inner_layernorm is not None: - inter = self.inner_layernorm(inter) - output = self.proj(inter, lora_runtime_params=mlp_proj_lora_params) - return output - - -class GatedMLP(MLP): - - def __init__( - self, - hidden_size, - ffn_hidden_size, - hidden_act, - bias=True, - dtype=None, - tp_group=None, - tp_size=1, - quant_mode=QuantMode(0), - inner_layernorm=False, - eps=1e-05, - is_expert=False, - ): - super().__init__(hidden_size, - ffn_hidden_size, - hidden_act, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - quant_mode=quant_mode, - inner_layernorm=inner_layernorm, - eps=eps, - is_expert=is_expert) - - self.hidden_size = hidden_size - self.ffn_hidden_size = ffn_hidden_size - self.tp_group = tp_group - self.tp_size = tp_size - - self.gate = ColumnLinear(hidden_size, - ffn_hidden_size, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - gather_output=False) - - def forward(self, - hidden_states, - lora_layer_params=None, - all_reduce_params: Optional[AllReduceParams] = None): - if lora_layer_params is not None: - assert lora_layer_params.get_runtime_params( - 0, "mlp_gate_up" - ) is None, f"LoRA module 'mlp_gate_up' is not supported in {self}" - - mlp_fc_lora_params = None - if lora_layer_params is not None: - mlp_fc_lora_params = lora_layer_params.get_runtime_params( - 0, "mlp_h_to_4h") - - mlp_gate_lora_params = None - if lora_layer_params is not None: - mlp_gate_lora_params = lora_layer_params.get_runtime_params( - 0, "mlp_gate") - - mlp_proj_lora_params = None - if lora_layer_params is not None: - mlp_proj_lora_params = lora_layer_params.get_runtime_params( - 0, "mlp_4h_to_h") - - inter = self.fc(hidden_states, mlp_fc_lora_params) - inter = ACT2FN[self.hidden_act](inter) - gate = self.gate(hidden_states, mlp_gate_lora_params) - intermediate = inter * gate - if self.inner_layernorm is not None: - intermediate = self.inner_layernorm(intermediate) - output = self.proj(intermediate, - lora_runtime_params=mlp_proj_lora_params, - all_reduce_params=all_reduce_params) - return output - - -class FusedGatedMLP(Module): - - def __init__( - self, - hidden_size, - ffn_hidden_size, - hidden_act, - bias=True, - dtype=None, - tp_group=None, - tp_size=1, - quant_mode=QuantMode(0), - inner_layernorm=False, - eps=1e-05, - is_expert=False, - ): - super().__init__() - self.hidden_size = hidden_size - self.ffn_hidden_size = ffn_hidden_size - self.hidden_act = hidden_act - self.bias = bias - self.dtype = dtype - self.tp_group = tp_group - self.tp_size = tp_size - self.quant_mode = quant_mode - - self.fused_fc = ColumnLinear( - self.hidden_size, - self.ffn_hidden_size * 2, - bias=self.bias, - dtype=self.dtype, - tp_group=self.tp_group, - tp_size=self.tp_size, - gather_output=False, - ) - self.inner_layernorm = LayerNorm(ffn_hidden_size, dtype=dtype, - eps=eps) if inner_layernorm else None - self.proj = RowLinear(ffn_hidden_size, - hidden_size, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - is_expert=is_expert) - - # see optimize_model's add_lora for LoRA initialization - self.lora = None # used for split up and gate proj - self.fused_gate_up_lora = None # used for merged up_gate proj - self.dora = None - self.fused_gate_up_dora = None - - def fc_gate_plugin(self, hidden_states, lora_layer_params=None): - # Combine the following pattern - # - # SiLU(FC(x)) * Gate(x) - # - # into: - # - # SwiGLU(FusedFC(x)) - if default_net( - ).plugin_config.low_latency_gemm_swiglu_plugin is not None: - p_dtype = default_net().plugin_config.low_latency_gemm_swiglu_plugin - else: - p_dtype = default_net().plugin_config.gemm_swiglu_plugin - use_fp8 = p_dtype == 'fp8' - assert use_fp8, "gemm_swiglu_plugin and low_latency_gemm_swiglu_plugin only supports fp8 now" - - if lora_layer_params is not None: - mlp_fc_lora_params = lora_layer_params.get_runtime_params( - 0, "mlp_h_to_4h") - mlp_gate_lora_params = lora_layer_params.get_runtime_params( - 0, "mlp_gate") - - if mlp_fc_lora_params is not None or mlp_gate_lora_params is not None: - raise NotImplementedError( - f"LoRA of splitting fc and gate is not yet implemented for gemm_swiglu_plugin" - ) - - if self.hidden_act != 'silu': - raise NotImplementedError( - f"Activation {self.hidden_act} not yet implemented for gemm_swiglu_plugin" - ) - - if self.bias: - raise NotImplementedError( - f"bias not yet implemented for gemm_swiglu_plugin fp8") - - assert isinstance( - self.fused_fc, - FP8Linear), "fp8 gemm_swiglu only supports fp8 weights" - assert isinstance( - self.proj, - FP8RowLinear), "fp8 gemm_swiglu only supports fp8 weights" - assert self.fused_fc.weight.shape == ( - self.hidden_size, self.ffn_hidden_size * 2 // - self.tp_size), "fp8 gemm_swiglu only supports (k, n) weights" - - scale_d0 = (self.fused_fc.weights_scaling_factor.raw_value.item() * - self.fused_fc.activation_scaling_factor.raw_value.item()) - scale_d1 = scale_d0 - scale_output = 1.0 / self.proj.activation_scaling_factor.raw_value.item( - ) - activation_scaling_factor = cast( - self.fused_fc.activation_scaling_factor.value, self.dtype) - if hidden_states.dtype != trt.fp8: - hidden_states = quantize(hidden_states, activation_scaling_factor, - 'fp8') - - if default_net( - ).plugin_config.low_latency_gemm_swiglu_plugin is not None: - inter = low_latency_gemm_swiglu(hidden_states, - self.fused_fc.weight.value, - scale_d0, scale_d1, scale_output) - else: - inter = gemm_swiglu(hidden_states, self.fused_fc.weight.value, None, - scale_d0, scale_d1, scale_output) - - lora_result = fc_gate_lora(hidden_states, self.lora, - self.fused_gate_up_lora, lora_layer_params) - if lora_result is not None: - inter = inter + lora_result - - return inter - - def fc_gate(self, hidden_states, lora_layer_params=None): - # Combine the following pattern - # - # SiLU(FC(x)) * Gate(x) - # - # into: - # - # SwiGLU(FusedFC(x)) - # - # Upside is we don't need to modify 4 different weight loading paths just to concat weights - - inter = self.fused_fc(hidden_states) - - lora_result = fc_gate_lora(hidden_states, self.lora, - self.fused_gate_up_lora, lora_layer_params) - if lora_result is not None: - inter = inter + lora_result - if self.dora is not None: - inter = fc_gate_dora(inter, self.dora, self.fused_gate_up_lora, - lora_layer_params) - - if self.hidden_act == 'silu': - inter = ACT2FN['swiglu'](inter) - elif self.hidden_act == 'gelu': - inter = ACT2FN['geglu'](inter) - else: - raise NotImplementedError( - f"Activation {self.hidden_act} not yet implemented for {self.__class__.__name__}." - ) - return inter - - def forward(self, - hidden_states, - lora_layer_params=None, - all_reduce_params: Optional[AllReduceParams] = None): - if default_net().plugin_config.gemm_swiglu_plugin or default_net( - ).plugin_config.low_latency_gemm_swiglu_plugin: - inter = self.fc_gate_plugin(hidden_states, lora_layer_params) - else: - inter = self.fc_gate(hidden_states, lora_layer_params) - - if self.inner_layernorm is not None: - inter = self.inner_layernorm(inter) - - mlp_proj_lora_params = None - if lora_layer_params is not None: - mlp_proj_lora_params = lora_layer_params.get_runtime_params( - 0, "mlp_4h_to_h") - output = self.proj(inter, - lora_runtime_params=mlp_proj_lora_params, - all_reduce_params=all_reduce_params) - return output - - -class LinearGELU(Module): - - def __init__(self, - dim_in: int, - dim_out: int, - approximate: str = 'tanh', - bias: bool = True, - mapping=Mapping(), - dtype=None): - super().__init__() - self.proj = ColumnLinear(dim_in, - dim_out, - bias=bias, - dtype=dtype, - tp_group=mapping.tp_group, - tp_size=mapping.tp_size) - if approximate != 'tanh': - raise NotImplementedError('GELU only support tanh now.') - - def forward(self, hidden_states): - hidden_states = self.proj(hidden_states) - hidden_states = ACT2FN['gelu_pytorch_tanh'](hidden_states) - return hidden_states - - -class LinearGEGLU(Module): - - def __init__(self, - dim_in: int, - dim_out: int, - approximate: str = 'tanh', - bias: bool = True, - mapping=Mapping(), - dtype=None): - super().__init__() - self.proj = ColumnLinear(dim_in, - dim_out * 2, - bias=bias, - dtype=dtype, - tp_group=mapping.tp_group, - tp_size=mapping.tp_size) - if approximate != 'tanh': - raise NotImplementedError('GELU only support tanh now.') - - def forward(self, hidden_states): - hidden_states = self.proj(hidden_states) - hidden_states, gate = chunk(hidden_states, - 2, - dim=(hidden_states.ndim() - 1)) - return hidden_states * ACT2FN['gelu_pytorch_tanh'](gate) - - -class LinearApproximateGELU(Module): - - def __init__(self, - dim_in: int, - dim_out: int, - bias: bool = True, - mapping=Mapping(), - dtype=None): - super().__init__() - self.proj = ColumnLinear(dim_in, - dim_out, - bias=bias, - dtype=dtype, - tp_group=mapping.tp_group, - tp_size=mapping.tp_size) - - def forward(self, x): - x = self.proj(x) - return x * ACT2FN['sigmoid'](1.702 * x) - - -class LinearSwiGLU(Module): - - def __init__(self, - dim_in: int, - dim_out: int, - bias: bool = True, - mapping=Mapping(), - dtype=None): - super().__init__() - - self.proj = ColumnLinear(dim_in, - dim_out * 2, - bias=bias, - dtype=dtype, - tp_group=mapping.tp_group, - tp_size=mapping.tp_size) - self.hidden_act = 'silu' - - def forward(self, hidden_states): - hidden_states = self.proj(hidden_states) - hidden_states, gate = chunk(hidden_states, - 2, - dim=(hidden_states.ndim() - 1)) - return hidden_states * ACT2FN[self.hidden_act](gate) - - -class LinearActivation(Module): - - def __init__(self, - dim_in: int, - dim_out: int, - bias: bool = True, - activation: str = "silu", - mapping=Mapping(), - dtype=None): - super().__init__() - - self.proj = ColumnLinear(dim_in, - dim_out, - bias=bias, - dtype=dtype, - tp_group=mapping.tp_group, - tp_size=mapping.tp_size) - self.hidden_act = activation - - def forward(self, hidden_states): - hidden_states = self.proj(hidden_states) - return ACT2FN[self.activation](hidden_states) diff --git a/tensorrt_llm/layers/moe.py b/tensorrt_llm/layers/moe.py deleted file mode 100755 index c6a0427b3a97..000000000000 --- a/tensorrt_llm/layers/moe.py +++ /dev/null @@ -1,1553 +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. -from dataclasses import asdict, dataclass -from enum import IntEnum -from typing import List, Optional, Type, Union - -import numpy as np -import tensorrt as trt -import torch - -from tensorrt_llm._torch.utils import ActivationType -from tensorrt_llm._utils import (get_init_params, str_dtype_to_torch, - str_dtype_to_trt) -from tensorrt_llm.layers.lora import LoraParams - -from .._common import default_net, default_trtnet -from .._utils import QuantModeWrapper, get_sm_version, int32_array -from ..functional import (AllReduceParams, SideStreamIDType, Tensor, - _add_plugin_info, _create_tensor, abs, allreduce, - cast, concat, constant, cuda_stream_sync, div, expand, - gather_nd, gt, is_gated_activation) -from ..functional import max as trt_max -from ..functional import (maximum, non_gated_version, nonzero, reduce_scatter, - repeat_interleave, scatter, scatter_nd, shape, - sigmoid, softmax, split, sub, sum, topk, unsqueeze, - where) -from ..mapping import Mapping -from ..module import Module, ModuleList -from ..parameter import Parameter -from ..plugin import TRT_LLM_PLUGIN_NAMESPACE -from ..quantization import GroupwiseQuantAlgo, QuantMode -from ..quantization.functional import (get_weight_scale_interleave_factor, - postprocess_weight_only, - preprocess_weights_for_mixed_gemm, - quantize) -from .linear import RowLinear -from .mlp import MLP, GatedMLP - -activation_str_to_int_map = { - # [WARNING] Keep the below in sync with cpp/tensorrt_llm/kernels/cutlass_kernels/include/common.h - "gelu": int(ActivationType.Gelu), - "gelu_new": int(ActivationType.Gelu), - "relu": int(ActivationType.Relu), - "silu": int(ActivationType.Silu), - "swiglu": int(ActivationType.Swiglu), - "geglu": int(ActivationType.Geglu), - "swiglu_bias": int(ActivationType.SwigluBias), - "identity": int(ActivationType.Identity), - "relu2": int(ActivationType.Relu2), -} - - -class MoeGroupwiseQuantParams(): - - def __init__(self, - group_size=-1, - zero=False, - pre_quant_scale=False, - use_w4a8_awq=False, - act_scale_1=None, - weight_scale_1=None, - weight_zero_1=None, - alpha_1=None, - act_scale_2=None, - weight_scale_2=None, - weight_zero_2=None, - alpha_2=None) -> None: - self.group_size = group_size - self.quant_algo = zero * GroupwiseQuantAlgo.ZERO + pre_quant_scale * GroupwiseQuantAlgo.PRE_QUANT_SCALE + use_w4a8_awq * GroupwiseQuantAlgo.W4A8_ALPHA - - self.quant_params = [] - - if group_size == -1: - return - - assert weight_scale_1 - assert weight_scale_2 - self.quant_params += [weight_scale_1, weight_scale_2] - if pre_quant_scale: - assert act_scale_1 - assert act_scale_2 - self.quant_params += [act_scale_1, act_scale_2] - if zero: - assert weight_zero_1 - assert weight_zero_2 - self.quant_params += [weight_zero_1, weight_zero_2] - if use_w4a8_awq: - assert alpha_1 - assert alpha_2 - self.quant_params += [alpha_1, alpha_2] - - -@dataclass -class MoeConfig: - - class ExpertScaleNormalizationMode(IntEnum): - NONE = 0 - RENORMALIZE = 1 - SPARSE_MIXER = 2 - DEVICE_LIMITED = 3 - DEVICE_LIMITED_RENORM = 4 - - num_experts: int = 0 - shared_expert_intermediate_size: int = 0 - - top_k: int = 0 - normalization_mode: ExpertScaleNormalizationMode = ExpertScaleNormalizationMode.RENORMALIZE - sparse_mixer_epsilon: float = 0.01 - tp_mode: int = 0 - - device_limited_n_group: int = 0 - device_limited_topk_group: int = 0 - device_limited_routed_scaling_factor: float = 1.0 - - def validate(self) -> "MoeConfig": - if (self.num_experts == 0) != (self.top_k == 0): - raise ValueError( - "Both or neither MoeConfig's num_experts and top_k must be set to 0" - ) - return self - - def has_moe(self) -> bool: - return self.num_experts > 1 - - @classmethod - def from_dict(cls, config: dict): - return cls(**config) - - def to_dict(self): - return asdict(self) - - -def _moe_plugin(moe_config, - hidden_states, - hidden_states_raw, - token_selected_experts, - token_final_scales, - expert_weights_1, - expert_weights_2, - expert_bias_1, - expert_bias_2, - expert_scale_1, - expert_scale_2, - expert_scale_3, - expert_scale_4, - expert_scale_5, - expert_scale_6, - groupwise_quant_params, - hidden_size, - ffn_hidden_size, - act_fn, - dtype, - weight_dtype, - output_dtype, - lora_params: LoraParams, - lora_max_low_rank, - quant_mode=QuantMode(0), - tp_size=1, - ep_size=1, - tp_rank=0, - ep_rank=0, - side_stream_id=SideStreamIDType.disable): - if isinstance(dtype, str): - dtype = str_dtype_to_trt(dtype) - - if isinstance(weight_dtype, str): - weight_dtype = str_dtype_to_trt(weight_dtype) - - if isinstance(output_dtype, str): - output_dtype = str_dtype_to_trt(output_dtype) - - def from_parameter(x): - if isinstance(x, Parameter): - return x.value - return x - - expert_weights_1 = from_parameter(expert_weights_1) - expert_weights_2 = from_parameter(expert_weights_2) - expert_bias_1 = from_parameter(expert_bias_1) - expert_bias_2 = from_parameter(expert_bias_2) - expert_scale_1 = from_parameter(expert_scale_1) - expert_scale_2 = from_parameter(expert_scale_2) - expert_scale_3 = from_parameter(expert_scale_3) - expert_scale_4 = from_parameter(expert_scale_4) - expert_scale_5 = from_parameter(expert_scale_5) - expert_scale_6 = from_parameter(expert_scale_6) - - # Create the plugin with our required state - num_experts = moe_config.num_experts - p_remove_input_padding = trt.PluginField( - "remove_input_padding", - np.array(np.int32(default_net().plugin_config.remove_input_padding), - dtype=np.int32), trt.PluginFieldType.INT32) - # We pass the full number of experts (not divided by ep_size) even for EP mode - p_num_experts = trt.PluginField("number_of_experts", - np.array(num_experts, dtype=np.int32), - trt.PluginFieldType.INT32) - p_experts_per_token = trt.PluginField( - "experts_per_token", np.array(moe_config.top_k, dtype=np.int32), - trt.PluginFieldType.INT32) - p_expert_hidden_size = trt.PluginField( - "expert_hidden_size", np.array(hidden_size, dtype=np.int32), - trt.PluginFieldType.INT32) - p_expert_inter_size = trt.PluginField( - "expert_inter_size", np.array(ffn_hidden_size, dtype=np.int32), - trt.PluginFieldType.INT32) - p_groupwise_quant_algo = trt.PluginField( - "groupwise_quant_algo", - np.array(groupwise_quant_params.quant_algo, dtype=np.int32), - trt.PluginFieldType.INT32) - p_group_size = trt.PluginField( - "group_size", np.array(groupwise_quant_params.group_size, - dtype=np.int32), trt.PluginFieldType.INT32) - p_activation_type = trt.PluginField( - "activation_type", - np.array(activation_str_to_int_map[act_fn], dtype=np.int32), - trt.PluginFieldType.INT32) - p_type_id = trt.PluginField("type_id", np.array([int(dtype)], - dtype=np.int32), - trt.PluginFieldType.INT32) - - p_weight_type_id = trt.PluginField( - "weight_type_id", np.array([int(weight_dtype)], dtype=np.int32), - trt.PluginFieldType.INT32) - p_output_type_id = trt.PluginField( - "output_type_id", np.array([int(output_dtype)], dtype=np.int32), - trt.PluginFieldType.INT32) - - if isinstance(quant_mode, QuantModeWrapper): - # We only need to get one quant mode here for specific moe layer - quant_mode = quant_mode[0] - p_quant_mode = trt.PluginField("quant_mode", - np.array([int(quant_mode)], dtype=np.int32), - trt.PluginFieldType.INT32) - p_use_final_scales = trt.PluginField( - "use_final_scales", - np.array([int(token_final_scales is not None)], dtype=np.int32), - trt.PluginFieldType.INT32) - p_use_bias = trt.PluginField( - "use_bias", np.array([int(expert_bias_1 is not None)], dtype=np.int32), - trt.PluginFieldType.INT32) - p_tp_size = trt.PluginField("tp_size", np.array(tp_size, dtype=np.int32), - trt.PluginFieldType.INT32) - p_tp_rank = trt.PluginField("tp_rank", np.array(tp_rank, dtype=np.int32), - trt.PluginFieldType.INT32) - p_ep_size = trt.PluginField("ep_size", np.array(ep_size, dtype=np.int32), - trt.PluginFieldType.INT32) - p_ep_rank = trt.PluginField("ep_rank", np.array(ep_rank, dtype=np.int32), - trt.PluginFieldType.INT32) - - p_force_determinism = trt.PluginField( - "force_determinism", np.array([int(False)], dtype=np.int32), - trt.PluginFieldType.INT32) - - p_side_stream_id = trt.PluginField("side_stream_id", - np.array(side_stream_id, dtype=np.int32), - trt.PluginFieldType.INT32) - - use_lora = default_net().plugin_config.lora_plugin is not None - p_use_lora = trt.PluginField("use_lora", np.array([int(use_lora)], - np.int32), - trt.PluginFieldType.INT32) - if use_lora: - p_lora_type_id = trt.PluginField( - "lora_type_id", - np.array([ - int(str_dtype_to_trt(default_net().plugin_config.lora_plugin)) - ], np.int32), trt.PluginFieldType.INT32) - p_max_low_rank = trt.PluginField( - "max_low_rank", np.array(lora_max_low_rank, dtype=np.int32), - trt.PluginFieldType.INT32) - - pfc_inputs = [ - p_remove_input_padding, p_num_experts, p_experts_per_token, - p_expert_hidden_size, p_expert_inter_size, p_groupwise_quant_algo, - p_group_size, p_activation_type, p_type_id, p_weight_type_id, - p_output_type_id, p_quant_mode, p_use_bias, p_use_final_scales, - p_tp_size, p_tp_rank, p_ep_size, p_ep_rank, p_force_determinism, - p_side_stream_id, p_use_lora - ] - - if use_lora: - pfc_inputs += [p_lora_type_id, p_max_low_rank] - - pfc = trt.PluginFieldCollection(pfc_inputs) - - # Create the plugin with our constant inputs to the constructor - plugin_creator = trt.get_plugin_registry().get_plugin_creator( - 'MixtureOfExperts', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert plugin_creator is not None - moe_plugin = plugin_creator.create_plugin("mixture_of_experts", pfc) - - # Instantiate the plugin with our specific inputs - plugin_inputs = [hidden_states, expert_weights_1, expert_weights_2] - plugin_inputs += [token_selected_experts] - - # Add conditional inputs - - # Final scales do a final rescale of the output of the experts - if token_final_scales is not None: - plugin_inputs += [token_final_scales] - - # Expert biases - if expert_bias_1: - assert expert_bias_2 - plugin_inputs += [expert_bias_1, expert_bias_2] - - # Add conditional inputs - if (quant_mode.is_weight_only() and not quant_mode.has_per_group_scaling()): - assert expert_scale_1 - assert expert_scale_2 - plugin_inputs += [expert_scale_1, expert_scale_2] - elif quant_mode.has_fp8_qdq(): - # FP8 always has scales 1-3 - assert expert_scale_1 - assert expert_scale_2 - assert expert_scale_3 - plugin_inputs += [expert_scale_1, expert_scale_2, expert_scale_3] - - if expert_scale_4 is not None: - assert output_dtype == trt.fp8 - plugin_inputs += [expert_scale_4] - - # Lora needs an extra parameter to be able to dequant the input back to backbone type - if use_lora: - assert expert_scale_5 - plugin_inputs += [expert_scale_5] - elif quant_mode.has_per_group_scaling(): - plugin_inputs += groupwise_quant_params.quant_params - # Lora needs an extra parameter to be able to dequant the input back to backbone type - if use_lora: - assert expert_scale_5 - plugin_inputs += [expert_scale_5] - elif quant_mode.has_nvfp4(): - assert expert_scale_1 - assert expert_scale_2 - assert expert_scale_3 - assert expert_scale_4 - assert expert_scale_5 - assert expert_scale_6 - plugin_inputs += [ - expert_scale_1, expert_scale_2, expert_scale_3, expert_scale_4, - expert_scale_5, expert_scale_6 - ] - - # Lora parameters - if use_lora: - # Check if lora_params is not None - moe_h_4h_params = lora_params.get_runtime_params(0, "moe_h_to_4h") - if moe_h_4h_params is not None: - moe_h_4h_weight_ptrs = moe_h_4h_params.lora_weights_pointers - moe_h_4h_lora_ranks = moe_h_4h_params.lora_ranks - plugin_inputs += (moe_h_4h_weight_ptrs + moe_h_4h_lora_ranks) - - moe_4h_h_params = lora_params.get_runtime_params(0, "moe_4h_to_h") - if moe_4h_h_params is not None: - moe_4h_h_weight_ptrs = moe_4h_h_params.lora_weights_pointers - moe_4h_h_lora_ranks = moe_4h_h_params.lora_ranks - plugin_inputs += (moe_4h_h_weight_ptrs + moe_4h_h_lora_ranks) - - if is_gated_activation(act_fn): - moe_gate_params = lora_params.get_runtime_params(0, "moe_gate") - if moe_gate_params is not None: - moe_gate_weight_ptrs = moe_gate_params.lora_weights_pointers - moe_gate_lora_ranks = moe_gate_params.lora_ranks - plugin_inputs += (moe_gate_weight_ptrs + moe_gate_lora_ranks) - - host_request_types = lora_params.host_request_types - plugin_inputs += [host_request_types] - - if default_net().plugin_config.remove_input_padding: - plugin_inputs += [lora_params.host_context_lengths] - - # A control flow tensor required to synchronize the side stream - if side_stream_id != SideStreamIDType.disable: - plugin_inputs += [hidden_states_raw] - - # Pass the inputs to the plugin - plugin_inputs = [i.trt_tensor for i in plugin_inputs] - layer = default_trtnet().add_plugin_v2(plugin_inputs, moe_plugin) - _add_plugin_info(layer, plugin_creator, "mixture_of_experts", pfc) - if not default_net().strongly_typed: - for ii in range(layer.num_inputs): - if layer.get_input(ii).dtype == str_dtype_to_trt("int8"): - layer.get_input(ii).set_dynamic_range(-127, 127) - - # Fetch the output tensor - output = _create_tensor(layer.get_output(0), layer) - - # If the side stream is enabled, also return the synchronization tensor for the side stream - if side_stream_id != SideStreamIDType.disable: - output = (output, _create_tensor(layer.get_output(1), layer)) - return output - - -def unpack_int32_into_int8(w_packed): - # Unpack inputs packed in int32/float32 into uint4 and store them in int8 format - w_packed_int4x2 = w_packed.contiguous().view(torch.uint8) - w_unpacked = torch.zeros(w_packed_int4x2.shape[0], - w_packed_int4x2.shape[1], - w_packed_int4x2.shape[2] * 2, - dtype=torch.int8) - w_unpacked[:, :, ::2] = w_packed_int4x2 % 16 - w_unpacked[:, :, 1::2] = w_packed_int4x2 // 16 - w_unpacked = w_unpacked.view(-1, 8)[:, [0, 4, 1, 5, 2, 6, 3, 7]].view( - w_unpacked.shape) - return w_unpacked.contiguous() - - -# This exists so that MOE can have the same name format as a regular MLP, just with different shaped weight tensors -class MOEWeightWrapper(Module): - - def __init__(self, in_features: int, out_features: int, - experts_per_node: int, quant_mode: QuantMode, - groupwise_quant_algo: int, group_size: int, - dtype: Union[str, - trt.DataType], weight_dtype: Union[str, - trt.DataType], - has_bias: bool, wrapper_tllm_to_externel_key_dict: dict, - tp_size: int, tp_dim: int): - super().__init__() - self.quant_mode = quant_mode - self.groupwise_quant_algo = groupwise_quant_algo - self.group_size = group_size - self.expert_shape = (experts_per_node, out_features, in_features) - self.dtype = dtype - self.weight_dtype = weight_dtype - self.has_bias = has_bias - self.tllm_to_externel_key_dict = wrapper_tllm_to_externel_key_dict - self.tp_size = tp_size - self.tp_dim = 1 - tp_dim if quant_mode.has_per_group_scaling( - ) else tp_dim - self.is_padded = False - - if quant_mode.is_weight_only( - ) and not quant_mode.has_per_group_scaling(): - bytes_per_col_scale = 2 if quant_mode.is_int4_weight_only() else 1 - # We use a different shape here because the quantized weights have their own layout - self.expert_shape = (experts_per_node, in_features, - out_features // bytes_per_col_scale) - self.per_channel_scale = Parameter(shape=(experts_per_node, - out_features), - dtype=dtype) - else: - self.register_parameter('per_channel_scale', None) - - if quant_mode.has_nvfp4(): - self.expert_shape = (experts_per_node, out_features, in_features) - weight_dtype = trt.fp4 - - if not quant_mode.has_per_group_scaling(): - self.weight = Parameter(shape=self.expert_shape, - dtype=weight_dtype, - prefer_managed=True) - - if has_bias: - self.bias = Parameter(shape=(experts_per_node, out_features), - dtype=dtype) - else: - self.register_parameter('bias', None) - - self.scaling_vector_size = 16 - if quant_mode.has_fp8_qdq(): - self.activation_scaling_factor = Parameter(shape=(1, ), - dtype=trt.float32) - self.weights_scaling_factor = Parameter(shape=(experts_per_node, 1), - dtype=trt.float32) - elif quant_mode.has_nvfp4(): - self.weights_block_scaling_factor_interleaved = Parameter( - shape=(experts_per_node, out_features, - in_features // self.scaling_vector_size), - dtype=trt.fp8) - self.weights_block_scaling_factor = Parameter( - shape=(experts_per_node, out_features, - in_features // self.scaling_vector_size), - dtype=trt.fp8) - self.activation_global_scaling_factor = Parameter(shape=(1, ), - dtype=trt.float32) - # alpha = 1.0 / (weight_global_scale * act_global_scale) - self.alpha = Parameter(shape=(experts_per_node, ), - dtype=trt.float32) - elif quant_mode.has_per_group_scaling(): - self.weight = Parameter( - shape=(experts_per_node, in_features, - out_features // 4), # int4 <--> fp16/bf16 - dtype=dtype) - if groupwise_quant_algo & GroupwiseQuantAlgo.W4A8_ALPHA: - scale_interleave_factor = get_weight_scale_interleave_factor( - in_features, group_size) - else: - scale_interleave_factor = 1 - scale_shape = (experts_per_node, - in_features // group_size // scale_interleave_factor, - out_features * scale_interleave_factor) - self.weights_scaling_factor = Parameter(shape=scale_shape, - dtype=dtype) - if groupwise_quant_algo & GroupwiseQuantAlgo.ZERO: - self.zero = Parameter(shape=scale_shape, dtype=dtype) - else: - self.register_parameter('zero', None) - if groupwise_quant_algo & GroupwiseQuantAlgo.PRE_QUANT_SCALE: - self.prequant_scaling_factor = Parameter(shape=(1, in_features), - dtype=dtype) - else: - self.register_parameter('prequant_scaling_factor', None) - if groupwise_quant_algo & GroupwiseQuantAlgo.W4A8_ALPHA: - self.alpha = Parameter(shape=(experts_per_node, 1), - dtype=trt.float32) - else: - self.register_parameter('alpha', None) - self.tllm_to_externel_key_dict.update( - {"weight": ["qweight", "qzeros", "scales"]}) - else: - self.register_parameter('weights_scaling_factor', None) - self.register_parameter('weights_block_scaling_factor', None) - self.register_parameter('weights_block_scaling_factor_interleaved', - None) - self.register_parameter('activation_scaling_factor', None) - self.register_parameter('activation_global_scaling_factor', None) - self.register_parameter('alpha', None) - self.register_parameter('zero', None) - self.register_parameter('prequant_scaling_factor', None) - - def postprocess(self, tllm_key, weights, **kwargs): - - def stack_weights(tllm_key, weights): - if "fc" in tllm_key: - weights = torch.cat([ - torch.stack(weights[:len(weights) // 2]), - torch.stack(weights[len(weights) // 2:]) - ], - dim=-2) - elif "proj" in tllm_key: - weights = torch.stack(weights) - return weights - - def postprocess_awq(tllm_key, weights): - if not tllm_key.endswith("weight"): - return {} - weights = [weights[i::3] for i in range(3)] - for idx, w in enumerate(weights): - if "fc" in tllm_key: - weights[idx] = torch.cat([ - torch.stack(w[:len(w) // 2]), - torch.stack(w[len(w) // 2:]) - ], - dim=-1) - elif "proj" in tllm_key: - weights[idx] = torch.stack(w) - qweight_int32, qzeros_int32, scales_fp16 = weights - qweight = unpack_int32_into_int8(qweight_int32) - 8 - qweight -= (qweight >> 4) << 4 - qweight = qweight.view(torch.uint8) - qweight = (qweight[:, :, 1::2] * 16 + qweight[:, :, ::2]).view( - torch.int8) - qweight = preprocess_weights_for_mixed_gemm( - qweight, torch.quint4x2, - torch.float16).view(str_dtype_to_torch(self.dtype)) - # zeros = zeros * scales - qzeros_unpacked_int32 = unpack_int32_into_int8(qzeros_int32) - zeros_x_scales_fp16 = (-qzeros_unpacked_int32 + 8) * scales_fp16 - zeros_x_scales_fp16 = zeros_x_scales_fp16.to( - str_dtype_to_torch(self.dtype)) - - results = { - tllm_key: qweight, - tllm_key.replace("weight", "weights_scaling_factor"): - scales_fp16, - tllm_key.replace("weight", "zero"): zeros_x_scales_fp16, - } - return results - - if self.quant_mode.has_per_group_scaling(): - return postprocess_awq(tllm_key, weights) - elif tllm_key.endswith("weight"): - if isinstance(weights, torch.Tensor): - weights = [weights] - else: - if self.quant_mode.has_fp8_qdq(): - experts_per_node = self.weights_scaling_factor.shape[0] - if 'fc' in tllm_key: - # Example weights: - # ['model.layers.0.block_sparse_moe.experts.0.w3.weight', - # 'model.layers.0.block_sparse_moe.experts.0.w3.weight_scale', - # ... - # 'model.layers.0.block_sparse_moe.experts.7.w3.weight', - # 'model.layers.0.block_sparse_moe.experts.7.w3.weight_scale', - # ... - # 'model.layers.0.block_sparse_moe.experts.0.w1.weight', - # 'model.layers.0.block_sparse_moe.experts.0.w1.weight_scale', - # ... - # 'model.layers.0.block_sparse_moe.experts.7.w1.weight', - # 'model.layers.0.block_sparse_moe.experts.7.w1.weight_scale'] - assert experts_per_node * 4 == len(weights) - - def w3_weight_idx(expert_id): - return 2 * expert_id - - def w3_weight_scale_idx(expert_id): - return 2 * expert_id + 1 - - def w1_weight_idx(expert_id): - return 2 * (expert_id + experts_per_node) - - def w1_weight_scale_idx(expert_id): - return 2 * (expert_id + experts_per_node) + 1 - - weights_requantized = [None] * (2 * experts_per_node) - # Since w1/w3 share weight_scale by picking the max, we need to requantize weight - for i in range(experts_per_node): - w3_weight = weights[w3_weight_idx(i)].float() - w3_weight_scale = weights[w3_weight_scale_idx( - i)].float() - w1_weight = weights[w1_weight_idx(i)].float() - w1_weight_scale = weights[w1_weight_scale_idx( - i)].float() - - max_weight_scale = max(w3_weight_scale, - w1_weight_scale) - - weights_requantized[i] = (w3_weight * - w3_weight_scale / - max_weight_scale).to( - torch.float8_e4m3fn) - weights_requantized[i + experts_per_node] = ( - w1_weight * w1_weight_scale / - max_weight_scale).to(torch.float8_e4m3fn) - - weights = weights_requantized - else: - assert 'proj' in tllm_key, f"tllm_key is {tllm_key}, which does not contain fc or proj" - # Example weights: - # ['model.layers.0.block_sparse_moe.experts.0.w2.weight', - # 'model.layers.0.block_sparse_moe.experts.0.w2.weight_scale', - # ... - # 'model.layers.0.block_sparse_moe.experts.7.w2.weight', - # 'model.layers.0.block_sparse_moe.experts.7.w2.weight_scale', - assert 2 * experts_per_node == len(weights) - - def w2_weight_idx(expert_id): - return 2 * expert_id - - # No need to requantize, simply skip weight_scale - weights = [ - weights[w2_weight_idx(i)] - for i in range(experts_per_node) - ] - - weights = stack_weights(tllm_key, weights) - - if not self.quant_mode.has_any_quant(): - # When each rank holds single expert, weights will be a list - if isinstance(weights, list): - weights = stack_weights(tllm_key, weights) - weights = weights.to(str_dtype_to_torch(self.dtype)) - - # FP8 scaling factors - if tllm_key.endswith("activation_scaling_factor"): - # Use max input range. - weights = max(weights).float().reshape((1, )) - - if tllm_key.endswith("weights_scaling_factor"): - if tllm_key.split('.')[-2] == 'fc': - # Example weights: - # ['model.layers.0.block_sparse_moe.experts.0.w3.weight_scale', - # ... - # 'model.layers.0.block_sparse_moe.experts.7.w3.weight_scale', - # 'model.layers.0.block_sparse_moe.experts.0.w1.weight_scale', - # ... - # 'model.layers.0.block_sparse_moe.experts.7.w1.weight_scale'] - experts_per_node = self.weights_scaling_factor.shape[0] - assert experts_per_node * 2 == len(weights) - - def w3_weight_scale_idx(expert_id): - return expert_id - - def w1_weight_scale_idx(expert_id): - return expert_id + experts_per_node - - # w1 and w3 share the weight scale by picking the max - weights = [ - max(weights[w3_weight_scale_idx(i)], - weights[w1_weight_scale_idx(i)]) - for i in range(experts_per_node) - ] - - weights = stack_weights(tllm_key, weights) - - # FP4 scaling factors - if tllm_key.endswith("weights_block_scaling_factor"): - weights = stack_weights(tllm_key, weights) - if tllm_key.endswith("weights_block_scaling_factor_interleaved"): - weights = stack_weights(tllm_key, weights) - weights = torch.ops.trtllm.block_scale_interleave( - weights.to(torch.float8_e4m3fn).view( - torch.uint8).cpu().contiguous()).reshape( - weights.shape).view(torch.float8_e4m3fn) - if tllm_key.endswith("activation_global_scaling_factor"): - # Use max input range. - weights = max(weights).float().reshape((1, )) - if tllm_key.endswith("alpha"): - # weights are: [e0_w3_weight_scale, e0_w3_input_scale, e1_w3_weight_scale, e1_w3_input_scale - # ..., e7_w3_weight_scale, e7_w3_input_scale, e0_w1_weight_scale, e0_w1_input_scale, ...] - weights_global_scale = weights[::2] - activation_global_scale = weights[1::2] - if 'fc' in tllm_key: - weights_global_scale = torch.stack( - weights_global_scale[:len(weights_global_scale) // 2]) - else: - weights_global_scale = torch.stack(weights_global_scale) - weights = (weights_global_scale * - max(activation_global_scale).float()).reshape((-1, )) - - # Weight only - if self.quant_mode.is_weight_only(): - if "per_channel_scale" in tllm_key: - return {} - weights = weights.to(str_dtype_to_torch(self.dtype)) - return postprocess_weight_only( - tllm_key, weights, torch.int8 if - self.quant_mode.is_int8_weight_only() else torch.quint4x2, self) - - return weights - - -class MixtureOfExperts(Module): - - def __init__(self, - moe_config: MoeConfig, - hidden_size: int, - ffn_hidden_size: int, - hidden_act: str, - mapping: Mapping = Mapping(), - bias: bool = True, - dtype=None, - tp_group: List[int] = None, - tp_size: int = 1, - quant_mode=QuantMode(0), - use_all_reduce=True, - pre_quant_scale=False, - zero=False, - use_w4a8_awq=False, - use_int8_weight=False, - group_size: int = -1, - static_routing=False): - super().__init__() - - self.moe_config = moe_config - self.num_experts = moe_config.num_experts - self.top_k = moe_config.top_k - - self.hidden_act = hidden_act - self.hidden_size = hidden_size - self.ffn_hidden_size = ffn_hidden_size - self.expert_inter_size = ffn_hidden_size - self.dtype = dtype - self.weight_dtype = dtype - self.tp_group = tp_group - self.tp_size = tp_size - self.mapping = mapping - self.quant_mode = quant_mode - self.bias = bias - self.use_all_reduce = use_all_reduce - self.zero = zero - self.pre_quant_scale = pre_quant_scale - self.use_w4a8_awq = use_w4a8_awq - self.use_int8_weight = use_int8_weight - self.group_size = group_size - - if self.use_int8_weight and self.group_size > 0: - raise NotImplementedError("INT8-GPTQ is not implemented for MoE.") - - self.static_routing = static_routing - - self.experts_per_node = self.num_experts - if self.mapping.has_moe_ep(): - if self.num_experts % self.mapping.moe_ep_size != 0: - raise ValueError( - f"MixtureOfExperts - Number of experts {self.num_experts} is not a multiple of EP size {self.mapping.moe_ep_size}" - ) - self.experts_per_node = self.experts_per_node // self.mapping.moe_ep_size - - if self.mapping.has_moe_tp(): - if self.ffn_hidden_size % self.mapping.moe_tp_size != 0: - raise ValueError( - f"MixtureOfExperts - FFN Hidden Size {self.ffn_hidden_size} is not a multiple of TP size {self.mapping.moe_tp_size}" - ) - self.expert_inter_size = self.ffn_hidden_size // self.mapping.moe_tp_size - - if quant_mode.has_fp8_rowwise(): - raise ValueError( - "MixtureOfExperts - MOE Does not support FP8 rowwise quantize") - - if quant_mode.has_fp8_qdq() and self.bias: - # TODO We will need to revisit this if we have a use case for it - raise ValueError( - f"MixtureOfExperts - Bias is not supported with FP8") - - if quant_mode.is_weight_only(): - self.weight_dtype = trt.int8 - elif quant_mode.has_fp8_qdq(): - self.weight_dtype = trt.fp8 - - rank_experts = self.mapping.ep_experts(self.num_experts) - self.wrapper_tllm_to_externel_key_dict = { - "mlp": - "block_sparse_moe", - "proj": [f"experts.{expert}.w2" for expert in rank_experts], - "fc": [f"experts.{expert}.w3" for expert in rank_experts] + - [f"experts.{expert}.w1" for expert in rank_experts] - } - - if quant_mode.has_fp8_qdq(): - self.wrapper_tllm_to_externel_key_dict.update({ - "weight": [ - "weight", "weight_scale" - ], # We need weight_scale to do requantization for w1/w3 fusion - "weights_scaling_factor": - "weight_scale", - "activation_scaling_factor": - "input_scale" - }) - - if quant_mode.has_nvfp4(): - self.wrapper_tllm_to_externel_key_dict.update({ - "weights_block_scaling_factor_interleaved": - "weight_scale", - "weights_block_scaling_factor": - "weight_scale", - "activation_global_scaling_factor": - "input_scale", - "alpha": ["weight_scale_2", "input_scale"], - }) - - # Since output dimension is usually low (in the order of 10s), no TP at - # all is more efficient as no allreduce required in the end. - # Note that if we see models that have large number of experts, we may - # need to consider add TP back here. - # TODO: Arctic has large # experts, we may need to add TP back here. - if not self.static_routing: - self.router = RowLinear( - hidden_size, - self.num_experts, - bias=False, - dtype= - "float32", # Routing is sensitive since it conditions what experts are used - tp_group=None, - tp_size=1, - strict_dtype=True) - self.router.tllm_to_externel_key_dict = { - "mlp": "block_sparse_moe", - "router": "gate" - } - - self.init_experts() - - self.max_low_rank = None - - def init_experts(self): - # Note we use horizontal fusion for gated activation to do the operation in one GEMM invocation - # The left matrix is a linear projection (no activation applied) - # The right matrix is the gating value (activation applied) - # The naming convention is the inverse of GatedMLP, but the same as `tensorrt_llm/functional.py` - fc_out_size = self.expert_inter_size * 2 if is_gated_activation( - self.hidden_act) else self.expert_inter_size - groupwise_quant_algo = self.zero * GroupwiseQuantAlgo.ZERO + self.pre_quant_scale * GroupwiseQuantAlgo.PRE_QUANT_SCALE + self.use_w4a8_awq * GroupwiseQuantAlgo.W4A8_ALPHA - - self.fc = MOEWeightWrapper(self.hidden_size, fc_out_size, - self.experts_per_node, self.quant_mode, - groupwise_quant_algo, self.group_size, - self.dtype, self.weight_dtype, self.bias, - self.wrapper_tllm_to_externel_key_dict, - self.mapping.moe_tp_size, 0) - self.proj = MOEWeightWrapper(self.expert_inter_size, self.hidden_size, - self.experts_per_node, self.quant_mode, - groupwise_quant_algo, self.group_size, - self.dtype, self.weight_dtype, self.bias, - self.wrapper_tllm_to_externel_key_dict, - self.mapping.moe_tp_size, 1) - - def default_routing(self, logits): - topk_values, topk_indices = topk(softmax(cast(logits, trt.float32), - dim=-1), - k=self.moe_config.top_k, - dim=-1) - return topk_indices, topk_values - - def renormalize(self, logits): - # Get top-k experts and renormalize their scores - token_scores, token_selected_experts = topk(cast(logits, trt.float32), - k=self.moe_config.top_k, - dim=-1) - token_final_scales = softmax(token_scores, dim=-1) - return token_selected_experts, token_final_scales - - def group_limited_greedy(self, logits): - n_group = self.moe_config.device_limited_n_group - scores = softmax(cast(logits, trt.float32), -1) - scores_shape = [shape(scores, i) for i in range(scores.ndim())] - group_scores = scores.view( - concat(scores_shape[:-1] + - [n_group, scores_shape[-1] // n_group])).max(dim=-1) - _, group_idx = topk(group_scores, - k=self.moe_config.device_limited_topk_group, - dim=-1) - group_mask = scatter(group_scores * 0, -1, group_idx, - cast(group_idx, group_scores.dtype) * 0 + 1) - score_mask = expand( - unsqueeze(group_mask, -1), - concat(scores_shape[:-1] + [n_group, scores_shape[-1] // n_group]), - ).view(concat(scores_shape)) - scores = scores * score_mask * \ - self.moe_config.device_limited_routed_scaling_factor - return scores - - def sparse_mixer(self, logits): - router_logits = cast(logits, trt.float32) - - topk_values = [] - topk_indices = [] - - assert self.top_k == 2, "Sparse mixer only supports top_k = 2" - - def mask_and_softmax(router_logits): - # Get max of remaining values - max_values = trt_max(router_logits, dim=-1, keepdim=True) - - # Calculate mask for epsilon condition - abs_values = abs(router_logits) - max_abs = maximum(abs_values, max_values) - diff = sub(max_values, router_logits) - ratio = div(diff, max_abs) - - # Apply epsilon mask - eps_mask = gt(ratio, 2 * self.moe_config.sparse_mixer_epsilon) - router_logits = where(eps_mask, -float('inf'), router_logits) - curr_values, curr_indices = topk(softmax(router_logits), - k=1, - dim=-1) - return curr_indices, curr_values - - curr_indices, curr_values = mask_and_softmax(router_logits) - topk_values.append(curr_values) - topk_indices.append(curr_indices) - - # Mask the last selected expert to -inf - router_logits = scatter(router_logits, -1, curr_indices, - curr_values * 0 - float('inf')) - - curr_indices, curr_values = mask_and_softmax(router_logits) - topk_values.append(curr_values) - topk_indices.append(curr_indices) - - # Concatenate results - values = concat(topk_values, dim=1) - indices = concat(topk_indices, dim=1) - - return indices, values - - def forward(self, - hidden_states, - lora_layer_params=None, - all_reduce_params: Optional[AllReduceParams] = None, - last_local_layer_residual=None, - side_stream_id: Optional[SideStreamIDType] = SideStreamIDType. - disable, - static_routing_input: Optional[Tensor] = None): - moe_router_lora_params = None - if lora_layer_params is not None: - moe_router_lora_params = lora_layer_params.get_runtime_params( - 0, "moe_router") - - if not self.static_routing: - routing_input = cast(hidden_states, trt.float32) - routing = self.router(routing_input, moe_router_lora_params) - else: - routing = None - - # token_selected_experts is shape (num_tokens, experts_per_token). - # It is a list of selected expert indices for each token - # token_final_scales is shape (num_tokens, experts_per_token). May be None - # It contains a final scaling/weighting factor applied to the output of each selected expert before summing the results - if self.static_routing: - token_selected_experts = static_routing_input - token_final_scales = None - elif self.moe_config.normalization_mode == MoeConfig.ExpertScaleNormalizationMode.DEVICE_LIMITED: - token_final_scales, token_selected_experts = topk( - self.group_limited_greedy(routing), - k=self.moe_config.top_k, - dim=-1) - elif self.moe_config.normalization_mode == MoeConfig.ExpertScaleNormalizationMode.DEVICE_LIMITED_RENORM: - token_final_scales, token_selected_experts = topk( - self.group_limited_greedy(routing), - k=self.moe_config.top_k, - dim=-1) - token_final_scales /= sum(token_final_scales, dim=-1, keepdim=True) - elif self.moe_config.normalization_mode == MoeConfig.ExpertScaleNormalizationMode.RENORMALIZE: - token_selected_experts, token_final_scales = self.renormalize( - routing) - elif self.moe_config.normalization_mode == MoeConfig.ExpertScaleNormalizationMode.SPARSE_MIXER: - token_selected_experts, token_final_scales = self.sparse_mixer( - routing) - else: - token_selected_experts, token_final_scales = self.default_routing( - routing) - - output = self.forward_experts(hidden_states, token_selected_experts, - token_final_scales, lora_layer_params, - side_stream_id) - if side_stream_id != SideStreamIDType.disable: - output, side_stream_sync_tensor = output - if self.use_all_reduce: - output = self.forward_allreduce(output, all_reduce_params, - last_local_layer_residual) - if side_stream_id != SideStreamIDType.disable: - # All tensors that the side channel receives as input must be synced - # on the main stream, to prevent their memory from being released or - # reused by the main stream before the side stream has finished. - tensors_to_sync = (side_stream_sync_tensor, hidden_states, - token_selected_experts, token_final_scales, - lora_layer_params) - tensors_to_sync = tuple(t for t in tensors_to_sync if t is not None) - output = (output, tensors_to_sync) - return output - - def forward_experts(self, hidden_states, token_selected_experts, - token_final_scales, lora_layer_params, side_stream_id): - - groupwise_quant_params = MoeGroupwiseQuantParams() - if self.quant_mode.has_fp8_qdq(): - assert self.fc.weight.value.dtype == trt.fp8, ( - "mlp fc weight dtype should be fp8 in the fp8 quantization mode." - ) - assert self.proj.weight.value.dtype == trt.fp8, ( - "mlp proj weight dtype should be fp8 in the fp8 quantization mode." - ) - hidden_states_quant = hidden_states - if hidden_states_quant.dtype != trt.fp8: - hidden_states_quant = quantize( - hidden_states, self.fc.activation_scaling_factor.value, - 'fp8') - - dtype_quant = trt.fp8 - weight_dtype_quant = trt.fp8 - - fc1_dequant = self.fc.weights_scaling_factor.value * self.fc.activation_scaling_factor.value - fc2_quant = div(1.0, self.proj.activation_scaling_factor.value) - fc2_dequant = self.proj.weights_scaling_factor.value * self.proj.activation_scaling_factor.value - fc1_act_dequant = self.fc.activation_scaling_factor.value - - scale_1 = fc1_dequant - scale_2 = fc2_quant - scale_3 = fc2_dequant - scale_4 = None - scale_5 = fc1_act_dequant - scale_6 = None - - output_dtype_quant = self.dtype - - if output_dtype_quant == trt.fp8 and scale_4 is None: - raise RuntimeError( - "Cannot output FP8 value without knowing quantization parameter" - ) - elif self.quant_mode.has_nvfp4(): - # We pass through the weights unchanged, the quantization is done in the plugin - hidden_states_quant = hidden_states - dtype_quant = trt.fp4 - weight_dtype_quant = trt.fp4 - output_dtype_quant = self.dtype - - scale_1 = div(1.0, self.fc.activation_global_scaling_factor.value) - scale_2 = self.fc.weights_block_scaling_factor_interleaved - scale_3 = self.fc.alpha - scale_4 = div(1.0, self.proj.activation_global_scaling_factor.value) - scale_5 = self.proj.weights_block_scaling_factor_interleaved - scale_6 = self.proj.alpha - elif self.quant_mode.has_per_group_scaling(): - hidden_states_quant = hidden_states - dtype_quant = trt.fp8 if self.use_w4a8_awq else self.dtype - weight_dtype_quant = self.weight_dtype - output_dtype_quant = self.dtype - - scale_1 = None - scale_2 = None - scale_3 = None - scale_4 = None - scale_5 = None - scale_6 = None - pre_quant_scale_1 = self.fc.prequant_scaling_factor.value if self.fc.prequant_scaling_factor else None - zero_1 = self.fc.zero.value if self.fc.zero else None - alpha_1 = self.fc.alpha.value if self.fc.alpha else None - pre_quant_scale_2 = self.proj.prequant_scaling_factor.value if self.proj.prequant_scaling_factor else None - zero_2 = self.proj.zero.value if self.proj.zero else None - alpha_2 = self.proj.alpha.value if self.proj.alpha else None - groupwise_quant_params = MoeGroupwiseQuantParams( - self.group_size, - self.zero, - self.pre_quant_scale, - self.use_w4a8_awq, - pre_quant_scale_1, - self.fc.weights_scaling_factor.value, - zero_1, - alpha_1, - pre_quant_scale_2, - self.proj.weights_scaling_factor.value, - zero_2, - alpha_2, - ) - else: - hidden_states_quant = hidden_states - dtype_quant = self.dtype - weight_dtype_quant = self.weight_dtype - output_dtype_quant = self.dtype - - scale_1 = self.fc.per_channel_scale - scale_2 = self.proj.per_channel_scale - scale_3 = None - scale_4 = None - scale_5 = None - scale_6 = None - output = _moe_plugin(self.moe_config, - hidden_states_quant, - hidden_states, - token_selected_experts, - token_final_scales, - expert_weights_1=self.fc.weight.value, - expert_weights_2=self.proj.weight.value, - expert_bias_1=self.fc.bias, - expert_bias_2=self.proj.bias, - expert_scale_1=scale_1, - expert_scale_2=scale_2, - expert_scale_3=scale_3, - expert_scale_4=scale_4, - expert_scale_5=scale_5, - expert_scale_6=scale_6, - groupwise_quant_params=groupwise_quant_params, - hidden_size=self.hidden_size, - ffn_hidden_size=self.expert_inter_size, - act_fn=self.hidden_act, - dtype=dtype_quant, - weight_dtype=weight_dtype_quant, - output_dtype=output_dtype_quant, - lora_params=lora_layer_params, - lora_max_low_rank=self.max_low_rank, - quant_mode=self.quant_mode, - tp_size=self.mapping.moe_tp_size, - tp_rank=self.mapping.moe_tp_rank, - ep_size=self.mapping.moe_ep_size, - ep_rank=self.mapping.moe_ep_rank, - side_stream_id=side_stream_id) - - return output - - def forward_allreduce(self, - output, - all_reduce_params: Optional[AllReduceParams], - last_local_layer_residual=None): - - if last_local_layer_residual is not None: - if self.mapping.tp_rank == 0: - output = output + last_local_layer_residual - else: - # we need to add this line here to minimize the numerical difference - output = output + 0 - # reshape to (-1) - output = output.view(concat([-1])) - if self.tp_size > 1 and self.tp_group is not None: - output = reduce_scatter(output, self.tp_group) - # reshape to (-1, hidden_size // tp_size) - output = output.view(concat([-1, self.hidden_size // self.tp_size])) - return output - if self.tp_size > 1 and self.tp_group is not None: - output = allreduce(output, - self.tp_group, - all_reduce_params=all_reduce_params) - return output - - def load_weights(self, moe: "MixtureOfExperts"): - ''' - Load weights from base MOE layer - ''' - raise NotImplementedError("Subclass shall override this") - - def to(self, - moe_cls: Type["MixtureOfExperts"], - quant_config=None) -> "MixtureOfExperts": - - if isinstance(moe_cls, MoeOOTB): - if self.moe_config.normalization_mode in [ - MoeConfig.ExpertScaleNormalizationMode.DEVICE_LIMITED, - MoeConfig.ExpertScaleNormalizationMode.DEVICE_LIMITED_RENORM - ]: - raise ValueError( - 'MoeOOTB doesn\'t support group_limited_greedy yet.') - from ..quantization.quantize import quantize - if isinstance(self, moe_cls): - return self - - new_moe = moe_cls(**get_init_params(self)) - # If config is not None, set quantization from config - if quant_config is not None: - quantize(new_moe, quant_config) - - new_moe.load_weights(self) - if not self.static_routing: - new_moe.router = self.router - return new_moe - - -MOE = MixtureOfExperts - - -# TODO: Support `group_limited_greedy` in MoeOOTB. -class MoeOOTB(MOE): - - def init_experts(self): - if self.quant_mode.is_weight_only(): - raise ValueError( - f"OOTB MOE does not support weight only quantization now, current quant mode: {self.quant_mode}" - ) - - if get_sm_version() >= 100: - raise RuntimeError( - "MoeOOTB does not support SM version >= 100, please use SM version < 100" - ) - - ClsMLP = GatedMLP if is_gated_activation(self.hidden_act) else MLP - - tp_size = 1 - tp_group = None - self.experts = ModuleList([ - ClsMLP(self.hidden_size, self.expert_inter_size, - non_gated_version(self.hidden_act), self.bias, self.dtype, - tp_group, tp_size, self.quant_mode) - for _ in range(self.experts_per_node) - ]) - - def moe_to_expert_lora_params(self, lora_layer_params, expert_idx): - - def get_params(module): - ranks = lora_layer_params.get_runtime_params(0, - module).lora_ranks[0] - weights_pointers = lora_layer_params.get_runtime_params( - 0, module).lora_weights_pointers[0] - return ranks, weights_pointers - - if lora_layer_params is None: - return None - fc_lora_ranks, fc_lora_weights_pointers = get_params("moe_h_to_4h") - proj_lora_ranks, proj_lora_weights_pointers = get_params("moe_4h_to_h") - gate_lora_ranks = None - gate_lora_weights_pointers = None - if is_gated_activation(self.hidden_act): - gate_lora_ranks, gate_lora_weights_pointers = get_params("moe_gate") - return LoraParams( - lora_ranks=[{ - "mlp_h_to_4h_lora_ranks": fc_lora_ranks, - "mlp_4h_to_h_lora_ranks": proj_lora_ranks, - "mlp_gate_lora_ranks": gate_lora_ranks, - }], - lora_weights_pointers=[{ - "mlp_h_to_4h_lora_weights_pointers": - fc_lora_weights_pointers, - "mlp_4h_to_h_lora_weights_pointers": - proj_lora_weights_pointers, - "mlp_gate_lora_weights_pointers": - gate_lora_weights_pointers, - }], - host_context_lengths=lora_layer_params.host_context_lengths, - max_encoder_context_length=lora_layer_params. - max_encoder_context_length, - host_request_types=lora_layer_params.host_request_types, - host_encoder_input_lengths=lora_layer_params. - host_encoder_input_lengths, - weight_index=expert_idx, - ) - - def forward_experts(self, hidden_states, token_selected_experts, - token_final_scales, lora_layer_params, side_stream_id): - assert side_stream_id == SideStreamIDType.disable, "MoeOOTB does not support using side stream" - # TODO: https://nvbugspro.nvidia.com/bug/4781396 after this nvbug is fixed, we will remove this check. - if lora_layer_params is not None: - for module in ["mlp_h_to_4h", "mlp_4h_to_h", "mlp_gate"]: - if lora_layer_params.get_runtime_params(0, module) is not None: - raise RuntimeError( - f"MoE OOTB does not support {module} LoRA module, please enable MoE plugin" - ) - - topk_indices = token_selected_experts - topk_values = token_final_scales - - hidden_size = shape(hidden_states, -1) - # [B*sq, hidden] - inputs_merged = hidden_states.view(concat([-1, hidden_size])) - flat_topk_indices = topk_indices.view( - concat([-1, shape(topk_indices, -1)])) - flat_topk_values = topk_values.view(concat([-1, - shape(topk_values, -1)])) - - # Create output space - zero_buffer = inputs_merged * 0.0 - output = zero_buffer - - expert_indices_stack = [] - indices_stack = [] - # When topk indices are equal to expert index, the expert will inference the tokens. - # Bundle all indices and experts index, then do mask once. - for i, expert in enumerate(self.experts): - if self.mapping.has_moe_ep(): - index = i + self.experts_per_node * self.mapping.moe_ep_rank - else: - index = i - expert_indices_stack.append( - flat_topk_indices.view(concat([1, shape(flat_topk_indices)]))) - - indices_stack.append(constant(int32_array(index))) - - all_expert_indices = concat(expert_indices_stack, dim=0) - indices = expand( - concat(indices_stack).view(concat([len(self.experts), 1, 1])), - shape(all_expert_indices)) - - # Create all experts mask - all_expert_mask = all_expert_indices == indices - - experts_weights = cast( - sum(flat_topk_values * - cast(all_expert_mask, flat_topk_values.dtype), - dim=-1, - keepdim=True), self.dtype) - - all_expert_mask = cast( - sum(cast(all_expert_mask, flat_topk_values.dtype), - dim=-1, - keepdim=True), 'bool') - all_expert_mask = repeat_interleave(all_expert_mask, shape(output, -1), - 2) - - # split the mask and weights for each expert - experts_mask = split(all_expert_mask, 1, dim=0) - expert_weights = split(experts_weights, 1, dim=0) - - for i, expert in enumerate(self.experts): - if self.mapping.has_moe_ep(): - index = i + self.experts_per_node * self.mapping.moe_ep_rank - else: - index = i - # get mask token index - non_zero_index = nonzero(experts_mask[i].view( - concat([-1, hidden_size]))) - non_zero_index = non_zero_index.transpose(1, 0) - input_for_expert = gather_nd(inputs_merged, non_zero_index, 0) - input_for_expert = input_for_expert.view(concat([-1, hidden_size]), - zero_is_placeholder=False) - - # Expert inference - expert_output = expert( - input_for_expert, - lora_layer_params=self.moe_to_expert_lora_params( - lora_layer_params, index)) - - # scatter expert output to real position - expert_finialized_output = zero_buffer - expert_finialized_output = scatter_nd( - expert_finialized_output, non_zero_index, - expert_output.view([-1])) * expert_weights[i] - - output += expert_finialized_output - - output = output.view(shape(hidden_states)) - - return output - - def load_weights(self, moe: MOE): - for i, expert in enumerate(self.experts): - is_gated_act = is_gated_activation(self.hidden_act) - # Gated weight pack in expert1 weights - # expert_weights_1 - experts_weight_1_raw = moe.fc.weight.raw_value - fc1_weight_scale = None - fc1_activation_scale = None - fc2_weight_scale = None - fc2_activation_scale = None - - if self.quant_mode.has_fp8_qdq(): - fc1_weight_scale = moe.fc.weights_scaling_factor.raw_value - fc1_activation_scale = moe.fc.activation_scaling_factor.raw_value - fc2_weight_scale = moe.proj.weights_scaling_factor.raw_value - fc2_activation_scale = moe.proj.activation_scaling_factor.raw_value - - if self.quant_mode.is_weight_only(): - expert.fc.weight.value = experts_weight_1_raw[ - i, :, -self.expert_inter_size:] - if is_gated_act: - expert.gate.weight.value = experts_weight_1_raw[ - i, :, :self.expert_inter_size] - else: - expert.fc.weight.value = experts_weight_1_raw[ - i, -self.expert_inter_size:, :] - if is_gated_act: - expert.gate.weight.value = experts_weight_1_raw[ - i, :self.expert_inter_size, :] - - if self.quant_mode.has_fp8_qdq(): - expert.fc.activation_scaling_factor.value = fc1_activation_scale - expert.fc.weights_scaling_factor.value = fc1_weight_scale[i] - expert.proj.activation_scaling_factor.value = fc2_activation_scale - expert.proj.weights_scaling_factor.value = fc2_weight_scale[i] - if is_gated_act: - expert.gate.activation_scaling_factor.value = fc1_activation_scale - expert.gate.weights_scaling_factor.value = fc1_weight_scale[ - i] - - if self.quant_mode.has_nvfp4(): - expert.fc.activation_global_scaling_factor.value = moe.fc.activation_global_scaling_factor.raw_value - expert.fc.weights_block_scaling_factor.value = moe.fc.weights_block_scaling_factor.raw_value[ - i, -self.expert_inter_size:, :] - expert.fc.weights_block_scaling_factor_interleaved.value = moe.fc.weights_block_scaling_factor_interleaved.raw_value[ - i, -self.expert_inter_size:, :] - expert.fc.alpha.value = np.array(moe.fc.alpha.raw_value[i]) - if is_gated_act: - expert.gate.activation_global_scaling_factor.value = moe.fc.activation_global_scaling_factor.raw_value - expert.gate.weights_block_scaling_factor.value = moe.fc.weights_block_scaling_factor.raw_value[ - i, :-self.expert_inter_size, :] - expert.gate.weights_block_scaling_factor_interleaved.value = moe.fc.weights_block_scaling_factor_interleaved.raw_value[ - i, :-self.expert_inter_size, :] - expert.gate.alpha.value = np.array( - moe.fc.alpha.raw_value[i]) - - expert.proj.activation_global_scaling_factor.value = moe.proj.activation_global_scaling_factor.raw_value - expert.proj.weights_block_scaling_factor.value = moe.proj.weights_block_scaling_factor.raw_value[ - i] - expert.proj.weights_block_scaling_factor_interleaved.value = moe.proj.weights_block_scaling_factor_interleaved.raw_value[ - i] - expert.proj.alpha.value = np.array(moe.proj.alpha.raw_value[i]) - - # expert_weights_2 - experts_weight_2_raw = moe.proj.weight.raw_value - expert.proj.weight.value = experts_weight_2_raw[i, :, :] - - has_bias = self.bias - if has_bias: - experts_bias_1_raw = moe.fc.bias.raw_value - expert.fc.bias.value = experts_bias_1_raw[ - i, -self.expert_inter_size:] - experts_bias_2_raw = moe.proj.bias.raw_value - expert.proj.bias.value = experts_bias_2_raw[i, :] - if is_gated_act: - expert.gate.bias.value = experts_bias_1_raw[ - i, :self.expert_inter_size] - - -# Add SharedMoE class -class SharedMoE(MOE): - - def __init__(self, - moe_config: MoeConfig, - hidden_size: int, - ffn_hidden_size: int, - hidden_act: str, - mapping: Mapping = Mapping(), - bias: bool = True, - dtype=None, - tp_group: List[int] = None, - tp_size: int = 1, - quant_mode=QuantMode(0), - use_shared_gate: bool = False, - use_side_stream: bool = False): - super().__init__( - moe_config=moe_config, - hidden_size=hidden_size, - ffn_hidden_size=ffn_hidden_size, - hidden_act=hidden_act, - mapping=mapping, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - quant_mode=quant_mode, - use_all_reduce=False, - ) - self.shared_expert = MLP( - hidden_size=hidden_size, - ffn_hidden_size=moe_config.shared_expert_intermediate_size, - hidden_act=hidden_act, - bias=False, - dtype=self.dtype, - tp_group=tp_group, - tp_size=tp_size, - quant_mode=self.quant_mode, - is_expert=True, - ) - self.use_shared_gate = use_shared_gate - if use_shared_gate: - self.shared_expert_gate = RowLinear( - hidden_size, - 1, - bias=False, - dtype=dtype, - tp_group=None, - tp_size=1, - ) - else: - self.shared_expert_gate = None - self.use_side_stream = use_side_stream - - def forward(self, hidden_states, lora_layer_params=None): - side_stream_id = SideStreamIDType.moe if self.use_side_stream else SideStreamIDType.disable - if self.use_side_stream: - routed_output, tensors_to_sync = super().forward( - hidden_states, - lora_layer_params=lora_layer_params, - side_stream_id=side_stream_id, - ) - else: - routed_output = super().forward( - hidden_states, - lora_layer_params=lora_layer_params, - ) - shared_output = self.shared_expert( - hidden_states, - lora_layer_params=lora_layer_params, - ) - if self.shared_expert_gate is not None: - gate_lora_params = None - if lora_layer_params is not None: - gate_lora_params = lora_layer_params.get_runtime_params( - 0, "mlp_router") - shared_output = sigmoid( - self.shared_expert_gate(hidden_states, - gate_lora_params)) * shared_output - if self.use_side_stream: - # tensors_to_sync are included in the inputs to ensure that their - # memory space is not reused for other tensors on the main stream - # until the side stream has finished - shared_output = cuda_stream_sync([shared_output, *tensors_to_sync], - side_stream_id) - hidden_states = routed_output + shared_output - if self.tp_size > 1 and self.tp_group is not None: - hidden_states = allreduce(hidden_states, self.tp_group) - return hidden_states diff --git a/tensorrt_llm/layers/normalization.py b/tensorrt_llm/layers/normalization.py deleted file mode 100644 index a09699c5f08f..000000000000 --- a/tensorrt_llm/layers/normalization.py +++ /dev/null @@ -1,341 +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. -from typing import Optional - -from ..functional import (ACT2FN, Tensor, chunk, group_norm, layer_norm, - rms_norm, unsqueeze) -from ..mapping import Mapping -from ..module import Module -from ..parameter import Parameter -from .embedding import CombinedTimestepLabelEmbeddings, Embedding -from .linear import Linear - - -class LayerNorm(Module): - - def __init__(self, - normalized_shape, - eps=1e-05, - elementwise_affine=True, - bias=True, - dtype=None, - tp_size=1, - tp_dim=-1): - super().__init__() - if isinstance(normalized_shape, int): - normalized_shape = (normalized_shape, ) - self.normalized_shape = tuple(normalized_shape) - self.elementwise_affine = elementwise_affine - if self.elementwise_affine: - self.weight = Parameter(shape=self.normalized_shape, dtype=dtype) - if bias: - self.bias = Parameter(shape=self.normalized_shape, dtype=dtype) - else: - self.register_parameter('bias', None) - else: - self.register_parameter('weight', None) - self.register_parameter('bias', None) - - self.eps = eps - self.dtype = dtype - self.tp_size = tp_size - self.tp_dim = tp_dim - - def forward(self, x, normalized_shape=None): - weight = 1. if self.weight is None else self.weight.value - bias = 0. if self.bias is None else self.bias.value - if normalized_shape is None: - normalized_shape = self.normalized_shape - return layer_norm(x, normalized_shape, weight, bias, self.eps) - - -class RmsNorm(Module): - - def __init__(self, - normalized_shape, - num_groups=1, - eps=1e-06, - elementwise_affine=True, - dtype=None): - super().__init__() - if isinstance(normalized_shape, int): - normalized_shape = (normalized_shape, ) - self.normalized_shape = tuple(normalized_shape) - self.elementwise_affine = elementwise_affine - self.num_groups = num_groups - num_channels = normalized_shape[-1] - if num_channels % num_groups != 0: - raise ValueError('num_channels must be divisible by num_groups') - if self.elementwise_affine: - self.weight = Parameter(shape=self.normalized_shape, dtype=dtype) - else: - self.register_parameter('weight', None) - - self.eps = eps - self.dtype = dtype - - def forward(self, x, normalized_shape=None): - weight = None if self.weight is None else self.weight.value - if normalized_shape is None: - normalized_shape = self.normalized_shape - return rms_norm(x, normalized_shape, self.num_groups, weight, self.eps) - - -class GroupNorm(Module): - - def __init__(self, - num_groups, - num_channels, - eps=1e-05, - affine=True, - dtype=None): - super().__init__() - - if num_channels % num_groups != 0: - raise ValueError('num_channels must be divisible by num_groups') - - self.num_groups = num_groups - self.num_channels = num_channels - self.affine = affine - - if self.affine: - self.weight = Parameter(shape=(self.num_channels, ), dtype=dtype) - self.bias = Parameter(shape=(self.num_channels, ), dtype=dtype) - else: - self.register_parameter('weight', None) - self.register_parameter('bias', None) - - self.eps = eps - - def forward(self, x): - weight = None if self.weight is None else self.weight.value - bias = None if self.bias is None else self.bias.value - return group_norm(x, self.num_groups, weight, bias, self.eps) - - -class AdaLayerNorm(Module): - - def __init__(self, - embedding_dim: int, - num_embeddings: Optional[int] = None, - output_dim: Optional[int] = None, - norm_elementwise_affine: bool = False, - norm_eps: float = 1e-5, - chunk_dim: int = 0, - mapping=Mapping(), - dtype=None): - super().__init__() - self.chunk_dim = chunk_dim - output_dim = output_dim or embedding_dim * 2 - if num_embeddings is not None: - self.emb = Embedding(num_embeddings, embedding_dim, dtype=dtype) - else: - self.emb = None - self.silu = ACT2FN['silu'] - self.linear = Linear(embedding_dim, - output_dim, - tp_group=mapping.tp_group, - tp_size=mapping.tp_size, - dtype=dtype) - self.norm = LayerNorm(output_dim // 2, - eps=norm_eps, - elementwise_affine=norm_elementwise_affine, - dtype=dtype) - - def forward(self, - x: Tensor, - timestep: Optional[Tensor] = None, - temb: Optional[Tensor] = None): - assert timestep is not None or temb is not None - if self.emb is not None and timestep is not None: - temb = self.emb(timestep) - temb = self.linear(self.silu(temb)) - if self.chunk_dim == 1: - shift, scale = chunk(temb, 2, dim=1) - shift = unsqueeze(shift, 1) - scale = unsqueeze(scale, 1) - else: - scale, shift = chunk(temb, 2, dim=0) - x = self.norm(x) * (1 + scale) + shift - return x - - -class AdaLayerNormZero(Module): - - def __init__(self, - embedding_dim: int, - num_embeddings: Optional[int] = None, - norm_type: str = "layer_norm", - bias: bool = True, - mapping=Mapping(), - dtype=None): - super().__init__() - if num_embeddings is not None: - self.emb = CombinedTimestepLabelEmbeddings(num_embeddings, - embedding_dim, - dtype=dtype) - else: - self.emb = None - - self.silu = ACT2FN['silu'] - self.linear = Linear(embedding_dim, - 6 * embedding_dim, - bias=bias, - tp_group=mapping.tp_group, - tp_size=mapping.tp_size, - dtype=dtype) - if norm_type == "layer_norm": - self.norm = LayerNorm(embedding_dim, - elementwise_affine=False, - eps=1e-6, - dtype=dtype) - elif norm_type == "fp32_layer_norm": - self.norm = LayerNorm(embedding_dim, - elementwise_affine=False, - bias=False, - dtype=dtype) - else: - raise ValueError( - f"Unsupported `norm_type` ({norm_type}) provided. Supported ones are: 'layer_norm', 'fp32_layer_norm'." - ) - - def forward(self, - x: Tensor, - timestep: Optional[Tensor] = None, - class_labels: Optional[Tensor] = None, - hidden_dtype: str = None, - emb: Optional[Tensor] = None): - assert emb is not None or self.emb is not None - if self.emb is not None: - emb = self.emb(timestep, class_labels, hidden_dtype=hidden_dtype) - emb = self.linear(self.silu(emb)) - shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = chunk( - emb, 6, dim=1) - x = self.norm(x) * (1 + unsqueeze(scale_msa, 1)) + unsqueeze( - shift_msa, 1) - return x, gate_msa, shift_mlp, scale_mlp, gate_mlp - - -class AdaLayerNormZeroSingle(Module): - - def __init__(self, - embedding_dim: int, - norm_type: str = "layer_norm", - bias: bool = True, - mapping=Mapping(), - dtype=None): - super().__init__() - - self.silu = ACT2FN['silu'] - self.linear = Linear(embedding_dim, - 3 * embedding_dim, - bias=bias, - tp_group=mapping.tp_group, - tp_size=mapping.tp_size, - dtype=dtype) - if norm_type == "layer_norm": - self.norm = LayerNorm(embedding_dim, - elementwise_affine=False, - eps=1e-6) - else: - raise ValueError( - f"Unsupported `norm_type` ({norm_type}) provided. Supported ones are: 'layer_norm', 'fp32_layer_norm'." - ) - - def forward(self, x: Tensor, emb: Optional[Tensor] = None): - emb = self.linear(self.silu(emb)) - shift_msa, scale_msa, gate_msa = chunk(emb, 3, dim=1) - x = self.norm(x) * (1 + unsqueeze(scale_msa, 1)) + unsqueeze( - shift_msa, 1) - return x, gate_msa - - -class AdaLayerNormContinuous(Module): - - def __init__(self, - embedding_dim: int, - conditioning_embedding_dim: int, - elementwise_affine: bool = True, - eps: float = 1e-5, - bias: bool = True, - norm_type: str = "layer_norm", - mapping=Mapping(), - dtype=None): - super().__init__() - self.silu = ACT2FN['silu'] - self.linear = Linear(conditioning_embedding_dim, - embedding_dim * 2, - bias=bias, - tp_group=mapping.tp_group, - tp_size=mapping.tp_size, - dtype=dtype) - if norm_type == "layer_norm": - self.norm = LayerNorm(embedding_dim, - eps=eps, - elementwise_affine=elementwise_affine, - bias=bias, - dtype=dtype) - elif norm_type == "rms_norm": - self.norm = RmsNorm(embedding_dim, - eps=eps, - elementwise_affine=elementwise_affine, - dtype=dtype) - else: - raise ValueError(f"unknown norm_type {norm_type}") - - def forward(self, x: Tensor, conditioning_embedding: Tensor): - # convert back to the original dtype in case `conditioning_embedding`` is upcasted to float32 (needed for hunyuanDiT) - emb = self.linear(self.silu(conditioning_embedding).cast(x.dtype)) - scale, shift = chunk(emb, 2, dim=1) - x = self.norm(x) * unsqueeze((1 + scale), 1) + unsqueeze(shift, 1) - return x - - -class SD35AdaLayerNormZeroX(Module): - - def __init__(self, - embedding_dim: int, - norm_type: str = "layer_norm", - bias: bool = True, - mapping=Mapping(), - dtype=None): - super().__init__() - self.silu = ACT2FN['silu'] - self.linear = Linear(embedding_dim, - 9 * embedding_dim, - bias=bias, - tp_group=mapping.tp_group, - tp_size=mapping.tp_size, - dtype=dtype) - if norm_type == "layer_norm": - self.norm = LayerNorm(embedding_dim, - elementwise_affine=False, - eps=1e-6, - dtype=dtype) - else: - raise ValueError( - f"Unsupported `norm_type` ({norm_type}) provided. Supported ones are: 'layer_norm'." - ) - - def forward(self, hidden_states: Tensor, emb: Tensor): - emb = self.linear(self.silu(emb).cast(hidden_states.dtype)) - shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp, shift_msa2, scale_msa2, gate_msa2 = chunk( - emb, 9, dim=1) - norm_hidden_states = self.norm(hidden_states) - hidden_states = norm_hidden_states * ( - 1 + unsqueeze(scale_msa, 1)) + unsqueeze(shift_msa, 1) - norm_hidden_states2 = norm_hidden_states * ( - 1 + unsqueeze(scale_msa2, 1)) + unsqueeze(shift_msa2, 1) - return hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp, norm_hidden_states2, gate_msa2 diff --git a/tensorrt_llm/layers/pooling.py b/tensorrt_llm/layers/pooling.py deleted file mode 100644 index 78153e6b1ef9..000000000000 --- a/tensorrt_llm/layers/pooling.py +++ /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. -from typing import Optional, Tuple - -from ..functional import avg_pool2d -from ..module import Module - - -class AvgPool2d(Module): - - def __init__(self, - kernel_size: Tuple[int], - stride: Optional[Tuple[int]] = None, - padding: Optional[Tuple[int]] = (0, 0), - ceil_mode: bool = False, - count_include_pad: bool = True) -> None: - super().__init__() - self.kernel_szie = kernel_size - self.stride = stride - self.padding = padding - self.ceil_mode = ceil_mode - self.count_include_pad = count_include_pad - - def forward(self, input): - return avg_pool2d(input, self.kernel_szie, self.stride, self.padding, - self.ceil_mode, self.count_include_pad) diff --git a/tensorrt_llm/layers/recurrent.py b/tensorrt_llm/layers/recurrent.py deleted file mode 100644 index fde8dc24eb24..000000000000 --- a/tensorrt_llm/layers/recurrent.py +++ /dev/null @@ -1,316 +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. - -from typing import Optional - -import torch - -from .._utils import set_obj_attrs -from ..functional import Tensor, allgather, cast, concat, matmul, rg_lru, shape -from ..mapping import Mapping -from ..module import Module -from ..parameter import Parameter -from .linear import ColumnLinear, RowLinear -from .ssm import MambaConv1d - - -class GroupedLinear(Module): - - def __init__(self, - in_features, - out_features, - num_blocks, - bias=True, - dtype=None, - use_fp8=False, - tp_group=None, - tp_size=1, - gather_output=True, - strict_dtype=False, - fuse_bias=False): - super().__init__() - assert in_features % num_blocks == 0 and out_features % num_blocks == 0 - assert num_blocks % tp_size == 0 - assert not (gather_output and fuse_bias) - self.in_features = in_features // tp_size - self.out_features = out_features // tp_size - self.num_blocks = num_blocks // tp_size - self.dtype = dtype - self.use_fp8 = use_fp8 - self.fuse_bias = fuse_bias - - self.weight = Parameter(shape=(self.num_blocks, - self.in_features // self.num_blocks, - self.out_features // self.num_blocks), - dtype=('fp8' if use_fp8 else dtype)) - set_obj_attrs(self.weight, { - "weight_loader": self.weight_loader, - }) - - self.tp_size = tp_size - self.tp_group = tp_group - self.gather_output = gather_output - self.strict_dtype = self.dtype if strict_dtype else None - - if bias: - self.bias = Parameter(shape=(self.num_blocks, - self.out_features // self.num_blocks), - dtype=dtype) - set_obj_attrs(self.bias, { - "weight_loader": self.weight_loader, - }) - else: - self.register_parameter('bias', None) - - def multiply_gather(self, x, weight): - grouped_shape = [] - out_shape = [] - ndim = x.ndim() - for i in range(x.ndim() - 1): - grouped_shape.append(shape(x, i)) - out_shape.append(shape(x, i)) - grouped_shape.extend( - [self.num_blocks, self.in_features // self.num_blocks]) - out_shape.append(self.out_features) - x = x.view(concat(grouped_shape)).permute([i for i in range(ndim - 2)] + - [-2, -3, -1]) - x = matmul(x, weight) - x = x.permute([i for i in range(ndim - 2)] + [-2, -3, -1]) - - if self.bias is not None and not self.fuse_bias: - bias = cast(self.bias.value, x.dtype) - x = x + bias - x = x.view(concat(out_shape)) - - if self.gather_output and self.tp_size > 1 and self.tp_group is not None: - # [dim0, local_dim] -> [dim0 * tp_size, local_dim] --> [dim0, local_dim * tp_size] - x = allgather(x, self.tp_group, gather_dim=-1) - - return x - - def forward(self, x): - return self.multiply_gather(x, self.weight.value) - - def weight_loader(self, mapping: Mapping, param: Parameter, - loaded_weight: torch.Tensor): - tp_rank = mapping.tp_rank - output_dim = 0 - shard_size = param._shape[output_dim] - start_idx = tp_rank * shard_size - loaded_weight = loaded_weight.narrow(output_dim, start_idx, shard_size) - param.value = loaded_weight - - -class RgLru(Module): - - def __init__(self, - lru_width, - num_heads=1, - dtype=None, - tp_group=None, - tp_size=1): - super().__init__() - self.lru_width = lru_width - self.dtype = dtype - self.num_heads = num_heads - self.tp_group = tp_group - self.tp_size = tp_size - - self.recurrent_param = Parameter(shape=(self.lru_width // - self.tp_size, ), - dtype=self.dtype) - self.input_gate = GroupedLinear(self.lru_width, - self.lru_width, - self.num_heads, - dtype=self.dtype, - tp_group=self.tp_group, - tp_size=self.tp_size, - gather_output=False, - fuse_bias=True) - self.recurrent_gate = GroupedLinear(self.lru_width, - self.lru_width, - self.num_heads, - dtype=self.dtype, - tp_group=self.tp_group, - tp_size=self.tp_size, - gather_output=False, - fuse_bias=True) - - def forward(self, - x: Tensor, - y: Tensor, - y_bias: Tensor, - lru_state: Tensor, - host_request_types: Tensor, - last_token_ids: Tensor, - slot_mapping: Optional[Tensor] = None): - gate_x = self.input_gate(x) - gate_a = self.recurrent_gate(x) - out, lru_state = rg_lru(input=x, - gate_x=gate_x, - gate_x_bias=self.input_gate.bias.value, - gate_a=gate_a, - gate_a_bias=self.recurrent_gate.bias.value, - y=y, - y_bias=y_bias, - state_or_ptr=lru_state, - A=self.recurrent_param.value, - host_request_types=host_request_types, - last_token_ids=last_token_ids, - dim=self.lru_width // self.tp_size, - dtype=self.dtype, - slot_mapping=slot_mapping) - return out, lru_state - - -class FusedRgLru(Module): - - def __init__(self, - lru_width, - num_heads=1, - dtype=None, - tp_group=None, - tp_size=1): - super().__init__() - self.lru_width = lru_width - self.tp_size = tp_size - self.dtype = dtype - self.dim = self.lru_width // self.tp_size - self.block_size = self.lru_width // num_heads - - self.recurrent_param = Parameter(shape=(self.lru_width // tp_size, ), - dtype=dtype) - self.gate = GroupedLinear(self.lru_width, - self.lru_width * 2, - num_heads, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - gather_output=False, - fuse_bias=True) - - def forward(self, - x: Tensor, - y: Tensor, - y_bias: Tensor, - lru_state: Tensor, - host_request_types: Tensor, - last_token_ids: Tensor, - slot_mapping: Optional[Tensor] = None): - gate = self.gate(x) - out, lru_state = rg_lru(input=x, - gate=gate, - gate_bias=self.gate.bias.value, - block_size=self.block_size, - y=y, - y_bias=y_bias, - state_or_ptr=lru_state, - A=self.recurrent_param.value, - host_request_types=host_request_types, - last_token_ids=last_token_ids, - dim=self.dim, - dtype=self.dtype, - slot_mapping=slot_mapping) - return out, lru_state - - -class Recurrent(Module): - - def __init__( - self, - width, - lru_width, - d_conv=4, - num_heads=1, - dtype=None, - tp_group=None, - tp_size=1, - ): - super().__init__() - self.width = width - self.lru_width = lru_width - self.d_conv = d_conv - self.dtype = dtype - - self.linear_x = ColumnLinear(self.width, - self.lru_width, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - gather_output=False) - self.linear_y = ColumnLinear(self.width, - self.lru_width, - bias=False, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - gather_output=False) - self.y_bias = Parameter(shape=(self.lru_width // tp_size, ), - dtype=dtype) - - self.conv1d = MambaConv1d(self.lru_width // tp_size, - self.d_conv, - dtype=self.dtype, - apply_silu=False) - - self.rg_lru = RgLru(self.lru_width, - num_heads=num_heads, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size) - - self.linear_out = RowLinear(self.lru_width, - self.width, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size) - - def forward(self, - hidden_states: Tensor, - conv_state: Tensor, - lru_state: Tensor, - host_request_types: Tensor, - last_token_ids: Tensor, - host_context_lengths: Optional[Tensor] = None, - slot_mapping: Optional[Tensor] = None, - conv_indices: Optional[Tensor] = None): - ''' - Parameters: - hidden_states: [B, L, D] or [T, D] - conv_state: [B, W, D] or [1] of type int64 for paged state - lru_state: [B, N] or [1] of type int64 for paged state - host_request_types: [B] - last_token_ids: [B] - host_context_lengths: [B] - slot_mapping: [B] - conv_indices: [B] - ''' - # y branch - y = self.linear_y(hidden_states) - - # x branch - x = self.linear_x(hidden_states) - x_conv, conv_state = self.conv1d(x, conv_state, host_request_types, - last_token_ids, host_context_lengths, - slot_mapping, conv_indices) - - # rg-lru - out, lru_state = self.rg_lru(x_conv, y, self.y_bias.value, lru_state, - host_request_types, last_token_ids, - slot_mapping) - - # linear out - out = self.linear_out(out) - return out, conv_state, lru_state diff --git a/tensorrt_llm/layers/ssm.py b/tensorrt_llm/layers/ssm.py deleted file mode 100644 index 41c7a5b8f2f3..000000000000 --- a/tensorrt_llm/layers/ssm.py +++ /dev/null @@ -1,358 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import math -from typing import Optional - -from .._common import default_net -from ..functional import (ACT2FN, Tensor, concat, conv2d, gather, mamba_conv1d, - permute, selective_scan, shape, split, view) -from ..module import Module -from ..parameter import Parameter -from .linear import ColumnLinear, Linear, RowLinear -from .normalization import RmsNorm - - -class MambaConv1d(Module): - - def __init__(self, - d_inner, - d_conv=4, - pre_stride=0, - post_stride=0, - dtype=None, - apply_silu=True): - super().__init__() - self.d_inner = d_inner - self.d_conv = d_conv - self.pre_stride = pre_stride - self.post_stride = post_stride - self.dtype = dtype - self.weight = Parameter(shape=(self.d_inner, 1, self.d_conv, 1), - dtype=dtype) - self.bias = Parameter(shape=(self.d_inner, ), dtype=dtype) - self.apply_silu = apply_silu - - def forward(self, - x: Tensor, - conv_state: Tensor, - host_request_types: Tensor, - last_token_ids: Tensor, - host_context_lengths: Optional[Tensor] = None, - slot_mapping: Optional[Tensor] = None, - conv_indices: Optional[Tensor] = None): - ''' - Parameters: - x: [B, L, D] or [T, D] - conv_state: [B, W, D] or [1] of type int64 for paged state - host_request_types: [B] - last_token_ids: [B] - host_context_lengths: [B] - slot_mapping: [B] - conv_indices: [B] - ''' - if default_net().plugin_config.mamba_conv1d_plugin: - transposed_weight = permute( - view(self.weight.value, shape=[self.d_inner, 1, self.d_conv]), - (1, 2, 0)) - x_conv, conv_state = mamba_conv1d( - x, conv_state, transposed_weight, self.bias.value, - host_request_types, last_token_ids, self.d_inner, self.d_conv, - self.dtype, self.pre_stride, self.post_stride, - host_context_lengths, slot_mapping, self.apply_silu) - else: - assert not default_net().plugin_config.paged_state - assert len( - x.shape - ) == 3, "remove_input_padding is not supported by OOTB for Mamba." - if self.pre_stride > 0: - _, x = split(x, - [self.pre_stride, self.d_inner + self.post_stride], - dim=-1) - if self.post_stride > 0: - x, _ = split(x, [self.d_inner, self.post_stride], dim=-1) - x = x.permute([0, 2, 1]) - - # In context phase, conv_state is a zero tensor, and it is used for padding - # In generation phase, conv_state is a tensor of the past x - x_pad = concat([conv_state, x], dim=2) - - # Update conv_state - conv_state = gather(x_pad, 2, conv_indices) - - # Convolution - x_pad = x_pad.view( - concat([shape(x_pad, 0), - shape(x_pad, 1), - shape(x_pad, 2), 1])) - x_conv = conv2d(x_pad, - self.weight.value, - self.bias.value, - groups=self.d_inner) - if self.apply_silu: - x_conv = ACT2FN['silu'](x_conv) - x_conv = x_conv.view( - concat([shape(x_conv, 0), - shape(x_conv, 1), - shape(x_conv, 2)])) - - # Get dt, B and C - x_conv = x_conv.permute([0, 2, 1]) - return x_conv, conv_state - - -class Mamba(Module): - - def __init__(self, - d_model, - d_inner, - d_state=16, - d_conv=4, - dt_rank="auto", - bias=False, - dtype=None): - super().__init__() - self.d_model = d_model - self.d_state = d_state - self.d_conv = d_conv - self.d_inner = d_inner - self.dt_rank = math.ceil(self.d_model / - 16) if dt_rank == "auto" else dt_rank - self.dtype = dtype - - self.A = Parameter(shape=(self.d_state, self.d_inner), dtype="float32") - - self.D = Parameter(shape=(self.d_inner, ), dtype="float32") - self.dt_bias = Parameter(shape=(self.d_inner, ), dtype="float32") - - self.in_proj_x = Linear(self.d_model, - self.d_inner, - bias=bias, - dtype=dtype, - gather_output=False) - self.in_proj_z = Linear(self.d_model, - self.d_inner, - bias=bias, - dtype=dtype, - gather_output=False) - - self.conv1d = MambaConv1d(self.d_inner, self.d_conv, dtype=self.dtype) - - self.x_proj = Linear(self.d_inner, - self.dt_rank + self.d_state * 2, - bias=False, - dtype=dtype, - gather_output=False) - - self.dt_proj = Linear(self.dt_rank, - self.d_inner, - bias=False, - dtype=dtype, - gather_output=False, - pad_lda=self.d_state * 2) - - self.out_proj = Linear(self.d_inner, - self.d_model, - bias=bias, - dtype=dtype, - gather_output=False) - - def forward(self, - hidden_states: Tensor, - conv_state: Tensor, - ssm_state: Tensor, - host_request_types: Tensor, - last_token_ids: Tensor, - host_context_lengths: Optional[Tensor] = None, - slot_mapping: Optional[Tensor] = None, - conv_indices: Optional[Tensor] = None): - ''' - Parameters: - hidden_states: [B, L, D] or [T, D] - conv_state: [B, W, D] or [1] of type int64 for paged state - ssm_state: [B, N, D] or [1] of type int64 for paged state - host_request_types: [B] - last_token_ids: [B] - host_context_lengths: [B] - slot_mapping: [B] - conv_indices: [B] - ''' - # in_proj - x = self.in_proj_x(hidden_states) - z = self.in_proj_z(hidden_states) - - x_conv, conv_state = self.conv1d(x, conv_state, host_request_types, - last_token_ids, host_context_lengths, - slot_mapping, conv_indices) - - # Get dt, B and C - x_dbl = self.x_proj(x_conv) - if default_net().plugin_config.gemm_plugin: - dt = self.dt_proj(x_dbl) - else: - dt, _ = split(x_dbl, [self.dt_rank, self.d_state * 2], dim=-1) - dt = self.dt_proj(dt) - - # selective scan - y, ssm_state = selective_scan(x_conv, - ssm_state, - dt, - self.dt_bias.value, - self.A.value, - x_dbl, - self.D.value, - host_request_types, - last_token_ids, - self.d_inner, - self.d_state, - self.dt_rank, - delta_softplus=True, - dtype=self.dtype, - z=z, - host_context_lengths=host_context_lengths, - slot_mapping=slot_mapping) - # out_proj - out = self.out_proj(y) - return out, conv_state, ssm_state - - -class Mamba2(Module): - - def __init__(self, - d_model, - d_inner, - d_state=16, - d_conv=4, - headdim=64, - ngroups=1, - chunk_size=256, - bias=False, - rmsnorm=True, - dtype=None, - tp_group=None, - tp_size=1): - super().__init__() - self.d_model = d_model - self.d_state = d_state - self.d_conv = d_conv - assert d_inner % tp_size == 0 - self.d_inner = d_inner // tp_size - self.headdim = headdim - assert ngroups % tp_size == 0 - self.ngroups = ngroups // tp_size - self.chunk_size = chunk_size - self.rmsnorm = rmsnorm - self.dtype = dtype - assert d_inner % headdim == 0 - nheads = d_inner // headdim - assert nheads % tp_size == 0 - self.nheads = nheads // tp_size - # conv1d needs alignment to 8 fp16s - self.pad_ldc = (self.nheads + 7) // 8 * 8 - self.nheads - pad_ldc = self.pad_ldc * tp_size - - self.A = Parameter(shape=(self.nheads, ), dtype="float32") - self.D = Parameter(shape=(self.nheads, ), dtype="float32") - self.dt_bias = Parameter(shape=(self.nheads, ), dtype="float32") - - d_in_proj = 2 * d_inner + 2 * ngroups * d_state + nheads - self.in_proj = ColumnLinear(d_model, - d_in_proj, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - gather_output=False, - pad_ldc=pad_ldc) - - self.conv_dim = (d_inner + 2 * ngroups * d_state) // tp_size - self.conv1d = MambaConv1d(self.conv_dim, - self.d_conv, - pre_stride=self.d_inner, - post_stride=self.nheads + self.pad_ldc, - dtype=self.dtype) - - if rmsnorm: - self.norm = RmsNorm(normalized_shape=self.d_inner, - num_groups=self.ngroups, - eps=1e-5, - dtype=dtype) - - self.out_proj = RowLinear(d_inner, - d_model, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size) - - def forward(self, - hidden_states: Tensor, - conv_state: Tensor, - ssm_state: Tensor, - host_request_types: Tensor, - last_token_ids: Tensor, - host_context_lengths: Optional[Tensor] = None, - slot_mapping: Optional[Tensor] = None, - conv_indices: Optional[Tensor] = None): - ''' - Parameters: - hidden_states: [B, L, D] or [T, D] - conv_state: [B, W, D_conv] or [1] of type int64 for paged state - ssm_state: [B, H, N, D] or [1] of type int64 for paged state - host_request_types: [B] - last_token_ids: [B] - host_context_lengths: [B] - slot_mapping: [B] - conv_indices: [B] - ''' - # in_proj - zxbcdt = self.in_proj(hidden_states) - - # conv1d - xbc_conv, conv_state = self.conv1d(zxbcdt, conv_state, - host_request_types, last_token_ids, - host_context_lengths, slot_mapping, - conv_indices) - - # mamba scan - y, ssm_state = selective_scan(xbc_conv, - ssm_state, - zxbcdt, - self.dt_bias.value, - self.A.value, - xbc_conv, - self.D.value, - host_request_types, - last_token_ids, - self.d_inner, - self.d_state, - dt_rank=0, - delta_softplus=True, - dtype=self.dtype, - z=zxbcdt, - host_context_lengths=host_context_lengths, - slot_mapping=slot_mapping, - nheads=self.nheads, - ngroups=self.ngroups, - chunk_size=self.chunk_size, - mamba_version='Mamba2') - - # norm - if self.rmsnorm: - y = self.norm(y) - - # out_proj - out = self.out_proj(y) - return out, conv_state, ssm_state diff --git a/tensorrt_llm/llmapi/__init__.py b/tensorrt_llm/llmapi/__init__.py index 1682eb51e71c..019b97913d38 100644 --- a/tensorrt_llm/llmapi/__init__.py +++ b/tensorrt_llm/llmapi/__init__.py @@ -4,7 +4,6 @@ from ..executor import CompletionOutput, LoRARequest, RequestError from ..sampling_params import GuidedDecodingParams, SamplingParams from ..scheduling_params import SchedulingParams -from .build_cache import BuildCacheConfig from .llm import LLM, RequestOutput # yapf: disable from .llm_args import (AttentionDpConfig, AutoDecodingConfig, BatchingType, @@ -24,9 +23,8 @@ SADecodingConfig, SAEnhancerConfig, SaveHiddenStatesDecodingConfig, SchedulerConfig, SkipSoftmaxAttentionConfig, TorchCompileConfig, - TorchLlmArgs, TrtLlmArgs, UserProvidedDecodingConfig) -from .llm_utils import (BuildConfig, KvCacheRetentionConfig, QuantAlgo, - QuantConfig) + TorchLlmArgs, UserProvidedDecodingConfig) +from .llm_utils import KvCacheRetentionConfig, QuantAlgo, QuantConfig from .mm_encoder import MultimodalEncoder from .mpi_session import MpiCommSession from .thinking_budget import (ThinkingBudgetLogitsProcessor, @@ -56,11 +54,9 @@ 'MTPDecodingConfig', 'SchedulerConfig', 'CapacitySchedulerPolicy', - 'BuildConfig', 'QuantConfig', 'QuantAlgo', 'CalibConfig', - 'BuildCacheConfig', 'RequestError', 'MpiCommSession', 'ExtendedRuntimePerfKnobConfig', @@ -78,7 +74,6 @@ 'DraftTargetDecodingConfig', 'LlmArgs', 'TorchLlmArgs', - 'TrtLlmArgs', 'AutoDecodingConfig', 'AttentionDpConfig', 'LoRARequest', diff --git a/tensorrt_llm/llmapi/build_cache.py b/tensorrt_llm/llmapi/build_cache.py deleted file mode 100644 index e689145863c9..000000000000 --- a/tensorrt_llm/llmapi/build_cache.py +++ /dev/null @@ -1,312 +0,0 @@ -import contextlib -import datetime -import enum -import hashlib -import json -import os -import shutil -from dataclasses import dataclass -from pathlib import Path -from typing import Any, List, Optional - -import filelock -from pydantic import Field, model_validator - -import tensorrt_llm -from tensorrt_llm.builder import BuildConfig -from tensorrt_llm.llmapi.utils import (StrictBaseModel, enable_llm_debug, - print_colored) -from tensorrt_llm.logger import logger - - -def get_build_cache_config_from_env() -> tuple[bool, str]: - """ - Get the build cache configuration from the environment variables - """ - build_cache_enabled = os.environ.get('TLLM_LLMAPI_BUILD_CACHE') == '1' - build_cache_root = os.environ.get( - 'TLLM_LLMAPI_BUILD_CACHE_ROOT', - '/tmp/.cache/tensorrt_llm/llmapi/') # nosec B108 - return build_cache_enabled, build_cache_root - - -class BuildCacheConfig(StrictBaseModel): - """ - Configuration for the build cache. - - Note: - The build-cache assumes the weights of the model are not changed during the execution. If the weights are - changed, you should remove the caches manually. - """ - - cache_root: Optional[Path] = Field( - default=None, - description= - "The root directory for the build cache. Falls back to env var if not provided." - ) - max_records: int = Field( - default=10, - gt=0, - description="The maximum number of records to store in the cache.") - max_cache_storage_gb: float = Field( - default=256.0, - description="The maximum amount of storage (in GB) to use for the cache." - ) - - @model_validator(mode="after") - def set_default_cache_root(self) -> "BuildCacheConfig": - """Set cache_root from environment variable if not provided.""" - if self.cache_root is None: - _, default_cache_root = get_build_cache_config_from_env() - self.cache_root = Path(default_cache_root) - return self - - -class BuildCache: - """ - The BuildCache class is a class that manages the intermediate products from the build steps. - - NOTE: currently, only engine-building is supported - TODO[chunweiy]: add support for other build steps, such as quantization, convert_checkpoint, etc. - """ - # The version of the cache, will be used to determine if the cache is compatible - CACHE_VERSION = 0 - - def __init__(self, config: Optional[BuildCacheConfig] = None): - config = config or BuildCacheConfig() - self.cache_root = config.cache_root - self.max_records = config.max_records - self.max_cache_storage_gb = config.max_cache_storage_gb - - def free_storage_in_gb(self) -> float: - ''' Get the free storage capacity of the cache. ''' - # measure the root directory - if self.cache_root.parent.exists(): - usage = shutil.disk_usage(self.cache_root.parent) - return usage.free / 1024**3 - return 0 - - def get_engine_building_cache_stage(self, - build_config: BuildConfig, - model_path: Optional[Path] = None, - force_rebuild: bool = False, - **kwargs) -> 'CachedStage': - ''' - Get the build step for engine building. - ''' - build_config_str = json.dumps(self.prune_build_config_for_cache_key( - build_config.model_dump(mode="json")), - sort_keys=True) - - kwargs_str = json.dumps(kwargs, sort_keys=True) - - return CachedStage(parent=self, - kind=CacheRecord.Kind.Engine, - cache_root=self.cache_root, - force_rebuild=force_rebuild, - inputs=[build_config_str, model_path, kwargs_str]) - - def prune_caches(self, has_incoming_record: bool = False): - ''' - Clean up the cache records to make sure the cache size is within the limit - - Args: - has_incoming_record (bool): If the cache has incoming record, the existing records will be further pruned to - reserve space for the incoming record - ''' - if not self.cache_root.exists(): - return - self._clean_up_cache_dir() - records = [] - for dir in self.cache_root.iterdir(): - records.append(self._load_cache_record(dir)) - records.sort(key=lambda x: x.time, reverse=True) - max_records = self.max_records - 1 if has_incoming_record else self.max_records - # prune the cache to meet max_records and max_cache_storage_gb limitation - while len(records) > max_records or sum( - r.storage_gb for r in records) > self.max_cache_storage_gb: - record = records.pop() - # remove the directory and its content - shutil.rmtree(record.path) - - @staticmethod - def prune_build_config_for_cache_key(build_config: dict) -> dict: - black_list = ['dry_run'] - dic = build_config.copy() - for key in black_list: - if key in dic: - dic.pop(key) - return dic - - def load_cache_records(self) -> List["CacheRecord"]: - ''' - Load all the cache records from the cache directory - ''' - records = [] - if not self.cache_root.exists(): - return records - - for dir in self.cache_root.iterdir(): - records.append(self._load_cache_record(dir)) - return records - - def _load_cache_record(self, cache_dir) -> "CacheRecord": - ''' - Get the cache record from the cache directory - ''' - metadata = json.loads((cache_dir / 'metadata.json').read_text()) - storage_gb = sum(f.stat().st_size for f in cache_dir.glob('**/*') - if f.is_file()) / 1024**3 - return CacheRecord(kind=CacheRecord.Kind.__members__[metadata['kind']], - storage_gb=storage_gb, - path=cache_dir, - time=datetime.datetime.fromisoformat( - metadata['datetime'])) - - def _clean_up_cache_dir(self): - ''' - Clean up the files in the cache directory, remove anything that is not in the cache - ''' - # get all the files and directies in the cache_root - if not self.cache_root.exists(): - return - for file_or_dir in self.cache_root.iterdir(): - if not self.is_cache_valid(file_or_dir): - logger.info(f"Removing invalid cache directory {dir}") - if file_or_dir.is_file(): - file_or_dir.unlink() - else: - shutil.rmtree(file_or_dir) - - def is_cache_valid(self, cache_dir: Path) -> bool: - ''' - Check if the cache directory is valid - ''' - if not cache_dir.exists(): - return False - - metadata_path = cache_dir / 'metadata.json' - if not metadata_path.exists(): - return False - - metadata = json.loads(metadata_path.read_text()) - if metadata.get('version') != BuildCache.CACHE_VERSION: - return False - - content = cache_dir / 'content' - if not content.exists(): - return False - - return True - - -@dataclass -class CachedStage: - ''' - CachedStage is a class that represents a stage in the build process, it helps to manage the intermediate product. - - The cache is organized as follows: - - this_cache_dir/ # name is like "engine-" - metadata.json # the metadata of the cache - content/ # the actual product of the build step, such trt-llm engine directory - ''' - # The parent should be kept alive by CachedStep instance - parent: BuildCache - cache_root: Path - # The inputs will be used to determine if the step needs to be re-run, so all the variables should be put here - inputs: List[Any] - kind: "CacheRecord.Kind" - # If force_rebuild is set to True, the cache will be ignored - force_rebuild: bool = False - - def get_hash_key(self): - lib_version = tensorrt_llm.__version__ - input_strs = [str(i) for i in self.inputs] - return hashlib.md5( - f"{lib_version}-{input_strs}".encode()).hexdigest() # nosec B324 - - def get_cache_path(self) -> Path: - ''' - The path to the product of the build step, will be overwritten if the step is re-run - ''' - return self.cache_root / f"{self.kind.value}-{self.get_hash_key()}" - - def get_engine_path(self) -> Path: - return self.get_cache_path() / 'content' - - def get_cache_metadata(self) -> dict: - res = { - "version": BuildCache.CACHE_VERSION, - "datetime": datetime.datetime.now().isoformat(), - "kind": self.kind.name, - } - return res - - def is_cached(self) -> bool: - ''' - Check if the product of the build step is in the cache - ''' - if self.force_rebuild: - return False - try: - if self.get_cache_path().exists(): - metadata = json.loads( - (self.get_cache_path() / 'metadata.json').read_text()) - if metadata["version"] == BuildCache.CACHE_VERSION: - return True - except: - pass - - return False - - @contextlib.contextmanager - def write_guard(self): - ''' Guard the cache writing process. - - The cache writing process should be atomic, so the filelock is used to protect the cache writing process. And - the cache metadata will be written to the cache directory. - - Args: - final_engien_dir: the final engine directory - ''' - self.parent.prune_caches(has_incoming_record=True) - - target_dir = self.get_cache_path() - if enable_llm_debug(): - print_colored(f"Writing cache to {target_dir}\n", "yellow") - - # To avoid the cache modification conflict, a dummy directory is used to write the cache, and then rename it to - # the target directory - dummy_target_dir = Path(f"{target_dir.parent}/{target_dir.name}.dummy") - - dummy_target_dir.mkdir(parents=True, exist_ok=True) - # TODO[chunweiy]: deal with the cache modification conflict - lock = filelock.FileLock(dummy_target_dir / '.filelock', timeout=10) - - with open(dummy_target_dir / 'metadata.json', 'w') as f: - f.write(json.dumps(self.get_cache_metadata())) - - with lock: - yield dummy_target_dir / 'content' - - # If engine building is successful, rename the dummy directory to the target directory - if target_dir.exists(): - shutil.rmtree(target_dir) - shutil.move(dummy_target_dir, target_dir) - - -@dataclass(unsafe_hash=True) -class CacheRecord: - ''' - CacheRecord is a class that represents a record in the cache directory. - ''' - - class Kind(enum.Enum): - Engine = 'engine' - Checkpoint = 'checkpoint' - - kind: Kind - storage_gb: float - path: Path - time: datetime.datetime diff --git a/tensorrt_llm/llmapi/llm.py b/tensorrt_llm/llmapi/llm.py index a8d4b5b3eada..b1ef5c01af67 100644 --- a/tensorrt_llm/llmapi/llm.py +++ b/tensorrt_llm/llmapi/llm.py @@ -1,9 +1,7 @@ import atexit import json import os -import shutil import socket -import tempfile import time import weakref from collections.abc import Mapping @@ -24,9 +22,7 @@ from tensorrt_llm.metrics.enums import MetricNames from .._utils import nvtx_range_debug -from ..bindings import executor as tllm from ..bindings import steady_clock_now -from ..builder import EngineConfig from ..conversation_params import ConversationParams from ..disaggregated_params import DisaggregatedParams from ..executor import (DetokenizedGenerationResultBase, GenerationExecutor, @@ -44,14 +40,12 @@ from ..logger import logger from ..sampling_params import LogitsProcessor, SamplingParams from ..scheduling_params import SchedulingParams -from .llm_args import (TORCH_LLMARGS_EXPLICIT_DOCSTRING, - TRT_LLMARGS_EXPLICIT_DOCSTRING, PeftCacheConfig, - PybindMirror, TorchLlmArgs, TrtLlmArgs) +from .llm_args import TORCH_LLMARGS_EXPLICIT_DOCSTRING, TorchLlmArgs from .llm_utils import (CachedModelLoader, KvCacheRetentionConfig, - LlmBuildStats, ModelLoader, _ModelRuntimeContext) + LlmBuildStats, ModelLoader) from .mpi_session import MpiPoolSession, external_mpi_comm_available from .thinking_budget import add_thinking_budget_logits_processor -from .tokenizer import TokenizerBase, _xgrammar_tokenizer_info +from .tokenizer import TokenizerBase # TODO[chunweiy]: move the following symbols back to utils scope, and remove the following import from .utils import (append_docstring, exception_handler, get_device_count, logger_debug, set_api_status) @@ -203,15 +197,6 @@ def _contains_bart_forced_tokens_logits_processor(processor: Any) -> bool: return False -TRT_LLM_DOCSTRING = TRT_LLMARGS_EXPLICIT_DOCSTRING + """ - - Attributes: - tokenizer (tensorrt_llm.llmapi.tokenizer.TokenizerBase, optional): The tokenizer loaded by LLM instance, if any. - workspace (pathlib.Path): The directory to store intermediate files. - llm_id (str): The unique ID of the LLM instance. - disaggregated_params (dict): The disaggregated parameters of the LLM instance. -""" - TORCH_LLM_DOCSTRING = TORCH_LLMARGS_EXPLICIT_DOCSTRING + """ Attributes: @@ -279,8 +264,9 @@ def __init__(self, LlmArgs as AutoDeployLlmArgs llm_args_cls = AutoDeployLlmArgs else: - logger.info("Using LLM with TensorRT backend") - llm_args_cls = TrtLlmArgs + raise ValueError( + f"Unknown backend: {backend!r}. Supported backends are " + "'pytorch' and '_autodeploy'.") # check the kwargs and raise ValueError directly valid_keys = set( @@ -356,17 +342,12 @@ def __init__(self, self._executor: Optional[GenerationExecutor] = None self._encode_only: bool = False self._encoder_executor = None - if self._on_trt_backend: - self._workspace = tempfile.TemporaryDirectory( - suffix="-llm-workspace", dir=self.args.workspace) - else: - self._workspace = None + self._workspace = None self._hf_model_dir: Optional[Path] = None self._hf_model_config = None self._generation_config = None - self.runtime_context: Optional[_ModelRuntimeContext] = None self.llm_build_stats = LlmBuildStats() self._build_model() @@ -652,11 +633,10 @@ def generate_async( # With pytorch backend, py_executor has logic to handle max_tokens of 1, # so set to 1 to avoid allocating unnecessary KV cache blocks for single request - # TODO: Also support for trt backend is_ctx_only = disaggregated_params is not None and disaggregated_params.request_type == "context_only" is_gen_only = disaggregated_params is not None and disaggregated_params.request_type == "generation_only" - if is_ctx_only and not self._on_trt_backend: + if is_ctx_only: sampling_params.max_tokens = 1 if isinstance(inputs, PreprocessedInputs): @@ -1434,10 +1414,6 @@ def _build_model(self): self.llm_build_stats)) self._engine_dir, self._hf_model_dir = model_loader() - @property - def _on_trt_backend(self) -> bool: - return isinstance(self.args, TrtLlmArgs) - def _try_load_tokenizer(self) -> Optional[TokenizerBase]: if self.args.skip_tokenizer_init: return None @@ -1446,9 +1422,6 @@ def _try_load_tokenizer(self) -> Optional[TokenizerBase]: assert isinstance(self.args.tokenizer, TokenizerBase) return self.args.tokenizer - if self.runtime_context is not None: - return self.runtime_context.tokenizer - # TODO smor- need to refine what is the desired behavior if lora is enabled # in terms of the tokenizer initialization process if hasattr(self.args, "backend") and self.args.backend in [ @@ -1552,199 +1525,6 @@ def __del__(self): self.shutdown() -@append_docstring(TRT_LLM_DOCSTRING) -class _TrtLLM(BaseLLM): - """LLM class is the main class for running a LLM model using TensorRT LLM backend. - - Parameters: - """ - - def __init__(self, - model: Union[str, Path], - tokenizer: Optional[Union[str, Path, TokenizerBase, - PreTrainedTokenizerBase]] = None, - tokenizer_mode: Literal['auto', 'slow'] = 'auto', - skip_tokenizer_init: bool = False, - trust_remote_code: bool = False, - tensor_parallel_size: int = 1, - dtype: str = "auto", - revision: Optional[str] = None, - tokenizer_revision: Optional[str] = None, - **kwargs: Any) -> None: - super().__init__(model, tokenizer, tokenizer_mode, skip_tokenizer_init, - trust_remote_code, tensor_parallel_size, dtype, - revision, tokenizer_revision, **kwargs) - - @property - def workspace(self) -> Path: - return Path(self._workspace.name) if self._on_trt_backend else None - - def save(self, engine_dir: str) -> None: - """Save the built engine to the given path. - - Args: - engine_dir (str): The path to save the engine. - """ - logger.info(f"Save model to {engine_dir}") - if self._engine_dir is None: - raise RuntimeError("The engine is not built yet.") - - if self._engine_dir.absolute() == os.path.abspath(engine_dir): - return - - if not self.mpi_session or not self.mpi_session.is_comm_session(): - shutil.copytree(self._engine_dir, engine_dir, dirs_exist_ok=True) - else: - # NFS is fragile, so we copy files one by one - target_engine_dir = Path(engine_dir) - target_engine_dir.mkdir(parents=True, exist_ok=True) - # copy files one by one - for file in self._engine_dir.iterdir(): - logger_debug( - f"Copying {file} to {target_engine_dir / file.name}\n") - shutil.copy(file, target_engine_dir / file.name) - - def _build_model(self): - super()._build_model() - # update the model_dir to a local dir for the runtime, such as tokenizer loading. - if self._engine_dir is not None: - self.args.model = self._engine_dir - - # Tokenizer and config loading should be after calling model_loader(), since model_loader() may download the model from HF hub. - # It should also be before bindings ExecutorConfig, which may depend on tokenizer info. - self._tokenizer = self._try_load_tokenizer() - # Load HF config from the original HF model dir when available, - # since self.args.model now points to the engine dir (whose - # config.json uses TRT-LLM schema, not HF schema). - if self._hf_model_dir is not None: - self._hf_model_config = ModelLoader.load_hf_model_config( - self._hf_model_dir, - trust_remote_code=self.args.trust_remote_code) - else: - self._hf_model_config = self._try_load_hf_model_config() - self._generation_config = self._try_load_generation_config() - - # Multimodal special handling: - # 1. Default load_tokenizer may fail because MM has different tokenizer configuration. Hence we initialize it inside input processor - # 2. May need to modify model weights for MM (e.g., resize vocab embedding). We must do such operation via input processor's __init__ - self.input_processor = create_input_processor( - self._hf_model_dir, - self.tokenizer, - trust_remote_code=self.args.trust_remote_code) - self._tokenizer = self.input_processor.tokenizer - - max_batch_size = self.args.max_batch_size - max_num_tokens = self.args.max_num_tokens - max_seq_len = self.args.max_seq_len - - build_config = self.args.build_config - - max_batch_size = max_batch_size or build_config.max_batch_size - max_num_tokens = max_num_tokens or build_config.max_num_tokens - max_seq_len = max_seq_len or build_config.max_seq_len - - self._executor_config = tllm.ExecutorConfig( - max_beam_width=self.args.max_beam_width, - scheduler_config=PybindMirror.maybe_to_pybind( - self.args.scheduler_config), - batching_type=PybindMirror.maybe_to_pybind(self.args.batching_type) - or tllm.BatchingType.INFLIGHT, - max_batch_size=max_batch_size, - max_num_tokens=max_num_tokens, - gather_generation_logits=self.args.gather_generation_logits, - fail_fast_on_attention_window_too_large=getattr( - self.args, 'fail_fast_on_attention_window_too_large', False)) - - # also set executor_config.max_seq_len in TRT workflow, to deduce default max_tokens - if max_seq_len is not None: - self._executor_config.max_seq_len = max_seq_len - else: - engine_config = EngineConfig.from_json_file(self._engine_dir / - "config.json") - self._executor_config.max_seq_len = engine_config.build_config.max_seq_len - - if self.args.kv_cache_config is not None: - self._executor_config.kv_cache_config = PybindMirror.maybe_to_pybind( - self.args.kv_cache_config) - if os.getenv("FORCE_DETERMINISTIC", "0") == "1": - # Disable KV cache reuse for deterministic mode - self._executor_config.kv_cache_config.enable_block_reuse = False - self._executor_config.kv_cache_config.enable_partial_reuse = False - if self.args.peft_cache_config is not None: - self._executor_config.peft_cache_config = PybindMirror.maybe_to_pybind( - self.args.peft_cache_config) - - lora_config = None - if self.args.build_config.plugin_config.lora_plugin: - engine_config = EngineConfig.from_json_file(self._engine_dir / - "config.json") - lora_config = engine_config.build_config.lora_config - if self.args.lora_config is not None: - logger.info( - "Overriding lora_config from engine with lora_config from LLM args" - ) - lora_config = self.args.lora_config - - max_lora_rank = lora_config.max_lora_rank - num_lora_modules = engine_config.pretrained_config.num_hidden_layers * \ - len(lora_config.lora_target_modules + lora_config.missing_qkv_modules) - - peft_cache_config_model = PeftCacheConfig.from_pybind( - self._executor_config.peft_cache_config - ) if self._executor_config.peft_cache_config is not None else PeftCacheConfig( - ) - if lora_config.max_loras is not None: - peft_cache_config_model.num_device_module_layer = \ - max_lora_rank * num_lora_modules * lora_config.max_loras - if lora_config.max_cpu_loras is not None: - peft_cache_config_model.num_host_module_layer = \ - max_lora_rank * num_lora_modules * lora_config.max_cpu_loras - self._executor_config.peft_cache_config = peft_cache_config_model._to_pybind( - ) - - if self.args.decoding_config is not None: - self._executor_config.decoding_config = self.args.decoding_config - if self.args.guided_decoding_backend == 'xgrammar': - self._executor_config.guided_decoding_config = tllm.GuidedDecodingConfig( - backend=tllm.GuidedDecodingConfig.GuidedDecodingBackend. - XGRAMMAR, - **_xgrammar_tokenizer_info(self.tokenizer)) - elif self.args.guided_decoding_backend is not None: - raise ValueError( - f"Unsupported guided decoding backend {self.args.guided_decoding_backend}" - ) - - self._executor_config.normalize_log_probs = self.args.normalize_log_probs - self._executor_config.enable_chunked_context = self.args.enable_chunked_prefill - self._executor_config.max_beam_width = self.args.max_beam_width or self.args.build_config.max_beam_width - if self.args.extended_runtime_perf_knob_config is not None: - self._executor_config.extended_runtime_perf_knob_config = PybindMirror.maybe_to_pybind( - self.args.extended_runtime_perf_knob_config) - if self.args.cache_transceiver_config is not None: - self._executor_config.cache_transceiver_config = PybindMirror.maybe_to_pybind( - self.args.cache_transceiver_config) - self._executor_config.llm_parallel_config = self.args.parallel_config - return_logits = (self.args.gather_generation_logits - or (self.args.build_config - and self.args.build_config.gather_context_logits)) - - self._executor = self._executor_cls.create( - self._engine_dir, - executor_config=self._executor_config, - batched_logits_processor=self.args.batched_logits_processor, - model_world_size=self.args.parallel_config.world_size, - mpi_session=self.mpi_session, - reuse_mpi_comm=external_mpi_comm_available( - self.args.parallel_config.world_size), - return_logits=return_logits, - postproc_worker_config=PostprocWorkerConfig( - num_postprocess_workers=self.args.num_postprocess_workers, - postprocess_tokenizer_dir=self.args.postprocess_tokenizer_dir, - post_processor_hook=self.args.post_processor_hook, - ), - is_llm_executor=True) - - @append_docstring(TORCH_LLM_DOCSTRING) class _TorchLLM(BaseLLM): """LLM class is the main class for running a LLM model using PyTorch backend. @@ -1767,7 +1547,7 @@ def __init__(self, backend = kwargs.pop("backend", "pytorch") - # Validate that users don't pass TrtLlmArgs-specific arguments + # Validate that only arguments supported by the PyTorch backend are passed. self._validate_args_for_torch_backend(kwargs) super().__init__(model, @@ -1897,23 +1677,20 @@ def _build_model(self): llm_args=self.args) def _validate_args_for_torch_backend(self, kwargs: dict) -> None: - """Validate that users don't pass TrtLlmArgs-specific arguments when using PyTorch backend. + """Validate that only arguments supported by the PyTorch backend are passed. """ - trtllm_fields = set(TrtLlmArgs.model_fields.keys()) torchllm_fields = set(TorchLlmArgs.model_fields.keys()) - trtllm_specific_fields = trtllm_fields - torchllm_fields - - # Check if any TrtLlmArgs-specific arguments are passed - trtllm_specific_args = [] - for key in kwargs: - if key in trtllm_specific_fields: - trtllm_specific_args.append(key) + # Check if any arguments not supported by the PyTorch backend are passed. + unsupported_args = [ + key for key in kwargs + if key not in torchllm_fields and key not in ('_mpi_session', + 'backend') + ] - if trtllm_specific_args: + if unsupported_args: raise ValueError( - f"The following arguments are specific to TensorRT backend and cannot be used with PyTorch backend: {trtllm_specific_args}.\n" - f"Please use 'from tensorrt_llm._tensorrt_engine import LLM' instead to use the TensorRT backend." + f"The following arguments are not supported by the PyTorch backend: {unsupported_args}." ) diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index faaeea64bc1c..cc6a66e53a08 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -70,16 +70,12 @@ # isort: on # yapf: enable -from ..builder import BuildConfig, EngineConfig from ..logger import logger from ..mapping import CpType, Mapping -from ..models.automodel import AutoConfig -from ..models.modeling_utils import (PretrainedConfig, QuantAlgo, QuantConfig, - SpeculativeDecodingMode) +from ..models.modeling_utils import QuantAlgo, QuantConfig from ..sampling_params import BatchedLogitsProcessor from ..usage.config import UsageContext # noqa: F401 from ..usage.config import TelemetryConfig, TelemetryField -from .build_cache import BuildCacheConfig from .tokenizer import TokenizerBase, tokenizer_factory from .utils import (StrictBaseModel, generate_api_docs_as_docstring, get_type_repr) @@ -1559,12 +1555,6 @@ class CalibConfig(StrictBaseModel): "The maximum sequence length to initialize tokenizer for calibration.") -class _ModelFormatKind(Enum): - HF = 0 - TLLM_CKPT = 1 - TLLM_ENGINE = 2 - - class DecodingBaseConfig(StrictBaseModel): max_draft_len: Optional[NonNegativeInt] = Field( default=None, description="The maximum number of draft tokens.") @@ -3806,7 +3796,7 @@ class DwdpConfig(StrictBaseModel): class BaseLlmArgs(StrictBaseModel): - """Base class for both TorchLlmArgs and TrtLlmArgs. It contains all the arguments that are common to both.""" + """Base class for the LLM arguments. It contains all the common arguments.""" model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") # Explicit arguments @@ -4078,8 +4068,7 @@ class BaseLlmArgs(StrictBaseModel): exclude_json_schema=True, # hide from API references validate_default=True, status="deprecated", - telemetry=TelemetryField.categorical('pytorch', 'tensorrt', - '_autodeploy')) + telemetry=TelemetryField.categorical('pytorch', '_autodeploy')) return_perf_metrics: bool = Field(default=False, description="Return perf metrics.", @@ -4130,16 +4119,11 @@ def coerce_env_overrides_to_str(cls, v): return {str(k): str(val) for k, val in v.items()} _parallel_config: Optional[_ParallelConfig] = PrivateAttr(default=None) - _model_format: Optional[_ModelFormatKind] = PrivateAttr(default=None) @property def parallel_config(self) -> _ParallelConfig: return self._parallel_config - @property - def model_format(self) -> _ModelFormatKind: - return self._model_format - @property def speculative_model(self) -> Optional[Union[str, Path]]: return self.speculative_config.speculative_model if self.speculative_config is not None else None @@ -4326,330 +4310,6 @@ def get_runtime_sizes(self, ) -> Tuple[int, int, int, int]: ) -class TrtLlmArgs(BaseLlmArgs): - enable_tqdm: bool = Field(default=False, - description="Enable tqdm for progress bar.") - - workspace: Optional[str] = Field(default=None, - description="The workspace for the model.") - - fail_fast_on_attention_window_too_large: bool = Field( - default=True, - description= - "Fail fast when attention window is too large to fit even a single sequence in the KV cache.", - status="deprecated") - - # Once set, the model will reuse the build_cache - enable_build_cache: Union[BuildCacheConfig, - bool] = Field(default=False, - description="Enable build cache.") - - extended_runtime_perf_knob_config: Optional[ - ExtendedRuntimePerfKnobConfig] = Field( - default=None, description="Extended runtime perf knob config.") - - # Quantization and calibration configurations - calib_config: CalibConfig = Field(default_factory=CalibConfig, - description="Calibration config.") - - quant_config: QuantConfig = Field(default_factory=QuantConfig, - description="Quantization config.") - - embedding_parallel_mode: Literal[ - 'NONE', 'SHARDING_ALONG_VOCAB', 'SHARDING_ALONG_HIDDEN'] = Field( - default='SHARDING_ALONG_VOCAB', - description="The embedding parallel mode.") - - fast_build: bool = Field(default=False, description="Enable fast build.") - - # BuildConfig is introduced to give users a familiar interface to configure the model building. - build_config: Optional[BuildConfig] = Field(default=None, - description="Build config.") - - # Prompt adapter arguments - enable_prompt_adapter: bool = Field(default=False, - description="Enable prompt adapter.") - - max_prompt_adapter_token: int = Field( - default=0, description="The maximum number of prompt adapter tokens.") - - batching_type: Optional[BatchingType] = Field(default=None, - description="Batching type.") - - normalize_log_probs: bool = Field( - default=False, description="Normalize log probabilities.") - - # Private attributes - # This is used to hold the options for convert_checkpoint - _convert_checkpoint_options: Dict[str, - Any] = PrivateAttr(default_factory=dict) - - @model_validator(mode="after") - def init_build_config(self): - """Creating a default BuildConfig if none is provided""" - build_config = getattr(self, "build_config", None) - if build_config is None: - kwargs = {} - if self.max_batch_size: - kwargs["max_batch_size"] = self.max_batch_size - if self.max_num_tokens: - kwargs["max_num_tokens"] = self.max_num_tokens - if self.max_seq_len: - kwargs["max_seq_len"] = self.max_seq_len - if self.max_beam_width: - kwargs["max_beam_width"] = self.max_beam_width - if self.max_input_len: - kwargs["max_input_len"] = self.max_input_len - self.build_config = BuildConfig(**kwargs) - return self - - @model_validator(mode="after") - def validate_build_config_remaining(self): - is_trt_llm_args = isinstance(self, TrtLlmArgs) - - # TODO: remove the checker when manage weights support all data types - if is_trt_llm_args and self.fast_build and (self.quant_config.quant_algo - is QuantAlgo.FP8): - self.build_config.plugin_config.manage_weights = True - - if self.parallel_config.world_size == 1 and self.build_config: - self.build_config.plugin_config.nccl_plugin = None - - if self.enable_lora and self.backend != 'pytorch': - self.build_config.plugin_config.lora_plugin = 'auto' - if self.lora_config is not None: - self.build_config.lora_config.max_lora_rank = self.lora_config.max_lora_rank - - if hasattr(self, - 'enable_prompt_adapter') and self.enable_prompt_adapter: - self.build_config.max_prompt_embedding_table_size = self.max_prompt_adapter_token * self.build_config.max_batch_size - - return self - - @model_validator(mode="after") - def validate_speculative_config(self): - if self.speculative_config: - if not self.speculative_config.supports_backend(self.backend): - raise ValueError( - f"Speculation type {self.speculative_config.decoding_type} does not " - f"support backend {self.backend}") - - # Below, we only need to set speculative_decoding_mode/decoding_config for speculation - # on the TRT backend. - if isinstance(self.speculative_config, LookaheadDecodingConfig): - max_draft_len = self.speculative_config.calculate_speculative_resource( - )[2] - assert max_draft_len > 0 - self.build_config.speculative_decoding_mode = SpeculativeDecodingMode.LOOKAHEAD_DECODING - self.build_config.max_draft_len = max( - self.build_config.max_draft_len, max_draft_len) - self.decoding_config = DecodingConfig( - decoding_mode=DecodingMode.Lookahead(), - lookahead_decoding_config=PybindMirror.maybe_to_pybind( - self.speculative_config)) - - elif isinstance(self.speculative_config, MedusaDecodingConfig): - assert self.speculative_config.max_draft_len > 0 - self.build_config.speculative_decoding_mode = SpeculativeDecodingMode.MEDUSA - self.build_config.max_draft_len = self.speculative_config.max_draft_len - self.decoding_config = DecodingConfig( - decoding_mode=DecodingMode.Medusa(), - medusa_choices=self.speculative_config.medusa_choices) - - elif isinstance(self.speculative_config, Eagle3DecodingConfig): - raise ValueError( - "speculative_config.decoding_type 'Eagle3' is only supported on the PyTorch backend. " - "Use decoding_type 'Eagle' for the TensorRT backend.") - - elif isinstance(self.speculative_config, EagleDecodingConfig): - assert self.speculative_config.max_draft_len > 0 - assert self.speculative_config.speculative_model is not None, "EAGLE draft model must be specified." - self.build_config.max_draft_len = self.speculative_config.max_draft_len - self.build_config.speculative_decoding_mode = SpeculativeDecodingMode.EAGLE - eagle_config = _EagleConfig( - self.speculative_config.eagle_choices, - self.speculative_config.greedy_sampling, - self.speculative_config.posterior_threshold, - self.speculative_config.use_dynamic_tree, - self.speculative_config.dynamic_tree_max_topK) - self.decoding_config = DecodingConfig( - decoding_mode=DecodingMode.Eagle(), - eagle_config=eagle_config) - elif isinstance(self.speculative_config, PARDDecodingConfig): - raise ValueError( - "speculative_config.decoding_type 'PARD' is only supported on the PyTorch backend." - ) - elif isinstance(self.speculative_config, DFlashDecodingConfig): - raise ValueError( - "speculative_config.decoding_type 'DFlash' is only supported on the PyTorch backend." - ) - else: - raise ValueError( - f"Unrecognized speculative config type {type(self.speculative_config)}" - ) - - else: - self.decoding_config = None - - return self - - def _load_config_from_engine(self, engine_dir: Path): - engine_config = EngineConfig.from_json_file(engine_dir / "config.json") - self._pretrained_config = engine_config.pretrained_config - self.build_config = engine_config.build_config - - # load and check parallel_config - mapping = self._pretrained_config.mapping - if self.parallel_config.tp_size not in (1, mapping.tp_size): - raise ValueError( - f"tp_size {self.parallel_config.tp_size} is not consistent with the engine's tp_size {mapping.tp_size}" - ) - if self.parallel_config.pp_size not in (1, mapping.pp_size): - raise ValueError( - f"pp_size {self.parallel_config.pp_size} is not consistent with the engine's pp_size {mapping.pp_size}" - ) - if self.parallel_config.cp_size not in (1, mapping.cp_size): - raise ValueError( - f"cp_size {self.parallel_config.cp_size} is not consistent with the engine's cp_size {mapping.cp_size}" - ) - self._parallel_config = _ParallelConfig( - tp_size=mapping.tp_size, - pp_size=mapping.pp_size, - cp_size=mapping.cp_size, - gpus_per_node=mapping.gpus_per_node, - moe_cluster_size=mapping.moe_cluster_size, - moe_tp_size=mapping.moe_tp_size, - moe_ep_size=mapping.moe_ep_size) - - def _load_config_from_ckpt(self, ckpt_dir: Path): - pretrained_config = PretrainedConfig.from_json_file(ckpt_dir / - "config.json") - tp_size = pretrained_config.mapping.tp_size - pp_size = pretrained_config.mapping.pp_size - cp_size = pretrained_config.mapping.cp_size - moe_cluster_size = pretrained_config.mapping.moe_cluster_size - moe_tp_size = pretrained_config.mapping.moe_tp_size - moe_ep_size = pretrained_config.mapping.moe_ep_size - gpus_per_node = pretrained_config.mapping.gpus_per_node - # load parallel_config - if self.parallel_config.tp_size != 1 and self.parallel_config.tp_size != tp_size: - raise ValueError( - f"tp_size {self.parallel_config.tp_size} is not consistent with the checkpoint's tp_size {tp_size}" - ) - if self.parallel_config.pp_size != 1 and self.parallel_config.pp_size != pp_size: - raise ValueError( - f"pp_size {self.parallel_config.pp_size} is not consistent with the checkpoint's pp_size {pp_size}" - ) - if self.parallel_config.cp_size != 1 and self.parallel_config.cp_size != cp_size: - raise ValueError( - f"cp_size {self.parallel_config.cp_size} is not consistent with the checkpoint's cp_size {cp_size}" - ) - self._parallel_config = _ParallelConfig( - tp_size=tp_size, - pp_size=pp_size, - cp_size=cp_size, - gpus_per_node=gpus_per_node, - moe_cluster_size=moe_cluster_size, - moe_tp_size=moe_tp_size, - moe_ep_size=moe_ep_size) - - @model_validator(mode="after") - def validate_model_format_misc(self): - """Load the model format, and do the following: - - 1. Load the build_config if got an engine. - 2. Load the parallel_config if got a checkpoint. - """ - model_obj = _ModelWrapper(self.model) - - if model_obj.is_local_model and self.backend not in [ - 'pytorch', '_autodeploy' - ]: - # Load parallel_config from the engine. - model_format = get_model_format( - self.model, trust_remote_code=self.trust_remote_code) - - if model_format is _ModelFormatKind.TLLM_ENGINE: - if self.build_config is not None: - logger.warning( - "The build_config is ignored for model format of TLLM_ENGINE." - ) - self._load_config_from_engine(model_obj.model_dir) - runtime_defaults = self._pretrained_config.runtime_defaults - if runtime_defaults: - self.kv_cache_config.fill_empty_fields_from_runtime_defaults( - runtime_defaults) - - # Load parallel_config from the checkpoint. - elif model_format is _ModelFormatKind.TLLM_CKPT: - # We need to create a temporary instance to call _load_config_from_ckpt - self._load_config_from_ckpt(model_obj.model_dir) - else: - model_format = _ModelFormatKind.HF - - # Store the model format in the values - self._model_format = model_format - return self - - @model_validator(mode="after") - def validate_build_config_with_runtime_params(self): - """Sync runtime parameters with build_config limits. - - This validator runs AFTER validate_model_format_misc so that when - loading from an engine, we have the real build_config loaded. - """ - if self.build_config is None: - raise ValueError("build_config is not initialized") - - # These can be lower than build_config limits - for field in ("max_batch_size", "max_num_tokens"): - runtime_val = getattr(self, field) - build_val = getattr(self.build_config, field) - if runtime_val is not None and runtime_val > build_val: - logger.warning( - f"{field} [{runtime_val}] clamped to build_config.{field} [{build_val}]" - ) - setattr(self, field, build_val) - - # These must match build_config exactly - for field in ("max_seq_len", "max_beam_width", "max_input_len"): - runtime_val = getattr(self, field) - build_val = getattr(self.build_config, field) - if runtime_val is not None and runtime_val != build_val: - logger.warning( - f"{field} [{runtime_val}] overridden by build_config.{field} [{build_val}]" - ) - setattr(self, field, build_val) - - return self - - @model_validator(mode="after") - def setup_embedding_parallel_mode(self): - if self.embedding_parallel_mode == 'NONE': - self._convert_checkpoint_options['use_parallel_embedding'] = False - elif self.embedding_parallel_mode == 'SHARDING_ALONG_VOCAB': - self._convert_checkpoint_options['use_parallel_embedding'] = True - self._convert_checkpoint_options['embedding_sharding_dim'] = 0 - elif self.embedding_parallel_mode == 'SHARDING_ALONG_HIDDEN': - self._convert_checkpoint_options['use_parallel_embedding'] = True - self._convert_checkpoint_options['embedding_sharding_dim'] = 1 - # No else clause needed since validation already happened - return self - - @model_validator(mode="after") - def validate_enable_build_cache(self): - if not self.enable_build_cache: - return self - self.enable_build_cache = BuildCacheConfig() if isinstance( - self.enable_build_cache, bool) else self.enable_build_cache - return self - - @model_validator(mode="after") - def validate_kv_cache_dtype(self): - assert self.kv_cache_config.dtype == "auto", "KvCacheConfig.dtype is not supported by the TensorRT backend." - return self - - class LoadFormat(Enum): AUTO = 0 # Initialize all weights randomly. @@ -5263,11 +4923,6 @@ def extra_resource_managers(self) -> Dict[str, object]: def extra_resource_managers(self, value: Dict[str, object]) -> None: self._extra_resource_managers = value - @model_validator(mode="after") - def set_model_format(self): - self._model_format = _ModelFormatKind.HF - return self - @model_validator(mode="after") def validate_encoder_modes(self) -> 'TorchLlmArgs': if self.encode_only and self.mm_encoder_only: @@ -5685,16 +5340,6 @@ def update_llm_args_with_extra_dict( "kv_cache_dtype": "dtype", "enable_block_reuse": "enable_block_reuse", } - # Scalars that live both at the top level of LlmArgs and inside - # `build_config`. The build_config patch propagates the winning source - # to the nested location. - build_config_dual_loc_keys = ( - "max_batch_size", - "max_num_tokens", - "max_beam_width", - "max_seq_len", - ) - explicit_cli_keys = explicit_cli_keys or set() if 'hf_revision' in llm_args_dict: @@ -5768,9 +5413,7 @@ def update_llm_args_with_extra_dict( field_mapping = { "quant_config": QuantConfig, "calib_config": CalibConfig, - "build_config": BuildConfig, "decoding_config": DecodingConfig, - "enable_build_cache": BuildCacheConfig, "lora_config": LoraConfig, "moe_config": MoeConfig, "nvfp4_gemm_config": Nvfp4GemmConfig, @@ -5790,36 +5433,6 @@ def update_llm_args_with_extra_dict( llm_args = llm_args | llm_args_dict - # build_config only works for TensorRT backend, it will be ignored in PyTorch backend - if "build_config" in llm_args: - # Ensure build_config is a BuildConfig object, not a dict - if isinstance(llm_args["build_config"], dict): - llm_args["build_config"] = BuildConfig(**llm_args["build_config"]) - - # Propagate dual-location scalars into build_config: explicit CLI flag - # wins; otherwise YAML's top-level scalar; otherwise leave alone. Warn - # only when the explicit CLI value actually differs from the YAML - # build_config value being replaced (a genuine override). - for key in build_config_dual_loc_keys: - if key in explicit_cli_keys and key in llm_args: - # Warn only on a genuine override of a YAML build_config value; - # otherwise just record where the value came from. - if getattr(llm_args["build_config"], key) != llm_args[key]: - logger.warning( - f"Explicit CLI flag --{key}={llm_args[key]} overrides " - f"the value set in the YAML build_config; CLI takes " - f"precedence.") - else: - logger.info( - f"build_config.{key} set to {llm_args[key]} from explicit CLI flag" - ) - setattr(llm_args["build_config"], key, llm_args[key]) - elif key in llm_args_dict: - setattr(llm_args["build_config"], key, llm_args_dict[key]) - logger.info( - f"build_config.{key} set to {llm_args_dict[key]} from YAML top-level scalar" - ) - return llm_args @@ -5838,40 +5451,8 @@ def update_llm_args_with_extra_options( return llm_args -def get_model_format(model_dir: str, - trust_remote_code: bool = False) -> _ModelFormatKind: - """Get the format of the model.""" - if not (Path(model_dir) / 'config.json').exists(): - raise ValueError( - f"Failed to infer model format because no config.json exists in {model_dir}" - ) - - with open(Path(model_dir) / 'config.json') as f: - config = json.load(f) - - try: - if 'pretrained_config' in config and 'build_config' in config: - model_format = _ModelFormatKind.TLLM_ENGINE - EngineConfig.from_json_file(Path(model_dir) / 'config.json') - elif 'architecture' in config and 'dtype' in config: - model_format = _ModelFormatKind.TLLM_CKPT - PretrainedConfig.from_checkpoint(model_dir) - else: - model_format = _ModelFormatKind.HF - AutoConfig.from_hugging_face(model_dir, - trust_remote_code=trust_remote_code) - except Exception as e: - raise ValueError( - f"Inferred model format {model_format}, but failed to load config.json: {e}" - ) - else: - return model_format - - LlmArgs = TorchLlmArgs -TRT_LLMARGS_EXPLICIT_DOCSTRING = generate_api_docs_as_docstring(TrtLlmArgs, - indent=' ' * 4) TORCH_LLMARGS_EXPLICIT_DOCSTRING = generate_api_docs_as_docstring(TorchLlmArgs, indent=' ' * 4) diff --git a/tensorrt_llm/llmapi/llm_utils.py b/tensorrt_llm/llmapi/llm_utils.py index b162d0bd4443..24288807f92c 100644 --- a/tensorrt_llm/llmapi/llm_utils.py +++ b/tensorrt_llm/llmapi/llm_utils.py @@ -1,39 +1,26 @@ import json import os -import shutil import tempfile -import time import weakref -from dataclasses import asdict, dataclass, field, is_dataclass +from dataclasses import dataclass, field from pathlib import Path from typing import Any, Callable, Dict, List, Optional, Tuple, Union -import torch import transformers -from pydantic import BaseModel -from tqdm import tqdm -from .._utils import (global_mpi_rank, local_mpi_rank, mpi_barrier, - mpi_broadcast, mpi_rank, release_gc) +from .._utils import global_mpi_rank, local_mpi_rank, mpi_rank # yapf: disable from ..bindings.executor import (BatchingType, CapacitySchedulerPolicy, ContextChunkingPolicy, ExecutorConfig, KvCacheRetentionConfig) # yapf: enable -from ..builder import BuildConfig, Engine, build -from ..llmapi.llm_args import TrtLlmArgs from ..logger import logger -from ..mapping import Mapping -from ..models.automodel import MODEL_MAP, AutoConfig, AutoModelForCausalLM -from ..models.modeling_utils import PretrainedConfig, QuantAlgo, QuantConfig +from ..models.modeling_utils import QuantAlgo, QuantConfig from ..models.quant_config_utils import \ update_quant_config_from_compressed_tensors -from ..module import Module from ..quantization.modelopt_config import (is_modelopt_quant_config, read_modelopt_quant_config, warn_if_inline_diverges) -from .build_cache import (BuildCache, BuildCacheConfig, CachedStage, - get_build_cache_config_from_env) # yapf: disable from .llm_args import (CalibConfig, CudaGraphConfig, DecodeCudaGraphConfig, DraftTargetDecodingConfig, Eagle3DecodingConfig, @@ -41,67 +28,14 @@ KvCacheConfig, LlmArgs, LookaheadDecodingConfig, MedusaDecodingConfig, MTPDecodingConfig, NGramDecodingConfig, SchedulerConfig, TorchLlmArgs, - UserProvidedDecodingConfig, _ModelFormatKind, - _ModelWrapper, _ParallelConfig, - update_llm_args_with_extra_dict, + UserProvidedDecodingConfig, _ModelWrapper, + _ParallelConfig, update_llm_args_with_extra_dict, update_llm_args_with_extra_options) # yapf: enable -from .mpi_session import MPINodeState, MpiSession +from .mpi_session import MpiSession from .tokenizer import TransformersTokenizer, load_hf_tokenizer # TODO[chunweiy]: move the following symbols back to utils scope, and remove the following import -from .utils import (download_hf_model, download_hf_pretrained_config, - enable_llm_debug, get_directory_size_in_gb, logger_debug, - print_colored, print_traceback_on_error) - - -@dataclass -class _ModelInfo: - dtype: Optional[str] = None - architecture: Optional[str] = None - - @property - def model_name(self) -> str: - if self.architecture is None: - raise RuntimeError("The architecture is not set yet.") - - return self.architecture - - @classmethod - def from_pretrained_config(cls, config: PretrainedConfig): - return cls(dtype=config.dtype, architecture=config.architecture) - - @classmethod - def from_builder_config_json(cls, config: dict): - if 'version' in config: - # The Dict format is { 'builder_config':..., 'plugin_config':...} - dtype = config['plugin_config']['gpt_attention_plugin'] - else: - dtype = config['pretrained_config']['dtype'] - - return cls(dtype=dtype, architecture=config['builder_config']['name']) - - @classmethod - def from_module(cls, module: Module): - raise NotImplementedError() - - -@dataclass -class _ModelRuntimeContext: - """_ModelRuntimeContext holds the minimum runtime resources for running a model. - - It could be a runtime cache in MPI nodes. - """ - engine: Optional[Engine] = None - mapping: Optional[Mapping] = None - model_info: Optional[_ModelInfo] = None - - # This is only used when build-cache is enabled - engine_path: Optional[str] = None - - @property - def model_arch(self) -> str: - # "LlamaForCausalLM" or "OPTForCausalLM" and so on - return self.engine.config.pretrained_config['architecture'] +from .utils import download_hf_model, print_traceback_on_error class ModelLoader: @@ -123,208 +57,16 @@ def __init__(self, self.llm_args.speculative_model ) if self.llm_args.speculative_model is not None else None - if isinstance(self.llm_args, TrtLlmArgs): - self.convert_checkpoint_options = self.llm_args._convert_checkpoint_options self.rank = mpi_rank() self.global_rank = global_mpi_rank() self.mapping = llm_args.parallel_config.to_mapping() - self._build_pipeline = [] - # For model from hub, the _model_dir is None, and will updated once downloaded self._model_dir: Optional[ Path] = self.model_obj.model_dir if self.model_obj.is_local_model else None self._speculative_model_dir: Optional[ Path] = self.speculative_model_obj.model_dir if self.speculative_model_obj is not None and self.speculative_model_obj.is_local_model else None - self._model_info: Optional[_ModelInfo] = None - self._model_format = self.llm_args.model_format - - if isinstance(self.llm_args, TrtLlmArgs): - assert self.llm_args.build_config - self.build_config = self.llm_args.build_config - - self._gather_build_steps() - - def _gather_build_steps(self): - # Prepare the model processing pipeline - if isinstance(self.llm_args.model, Module): - # Build engine from user provided model - self._build_pipeline.append( - ("Build TensorRT LLM engine", - self._build_engine_from_inmemory_model)) - return - - if (self.model_obj.is_hub_model - and self._model_format is not _ModelFormatKind.TLLM_ENGINE): - # Download HF model if necessary - if self.model_obj.model_name is None: - raise ValueError( - "Either model_dir or model should be provided to ModelConfig." - ) - self._build_pipeline.append( - ("Downloading HF model", self._download_hf_model)) - - if self._model_format is _ModelFormatKind.HF: - # HF -> TRT checkpoints -> engine - self._build_pipeline.append( - ("Loading HF model to memory", self._load_model_from_hf)) - self._build_pipeline.append( - ("Building TRT-LLM engine", self._build_engine)) - elif self._model_format is _ModelFormatKind.TLLM_CKPT: - # TRT checkpoints -> engine - self._build_pipeline.append(("Loading TRT checkpoints to memory", - self._load_model_from_ckpt)) - self._build_pipeline.append( - ("Build TRT-LLM engine", self._build_engine)) - elif self._model_format is _ModelFormatKind.TLLM_ENGINE: - # Nothing need to do - pass - else: - raise ValueError(f"Unknown model format {self._model_format}") - - class BuildPipeline: - - def __init__(self, enable_tqdm: bool, labels: List[str], - step_handlers: List[Callable], - llm_build_stats: "LlmBuildStats"): - assert len(labels) == len(step_handlers) - self.labels = labels - self.step_handlers = step_handlers - self.llm_build_stats = llm_build_stats - - self.to_log = mpi_rank() == 0 - self.counter = 0 - - self.progress_bar = tqdm( - total=len(labels)) if enable_tqdm and self.to_log else None - - def __call__(self): - start_time = time.time() - - for i in range(len(self.labels)): - self.step_forward() - - if self.to_log: - if self.progress_bar: - self.progress_bar.close() - else: - overall_latency = time.time() - start_time - print_colored("Loading model done.\n", 'bold_green') - print_colored( - 'Total latency: {:.3f}s\n'.format(overall_latency), - 'grey') - - def step_forward(self): - n_steps = len(self.labels) - - label = self.labels[self.counter] - - # display step information - if self.to_log: - if self.progress_bar: - self.progress_bar.set_description(self.labels[self.counter]) - else: - print_colored("Loading Model: ") - print_colored(f"[{self.counter+1}/{n_steps}]\t", - 'bold_green') - print_colored(f"{label}\n") - - # execute the step - start_time = time.time() - self.step_handlers[self.counter]() - # release resource after each step - release_gc() - - if self.progress_bar: - self.progress_bar.update(1) - - latency = time.time() - start_time - if self.to_log and not self.progress_bar: - print_colored("Time: {:.3f}s\n".format(latency), 'grey') - - self.llm_build_stats.build_steps_info.append((label, latency)) - - self.counter += 1 - - def __call__(self, engine_dir: Optional[Path] = None) -> Path: - """The engine_dir is the path to save the built engine. - """ - if self.llm_args.model_format is _ModelFormatKind.TLLM_ENGINE: - return self.model_obj.model_dir - - if self.llm_args.parallel_config.is_multi_gpu: - torch.cuda.set_device(self.global_rank % self.mapping.gpus_per_node) - - pipeline = ModelLoader.BuildPipeline( - self.llm_args.enable_tqdm, - [label for label, _ in self._build_pipeline], - [handler for _, handler in self._build_pipeline], - llm_build_stats=self.llm_build_stats, - ) - pipeline() - - assert engine_dir - - runtime_context = _ModelRuntimeContext( - engine=self._engine, - mapping=self.mapping, - model_info=self._model_info, - ) - self.save(runtime_context, self.model_obj.model_dir, engine_dir) - return engine_dir - - def __enter__(self): - return self - - def __exit__(self, exc_type, exc_value, traceback): - for attr_name in dir(self): - if not callable(getattr( - self, attr_name)) and not attr_name.startswith("__"): - if attr_name not in ('model_format', 'workspace'): - setattr(self, attr_name, None) - - release_gc() - - @property - def workspace(self) -> str: - return self._workspace - - @property - def model_format(self) -> _ModelFormatKind: - return self._model_format - - def save( - self, - model: _ModelRuntimeContext, - model_dir: str, - engine_dir: str, - ): - """Save the built engine on a single GPU to the given path.""" - model.engine.save(engine_dir) - if model.mapping.rank == 0: - tokenizer = ModelLoader.load_hf_tokenizer( - model_dir, - trust_remote_code=self.llm_args.trust_remote_code, - use_fast=self.llm_args.tokenizer_mode != 'slow') - if tokenizer is not None: - tokenizer.save_pretrained(engine_dir) - - def _download_hf_model(self): - """Download HF model from third-party model hub like www.modelscope.cn or huggingface.""" - model_dir = None - # Only the rank0 are allowed to download model - if mpi_rank() == 0: - assert self._workspace is not None - assert isinstance(self.model_obj.model_name, str) - # this will download only once when multiple MPI processes are running - model_dir = download_hf_model(self.model_obj.model_name, - revision=self.llm_args.revision) - print_colored(f"Downloaded model to {model_dir}\n", 'grey') - # Make all the processes got the same model_dir - self._model_dir = mpi_broadcast(model_dir, root=0) - self.model_obj.model_dir = self._model_dir # mark as a local model - assert self.model_obj.is_local_model def _apply_modelopt_quant_config(self, hf_quant_config: Dict[str, Any], explicit_kv_cache_quant_algo) -> None: @@ -541,136 +283,6 @@ def _update_from_hf_quant_config(self) -> bool: return False - def _load_model_from_hf(self): - """Load a TRT-LLM model from a HF model.""" - assert self._model_dir is not None - - model_cls = AutoModelForCausalLM.get_trtllm_model_class( - self._model_dir, self.llm_args.trust_remote_code, - self.llm_args.decoding_config.decoding_mode - if hasattr(self.llm_args, "speculative_model") - and self.llm_args.speculative_model else None) - - prequantized = self._update_from_hf_quant_config() - - # FP4 Gemm force to use plugin. - if self.llm_args.quant_config.quant_mode.has_nvfp4(): - self.llm_args.build_config.plugin_config.gemm_plugin = "nvfp4" - - if self.llm_args.load_format == 'dummy': - config = model_cls.config_class.from_hugging_face( - str(self._model_dir), - dtype=self.llm_args.dtype, - mapping=self.mapping, - quant_config=self.llm_args.quant_config, - **self.convert_checkpoint_options, - ) - self.model = model_cls(config) - elif self.llm_args.quant_config._requires_calibration and not prequantized: - assert self.workspace is not None - checkpoint_dir = f"{self.workspace}/quantized-checkpoint" - if self.rank == 0: - model_cls.quantize( - self._model_dir, - checkpoint_dir, - dtype=self.llm_args.dtype, - mapping=self.mapping, - quant_config=self.llm_args.quant_config, - **self.llm_args.calib_config.model_dump(), - trust_remote_code=self.llm_args.trust_remote_code, - ) - if self.llm_args.parallel_config.is_multi_gpu: - mpi_barrier() - self.model = model_cls.from_checkpoint(checkpoint_dir, - rank=self.mapping.rank) - else: - self.model = model_cls.from_hugging_face( - str(self._model_dir), - dtype=self.llm_args.dtype, - mapping=self.mapping, - quant_config=self.llm_args.quant_config, - load_model_on_cpu= - True, # TODO:TRTLLM-195 to enhance the weights loading memory usage and choose best location - trust_remote_code=self.llm_args.trust_remote_code, - speculative_model_dir=self._speculative_model_dir, - speculative_config=self.llm_args.speculative_config - if not isinstance(self.llm_args.speculative_config, - LookaheadDecodingConfig) else None, - **self.convert_checkpoint_options, - ) - - self.pretrained_config = self.model.config - self._model_info = _ModelInfo.from_pretrained_config( - self.pretrained_config) - - @print_traceback_on_error - def _load_model_from_ckpt(self): - """Load a TRT-LLM model from checkpoint.""" - self.pretrained_config = PretrainedConfig.from_json_file( - os.path.join(self._model_dir, 'config.json')) - self.pretrained_config.mapping = self.mapping - - #TODO: TRTLLM-1091, change the architecture in the checkpoint to TRT-LLM one, not HF one. - architecture = self.pretrained_config.architecture - assert architecture in MODEL_MAP, \ - f"Unsupported model architecture: {architecture}" - model_cls = MODEL_MAP[architecture] - - if self.llm_args.load_format == 'dummy': - self.model = model_cls(self.pretrained_config) - else: - self.model = model_cls.from_checkpoint( - self._model_dir, config=self.pretrained_config) - self._model_info = _ModelInfo.from_pretrained_config( - self.pretrained_config) - - # load parallel embedding related options - self.convert_checkpoint_options[ - 'use_parallel_embedding'] = self.pretrained_config.use_parallel_embedding - - def _build_engine_from_inmemory_model(self): - assert isinstance(self.llm_args.model, Module) - self._model_info = _ModelInfo.from_module(self.model) - - @print_traceback_on_error - def _build_engine(self): - assert isinstance( - self.build_config, - BuildConfig), f"build_config is not set yet: {self.build_config}" - - logger_debug(f"rank{mpi_rank()} begin to build engine...\n", "green") - - # avoid side effects by copying the original build_config - copied_build_config = self.build_config.model_copy(deep=True) - - copied_build_config.update_kv_cache_type(self._model_info.architecture) - assert self.model is not None, "model has not been loaded yet." - - self._engine = build(self.model, copied_build_config) - self.mapping = self.model.config.mapping - - # delete the model explicitly to free all the build-time resources - self.model = None - logger_debug(f"rank{mpi_rank()} build engine done\n", "green") - - def _save_engine_for_runtime(self): - """Persist the engine to disk for the cpp runtime. - - Currently, the cpp runtime can accept an engine path, - that requires the engine should always be saved to disk. - - This explicit saving will be removed in the future when the cpp runtime can accept the engine buffer directly. - But this is necessary for a build cache, but it can be optimized to async IO. - """ - if self.build_cache_enabled: - self._model_dir = self.engine_cache_stage.cache_dir - self._model_format = _ModelFormatKind.TLLM_ENGINE - return - - def _load_engine_buffer(self): - # Load engine buffer from disk - self._engine = Engine.from_dir(self._model_dir) - @staticmethod def load_hf_tokenizer(model_dir, trust_remote_code: bool = True, @@ -777,9 +389,6 @@ def _download_hf_model_if_needed(self, def __call__(self) -> Tuple[Path, Union[Path, None]]: - if self.llm_args.model_format is _ModelFormatKind.TLLM_ENGINE: - return Path(self.llm_args.model), None - # Download speculative model from HuggingFace if needed (all backends) if (self.llm_args.speculative_config is not None and self.llm_args.speculative_config.speculative_model is not None): @@ -792,196 +401,25 @@ def __call__(self) -> Tuple[Path, Union[Path, None]]: if self.llm_args.backend == "_autodeploy": return None, "" - self.engine_cache_stage: Optional[CachedStage] = None self._hf_model_dir = None self.model_loader = ModelLoader(self.llm_args) - if self.llm_args.backend is not None: - if self.llm_args.backend not in ["pytorch", "_autodeploy"]: - raise ValueError( - f'backend {self.llm_args.backend} is not supported.') - - self._hf_model_dir = self._download_hf_model_if_needed( - self.model_loader.model_obj, revision=self.llm_args.revision) - - if self.llm_args.quant_config.quant_algo is not None: - logger.warning( - "QuantConfig for pytorch backend is ignored. You can load " - "quantized model with hf_quant_config.json directly.") - # Currently, this is to make updated quant_config visible by llm.args.quant_config - # TODO: Unify the logics with those in tensorrt_llm/_torch/model_config.py - self.model_loader._update_from_hf_quant_config() - - return None, self._hf_model_dir - - if self.model_loader.model_obj.is_hub_model: - # This will download the config.json from HF model hub, this helps to create a PretrainedConfig for - # cache key. - self._hf_model_dir = download_hf_pretrained_config( - self.model_loader.model_obj.model_name, - revision=self.llm_args.revision) - - elif self.model_loader.model_obj.is_local_model: - self._hf_model_dir = self.model_loader.model_obj.model_dir if self.llm_args.model_format is _ModelFormatKind.HF else None - - if self.build_cache_enabled: - print_colored("Build cache is enabled.\n", 'yellow') - - self.engine_cache_stage = self._get_engine_cache_stage() - if self.engine_cache_stage.is_cached(): - self.llm_build_stats.cache_hitted = True - print_colored( - f"Reusing cached engine in {self.engine_cache_stage.get_engine_path()}\n\n", - 'grey') - self.model_loader.model_obj.model_dir = self.engine_cache_stage.get_engine_path( - ) - self.llm_build_stats.engine_dir = self.model_loader.model_obj.model_dir - return self.llm_build_stats.engine_dir, self._hf_model_dir - - return self._build_model(), self._hf_model_dir - - def get_engine_dir(self) -> Path: - if self.llm_args.model_format is _ModelFormatKind.TLLM_ENGINE: - return self.model_obj.model_dir - - # generate a new path for writing the engine - if self.build_cache_enabled: - cache_stage = self._get_engine_cache_stage() - return cache_stage.get_engine_path() - - return self.workspace / "tmp.engine" - - @property - def build_cache_enabled(self) -> bool: - _enable_build_cache, _ = get_build_cache_config_from_env() - - return (self.llm_args.enable_build_cache - or _enable_build_cache) and (self.llm_args.model_format - is _ModelFormatKind.HF) - - def _get_engine_cache_stage(self) -> CachedStage: - """Get the cache stage for engine building.""" - build_cache = BuildCache(self.llm_args.enable_build_cache) - - assert self._hf_model_dir is not None, "HF model dir is required for cache key." - - def serialize(d) -> str: - if hasattr(d, "to_dict"): - dic = d.to_dict() - elif is_dataclass(d): - dic = asdict(d) - elif isinstance(d, BaseModel): - dic = d.model_dump(mode="json") - else: - raise ValueError(f"Could not serialize type: {type(d)}") - return json.dumps(dic, sort_keys=True) - - parallel_config = self.llm_args.parallel_config - - force_rebuild = False - if self.llm_args.model_format is not _ModelFormatKind.HF: - force_rebuild = True - - return build_cache.get_engine_building_cache_stage( - build_config=self.llm_args.build_config, - model_path=self._hf_model_dir, - force_rebuild=force_rebuild, - # Other configs affecting the engine building - parallel_config=serialize(parallel_config), - pretrained_config=serialize(self.get_pretrained_config()), - quant_config=serialize(self.llm_args.quant_config), - ) - - def get_pretrained_config(self) -> PretrainedConfig: - """Get the PretrainedConfig for cache key. - - NOTE, this is not the HF model's config, but the TRT-LLM's config. We use this as a generic information for - HF and other models. - """ - assert self._hf_model_dir is not None - return AutoConfig.from_hugging_face( - self._hf_model_dir, - mapping=self.llm_args.parallel_config.to_mapping(), - quant_config=self.llm_args.quant_config, - dtype=self.llm_args.dtype, - trust_remote_code=self.llm_args.trust_remote_code) - - def _build_model(self) -> Path: - model_format = self.llm_args.model_format - - def build_task(engine_dir: Path): - if model_format is not _ModelFormatKind.TLLM_ENGINE: - model_loader_kwargs = { - 'llm_args': self.llm_args, - 'workspace': str(self.workspace), - 'llm_build_stats': self.llm_build_stats, - } - - if self.llm_args.parallel_config.is_multi_gpu: - assert self.mpi_session - - #mpi_session cannot be pickled so remove from self.llm_args - if self.llm_args.mpi_session: - del self.llm_args.mpi_session - - # The engine_dir:Path will be stored to MPINodeState.state - build_infos = self.mpi_session.submit_sync( - CachedModelLoader._node_build_task, - engine_dir=engine_dir, - **model_loader_kwargs) - self.llm_build_stats.build_steps_info = build_infos[0] - - else: # single-gpu - with ModelLoader(**model_loader_kwargs) as model_loader: - model_loader(engine_dir=engine_dir) - - release_gc() - - has_storage = True - if self.build_cache_enabled: - try: - # TODO[chunweiy]: Cover the case when the model is from HF model hub. - if self.model_loader.model_obj.is_local_model: - # This is not perfect, but will make build-cache much more robust. - free_storage = self.engine_cache_stage.parent.free_storage_in_gb( - ) - model_size = get_directory_size_in_gb( - self.model_loader.model_obj.model_dir) - require_size = model_size * 1.3 - has_storage = free_storage >= require_size - - if not has_storage: - print_colored( - "Build cache is disabled since the cache storage is too small.\n ", - 'yellow') - print_colored( - f"Free storage: {free_storage}GB, Required storage: {require_size}GB\n", - 'grey') - except ValueError: - has_storage = False - except Exception as e: - logger.error(e) - has_storage = False - - if enable_llm_debug(): - print_colored(f"Has cache storage: {has_storage}\n", 'yellow') - - if has_storage: - with self.engine_cache_stage.write_guard() as engine_dir: - build_task(engine_dir) - self.llm_build_stats.cache_hitted = True + if self.llm_args.backend not in ["pytorch", "_autodeploy"]: + raise ValueError( + f'backend {self.llm_args.backend} is not supported.') - else: - print_colored( - "The cache directory is too small, build-cache is disabled.\n", - 'grey') - self.llm_build_stats.cache_hitted = False - self.llm_build_stats.cache_info = "The cache root directory is too small." + self._hf_model_dir = self._download_hf_model_if_needed( + self.model_loader.model_obj, revision=self.llm_args.revision) - if not (has_storage and self.build_cache_enabled): - build_task(self.get_engine_dir()) + if self.llm_args.quant_config.quant_algo is not None: + logger.warning( + "QuantConfig for pytorch backend is ignored. You can load " + "quantized model with hf_quant_config.json directly.") + # Currently, this is to make updated quant_config visible by llm.args.quant_config + # TODO: Unify the logics with those in tensorrt_llm/_torch/model_config.py + self.model_loader._update_from_hf_quant_config() - return self.get_engine_dir() + return None, self._hf_model_dir @print_traceback_on_error @staticmethod @@ -994,27 +432,6 @@ def _node_download_hf_model( else: return None - @print_traceback_on_error - @staticmethod - def _node_build_task( - llm_args: LlmArgs, - workspace: Optional[str | tempfile.TemporaryDirectory] = None, - llm_build_stats: Optional['LlmBuildStats'] = None, - engine_dir: Optional[Path] = None, - ): - if MPINodeState.is_initialized(): - raise RuntimeError("The MPI node is already initialized.") - - with ModelLoader(llm_args, - workspace=workspace, - llm_build_stats=llm_build_stats) as model_loader: - model_loader(engine_dir=engine_dir) - return model_loader.llm_build_stats.build_steps_info - - def save(self, engine_dir: Path): - # copy the engine directory to the target directory - shutil.copytree(self.get_engine_dir(), engine_dir) - @dataclass class LlmBuildStats: @@ -1038,10 +455,7 @@ class LlmBuildStats: 'LlmArgs', 'LlmBuildStats', 'ModelLoader', - '_ModelRuntimeContext', - '_ModelInfo', '_ParallelConfig', - '_ModelFormatKind', '_ModelWrapper', 'BatchingType', 'ExecutorConfig', @@ -1055,8 +469,6 @@ class LlmBuildStats: 'UserProvidedDecodingConfig', 'ContextChunkingPolicy', 'CapacitySchedulerPolicy', - 'BuildConfig', - 'BuildCacheConfig', 'QuantConfig', 'CalibConfig', 'CudaGraphConfig', diff --git a/tensorrt_llm/logger.py b/tensorrt_llm/logger.py index 49bd2a415342..7ed205e8f24f 100644 --- a/tensorrt_llm/logger.py +++ b/tensorrt_llm/logger.py @@ -17,8 +17,6 @@ import sys from typing import Dict, Optional -import tensorrt as trt - try: from polygraphy.logger import G_LOGGER except ImportError: @@ -75,7 +73,6 @@ def _extract_module(qualname: str) -> str: "flash_mla": "flashmla", "quantization": "quantize", "scaffolding": "scaffold", - "_tensorrt_engine": "trt_engn", "visual_gen": "vis_gen", "__pycache__": "pycache", "tokenizer": "tokenizr", @@ -186,7 +183,6 @@ def __init__(self): min_severity = self.DEFAULT_LEVEL self._min_severity = min_severity - self._trt_logger = trt.Logger(severity_map[min_severity][0]) self._logger = logging.getLogger(self.PREFIX) self._logger.propagate = False handler = logging.StreamHandler(stream=sys.stdout) @@ -204,7 +200,7 @@ def __init__(self): # Set the underlying Python logger to the minimum of all configured # levels so that per-module overrides more verbose than the global # level are not silently dropped by Python's logging framework. - global_py_level = severity_map[min_severity][1] + global_py_level = severity_map[min_severity][0] if self._module_levels: # Map our numeric levels to Python logging levels. _numeric_to_py = { @@ -224,7 +220,7 @@ def __init__(self): self._polygraphy_logger = G_LOGGER if self._polygraphy_logger is not None: - self._polygraphy_logger.module_severity = severity_map[min_severity][2] + self._polygraphy_logger.module_severity = severity_map[min_severity][1] # For log_once self._appeared_keys = set() @@ -268,10 +264,6 @@ def _func_wrapper(self, severity): else: raise AttributeError(f"No such severity: {severity}") - @property - def trt_logger(self) -> trt.ILogger: - return self._trt_logger - def log(self, severity, *msg): module = _get_caller_module() if not self.is_severity_enabled(severity, module): @@ -338,20 +330,21 @@ def set_level(self, min_severity): ) return self._min_severity = min_severity - self._trt_logger.min_severity = severity_map[min_severity][0] - self._logger.setLevel(severity_map[min_severity][1]) + self._logger.setLevel(severity_map[min_severity][0]) if self._polygraphy_logger is not None: - self._polygraphy_logger.module_severity = severity_map[min_severity][2] + self._polygraphy_logger.module_severity = severity_map[min_severity][1] severity_map = { - "internal_error": [trt.Logger.INTERNAL_ERROR, logging.CRITICAL], - "error": [trt.Logger.ERROR, logging.ERROR], - "warning": [trt.Logger.WARNING, logging.WARNING], - "info": [trt.Logger.INFO, logging.INFO], - "verbose": [trt.Logger.VERBOSE, logging.DEBUG], - "debug": [trt.Logger.VERBOSE, logging.DEBUG], - "trace": [trt.Logger.VERBOSE, logging.DEBUG], + # [0] Python logging level; [1] polygraphy level (appended below when + # polygraphy is installed). + "internal_error": [logging.CRITICAL], + "error": [logging.ERROR], + "warning": [logging.WARNING], + "info": [logging.INFO], + "verbose": [logging.DEBUG], + "debug": [logging.DEBUG], + "trace": [logging.DEBUG], } if G_LOGGER is not None: diff --git a/tensorrt_llm/lora_helper.py b/tensorrt_llm/lora_helper.py index f2222505cfe3..a2f85efb9359 100644 --- a/tensorrt_llm/lora_helper.py +++ b/tensorrt_llm/lora_helper.py @@ -67,27 +67,6 @@ def get_default_trtllm_modules_to_hf_modules(): } -def use_lora( - model, - lora_config: "LoraConfig", - trtllm_modules_to_hf_modules: Optional[Dict[str, str]] = None, -): - """Use LoRA with the given model and configuration. - - This function is a wrapper that delegates to the appropriate loading function - based on the LoRA checkpoint source. - """ - if lora_config.lora_ckpt_source == "nemo": - from .lora_manager import load_nemo_lora - load_nemo_lora(model, lora_config) - elif lora_config.lora_ckpt_source == "hf": - from .lora_manager import load_hf_lora - load_hf_lora(model, lora_config, trtllm_modules_to_hf_modules) - else: - raise ValueError( - f"Unsupported lora_ckpt_source: {lora_config.lora_ckpt_source}") - - class LoraConfig(StrictBaseModel): lora_dir: List[str] = Field( default_factory=list, diff --git a/tensorrt_llm/lora_manager.py b/tensorrt_llm/lora_manager.py index 933c63d1f615..f703a85769d5 100644 --- a/tensorrt_llm/lora_manager.py +++ b/tensorrt_llm/lora_manager.py @@ -17,15 +17,14 @@ from tensorrt_llm.bindings import internal as tb_internal -from ._utils import pad_vocab_size, release_gc, str_dtype_to_torch, torch_to_numpy -from .layers.linear import ColumnLinear +from ._utils import release_gc, str_dtype_to_torch, torch_to_numpy from .lora_helper import ( LoraConfig, get_default_trtllm_modules_to_hf_modules, get_missing_qkv_modules_from_lora_modules, ) from .mapping import Mapping -from .models.convert_utils import get_model_path, load_state_dict, split_matrix_tp +from .models.convert_utils import get_model_path, load_state_dict if TYPE_CHECKING: from .runtime import ModelConfig @@ -426,23 +425,13 @@ def get_target_modules(self): return self.lora_target_modules -def load_nemo_lora(model, lora_config: LoraConfig): - lora_loader = NemoLoraLoader(lora_config.lora_dir) - - if not lora_loader.is_valid: - raise ValueError(f"Failed to load NeMo LoRA from {lora_config.lora_dir}") - - if len(lora_config.lora_target_modules) == 0: - lora_config.lora_target_modules = lora_loader.lora_target_modules - - def load_torch_hf_lora(lora_config: LoraConfig): - """This is a shortned version of load_hf_lora that is used for torch models. + """Load an HF LoRA checkpoint for the PyTorch workflow. - Main problem is model.config in legacy code is custom (defined in the legacy code) whereas - pivot model config is the transformer's one. + Populates lora_config (trtllm_modules_to_hf_modules and inferred + lora_target_modules) from the HF adapter directory. The actual weights are + loaded later by LoraManager when requests arrive with LoRA UIDs. """ - # TODO smor- need to comibe with load_hf_lora if not lora_config.trtllm_modules_to_hf_modules: lora_config.trtllm_modules_to_hf_modules = get_default_trtllm_modules_to_hf_modules() @@ -530,94 +519,6 @@ def load_torch_lora(lora_config: LoraConfig): ) -def load_hf_lora( - model, - lora_config: LoraConfig, - trtllm_modules_to_hf_modules: Optional[Dict[str, str]] = None, -): - trtllm_modules_to_hf_modules = ( - trtllm_modules_to_hf_modules or get_default_trtllm_modules_to_hf_modules() - ) - lora_config.trtllm_modules_to_hf_modules = trtllm_modules_to_hf_modules - - lora_loader = HfLoraLoader(lora_config.lora_dir) - - if len(lora_config.lora_target_modules) == 0: - lora_config.lora_target_modules = lora_loader.get_target_modules( - trtllm_modules_to_hf_modules - ) - if len(lora_config.lora_target_modules) == 0: - raise ValueError( - "lora_target_modules is empty. " - "Please specify lora_target_modules or provide lora_dir to infer lora_target_modules." - ) - - missing_qkv_modules = LoraManager.get_missing_qkv_modules(lora_config.lora_target_modules) - lora_config.lora_target_modules.extend(missing_qkv_modules) - - if lora_loader.is_valid: - config = model.config - torch_dtype = str_dtype_to_torch(config.dtype) - # the lora checkpoint might finetune the embedding - if lora_loader.vocab_size != 0: - config.vocab_size = lora_loader.vocab_size - mapping = config.mapping - if mapping.is_first_pp_rank() and lora_loader.embed_tokens is not None: - weight = lora_loader.embed_tokens - if config.use_parallel_embedding: - weight = split_matrix_tp( - weight, - mapping.tp_size, - mapping.tp_rank, - dim=config.embedding_sharding_dim, - ) - if model.transformer.vocab_embedding.weight.raw_value.shape != weight.shape: - model.transformer.vocab_embedding = model.transformer.vocab_embedding.__class__( - num_embeddings=config.vocab_size, - embedding_dim=config.hidden_size, - dtype=config.dtype, - tp_size=mapping.tp_size if config.use_parallel_embedding else 1, - tp_group=mapping.tp_group if config.use_parallel_embedding else None, - sharding_dim=config.embedding_sharding_dim, - tp_rank=mapping.tp_rank, - ) - model.transformer.vocab_embedding.weight.value = weight.to(torch_dtype) - if mapping.is_last_pp_rank() and lora_loader.lm_head is not None: - weight = lora_loader.lm_head - vocab_size = lora_loader.vocab_size - if vocab_size % mapping.tp_size != 0: - # padding - vocab_size_padded = pad_vocab_size(vocab_size, mapping.tp_size) - pad_width = vocab_size_padded - vocab_size - - weight = torch.from_numpy( - np.pad( - torch_to_numpy(weight), - ((0, pad_width), (0, 0)), - "constant", - constant_values=0, - ) - ) - else: - vocab_size_padded = vocab_size - if model.lm_head.weight.raw_value.shape != weight.shape: - model.lm_head = ColumnLinear( - config.hidden_size, - vocab_size_padded, - bias=False, - dtype=config.dtype, - tp_group=mapping.tp_group, - tp_size=mapping.tp_size, - gather_output=True, - ) - model.lm_head.weight.value = split_matrix_tp( - weight, - mapping.tp_size, - mapping.tp_rank, - dim=0, - ).to(torch_dtype) - - def unpack_nemo_weights(nemo_archive_path: str) -> Tuple[Dict, Dict[str, torch.Tensor]]: """Unpack model config and weights from a NeMo .nemo archive file. diff --git a/tensorrt_llm/models/__init__.py b/tensorrt_llm/models/__init__.py index 96bd4eff96d2..3af378d4678f 100755 --- a/tensorrt_llm/models/__init__.py +++ b/tensorrt_llm/models/__init__.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# 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"); @@ -12,212 +12,19 @@ # WITHOUT 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 .baichuan.model import BaichuanForCausalLM -from .bert.model import (BertForQuestionAnswering, - BertForSequenceClassification, BertModel, - RobertaForQuestionAnswering, - RobertaForSequenceClassification, RobertaModel) -from .bloom.model import BloomForCausalLM, BloomModel -from .chatglm.config import ChatGLMConfig -from .chatglm.model import ChatGLMForCausalLM, ChatGLMModel -from .clip.model import CLIPVisionTransformer -from .cogvlm.config import CogVLMConfig -from .cogvlm.model import CogVLMForCausalLM -from .commandr.model import CohereForCausalLM -from .dbrx.config import DbrxConfig -from .dbrx.model import DbrxForCausalLM -from .deepseek_v1.model import DeepseekForCausalLM -from .deepseek_v2.model import DeepseekV2ForCausalLM -from .dit.model import DiT -from .eagle.model import EagleForCausalLM -from .enc_dec.model import DecoderModel, EncoderModel, WhisperEncoder -from .falcon.config import FalconConfig -from .falcon.model import FalconForCausalLM, FalconModel -from .gemma.config import (GEMMA2_ARCHITECTURE, GEMMA3_ARCHITECTURE, - GEMMA_ARCHITECTURE, GemmaConfig) -from .gemma.model import GemmaForCausalLM -from .gpt.config import GPTConfig -from .gpt.model import GPTForCausalLM, GPTModel -from .gptj.config import GPTJConfig -from .gptj.model import GPTJForCausalLM, GPTJModel -from .gptneox.model import GPTNeoXForCausalLM, GPTNeoXModel -from .grok.model import GrokForCausalLM -from .llama.config import LLaMAConfig -from .llama.model import LLaMAForCausalLM, LLaMAModel -from .mamba.model import MambaForCausalLM -from .medusa.config import MedusaConfig -from .medusa.model import MedusaForCausalLm -from .mllama.model import MLLaMAForCausalLM -from .mmdit_sd3.model import SD3Transformer2DModel -from .modeling_utils import (PretrainedConfig, PretrainedModel, - SpeculativeDecodingMode) -from .mpt.model import MPTForCausalLM, MPTModel -from .multimodal_encoders.config import LlavaNextVisionConfig -from .multimodal_encoders.model import LlavaNextVisionWrapper -from .nemotron_nas.model import DeciLMForCausalLM -from .opt.model import OPTForCausalLM, OPTModel -from .phi3.model import Phi3ForCausalLM, Phi3Model -from .phi.model import PhiForCausalLM, PhiModel -from .qwen.model import QWenForCausalLM -from .recurrentgemma.model import RecurrentGemmaForCausalLM -from .redrafter.model import ReDrafterForLLaMALM, ReDrafterForQWenLM -from .stdit.model import STDiT3Model + +from .modeling_utils import (LayerQuantConfig, PretrainedConfig, QuantAlgo, + QuantConfig, SpeculativeDecodingMode) + +# Architecture registry is intentionally empty: the PyTorch backend resolves +# model classes via tensorrt_llm._torch.models. +MODEL_MAP = {} __all__ = [ - 'BertModel', - 'BertForQuestionAnswering', - 'BertForSequenceClassification', - 'RobertaModel', - 'RobertaForQuestionAnswering', - 'RobertaForSequenceClassification', - 'BloomModel', - 'BloomForCausalLM', - 'CLIPVisionTransformer', - 'DiT', - 'SD3Transformer2DModel', - 'STDiT3', - 'DeepseekForCausalLM', - 'FalconConfig', - 'DeepseekV2ForCausalLM', - 'FalconForCausalLM', - 'FalconModel', - 'GPTConfig', - 'GPTModel', - 'GPTForCausalLM', - 'OPTForCausalLM', - 'OPTModel', - 'LLaMAConfig', - 'LLaMAForCausalLM', - 'LLaMAModel', - 'LlavaNextVisionWrapper', - 'LlavaNextVisionConfig', - 'MedusaConfig', - 'MedusaForCausalLm', - 'ReDrafterForLLaMALM', - 'ReDrafterForQWenLM', - 'GPTJConfig', - 'GPTJModel', - 'GPTJForCausalLM', - 'GPTNeoXModel', - 'GPTNeoXForCausalLM', - 'PhiModel', - 'PhiConfig', - 'Phi3Model', - 'Phi3Config', - 'PhiForCausalLM', - 'Phi3ForCausalLM', - 'ChatGLMConfig', - 'ChatGLMForCausalLM', - 'ChatGLMModel', - 'BaichuanForCausalLM', - 'QWenConfig' - 'QWenForCausalLM', - 'QWenModel', - 'EncoderModel', - 'DecoderModel', 'PretrainedConfig', - 'PretrainedModel', - 'WhisperEncoder', - 'MambaForCausalLM', - 'MambaConfig', - 'MPTForCausalLM', - 'MPTModel', - 'SkyworkForCausalLM', - 'GemmaConfig', - 'GemmaForCausalLM', - 'DbrxConfig', - 'DbrxForCausalLM', - 'RecurrentGemmaForCausalLM', - 'CogVLMConfig', - 'CogVLMForCausalLM', - 'EagleForCausalLM', 'SpeculativeDecodingMode', - 'CohereForCausalLM', - 'MLLaMAForCausalLM', + 'QuantConfig', + 'LayerQuantConfig', + 'QuantAlgo', + 'MODEL_MAP', ] - -MODEL_MAP = { - 'GPT2LMHeadModel': GPTForCausalLM, - 'GPT2LMHeadCustomModel': GPTForCausalLM, - 'GPTBigCodeForCausalLM': GPTForCausalLM, - 'Starcoder2ForCausalLM': GPTForCausalLM, - 'FuyuForCausalLM': GPTForCausalLM, - 'Kosmos2ForConditionalGeneration': GPTForCausalLM, - 'JAISLMHeadModel': GPTForCausalLM, - 'GPTForCausalLM': GPTForCausalLM, - 'NemotronForCausalLM': GPTForCausalLM, - 'OPTForCausalLM': OPTForCausalLM, - 'BloomForCausalLM': BloomForCausalLM, - 'RWForCausalLM': FalconForCausalLM, - 'FalconForCausalLM': FalconForCausalLM, - 'PhiForCausalLM': PhiForCausalLM, - 'Phi3ForCausalLM': Phi3ForCausalLM, - 'Phi3VForCausalLM': Phi3ForCausalLM, - 'Phi3SmallForCausalLM': Phi3ForCausalLM, - 'PhiMoEForCausalLM': Phi3ForCausalLM, - 'Phi4MMForCausalLM': Phi3ForCausalLM, - 'MambaForCausalLM': MambaForCausalLM, - 'GPTNeoXForCausalLM': GPTNeoXForCausalLM, - 'GPTJForCausalLM': GPTJForCausalLM, - 'MptForCausalLM': MPTForCausalLM, - 'MPTForCausalLM': MPTForCausalLM, - 'GLMModel': ChatGLMForCausalLM, - 'ChatGLMModel': ChatGLMForCausalLM, - 'ChatGLMForCausalLM': ChatGLMForCausalLM, - 'ChatGLMForConditionalGeneration': ChatGLMForCausalLM, - 'LlamaForCausalLM': LLaMAForCausalLM, - 'LlavaLlamaModel': LLaMAForCausalLM, - 'LlavaNextForConditionalGeneration': LlavaNextVisionWrapper, - 'ExaoneForCausalLM': LLaMAForCausalLM, - 'MistralForCausalLM': LLaMAForCausalLM, - 'MixtralForCausalLM': LLaMAForCausalLM, - 'ArcticForCausalLM': LLaMAForCausalLM, - 'Grok1ModelForCausalLM': GrokForCausalLM, - 'InternLMForCausalLM': LLaMAForCausalLM, - 'InternLM2ForCausalLM': LLaMAForCausalLM, - 'InternLMXComposer2ForCausalLM': LLaMAForCausalLM, - 'GraniteForCausalLM': LLaMAForCausalLM, - 'GraniteMoeForCausalLM': LLaMAForCausalLM, - 'MedusaForCausalLM': MedusaForCausalLm, - 'MedusaLlamaForCausalLM': MedusaForCausalLm, - 'ReDrafterForLLaMALM': ReDrafterForLLaMALM, - 'ReDrafterForQWenLM': ReDrafterForQWenLM, - 'BaichuanForCausalLM': BaichuanForCausalLM, - 'BaiChuanForCausalLM': BaichuanForCausalLM, - 'SkyworkForCausalLM': LLaMAForCausalLM, - GEMMA_ARCHITECTURE: GemmaForCausalLM, - GEMMA2_ARCHITECTURE: GemmaForCausalLM, - GEMMA3_ARCHITECTURE: GemmaForCausalLM, - 'QWenLMHeadModel': QWenForCausalLM, - 'QWenForCausalLM': QWenForCausalLM, - 'Qwen2ForCausalLM': QWenForCausalLM, - 'Qwen2MoeForCausalLM': QWenForCausalLM, - 'Qwen2ForSequenceClassification': QWenForCausalLM, - 'Qwen2VLForConditionalGeneration': QWenForCausalLM, - 'Qwen2VLModel': QWenForCausalLM, - 'Qwen3ForCausalLM': QWenForCausalLM, - 'Qwen3MoeForCausalLM': QWenForCausalLM, - 'WhisperEncoder': WhisperEncoder, - 'EncoderModel': EncoderModel, - 'DecoderModel': DecoderModel, - 'DbrxForCausalLM': DbrxForCausalLM, - 'RecurrentGemmaForCausalLM': RecurrentGemmaForCausalLM, - 'CogVLMForCausalLM': CogVLMForCausalLM, - 'DiT': DiT, - 'SD3Transformer2DModel': SD3Transformer2DModel, - 'STDiT3': STDiT3Model, - 'DeepseekForCausalLM': DeepseekForCausalLM, - 'DeciLMForCausalLM': DeciLMForCausalLM, - 'DeepseekV2ForCausalLM': DeepseekV2ForCausalLM, - 'EagleForCausalLM': EagleForCausalLM, - 'CohereForCausalLM': CohereForCausalLM, - 'MLLaMAModel': MLLaMAForCausalLM, # For modelopt - 'MllamaForConditionalGeneration': - MLLaMAForCausalLM, # For mllama load by Auto - 'BertForQuestionAnswering': BertForQuestionAnswering, - 'BertForSequenceClassification': BertForSequenceClassification, - 'BertModel': BertModel, - 'RobertaModel': RobertaModel, - 'RobertaForQuestionAnswering': RobertaForQuestionAnswering, - 'RobertaForSequenceClassification': RobertaForSequenceClassification, -} diff --git a/tensorrt_llm/models/baichuan/__init__.py b/tensorrt_llm/models/baichuan/__init__.py deleted file mode 100644 index 71bf6d298d42..000000000000 --- a/tensorrt_llm/models/baichuan/__init__.py +++ /dev/null @@ -1,14 +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. diff --git a/tensorrt_llm/models/baichuan/config.py b/tensorrt_llm/models/baichuan/config.py deleted file mode 100644 index 599b98089f75..000000000000 --- a/tensorrt_llm/models/baichuan/config.py +++ /dev/null @@ -1,92 +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. -from typing import Optional, Union - -from ...logger import logger -from ...mapping import Mapping -from ..convert_utils import infer_dtype -from ..modeling_utils import PretrainedConfig, QuantConfig - - -class BaichuanConfig(PretrainedConfig): - - def __init__(self, model_version: Optional[str] = None, **kwargs): - super().__init__(**kwargs) - if model_version is None: - model_version = BaichuanConfig.guess_model_version(self) - self.model_version = model_version - - def to_dict(self): - output = super().to_dict() - # Serialize the fields added in BaichuanConfig - output['model_version'] = self.model_version - return output - - @staticmethod - def guess_model_version( - config: Union['transformers.PretrainedConfig', - 'BaichuanConfig']) -> str: - logger.warning( - "Model version is not set, trying to guess from loaded config") - size = '7' if config.num_attention_heads == 32 else '13' - version = '1' if config.vocab_size == 64000 else '2' - return f'v{version}_{size}b' - - @classmethod - def from_hugging_face( - cls, - hf_config_or_dir: Union[str, 'transformers.PretrainedConfig'], - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - import transformers - trust_remote_code = kwargs.pop('trust_remote_code', True) - if isinstance(hf_config_or_dir, transformers.PretrainedConfig): - hf_config = hf_config_or_dir - else: - hf_config_dir = str(hf_config_or_dir) - hf_config = transformers.AutoConfig.from_pretrained( - hf_config_dir, trust_remote_code=trust_remote_code) - - model_version = kwargs.pop('model_version', None) - if model_version is None: - model_version = BaichuanConfig.guess_model_version(hf_config) - - if model_version == 'v1_7b' or model_version == 'v2_7b': - position_embedding_type = 'rope_gpt_neox' - max_position_embeddings = hf_config.max_position_embeddings - else: - position_embedding_type = 'alibi' - max_position_embeddings = hf_config.model_max_length - - dtype = infer_dtype(dtype, getattr(hf_config, 'torch_dtype', None)) - - return cls(architecture='BaichuanForCausalLM', - dtype=dtype, - vocab_size=hf_config.vocab_size, - max_position_embeddings=max_position_embeddings, - hidden_size=hf_config.hidden_size, - num_hidden_layers=hf_config.num_hidden_layers, - num_attention_heads=hf_config.num_attention_heads, - num_key_value_heads=hf_config.num_attention_heads, - hidden_act=hf_config.hidden_act, - intermediate_size=hf_config.intermediate_size, - norm_epsilon=hf_config.rms_norm_eps, - position_embedding_type=position_embedding_type, - model_version=model_version, - mapping=mapping, - quantization=quant_config, - **kwargs) diff --git a/tensorrt_llm/models/baichuan/convert.py b/tensorrt_llm/models/baichuan/convert.py deleted file mode 100644 index 576d5e1a34c8..000000000000 --- a/tensorrt_llm/models/baichuan/convert.py +++ /dev/null @@ -1,931 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -import copy -import functools -import os -import time -from collections import defaultdict -from typing import Any, Dict, List, Optional - -import numpy as np -import safetensors -import torch -import torch.nn as nn -from datasets import Dataset -from tqdm import tqdm -from transformers import AutoModelForCausalLM, AutoTokenizer -from transformers.pytorch_utils import Conv1D - -from ...logger import logger -from ...quantization import QuantAlgo, QuantMode -from ..convert_utils import (load_calib_dataset, smooth_gemm, - smooth_gemm_fc1_gate, weight_only_quantize_dict) -from .config import BaichuanConfig - - -def generate_int8( - weights: torch.Tensor, - act_range: Dict[str, torch.Tensor], - is_qkv: bool = False, - multi_query_mode: bool = False, -): - """ - This function has two purposes: - - compute quantized weights, scaled either per-tensor or per-column - - compute scaling factors - - Depending on the GEMM API (CUTLASS/CUBLAS) the required scaling factors differ. - CUTLASS uses two sets of scaling factors. One for the activation X, one for the weight W. - CUBLAS only has one (we can't do per-row scaling). So we must provide pre-multiplied scaling factor. - - Here is the list of what we need (T means per-tensor, C per-column): - - scale_x_orig_quant puts fp activation into the quantized range (i.e. [-128, 127], for int8). Used before the GEMM. (T) - - scale_y_quant_orig puts quantized activation into the fp range. Used if the GEMM outputs int8. (T) - - scale_w_quant_orig puts weights from quant range to fp range (used with CUTLASS) (T, C) - - scale_y_accum_quant puts the GEMM result (XW) from accumulation range (int32) - to quant range (int8) (used for CUBLAS) (T, C) - - Note that we don't do anything special about row-parallel GEMM. Theoretically, we could have per-GPU scaling factors too, - but then the model would change depending on the number of GPUs used. - - For QKV projection, the behavior is special. Even if we have a single matrix to perform QKV projection, we consider it - as three different matrices: Q, K, and V. So per-tensor actually means one scaling factor for each Q, K and V. - For our GEMM implementation to respect this behavior, we use per-column mode and replicate values along columns. - """ - - # compute weight scaling factors for fp->int8 and int8->fp - if is_qkv and not multi_query_mode: - scale_w_orig_quant_t = 127. / act_range["w"].reshape(3, -1).max( - dim=-1, keepdims=True)[0].cpu().numpy() - scale_w_orig_quant_c = 127. / act_range["w"].reshape(3, - -1).cpu().numpy() - elif is_qkv and multi_query_mode: - hidden_dim = weights.shape[0] - local_dim = act_range["w"].shape[0] - kv_dim = (local_dim - hidden_dim) // 2 - scale_w_q = act_range["w"][0:hidden_dim] - scale_w_k = act_range["w"][hidden_dim:hidden_dim + kv_dim] - scale_w_v = act_range["w"][-kv_dim:] - - scale_w_qkv_t = torch.concat([ - scale_w_q.max(dim=0, keepdim=True)[0], - scale_w_k.max(dim=0, keepdim=True)[0], - scale_w_v.max(dim=0, keepdim=True)[0] - ]) - - scale_w_orig_quant_t = 127. / scale_w_qkv_t.cpu().numpy() - scale_w_orig_quant_c = 127. / act_range["w"].cpu().numpy() - else: - scale_w_orig_quant_t = 127. / act_range["w"].max().cpu().numpy() - scale_w_orig_quant_c = 127. / act_range["w"].cpu().numpy() - scale_w_quant_orig_t = 1.0 / scale_w_orig_quant_t - scale_w_quant_orig_c = 1.0 / scale_w_orig_quant_c - - # compute the rest of needed scaling factors - scale_x_orig_quant_t = np.array(127. / act_range["x"].max().item()) - scale_y_orig_quant_t = np.array(127. / act_range["y"].max().item()) - scale_y_quant_orig_t = np.array(act_range["y"].max().item() / 127.) - scale_y_accum_quant_t = scale_y_orig_quant_t / (scale_x_orig_quant_t * - scale_w_orig_quant_t) - scale_y_accum_quant_c = scale_y_orig_quant_t / (scale_x_orig_quant_t * - scale_w_orig_quant_c) - if is_qkv and not multi_query_mode: - scale_y_accum_quant_t = np.broadcast_to(scale_y_accum_quant_t, - scale_w_orig_quant_c.shape) - scale_w_quant_orig_t = np.broadcast_to(scale_w_quant_orig_t, - scale_w_orig_quant_c.shape) - if is_qkv and multi_query_mode: - scale_q_y_accum_t = np.broadcast_to(scale_y_accum_quant_t[0], - scale_w_q.shape) - scale_k_y_accum_t = np.broadcast_to(scale_y_accum_quant_t[1], - scale_w_k.shape) - scale_v_y_accum_t = np.broadcast_to(scale_y_accum_quant_t[2], - scale_w_v.shape) - scale_y_accum_quant_t = np.concatenate( - [scale_q_y_accum_t, scale_k_y_accum_t, scale_v_y_accum_t]) - scale_w_quant_orig_t = np.concatenate([ - np.broadcast_to(scale_w_quant_orig_t[0], scale_w_q.shape), - np.broadcast_to(scale_w_quant_orig_t[1], scale_w_k.shape), - np.broadcast_to(scale_w_quant_orig_t[2], scale_w_v.shape) - ]) - - to_i8 = lambda x: x.round().clip(-127, 127).astype(np.int8) - - if is_qkv and multi_query_mode: - scale_w_quant_orig_t_expand = np.ones([weights.shape[-1]]) - scale_w_quant_orig_t_expand[:hidden_dim] = scale_w_quant_orig_t[0] - scale_w_quant_orig_t_expand[hidden_dim:hidden_dim + - kv_dim] = scale_w_quant_orig_t[1] - scale_w_quant_orig_t_expand[-kv_dim:] = scale_w_quant_orig_t[2] - weight_int8 = to_i8(weights * scale_w_quant_orig_t_expand) - else: - weight_int8 = to_i8(weights * scale_w_orig_quant_t) - return { - "weight.int8": weight_int8, - "weight.int8.col": to_i8(weights * scale_w_orig_quant_c), - "scale_x_orig_quant": scale_x_orig_quant_t.astype(np.float32), - "scale_w_quant_orig": scale_w_quant_orig_t.astype(np.float32), - "scale_w_quant_orig.col": scale_w_quant_orig_c.astype(np.float32), - "scale_y_accum_quant": scale_y_accum_quant_t.astype(np.float32), - "scale_y_accum_quant.col": scale_y_accum_quant_c.astype(np.float32), - "scale_y_quant_orig": scale_y_quant_orig_t.astype(np.float32), - } - - -@torch.no_grad() -def capture_activation_range( - model: AutoModelForCausalLM, - tokenizer: AutoTokenizer, - dataset: Dataset, - num_samples: int = 512, -): - model.eval() - device = next(model.parameters()).device - act_scales: defaultdict[Any, Dict[str, - torch.Tensor]] = defaultdict(lambda: { - "x": None, - "y": None, - "w": None - }) - - test_token_num = 923 - tokenizer.pad_token = tokenizer.eos_token - - def stat_tensor( - name: str, - tensor: torch.Tensor, - act_scales: defaultdict[Any, Dict[str, torch.Tensor]], - key: str, - ): - hidden_dim = tensor.shape[-1] - tensor = tensor.view(-1, hidden_dim).abs().detach() - comming_max = torch.max(tensor, dim=0)[0].float() - - if act_scales[name][key] is None: - act_scales[name][key] = comming_max - else: - act_scales[name][key] = torch.max(act_scales[name][key], - comming_max) - - def stat_input_hook( - m: nn.Module, - x: torch.Tensor, - y: torch.Tensor, - name: str, - ): - if isinstance(x, tuple): - x = x[0] - stat_tensor(name, x, act_scales, "x") - stat_tensor(name, y, act_scales, "y") - - if act_scales[name]["w"] is None: - act_scales[name]["w"] = m.weight.abs().clip(1e-8, - None).max(dim=1)[0] - - hooks = [] - for name, m in model.named_modules(): - if isinstance(m, nn.Linear) or isinstance(m, Conv1D): - hooks.append( - m.register_forward_hook( - functools.partial(stat_input_hook, name=name))) - - for i in tqdm(range(num_samples), desc="calibrating model"): - datapoint = dataset[i:i + 1] - line = copy.copy(datapoint) - line[0] = line[0] + ' TL;DR: ' - line[0] = line[0].strip() - line[0] = line[0].replace(" n't", "n't") - line_encoded = tokenizer(line, - return_tensors="pt", - max_length=test_token_num, - padding=True, - truncation=True).input_ids.to(device) - model(line_encoded) - - for h in hooks: - h.remove() - - return act_scales - - -@torch.no_grad() -def smooth_baichuan_model( - model: AutoModelForCausalLM, - scales: Dict[Any, Dict[str, torch.Tensor]], - alpha: float, - baichuan_smoother: Dict[str, torch.Tensor], -): - # Smooth the activation and weights with smoother = $\diag{s}$ - for name, module in model.named_modules(): - class_name = module.__class__.__name__ - if 'Layer' not in class_name: - continue - print(f'smoothing module: {name}, class_name: {class_name}') - # qkv_proj - layer_name_qkv = name + ".self_attn.W_pack" - - smoother = smooth_gemm(module.self_attn.W_pack.weight, - scales[layer_name_qkv]["x"], - module.input_layernorm.weight, None, alpha) - - scales[layer_name_qkv]["x"] = scales[layer_name_qkv]["x"] / smoother - scales[layer_name_qkv]["w"] = module.self_attn.W_pack.weight.abs().max( - dim=1)[0].float() - - # ================================================================= - layer_name = name + ".self_attn.o_proj" - smoother = smooth_gemm(module.self_attn.o_proj.weight, - scales[layer_name]["x"], None, None, alpha) - baichuan_smoother[layer_name] = smoother.float() - - scales[layer_name]["x"] = scales[layer_name]["x"] / smoother - scales[layer_name]["w"] = module.self_attn.o_proj.weight.abs().max( - dim=1)[0].float() - - # ================================================================== - fc1_layer_name = name + ".mlp.gate_proj" - gate_layer_name = name + ".mlp.up_proj" - - smoother = smooth_gemm_fc1_gate(module.mlp.gate_proj.weight, - module.mlp.up_proj.weight, - scales[fc1_layer_name]["x"], - module.post_attention_layernorm.weight, - None, alpha) - - scales[fc1_layer_name]["x"] = scales[fc1_layer_name]["x"] / smoother - scales[fc1_layer_name]["w"] = module.mlp.gate_proj.weight.abs().max( - dim=1)[0].float() - - scales[gate_layer_name]["x"] = scales[gate_layer_name]["x"] / smoother - scales[gate_layer_name]["w"] = module.mlp.up_proj.weight.abs().max( - dim=1)[0].float() - - # ================================================================== - layer_name = name + ".mlp.down_proj" - smoother = smooth_gemm(module.mlp.down_proj.weight, - scales[layer_name]["x"], None, None, alpha) - baichuan_smoother[layer_name] = smoother.float() - scales[layer_name]["x"] = scales[layer_name]["x"] / smoother - scales[layer_name]["w"] = module.mlp.down_proj.weight.abs().max( - dim=1)[0].float() - - -def get_tllm_linear_sq_weight( - vals: Dict[str, np.ndarray], - prefix: str, - shape: List[int], - tensor_parallel: int, - quant_mode: QuantMode, - is_qkv: bool = False, - last_prefix: Optional[str] = None, - bias: Optional[torch.Tensor] = None, - smoother_value=None, - smoother_shape=None, - rank: int = 0, - cat_dim: int = 0, - multi_query_mode: bool = False, -): - per_token = quant_mode.has_per_token_dynamic_scaling() - per_channel = quant_mode.has_per_channel_scaling() - results: Dict[str, torch.Tensor] = {} - - def multi_query_split( - data: np.array, - start: int, - size: int, - tp_size: int, - cur_rank: int, - ): - q, k, v = np.split(data, [start, start + size], axis=-1) - q_split = np.split(q, tp_size, axis=-1) - k_split = np.split(k, tp_size, axis=-1) - v_split = np.split(v, tp_size, axis=-1) - return [ - np.concatenate((q_split[ii], k_split[ii], v_split[ii]), axis=-1) - for ii in range(tp_size) - ][cur_rank] - - col_shape = shape if (is_qkv or per_channel) else [1, 1] - - if per_token: - if per_channel: - original_weights = np.array(vals["weight.int8.col"]) - else: - original_weights = np.array(vals["weight.int8"]) - local_dim = original_weights.shape[0] - head_size = (original_weights.shape[1] - local_dim) // 2 - - if multi_query_mode: - cur_weights = multi_query_split(original_weights, local_dim, - head_size, tensor_parallel, rank) - else: - cur_weights = np.split(original_weights, - tensor_parallel, - axis=cat_dim)[rank] - if is_qkv: - hidden_dim = cur_weights.shape[0] - cur_weights = cur_weights.reshape(hidden_dim, -1) - results[prefix + - 'weight'] = torch.from_numpy(cur_weights).t().contiguous() - if smoother_value is None: - results[last_prefix] = torch.from_numpy( - np.array([1.0], dtype=np.float32)) - - if per_channel: - cur_per_channel_value = vals["scale_w_quant_orig.col"] - if smoother_value is None: - if multi_query_mode: - cur_per_channel_value = multi_query_split( - vals["scale_w_quant_orig.col"], local_dim, head_size, - tensor_parallel, rank) - else: - cur_per_channel_value = np.split( - vals["scale_w_quant_orig.col"], - tensor_parallel, - axis=cat_dim)[rank] - else: - cur_per_channel_value = vals["scale_w_quant_orig"] - if is_qkv: - if multi_query_mode: - cur_per_channel_value = multi_query_split( - vals["scale_w_quant_orig"], local_dim, head_size, - tensor_parallel, rank) - else: - cur_per_channel_value = np.split(vals["scale_w_quant_orig"], - tensor_parallel, - axis=cat_dim)[rank] - - results[prefix + 'per_channel_scale'] = torch.from_numpy( - np.array(cur_per_channel_value, - dtype=np.float32).reshape(col_shape)).contiguous() - else: - if per_channel: - original_weights = np.array(vals["weight.int8.col"]) - else: - original_weights = np.array(vals["weight.int8"]) - local_dim = original_weights.shape[0] - head_size = (original_weights.shape[1] - local_dim) // 2 - - if multi_query_mode: - cur_weights = multi_query_split(original_weights, local_dim, - head_size, tensor_parallel, rank) - else: - cur_weights = np.split(original_weights, - tensor_parallel, - axis=cat_dim)[rank] - if is_qkv: - hidden_dim = cur_weights.shape[0] - cur_weights = cur_weights.reshape(hidden_dim, -1) - results[prefix + - 'weight'] = torch.from_numpy(cur_weights).t().contiguous() - - if per_channel: - cur_per_channel_value = vals["scale_y_accum_quant.col"] - if smoother_value is None: - if multi_query_mode: - cur_per_channel_value = multi_query_split( - vals["scale_y_accum_quant.col"], local_dim, head_size, - tensor_parallel, rank) - else: - cur_per_channel_value = np.split( - vals["scale_y_accum_quant.col"], - tensor_parallel, - axis=cat_dim)[rank] - else: - cur_per_channel_value = vals["scale_y_accum_quant"] - # QKV is always per_channel - if is_qkv: - if multi_query_mode: - cur_per_channel_value = multi_query_split( - vals["scale_y_accum_quant"], local_dim, head_size, - tensor_parallel, rank) - else: - cur_per_channel_value = np.split( - vals["scale_y_accum_quant"], - tensor_parallel, - axis=cat_dim)[rank] - - results[prefix + 'per_channel_scale'] = torch.from_numpy( - np.array([cur_per_channel_value], - dtype=np.float32).reshape(col_shape)).contiguous() - - results[last_prefix] = torch.from_numpy( - np.array([vals['scale_x_orig_quant']], - dtype=np.float32)).contiguous() - - results[prefix + 'act_scale'] = torch.from_numpy( - np.array([[vals["scale_y_quant_orig"]]], - dtype=np.float32)).contiguous() - - if smoother_value is not None: - cur_smoother_value = np.split(smoother_value, - tensor_parallel, - axis=cat_dim)[rank] - results[prefix + 'smoother'] = cur_smoother_value.reshape( - smoother_shape).contiguous().to(torch.float32) - - if bias is not None: - results[prefix + 'bias'] = bias - - return results - - -def split(weight: torch.Tensor, - tp_size: int, - rank: int = 0, - dim: int = 0) -> torch.Tensor: - if tp_size == 1: - return weight - elif weight.ndim == 1: - return torch.chunk(weight, tp_size)[rank].contiguous() - else: - return torch.chunk(weight, tp_size, dim=dim)[rank].contiguous() - - -def split_matrix(weight: torch.Tensor, tp_size: int, rank: int, - dim: int) -> torch.Tensor: - return split(weight, tp_size, rank, dim=dim) - - -def get_weight(params: Dict[str, torch.Tensor], prefix: str, - dtype: torch.dtype) -> Optional[torch.Tensor]: - if f'{prefix}.weight' not in params: - return None - return params[f'{prefix}.weight'].to(dtype).detach().cpu() - - -def quantize(hf_model_dir: str, - output_dir: str, - config: BaichuanConfig, - device: str = 'cuda', - calib_dataset: str = 'ccdv/cnn_dailymail', - trust_remote_code: bool = True): - os.makedirs(output_dir, exist_ok=True) - config.to_json_file(os.path.join(output_dir, 'config.json')) - - mapping = config.mapping - assert mapping.rank == 0, "quantize should be called at rank 0 only" - - hf_model = AutoModelForCausalLM.from_pretrained( - hf_model_dir, - device_map='auto' if device != 'cpu' else 'cpu', - dtype='auto' - if not config.quantization._use_plugin_sq else torch.float16, - trust_remote_code=trust_remote_code) - tokenizer = AutoTokenizer.from_pretrained( - hf_model_dir, use_fast=False, trust_remote_code=trust_remote_code) - dataset = load_calib_dataset(calib_dataset) - - act_range = capture_activation_range(hf_model, tokenizer, dataset) - smoother = {} - if config.quantization._use_plugin_sq: - smooth_baichuan_model(hf_model, act_range, - config.quantization.smoothquant_val, smoother) - - quant_mode = config.quantization.quant_mode - model_params = dict(hf_model.named_parameters()) - dtype = getattr(torch, config.dtype) - num_attention_heads = config.num_attention_heads - hidden_size = config.hidden_size - inter_size = config.intermediate_size - num_key_value_heads = config.num_key_value_heads - multi_query_mode = (num_key_value_heads != num_attention_heads) - - for rank in range(config.mapping.world_size): - # To avoid changing the mapping arg in-place, also the given mapping from caller is rank agnostic, since quantize is called from only one rank - config = copy.deepcopy(config) - config.set_rank(rank) - weights = {} - tensor_parallel = config.mapping.tp_size - - for l in config.mapping.pp_layers(config.num_hidden_layers): - prefix = f'model.layers.{l}.' - tllm_prex = f'transformer.layers.{l}.' - - # self_attn.W_pack -> attention.qkv - qkv_weight = get_weight(model_params, prefix + 'self_attn.W_pack', - dtype) - qkv_weight = qkv_weight.t().numpy() - qkv_out_dim = qkv_weight.shape[1] - if not multi_query_mode: - qkv_weight = qkv_weight.reshape(hidden_size, 3, hidden_size) - int8_weights = generate_int8(qkv_weight, - act_range.get(prefix + - 'self_attn.W_pack'), - is_qkv=True, - multi_query_mode=multi_query_mode) - weights.update( - get_tllm_linear_sq_weight(int8_weights, - tllm_prex + 'attention.qkv.', - [1, qkv_out_dim // tensor_parallel], - tensor_parallel, - quant_mode, - is_qkv=True, - last_prefix=tllm_prex + - 'input_layernorm.scale_to_int', - smoother_value=None, - smoother_shape=None, - rank=rank, - cat_dim=-1, - multi_query_mode=multi_query_mode)) - - if config.quantization.kv_cache_quant_algo == QuantAlgo.INT8: - qkv_weight = get_weight(model_params, - prefix + 'self_attn.W_pack', dtype) - qkv_weight = qkv_weight.t().numpy() - if not multi_query_mode: - qkv_weight = qkv_weight.reshape(hidden_size, 3, hidden_size) - int8_weights = generate_int8(qkv_weight, - act_range.get(prefix + - 'self_attn.W_pack'), - is_qkv=True, - multi_query_mode=multi_query_mode) - weights[tllm_prex + - 'attention.kv_cache_scaling_factor'] = torch.from_numpy( - np.array([int8_weights['scale_y_quant_orig']], - dtype=np.float32)).contiguous() - - # attn.out_proj -> attention.dense - attn_dense_weight = get_weight(model_params, - prefix + 'self_attn.o_proj', dtype) - attn_dense_weight = attn_dense_weight.t().numpy() - int8_weights = generate_int8( - attn_dense_weight, act_range.get(prefix + 'self_attn.o_proj')) - weights.update( - get_tllm_linear_sq_weight( - int8_weights, - tllm_prex + 'attention.dense.', [1, hidden_size], - tensor_parallel, - quant_mode, - is_qkv=False, - last_prefix=tllm_prex + - 'attention.quantization_scaling_factor', - smoother_value=smoother[(prefix + 'self_attn.o_proj')], - smoother_shape=[1, hidden_size // tensor_parallel], - rank=rank, - cat_dim=0)) - - # mlp.gate_proj -> mlp.fc - mlp_fc_weight = get_weight(model_params, prefix + 'mlp.gate_proj', - dtype) - mlp_fc_weight = mlp_fc_weight.t().numpy() - int8_weights = generate_int8( - mlp_fc_weight, act_range.get(prefix + 'mlp.gate_proj')) - weights.update( - get_tllm_linear_sq_weight( - int8_weights, - tllm_prex + 'mlp.fc.', [1, inter_size // tensor_parallel], - tensor_parallel, - quant_mode, - is_qkv=False, - last_prefix=tllm_prex + 'post_layernorm.scale_to_int', - smoother_value=None, - smoother_shape=None, - rank=rank, - cat_dim=-1)) - - # mlp.down_proj -> mlp.proj - mlp_proj_weight = get_weight(model_params, prefix + 'mlp.down_proj', - dtype) - mlp_proj_weight = mlp_proj_weight.t().numpy() - int8_weights = generate_int8( - mlp_proj_weight, act_range.get(prefix + 'mlp.down_proj')) - weights.update( - get_tllm_linear_sq_weight( - int8_weights, - tllm_prex + 'mlp.proj.', [1, hidden_size], - tensor_parallel, - quant_mode, - is_qkv=False, - last_prefix=tllm_prex + 'mlp.quantization_scaling_factor', - smoother_value=smoother[prefix + 'mlp.down_proj'], - smoother_shape=[1, inter_size // tensor_parallel], - rank=rank, - cat_dim=0)) - - # mlp.up_proj -> mlp.gate - mlp_gate_weight = get_weight(model_params, prefix + 'mlp.up_proj', - dtype) - mlp_gate_weight = mlp_gate_weight.t().numpy() - int8_weights = generate_int8(mlp_gate_weight, - act_range.get(prefix + 'mlp.up_proj')) - weights.update( - get_tllm_linear_sq_weight( - int8_weights, - tllm_prex + 'mlp.gate.', [1, inter_size // tensor_parallel], - tensor_parallel, - quant_mode, - is_qkv=False, - last_prefix=tllm_prex + 'post_layernorm.scale_to_int', - smoother_value=None, - smoother_shape=None, - rank=rank, - cat_dim=-1)) - - # input layer_norm - input_ln_weight = get_weight(model_params, - prefix + 'input_layernorm', dtype) - weights[tllm_prex + 'input_layernorm.weight'] = input_ln_weight - - # post layer_norm - post_ln_weight = get_weight(model_params, - prefix + 'post_attention_layernorm', - dtype) - weights[tllm_prex + 'post_layernorm.weight'] = post_ln_weight - - embed_w = get_weight(model_params, 'model.embed_tokens', dtype) - if config.mapping.is_first_pp_rank(): - # Embedding - if config.use_parallel_embedding: - embed_w = split_matrix(embed_w, - config.mapping.tp_size, - config.mapping.tp_rank, - dim=config.embedding_sharding_dim) - weights['transformer.vocab_embedding.weight'] = embed_w - - lm_head_w = get_weight(model_params, 'lm_head', dtype) - if config.mapping.is_last_pp_rank(): - # lm_head weight and bias - weights['lm_head.weight'] = split_matrix(lm_head_w.clone(), - config.mapping.tp_size, - config.mapping.tp_rank, - dim=0) - ln_f_w = get_weight(model_params, 'model.norm', dtype) - # ln_f weight and bias - weights['transformer.ln_f.weight'] = ln_f_w - - safetensors.torch.save_file( - weights, os.path.join(output_dir, f'rank{rank}.safetensors')) - del weights - - -def load_weights_from_hf_model(hf_model: AutoModelForCausalLM, - config: BaichuanConfig): - weights = {} - tik = time.time() - - model_params = dict(hf_model.named_parameters()) - dtype = getattr(torch, config.dtype) - num_hidden_layers = config.num_hidden_layers - hf_key = [ - "model.embed_tokens.weight", # vocab_embedding - "lm_head.weight", # lm_head - "model.norm.weight", # ln_f - "self_attn.W_pack.weight", # attention.qkv - "self_attn.o_proj.weight", # attention.dense - "mlp.up_proj.weight", # mlp.gate - "mlp.down_proj.weight", # mlp.proj - "mlp.gate_proj.weight", # mlp.fc - "input_layernorm.weight", # input_layernorm - "post_attention_layernorm.weight", # post_layernorm - ] - - def load(key_id: int, layer_idx: int = -1, tp_dim: int = -1): - layer_prefix = "" if layer_idx == -1 else f"model.layers.{layer_idx}." - v: torch.Tensor = model_params[layer_prefix + hf_key[key_id]] - if key_id == 3: - q_emb = v.shape[0] // 3 - model_emb = v.shape[1] - v = v.reshape(3, q_emb, model_emb) - if v.shape[1] % config.mapping.tp_size != 0: - logger.error( - "Current weight shape is invalid for mapping.tp_size=" + - str(config.mapping.tp_size)) - v = v.split(v.shape[1] // config.mapping.tp_size, - dim=1)[config.mapping.tp_rank] - v = v.reshape(3 * (q_emb // config.mapping.tp_size), model_emb) - if tp_dim >= 0: - if v.shape[tp_dim] % config.mapping.tp_size != 0: - logger.error( - "Current weight shape is invalid for mapping.tp_size=" + - str(config.mapping.tp_size)) - v = v.split(v.shape[tp_dim] // config.mapping.tp_size, - dim=tp_dim)[config.mapping.tp_rank] - v = v.to(dtype).contiguous().detach().cpu() - return v - - # Convert vocab_embedding - if config.mapping.is_first_pp_rank(): - embed_w = load(0) - if config.use_parallel_embedding: - embed_w = split_matrix(embed_w, - config.mapping.tp_size, - config.mapping.tp_rank, - dim=config.embedding_sharding_dim) - weights['transformer.vocab_embedding.weight'] = embed_w - - # Convert lm_head - v = load(1, -1, 0) - if config.model_version.startswith('v2'): - v = torch.nn.functional.normalize(v) - if config.mapping.is_last_pp_rank(): - weights['lm_head.weight'] = v - - # Convert ln_f - if config.mapping.is_last_pp_rank(): - weights['transformer.ln_f.weight'] = load(2) - - # Convert layers - layers_range = config.mapping.pp_layers(num_hidden_layers) - for l in layers_range: - prefix = f"transformer.layers.{l}." - weights[prefix + 'attention.qkv.weight'] = load(3, l) - weights[prefix + 'attention.dense.weight'] = load(4, l, 1) - weights[prefix + 'mlp.gate.weight'] = load(5, l, 0) - weights[prefix + 'mlp.proj.weight'] = load(6, l, 1) - weights[prefix + 'mlp.fc.weight'] = load(7, l, 0) - weights[prefix + 'input_layernorm.weight'] = load(8, l) - weights[prefix + 'post_layernorm.weight'] = load(9, l) - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Weights loaded. Total time: {t}') - - return weight_only_quantize_dict(weights, - quant_algo=config.quantization.quant_algo, - plugin=True) - - -def load_weights_from_gptq(config: BaichuanConfig, quant_ckpt_path: str): - logger.info('Loading weights from groupwise GPTQ Baichuan safetensors...') - weights = {} - tik = time.time() - - gptq_baichuan = safetensors.safe_open(quant_ckpt_path, - framework="pt", - device=0) - gptq_prefix = "model." - gptq_suffix_list = [".qweight", ".qzeros", ".scales"] - gptq_key_list = [ - "embed_tokens.weight", # vocab_embedding - "lm_head.weight", # lm_head - "norm.weight", # ln_f - "self_attn.W_pack", # attention.qkv - "_proj", # - "self_attn.o_proj", # attention.dense - "mlp.up_proj", # mlp.gate - "mlp.down_proj", # mlp.proj - "mlp.gate_proj", # mlp.fc - "input_layernorm.weight", # input_layernorm - "post_attention_layernorm.weight", # post_layernorm - ] - - packer = torch.ops.trtllm.pack_int8_tensor_to_packed_int4 - preprocessor = torch.ops.trtllm.preprocess_weights_for_mixed_gemm - torch_dtype = getattr(torch, config.dtype) - - def load(key: str, no_prefix: bool = False): - if no_prefix: - return gptq_baichuan.get_tensor(key) - else: - return gptq_baichuan.get_tensor(gptq_prefix + key) - - def torch_split(tensor: torch.Tensor, dim: int): - if tensor.shape[dim] % config.mapping.tp_size != 0: - logger.error( - "Current weight shape is invalid for mapping.tp_size=" + - str(config.mapping.tp_size)) - assert False, "Invalid TP size" - return tensor.split(tensor.shape[dim] // config.mapping.tp_size, - dim=dim)[config.mapping.tp_rank] - - def unpack_int32_into_int8(w_packed: torch.Tensor): - # Unpack inputs packed in int32/float32 into uint4 and store them in int8 format - w_packed_int4x2 = w_packed.contiguous().view(torch.uint8) - w_unpacked = torch.zeros(w_packed_int4x2.shape[0], - w_packed_int4x2.shape[1] * 2, - dtype=torch.int8) - w_unpacked[:, ::2] = w_packed_int4x2 % 16 - w_unpacked[:, 1::2] = w_packed_int4x2 // 16 - return w_unpacked.contiguous() - - def process_and_assign_weight(prefix: str, - tensors: List[torch.Tensor], - tp_dim: int = -1): - if tp_dim == -1: - qweight_int32, qzeros_int32, scales_fp16 = [ - item.cpu() for item in tensors - ] - else: - qweight_int32, qzeros_int32, scales_fp16 = [ - torch_split(item, tp_dim).cpu() for item in tensors - ] - - USE_UINT4_INPUT = 1 # Set to true if checkpoint store UINT4 weights - USE_GPTQ_FOR_LLAMA = 1 # GPTQ-for-LLaMA added 1 to zeros - - qweight_unpacked_int8 = unpack_int32_into_int8( - qweight_int32.T).T.contiguous() - 8 - qweight_interleaved = preprocessor(packer(qweight_unpacked_int8), - torch.quint4x2, - torch.float16).view(torch.float16) - # zeros = zeros * scales - qzeros_unpacked_int32 = unpack_int32_into_int8(qzeros_int32) - zeros_x_scales_fp16 = (-qzeros_unpacked_int32 + 8 * USE_UINT4_INPUT - - USE_GPTQ_FOR_LLAMA) * scales_fp16 - zeros_x_scales_fp16 = zeros_x_scales_fp16.half() - - # return processed interleaved weight, original scales and zeros * scales - weights[prefix + ".weight"] = qweight_interleaved - weights[prefix + ".weights_scaling_factor"] = scales_fp16 - weights[prefix + ".zero"] = zeros_x_scales_fp16 - - # Load weights from GPTQ checkpoint into TRT-LLM module - # 1. vocab_embedding - v = load(gptq_key_list[0]) - if config.mapping.is_first_pp_rank(): - if config.use_parallel_embedding: - v = split_matrix(v, - config.mapping.tp_size, - config.mapping.tp_rank, - dim=config.embedding_sharding_dim) - weights['transformer.vocab_embedding.weight'] = v.to(torch_dtype) - - # 2. lm_head - original_v = load(gptq_key_list[1], True) - if config.model_version.startswith('v2'): - # baichuan v2 models use NormHead - logger.info(f'Normalizing lm_head.weight for {config.model_version}') - v = torch_split(torch.nn.functional.normalize(original_v), 0) - else: - v = torch_split(original_v, 0) - if config.mapping.is_last_pp_rank(): - weights['lm_head.weight'] = v.to(torch_dtype) - - # 3. ln_f - v = load(gptq_key_list[2]) - if config.mapping.is_last_pp_rank(): - weights['transformer.ln_f.weight'] = v.to(torch_dtype) - - # 4. Weights inside each layer - num_hidden_layers = config.num_hidden_layers - layers_range = config.mapping.pp_layers(num_hidden_layers) - for l in layers_range: - layer_idx = l - layers_range[0] - hf_prefix = f"layers.{l}." - tllm_prefix = f"transformer.layers.{l}." - logger.info(f'Process weights in layer: {layer_idx}') - - # 4.1 attention.qkv - qkv_weight_list = [] - for suf in gptq_suffix_list: - qkv_list = [] - comp_part = load(hf_prefix + gptq_key_list[3] + suf) - qkv = torch.chunk(comp_part, 3, 1) - for i in range(3): - comp_part = qkv[i] - comp_part = torch_split(comp_part, 1) - qkv_list.append(comp_part) - qkv_weight_list.append(torch.cat(qkv_list, dim=1)) - - process_and_assign_weight(tllm_prefix + "attention.qkv", - qkv_weight_list) - - # 4.2 attention.dense - v = [ - load(hf_prefix + gptq_key_list[5] + suf) for suf in gptq_suffix_list - ] - process_and_assign_weight(tllm_prefix + "attention.dense", v, 0) - - # 4.3 mlp.gate - v = [ - load(hf_prefix + gptq_key_list[6] + suf) for suf in gptq_suffix_list - ] - process_and_assign_weight(tllm_prefix + "mlp.gate", v, 1) - - # 4.4 mlp.proj - v = [ - load(hf_prefix + gptq_key_list[7] + suf) for suf in gptq_suffix_list - ] - process_and_assign_weight(tllm_prefix + "mlp.proj", v, 0) - - # 4.5 mlp.fc - v = [ - load(hf_prefix + gptq_key_list[8] + suf) for suf in gptq_suffix_list - ] - process_and_assign_weight(tllm_prefix + "mlp.fc", v, 1) - - # 4.6 input_layernorm - v = load(hf_prefix + gptq_key_list[9]) - weights[tllm_prefix + 'input_layernorm.weight'] = v.to(torch_dtype) - - # 4.7 pst_layernorm - v = load(hf_prefix + gptq_key_list[10]) - weights[tllm_prefix + 'post_layernorm.weight'] = v.to(torch_dtype) - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - logger.info(f'Weights loaded. Total time: {t}') - return weights diff --git a/tensorrt_llm/models/baichuan/model.py b/tensorrt_llm/models/baichuan/model.py deleted file mode 100644 index 00dab61b69ec..000000000000 --- a/tensorrt_llm/models/baichuan/model.py +++ /dev/null @@ -1,253 +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. -from typing import Optional, Union - -from ..._utils import pad_vocab_size -from ...functional import Tensor -from ...layers import (Attention, AttentionMaskType, ColumnLinear, Embedding, - GatedMLP, RmsNorm) -from ...mapping import Mapping -from ...module import Module -from ..modeling_utils import (DecoderLayerList, DecoderModelForCausalLM, - PretrainedConfig, QuantConfig) -from .config import BaichuanConfig -from .convert import load_weights_from_hf_model - - -class BaichuanDecoderLayer(Module): - - def __init__(self, config: PretrainedConfig, layer_idx): - super().__init__() - self.layer_idx = layer_idx - self.config = config - hidden_size = config.hidden_size - dtype = config.dtype - position_embedding_type = config.position_embedding_type - tp_group = config.mapping.tp_group - tp_size = config.mapping.tp_size - tp_rank = config.mapping.tp_rank - quant_mode = config.quant_mode - - self.input_layernorm = RmsNorm(normalized_shape=hidden_size, - eps=config.norm_epsilon, - dtype=dtype) - layers_range = config.mapping.pp_layers(config.num_hidden_layers) - local_layer_idx = layer_idx - layers_range[0] - self.attention = Attention( - local_layer_idx=local_layer_idx, - hidden_size=hidden_size, - num_attention_heads=config.num_attention_heads, - num_kv_heads=config.num_key_value_heads, - max_position_embeddings=config.max_position_embeddings, - dtype=dtype, - attention_mask_type=AttentionMaskType.causal, - bias=False, - position_embedding_type=position_embedding_type, - tp_group=tp_group, - tp_size=tp_size, - tp_rank=tp_rank, - quant_mode=quant_mode) - - self.mlp = GatedMLP(hidden_size=hidden_size, - ffn_hidden_size=config.intermediate_size, - hidden_act=config.hidden_act, - dtype=dtype, - bias=False, - tp_group=tp_group, - tp_size=tp_size, - quant_mode=quant_mode) - self.post_layernorm = RmsNorm(normalized_shape=hidden_size, - eps=config.norm_epsilon, - dtype=dtype) - - def forward(self, - hidden_states: Tensor, - attention_mask=None, - use_cache=False, - kv_cache_params=None, - attention_params=None): - residual = hidden_states - hidden_states = self.input_layernorm(hidden_states) - - attention_output = self.attention(hidden_states, - attention_mask=attention_mask, - use_cache=use_cache, - kv_cache_params=kv_cache_params, - attention_params=attention_params) - - if use_cache: - attention_output, presents = attention_output - - hidden_states = residual + attention_output - - residual = hidden_states - hidden_states = self.post_layernorm(hidden_states) - - hidden_states = self.mlp(hidden_states) - - hidden_states = residual + hidden_states - if use_cache: - return (hidden_states, presents) - return hidden_states - - -class BaichuanModel(Module): - - def __init__(self, config: PretrainedConfig): - super().__init__() - hidden_size = config.hidden_size - - self.vocab_embedding = Embedding(config.vocab_size, - config.hidden_size, - dtype=config.dtype) - - self.layers = DecoderLayerList(BaichuanDecoderLayer, config) - self.ln_f = RmsNorm(normalized_shape=hidden_size, - eps=config.norm_epsilon, - dtype=config.dtype) - - def forward(self, - input_ids: Tensor, - position_ids=None, - use_cache=False, - attention_mask=None, - kv_cache_params=None, - attention_params=None, - prompt_embedding_table=None, - prompt_tasks=None, - prompt_vocab_size=None): - args = [prompt_embedding_table, prompt_tasks, prompt_vocab_size - ] if prompt_embedding_table is not None else [] - hidden_states = self.vocab_embedding(input_ids, *args) - - hidden_states = self.layers(hidden_states, - use_cache=use_cache, - attention_mask=attention_mask, - kv_cache_params=kv_cache_params, - attention_params=attention_params) - - if use_cache: - hidden_states, presents = hidden_states - - hidden_states = self.ln_f(hidden_states) - - if use_cache: - return (hidden_states, tuple(presents)) - return hidden_states - - -class BaichuanForCausalLM(DecoderModelForCausalLM): - config_class = BaichuanConfig - - def __init__(self, config: PretrainedConfig): - transformer = BaichuanModel(config) - vocab_size_padded = pad_vocab_size(config.vocab_size, - config.mapping.tp_size) - lm_head = ColumnLinear(config.hidden_size, - vocab_size_padded, - bias=False, - dtype=config.dtype, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size, - gather_output=True) - super().__init__(config, transformer, lm_head) - - @classmethod - def from_hugging_face( - cls, - hf_model_or_dir: Union[str, 'transformers.PreTrainedModel'], - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - ''' Create a BaichuanForCausalLM object from give parameters - ''' - import transformers - - assert hf_model_or_dir is not None - if isinstance(hf_model_or_dir, transformers.PreTrainedModel): - hf_model = hf_model_or_dir - hf_config_or_dir = hf_model.config - else: - trust_remote_code = kwargs.pop('trust_remote_code', True) - - hf_model = transformers.AutoModelForCausalLM.from_pretrained( - hf_model_or_dir, - trust_remote_code=trust_remote_code, - dtype='auto') - hf_config_or_dir = hf_model_or_dir - - config = BaichuanConfig.from_hugging_face(hf_config_or_dir, - dtype=dtype, - mapping=mapping, - quant_config=quant_config, - **kwargs) - - weights = load_weights_from_hf_model(hf_model, config) - - model = cls(config) - model.load(weights) - return model - - @classmethod - def quantize( - cls, - hf_model_dir: str, - output_dir: str, - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - *, - device: str = 'cuda', - calib_dataset: str = 'cnn_dailymail', - calib_batches: int = 512, - calib_batch_size: int = 1, - calib_max_seq_length: int = 512, - random_seed: int = 1234, - tokenizer_max_seq_length: int = 2048, - **kwargs, - ): - if quant_config._requires_modelopt_quantization: - # modelopt quantization flow - super().quantize(hf_model_dir, - output_dir, - dtype=dtype, - mapping=mapping, - quant_config=quant_config, - device=device, - calib_dataset=calib_dataset, - calib_batches=calib_batches, - calib_batch_size=calib_batch_size, - calib_max_seq_length=calib_max_seq_length, - random_seed=random_seed, - tokenizer_max_seq_length=tokenizer_max_seq_length) - elif quant_config._requires_calibration: - # non-modelopt quantization flow - from .convert import quantize - - config = BaichuanConfig.from_hugging_face(hf_model_dir, - dtype=dtype, - mapping=mapping, - quant_config=quant_config, - **kwargs) - quantize(hf_model_dir, - output_dir, - config=config, - device=device, - calib_dataset=calib_dataset) - else: - raise ValueError( - f"The quant_config ({quant_config}) does not require calibration, try {cls.__name__}.from_hugging_face instead." - ) diff --git a/tensorrt_llm/models/bert/__init__.py b/tensorrt_llm/models/bert/__init__.py deleted file mode 100644 index 71bf6d298d42..000000000000 --- a/tensorrt_llm/models/bert/__init__.py +++ /dev/null @@ -1,14 +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. diff --git a/tensorrt_llm/models/bert/config.py b/tensorrt_llm/models/bert/config.py deleted file mode 100644 index 5c2627aa7e23..000000000000 --- a/tensorrt_llm/models/bert/config.py +++ /dev/null @@ -1,117 +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. -from typing import Optional, Union - -import torch -import transformers - -from ..._utils import torch_dtype_to_str -from ...mapping import Mapping -from ..modeling_utils import PretrainedConfig, QuantConfig - - -class BERTConfig(PretrainedConfig): - - def __init__(self, - *, - is_roberta: bool = False, - type_vocab_size, - pad_token_id=None, - num_labels=None, - **kwargs): - self.is_roberta = is_roberta - self.type_vocab_size = type_vocab_size - self.pad_token_id = pad_token_id - self.num_labels = num_labels - - super().__init__(**kwargs) - - def to_dict(self): - output = super().to_dict() - output['is_roberta'] = self.is_roberta - output['type_vocab_size'] = self.type_vocab_size - output['pad_token_id'] = self.pad_token_id - output['num_labels'] = self.num_labels - - return output - - @classmethod - def from_hugging_face( - cls, - hf_config_or_dir: Union[str, 'transformers.PretrainedConfig'], - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - import transformers - - if isinstance(hf_config_or_dir, transformers.PretrainedConfig): - hf_config = hf_config_or_dir - else: - hf_config_dir = str(hf_config_or_dir) - - hf_config = transformers.AutoConfig.from_pretrained(hf_config_dir) - - num_key_value_heads = getattr(hf_config, "num_key_value_heads", - hf_config.num_attention_heads) - head_dim = getattr( - hf_config, "head_dim", - hf_config.hidden_size // hf_config.num_attention_heads) - head_size = getattr(hf_config, "kv_channels", head_dim) - num_labels = getattr(hf_config, "num_labels", None) - - if (hf_config.position_embedding_type == 'absolute'): - position_embedding_type = 'learned_absolute' - else: - raise NotImplementedError( - f"{hf_config.position_embedding_type} hasn't been supported") - - if hf_config.model_type == "bert": - is_roberta = False - else: - is_roberta = True - - if dtype == 'auto': - dtype = getattr(hf_config, 'torch_dtype', None) - if dtype is None: - dtype = 'float16' - if isinstance(dtype, torch.dtype): - dtype = torch_dtype_to_str(dtype) - if dtype == 'float32': - dtype = 'float16' - - return cls( - architecture=hf_config.architectures[0], - dtype=dtype, - hidden_size=hf_config.hidden_size, - num_hidden_layers=hf_config.num_hidden_layers, - num_attention_heads=hf_config.num_attention_heads, - vocab_size=hf_config.vocab_size, - hidden_act=hf_config.hidden_act, - logits_dtype='float32', - norm_epsilon=hf_config.layer_norm_eps, - position_embedding_type=position_embedding_type, - max_position_embeddings=hf_config.max_position_embeddings, - num_key_value_heads=num_key_value_heads, - intermediate_size=hf_config.intermediate_size, - head_size=head_size, - quantization=quant_config, - mapping=mapping, - #BERT model args - is_roberta=is_roberta, - type_vocab_size=hf_config.type_vocab_size, - pad_token_id=hf_config.pad_token_id, - num_labels=num_labels, - **kwargs) diff --git a/tensorrt_llm/models/bert/convert.py b/tensorrt_llm/models/bert/convert.py deleted file mode 100644 index 90fe2a62bfe6..000000000000 --- a/tensorrt_llm/models/bert/convert.py +++ /dev/null @@ -1,309 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -import time -from typing import Union - -import torch - -# isort: off -from transformers import (AutoModel, AutoModelForQuestionAnswering, - AutoModelForSequenceClassification) -from transformers import (BertPreTrainedModel, RobertaPreTrainedModel) -# isort: on -from ...logger import logger -from ..convert_utils import split, split_qkv_bias_tp, split_qkv_tp -from .config import BERTConfig - - -def extract_layer_idx(name): - ss = name.split('.') - for s in ss: - if s.isdigit(): - return s - return None - - -def _load_weights_from_hf_bert_model(hf_model: Union[BertPreTrainedModel, - RobertaPreTrainedModel], - model_config: BERTConfig, - torch_dtype: torch.dtype = torch.float16): - weights = {} - no_match = {} - mapping = model_config.mapping - # use different prefix because BertModel is used both individually and as part of model - trtllm_prefix = "" if (model_config.architecture - in ["BertModel", "RobertaModel"]) else "bert." - for k, v in hf_model.state_dict().items(): - key = None - v = v.to(torch_dtype).cpu() - if 'embeddings.word_embeddings.weight' in k: - key = f'{trtllm_prefix}embedding.vocab_embedding.weight' - elif 'embeddings.position_embeddings.weight' in k: - key = f'{trtllm_prefix}embedding.position_embedding.weight' - elif 'embeddings.token_type_embeddings.weight' in k: - key = f'{trtllm_prefix}embedding.token_embedding.weight' - elif 'embeddings.LayerNorm.weight' in k: - key = f'{trtllm_prefix}embedding.embedding_ln.weight' - elif 'embeddings.LayerNorm.bias' in k: - key = f'{trtllm_prefix}embedding.embedding_ln.bias' - else: - layer_idx = extract_layer_idx(k) - if layer_idx is None: - no_match[k] = v - continue - idx = int(layer_idx) - if 'attention.output.dense.weight' in k: - #TODO: add TP support - key = f'{trtllm_prefix}layers.{idx}.attention.dense.weight' - v_clone = v.clone() - v = split(v=v_clone, - tp_size=mapping.tp_size, - idx=mapping.tp_rank, - dim=1) - elif 'attention.output.dense.bias' in k: - key = f'{trtllm_prefix}layers.{idx}.attention.dense.bias' - elif 'attention.output.LayerNorm.weight' in k: - key = f'{trtllm_prefix}layers.{idx}.input_layernorm.weight' - elif 'attention.output.LayerNorm.bias' in k: - key = f'{trtllm_prefix}layers.{idx}.input_layernorm.bias' - elif 'intermediate.dense.weight' in k: - key = f'{trtllm_prefix}layers.{idx}.mlp.fc.weight' - v_clone = v.clone() - v = split(v=v_clone, - tp_size=mapping.tp_size, - idx=mapping.tp_rank, - dim=0) - elif 'intermediate.dense.bias' in k: - key = f'{trtllm_prefix}layers.{idx}.mlp.fc.bias' - v_clone = v.clone() - v = split(v=v_clone, - tp_size=mapping.tp_size, - idx=mapping.tp_rank, - dim=0) - elif 'output.dense.weight' in k: - key = f'{trtllm_prefix}layers.{idx}.mlp.proj.weight' - v_clone = v.clone() - v = split(v=v_clone, - tp_size=mapping.tp_size, - idx=mapping.tp_rank, - dim=1) - elif 'output.dense.bias' in k: - key = f'{trtllm_prefix}layers.{idx}.mlp.proj.bias' - elif 'output.LayerNorm.weight' in k: - key = f'{trtllm_prefix}layers.{idx}.post_layernorm.weight' - elif 'output.LayerNorm.bias' in k: - key = f'{trtllm_prefix}layers.{idx}.post_layernorm.bias' - elif 'attention.self.query.weight' in k: - key = f'{trtllm_prefix}layers.{idx}.attention.q.weight' - elif 'attention.self.query.bias' in k: - key = f'{trtllm_prefix}layers.{idx}.attention.q.bias' - elif 'attention.self.key.weight' in k: - key = f'{trtllm_prefix}layers.{idx}.attention.k.weight' - elif 'attention.self.key.bias' in k: - key = f'{trtllm_prefix}layers.{idx}.attention.k.bias' - elif 'attention.self.value.weight' in k: - key = f'{trtllm_prefix}layers.{idx}.attention.v.weight' - elif 'attention.self.value.bias' in k: - key = f'{trtllm_prefix}layers.{idx}.attention.v.bias' - else: - no_match[k] = v - continue - weights[key] = v - - for idx in range(model_config.num_hidden_layers): - qkv_key = f'{trtllm_prefix}layers.{idx}.attention.qkv' - q_key = f'{trtllm_prefix}layers.{idx}.attention.q' - k_key = f'{trtllm_prefix}layers.{idx}.attention.k' - v_key = f'{trtllm_prefix}layers.{idx}.attention.v' - for postfix in ['weight', 'bias']: - v = torch.cat( - (weights[f'{q_key}.{postfix}'], weights[f'{k_key}.{postfix}'], - weights[f'{v_key}.{postfix}']), - dim=0) - v_clone = v.clone() - split_v = v_clone - if postfix == 'weight': - split_v = split_qkv_tp(v_clone, - model_config.num_attention_heads, - model_config.hidden_size, - mapping.tp_size, mapping.tp_rank) - - elif postfix == 'bias': - split_v = split_qkv_bias_tp(v_clone, - model_config.num_attention_heads, - model_config.hidden_size, - mapping.tp_size, mapping.tp_rank) - else: - assert True, f"Unknown postfix={postfix}!" - #add qkv weight/bias - weights[f'{qkv_key}.{postfix}'] = split_v - #remove separate q, k , v - del weights[f'{q_key}.{postfix}'] - del weights[f'{k_key}.{postfix}'] - del weights[f'{v_key}.{postfix}'] - return (weights, no_match) - - -def _load_weights_from_hf_bert_qa_model( - hf_model: Union[BertPreTrainedModel, RobertaPreTrainedModel], - model_config: BERTConfig, - torch_dtype: torch.dtype = torch.float16): - weights, no_match = _load_weights_from_hf_bert_model( - hf_model, model_config, torch_dtype) - - weights['qa_outputs.weight'] = no_match['qa_outputs.weight'] - - weights['qa_outputs.bias'] = no_match['qa_outputs.bias'] - del no_match['qa_outputs.weight'] - del no_match['qa_outputs.bias'] - - return (weights, no_match) - - -def _load_weights_from_hf_bert_cls_model( - hf_model: Union[BertPreTrainedModel, RobertaPreTrainedModel], - model_config: BERTConfig, - torch_dtype: torch.dtype = torch.float16): - - weights, no_match = _load_weights_from_hf_bert_model( - hf_model, model_config, torch_dtype) - - if model_config.is_roberta: - # roberta Version - weights['classifier.dense.weight'] = no_match['classifier.dense.weight'] - weights['classifier.dense.bias'] = no_match['classifier.dense.bias'] - weights['classifier.out_proj.weight'] = no_match[ - 'classifier.out_proj.weight'] - weights['classifier.out_proj.bias'] = no_match[ - 'classifier.out_proj.bias'] - del no_match['classifier.dense.weight'] - del no_match['classifier.dense.bias'] - del no_match['classifier.out_proj.weight'] - del no_match['classifier.out_proj.bias'] - else: - weights['pooler.dense.weight'] = no_match['bert.pooler.dense.weight'] - weights['pooler.dense.bias'] = no_match['bert.pooler.dense.bias'] - weights['classifier.weight'] = no_match['classifier.weight'] - weights['classifier.bias'] = no_match['classifier.bias'] - del no_match['bert.pooler.dense.weight'] - del no_match['bert.pooler.dense.bias'] - del no_match['classifier.weight'] - del no_match['classifier.bias'] - - return (weights, no_match) - - -def load_hf_bert_base(model_dir: str, - load_model_on_cpu: bool = False, - dtype: torch.dtype = torch.float16): - """ - load huggingface BertModel and RobertaModel model - """ - model = AutoModel.from_pretrained( - model_dir, - trust_remote_code=True, - ) - if not load_model_on_cpu: - model.cuda().to(dtype) - model.eval() - return model - - -def load_hf_bert_qa(model_dir: str, - load_model_on_cpu: bool = False, - dtype: torch.dtype = torch.float16): - """ - load huggingface BertForQuestionAnswering and RobertaForQuestionAnswering - """ - model = AutoModelForQuestionAnswering.from_pretrained( - model_dir, - trust_remote_code=True, - ) - if not load_model_on_cpu: - model.cuda().to(dtype) - model.eval() - return model - - -def load_hf_bert_cls(model_dir: str, - load_model_on_cpu: bool = False, - dtype: torch.dtype = torch.float16): - """ - load huggingface BertForSequenceClassification and RobertaForSequenceClassification - """ - model = AutoModelForSequenceClassification.from_pretrained( - model_dir, - trust_remote_code=True, - ) - if not load_model_on_cpu: - model.cuda().to(dtype) - model.eval() - return model - - -def load_weights_from_hf_model( - hf_model, - config: BERTConfig, -): - """ - load trtllm weights from hf model - - return a dict of weights, with trtllm weights naming - - """ - #TODO: add quantization support - weights = {} - tik = time.time() - - torch_dtype = getattr(torch, config.dtype) - - #NOTE: Bert - no_match = None - if config.architecture in [ - "BertForQuestionAnswering", "RobertaForQuestionAnswering" - ]: - weights, no_match = _load_weights_from_hf_bert_qa_model( - hf_model=hf_model, model_config=config, torch_dtype=torch_dtype) - elif config.architecture in ["BertModel", "RobertaModel"]: - weights, no_match = _load_weights_from_hf_bert_model( - hf_model=hf_model, model_config=config, torch_dtype=torch_dtype) - elif config.architecture in [ - "BertForSequenceClassification", "RobertaForSequenceClassification" - ]: - weights, no_match = _load_weights_from_hf_bert_cls_model( - hf_model=hf_model, model_config=config, torch_dtype=torch_dtype) - else: - assert False, f"Unknown BERT model {config.architecture}" - - if no_match is not None: - logger.warning( - f"These weights from huggingface model are not used:\n {[key for key in no_match.keys()]}" - ) - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Weights loaded. Total time: {t}') - return weights - - -def quantize(hf_model_dir: str, - output_dir: str, - config: BERTConfig, - device: str = 'cuda', - calib_dataset: str = 'cnn_dailymail'): - ''' - Quantize the save the model as TRT-LLM checkpoint to output_dir - ''' - logger.warning(f"FP8 Support for Bert will come soon!") diff --git a/tensorrt_llm/models/bert/model.py b/tensorrt_llm/models/bert/model.py deleted file mode 100644 index a276ac8edf62..000000000000 --- a/tensorrt_llm/models/bert/model.py +++ /dev/null @@ -1,548 +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. -from typing import Optional, OrderedDict, Union - -import numpy as np -import tensorrt as trt -import torch -import transformers - -from tensorrt_llm.models.modeling_utils import PretrainedModel - -from ..._common import default_net -from ...functional import (ACT2FN, Tensor, concat, constant, cumsum, expand, - index_select, select, shape, slice, unsqueeze) -from ...layers import MLP, BertAttention, Embedding, LayerNorm, Linear -from ...mapping import Mapping -from ...module import Module, ModuleList -from ..modeling_utils import QuantConfig -from .config import BERTConfig -from .convert import (load_hf_bert_base, load_hf_bert_cls, load_hf_bert_qa, - load_weights_from_hf_model) - - -class BertEmbedding(Module): - - def __init__(self, - vocab_size, - hidden_size, - max_position_embeddings, - type_vocab_size, - dtype=None): - super().__init__() - self.vocab_embedding = Embedding(vocab_size, hidden_size, dtype=dtype) - self.position_embedding = Embedding(max_position_embeddings, - hidden_size, - dtype=dtype) - self.token_embedding = Embedding(type_vocab_size, - hidden_size, - dtype=dtype) - self.max_position_embeddings = max_position_embeddings - - self.embedding_ln = LayerNorm(normalized_shape=hidden_size, dtype=dtype) - - def forward(self, input_ids, position_ids, token_type_ids): - x = self.vocab_embedding(input_ids) - x = x + self.position_embedding(position_ids) - x = x + self.token_embedding(token_type_ids) - x = self.embedding_ln(x) - return x - - -class BertEncoderLayer(Module): - - def __init__(self, - hidden_size, - num_attention_heads, - max_position_embeddings, - hidden_act='relu', - tp_group=None, - tp_size=1, - dtype=None): - super().__init__() - self.input_layernorm = LayerNorm(normalized_shape=hidden_size, - dtype=dtype) - - self.attention = BertAttention( - hidden_size=hidden_size, - num_attention_heads=num_attention_heads, - max_position_embeddings=max_position_embeddings, - tp_group=tp_group, - tp_size=tp_size, - dtype=dtype) - self.mlp = MLP(hidden_size=hidden_size, - ffn_hidden_size=hidden_size * 4, - hidden_act=hidden_act, - tp_group=tp_group, - tp_size=tp_size, - dtype=dtype) - self.post_layernorm = LayerNorm(normalized_shape=hidden_size, - dtype=dtype) - - def forward(self, - hidden_states, - attention_mask=None, - input_lengths=None, - max_input_length=None): - residual = hidden_states - - attention_output = self.attention(hidden_states, - attention_mask=attention_mask, - input_lengths=input_lengths, - max_input_length=max_input_length) - - hidden_states = residual + attention_output - - hidden_states = self.input_layernorm(hidden_states) - - residual = hidden_states - - hidden_states = self.mlp(hidden_states) - - hidden_states = residual + hidden_states - - hidden_states = self.post_layernorm(hidden_states) - - return hidden_states - - -class BertBase(PretrainedModel): - ''' - Base class that provides from_huggingface() and prepare_inputs() methods - ''' - config_class = BERTConfig - - def __init__(self, config: BERTConfig): - super().__init__(config) - - @classmethod - def load_hf_bert(cls, model_dir: str, load_model_on_cpu: bool, - dtype: torch.dtype): - """ - Use as the abstractmethod, load corresponding HF model. - Subclass must implement this method! - """ - - assert cls.__name__ != "BertBase", f"Never call from BertBase class!" - - if cls.__name__ == "BertModel": - return load_hf_bert_base(model_dir, load_model_on_cpu, dtype) - elif cls.__name__ == "BertForQuestionAnswering": - return load_hf_bert_qa(model_dir, load_model_on_cpu, dtype) - elif cls.__name__ == "BertForSequenceClassification": - return load_hf_bert_cls(model_dir, load_model_on_cpu, dtype) - else: - assert False, f"Unknown class {cls.__name__}!" - - @classmethod - def from_hugging_face( - cls, - hf_model_or_dir: Union[str, 'transformers.PreTrainedModel'], - dtype: str = 'float16', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - """ - Create a BertModel object from give parameters - """ - import transformers - - assert hf_model_or_dir is not None - use_preloading = isinstance(hf_model_or_dir, - transformers.PreTrainedModel) - if use_preloading: - hf_model = hf_model_or_dir - hf_config_or_dir = hf_model.config - else: - hf_model_dir = hf_model_or_dir - hf_config_or_dir = hf_model_or_dir - - load_model_on_cpu = kwargs.pop('load_model_on_cpu', False) - tllm_config = BERTConfig.from_hugging_face( - hf_config_or_dir=hf_config_or_dir, - dtype=dtype, - mapping=mapping, - quant_config=quant_config, - **kwargs) - #NOTE: override architecture info - RobertaCls_mapping = { - "BertModel": "RobertaModel", - "BertForQuestionAnswering": "RobertaForQuestionAnswering", - "BertForSequenceClassification": "RobertaForSequenceClassification", - } - if tllm_config.is_roberta: - setattr(tllm_config, 'architecture', - RobertaCls_mapping[cls.__name__]) - else: - setattr(tllm_config, 'architecture', cls.__name__) - - torch_dtype = torch.float16 if dtype == 'float16' else torch.float32 - if not use_preloading: - hf_model = cls.load_hf_bert(model_dir=hf_model_dir, - load_model_on_cpu=load_model_on_cpu, - dtype=torch_dtype) - weights = load_weights_from_hf_model(hf_model=hf_model, - config=tllm_config) - model = cls(tllm_config) - model.load(weights) - - return model - - # Override the PretrainedModel's meothd, can unify in the future. - def prepare_inputs(self, max_batch_size, max_input_len, **kwargs): - remove_input_padding = default_net().plugin_config.remove_input_padding - # opt_shape is set to half of max batch_size and seq_len by default - # tune this according to real data distribution - bs_range = [1, (max_batch_size + 1) // 2, max_batch_size] - inlen_range = [1, (max_input_len + 1) // 2, max_input_len] - num_tokens_range = [ - 1, - (max_input_len * max_batch_size + 1) // 2, - max_input_len * max_batch_size, - ] - if not remove_input_padding: - input_ids = Tensor( - name='input_ids', - dtype=trt.int32, - shape=[-1, -1], - dim_range=OrderedDict([('batch_size', [bs_range]), - ('input_len', [inlen_range])]), - ) - # also called segment_ids - token_type_ids = Tensor( - name='token_type_ids', - dtype=trt.int32, - shape=[-1, -1], - dim_range=OrderedDict([('batch_size', [bs_range]), - ('input_len', [inlen_range])]), - ) - else: - input_ids = Tensor( - name="input_ids", - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([("num_tokens", [num_tokens_range])]), - ) - token_type_ids = Tensor( - name='token_type_ids', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([('num_tokens', [num_tokens_range])]), - ) - position_ids = Tensor( - name='position_ids', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([('num_tokens', [num_tokens_range])]), - ) - max_input_length = Tensor( - name="max_input_length", - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([("max_input_length", [inlen_range])]), - ) - input_lengths = Tensor(name='input_lengths', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([('batch_size', [bs_range]) - ])) - - inputs = { - 'input_ids': input_ids, - 'input_lengths': input_lengths, - 'token_type_ids': token_type_ids, - } - - if remove_input_padding: - inputs['position_ids'] = position_ids - inputs['max_input_length'] = max_input_length - - return inputs - - -class BertModel(BertBase): - - def __init__(self, config: BERTConfig): - super().__init__(config) - - self.config = config - self.max_position_embeddings = config.max_position_embeddings - self.padding_idx = config.pad_token_id - self.is_roberta = config.is_roberta - self.embedding = BertEmbedding( - vocab_size=config.vocab_size, - hidden_size=config.hidden_size, - max_position_embeddings=config.max_position_embeddings, - type_vocab_size=config.type_vocab_size, - dtype=config.dtype) - - self.layers = ModuleList([ - BertEncoderLayer( - hidden_size=config.hidden_size, - num_attention_heads=config.num_attention_heads, - max_position_embeddings=config.max_position_embeddings, - hidden_act=config.hidden_act, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size, - dtype=config.dtype) for _ in range(config.num_hidden_layers) - ]) - - def forward(self, - input_ids=None, - input_lengths=None, - position_ids=None, - token_type_ids=None, - hidden_states=None, - max_input_length=None): - # remove_input_padding requires these fields as explicit input - mask = None - if not default_net().plugin_config.remove_input_padding: - seq_len_2d = concat([1, shape(input_ids, 1)]) - - # create position ids - position_ids_buffer = constant( - np.expand_dims( - np.arange(self.max_position_embeddings).astype(np.int32), - 0)) - tmp_position_ids = slice(position_ids_buffer, - starts=[0, 0], - sizes=seq_len_2d) - tmp_position_ids = expand(tmp_position_ids, shape(input_ids)) #BxL - tmp_input_lengths = unsqueeze(input_lengths, 1) #Bx1 - tmp_input_lengths = expand(tmp_input_lengths, - shape(input_ids)) #BxL - mask = tmp_position_ids < tmp_input_lengths # BxL - mask = mask.cast('int32') - - if position_ids is None: - if self.is_roberta: - # see create_position_ids_from_input_ids() in https://github.com/huggingface/transformers/blob/main/src/transformers/models/roberta/modeling_roberta.py - position_ids = (tmp_position_ids + 1) * mask - position_ids = position_ids + self.padding_idx - else: - position_ids = slice(position_ids_buffer, - starts=[0, 0], - sizes=seq_len_2d) - position_ids = expand(position_ids, shape(input_ids)) - - # create token_type_ids - if token_type_ids is None: - token_type_ids_buffer = constant( - np.expand_dims( - np.zeros(self.max_position_embeddings).astype(np.int32), - 0)) - token_type_ids = slice(token_type_ids_buffer, - starts=[0, 0], - sizes=seq_len_2d) - token_type_ids = expand(token_type_ids, shape(input_ids)) - - hidden_states = self.embedding(input_ids, position_ids, token_type_ids) - self.register_network_output('embedding_output', hidden_states) - - for idx, layer in enumerate(self.layers): - hidden_states = layer(hidden_states=hidden_states, - input_lengths=input_lengths, - attention_mask=mask, - max_input_length=max_input_length) - # keep the last layer output name as hidden_states - if ((idx == (self.config.num_hidden_layers - 1)) and - (self.config.architecture in ["BertModel", "RobertaModel"])): - hidden_states.mark_output('hidden_states', self.config.dtype) - else: - self.register_network_output(f"layer_{idx}_output", - hidden_states) - - return hidden_states - - -RobertaModel = BertModel - - -class BertForQuestionAnswering(BertBase): - - def __init__(self, config: BERTConfig): - super().__init__(config) - self.bert = BertModel(config) - self.num_labels = config.num_labels - self.qa_outputs = Linear(config.hidden_size, - config.num_labels, - dtype=config.dtype) - - def forward(self, - input_ids=None, - input_lengths=None, - token_type_ids=None, - position_ids=None, - hidden_states=None, - max_input_length=None): - - remove_input_padding = default_net().plugin_config.remove_input_padding - if remove_input_padding: - assert token_type_ids is not None and \ - position_ids is not None and \ - max_input_length is not None, \ - "token_type_ids, position_ids, max_input_length is required " \ - "in remove_input_padding mode" - hidden_states = self.bert.forward(input_ids=input_ids, - input_lengths=input_lengths, - token_type_ids=token_type_ids, - position_ids=position_ids, - hidden_states=hidden_states, - max_input_length=max_input_length) - - logits = self.qa_outputs(hidden_states) - logits.mark_output('logits', self.config.logits_dtype) - - return logits - - -RobertaForQuestionAnswering = BertForQuestionAnswering - - -class BertPooler(Module): - - def __init__(self, hidden_size, dtype): - super().__init__() - self.dense = Linear(hidden_size, hidden_size, dtype=dtype) - self.activation = ACT2FN['tanh'] - - def forward(self, hidden_states, input_lengths, remove_input_padding): - if not remove_input_padding: - # We "pool" the model by simply taking the hidden state corresponding - # to the first token. - first_token_tensor = select(hidden_states, 1, 0) - else: - # when remove_input_padding is enabled, the shape of hidden_states is [num_tokens, hidden_size] - # We can take the first token of each sequence according to input_lengths, - # and then do pooling similar to padding mode. - # For example, if input_lengths is [8, 5, 6], then the indices of first tokens - # should be [0, 8, 13] - first_token_indices = cumsum( - concat([ - 0, - slice(input_lengths, - starts=[0], - sizes=(shape(input_lengths) - - constant(np.array([1], dtype=np.int32)))) - ]), 0) - first_token_tensor = index_select(hidden_states, 0, - first_token_indices) - - pooled_output = self.dense(first_token_tensor) - pooled_output = self.activation(pooled_output) - return pooled_output - - -class RobertaClassificationHead(Module): - """Head for sentence-level classification tasks.""" - - def __init__(self, hidden_size, dtype, num_labels): - super().__init__() - self.dense = Linear(hidden_size, hidden_size, dtype=dtype) - self.out_proj = Linear(hidden_size, num_labels) - - def forward(self, hidden_states, input_lengths, remove_input_padding): - - if not remove_input_padding: - # We "pool" the model by simply taking the hidden state corresponding - # to the first token. - first_token_tensor = select(hidden_states, 1, 0) - else: - # when remove_input_padding is enabled, the shape of hidden_states is [num_tokens, hidden_size] - # We can take the first token of each sequence according to input_lengths, - # and then do pooling similar to padding mode. - # For example, if input_lengths is [8, 5, 6], then the indices of first tokens - # should be [0, 8, 13] - first_token_indices = cumsum( - concat([ - 0, - slice(input_lengths, - starts=[0], - sizes=(shape(input_lengths) - - constant(np.array([1], dtype=np.int32)))) - ]), 0) - first_token_tensor = index_select(hidden_states, 0, - first_token_indices) - - x = self.dense(first_token_tensor) - x = ACT2FN['tanh'](x) - x = self.out_proj(x) - return x - - -class BertForSequenceClassification(BertBase): - - def __init__(self, config: BERTConfig): - super().__init__(config) - - self.config = config - self.is_roberta = config.is_roberta - self.bert = BertModel(config) - self.num_labels = config.num_labels - - if not config.is_roberta: - self.pooler = BertPooler(hidden_size=config.hidden_size, - dtype=config.dtype) - self.classifier = Linear(config.hidden_size, - config.num_labels, - dtype=config.dtype) - else: - self.classifier = RobertaClassificationHead( - hidden_size=config.hidden_size, - num_labels=config.num_labels, - dtype=config.dtype) - - def forward(self, - input_ids, - input_lengths, - token_type_ids=None, - position_ids=None, - hidden_states=None, - max_input_length=None): - - remove_input_padding = default_net().plugin_config.remove_input_padding - - # required as explicit input in remove_input_padding mode - # see examples/models/core/bert/run_remove_input_padding.py for how to create them from input_ids and input_lengths - if remove_input_padding: - assert token_type_ids is not None and \ - position_ids is not None and \ - max_input_length is not None, \ - "token_type_ids, position_ids, max_input_length is required " \ - "in remove_input_padding mode" - - hidden_states = self.bert.forward(input_ids=input_ids, - input_lengths=input_lengths, - token_type_ids=token_type_ids, - position_ids=position_ids, - hidden_states=hidden_states, - max_input_length=max_input_length) - - if not self.is_roberta: - pooled_output = self.pooler( - hidden_states=hidden_states, - input_lengths=input_lengths, - remove_input_padding=remove_input_padding) - logits = self.classifier(pooled_output) - else: - logits = self.classifier(hidden_states=hidden_states, - input_lengths=input_lengths, - remove_input_padding=remove_input_padding) - - logits.mark_output('logits', self.config.logits_dtype) - return logits - - -RobertaForSequenceClassification = BertForSequenceClassification diff --git a/tensorrt_llm/models/bloom/__init__.py b/tensorrt_llm/models/bloom/__init__.py deleted file mode 100644 index 71bf6d298d42..000000000000 --- a/tensorrt_llm/models/bloom/__init__.py +++ /dev/null @@ -1,14 +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. diff --git a/tensorrt_llm/models/bloom/model.py b/tensorrt_llm/models/bloom/model.py deleted file mode 100644 index a189f7c3b740..000000000000 --- a/tensorrt_llm/models/bloom/model.py +++ /dev/null @@ -1,165 +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. - -from ..._utils import pad_vocab_size -from ...functional import Tensor -from ...layers import (MLP, Attention, AttentionMaskType, ColumnLinear, - Embedding, LayerNorm, PositionEmbeddingType) -from ...module import Module -from ..modeling_utils import (DecoderLayerList, DecoderModelForCausalLM, - PretrainedConfig) - - -class BloomDecoderLayer(Module): - - def __init__(self, config: PretrainedConfig, layer_idx: int): - super().__init__() - self.layer_idx = layer_idx - self.config = config - - hidden_size = config.hidden_size - dtype = config.dtype - tp_group = config.mapping.tp_group - tp_size = config.mapping.tp_size - tp_rank = config.mapping.tp_rank - self.input_layernorm = LayerNorm(normalized_shape=hidden_size, - dtype=dtype) - - layers_range = config.mapping.pp_layers(config.num_hidden_layers) - local_layer_idx = layer_idx - layers_range[0] - self.attention = Attention( - local_layer_idx=local_layer_idx, - hidden_size=hidden_size, - num_attention_heads=config.num_attention_heads, - num_kv_heads=config.num_key_value_heads, - num_layers=config.num_hidden_layers, - dtype=dtype, - attention_mask_type=AttentionMaskType.causal, - position_embedding_type=PositionEmbeddingType.alibi, - bias=True, - tp_group=tp_group, - tp_size=tp_size, - tp_rank=tp_rank, - reorder=True, - quant_mode=config.quant_mode) - - mlp_hidden_size = hidden_size * 4 if config.intermediate_size is None else config.intermediate_size - - self.mlp = MLP(hidden_size=hidden_size, - ffn_hidden_size=mlp_hidden_size, - hidden_act='gelu', - dtype=dtype, - bias=True, - tp_group=tp_group, - tp_size=tp_size, - quant_mode=config.quant_mode) - self.post_layernorm = LayerNorm(normalized_shape=hidden_size, - dtype=dtype) - - def forward(self, - hidden_states: Tensor, - attention_mask=None, - use_cache=False, - kv_cache_params=None, - attention_params=None): - - assert isinstance(hidden_states, Tensor) - - residual = hidden_states - - hidden_states = self.input_layernorm(hidden_states) - - attention_output = self.attention(hidden_states, - attention_mask=attention_mask, - use_cache=use_cache, - kv_cache_params=kv_cache_params, - attention_params=attention_params) - - if use_cache: - attention_output, presents = attention_output - - hidden_states = residual + attention_output - - residual = hidden_states - hidden_states = self.post_layernorm(hidden_states) - - hidden_states = self.mlp(hidden_states) - - hidden_states = residual + hidden_states - - if use_cache: - return (hidden_states, presents) - return hidden_states - - -class BloomModel(Module): - - def __init__(self, config: PretrainedConfig): - super().__init__() - self.vocab_embedding = Embedding(config.vocab_size, - config.hidden_size, - dtype=config.dtype) - self.ln_embed = LayerNorm(normalized_shape=config.hidden_size, - dtype=config.dtype) - self.layers = DecoderLayerList(BloomDecoderLayer, config) - self.ln_f = LayerNorm(normalized_shape=config.hidden_size, - dtype=config.dtype) - - def forward(self, - input_ids: Tensor, - position_ids=None, - use_cache=False, - attention_mask=None, - kv_cache_params=None, - prompt_embedding_table=None, - prompt_tasks=None, - prompt_vocab_size=None, - attention_params=None): - hidden_states = self.vocab_embedding(input_ids) - - hidden_states = self.ln_embed(hidden_states) - - hidden_states = self.layers(hidden_states, - use_cache=use_cache, - attention_mask=attention_mask, - kv_cache_params=kv_cache_params, - attention_params=attention_params) - - if use_cache: - hidden_states, presents = hidden_states - - hidden_states = self.ln_f(hidden_states) - - if use_cache: - return (hidden_states, tuple(presents)) - return hidden_states - - -class BloomForCausalLM(DecoderModelForCausalLM): - - def __init__(self, config: PretrainedConfig): - transformer = BloomModel(config) - vocab_size_padded = pad_vocab_size(config.vocab_size, - config.mapping.tp_size) - - lm_head = ColumnLinear(config.hidden_size, - vocab_size_padded, - bias=False, - dtype=config.dtype, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size, - gather_output=True) - - super().__init__(config, transformer, lm_head) diff --git a/tensorrt_llm/models/chatglm/__init__.py b/tensorrt_llm/models/chatglm/__init__.py deleted file mode 100644 index 71bf6d298d42..000000000000 --- a/tensorrt_llm/models/chatglm/__init__.py +++ /dev/null @@ -1,14 +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. diff --git a/tensorrt_llm/models/chatglm/config.py b/tensorrt_llm/models/chatglm/config.py deleted file mode 100644 index 9ae0be3152dc..000000000000 --- a/tensorrt_llm/models/chatglm/config.py +++ /dev/null @@ -1,182 +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. -from typing import Optional, Union - -from ...mapping import Mapping -from ..convert_utils import infer_dtype -from ..modeling_utils import PretrainedConfig, QuantConfig - -GLM_VERSIONS = ['glm4', 'chatglm3', 'chatglm2', 'chatglm', 'glm'] -GLM_ARCH1_VERSIONS = ['chatglm', 'glm'] -GLM_ARCH2_VERSIONS = ['glm4', 'chatglm3', 'chatglm2'] - - -class ChatGLMConfig(PretrainedConfig): - - def __init__(self, - *, - chatglm_version: str = 'chatglm3', - add_bias_linear: bool = False, - add_qkv_bias: bool = True, - apply_query_key_layer_scaling: bool = False, - apply_residual_connection_post_layernorm: bool = False, - rmsnorm: bool = True, - rotary_pct: float = 0.5, - rotary_base: float = 10000.0, - rotary_scaling: Optional[dict] = None, - **kwargs): - self.chatglm_version = chatglm_version - self.add_bias_linear = add_bias_linear - self.add_qkv_bias = add_qkv_bias - self.apply_query_key_layer_scaling = apply_query_key_layer_scaling - self.apply_residual_connection_post_layernorm = apply_residual_connection_post_layernorm - self.rmsnorm = rmsnorm - self.rotary_pct = rotary_pct - self.rotary_base = rotary_base - self.rotary_scaling = rotary_scaling - super().__init__(**kwargs) - - def to_dict(self): - output = super().to_dict() - # Serialize the fields added in ChatGLMConfig - output['chatglm_version'] = self.chatglm_version - output['add_bias_linear'] = self.add_bias_linear - output['add_qkv_bias'] = self.add_qkv_bias - output[ - 'apply_query_key_layer_scaling'] = self.apply_query_key_layer_scaling - output[ - 'apply_residual_connection_post_layernorm'] = self.apply_residual_connection_post_layernorm - output['rmsnorm'] = self.rmsnorm - output['rotary_pct'] = self.rotary_pct - output['rotary_base'] = self.rotary_base - output['rotary_scaling'] = self.rotary_scaling - return output - - @classmethod - def from_hugging_face( - cls, - hf_config_or_dir: Union[str, 'transformers.PretrainedConfig'], - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - import transformers - trust_remote_code = kwargs.pop('trust_remote_code', True) - - # load hugging face config - if isinstance(hf_config_or_dir, transformers.PretrainedConfig): - hf_config = hf_config_or_dir - else: - hf_config_dir = str(hf_config_or_dir) - hf_config = transformers.AutoConfig.from_pretrained( - hf_config_dir, trust_remote_code=trust_remote_code) - - logits_dtype = kwargs.pop('logits_dtype', 'float32') - use_parallel_embedding = kwargs.pop('use_parallel_embedding', False) - embedding_sharding_dim = kwargs.pop('embedding_sharding_dim', 0) - chatglm_version = kwargs.pop('chatglm_version', None) - - # get chatglm version - if chatglm_version is None: - print("Inferring chatglm version from path...") - for v in GLM_VERSIONS: - if v in hf_config._name_or_path: - chatglm_version = v - break - if 'glm_4' in hf_config._name_or_path.replace("-", "_"): - chatglm_version = 'glm4' - assert chatglm_version in GLM_VERSIONS - print(f"Chatglm version: {chatglm_version}") - - if chatglm_version == 'glm': - hf_config.num_kv_heads = hf_config.num_attention_heads - hf_config.ffn_hidden_size = hf_config.hidden_size * 4 - hf_config.hidden_act = 'gelu' - hf_config.layernorm_epsilon = 1e-5 - hf_config.max_position_embeddings = hf_config.max_sequence_length - hf_config.add_bias_linear = True - hf_config.add_qkv_bias = True - hf_config.apply_query_key_layer_scaling = False - hf_config.apply_residual_connection_post_layernorm = False - hf_config.rmsnorm = False - hf_config.rope_ratio = 1.0 - elif chatglm_version == 'chatglm': - hf_config.num_kv_heads = hf_config.num_attention_heads - hf_config.ffn_hidden_size = hf_config.inner_hidden_size - hf_config.hidden_act = 'gelu' - hf_config.max_position_embeddings = hf_config.max_sequence_length - hf_config.add_bias_linear = True - hf_config.add_qkv_bias = True - hf_config.apply_query_key_layer_scaling = False - hf_config.apply_residual_connection_post_layernorm = False - hf_config.rmsnorm = False - hf_config.rope_ratio = 1.0 - else: - hf_config.vocab_size = hf_config.padded_vocab_size - hf_config.num_kv_heads = hf_config.multi_query_group_num - hf_config.hidden_act = 'swiglu' - hf_config.max_position_embeddings = hf_config.seq_length - hf_config.rmsnorm = getattr(hf_config, 'rmsnorm', 1.0) - hf_config.rope_ratio = getattr(hf_config, 'rope_ratio', 1.0) - - if chatglm_version == 'glm': - position_embedding_type = 'learned_absolute' - elif chatglm_version == 'chatglm': - position_embedding_type = 'chatglm' - elif chatglm_version in GLM_ARCH2_VERSIONS: - position_embedding_type = 'rope_gptj' - - rotary_base = 10000.0 - rotary_embedding_scaling = None - if chatglm_version == 'chatglm2': - if hf_config.rope_ratio > 1: - rotary_embedding_scaling = { - 'type': 'linear', - 'factor': hf_config.rope_ratio - } - elif chatglm_version == 'chatglm3' or chatglm_version == 'glm4': - rotary_base *= hf_config.rope_ratio - - dtype = infer_dtype(dtype, getattr(hf_config, 'torch_dtype', None)) - - return cls( - architecture=hf_config.architectures[0], - dtype=dtype, - logits_dtype=logits_dtype, - num_hidden_layers=hf_config.num_layers, - num_attention_heads=hf_config.num_attention_heads, - num_key_value_heads=hf_config.num_kv_heads, - hidden_size=hf_config.hidden_size, - intermediate_size=hf_config.ffn_hidden_size, - norm_epsilon=hf_config.layernorm_epsilon, - vocab_size=hf_config.vocab_size, - position_embedding_type=position_embedding_type, - max_position_embeddings=hf_config.max_position_embeddings, - rotary_pct=0.5, - rotary_base=rotary_base, - rotary_scaling=rotary_embedding_scaling, - hidden_act=hf_config.hidden_act, - use_parallel_embedding=use_parallel_embedding, - embedding_sharding_dim=embedding_sharding_dim, - quantization=quant_config, - mapping=mapping, - chatglm_version=chatglm_version, - add_bias_linear=hf_config.add_bias_linear, - add_qkv_bias=hf_config.add_qkv_bias, - apply_query_key_layer_scaling=False, - apply_residual_connection_post_layernorm=hf_config. - apply_residual_connection_post_layernorm, - rmsnorm=hf_config.rmsnorm, - ) diff --git a/tensorrt_llm/models/chatglm/convert.py b/tensorrt_llm/models/chatglm/convert.py deleted file mode 100644 index 6113de88300d..000000000000 --- a/tensorrt_llm/models/chatglm/convert.py +++ /dev/null @@ -1,727 +0,0 @@ -import copy -import functools -import os -import time -from collections import defaultdict -from typing import Dict, Optional - -import numpy as np -import safetensors -import torch -from tqdm import tqdm -from transformers import AutoModel, AutoTokenizer - -from tensorrt_llm._utils import pad_vocab_size -from tensorrt_llm.models import ChatGLMConfig -from tensorrt_llm.models.convert_utils import (generate_int8, get_weight, - get_weight_and_bias, - load_calib_dataset, smooth_gemm) -from tensorrt_llm.quantization import QuantAlgo - -from .config import GLM_ARCH1_VERSIONS, GLM_ARCH2_VERSIONS - - -def split(weight: torch.Tensor, - tp_size: int, - rank: int = 0, - dim: int = 0) -> torch.Tensor: - if tp_size == 1: - return weight - elif weight.ndim == 1: - return torch.chunk(weight, tp_size)[rank].contiguous() - else: - return torch.chunk(weight, tp_size, dim=dim)[rank].contiguous() - - -def tile_kv_weight_bias(v: torch.Tensor, kv_num_head: int, tp_size: int): - head_size = v.shape[0] // kv_num_head - reps = tp_size // kv_num_head - if v.ndim == 1: - v = v.reshape(kv_num_head, head_size)[:, None, :] - v = v.expand(kv_num_head, reps, head_size).reshape(-1).clone() - else: - hidden_size = v.shape[1] - v = v.reshape(kv_num_head, head_size, hidden_size)[:, None, :, :] - v = v.expand(kv_num_head, reps, head_size, - hidden_size).reshape(-1, hidden_size).clone() - return v - - -def split_qkv(v: torch.Tensor, tp_size: int, rank: int, hidden_size: int, - num_heads: int, num_kv_heads: int): - head_size = hidden_size // num_heads - if tp_size == 1: - return v - - assert v.shape[0] == hidden_size + head_size * num_kv_heads * 2 - query = v[:hidden_size] - key = v[hidden_size:hidden_size + head_size * num_kv_heads] - value = v[hidden_size + head_size * num_kv_heads:hidden_size + - head_size * num_kv_heads * 2] - - if num_kv_heads < tp_size: - key = tile_kv_weight_bias(key, num_kv_heads, tp_size) - value = tile_kv_weight_bias(value, num_kv_heads, tp_size) - assert (key.shape[0] % (tp_size * head_size)) == 0 - assert (value.shape[0] % (tp_size * head_size)) == 0 - - q_tmp = torch.chunk(query, tp_size, dim=0)[rank] - k_tmp = torch.chunk(key, tp_size, dim=0)[rank] - v_tmp = torch.chunk(value, tp_size, dim=0)[rank] - return torch.concatenate([q_tmp, k_tmp, v_tmp], dim=0).contiguous() - - -def split_embedding( - param: torch.Tensor, - tp_size: int, - tp_rank: int, - use_parallel_embedding: bool = False, - sharding_dim: int = 0, -) -> torch.Tensor: - if param is None: - return None - if not use_parallel_embedding: - return param - - vocab_size, hidden_size = param.size() - if sharding_dim == 0: - if vocab_size % tp_size != 0: - vocab_size_padded = pad_vocab_size(vocab_size, tp_size) - pad_width = vocab_size_padded - vocab_size - param = torch.nn.functional.pad(param, (0, 0, 0, pad_width), - value=0) - else: - assert hidden_size % tp_size == 0 - return split(param, tp_size, tp_rank, dim=sharding_dim) - - -def swap_and_split_mlp(weight: torch.Tensor, - tp_size: int, - tp_rank: int, - dim: int = 0) -> torch.Tensor: - """Swap the positions of gate and fc weights, and split weights for tensor parallel. - """ - gate_weight, fc_weight = torch.chunk(weight, 2, dim=dim) - fc_w = split(fc_weight, tp_size, tp_rank, dim=dim) - gate_w = split(gate_weight, tp_size, tp_rank, dim=dim) - return torch.cat([fc_w, gate_w], dim=dim).contiguous() - - -def get_tllm_linear_weight( - weight: torch.Tensor, - prefix: str, - bias: Optional[torch.Tensor] = None, - use_weight_only: bool = False, - plugin_weight_only_quant_type: torch.dtype = torch.int8 -) -> Dict[str, torch.Tensor]: - results = {} - if use_weight_only: - v = weight.t().contiguous() - processed_torch_weights, torch_weight_scales = \ - torch.ops.trtllm.symmetric_quantize_last_axis_of_batched_matrix( - v, plugin_weight_only_quant_type) - results[f'{prefix}.weight'] = processed_torch_weights - results[f'{prefix}.per_channel_scale'] = torch_weight_scales - else: - results[f'{prefix}.weight'] = weight.contiguous() - - if bias is not None: - results[f'{prefix}.bias'] = bias - - return results - - -@torch.no_grad() -def capture_activation_range( - model, - tokenizer, - dataset, - num_samples=64, - seq_len=512, -): - - model.eval() - device = next(model.parameters()).device - scales = defaultdict(lambda: {"x": None, "y": None, "w": None}) - - def stat_tensor(name, tensor, key): - tensor = tensor.view(-1, tensor.shape[-1]).detach() - comming_max = tensor.abs().max(dim=0)[0].float() - if scales[name][key] is None: - scales[name][key] = comming_max - else: - scales[name][key] = torch.max(scales[name][key], comming_max) - - def stat_input_hook(m, x, y, name): - if isinstance(x, tuple): - x = x[0] - stat_tensor(name, x, "x") - stat_tensor(name, y, "y") - # TODO: we don't need to do it every forward because inference does not change weight - if scales[name]["w"] is None: - scales[name]["w"] = m.weight.abs().clip(1e-8, None).max(dim=1)[0] - - hooks = [] - for name, m in model.named_modules(): - if isinstance(m, torch.nn.Linear): - hooks.append( - m.register_forward_hook( - functools.partial(stat_input_hook, name=name))) - - for i in tqdm(range(num_samples), desc="Calibration"): - input_ids = tokenizer( - dataset[i], - return_tensors="pt", - max_length=seq_len, - truncation=True, - ) - model(input_ids.input_ids.to(device)) - - for h in hooks: - h.remove() - - return scales - - -@torch.no_grad() -def smooth_chatglm_model( - model, - act_range, - alpha, - model_smoother, -): - for name, module in model.named_modules(): - if not module._get_name() == "GLMBlock": - continue - - # QKV multiplication weight - layer_name = name + '.self_attention.query_key_value' - print(f'Smoothing module: {layer_name}') - weight = module.self_attention.query_key_value.weight - smoother = smooth_gemm( - weight, - act_range[layer_name]["x"], - module.input_layernorm.weight, - None, - alpha, - ) - act_range[layer_name]["x"] = act_range[layer_name]["x"] / smoother - act_range[layer_name]["w"] = weight.abs().max(dim=1)[0] - - # Dense multiplication weight - layer_name = name + ".self_attention.dense" - print(f'Smoothing module: {layer_name}') - weight = module.self_attention.dense.weight - smoother = smooth_gemm( - weight, - act_range[layer_name]["x"], - None, - None, - alpha, - ) - model_smoother[layer_name] = smoother.float() - act_range[layer_name]["x"] = act_range[layer_name]["x"] / smoother - act_range[layer_name]["w"] = weight.abs().max(dim=1)[0] - - # Multilayer perceptron h -> 4h weight - layer_name = name + ".mlp.dense_h_to_4h" - print(f'Smoothing module: {layer_name}') - weight = module.mlp.dense_h_to_4h.weight - smoother = smooth_gemm( - weight, - act_range[layer_name]["x"], - module.post_attention_layernorm.weight, - None, - alpha, - ) - act_range[layer_name]["x"] = act_range[layer_name]["x"] / smoother - act_range[layer_name]["w"] = weight.abs().max(dim=1)[0] - - # Multilayer perceptron 4h -> h weight - layer_name = name + ".mlp.dense_4h_to_h" - print(f'Smoothing module: {layer_name}') - weight = module.mlp.dense_4h_to_h.weight - smoother = smooth_gemm( - weight, - act_range[layer_name]["x"], - None, - None, - alpha, - ) - model_smoother[layer_name] = smoother.float() - act_range[layer_name]["x"] = act_range[layer_name]["x"] / smoother - act_range[layer_name]["w"] = weight.abs().max(dim=1)[0] - - -def get_tllm_linear_sq_weight(vals, - prefix, - shape, - is_qkv=False, - per_token=False, - per_channel=False, - last_prefix=None, - smoother_value=None, - smoother_shape=None): - results = {} - col_shape = shape if (is_qkv or per_channel) else [1, 1] - - if per_token: - if per_channel: - original_weights = np.array(vals["weight.int8.col"]) - else: - original_weights = np.array(vals["weight.int8"]) - - cur_weights = original_weights - if is_qkv: - hidden_dim = cur_weights.shape[0] - cur_weights = cur_weights.reshape(hidden_dim, -1) - results[prefix + - 'weight'] = torch.from_numpy(cur_weights).t().contiguous() - if smoother_value is None: - results[last_prefix] = torch.from_numpy( - np.array([1.0], dtype=np.float16)) - - if per_channel: - cur_per_channel_value = vals["scale_w_quant_orig.col"] - else: - cur_per_channel_value = vals["scale_w_quant_orig"] - - results[prefix + - 'per_channel_scale'] = cur_per_channel_value.reshape(col_shape) - else: - if per_channel: - original_weights = np.array(vals["weight.int8.col"]) - else: - original_weights = np.array(vals["weight.int8"]) - cur_weights = original_weights - - if is_qkv: - hidden_dim = cur_weights.shape[0] - cur_weights = cur_weights.reshape(hidden_dim, -1) - results[prefix + - 'weight'] = torch.from_numpy(cur_weights).t().contiguous() - - if per_channel: - cur_per_channel_value = vals["scale_y_accum_quant.col"] - else: - cur_per_channel_value = vals["scale_y_accum_quant"] - - results[prefix + 'per_channel_scale'] = torch.from_numpy( - np.array([cur_per_channel_value], - dtype=np.float32).reshape(col_shape)).contiguous() - - results[last_prefix] = torch.from_numpy( - np.array([vals['scale_x_orig_quant']], - dtype=np.float32)).contiguous() - - results[prefix + 'act_scale'] = torch.from_numpy( - np.array([[vals["scale_y_quant_orig"]]], - dtype=np.float32)).contiguous() - - if smoother_value is not None: - results[prefix + 'smoother'] = smoother_value.reshape( - smoother_shape).contiguous().to(torch.float32) - - return results - - -def load_weights_from_hf_model(hf_model: AutoModel, - config: ChatGLMConfig, - act_range: Optional[dict] = None, - smoother: Optional[dict] = None): - weights = {} - tik = time.time() - - model_params = dict(hf_model.named_parameters()) - dtype = getattr(torch, config.dtype) - num_attention_heads = config.num_attention_heads - hidden_size = config.hidden_size - num_kv_heads = config.num_key_value_heads - num_hidden_layers = config.num_hidden_layers - - chatglm_version = config.chatglm_version - mapping = config.mapping - use_parallel_embedding = config.use_parallel_embedding - sharding_dim = config.embedding_sharding_dim - - quant_algo = config.quantization.quant_algo - use_weight_only = quant_algo in [QuantAlgo.W8A16, QuantAlgo.W4A16] - use_smooth_quant = config.quantization._use_plugin_sq - per_channel = use_smooth_quant and 'PER_CHANNEL' in quant_algo - per_token = use_smooth_quant and 'PER_TOKEN' in quant_algo - int8_kv_cache = config.quantization.kv_cache_quant_algo == QuantAlgo.INT8 - - if use_weight_only: - if quant_algo == QuantAlgo.W8A16: - plugin_weight_only_quant_type = torch.int8 - elif quant_algo == QuantAlgo.W4A16: - plugin_weight_only_quant_type = torch.quint4x2 - else: - plugin_weight_only_quant_type = None - - layers_range = mapping.pp_layers(num_hidden_layers) - for l in layers_range: - if chatglm_version in GLM_ARCH1_VERSIONS: - prefix = f'transformer.layers.{l}' - elif chatglm_version in GLM_ARCH2_VERSIONS: - prefix = f'transformer.encoder.layers.{l}' - tllm_prex = f'transformer.layers.{l-layers_range[0]}' - - # Attention QKV - attention_attr_name = '' - if chatglm_version in GLM_ARCH1_VERSIONS: - attention_attr_name = 'attention' - elif chatglm_version in GLM_ARCH2_VERSIONS: - attention_attr_name = 'self_attention' - qkv_weight, qkv_bias = get_weight_and_bias( - model_params, f'{prefix}.{attention_attr_name}.query_key_value', - dtype) - - if use_smooth_quant: - qkv_act_range = act_range.get( - f'{prefix}.{attention_attr_name}.query_key_value') - qkv_vals_int8 = generate_int8(qkv_weight.t(), - qkv_act_range, - is_qkv=True, - multi_query_mode=True) - weights.update( - get_tllm_linear_sq_weight( - vals=qkv_vals_int8, - prefix=f'{tllm_prex}.attention.qkv.', - shape=[1, qkv_weight.size(0)], - is_qkv=True, - per_token=per_token, - per_channel=per_channel, - last_prefix=f'{tllm_prex}.input_layernorm.scale_to_int', - smoother_value=None, - smoother_shape=None)) - if qkv_bias is not None: - qkv_b = split_qkv(qkv_bias, - mapping.tp_size, - mapping.tp_rank, - hidden_size, - num_attention_heads, - num_kv_heads=num_kv_heads) - weights[f'{tllm_prex}.attention.qkv.bias'] = qkv_b - else: - qkv_w = split_qkv(qkv_weight, - mapping.tp_size, - mapping.tp_rank, - hidden_size, - num_attention_heads, - num_kv_heads=num_kv_heads) - if qkv_bias is None: - qkv_b = None - else: - qkv_b = split_qkv(qkv_bias, - mapping.tp_size, - mapping.tp_rank, - hidden_size, - num_attention_heads, - num_kv_heads=num_kv_heads) - - weights.update( - get_tllm_linear_weight(qkv_w, f'{tllm_prex}.attention.qkv', - qkv_b, use_weight_only, - plugin_weight_only_quant_type)) - - if int8_kv_cache: - qkv_act_range = act_range.get( - f'{prefix}.{attention_attr_name}.query_key_value') - qkv_vals_int8 = generate_int8(qkv_weight.t(), - qkv_act_range, - is_qkv=True, - multi_query_mode=True) - weights[ - f'{tllm_prex}.attention.kv_cache_scaling_factor'] = qkv_vals_int8[ - 'scale_y_quant_orig'].contiguous() - - # Attention dense - attn_dense_weight, attn_dense_bias = get_weight_and_bias( - model_params, f'{prefix}.{attention_attr_name}.dense', dtype) - - if use_smooth_quant: - dense_act_range = act_range.get( - f'{prefix}.{attention_attr_name}.dense') - dense_smoother = smoother.get( - f'{prefix}.{attention_attr_name}.dense') - dense_vals_int8 = generate_int8(attn_dense_weight.t(), - dense_act_range, - is_qkv=False, - multi_query_mode=True) - weights.update( - get_tllm_linear_sq_weight( - vals=dense_vals_int8, - prefix=f'{tllm_prex}.attention.dense.', - shape=[1, hidden_size], - is_qkv=False, - per_token=per_token, - per_channel=per_channel, - last_prefix= - f'{tllm_prex}.attention.quantization_scaling_factor', - smoother_value=dense_smoother, - smoother_shape=[1, hidden_size])) - if attn_dense_bias is not None: - weights[f'{tllm_prex}.attention.dense.bias'] = attn_dense_bias - else: - attn_dense_w = split(attn_dense_weight, - mapping.tp_size, - mapping.tp_rank, - dim=1) - weights.update( - get_tllm_linear_weight(attn_dense_w, - f'{tllm_prex}.attention.dense', - attn_dense_bias, use_weight_only, - plugin_weight_only_quant_type)) - - # MLP FC - mlp_fc_weight, mlp_fc_bias = get_weight_and_bias( - model_params, f'{prefix}.mlp.dense_h_to_4h', dtype) - - if use_smooth_quant: - fc_act_range = act_range.get(f'{prefix}.mlp.dense_h_to_4h') - fc_vals_int8 = generate_int8(mlp_fc_weight.t(), - fc_act_range, - is_qkv=False, - multi_query_mode=True) - cur_weights = get_tllm_linear_sq_weight( - vals=fc_vals_int8, - prefix=f'{tllm_prex}.mlp.fc.', - shape=[1, mlp_fc_weight.size(0)], - is_qkv=False, - per_token=per_token, - per_channel=per_channel, - last_prefix=f'{tllm_prex}.post_layernorm.scale_to_int', - smoother_value=None, - smoother_shape=None, - ) - cur_weights[f'{tllm_prex}.mlp.fc.weight'] = swap_and_split_mlp( - cur_weights[f'{tllm_prex}.mlp.fc.weight'], - mapping.tp_size, - mapping.tp_rank, - dim=0, - ) - if per_channel: - cur_weights[ - f'{tllm_prex}.mlp.fc.per_channel_scale'] = swap_and_split_mlp( - cur_weights[f'{tllm_prex}.mlp.fc.per_channel_scale'], - mapping.tp_size, - mapping.tp_rank, - dim=1, - ) - weights.update(cur_weights) - if chatglm_version in GLM_ARCH1_VERSIONS: - if mlp_fc_bias is not None: - mlp_fc_b = split(mlp_fc_bias, - mapping.tp_size, - mapping.tp_rank, - dim=0) - weights[f'{tllm_prex}.mlp.fc.bias'] = mlp_fc_b - elif chatglm_version in GLM_ARCH2_VERSIONS: - if mlp_fc_bias is not None: - mlp_fc_b = swap_and_split_mlp(mlp_fc_bias, mapping.tp_size, - mapping.tp_rank) - weights[f'{tllm_prex}.mlp.fc.bias'] = mlp_fc_b - else: - if chatglm_version in GLM_ARCH1_VERSIONS: - mlp_fc_w = split(mlp_fc_weight, - mapping.tp_size, - mapping.tp_rank, - dim=0) - if mlp_fc_bias is None: - mlp_fc_b = None - else: - mlp_fc_b = split(mlp_fc_bias, - mapping.tp_size, - mapping.tp_rank, - dim=0) - elif chatglm_version in GLM_ARCH2_VERSIONS: - mlp_fc_w = swap_and_split_mlp(mlp_fc_weight, mapping.tp_size, - mapping.tp_rank) - if mlp_fc_bias is None: - mlp_fc_b = None - else: - mlp_fc_b = swap_and_split_mlp(mlp_fc_bias, mapping.tp_size, - mapping.tp_rank) - weights.update( - get_tllm_linear_weight(mlp_fc_w, f'{tllm_prex}.mlp.fc', - mlp_fc_b, use_weight_only, - plugin_weight_only_quant_type)) - - # MLP Proj - mlp_proj_weight, mlp_proj_bias = get_weight_and_bias( - model_params, f'{prefix}.mlp.dense_4h_to_h', dtype) - - if use_smooth_quant: - proj_act_range = act_range.get(f'{prefix}.mlp.dense_4h_to_h') - proj_smoother = smoother.get(f'{prefix}.mlp.dense_4h_to_h') - proj_vals_int8 = generate_int8(mlp_proj_weight.t(), - proj_act_range, - is_qkv=False, - multi_query_mode=True) - weights.update( - get_tllm_linear_sq_weight( - vals=proj_vals_int8, - prefix=f'{tllm_prex}.mlp.proj.', - shape=[1, hidden_size], - is_qkv=False, - per_token=per_token, - per_channel=per_channel, - last_prefix=f'{tllm_prex}.mlp.quantization_scaling_factor', - smoother_value=proj_smoother, - smoother_shape=[1, config.intermediate_size])) - if mlp_proj_bias is not None: - weights[f'{tllm_prex}.mlp.proj.bias'] = mlp_proj_bias - else: - mlp_proj_w = split(mlp_proj_weight, - mapping.tp_size, - mapping.tp_rank, - dim=1) - weights.update( - get_tllm_linear_weight(mlp_proj_w, f'{tllm_prex}.mlp.proj', - mlp_proj_bias, use_weight_only, - plugin_weight_only_quant_type)) - - input_ln_weight, input_ln_bias = get_weight_and_bias( - model_params, f'{prefix}.input_layernorm', dtype) - weights[f'{tllm_prex}.input_layernorm.weight'] = input_ln_weight - if input_ln_bias is not None: - weights[f'{tllm_prex}.input_layernorm.bias'] = input_ln_bias - - post_ln_weight, post_ln_bias = get_weight_and_bias( - model_params, f'{prefix}.post_attention_layernorm', dtype) - weights[f'{tllm_prex}.post_layernorm.weight'] = post_ln_weight - if post_ln_bias is not None: - weights[f'{tllm_prex}.post_layernorm.bias'] = post_ln_bias - - if mapping.is_first_pp_rank(): - if chatglm_version == 'glm': - embed_w = get_weight(model_params, 'word_embeddings', dtype) - pos_embed_w = get_weight(model_params, - 'transformer.position_embeddings', dtype) - weights['transformer.position_embedding.weight'] = split_embedding( - pos_embed_w, - tp_size=mapping.tp_size, - tp_rank=mapping.tp_rank, - use_parallel_embedding=use_parallel_embedding, - sharding_dim=sharding_dim) - block_embed_w = get_weight(model_params, - 'transformer.block_position_embeddings', - dtype) - weights['transformer.block_embedding.weight'] = split_embedding( - block_embed_w, - tp_size=mapping.tp_size, - tp_rank=mapping.tp_rank, - use_parallel_embedding=use_parallel_embedding, - sharding_dim=sharding_dim) - elif chatglm_version == 'chatglm': - embed_w = get_weight(model_params, 'transformer.word_embeddings', - dtype) - elif chatglm_version in GLM_ARCH2_VERSIONS: - embed_w = get_weight(model_params, - 'transformer.embedding.word_embeddings', dtype) - - weights['transformer.vocab_embedding.weight'] = split_embedding( - embed_w, - tp_size=mapping.tp_size, - tp_rank=mapping.tp_rank, - use_parallel_embedding=use_parallel_embedding, - sharding_dim=sharding_dim) - - if mapping.is_last_pp_rank(): - if chatglm_version == 'glm': - lm_head_weight = get_weight(model_params, 'word_embeddings', - dtype).clone() - elif chatglm_version == 'chatglm': - lm_head_weight = get_weight(model_params, - 'transformer.word_embeddings', - dtype).clone() - elif chatglm_version in GLM_ARCH2_VERSIONS: - lm_head_weight = get_weight(model_params, - 'transformer.output_layer', dtype) - - weights['lm_head.weight'] = split(lm_head_weight, - mapping.tp_size, - mapping.tp_rank, - dim=0) - - if chatglm_version in GLM_ARCH1_VERSIONS: - ln_f_w, ln_f_b = get_weight_and_bias(model_params, - 'transformer.final_layernorm', - dtype) - elif chatglm_version in GLM_ARCH2_VERSIONS: - ln_f_w, ln_f_b = get_weight_and_bias( - model_params, 'transformer.encoder.final_layernorm', dtype) - weights['transformer.ln_f.weight'] = ln_f_w - if ln_f_b is not None: - weights['transformer.ln_f.bias'] = ln_f_b - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Weights loaded. Total time: {t}') - return weights - - -def quantize(hf_model_dir: str, - output_dir: str, - config: ChatGLMConfig, - calib_dataset: str = 'cnn_dailymail', - device: str = 'auto', - trust_remote_code: bool = True): - ''' - Quantize the save the model as TRT-LLM checkpoint to output_dir - ''' - os.makedirs(output_dir, exist_ok=True) - config.to_json_file(os.path.join(output_dir, 'config.json')) - - mapping = config.mapping - assert mapping.rank == 0, "quantize should be called at rank 0 only" - - quant_config = config.quantization - use_smooth_quant = quant_config._use_plugin_sq - int8_kv_cache = quant_config.kv_cache_quant_algo == QuantAlgo.INT8 - - assert use_smooth_quant or int8_kv_cache, "Call from_hugging_face when there is no quantization" - assert hf_model_dir is not None - ## only load and call smooth quant routine once for all ranks - if config.chatglm_version == 'glm': - device_map = 'cuda' if device != "cpu" else 'cpu' - else: - device_map = 'auto' if device != "cpu" else 'cpu' - hf_model = AutoModel.from_pretrained( - hf_model_dir, - trust_remote_code=trust_remote_code, - dtype='auto' if config.chatglm_version != 'glm' else getattr( - torch, config.dtype), - device_map=device_map) - - os.environ["TOKENIZERS_PARALLELISM"] = os.environ.get( - "TOKENIZERS_PARALLELISM", "false") - tokenizer = AutoTokenizer.from_pretrained( - hf_model_dir, - trust_remote_code=trust_remote_code, - ) - dataset = load_calib_dataset(calib_dataset) - - act_range = capture_activation_range(hf_model, - tokenizer, - dataset, - num_samples=64) - smoother = {} - if use_smooth_quant: - smooth_chatglm_model(hf_model, act_range, quant_config.smoothquant_val, - smoother) - - for rank in range(mapping.world_size): - # To avoid changing the mapping arg in-place, also the given mapping from caller is rank agnostic, since quantize is called from only one rank - config = copy.deepcopy(config) - config.set_rank(rank) - weights = load_weights_from_hf_model( - hf_model, - config=config, - act_range=act_range, - smoother=smoother, - ) - safetensors.torch.save_file( - weights, os.path.join(output_dir, f'rank{rank}.safetensors')) - del weights diff --git a/tensorrt_llm/models/chatglm/model.py b/tensorrt_llm/models/chatglm/model.py deleted file mode 100644 index 8382419e8eb1..000000000000 --- a/tensorrt_llm/models/chatglm/model.py +++ /dev/null @@ -1,372 +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. - -from typing import Optional, Union - -import torch -from transformers import AutoModel - -from ..._common import default_net -from ..._utils import pad_vocab_size -from ...functional import Tensor, concat, shape -from ...layers import (MLP, Attention, AttentionMaskType, AttentionParams, - ColumnLinear, Embedding, KeyValueCacheParams, LayerNorm, - RmsNorm) -from ...mapping import Mapping -from ...module import Module -from ..modeling_utils import (DecoderLayerList, DecoderModelForCausalLM, - QuantConfig) -from .config import GLM_ARCH1_VERSIONS, GLM_ARCH2_VERSIONS, ChatGLMConfig -from .convert import load_weights_from_hf_model - - -class ChatGLMDecoderLayer(Module): - - def __init__(self, config: ChatGLMConfig, layer_idx: int): - super().__init__() - self.layer_idx = layer_idx - self.config = config - self.chatglm_version = config.chatglm_version - - hidden_size = config.hidden_size - dtype = config.dtype - tp_group = config.mapping.tp_group - tp_size = config.mapping.tp_size - tp_rank = config.mapping.tp_rank - layernorm_epsilon = config.norm_epsilon - - self.apply_residual_connection_post_layernorm = config.apply_residual_connection_post_layernorm - self.alpha = (2 * config.num_hidden_layers)**0.5 - norm_cls = RmsNorm if config.rmsnorm else LayerNorm - - if config.chatglm_version == 'glm': - attention_mask_type = AttentionMaskType.bidirectionalglm - elif config.chatglm_version == 'chatglm': - attention_mask_type = AttentionMaskType.bidirectional - elif config.chatglm_version in GLM_ARCH2_VERSIONS: - attention_mask_type = AttentionMaskType.causal - - self.input_layernorm = norm_cls( - normalized_shape=hidden_size, - eps=layernorm_epsilon, - elementwise_affine=True, - dtype=dtype, - ) - - layers_range = config.mapping.pp_layers(config.num_hidden_layers) - local_layer_idx = layer_idx - layers_range[0] - self.attention = Attention( - local_layer_idx=local_layer_idx, - hidden_size=hidden_size, - num_attention_heads=config.num_attention_heads, - num_kv_heads=config.num_key_value_heads, - max_position_embeddings=config.max_position_embeddings, - num_layers=config.num_hidden_layers, - apply_query_key_layer_scaling=config.apply_query_key_layer_scaling, - attention_mask_type=attention_mask_type, - bias=config.add_qkv_bias, - dense_bias=config.add_bias_linear, - dtype=config.dtype, - position_embedding_type=config.position_embedding_type, - rotary_embedding_base=config.rotary_base, - rotary_embedding_scaling=config.rotary_scaling, - rotary_embedding_percentage=config.rotary_pct, - tp_group=tp_group, - tp_size=tp_size, - tp_rank=tp_rank, - quant_mode=config.quant_mode, - q_scaling=1.0, - cross_attention=False, - relative_attention=False, - max_distance=0, - num_buckets=0, - cp_rank=config.mapping.cp_rank, - cp_size=config.mapping.cp_size, - cp_group=config.mapping.cp_group, - ) - - mlp_hidden_size = hidden_size * 4 if config.intermediate_size is None else config.intermediate_size - - self.mlp = MLP( - hidden_size=hidden_size, - ffn_hidden_size=mlp_hidden_size, - hidden_act=config.hidden_act, - bias=config.add_bias_linear, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - quant_mode=config.quant_mode, - ) - - self.post_layernorm = norm_cls( - normalized_shape=hidden_size, - eps=layernorm_epsilon, - elementwise_affine=True, - dtype=dtype, - ) - - def forward( - self, - hidden_states: Tensor, - attention_mask: Tensor = None, - position_ids: Tensor = None, # only used in ChatGLM-6B - use_cache: bool = False, - kv_cache_params: KeyValueCacheParams = None, - attention_params: AttentionParams = None, - ): - norm_output = self.input_layernorm(hidden_states) - - attention_output = self.attention( - hidden_states=norm_output, - attention_mask=attention_mask, - use_cache=use_cache, - kv_cache_params=kv_cache_params, - attention_params=attention_params, - encoder_output=None, - position_embedding=position_ids, - ) - - if use_cache: - attention_output, presents = attention_output - - if self.chatglm_version == 'chatglm': - residual = norm_output - - norm_input = residual * self.alpha + attention_output - - norm_output = self.post_layernorm(norm_input) - - mlp_output = self.mlp(norm_output) - - residual = norm_output - - output = residual * self.alpha + mlp_output - - else: - residual = norm_output if self.apply_residual_connection_post_layernorm else hidden_states - - norm_input = residual + attention_output - - norm_output = self.post_layernorm(norm_input) - - mlp_output = self.mlp(norm_output) - - residual = norm_output if self.apply_residual_connection_post_layernorm else norm_input - - output = residual + mlp_output - - if use_cache: - return (output, presents) - return output - - -class ChatGLMModel(Module): - - def __init__(self, config: ChatGLMConfig): - super().__init__() - self.chatglm_version = config.chatglm_version - norm_cls = RmsNorm if config.rmsnorm else LayerNorm - - self.vocab_embedding = Embedding(config.vocab_size, - config.hidden_size, - dtype=config.dtype) - - if config.chatglm_version == 'glm': - self.position_embedding = Embedding( - config.max_position_embeddings + 1, - config.hidden_size, - dtype=config.dtype, - ) - self.block_embedding = Embedding( - config.max_position_embeddings + 1, - config.hidden_size, - dtype=config.dtype, - ) - - self.layers = DecoderLayerList(ChatGLMDecoderLayer, config) - - self.ln_f = norm_cls( - normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - elementwise_affine=True, - dtype=config.dtype, - ) - - def forward( - self, - input_ids: Tensor = None, - position_ids: Tensor = None, # only used in ChatGLM-6B - use_cache: bool = False, - attention_mask: Tensor = None, - kv_cache_params: KeyValueCacheParams = None, - attention_params: AttentionParams = None, - ): - hidden_states = self.vocab_embedding(input_ids) - - if self.chatglm_version == 'glm': - if default_net().plugin_config.remove_input_padding: - position_ids_list = position_ids.split(1, dim=0) - else: - position_ids_list = position_ids.split(1, dim=1) - - position_embedding = self.position_embedding(position_ids_list[0]) - block_embedding = self.block_embedding(position_ids_list[1]) - position_embedding = position_embedding + block_embedding - - if default_net().plugin_config.remove_input_padding: - position_embedding = position_embedding.view( - concat([ - shape(position_embedding, 1), - shape(position_embedding, 2) - ])) - else: - position_embedding = position_embedding.view( - concat([ - shape(position_embedding, 0), - shape(position_embedding, 2), - shape(position_embedding, 3), - ])) - - hidden_states = hidden_states + position_embedding - - hidden_states = self.layers(hidden_states, - use_cache=use_cache, - attention_mask=attention_mask, - kv_cache_params=kv_cache_params, - attention_params=attention_params, - position_ids=position_ids) - - if use_cache: - hidden_states, presents = hidden_states - - hidden_states = self.ln_f(hidden_states) - - if use_cache: - return (hidden_states, tuple(presents)) - return hidden_states - - -class ChatGLMForCausalLM(DecoderModelForCausalLM): - config_class = ChatGLMConfig - - def __init__(self, config: ChatGLMConfig): - transformer = ChatGLMModel(config) - vocab_size_padded = pad_vocab_size(config.vocab_size, - config.mapping.tp_size) - - lm_head = ColumnLinear(config.hidden_size, - vocab_size_padded, - bias=False, - dtype=config.dtype, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size, - gather_output=True) - super().__init__(config, transformer, lm_head) - - @classmethod - def from_hugging_face( - cls, - hf_model_or_dir: Union[str, 'transformers.PreTrainedModel'], - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - ''' Create a LLaMAForCausalLM object from give parameters - ''' - load_model_on_cpu = kwargs.pop('load_model_on_cpu', False) - trust_remote_code = kwargs.pop('trust_remote_code', True) - - config = ChatGLMConfig.from_hugging_face(hf_model_or_dir, - dtype=dtype, - mapping=mapping, - quant_config=quant_config, - **kwargs) - if config.chatglm_version == 'glm': - device_map = 'cuda' if not load_model_on_cpu else 'cpu' - else: - device_map = 'auto' if not load_model_on_cpu else 'cpu' - hf_model = AutoModel.from_pretrained( - hf_model_or_dir, - trust_remote_code=trust_remote_code, - dtype='auto' if config.chatglm_version != 'glm' else getattr( - torch, config.dtype), - device_map=device_map) - weights = load_weights_from_hf_model(hf_model, config) - - model = cls(config) - model.load(weights) - return model - - @classmethod - def quantize( - cls, - hf_model_dir: str, - output_dir: str, - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - *, - device: str = 'cuda', - calib_dataset: str = 'cnn_dailymail', - calib_batches: int = 512, - calib_batch_size: int = 1, - calib_max_seq_length: int = 512, - random_seed: int = 1234, - tokenizer_max_seq_length: int = 2048, - **kwargs, - ): - if quant_config._requires_modelopt_quantization: - # modelopt quantization flow - super().quantize(hf_model_dir, - output_dir, - dtype=dtype, - mapping=mapping, - quant_config=quant_config, - device=device, - calib_dataset=calib_dataset, - calib_batches=calib_batches, - calib_batch_size=calib_batch_size, - calib_max_seq_length=calib_max_seq_length, - random_seed=random_seed, - tokenizer_max_seq_length=tokenizer_max_seq_length) - elif quant_config._requires_calibration: - # non-modelopt quantization flow - from . import convert - - config = ChatGLMConfig.from_hugging_face(hf_model_dir, - dtype=dtype, - mapping=mapping, - quant_config=quant_config, - **kwargs) - convert.quantize(hf_model_dir, - output_dir, - config=config, - calib_dataset=calib_dataset, - device=device) - else: - raise ValueError( - f"The quant_config ({quant_config}) does not require calibration, try {cls.__name__}.from_hugging_face instead." - ) - - def prepare_inputs(self, *args, **kwargs): - """See `PretrainedModel.prepare_inputs` for the detailed parameter list. - """ - if self.transformer.chatglm_version in GLM_ARCH1_VERSIONS: - position_encoding_2d = True - else: - position_encoding_2d = False - return super().prepare_inputs(*args, - **kwargs, - position_encoding_2d=position_encoding_2d) diff --git a/tensorrt_llm/models/clip/__init__.py b/tensorrt_llm/models/clip/__init__.py deleted file mode 100644 index 71bf6d298d42..000000000000 --- a/tensorrt_llm/models/clip/__init__.py +++ /dev/null @@ -1,14 +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. diff --git a/tensorrt_llm/models/clip/model.py b/tensorrt_llm/models/clip/model.py deleted file mode 100644 index df131434b122..000000000000 --- a/tensorrt_llm/models/clip/model.py +++ /dev/null @@ -1,213 +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. - -from ...functional import arange, concat, expand, expand_dims, shape -from ...layers import MLP, BertAttention, Conv2d, Embedding, LayerNorm -from ...mapping import Mapping -from ...module import Module, ModuleList -from ...parameter import Parameter - - -# Adapted from https://github.com/huggingface/transformers/blob/v4.39.0/src/transformers/models/clip/modeling_clip.py#L164 -class CLIPVisionEmbeddings(Module): - - def __init__(self, image_size, num_channels, patch_size, hidden_size, - dtype): - super().__init__() - self.image_size = image_size - self.num_channels = num_channels - self.patch_size = patch_size - self.embed_dim = hidden_size - self.dtype = dtype - - self.class_embedding = Parameter(shape=[ - self.embed_dim, - ], - dtype=self.dtype) - - self.patch_embedding = Conv2d(in_channels=self.num_channels, - out_channels=self.embed_dim, - kernel_size=(self.patch_size, - self.patch_size), - stride=(self.patch_size, self.patch_size), - bias=False, - dtype=self.dtype) - - self.num_patches = (self.image_size // self.patch_size)**2 - self.num_positions = self.num_patches + 1 - self.position_embedding = Embedding(self.num_positions, - self.embed_dim, - dtype=self.dtype) - - def forward(self, pixel_values): - batch_size = shape(pixel_values, 0) - target_dtype = self.patch_embedding.weight.dtype - patch_embeds = self.patch_embedding( - pixel_values.cast( - dtype=target_dtype)) # shape = [*, width, grid, grid] - patch_embeds = patch_embeds.flatten(2).transpose(1, 2) - class_embeds = expand_dims(expand_dims(self.class_embedding.value, 0), - 0) - expand_shape = concat( - [batch_size, - shape(class_embeds, -2), - shape(class_embeds, -1)]) - class_embeds = expand(class_embeds, - expand_shape) # shape = [*, 1, grid, grid] - embeddings = concat([class_embeds, patch_embeds], - dim=1) # shape = [*, width + 1, grid, grid] - position_ids = arange(0, self.num_positions, dtype='int32') - position_embeds = self.position_embedding(position_ids) - position_embeds = expand_dims(position_embeds, 0) - expand_shape = concat([ - batch_size, - shape(position_embeds, -2), - shape(position_embeds, -1) - ]) - position_embeds = expand( - position_embeds, expand_shape) # shape = [*, width + 1, grid, grid] - embeddings = embeddings + position_embeds - return embeddings - - -class CLIPEncoderLayer(Module): - - def __init__(self, hidden_size, num_attention_heads, - max_position_embeddings, norm_epsilon, intermediate_size, - hidden_act, mapping: Mapping, dtype): - super().__init__() - self.hidden_size = hidden_size - self.dtype = dtype - self.mapping = mapping - - self.input_layernorm = LayerNorm(normalized_shape=self.hidden_size, - eps=norm_epsilon, - dtype=self.dtype) - - self.attention = BertAttention( - hidden_size=self.hidden_size, - num_attention_heads=num_attention_heads, - max_position_embeddings=max_position_embeddings, - attention_head_size=self.hidden_size // num_attention_heads, - num_kv_heads=num_attention_heads, - dtype=self.dtype, - tp_group=self.mapping.tp_group, - tp_size=self.mapping.tp_size, - tp_rank=self.mapping.tp_rank, - cp_group=self.mapping.cp_group, - cp_size=self.mapping.cp_size) - - self.post_layernorm = LayerNorm(normalized_shape=self.hidden_size, - eps=norm_epsilon, - dtype=self.dtype) - - self.mlp = MLP(hidden_size=self.hidden_size, - ffn_hidden_size=intermediate_size, - hidden_act=hidden_act, - dtype=self.dtype, - tp_group=self.mapping.tp_group, - tp_size=self.mapping.tp_size) - - def forward(self, hidden_states): - - residual = hidden_states - hidden_states = self.input_layernorm(hidden_states) - hidden_states = self.attention(hidden_states) - hidden_states = residual + hidden_states - - residual = hidden_states - hidden_states = self.post_layernorm(hidden_states) - hidden_states = self.mlp(hidden_states) - hidden_states = residual + hidden_states - return hidden_states - - -class CLIPEncoder(Module): - - def __init__(self, hidden_size, num_attention_heads, - max_position_embeddings, norm_epsilon, intermediate_size, - hidden_act, num_hidden_layers, mapping: Mapping, dtype): - super().__init__() - self.hidden_size = hidden_size - self.dtype = dtype - self.mapping = mapping - - self.layers = ModuleList([ - CLIPEncoderLayer(hidden_size=self.hidden_size, - num_attention_heads=num_attention_heads, - max_position_embeddings=max_position_embeddings, - norm_epsilon=norm_epsilon, - intermediate_size=intermediate_size, - hidden_act=hidden_act, - mapping=self.mapping, - dtype=self.dtype) for _ in range(num_hidden_layers) - ]) - - def forward(self, inputs_embeds): - - hidden_states = inputs_embeds - for layer in self.layers: - hidden_states = layer(hidden_states) - - return hidden_states - - -class CLIPVisionTransformer(Module): - - def __init__(self, image_size, num_channels, patch_size, hidden_size, - num_attention_heads, max_position_embeddings, norm_epsilon, - intermediate_size, hidden_act, num_hidden_layers, require_ln_f, - mapping: Mapping, dtype) -> None: - super().__init__() - self.hidden_size = hidden_size - self.dtype = dtype - self.mapping = mapping - - self.embeddings = CLIPVisionEmbeddings(image_size=image_size, - num_channels=num_channels, - patch_size=patch_size, - hidden_size=hidden_size, - dtype=self.dtype) - self.pre_layernorm = LayerNorm(normalized_shape=self.hidden_size, - eps=norm_epsilon, - dtype=self.dtype) - - self.encoder = CLIPEncoder( - hidden_size=self.hidden_size, - num_attention_heads=num_attention_heads, - max_position_embeddings=max_position_embeddings, - norm_epsilon=norm_epsilon, - intermediate_size=intermediate_size, - hidden_act=hidden_act, - num_hidden_layers=num_hidden_layers, - mapping=self.mapping, - dtype=self.dtype) - - self.ln_f = None - - if require_ln_f: - self.ln_f = LayerNorm(normalized_shape=self.hidden_size, - eps=norm_epsilon, - dtype=self.dtype) - - def forward(self, pixel_values): - hidden_states = self.embeddings(pixel_values) - hidden_states = self.pre_layernorm(hidden_states) - hidden_states = self.encoder(inputs_embeds=hidden_states) - - if self.ln_f is None: - return hidden_states - - return self.ln_f(hidden_states) diff --git a/tensorrt_llm/models/cogvlm/__init__.py b/tensorrt_llm/models/cogvlm/__init__.py deleted file mode 100644 index 71bf6d298d42..000000000000 --- a/tensorrt_llm/models/cogvlm/__init__.py +++ /dev/null @@ -1,14 +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. diff --git a/tensorrt_llm/models/cogvlm/config.py b/tensorrt_llm/models/cogvlm/config.py deleted file mode 100644 index 1588a95f6a67..000000000000 --- a/tensorrt_llm/models/cogvlm/config.py +++ /dev/null @@ -1,44 +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. - -from typing import Optional - -from ..modeling_utils import PretrainedConfig - - -class CogVLMConfig(PretrainedConfig): - - def __init__(self, - *, - mlp_bias: bool = False, - attn_bias: bool = False, - rotary_base: float = 10000.0, - rotary_scaling: Optional[dict] = None, - **kwargs): - self.mlp_bias = mlp_bias - self.attn_bias = attn_bias - self.rotary_base = rotary_base - self.rotary_scaling = rotary_scaling - self.position_embedding_type = 'rope_gpt_neox' - super().__init__(**kwargs) - - def to_dict(self): - output = super().to_dict() - # Serialize the fields added in CogVLMConfig - output['mlp_bias'] = self.mlp_bias - output['attn_bias'] = self.attn_bias - output['rotary_base'] = self.rotary_base - output['rotary_scaling'] = self.rotary_scaling - return output diff --git a/tensorrt_llm/models/cogvlm/convert.py b/tensorrt_llm/models/cogvlm/convert.py deleted file mode 100644 index 68413e4e224a..000000000000 --- a/tensorrt_llm/models/cogvlm/convert.py +++ /dev/null @@ -1,240 +0,0 @@ -import time - -import numpy as np -import torch - -from tensorrt_llm.logger import logger - -from ..._utils import pad_vocab_size -from ..llama.convert import (get_tllm_linear_weight, get_weight, split, - split_matrix_tp, split_qkv_tp) - - -def convert_hf_cogvlm(hf_model, - mapping, - vocab_size=32000, - dtype='float32', - use_parallel_embedding=False, - sharding_dim=0, - use_weight_only=False, - use_gemm_woq_plugin=False, - plugin_weight_only_quant_type=torch.int8, - use_smooth_quant=False, - per_channel=False, - per_token=False, - int8_kv_cache=False, - act_range=[], - qkv_para=[], - smoother=[]): - - weights = {} - tik = time.time() - tensor_parallel = mapping.tp_size - model_params = dict(hf_model.named_parameters()) - dtype = getattr(torch, dtype) - num_attention_heads = hf_model.config.num_attention_heads - hidden_size = hf_model.config.hidden_size - if hasattr(hf_model.config, "num_key_value_heads"): - num_key_value_heads = hf_model.config.num_key_value_heads - else: - num_key_value_heads = num_attention_heads - mha_mode = (num_key_value_heads == num_attention_heads) - layers_range = mapping.pp_layers(hf_model.config.num_hidden_layers) - assert mha_mode, "CogVLM only supports mha mode" - assert not use_smooth_quant, "CogVLM currently doesn't support smooth quant" - assert not int8_kv_cache, "CogVLM currently doesn't support int8 kv cache" - - for l in layers_range: - prefix = f'model.layers.{l}.' - tllm_prex = f'transformer.layers.{l - layers_range[0]}.' - - qkv_weight = get_weight( - model_params, prefix + 'self_attn.language_expert_query_key_value', - dtype) - split_v = split_qkv_tp(qkv_weight, num_attention_heads, hidden_size, - tensor_parallel, mapping.tp_rank) - weights.update( - get_tllm_linear_weight(split_v, tllm_prex + 'attention.qkv.', None, - use_weight_only, - plugin_weight_only_quant_type, dtype, - use_gemm_woq_plugin)) - - vis_qkv_weight = get_weight( - model_params, prefix + 'self_attn.vision_expert_query_key_value', - dtype) - split_v = split_qkv_tp(vis_qkv_weight, num_attention_heads, hidden_size, - tensor_parallel, mapping.tp_rank) - weights.update( - get_tllm_linear_weight(split_v, tllm_prex + 'attention.vis_qkv.', - None, use_weight_only, - plugin_weight_only_quant_type, dtype, - use_gemm_woq_plugin)) - - attn_dense_weight = get_weight( - model_params, prefix + 'self_attn.language_expert_dense', dtype) - split_v = split_matrix_tp(attn_dense_weight, - tensor_parallel, - mapping.tp_rank, - dim=1) - weights.update( - get_tllm_linear_weight(split_v, tllm_prex + 'attention.dense.', - None, use_weight_only, - plugin_weight_only_quant_type, dtype, - use_gemm_woq_plugin)) - - attn_vision_dense_weight = get_weight( - model_params, prefix + 'self_attn.vision_expert_dense', dtype) - split_v = split_matrix_tp(attn_vision_dense_weight, - tensor_parallel, - mapping.tp_rank, - dim=1) - weights.update( - get_tllm_linear_weight(split_v, tllm_prex + 'attention.vis_dense.', - None, use_weight_only, - plugin_weight_only_quant_type, dtype, - use_gemm_woq_plugin)) - - mlp_gate_weight = get_weight(model_params, - prefix + 'mlp.language_mlp.up_proj', dtype) - split_v = split_matrix_tp(mlp_gate_weight, - tensor_parallel, - mapping.tp_rank, - dim=0) - weights.update( - get_tllm_linear_weight(split_v, tllm_prex + 'mlp.gate.', None, - use_weight_only, - plugin_weight_only_quant_type, dtype, - use_gemm_woq_plugin)) - - vision_mlp_gate_weight = get_weight(model_params, - prefix + 'mlp.vision_mlp.up_proj', - dtype) - split_v = split_matrix_tp(vision_mlp_gate_weight, - tensor_parallel, - mapping.tp_rank, - dim=0) - weights.update( - get_tllm_linear_weight(split_v, tllm_prex + 'vis_mlp.gate.', None, - use_weight_only, - plugin_weight_only_quant_type, dtype, - use_gemm_woq_plugin)) - - mlp_fc_weight = get_weight(model_params, - prefix + 'mlp.language_mlp.gate_proj', dtype) - split_v = split_matrix_tp(mlp_fc_weight, - tensor_parallel, - mapping.tp_rank, - dim=0) - weights.update( - get_tllm_linear_weight(split_v, tllm_prex + 'mlp.fc.', None, - use_weight_only, - plugin_weight_only_quant_type, dtype, - use_gemm_woq_plugin)) - - vision_mlp_fc_weight = get_weight(model_params, - prefix + 'mlp.vision_mlp.gate_proj', - dtype) - split_v = split_matrix_tp(vision_mlp_fc_weight, - tensor_parallel, - mapping.tp_rank, - dim=0) - weights.update( - get_tllm_linear_weight(split_v, tllm_prex + 'vis_mlp.fc.', None, - use_weight_only, - plugin_weight_only_quant_type, dtype, - use_gemm_woq_plugin)) - - mlp_proj_weight = get_weight(model_params, - prefix + 'mlp.language_mlp.down_proj', - dtype) - split_v = split_matrix_tp(mlp_proj_weight, - tensor_parallel, - mapping.tp_rank, - dim=1) - weights.update( - get_tllm_linear_weight(split_v, tllm_prex + 'mlp.proj.', None, - use_weight_only, - plugin_weight_only_quant_type, dtype, - use_gemm_woq_plugin)) - - vision_mlp_proj_weight = get_weight(model_params, - prefix + 'mlp.vision_mlp.down_proj', - dtype) - split_v = split_matrix_tp(vision_mlp_proj_weight, - tensor_parallel, - mapping.tp_rank, - dim=1) - weights.update( - get_tllm_linear_weight(split_v, tllm_prex + 'vis_mlp.proj.', None, - use_weight_only, - plugin_weight_only_quant_type, dtype, - use_gemm_woq_plugin)) - - # Layer norms do not use tensor parallelism - input_ln_weight = get_weight(model_params, prefix + 'input_layernorm', - dtype) - weights[tllm_prex + 'input_layernorm.weight'] = input_ln_weight - - post_ln_weight = get_weight(model_params, - prefix + 'post_attention_layernorm', dtype) - weights[tllm_prex + 'post_layernorm.weight'] = post_ln_weight - cur_block_weights = [ - weight_name for weight_name in model_params - if weight_name.find(prefix) != -1 - ] - for weight_name in cur_block_weights: - model_params[weight_name] = None - - if torch.cuda.is_available(): - torch.cuda.empty_cache() - torch.cuda.ipc_collect() - - v = get_weight(model_params, 'model.embed_tokens', dtype) - if hf_model.config.tie_word_embeddings: - # lm_head.weight has the same weights as embedding - if mapping.is_last_pp_rank(): - if vocab_size % mapping.tp_size != 0: - # padding - vocab_size_padded = pad_vocab_size(vocab_size, mapping.tp_size) - pad_width = vocab_size_padded - vocab_size - - v = torch.from_numpy( - np.pad(v.detach().cpu().numpy(), ((0, pad_width), (0, 0)), - 'constant', - constant_values=0)) - weights['lm_head.weight'] = split(v, mapping.tp_size, - mapping.tp_rank) - - if use_parallel_embedding: - v = split_matrix_tp(v, - mapping.tp_size, - mapping.tp_rank, - dim=sharding_dim) - - if mapping.is_first_pp_rank(): - weights['transformer.vocab_embedding.weight'] = v - - lm_head_weights = get_weight(model_params, 'lm_head', dtype) - - if mapping.is_last_pp_rank(): - if vocab_size % mapping.tp_size != 0: - # padding - vocab_size_padded = pad_vocab_size(vocab_size, mapping.tp_size) - pad_width = vocab_size_padded - vocab_size - - lm_head_weights = torch.from_numpy( - np.pad(lm_head_weights.detach().cpu().numpy(), - ((0, pad_width), (0, 0)), - 'constant', - constant_values=0)) - weights['lm_head.weight'] = split_matrix_tp(lm_head_weights, - tensor_parallel, - mapping.tp_rank, - dim=0) - ln_f_w = get_weight(model_params, 'model.norm', dtype) - weights['transformer.ln_f.weight'] = ln_f_w - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - logger.info(f'Weights loaded. Total time: {t}') - return weights diff --git a/tensorrt_llm/models/cogvlm/model.py b/tensorrt_llm/models/cogvlm/model.py deleted file mode 100644 index c5c8ea1d02d1..000000000000 --- a/tensorrt_llm/models/cogvlm/model.py +++ /dev/null @@ -1,291 +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. -from typing import Optional - -import numpy as np - -from ..._common import default_net -from ..._utils import pad_vocab_size -from ...functional import (Tensor, concat, constant, expand, op_and, recv, send, - shape, slice, unsqueeze, where) -from ...layers import (AttentionMaskType, CogVLMAttention, ColumnLinear, - Embedding, GatedMLP, PromptTuningEmbedding, RmsNorm) -from ...mapping import Mapping -from ...module import Module -# this is to use to module global algo string with a quant_algo prefix -from ...quantization import QuantMode -from ...top_model_mixin import TopModelMixin -from ..modeling_utils import (DecoderLayerList, DecoderModelForCausalLM, - QuantConfig) -from .config import CogVLMConfig - - -class CogvlmDecoderLayer(Module): - - def __init__(self, config: CogVLMConfig, layer_idx: int): - super().__init__() - self.layer_idx = layer_idx - self.config = config - - self.input_layernorm = RmsNorm(normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - dtype=config.dtype) - - layers_range = config.mapping.pp_layers(config.num_hidden_layers) - local_layer_idx = layer_idx - layers_range[0] - self.attention = CogVLMAttention( - local_layer_idx=local_layer_idx, - hidden_size=config.hidden_size, - num_attention_heads=config.num_attention_heads, - num_kv_heads=config.num_key_value_heads, - max_position_embeddings=config.max_position_embeddings, - dtype=config.dtype, - attention_mask_type=AttentionMaskType.causal, - bias=config.attn_bias, - position_embedding_type=config.position_embedding_type, - rotary_embedding_base=config.rotary_base, - rotary_embedding_scaling=config.rotary_scaling, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size, - tp_rank=config.mapping.tp_rank, - quant_mode=config.quant_mode) - - mlp_hidden_size = config.hidden_size * 4 if config.intermediate_size is None else config.intermediate_size - - self.hidden_size = config.hidden_size - self.mlp = GatedMLP(hidden_size=config.hidden_size, - ffn_hidden_size=mlp_hidden_size, - hidden_act=config.hidden_act, - dtype=config.dtype, - bias=config.mlp_bias, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size, - quant_mode=config.quant_mode) - self.vis_mlp = GatedMLP(hidden_size=config.hidden_size, - ffn_hidden_size=mlp_hidden_size, - hidden_act=config.hidden_act, - dtype=config.dtype, - bias=config.mlp_bias, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size, - quant_mode=config.quant_mode) - self.post_layernorm = RmsNorm(normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - dtype=config.dtype) - - def forward( - self, - hidden_states, - attention_mask=None, - use_cache=False, - kv_cache_params=None, - attention_params=None, - lora_layer_params=None, - vision_token_mask=None, - position_ids=None, - ): - residual = hidden_states - hidden_states = self.input_layernorm(hidden_states) - - attention_output = self.attention(hidden_states, - use_cache=use_cache, - kv_cache_params=kv_cache_params, - attention_params=attention_params, - vision_token_mask=vision_token_mask, - position_embedding=position_ids) - - if use_cache: - attention_output, presents = attention_output - - hidden_states = residual + attention_output - - residual = hidden_states - hidden_states = self.post_layernorm(hidden_states) - - vision_mlp_out = self.vis_mlp(hidden_states) - language_mlp_out = self.mlp(hidden_states) - hidden_states = where(vision_token_mask, vision_mlp_out, - language_mlp_out) - - # hidden_states = self.mlp(hidden_states, - # lora_layer_params=lora_layer_params) - - hidden_states = residual + hidden_states - if use_cache: - return (hidden_states, presents) - return hidden_states - - -class CogvlmModel(Module): - - def __init__(self, config: CogVLMConfig) -> None: - super().__init__() - - self.mapping = config.mapping - self.use_prompt_tuning = config.use_prompt_tuning - self.vocab_size = config.vocab_size - EmbeddingCls = PromptTuningEmbedding if config.use_prompt_tuning else Embedding - if self.mapping.is_first_pp_rank(): - self.vocab_embedding = EmbeddingCls( - num_embeddings=config.vocab_size, - embedding_dim=config.hidden_size, - dtype=config.dtype, - tp_size=self.mapping.tp_size - if config.use_parallel_embedding else 1, - tp_group=self.mapping.tp_group - if config.use_parallel_embedding else None, - sharding_dim=config.embedding_sharding_dim, - tp_rank=self.mapping.tp_rank, - ) - - self.layers = DecoderLayerList(CogvlmDecoderLayer, config) - - if self.mapping.is_last_pp_rank(): - self.ln_f = RmsNorm(normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - dtype=config.dtype) - - def forward(self, - input_ids, - position_ids=None, - use_cache=False, - attention_mask=None, - kv_cache_params=None, - attention_params=None, - hidden_states=None, - prompt_embedding_table: Optional[Tensor] = None, - prompt_tasks: Optional[Tensor] = None, - prompt_vocab_size: Optional[Tensor] = None, - lora_params=None): - - kv_cache_params.fill_none_tensor_list(len(self.layers)) - - if use_cache: - presents = [] - - ptuning_args = [ - prompt_embedding_table, prompt_tasks, prompt_vocab_size - ] if self.use_prompt_tuning else [] - - if self.mapping.is_first_pp_rank(): - hidden_states = self.vocab_embedding(input_ids, *ptuning_args) - else: - hidden_states = recv(hidden_states, self.mapping.prev_pp_rank()) - - vision_mask = input_ids > (self.vocab_size - 1) - - if default_net().plugin_config.remove_input_padding: - seq_length = shape(vision_mask, 0) # lvvvvvvllvvvvlll - zero = constant(np.ascontiguousarray(np.zeros([1], dtype=bool))) - one = constant(np.ascontiguousarray(np.ones([1], dtype=bool))) - - t1 = slice(vision_mask, [0], seq_length - 1) - t2 = slice(vision_mask, [1], seq_length - 1) - vision_token_mask = concat([op_and(t1 == one, t2 == one), - zero]) # 0111110001110000 - vision_token_mask = unsqueeze(vision_token_mask, - -1) # [num_tokens, 1] - else: - seq_length = shape(vision_mask, - 1) # lvvvvvvllvvvvlll, lvvvvvvllvvvvlll - batch_size = shape(vision_mask, 0) - t1 = slice(vision_mask, [0, 0], concat([batch_size, - seq_length - 1])) - t2 = slice(vision_mask, [0, 1], concat([batch_size, - seq_length - 1])) - zero = expand( - constant(np.ascontiguousarray(np.zeros([1, 1], dtype=bool))), - concat([batch_size, 1])) - one = constant(np.ascontiguousarray(np.ones([1, 1], dtype=bool))) - - vision_token_mask = concat([op_and(t1 == one, t2 == one), zero], - dim=1) # 0111110001110000 [bs, seqlen] - vision_token_mask = unsqueeze(vision_token_mask, - -1) # [bs, seqlen, 1] - - hidden_states = self.layers.forward(hidden_states, - use_cache=use_cache, - attention_mask=attention_mask, - kv_cache_params=kv_cache_params, - attention_params=attention_params, - lora_params=lora_params, - vision_token_mask=vision_token_mask, - position_ids=position_ids) - - if use_cache: - hidden_states, presents = hidden_states - - if self.mapping.is_last_pp_rank(): - hidden_states = self.ln_f(hidden_states) - else: - hidden_states = send(hidden_states, self.mapping.next_pp_rank()) - - if use_cache: - return (hidden_states, tuple(presents)) - return hidden_states - - -class CogVLMForCausalLM(DecoderModelForCausalLM, TopModelMixin): - config_class = CogVLMConfig - - def __init__(self, config: CogVLMConfig): - transformer = CogvlmModel(config) - vocab_size_padded = pad_vocab_size(config.vocab_size, - config.mapping.tp_size) - if config.mapping.is_last_pp_rank(): - lm_head = ColumnLinear(config.hidden_size, - vocab_size_padded, - bias=False, - dtype=config.dtype, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size, - gather_output=True) - else: - lm_head = None - self.quant_mode = config.quant_mode - self.mapping = config.mapping - super().__init__(config, transformer, lm_head) - - @classmethod - def from_hugging_face(cls, - hf_model_dir, - dtype='float16', - mapping: Optional[Mapping] = None, - quant_mode: Optional[QuantMode] = None, - **kwargs): - pass - - def default_plugin_config(self, **kwargs): - plugin_config = super().default_plugin_config(**kwargs) - if self.quant_mode.is_int4_weight_only_per_group(): - plugin_config.weight_only_groupwise_quant_matmul_plugin = 'auto' - return plugin_config - - @classmethod - def quantize( - cls, - hf_model_dir, - output_dir, - quant_config: QuantConfig, - *, - dtype='float16', - mapping: Optional[Mapping] = None, - calib_batches=512, - calib_batch_size=1, - random_seed=1234, - tokenizer_max_seq_length=2048, - **kwargs, - ): - pass diff --git a/tensorrt_llm/models/commandr/__init__.py b/tensorrt_llm/models/commandr/__init__.py deleted file mode 100644 index 71bf6d298d42..000000000000 --- a/tensorrt_llm/models/commandr/__init__.py +++ /dev/null @@ -1,14 +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. diff --git a/tensorrt_llm/models/commandr/config.py b/tensorrt_llm/models/commandr/config.py deleted file mode 100644 index 511640c22494..000000000000 --- a/tensorrt_llm/models/commandr/config.py +++ /dev/null @@ -1,88 +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. -from typing import Optional, Union - -import transformers - -from ..._utils import get_hf_rope_theta -from ...mapping import Mapping -from ..convert_utils import infer_dtype -from ..modeling_utils import PretrainedConfig, QuantConfig - - -class CohereConfig(PretrainedConfig): - - def __init__(self, - *, - output_multiplier_scale: float = 0.0625, - rotary_base: float = 10000.0, - attn_bias: bool = False, - **kwargs): - self.output_multiplier_scale = output_multiplier_scale - self.rotary_base = rotary_base - self.attn_bias = attn_bias - super().__init__(**kwargs) - - def to_dict(self): - output = super().to_dict() - # Serialize the fields added in CohereConfig - output['output_multiplier_scale'] = self.output_multiplier_scale - output['rotary_base'] = self.rotary_base - output['attn_bias'] = self.attn_bias - return output - - @classmethod - def from_hugging_face( - cls, - hf_config_or_dir: Union[str, 'transformers.PretrainedConfig'], - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - if isinstance(hf_config_or_dir, transformers.PretrainedConfig): - hf_config = hf_config_or_dir - else: - hf_config = transformers.AutoConfig.from_pretrained( - hf_config_or_dir, trust_remote_code=True) - - head_size = hf_config.hidden_size // hf_config.num_attention_heads - - dtype = infer_dtype(dtype, getattr(hf_config, 'torch_dtype', None)) - - if hf_config.tie_word_embeddings: - kwargs['use_parallel_embedding'] = True - kwargs['embedding_sharding_dim'] = 0 - - return CohereConfig( - architecture=hf_config.architectures[0], - dtype=dtype, - num_hidden_layers=hf_config.num_hidden_layers, - num_attention_heads=hf_config.num_attention_heads, - hidden_size=hf_config.hidden_size, - intermediate_size=hf_config.intermediate_size, - num_key_value_heads=hf_config.num_key_value_heads, - head_size=head_size, - vocab_size=hf_config.vocab_size, - position_embedding_type='rope_gptj', # different rope type - max_position_embeddings=hf_config.max_position_embeddings, - hidden_act=hf_config.hidden_act, - norm_epsilon=hf_config.layer_norm_eps, - output_multiplier_scale=hf_config.logit_scale, - rotary_base=get_hf_rope_theta(hf_config, 10000.0), - attn_bias=hf_config.attention_bias, - qk_layernorm=hf_config.use_qk_norm, - mapping=mapping, - quantization=quant_config, - **kwargs) diff --git a/tensorrt_llm/models/commandr/model.py b/tensorrt_llm/models/commandr/model.py deleted file mode 100644 index ac0fe67dd754..000000000000 --- a/tensorrt_llm/models/commandr/model.py +++ /dev/null @@ -1,195 +0,0 @@ -from typing import Optional - -from ..._common import default_net -from ..._utils import pad_vocab_size -from ...functional import recv, send -from ...layers import (Attention, AttentionMaskType, ColumnLinear, Embedding, - GatedMLP, LayerNorm, PositionEmbeddingType) -from ...mapping import Mapping -from ...module import Module -from ..model_weights_loader import ModelWeightsLoader -from ..modeling_utils import (DecoderLayerList, DecoderModelForCausalLM, - QuantConfig) -from .config import CohereConfig - - -class CohereDecoderLayer(Module): - - def __init__(self, config: CohereConfig, layer_idx: int): - super().__init__() - self.layer_idx = layer_idx - self.config = config - - self.input_layernorm = LayerNorm(normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - bias=False, - dtype=config.dtype) - - layers_range = config.mapping.pp_layers(config.num_hidden_layers) - self.local_layer_idx = layer_idx - layers_range[0] - self.attention = Attention( - local_layer_idx=self.local_layer_idx, - hidden_size=config.hidden_size, - attention_head_size=config.head_size, - num_attention_heads=config.num_attention_heads, - num_kv_heads=config.num_key_value_heads, - max_position_embeddings=config.max_position_embeddings, - dtype=config.dtype, - attention_mask_type=AttentionMaskType.causal, - bias=config.attn_bias, - position_embedding_type=PositionEmbeddingType.rope_gptj, - rotary_embedding_base=config.rotary_base, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size, - tp_rank=config.mapping.tp_rank, - qk_layernorm=config.qk_layernorm, - layernorm_share=False, - eps=config.norm_epsilon, - quant_mode=config.quant_mode) - - self.mlp = GatedMLP(hidden_size=config.hidden_size, - ffn_hidden_size=config.intermediate_size, - hidden_act=config.hidden_act, - dtype=config.dtype, - bias=False, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size, - quant_mode=config.quant_mode) - - def forward(self, - hidden_states, - attention_mask=None, - use_cache=False, - spec_decoding_params=None, - kv_cache_params=None, - attention_params=None): - assert not ( - default_net().plugin_config.reduce_fusion - ), "Custom all reduce and residual mlp can't be enabled at the same time." - - residual = hidden_states - hidden_states = self.input_layernorm(hidden_states) - - attention_output = self.attention( - hidden_states, - attention_mask=attention_mask, - use_cache=use_cache, - spec_decoding_params=spec_decoding_params, - kv_cache_params=kv_cache_params, - attention_params=attention_params) - - if use_cache: - attention_output, presents = attention_output - - mlp_output = self.mlp(hidden_states) - hidden_states = residual + attention_output + mlp_output - - if use_cache: - return (hidden_states, presents) - return hidden_states - - -class CohereModel(Module): - - def __init__(self, config: CohereConfig) -> None: - super().__init__() - - self.mapping = config.mapping - if self.mapping.is_first_pp_rank(): - self.vocab_embedding = Embedding(config.vocab_size, - config.hidden_size, - dtype=config.dtype, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size, - tp_rank=config.mapping.tp_rank) - - self.layers = DecoderLayerList(CohereDecoderLayer, config) - - if self.mapping.is_last_pp_rank(): - self.ln_f = LayerNorm(normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - bias=False, - dtype=config.dtype) - - def forward( - self, - input_ids=None, - position_ids=None, - use_cache=False, - attention_mask=None, - spec_decoding_params=None, - kv_cache_params=None, - attention_params=None, - hidden_states=None, - ): - if self.mapping.is_first_pp_rank(): - hidden_states = self.vocab_embedding(input_ids) - else: - hidden_states = recv(hidden_states, self.mapping.prev_pp_rank()) - - hidden_states = self.layers.forward( - hidden_states, - use_cache=use_cache, - attention_mask=attention_mask, - kv_cache_params=kv_cache_params, - attention_params=attention_params, - spec_decoding_params=spec_decoding_params) - - if use_cache: - hidden_states, presents = hidden_states - - if self.mapping.is_last_pp_rank(): - hidden_states = self.ln_f(hidden_states) - else: - hidden_states = send(hidden_states, self.mapping.next_pp_rank()) - - if use_cache: - return (hidden_states, presents) - return hidden_states - - -class CohereForCausalLM(DecoderModelForCausalLM): - config_class = CohereConfig - - def __init__(self, config: CohereConfig): - transformer = CohereModel(config) - vocab_size_padded = pad_vocab_size(config.vocab_size, - config.mapping.tp_size) - if config.mapping.is_last_pp_rank(): - lm_head = ColumnLinear(config.hidden_size, - vocab_size_padded, - bias=False, - dtype=config.dtype, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size, - gather_output=True) - else: - lm_head = None - self.quant_mode = config.quant_mode - self.mapping = config.mapping - super().__init__(config, transformer, lm_head) - - @classmethod - def from_hugging_face(cls, - hf_model_or_dir: str, - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - ''' Create a CohereForCausalLM object from give parameters - ''' - - config = CohereConfig.from_hugging_face(hf_model_or_dir, - dtype=dtype, - mapping=mapping, - quant_config=quant_config, - **kwargs) - model = cls(config) - custom_dict = { - 'q_layernorm': 'q_norm', - 'k_layernorm': 'k_norm', - } - loader = ModelWeightsLoader(hf_model_or_dir, custom_dict) - loader.generate_tllm_weights(model) - - return model diff --git a/tensorrt_llm/models/dbrx/__init__.py b/tensorrt_llm/models/dbrx/__init__.py deleted file mode 100644 index 71bf6d298d42..000000000000 --- a/tensorrt_llm/models/dbrx/__init__.py +++ /dev/null @@ -1,14 +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. diff --git a/tensorrt_llm/models/dbrx/config.py b/tensorrt_llm/models/dbrx/config.py deleted file mode 100644 index d97311ff88ca..000000000000 --- a/tensorrt_llm/models/dbrx/config.py +++ /dev/null @@ -1,59 +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. - -from typing import Optional, Union - -from ...layers import MoeConfig -from ..modeling_utils import PretrainedConfig - - -class DbrxConfig(PretrainedConfig): - - def __init__(self, - *, - bias: bool = False, - clip_qkv: Optional[float] = None, - rotary_base: float = 500000.0, - rotary_scaling: Optional[dict] = None, - moe: Optional[Union[MoeConfig, dict]] = None, - **kwargs): - self.bias = bias - self.clip_qkv = clip_qkv - self.rotary_base = rotary_base - self.rotary_scaling = rotary_scaling - if moe is None: - # Legacy MOE config fields - moe = MoeConfig( - num_experts=kwargs.pop('moe_num_experts', 0), - top_k=kwargs.pop('moe_top_k', 0), - normalization_mode=kwargs.pop( - 'moe_normalization_mode', - MoeConfig.ExpertScaleNormalizationMode.RENORMALIZE)) - elif isinstance(moe, dict): - moe = MoeConfig.from_dict(moe) - assert isinstance(moe, MoeConfig) - self.moe = moe.validate() - - super().__init__(**kwargs) - - def to_dict(self): - output = super().to_dict() - # Serialize the fields added in DbrxConfig - output['bias'] = self.bias - output['clip_qkv'] = self.clip_qkv - output['rotary_base'] = self.rotary_base - output['rotary_scaling'] = self.rotary_scaling - output['moe'] = self.moe.to_dict() - return output diff --git a/tensorrt_llm/models/dbrx/model.py b/tensorrt_llm/models/dbrx/model.py deleted file mode 100644 index 14617e0188aa..000000000000 --- a/tensorrt_llm/models/dbrx/model.py +++ /dev/null @@ -1,188 +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. - -from ..._utils import pad_vocab_size -from ...functional import Tensor, recv, send -from ...layers import (MOE, Attention, AttentionMaskType, ColumnLinear, - Embedding, GatedMLP, LayerNorm) -from ...module import Module -from ..modeling_utils import DecoderLayerList, DecoderModelForCausalLM -from .config import DbrxConfig - - -class DbrxDecoderLayer(Module): - - def __init__(self, config: DbrxConfig, layer_idx: int): - super().__init__() - self.layer_idx = layer_idx - self.config = config - - self.input_layernorm = LayerNorm(normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - bias=False, - dtype=config.dtype) - - layers_range = config.mapping.pp_layers(config.num_hidden_layers) - local_layer_idx = layer_idx - layers_range[0] - self.attention = Attention( - local_layer_idx=local_layer_idx, - hidden_size=config.hidden_size, - num_attention_heads=config.num_attention_heads, - num_kv_heads=config.num_key_value_heads, - max_position_embeddings=config.max_position_embeddings, - dtype=config.dtype, - attention_mask_type=AttentionMaskType.causal, - bias=config.bias, - position_embedding_type=config.position_embedding_type, - rotary_embedding_base=config.rotary_base, - rotary_embedding_scaling=config.rotary_scaling, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size, - quant_mode=config.quant_mode, - clip_qkv=config.clip_qkv) - - ClsMLP = GatedMLP - mlp_kwargs = {} - if config.moe.has_moe(): - ClsMLP = MOE - mlp_kwargs = { - "moe_config": config.moe, - "mapping": config.mapping, - } - - self.mlp = ClsMLP(hidden_size=config.hidden_size, - ffn_hidden_size=config.intermediate_size, - hidden_act=config.hidden_act, - dtype=config.dtype, - bias=config.bias, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size, - quant_mode=config.quant_mode, - **mlp_kwargs) - self.post_layernorm = LayerNorm(normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - bias=False, - dtype=config.dtype) - - def forward(self, - hidden_states: Tensor, - attention_mask=None, - use_cache=False, - kv_cache_params=None, - attention_params=None): - - assert isinstance(hidden_states, Tensor) - - residual = hidden_states - - hidden_states = self.input_layernorm(hidden_states) - - attention_output = self.attention(hidden_states, - attention_mask=attention_mask, - use_cache=use_cache, - kv_cache_params=kv_cache_params, - attention_params=attention_params) - - if use_cache: - attention_output, presents = attention_output - - hidden_states = residual + attention_output - - residual = hidden_states - hidden_states = self.post_layernorm(hidden_states) - - hidden_states = self.mlp(hidden_states) - - hidden_states = residual + hidden_states - - if use_cache: - return (hidden_states, presents) - return hidden_states - - -class DbrxModel(Module): - - def __init__(self, config: DbrxConfig): - super().__init__() - self.config = config - - if config.mapping.is_first_pp_rank(): - self.vocab_embedding = Embedding(config.vocab_size, - config.hidden_size, - dtype=config.dtype) - - self.layers = DecoderLayerList(DbrxDecoderLayer, config) - - if config.mapping.is_last_pp_rank(): - self.ln_f = LayerNorm(normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - bias=False, - dtype=config.dtype) - - def forward(self, - input_ids, - position_ids, - use_cache=False, - attention_mask=None, - kv_cache_params=None, - attention_params=None, - hidden_states=None): - - if self.config.mapping.is_first_pp_rank(): - hidden_states = self.vocab_embedding(input_ids) - else: - hidden_states = recv(hidden_states, - self.config.mapping.prev_pp_rank()) - - hidden_states = self.layers(hidden_states, - use_cache=use_cache, - attention_mask=attention_mask, - kv_cache_params=kv_cache_params, - attention_params=attention_params) - - if use_cache: - hidden_states, presents = hidden_states - - if self.config.mapping.is_last_pp_rank(): - hidden_states = self.ln_f(hidden_states) - else: - hidden_states = send(hidden_states, - self.config.mapping.next_pp_rank()) - - if use_cache: - return (hidden_states, tuple(presents)) - return hidden_states - - -class DbrxForCausalLM(DecoderModelForCausalLM): - config_class = DbrxConfig - - def __init__(self, config: DbrxConfig): - transformer = DbrxModel(config) - vocab_size_padded = pad_vocab_size(config.vocab_size, - config.mapping.tp_size) - if config.mapping.is_last_pp_rank(): - lm_head = ColumnLinear(config.hidden_size, - vocab_size_padded, - bias=config.bias, - dtype=config.dtype, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size, - gather_output=True) - else: - lm_head = None - self.quant_mode = config.quant_mode - self.mapping = config.mapping - super().__init__(config, transformer, lm_head) diff --git a/tensorrt_llm/models/deepseek_v1/__init__.py b/tensorrt_llm/models/deepseek_v1/__init__.py deleted file mode 100644 index 71bf6d298d42..000000000000 --- a/tensorrt_llm/models/deepseek_v1/__init__.py +++ /dev/null @@ -1,14 +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. diff --git a/tensorrt_llm/models/deepseek_v1/config.py b/tensorrt_llm/models/deepseek_v1/config.py deleted file mode 100755 index e7bff0d9aab8..000000000000 --- a/tensorrt_llm/models/deepseek_v1/config.py +++ /dev/null @@ -1,105 +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. - -from typing import Optional, Union - -from ..._utils import get_hf_rope_theta -from ...layers import MoeConfig -from ...mapping import Mapping -from ..convert_utils import infer_dtype -from ..modeling_utils import PretrainedConfig, QuantConfig - - -class DeepSeekV1Config(PretrainedConfig): - - def __init__(self, - *, - rotary_base: float = 10000.0, - rotary_scaling: Optional[dict] = None, - moe: Optional[Union[MoeConfig, dict]] = None, - **kwargs): - - self.rotary_base = rotary_base - self.rotary_scaling = rotary_scaling - if isinstance(moe, dict): - moe = MoeConfig.from_dict(moe) - assert isinstance(moe, MoeConfig) - self.moe = moe.validate() - super().__init__(**kwargs) - - def to_dict(self): - output = super().to_dict() - # Serialize the fields added in DeepSeekV1Config - - output['rotary_base'] = self.rotary_base - output['rotary_scaling'] = self.rotary_scaling - output['moe'] = self.moe.to_dict() - - return output - - @classmethod - def from_hugging_face( - cls, - hf_config_or_dir: Union[str, 'transformers.PretrainedConfig'], - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - import transformers - trust_remote_code = kwargs.pop('trust_remote_code', True) - - if isinstance(hf_config_or_dir, transformers.PretrainedConfig): - hf_config = hf_config_or_dir - else: - hf_config_dir = str(hf_config_or_dir) - - hf_config = transformers.AutoConfig.from_pretrained( - hf_config_dir, trust_remote_code=trust_remote_code) - - num_key_value_heads = getattr(hf_config, "num_key_value_heads", - hf_config.num_attention_heads) - rotary_scaling = getattr(hf_config, "rope_scaling", None) - rotary_base = get_hf_rope_theta(hf_config, 10000.0) - dtype = infer_dtype(dtype, getattr(hf_config, 'torch_dtype', None)) - moe_config = MoeConfig( - num_experts=getattr(hf_config, 'n_routed_experts', 0), - shared_expert_intermediate_size=getattr(hf_config, - 'n_shared_experts', 0) * - getattr(hf_config, "moe_intermediate_size", 0), - top_k=getattr(hf_config, 'num_experts_per_tok', 0), - normalization_mode=getattr( - hf_config, 'moe_normalization_mode', - MoeConfig.ExpertScaleNormalizationMode.NONE), - ) - moe_config.validate() - return cls(architecture=hf_config.architectures[0], - dtype=dtype, - num_hidden_layers=hf_config.num_hidden_layers, - num_attention_heads=hf_config.num_attention_heads, - hidden_size=hf_config.hidden_size, - intermediate_size=hf_config.intermediate_size, - num_key_value_heads=num_key_value_heads, - vocab_size=hf_config.vocab_size, - position_embedding_type='rope_gpt_neox', - max_position_embeddings=hf_config.max_position_embeddings, - hidden_act='swiglu', - rotary_base=rotary_base, - rotary_scaling=rotary_scaling, - norm_epsilon=hf_config.rms_norm_eps, - mapping=mapping, - quantization=quant_config, - moe=moe_config, - moe_intermediate_size=hf_config.moe_intermediate_size, - **kwargs) diff --git a/tensorrt_llm/models/deepseek_v1/convert.py b/tensorrt_llm/models/deepseek_v1/convert.py deleted file mode 100755 index 6871ebbdad2f..000000000000 --- a/tensorrt_llm/models/deepseek_v1/convert.py +++ /dev/null @@ -1,290 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import time - -import torch -from transformers import AutoModelForCausalLM - -from ..._utils import pad_vocab_size, release_gc - - -## Get HF model -def load_hf_deepseek(model_dir): - model = AutoModelForCausalLM.from_pretrained(model_dir, - device_map='auto', - dtype='auto', - trust_remote_code=True) - return model - - -## Prepare weights for TP -def split(v, tp_size, idx, dim=0): - if tp_size == 1: - return v - if len(v.shape) == 1: - return torch.chunk(v, tp_size)[idx].contiguous() - else: - return torch.chunk(v, tp_size, dim=dim)[idx].contiguous() - - -def split_qkv_tp(v, n_head, n_hidden, tensor_parallel, rank): - """ - Splits the QKV matrix according to tensor parallelism - """ - v = v.reshape(3, n_hidden, n_hidden) - split_v = split(v, tensor_parallel, rank, dim=1) - split_v = split_v.reshape(3 * (n_hidden // tensor_parallel), n_hidden) - return split_v.contiguous() - - -def split_matrix_tp(v, tensor_parallel, rank, dim): - return split(v, tensor_parallel, rank, dim=dim) - - -def get_weight(config, prefix, dtype, postfix='.weight'): - if config[prefix + postfix].dtype != dtype: - config[prefix + postfix].data = config[prefix + postfix].to(dtype) - return config[prefix + postfix].detach().cpu() - - -def get_trtllm_linear_weight(weight, prefix, postfix='weight'): - results = {} - results[prefix + postfix] = weight - - return results - - -def convert_deepseek(hf_model, - config, - mapping, - dtype='float32', - use_parallel_embedding=False, - sharding_dim=0): - - weights = {} - tik = time.time() - mapping.tp_size - model_params = dict(hf_model.named_parameters()) - dtype = getattr(torch, dtype) - moe_config = config.moe - layers_range = mapping.pp_layers(config.num_hidden_layers) - - def convert_layer(l): - prefix = f'model.layers.{l}.' - print(prefix) - trtllm_prex = f'transformer.layers.{l - layers_range[0]}.' - q_weight = get_weight(model_params, prefix + 'self_attn.q_proj', dtype) - k_weight = get_weight(model_params, prefix + 'self_attn.k_proj', dtype) - v_weight = get_weight(model_params, prefix + 'self_attn.v_proj', dtype) - - qkv_weight = torch.cat([q_weight, k_weight, v_weight], dim=0) - - split_v = split_qkv_tp(qkv_weight, config.num_attention_heads, - config.hidden_size, mapping.tp_size, - mapping.tp_rank) - - weights.update( - get_trtllm_linear_weight(split_v, trtllm_prex + 'attention.qkv.')) - - attn_dense_weight = get_weight(model_params, - prefix + 'self_attn.o_proj', dtype) - split_v = split_matrix_tp(attn_dense_weight, - mapping.tp_size, - mapping.tp_rank, - dim=1) - - weights.update( - get_trtllm_linear_weight(split_v, trtllm_prex + 'attention.dense.')) - - if moe_config.has_moe() and l > 0: - rank_experts = list(range(moe_config.num_experts)) - if mapping.has_moe_ep(): - rank_experts = mapping.ep_experts(moe_config.num_experts) - for suffix in ["gate_proj", "down_proj", "up_proj"]: - model_params[f'model.layers.{l}.mlp.experts.{suffix}.weight'] = \ - torch.stack([model_params[f'model.layers.{l}.mlp.experts.{expert}.{suffix}.weight'].detach().cpu() - for expert in rank_experts]) - - gate_proj = model_params[ - f'model.layers.{l}.mlp.experts.gate_proj.weight'] - down_proj = model_params[ - f'model.layers.{l}.mlp.experts.down_proj.weight'] - up_proj = model_params[ - f'model.layers.{l}.mlp.experts.up_proj.weight'] - if mapping.has_moe_tp(): - gate_proj = split(gate_proj, - mapping.tp_size, - mapping.tp_rank, - dim=1) - down_proj = split(down_proj, - mapping.tp_size, - mapping.tp_rank, - dim=2) - up_proj = split(up_proj, - mapping.tp_size, - mapping.tp_rank, - dim=1) - - model_params[ - f'model.layers.{l}.mlp.experts.up_gate_proj.weight'] = torch.concat( - [up_proj, gate_proj], dim=-2) - model_params[ - f'model.layers.{l}.mlp.experts.down_proj.weight'] = down_proj - - ## mlp.experts.down_proj.weight - moe_experts_down_proj_weights = get_weight( - model_params, prefix + 'mlp.experts.down_proj', dtype) - weights.update( - get_trtllm_linear_weight(moe_experts_down_proj_weights, - trtllm_prex + 'mlp.proj.')) - ##mlp.experts.up_gate.weight - moe_experts_up_gate_proj_weights = get_weight( - model_params, prefix + 'mlp.experts.up_gate_proj', dtype) - weights.update( - get_trtllm_linear_weight(moe_experts_up_gate_proj_weights, - trtllm_prex + 'mlp.fc.')) - ## MOE hardcoded routing_input into trt.float32, please refer to moe.py line 397 - moe_experts_gate_weights = get_weight(model_params, - prefix + 'mlp.gate', - torch.float32) - weights.update( - get_trtllm_linear_weight(moe_experts_gate_weights, - trtllm_prex + 'mlp.router.')) - - if moe_config.shared_expert_intermediate_size > 0: - shared_moe_up_proj_weights = get_weight( - model_params, prefix + 'mlp.shared_experts.up_proj', dtype) - shared_moe_up_proj_weights = split_matrix_tp( - shared_moe_up_proj_weights, - mapping.tp_size, - mapping.tp_rank, - dim=0) - shared_moe_down_proj_weights = get_weight( - model_params, prefix + 'mlp.shared_experts.down_proj', - dtype) - shared_moe_down_proj_weights = split_matrix_tp( - shared_moe_down_proj_weights, - mapping.tp_size, - mapping.tp_rank, - dim=1) - shared_moe_gate_proj_weights = get_weight( - model_params, prefix + 'mlp.shared_experts.gate_proj', - dtype) - shared_moe_gate_proj_weights = split_matrix_tp( - shared_moe_gate_proj_weights, - mapping.tp_size, - mapping.tp_rank, - dim=0) - shared_moe_gate_up_proj_weights = torch.concat( - [shared_moe_up_proj_weights, shared_moe_gate_proj_weights], - dim=-2) - - ## mlp.shared_experts.gate_up_proj.weight - weights.update( - get_trtllm_linear_weight( - shared_moe_gate_up_proj_weights, - trtllm_prex + 'mlp.shared_expert.fc.')) - - ## mlp.shared_experts.down_proj.weight - weights.update( - get_trtllm_linear_weight( - shared_moe_down_proj_weights, - trtllm_prex + 'mlp.shared_expert.proj.')) - - else: - ## Current deepseek model has one MLP layer only, if it goes large consider to do fuse - mlp_gate_weight = get_weight(model_params, prefix + 'mlp.up_proj', - dtype) - split_gate = split_matrix_tp(mlp_gate_weight, - mapping.tp_size, - mapping.tp_rank, - dim=0) - weights.update( - get_trtllm_linear_weight(split_gate, trtllm_prex + 'mlp.gate.')) - - mlp_fc_weight = get_weight(model_params, prefix + 'mlp.gate_proj', - dtype) - split_fc = split_matrix_tp(mlp_fc_weight, - mapping.tp_size, - mapping.tp_rank, - dim=0) - weights.update( - get_trtllm_linear_weight(split_fc, trtllm_prex + 'mlp.fc.')) - - mlp_proj_weight = get_weight(model_params, prefix + 'mlp.down_proj', - dtype) - split_proj = split_matrix_tp(mlp_proj_weight, - mapping.tp_size, - mapping.tp_rank, - dim=1) - weights.update( - get_trtllm_linear_weight(split_proj, trtllm_prex + 'mlp.proj.')) - - # Layer norms do not use tensor parallelism - input_ln_weight = get_weight(model_params, prefix + 'input_layernorm', - dtype) - weights[trtllm_prex + 'input_layernorm.weight'] = input_ln_weight - post_ln_weight = get_weight(model_params, - prefix + 'post_attention_layernorm', dtype) - weights[trtllm_prex + 'post_layernorm.weight'] = post_ln_weight - - for l in layers_range: - convert_layer(l) - release_gc() - - v = get_weight(model_params, 'model.embed_tokens', dtype) - if hf_model.config.tie_word_embeddings: - # lm_head.weight has the same weights as embedding - if mapping.is_last_pp_rank(): - if config.vocab_size % mapping.tp_size != 0: - # padding - vocab_size_padded = pad_vocab_size(config.vocab_size, - mapping.tp_size) - pad_width = vocab_size_padded - config.vocab_size - v = torch.nn.functional.pad(v, (0, 0, 0, pad_width), 'constant', - 0) - weights['lm_head.weight'] = split(v, mapping.tp_size, - mapping.tp_rank) - if use_parallel_embedding: - v = split_matrix_tp(v, - mapping.tp_size, - mapping.tp_rank, - dim=config.embedding_sharding_dim) - if mapping.is_first_pp_rank(): - weights['transformer.vocab_embedding.weight'] = v - lm_head_weights = get_weight(model_params, 'lm_head', dtype) - - if mapping.is_last_pp_rank(): - if config.vocab_size % mapping.tp_size != 0: - # padding - vocab_size_padded = pad_vocab_size(config.vocab_size, - mapping.tp_size) - pad_width = vocab_size_padded - config.vocab_size - lm_head_weights = torch.nn.functional.pad(lm_head_weights, - (0, 0, 0, pad_width), - 'constant', - value=0) - weights['lm_head.weight'] = split_matrix_tp(lm_head_weights, - mapping.tp_size, - mapping.tp_rank, - dim=0) - ln_f_w = get_weight(model_params, 'model.norm', dtype) - weights['transformer.ln_f.weight'] = ln_f_w - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Weights loaded. Total time: {t}') - #print(set(weights.keys())) - return weights diff --git a/tensorrt_llm/models/deepseek_v1/model.py b/tensorrt_llm/models/deepseek_v1/model.py deleted file mode 100755 index 164f90387a92..000000000000 --- a/tensorrt_llm/models/deepseek_v1/model.py +++ /dev/null @@ -1,279 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import os -from typing import Optional - -from ..._utils import pad_vocab_size -from ...functional import Tensor, non_gated_version, recv, send -from ...layers import (Attention, AttentionMaskType, ColumnLinear, Embedding, - GatedMLP, PositionEmbeddingType, RmsNorm, SharedMoE) -from ...mapping import Mapping -from ...module import Module -from ...plugin import init_all_reduce_helper -from ..model_weights_loader import ModelWeightsLoader -from ..modeling_utils import (DecoderLayerList, DecoderModelForCausalLM, - PretrainedConfig) -from .config import DeepSeekV1Config -from .convert import convert_deepseek, load_hf_deepseek - - -class DeepseekDecoderLayer(Module): - - def __init__(self, config: PretrainedConfig, layer_idx: int): - super().__init__() - self.layer_idx = layer_idx - self.config = config - - ### Input layernorm in Deepseek v1 is same as Llama - self.input_layernorm = RmsNorm(normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - dtype=config.dtype) - - layers_range = config.mapping.pp_layers(config.num_hidden_layers) - local_layer_idx = layer_idx - layers_range[0] - ### Deepseek v1 model with standard attention - self.attention = Attention( - local_layer_idx=local_layer_idx, - hidden_size=config.hidden_size, - attention_head_size=config.head_size, - num_attention_heads=config.num_attention_heads, - num_kv_heads=config.num_key_value_heads, - max_position_embeddings=config.max_position_embeddings, - dtype=config.dtype, - attention_mask_type=AttentionMaskType.causal, - bias=False, - position_embedding_type=PositionEmbeddingType.rope_gpt_neox, - rotary_embedding_base=config.rotary_base, - rotary_embedding_scaling=config.rotary_scaling, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size, - tp_rank=config.mapping.tp_rank, - quant_mode=config.quant_mode) - - ClsMLP = GatedMLP - moe_config = config.moe - - if moe_config.num_experts > 0 and layer_idx > 0: - mlp_hidden_size = config.moe_intermediate_size - hidden_act = config.hidden_act - mlp_kwargs = {'moe_config': moe_config, 'mapping': config.mapping} - if moe_config.shared_expert_intermediate_size > 0: - ClsMLP = SharedMoE - mlp_kwargs['use_shared_gate'] = False - mlp_kwargs['use_side_stream'] = False - else: - ClsMLP = MOE - else: - ClsMLP = GatedMLP - mlp_hidden_size = config.intermediate_size - hidden_act = non_gated_version( - config.hidden_act) # back to non gated for dense layers - mlp_kwargs = {} - - self.mlp = ClsMLP(hidden_size=config.hidden_size, - ffn_hidden_size=mlp_hidden_size, - hidden_act=hidden_act, - dtype=config.dtype, - bias=False, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size, - quant_mode=config.quant_mode, - **mlp_kwargs) - - ### Pose layernorm in Deepseek v1 is same as Llama ) - self.post_layernorm = RmsNorm(normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - dtype=config.dtype) - - def forward(self, - hidden_states, - attention_mask=None, - use_cache=False, - spec_decoding_params=None, - kv_cache_params=None, - attention_params=None): - residual = hidden_states - hidden_states = self.input_layernorm(hidden_states) - - attention_output = self.attention( - hidden_states, - attention_mask=attention_mask, - use_cache=use_cache, - spec_decoding_params=spec_decoding_params, - kv_cache_params=kv_cache_params, - attention_params=attention_params) - if use_cache: - attention_output, presents = attention_output - - hidden_states = residual + attention_output - - residual = hidden_states - - hidden_states = self.post_layernorm(hidden_states) - hidden_states = self.mlp(hidden_states) - hidden_states = residual + hidden_states - if use_cache: - return (hidden_states, presents) - return hidden_states - - -class DeepseekModel(Module): - - def __init__(self, config: PretrainedConfig) -> None: - super().__init__() - init_all_reduce_helper() # enable use_customer_all_reduce - - self.mapping = config.mapping - if self.mapping.is_first_pp_rank(): - self.vocab_embedding = Embedding(config.vocab_size, - config.hidden_size, - dtype=config.dtype) - - self.layers = DecoderLayerList(DeepseekDecoderLayer, config) - - if self.mapping.is_last_pp_rank(): - self.ln_f = RmsNorm(normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - dtype=config.dtype) - - def forward(self, - input_ids, - position_ids=None, - use_cache=False, - attention_mask=None, - spec_decoding_params=None, - kv_cache_params=None, - attention_params=None, - hidden_states=None, - prompt_embedding_table: Optional[Tensor] = None, - prompt_tasks: Optional[Tensor] = None, - prompt_vocab_size: Optional[Tensor] = None): - - ptuning_args = [ - prompt_embedding_table, prompt_tasks, prompt_vocab_size - ] if prompt_embedding_table is not None else [] - - if self.mapping.is_first_pp_rank(): - hidden_states = self.vocab_embedding(input_ids, *ptuning_args) - else: - hidden_states = recv(hidden_states, self.mapping.prev_pp_rank()) - - hidden_states = self.layers.forward( - hidden_states, - use_cache=use_cache, - attention_mask=attention_mask, - kv_cache_params=kv_cache_params, - attention_params=attention_params, - spec_decoding_params=spec_decoding_params) - - if use_cache: - hidden_states, presents = hidden_states - - if self.mapping.is_last_pp_rank(): - hidden_states = self.ln_f(hidden_states) - else: - hidden_states = send(hidden_states, self.mapping.next_pp_rank()) - - if use_cache: - return (hidden_states, tuple(presents)) - return hidden_states - - -class DeepseekForCausalLM(DecoderModelForCausalLM): - config_class = DeepSeekV1Config - - def __init__(self, config: PretrainedConfig): - transformer = DeepseekModel(config) - vocab_size_padded = pad_vocab_size(config.vocab_size, - config.mapping.tp_size) - if config.mapping.is_last_pp_rank(): - lm_head = ColumnLinear(config.hidden_size, - vocab_size_padded, - bias=False, - dtype=config.dtype, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size, - gather_output=True) - else: - lm_head = None - self.mapping = config.mapping - super().__init__(config, transformer, lm_head) - - @classmethod - def from_hugging_face(cls, - model_dir, - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - override_fields={}, - **kwargs): - if mapping is None: - mapping = Mapping() - - pretrained_config = DeepSeekV1Config.from_hugging_face(model_dir, - dtype=dtype, - mapping=mapping, - **kwargs) - deepseek = cls.from_config(pretrained_config) - if os.environ.get("TRTLLM_DISABLE_UNIFIED_CONVERTER") is None: - - custom_dict = {} - - rank_experts = mapping.ep_experts(pretrained_config.moe.num_experts) - for index, module in enumerate(deepseek.transformer.layers): - if index > 0: - - module.mlp.shared_expert.fc.tllm_to_externel_key_dict = { - "fc": ["up_proj", "gate_proj"], - "shared_expert": "shared_experts" - } - module.mlp.shared_expert.proj.tllm_to_externel_key_dict = { - "shared_expert": "shared_experts" - } - module.mlp.fc.tllm_to_externel_key_dict = { - "fc": [ - f"experts.{expert}.up_proj" - for expert in rank_experts - ] + [ - f"experts.{expert}.gate_proj" - for expert in rank_experts - ] - } - module.mlp.proj.tllm_to_externel_key_dict = { - "proj": [ - f"experts.{expert}.down_proj" - for expert in rank_experts - ] - } - module.mlp.router.tllm_to_externel_key_dict = { - "mlp": "mlp", - "router": "gate" - } - - loader = ModelWeightsLoader(model_dir, custom_dict) - loader.generate_tllm_weights(deepseek) - return deepseek - else: - - hf_model = load_hf_deepseek(model_dir) - weights = convert_deepseek( - hf_model, - pretrained_config, - mapping=pretrained_config.mapping, - dtype=pretrained_config.dtype, - use_parallel_embedding=pretrained_config.use_parallel_embedding, - sharding_dim=pretrained_config.embedding_sharding_dim) - deepseek.load(weights) - return deepseek diff --git a/tensorrt_llm/models/deepseek_v2/__init__.py b/tensorrt_llm/models/deepseek_v2/__init__.py deleted file mode 100644 index 71bf6d298d42..000000000000 --- a/tensorrt_llm/models/deepseek_v2/__init__.py +++ /dev/null @@ -1,14 +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. diff --git a/tensorrt_llm/models/deepseek_v2/config.py b/tensorrt_llm/models/deepseek_v2/config.py deleted file mode 100644 index c110df0d53f1..000000000000 --- a/tensorrt_llm/models/deepseek_v2/config.py +++ /dev/null @@ -1,148 +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. - -from typing import Optional, Union - -from transformers import AutoConfig - -from ..._utils import get_hf_rope_theta -from ...layers import MoeConfig -from ...mapping import Mapping -from ..modeling_utils import PretrainedConfig, QuantConfig - - -class DeepSeekV2Config(PretrainedConfig): - - def __init__(self, - *, - rotary_base: float = 10000.0, - rotary_scaling: Optional[dict] = None, - moe: Optional[Union[MoeConfig, dict]] = None, - **kwargs): - self.rotary_base = rotary_base - self.rotary_scaling = rotary_scaling - if isinstance(moe, dict): - moe = MoeConfig.from_dict(moe) - assert isinstance(moe, - MoeConfig), "moe must be a MoeConfig or a dictionary" - self.moe = moe.validate() - super().__init__(**kwargs) - - def to_dict(self): - output = super().to_dict() - # Serialize the fields added in DeepSeekV2Config - output['rotary_base'] = self.rotary_base - output['rotary_scaling'] = self.rotary_scaling - output['moe'] = self.moe.to_dict() - return output - - @classmethod - def from_hugging_face( - cls, - hf_config_or_dir: Union[str, 'transformers.PretrainedConfig'], - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - import transformers - trust_remote_code = kwargs.pop('trust_remote_code', True) - - if isinstance(hf_config_or_dir, transformers.PretrainedConfig): - hf_config = hf_config_or_dir - else: - hf_config_dir = str(hf_config_or_dir) - - hf_config = AutoConfig.from_pretrained( - hf_config_dir, trust_remote_code=trust_remote_code) - - moe_routed_scaling_factor = hf_config.routed_scaling_factor - moe_top_k = hf_config.num_experts_per_tok - assert moe_routed_scaling_factor > 0, 'routed_scaling_factor should be greater than 0' - if hf_config.topk_method == 'group_limited_greedy': - if moe_top_k > 1 and hf_config.norm_topk_prob: - moe_renorm_mode = MoeConfig.ExpertScaleNormalizationMode.DEVICE_LIMITED_RENORM - else: - moe_renorm_mode = MoeConfig.ExpertScaleNormalizationMode.DEVICE_LIMITED - elif hf_config.topk_method == 'greedy': - assert moe_routed_scaling_factor == 1.0, 'The combination of topk_method == greedy and routed_scaling_factor != 1.0 is not supported' - if moe_top_k > 1 and hf_config.norm_topk_prob: - moe_renorm_mode = MoeConfig.ExpertScaleNormalizationMode.RENORMALIZE - else: - moe_renorm_mode = MoeConfig.ExpertScaleNormalizationMode.NONE - else: - raise AssertionError( - f'Unsupported topk_method in hf_config: {hf_config.topk_method}' - ) - - rotary_scaling = None - if hf_config.rope_scaling is not None: - rotary_scaling = { - 'beta_fast': - hf_config.rope_scaling['beta_fast'], - 'beta_slow': - hf_config.rope_scaling['beta_slow'], - 'factor': - hf_config.rope_scaling['factor'], - 'mscale': - hf_config.rope_scaling['mscale'], - 'mscale_all_dim': - hf_config.rope_scaling['mscale_all_dim'], - 'original_max_position_embeddings': - hf_config.rope_scaling['original_max_position_embeddings'], - 'type': - hf_config.rope_scaling['type'] - } - - moe_config = MoeConfig( - num_experts=hf_config.n_routed_experts, - shared_expert_intermediate_size=hf_config.n_shared_experts * - hf_config.moe_intermediate_size, - top_k=hf_config.num_experts_per_tok, - normalization_mode=moe_renorm_mode, - device_limited_n_group=hf_config.n_group, - device_limited_topk_group=hf_config.topk_group, - device_limited_routed_scaling_factor=hf_config.routed_scaling_factor - ) - - moe_config.validate() - - return cls(architecture=hf_config.architectures[0], - dtype=dtype, - num_hidden_layers=hf_config.num_hidden_layers, - num_attention_heads=hf_config.num_attention_heads, - hidden_size=hf_config.hidden_size, - intermediate_size=hf_config.intermediate_size, - num_key_value_heads=hf_config.num_key_value_heads, - vocab_size=hf_config.vocab_size, - position_embedding_type='rope_gpt_neox', - max_position_embeddings=hf_config.max_position_embeddings, - hidden_act='swiglu', - norm_epsilon=hf_config.rms_norm_eps, - rotary_base=get_hf_rope_theta(hf_config, 10000.0), - rotary_scaling=rotary_scaling, - moe_inter_size=hf_config.moe_intermediate_size, - moe=moe_config, - mapping=mapping, - quantization=quant_config, - kv_lora_rank=hf_config.kv_lora_rank, - q_lora_rank=hf_config.q_lora_rank, - qk_nope_head_dim=hf_config.qk_nope_head_dim, - qk_rope_head_dim=hf_config.qk_rope_head_dim, - v_head_dim=hf_config.v_head_dim, - topk_method=hf_config.topk_method, - first_k_dense_replace=hf_config.first_k_dense_replace, - moe_layer_freq=hf_config.moe_layer_freq, - scoring_func=hf_config.scoring_func, - **kwargs) diff --git a/tensorrt_llm/models/deepseek_v2/convert.py b/tensorrt_llm/models/deepseek_v2/convert.py deleted file mode 100755 index 5a23130fc528..000000000000 --- a/tensorrt_llm/models/deepseek_v2/convert.py +++ /dev/null @@ -1,929 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import time - -import torch -from transformers import AutoConfig, AutoModelForCausalLM - -from tensorrt_llm.layers import MoeConfig - -from ..._utils import (get_hf_rope_theta, pad_vocab_size, release_gc, - str_dtype_to_torch) -from ...logger import logger -from ...mapping import Mapping -from ..convert_utils import get_tllm_linear_weight - -# `Override num_hidden_layers` used for reduce number of hidden layers in DeepseekV2ForCausalLM for debug purpose -OVERRIDE_HIDDEN_LAYERS = None # 2 - - -# Convert config parameters to dict TODO: remove this function and change CI test -def create_trt_config_from_hf(model_dir, - dtype, - mapping: Mapping, - override_fields: dict = {}): - config = {} - assert isinstance(model_dir, str) - hf_config = AutoConfig.from_pretrained(model_dir, trust_remote_code=True) - # Override num_hidden_layers - if OVERRIDE_HIDDEN_LAYERS is not None: - hf_config.num_hidden_layers = OVERRIDE_HIDDEN_LAYERS - print( - f'Override hidden layers to {hf_config.num_hidden_layers} for DeepseekV2ForCausalLM' - ) - dtype = dtype - n_layer = hf_config.num_hidden_layers - n_head = hf_config.num_attention_heads - n_embd = hf_config.hidden_size - inter_size = hf_config.intermediate_size - n_kv_head = hf_config.num_key_value_heads - vocab_size = hf_config.vocab_size - n_positions = hf_config.max_position_embeddings - hidden_act = 'swiglu' # TRT-LLM request make gated activation explicit for MOE implementation - rotary_base = get_hf_rope_theta(hf_config, 10000.0) - rms_norm_eps = hf_config.rms_norm_eps - rotary_scaling_beta_fast = hf_config.rope_scaling['beta_fast'] - rotary_scaling_beta_slow = hf_config.rope_scaling['beta_slow'] - rotary_scaling_factor = hf_config.rope_scaling['factor'] - rotary_scaling_mscale = hf_config.rope_scaling['mscale'] - rotary_scaling_mscale_all_dim = hf_config.rope_scaling['mscale_all_dim'] - rotary_scaling_original_max_position_embeddings = hf_config.rope_scaling[ - 'original_max_position_embeddings'] - rotary_scaling_type = 'yarn' - kv_lora_rank = hf_config.kv_lora_rank - q_lora_rank = hf_config.q_lora_rank - qk_nope_head_dim = hf_config.qk_nope_head_dim - qk_rope_head_dim = hf_config.qk_rope_head_dim - v_head_dim = hf_config.v_head_dim - moe_num_experts = hf_config.n_routed_experts - moe_inter_size = hf_config.moe_intermediate_size - moe_num_shared_experts = hf_config.n_shared_experts - moe_top_k = hf_config.num_experts_per_tok - moe_n_group = hf_config.n_group - moe_topk_group = hf_config.topk_group - moe_routed_scaling_factor = hf_config.routed_scaling_factor - assert moe_routed_scaling_factor > 0, 'routed_scaling_factor should be greater than 0' - if hf_config.topk_method == 'group_limited_greedy': - if moe_top_k > 1 and hf_config.norm_topk_prob: - moe_renorm_mode = MoeConfig.ExpertScaleNormalizationMode.DEVICE_LIMITED_RENORM - else: - moe_renorm_mode = MoeConfig.ExpertScaleNormalizationMode.DEVICE_LIMITED - elif hf_config.topk_method == 'greedy': - assert moe_routed_scaling_factor == 1.0, 'The combination of topk_method == greedy and routed_scaling_factor != 1.0 is not supported' - if moe_top_k > 1 and hf_config.norm_topk_prob: - moe_renorm_mode = MoeConfig.ExpertScaleNormalizationMode.RENORMALIZE - else: - moe_renorm_mode = MoeConfig.ExpertScaleNormalizationMode.NONE - else: - raise AssertionError( - 'Unsupported topk_method in hf_config: {hf_config.topk_method}') - - config = { - 'architecture': 'DeepseekV2ForCausalLM', - 'dtype': dtype, - 'logits_type': 'float32', - 'num_hidden_layers': n_layer, - 'num_attention_heads': n_head, - 'hidden_size': n_embd, - 'intermediate_size': inter_size, - 'num_key_value_heads': n_kv_head, - 'vocab_size': vocab_size, - 'position_embedding_type': 'rope_gpt_neox', - 'max_position_embeddings': n_positions, - 'hidden_act': hidden_act, - 'rotary_base': rotary_base, - 'norm_epsilon': rms_norm_eps, - 'rotary_scaling': { - 'beta_fast': rotary_scaling_beta_fast, - 'beta_slow': rotary_scaling_beta_slow, - 'factor': rotary_scaling_factor, - 'mscale': rotary_scaling_mscale, - 'mscale_all_dim': rotary_scaling_mscale_all_dim, - 'original_max_position_embeddings': - rotary_scaling_original_max_position_embeddings, - 'type': rotary_scaling_type, - }, - 'mapping': { - 'world_size': mapping.tp_size * mapping.pp_size, - 'tp_size': mapping.tp_size, - 'pp_size': mapping.pp_size, - 'moe_tp_size': mapping.moe_tp_size, - 'moe_ep_size': mapping.moe_ep_size, - }, - 'kv_lora_rank': kv_lora_rank, - 'q_lora_rank': q_lora_rank, - 'qk_nope_head_dim': qk_nope_head_dim, - 'qk_rope_head_dim': qk_rope_head_dim, - 'v_head_dim': v_head_dim, - 'moe_num_experts': moe_num_experts, - 'moe_inter_size': moe_inter_size, - 'moe_num_shared_experts': moe_num_shared_experts, - 'moe_top_k': moe_top_k, - 'moe_renorm_mode': moe_renorm_mode, - 'moe_n_group': moe_n_group, - 'moe_topk_group': moe_topk_group, - 'moe_routed_scaling_factor': moe_routed_scaling_factor, - } - - config.update(override_fields) - - moe_config = MoeConfig( - num_experts=config['moe_num_experts'], - shared_expert_intermediate_size=config['moe_num_shared_experts'] * - config['moe_inter_size'], - top_k=config['moe_top_k'], - normalization_mode=config['moe_renorm_mode'], - device_limited_n_group=config['moe_n_group'], - device_limited_topk_group=config['moe_topk_group'], - device_limited_routed_scaling_factor=config['moe_routed_scaling_factor'] - ) - moe_config.validate() - - return config - - -# Get HF model -def load_hf_deepseek(model_dir, load_model_on_cpu=False): - hf_config = AutoConfig.from_pretrained(model_dir, trust_remote_code=True) - if OVERRIDE_HIDDEN_LAYERS is not None: - hf_config.num_hidden_layers = OVERRIDE_HIDDEN_LAYERS - print( - f'Override hidden layers to {hf_config.num_hidden_layers} for DeepseekV2ForCausalLM' - ) - - if load_model_on_cpu: - # Skip setting max_memory when loading on CPU, you might have OOM. - model = AutoModelForCausalLM.from_pretrained(model_dir, - config=hf_config, - device_map='cpu', - dtype='auto', - trust_remote_code=True) - else: - # Deepseek-v2 236B parameters with FP16 dtype need at least 472G GPU memory - # (official suggest at least 8x80G GPUs, see https://huggingface.co/deepseek-ai/DeepSeek-V2) - - max_memory = None - device_map = 'auto' - - gpu_memory = torch.cuda.get_device_properties(0).total_memory - gpu_count = torch.cuda.device_count() - - if gpu_memory < 90_000_000_000 and gpu_count == 8: - # WAR OOM loading on 8*80G GPUs - max_memory = {i: "76GB" for i in range(8)} - device_map = 'sequential' - elif gpu_memory < 180_000_000_000 and gpu_count == 4: - # WAR OOM loading on 4*141G GPUs - max_memory = {i: "128GB" for i in range(4)} - device_map = 'sequential' - elif gpu_memory < 180_000_000_000 and gpu_count == 8: - # WAR OOM loading on 8*141G GPUs - max_memory = {i: "128GB" for i in range(8)} - device_map = 'sequential' - - model = AutoModelForCausalLM.from_pretrained(model_dir, - config=hf_config, - device_map=device_map, - max_memory=max_memory, - dtype='auto', - trust_remote_code=True) - - return model - - -# Prepare weights for TP -def split(v, tp_size, idx, dim=0): - if tp_size == 1: - return v - if len(v.shape) == 1: - return torch.chunk(v, tp_size)[idx].contiguous() - else: - return torch.chunk(v, tp_size, dim=dim)[idx].contiguous() - - -def split_matrix_tp(v, tensor_parallel, rank, dim): - return split(v, tensor_parallel, rank, dim=dim) - - -def get_weight(config, prefix, dtype, postfix='.weight'): - if config[prefix + postfix].dtype != dtype: - config[prefix + postfix].data = config[prefix + postfix].to(dtype) - return config[prefix + postfix].detach().cpu() - - -def get_param_weight(weight, prefix): - results = {} - results[prefix] = weight - - return results - - -def convert_deepseekv2(hf_model, - config, - mapping, - dtype='float32', - use_parallel_embedding=False, - sharding_dim=0): - - weights = {} - tik = time.time() - model_params = dict(hf_model.named_parameters()) - dtype = getattr(torch, dtype) - moe_config = config.moe - - layers_range = mapping.pp_layers(config.num_hidden_layers) - - def convert_layer(l): - prefix = f'model.layers.{l}.' - trtllm_prefix = f'transformer.layers.{l - layers_range[0]}.' - # Fuse matrices for compression - # Split matrices for decompression - q_lora_rank = config.q_lora_rank - kv_lora_rank = config.kv_lora_rank - num_heads = config.num_attention_heads - qk_nope_head_dim = config.qk_nope_head_dim - qk_rope_head_dim = config.qk_rope_head_dim - v_head_dim = config.v_head_dim - hidden_size = config.hidden_size - - if q_lora_rank is not None: - q_a_proj_weight = get_weight(model_params, - prefix + 'self_attn.q_a_proj', dtype) - # Layer normalization - q_a_layernorm_weight = get_weight( - model_params, - prefix + 'self_attn.q_a_layernorm', - dtype, - ) - - kv_a_proj_with_mqa_weight = get_weight( - model_params, prefix + 'self_attn.kv_a_proj_with_mqa', dtype) - - kv_a_layernorm_weight = get_weight( - model_params, - prefix + 'self_attn.kv_a_layernorm', - dtype, - ) - - if q_lora_rank is not None: - fused_a_weight = torch.cat( - [q_a_proj_weight, kv_a_proj_with_mqa_weight], - dim=0, - ) - - q_b_proj_weight = get_weight( - model_params, prefix + 'self_attn.q_b_proj', dtype).unflatten( - 0, - [ - num_heads, - qk_nope_head_dim + qk_rope_head_dim, - ], - ) - else: - fused_a_weight = kv_a_proj_with_mqa_weight - - q_b_proj_weight = get_weight( - model_params, prefix + 'self_attn.q_proj', dtype).unflatten( - 0, - [ - num_heads, - qk_nope_head_dim + qk_rope_head_dim, - ], - ) - - kv_b_proj_weight = get_weight(model_params, - prefix + 'self_attn.kv_b_proj', - dtype).unflatten( - 0, - [ - num_heads, - qk_nope_head_dim + v_head_dim, - ], - ) - - o_proj_weight = get_weight(model_params, prefix + 'self_attn.o_proj', - dtype) - - q_b_proj_weight = split_matrix_tp( - q_b_proj_weight, - mapping.tp_size, - mapping.tp_rank, - dim=0, - ) - kv_b_proj_weight = split_matrix_tp( - kv_b_proj_weight, - mapping.tp_size, - mapping.tp_rank, - dim=0, - ) - o_proj_weight = split_matrix_tp( - o_proj_weight, - mapping.tp_size, - mapping.tp_rank, - dim=1, - ) - - q_nope_weight, q_pe_weight = q_b_proj_weight.split( - [qk_nope_head_dim, qk_rope_head_dim], - dim=1, - ) - k_nope_weight, v_weight = kv_b_proj_weight.split( - [qk_nope_head_dim, v_head_dim], - dim=1, - ) - - if q_lora_rank is None: - q_b_proj_weight = q_b_proj_weight.reshape( - num_heads * (qk_nope_head_dim + qk_rope_head_dim) // - mapping.tp_size, hidden_size) - else: - q_b_proj_weight = q_b_proj_weight.reshape( - num_heads * (qk_nope_head_dim + qk_rope_head_dim) // - mapping.tp_size, q_lora_rank) - - kv_b_proj_weight = torch.concat([ - k_nope_weight.reshape( - num_heads * qk_nope_head_dim // mapping.tp_size, kv_lora_rank), - v_weight.reshape(num_heads * v_head_dim // mapping.tp_size, - kv_lora_rank) - ], - dim=0) - - # Fuse matrices for decompression - fused_q_nope_weight = torch.einsum( - 'hdq,hdk->hkq', - q_nope_weight, - k_nope_weight, - ) - fused_q_weight = torch.cat( - [fused_q_nope_weight, q_pe_weight], - dim=1, - ).flatten(start_dim=0, end_dim=1) - - weights.update( - get_tllm_linear_weight(fused_a_weight, - trtllm_prefix + 'attention.fused_a.')) - weights.update( - get_tllm_linear_weight(kv_a_layernorm_weight, - trtllm_prefix + 'attention.kv_a_layernorm.')) - weights.update( - get_param_weight(fused_q_weight, - trtllm_prefix + 'attention.fused_q_proj')) - weights.update( - get_param_weight(q_b_proj_weight, - trtllm_prefix + 'attention.q_b_proj')) - weights.update( - get_param_weight(kv_b_proj_weight, - trtllm_prefix + 'attention.kv_b_proj')) - weights.update( - get_tllm_linear_weight(o_proj_weight, - trtllm_prefix + 'attention.dense.')) - - if q_lora_rank is not None: - weights.update( - get_tllm_linear_weight( - q_a_layernorm_weight, - trtllm_prefix + 'attention.q_a_layernorm.')) - - if moe_config.has_moe() and l > 0: - rank_experts = list(range(moe_config.num_experts)) - if mapping.has_moe_ep(): - rank_experts = mapping.ep_experts(moe_config.num_experts) - for suffix in ["gate_proj", "down_proj", "up_proj"]: - model_params[f'model.layers.{l}.mlp.experts.{suffix}.weight'] = \ - torch.stack([model_params[f'model.layers.{l}.mlp.experts.{expert}.{suffix}.weight'].detach().cpu() - for expert in rank_experts]) - - gate_proj = model_params[ - f'model.layers.{l}.mlp.experts.gate_proj.weight'] - down_proj = model_params[ - f'model.layers.{l}.mlp.experts.down_proj.weight'] - up_proj = model_params[ - f'model.layers.{l}.mlp.experts.up_proj.weight'] - if mapping.has_moe_tp(): - gate_proj = split(gate_proj, - mapping.tp_size, - mapping.tp_rank, - dim=1) - down_proj = split(down_proj, - mapping.tp_size, - mapping.tp_rank, - dim=2) - up_proj = split(up_proj, - mapping.tp_size, - mapping.tp_rank, - dim=1) - - model_params[ - f'model.layers.{l}.mlp.experts.up_gate_proj.weight'] = torch.concat( - [up_proj, gate_proj], dim=-2) - model_params[ - f'model.layers.{l}.mlp.experts.down_proj.weight'] = down_proj - - # mlp.experts.down_proj.weight - moe_experts_down_proj_weights = get_weight( - model_params, prefix + 'mlp.experts.down_proj', dtype) - weights.update( - get_tllm_linear_weight(moe_experts_down_proj_weights, - trtllm_prefix + 'mlp.proj.')) - # mlp.experts.up_gate.weight - moe_experts_up_gate_proj_weights = get_weight( - model_params, prefix + 'mlp.experts.up_gate_proj', dtype) - weights.update( - get_tllm_linear_weight(moe_experts_up_gate_proj_weights, - trtllm_prefix + 'mlp.fc.')) - # MOE hardcoded routing_input into trt.float32, please refer to moe.py line 397 - moe_experts_gate_weights = get_weight(model_params, - prefix + 'mlp.gate', - torch.float32) - weights.update( - get_tllm_linear_weight(moe_experts_gate_weights, - trtllm_prefix + 'mlp.router.')) - - if moe_config.shared_expert_intermediate_size > 0: - shared_moe_up_proj_weights = get_weight( - model_params, prefix + 'mlp.shared_experts.up_proj', dtype) - shared_moe_up_proj_weights = split_matrix_tp( - shared_moe_up_proj_weights, - mapping.tp_size, - mapping.tp_rank, - dim=0) - shared_moe_down_proj_weights = get_weight( - model_params, prefix + 'mlp.shared_experts.down_proj', - dtype) - shared_moe_down_proj_weights = split_matrix_tp( - shared_moe_down_proj_weights, - mapping.tp_size, - mapping.tp_rank, - dim=1) - shared_moe_gate_proj_weights = get_weight( - model_params, prefix + 'mlp.shared_experts.gate_proj', - dtype) - shared_moe_gate_proj_weights = split_matrix_tp( - shared_moe_gate_proj_weights, - mapping.tp_size, - mapping.tp_rank, - dim=0) - shared_moe_gate_up_proj_weights = torch.concat( - [shared_moe_up_proj_weights, shared_moe_gate_proj_weights], - dim=-2) - - # mlp.shared_experts.gate_up_proj.weight - weights.update( - get_tllm_linear_weight( - shared_moe_gate_up_proj_weights, - trtllm_prefix + 'mlp.shared_expert.fc.')) - - # mlp.shared_experts.down_proj.weight - weights.update( - get_tllm_linear_weight( - shared_moe_down_proj_weights, - trtllm_prefix + 'mlp.shared_expert.proj.')) - - else: - # Current MLP layer is only one, if it goes large consider to do fuse - mlp_gate_weight = get_weight(model_params, prefix + 'mlp.up_proj', - dtype) - split_gate = split_matrix_tp(mlp_gate_weight, - mapping.tp_size, - mapping.tp_rank, - dim=0) - weights.update( - get_tllm_linear_weight(split_gate, trtllm_prefix + 'mlp.gate.')) - - mlp_fc_weight = get_weight(model_params, prefix + 'mlp.gate_proj', - dtype) - split_fc = split_matrix_tp(mlp_fc_weight, - mapping.tp_size, - mapping.tp_rank, - dim=0) - weights.update( - get_tllm_linear_weight(split_fc, trtllm_prefix + 'mlp.fc.')) - - mlp_proj_weight = get_weight(model_params, prefix + 'mlp.down_proj', - dtype) - split_proj = split_matrix_tp(mlp_proj_weight, - mapping.tp_size, - mapping.tp_rank, - dim=1) - weights.update( - get_tllm_linear_weight(split_proj, trtllm_prefix + 'mlp.proj.')) - - # Layer norms do not use tensor parallelism - input_ln_weight = get_weight(model_params, prefix + 'input_layernorm', - dtype) - weights[trtllm_prefix + 'input_layernorm.weight'] = input_ln_weight - post_ln_weight = get_weight(model_params, - prefix + 'post_attention_layernorm', dtype) - weights[trtllm_prefix + 'post_layernorm.weight'] = post_ln_weight - - for l in layers_range: - convert_layer(l) - release_gc() - - v = get_weight(model_params, 'model.embed_tokens', dtype) - if hf_model.config.tie_word_embeddings: - # lm_head.weight has the same weights as embedding - if mapping.is_last_pp_rank(): - if config.vocab_size % mapping.tp_size != 0: - # padding - vocab_size_padded = pad_vocab_size(config.vocab_size, - mapping.tp_size) - pad_width = vocab_size_padded - config.vocab_size - v = torch.nn.functional.pad(v, (0, 0, 0, pad_width), 'constant', - 0) - weights['lm_head.weight'] = split(v, mapping.tp_size, - mapping.tp_rank) - if use_parallel_embedding: - v = split_matrix_tp(v, - mapping.tp_size, - mapping.tp_rank, - dim=config.embedding_sharding_dim) - if mapping.is_first_pp_rank(): - weights['transformer.vocab_embedding.weight'] = v - lm_head_weights = get_weight(model_params, 'lm_head', dtype) - - if mapping.is_last_pp_rank(): - if config.vocab_size % mapping.tp_size != 0: - # padding - vocab_size_padded = pad_vocab_size(config.vocab_size, - mapping.tp_size) - pad_width = vocab_size_padded - config.vocab_size - lm_head_weights = torch.nn.functional.pad(lm_head_weights, - (0, 0, 0, pad_width), - 'constant', - value=0) - weights['lm_head.weight'] = split_matrix_tp(lm_head_weights, - mapping.tp_size, - mapping.tp_rank, - dim=0) - ln_f_w = get_weight(model_params, 'model.norm', dtype) - weights['transformer.ln_f.weight'] = ln_f_w - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Weights loaded. Total time: {t}') - #print(set(weights.keys())) - return weights - - -def load_weights_from_hf_safetensors(model_dir, - config, - mapping, - use_parallel_embedding=False, - sharding_dim=0): - logger.info('Loading weights from Huggingface safetensors...') - weights = {} - tik = time.time() - import json - import os - - import safetensors - - model_dir = model_dir if model_dir.endswith("/") else model_dir + "/" - safetensors_map = {} - has_safetensor_index_json = True - try: - with open(model_dir + "model.safetensors.index.json", 'r') as fr: - sharding_map = json.load(fr) - # safetensors_map structure: - # key: e.g., model.layers.0.self_attn.q_a_proj.weight - # value: e.g., model-00001-of-000055.safetensors - # int(value[6:11]) -> safetensors_idx - for k, v in sharding_map['weight_map'].items(): - safetensors_map[k] = int(v[6:11]) - 1 - except FileNotFoundError: - has_safetensor_index_json = False - # shard_files purpose to verify all .safetensors is aligned with .index.json - shard_files = [] - for name in os.listdir(model_dir): - if name.endswith(".safetensors"): - if has_safetensor_index_json and name not in sharding_map[ - 'weight_map'].values(): - continue - shard_files.append(name) - shard_files.sort() - # Create all .safetensors memory address - safetensors_ptrs = [ - safetensors.safe_open(model_dir + shard_file, - framework="pt", - device="cpu") for shard_file in shard_files - ] - - moe_config = MoeConfig( - num_experts=config.moe.num_experts, - shared_expert_intermediate_size=config.moe. - shared_expert_intermediate_size, - top_k=config.moe.top_k, - normalization_mode=config.moe.normalization_mode, - device_limited_n_group=config.moe.device_limited_n_group, - device_limited_topk_group=config.moe.device_limited_topk_group, - device_limited_routed_scaling_factor=config.moe. - device_limited_routed_scaling_factor) - - torch_dtype = str_dtype_to_torch(config.dtype) - - def load_weights(key, dtype): - assert key in safetensors_map, f"'{key}' not found in safetensors_map" - ptr_idx = safetensors_map[key] - - assert key in safetensors_ptrs[ptr_idx].keys( - ), f"'{key}' not found in safetensors file {ptr_idx}" - tensor_slice = safetensors_ptrs[ptr_idx].get_slice(key) - tensor_slice.get_shape() - res = tensor_slice[:] - if res.dtype != dtype: - res = res.to(dtype) - return res.contiguous().detach().cpu() - - layers_range = mapping.pp_layers(config.num_hidden_layers) - - def convert_layer(l): - prefix = f'model.layers.{l}.' - trtllm_prefix = f'transformer.layers.{l - layers_range[0]}.' - # Fuse matrices for compression - # Split matrices for decompression - q_lora_rank = config.q_lora_rank - kv_lora_rank = config.kv_lora_rank - num_heads = config.num_attention_heads - qk_nope_head_dim = config.qk_nope_head_dim - qk_rope_head_dim = config.qk_rope_head_dim - v_head_dim = config.v_head_dim - hidden_size = config.hidden_size - - # MLA: - # q_lora_rank used for different deepseek-v2 and deepseek-v2-lite - if q_lora_rank is not None: - q_a_proj_weight = load_weights(prefix + 'self_attn.q_a_proj.weight', - torch_dtype) - q_a_layernorm_weight = load_weights( - prefix + 'self_attn.q_a_layernorm.weight', torch_dtype) - - kv_a_proj_with_mqa_weight = load_weights( - prefix + 'self_attn.kv_a_proj_with_mqa.weight', torch_dtype) - kv_a_layernorm_weight = load_weights( - prefix + 'self_attn.kv_a_layernorm.weight', torch_dtype) - - if q_lora_rank is not None: - fused_a_weight = torch.cat( - [q_a_proj_weight, kv_a_proj_with_mqa_weight], dim=0) - q_b_proj_weight = load_weights( - prefix + 'self_attn.q_b_proj.weight', torch_dtype).unflatten( - 0, [num_heads, qk_nope_head_dim + qk_rope_head_dim]) - else: - fused_a_weight = kv_a_proj_with_mqa_weight - q_b_proj_weight = load_weights( - prefix + 'self_attn.q_proj.weight', torch_dtype).unflatten( - 0, [num_heads, qk_nope_head_dim + qk_rope_head_dim]) - - kv_b_proj_weight = load_weights( - prefix + 'self_attn.kv_b_proj.weight', - torch_dtype).unflatten(0, - [num_heads, qk_nope_head_dim + v_head_dim]) - o_proj_weight = load_weights(prefix + 'self_attn.o_proj.weight', - torch_dtype) - - q_b_proj_weight = split_matrix_tp(q_b_proj_weight, - mapping.tp_size, - mapping.tp_rank, - dim=0) - kv_b_proj_weight = split_matrix_tp(kv_b_proj_weight, - mapping.tp_size, - mapping.tp_rank, - dim=0) - o_proj_weight = split_matrix_tp(o_proj_weight, - mapping.tp_size, - mapping.tp_rank, - dim=1) - - q_nope_weight, q_pe_weight = q_b_proj_weight.split( - [qk_nope_head_dim, qk_rope_head_dim], dim=1) - k_nope_weight, v_weight = kv_b_proj_weight.split( - [qk_nope_head_dim, v_head_dim], dim=1) - - if q_lora_rank is None: - q_b_proj_weight = q_b_proj_weight.reshape( - num_heads * (qk_nope_head_dim + qk_rope_head_dim) // - mapping.tp_size, hidden_size) - else: - q_b_proj_weight = q_b_proj_weight.reshape( - num_heads * (qk_nope_head_dim + qk_rope_head_dim) // - mapping.tp_size, q_lora_rank) - - kv_b_proj_weight = torch.concat([ - k_nope_weight.reshape( - num_heads * qk_nope_head_dim // mapping.tp_size, kv_lora_rank), - v_weight.reshape(num_heads * v_head_dim // mapping.tp_size, - kv_lora_rank) - ], - dim=0) - - # Fuse matrices for decompression - fused_q_nope_weight = torch.einsum('hdq,hdk->hkq', q_nope_weight, - k_nope_weight) - fused_q_weight = torch.cat([fused_q_nope_weight, q_pe_weight], - dim=1).flatten(start_dim=0, end_dim=1) - - weights.update( - get_tllm_linear_weight(fused_a_weight, - trtllm_prefix + 'attention.fused_a.')) - weights.update( - get_tllm_linear_weight(kv_a_layernorm_weight, - trtllm_prefix + 'attention.kv_a_layernorm.')) - weights.update( - get_param_weight(fused_q_weight, - trtllm_prefix + 'attention.fused_q_proj')) - weights.update( - get_param_weight(q_b_proj_weight, - trtllm_prefix + 'attention.q_b_proj')) - weights.update( - get_param_weight(kv_b_proj_weight, - trtllm_prefix + 'attention.kv_b_proj')) - weights.update( - get_tllm_linear_weight(o_proj_weight, - trtllm_prefix + 'attention.dense.')) - - if q_lora_rank is not None: - weights.update( - get_tllm_linear_weight( - q_a_layernorm_weight, - trtllm_prefix + 'attention.q_a_layernorm.')) - - # MOE: - if moe_config.has_moe() and l > 0: - rank_experts = list(range(moe_config.num_experts)) - if mapping.has_moe_ep(): - rank_experts = mapping.ep_experts(moe_config.num_experts) - expert_weight_dict = {} - for suffix in ["gate_proj", "down_proj", "up_proj"]: - expert_weight_dict[f'model.layers.{l}.mlp.experts.{suffix}.weight'] = \ - torch.stack([load_weights(prefix + f'mlp.experts.{expert}.{suffix}.weight', torch_dtype) for expert in rank_experts]) - - gate_proj = expert_weight_dict[ - f'model.layers.{l}.mlp.experts.gate_proj.weight'] - down_proj = expert_weight_dict[ - f'model.layers.{l}.mlp.experts.down_proj.weight'] - up_proj = expert_weight_dict[ - f'model.layers.{l}.mlp.experts.up_proj.weight'] - - if mapping.has_moe_tp(): - gate_proj = split(gate_proj, - mapping.tp_size, - mapping.tp_rank, - dim=1) - down_proj = split(down_proj, - mapping.tp_size, - mapping.tp_rank, - dim=2) - up_proj = split(up_proj, - mapping.tp_size, - mapping.tp_rank, - dim=1) - - expert_weight_dict[ - f'model.layers.{l}.mlp.experts.up_gate_proj.weight'] = torch.concat( - [up_proj, gate_proj], dim=-2) - expert_weight_dict[ - f'model.layers.{l}.mlp.experts.down_proj.weight'] = down_proj - - # mlp.experts.down_proj.weight - moe_experts_down_proj_weights = expert_weight_dict[ - f'model.layers.{l}.mlp.experts.down_proj.weight'] - weights.update( - get_tllm_linear_weight(moe_experts_down_proj_weights, - trtllm_prefix + 'mlp.proj.')) - - # mlp.experts.up_gate.weight - moe_experts_up_gate_proj_weights = expert_weight_dict[ - f'model.layers.{l}.mlp.experts.up_gate_proj.weight'] - weights.update( - get_tllm_linear_weight(moe_experts_up_gate_proj_weights, - trtllm_prefix + 'mlp.fc.')) - - # MOE hardcoded routing_input into trt.float32, please refer to moe.py line 397 - moe_experts_gate_weights = load_weights(prefix + 'mlp.gate.weight', - torch.float32) - weights.update( - get_tllm_linear_weight(moe_experts_gate_weights, - trtllm_prefix + 'mlp.router.')) - - if moe_config.shared_expert_intermediate_size > 0: - shared_moe_up_proj_weights = load_weights( - prefix + 'mlp.shared_experts.up_proj.weight', torch_dtype) - shared_moe_up_proj_weights = split_matrix_tp( - shared_moe_up_proj_weights, - mapping.tp_size, - mapping.tp_rank, - dim=0) - shared_moe_down_proj_weights = load_weights( - prefix + 'mlp.shared_experts.down_proj.weight', torch_dtype) - shared_moe_down_proj_weights = split_matrix_tp( - shared_moe_down_proj_weights, - mapping.tp_size, - mapping.tp_rank, - dim=1) - shared_moe_gate_proj_weights = load_weights( - prefix + 'mlp.shared_experts.gate_proj.weight', torch_dtype) - shared_moe_gate_proj_weights = split_matrix_tp( - shared_moe_gate_proj_weights, - mapping.tp_size, - mapping.tp_rank, - dim=0) - shared_moe_gate_up_proj_weights = torch.concat( - [shared_moe_up_proj_weights, shared_moe_gate_proj_weights], - dim=-2) - - # mlp.shared_experts.gate_up_proj.weight - weights.update( - get_tllm_linear_weight( - shared_moe_gate_up_proj_weights, - trtllm_prefix + 'mlp.shared_expert.fc.')) - - # mlp.shared_experts.down_proj.weight - weights.update( - get_tllm_linear_weight( - shared_moe_down_proj_weights, - trtllm_prefix + 'mlp.shared_expert.proj.')) - - else: - # Current MLP layer is only one (layer 0), if it goes large consider to do fuse - mlp_gate_weight = load_weights(prefix + 'mlp.up_proj.weight', - torch_dtype) - split_gate = split_matrix_tp(mlp_gate_weight, - mapping.tp_size, - mapping.tp_rank, - dim=0) - weights.update( - get_tllm_linear_weight(split_gate, trtllm_prefix + 'mlp.gate.')) - - mlp_fc_weight = load_weights(prefix + 'mlp.gate_proj.weight', - torch_dtype) - split_fc = split_matrix_tp(mlp_fc_weight, - mapping.tp_size, - mapping.tp_rank, - dim=0) - weights.update( - get_tllm_linear_weight(split_fc, trtllm_prefix + 'mlp.fc.')) - - mlp_proj_weight = load_weights(prefix + 'mlp.down_proj.weight', - torch_dtype) - split_proj = split_matrix_tp(mlp_proj_weight, - mapping.tp_size, - mapping.tp_rank, - dim=1) - weights.update( - get_tllm_linear_weight(split_proj, trtllm_prefix + 'mlp.proj.')) - - # Layer norms do not use tensor parallelism - input_ln_weight = load_weights(prefix + 'input_layernorm.weight', - torch_dtype) - weights[trtllm_prefix + 'input_layernorm.weight'] = input_ln_weight - post_ln_weight = load_weights( - prefix + 'post_attention_layernorm.weight', torch_dtype) - weights[trtllm_prefix + 'post_layernorm.weight'] = post_ln_weight - - for l in layers_range: - convert_layer(l) - release_gc() - - v = load_weights('model.embed_tokens.weight', torch_dtype) - if use_parallel_embedding: - v = split_matrix_tp(v, - mapping.tp_size, - mapping.tp_rank, - dim=config.embedding_sharding_dim) - if mapping.is_first_pp_rank(): - weights['transformer.vocab_embedding.weight'] = v - lm_head_weights = load_weights('lm_head.weight', torch_dtype) - - if mapping.is_last_pp_rank(): - if config.vocab_size % mapping.tp_size != 0: - # padding - vocab_size_padded = pad_vocab_size(config.vocab_size, - mapping.tp_size) - pad_width = vocab_size_padded - config.vocab_size - lm_head_weights = torch.nn.functional.pad(lm_head_weights, - (0, 0, 0, pad_width), - 'constant', 0) - weights['lm_head.weight'] = split_matrix_tp(lm_head_weights, - mapping.tp_size, - mapping.tp_rank, - dim=0) - ln_f_w = load_weights('model.norm.weight', torch_dtype) - weights['transformer.ln_f.weight'] = ln_f_w - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Weights loaded. Total time: {t}') - return weights diff --git a/tensorrt_llm/models/deepseek_v2/model.py b/tensorrt_llm/models/deepseek_v2/model.py deleted file mode 100755 index c075d6fc7a95..000000000000 --- a/tensorrt_llm/models/deepseek_v2/model.py +++ /dev/null @@ -1,361 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import os -from typing import Optional - -import torch -import transformers - -from ..._utils import pad_vocab_size, torch_dtype_to_str -from ...functional import Tensor, non_gated_version, recv, send -from ...layers import (MOE, AttentionMaskType, ColumnLinear, - DeepseekV2Attention, Embedding, GatedMLP, MoeConfig, - PositionEmbeddingType, RmsNorm, SharedMoE) -from ...mapping import Mapping -from ...module import Module -from ...plugin import init_all_reduce_helper -from ..model_weights_loader import ModelWeightsLoader -from ..modeling_utils import (DecoderLayerList, DecoderModelForCausalLM, - PretrainedConfig) -from .config import DeepSeekV2Config -from .convert import convert_deepseekv2, load_weights_from_hf_safetensors - - -class DeepseekV2DecoderLayer(Module): - - def __init__(self, config: PretrainedConfig, layer_idx: int): - super().__init__() - self.layer_idx = layer_idx - self.config = config - - # Input layernorm in Deepseek v2 is same as Llama - self.input_layernorm = RmsNorm(normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - dtype=config.dtype) - - layers_range = config.mapping.pp_layers(config.num_hidden_layers) - local_layer_idx = layer_idx - layers_range[0] - - self.attention = DeepseekV2Attention( - local_layer_idx=local_layer_idx, - hidden_size=config.hidden_size, - num_attention_heads=config.num_attention_heads, - q_lora_rank=config.q_lora_rank, - kv_lora_rank=config.kv_lora_rank, - qk_nope_head_dim=config.qk_nope_head_dim, - qk_rope_head_dim=config.qk_rope_head_dim, - v_head_dim=config.v_head_dim, - max_position_embeddings=config.max_position_embeddings, - eps=config.norm_epsilon, - attention_mask_type=AttentionMaskType.causal, - dtype=config.dtype, - position_embedding_type=PositionEmbeddingType.learned_absolute, - rotary_embedding_base=config.rotary_base, - rotary_embedding_scaling=None, - rotary_embedding_beta_fast=config.rotary_scaling['beta_fast'], - rotary_embedding_beta_slow=config.rotary_scaling['beta_slow'], - rotary_embedding_mscale=config.rotary_scaling['mscale'], - rotary_embedding_mscale_all_dim=config. - rotary_scaling['mscale_all_dim'], - rotary_embedding_origin_max_position=config. - rotary_scaling['original_max_position_embeddings'], - rotary_scaling=config.rotary_scaling, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size, - tp_rank=config.mapping.tp_rank) - - # Added deepseek MoE and shared_experts - # First decoder layer: MLA + dense MLP + input_layernorm(RMSNorm) + post_attention_layernorm(RMSNorm) - # Rest decoder layer: MLA + MoE MLP + MoE Gate + shared_experts(MLP) + input_layernorm(RMSNorm) + post_attention_layernorm(RMSNorm) - # Added MLA in co-testing phase, use standard attention for MoE testing - - # Distinguish dense MLP and MoE MLP - # dense_config = DenseConfig(intermediate_size=config.intermediate_size) - moe_config = config.moe - # In case of moe_config is a dict - if isinstance(moe_config, dict): - moe_config = MoeConfig.from_dict(moe_config) - - if moe_config.num_experts > 0 and layer_idx > 0: - hidden_act = config.hidden_act - mlp_hidden_size = config.moe_inter_size - mlp_kwargs = {'moe_config': moe_config, 'mapping': config.mapping} - if moe_config.shared_expert_intermediate_size > 0: - ClsMLP = SharedMoE - mlp_kwargs['use_shared_gate'] = False - mlp_kwargs['use_side_stream'] = False - else: - ClsMLP = MOE - else: - ClsMLP = GatedMLP - mlp_hidden_size = config.intermediate_size - hidden_act = non_gated_version( - config.hidden_act) # back to non gated for dense layers - mlp_kwargs = {} - - self.mlp = ClsMLP(hidden_size=config.hidden_size, - ffn_hidden_size=mlp_hidden_size, - hidden_act=hidden_act, - dtype=config.dtype, - bias=False, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size, - quant_mode=config.quant_mode, - **mlp_kwargs) - - # Pose layernorm in Deepseek v2 is same as Llama - self.post_layernorm = RmsNorm(normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - dtype=config.dtype) - - def forward(self, - hidden_states, - attention_mask=None, - use_cache=False, - spec_decoding_params=None, - kv_cache_params=None, - attention_params=None): - - residual = hidden_states - hidden_states = self.input_layernorm(hidden_states) - - attention_output = self.attention( - hidden_states=hidden_states, - use_cache=use_cache, - spec_decoding_params=spec_decoding_params, - kv_cache_params=kv_cache_params, - attention_params=attention_params) - if use_cache: - attention_output, presents = attention_output - - hidden_states = residual + attention_output - - residual_attn = hidden_states - - hidden_states = self.post_layernorm(hidden_states) - hidden_states = self.mlp(hidden_states) - hidden_states = residual_attn + hidden_states - if use_cache: - return (hidden_states, presents) - return hidden_states - - -class DeepseekV2Model(Module): - - def __init__(self, config: PretrainedConfig) -> None: - super().__init__() - init_all_reduce_helper() # enable use_customer_all_reduce - self.dtype = config.dtype - self.mapping = config.mapping - if self.mapping.is_first_pp_rank(): - self.vocab_embedding = Embedding(config.vocab_size, - config.hidden_size, - dtype=config.dtype) - self.layers = DecoderLayerList(DeepseekV2DecoderLayer, config) - - if self.mapping.is_last_pp_rank(): - self.ln_f = RmsNorm(normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - dtype=config.dtype) - - self.head_num = config.num_attention_heads - self.head_size = config.qk_nope_head_dim + config.qk_rope_head_dim - - def forward(self, - input_ids, - position_ids=None, - use_cache=False, - attention_mask=None, - spec_decoding_params=None, - kv_cache_params=None, - attention_params=None, - hidden_states=None, - prompt_embedding_table: Optional[Tensor] = None, - prompt_tasks: Optional[Tensor] = None, - prompt_vocab_size: Optional[Tensor] = None): - - ptuning_args = [ - prompt_embedding_table, prompt_tasks, prompt_vocab_size - ] if prompt_embedding_table is not None else [] - - if self.mapping.is_first_pp_rank(): - hidden_states = self.vocab_embedding(input_ids, *ptuning_args) - else: - hidden_states = recv(hidden_states, self.mapping.prev_pp_rank()) - - hidden_states = self.layers.forward( - hidden_states, - use_cache=use_cache, - kv_cache_params=kv_cache_params, - attention_params=attention_params, - spec_decoding_params=spec_decoding_params) - - if use_cache: - hidden_states, presents = hidden_states - - if self.mapping.is_last_pp_rank(): - hidden_states = self.ln_f(hidden_states) - else: - hidden_states = send(hidden_states, self.mapping.next_pp_rank()) - - if use_cache: - return (hidden_states, tuple(presents)) - return hidden_states - - -class DeepseekV2ForCausalLM(DecoderModelForCausalLM): - config_class = DeepSeekV2Config - - def __init__(self, config: PretrainedConfig): - transformer = DeepseekV2Model(config) - vocab_size_padded = pad_vocab_size(config.vocab_size, - config.mapping.tp_size) - if config.mapping.is_last_pp_rank(): - lm_head = ColumnLinear(config.hidden_size, - vocab_size_padded, - bias=False, - dtype=config.dtype, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size, - gather_output=True) - else: - lm_head = None - self.mapping = config.mapping - super().__init__(config, transformer, lm_head) - - @classmethod - def from_hugging_face( - cls, - model_dir, - dtype: str = 'auto', - hf_model: Optional[transformers.PreTrainedModel] = None, - use_preloading: bool = False, - use_safetensors_loading: bool = False, - mapping: Optional[Mapping] = None, - override_fields={}, - **kwargs): - - if mapping is None: - mapping = Mapping() - pretrained_config = DeepSeekV2Config.from_hugging_face(model_dir, - dtype=dtype, - mapping=mapping, - **kwargs) - if dtype == 'auto': - dtype = getattr(pretrained_config, 'torch_dtype', None) - if dtype is None: - dtype = 'float16' - if isinstance(dtype, torch.dtype): - dtype = torch_dtype_to_str(dtype) - if dtype == 'float32': # should remove "float32" - dtype = 'float16' - if dtype == 'bfloat16' and torch.cuda.get_device_properties( - 0).major < 8: - logger.warning( - "Pre SM 80 GPUs do not support bfloat16, fallback to float16") - dtype = 'float16' - - deepseek = cls.from_config(pretrained_config) - - # If use_preloading is True, load the model from hf_model - # If use_safetensors_loading is True, load the model from safetensors - # if TRTLLM_DISABLE_UNIFIED_CONVERTER is not set, load the model use unified converter (recommended and default) - if use_preloading: - weights = convert_deepseekv2( - hf_model, - pretrained_config, - mapping, - dtype=dtype, - use_parallel_embedding=pretrained_config.use_parallel_embedding, - sharding_dim=pretrained_config.embedding_sharding_dim) - deepseek.load(weights) - return deepseek - - if os.environ.get("TRTLLM_DISABLE_UNIFIED_CONVERTER") is None: - - custom_dict = {} - rank_experts = mapping.ep_experts(pretrained_config.moe.num_experts) - for index, module in enumerate(deepseek.transformer.layers): - - if pretrained_config.q_lora_rank is not None: - module.attention.tllm_to_externel_key_dict = { - "fused_q_proj": ["q_b_proj.weight", "kv_b_proj.weight"], - "q_b_proj": "q_b_proj.weight", #v2 - "q_a_proj": "q_a_proj.weight", #v2 - "kv_b_proj": "kv_b_proj.weight", - "q_a_layernorm": "q_a_layernorm" - } - module.attention.fused_a.tllm_to_externel_key_dict = { - "fused_a": ["q_a_proj", "kv_a_proj_with_mqa"] - } #v2 - else: - module.attention.tllm_to_externel_key_dict = { - "fused_q_proj": ["q_proj.weight", - "kv_b_proj.weight"], #v2 lite - "q_b_proj": "q_proj.weight", #v2 lite - "kv_b_proj": "kv_b_proj.weight", - "q_a_layernorm": "q_a_layernorm" - } - module.attention.fused_a.tllm_to_externel_key_dict = { - "fused_a": "kv_a_proj_with_mqa" - } # v2 lite - - module.attention.kv_a_layernorm.tllm_to_externel_key_dict = { - 'kv_a_layernorm': 'kv_a_layernorm' - } - - if index > 0: - - module.mlp.shared_expert.fc.tllm_to_externel_key_dict = { - "fc": ["up_proj", "gate_proj"], - "shared_expert": "shared_experts" - } - module.mlp.shared_expert.proj.tllm_to_externel_key_dict = { - "shared_expert": "shared_experts" - } - module.mlp.fc.tllm_to_externel_key_dict = { - "fc": [ - f"experts.{expert}.up_proj" - for expert in rank_experts - ] + [ - f"experts.{expert}.gate_proj" - for expert in rank_experts - ] - } - module.mlp.proj.tllm_to_externel_key_dict = { - "proj": [ - f"experts.{expert}.down_proj" - for expert in rank_experts - ] - } - module.mlp.router.tllm_to_externel_key_dict = { - "mlp": "mlp", - "router": "gate" - } - - loader = ModelWeightsLoader(model_dir, custom_dict) - loader.generate_tllm_weights(deepseek) - return deepseek - - if use_safetensors_loading: - weights = load_weights_from_hf_safetensors( - model_dir, - pretrained_config, - mapping, - use_parallel_embedding=pretrained_config.use_parallel_embedding, - sharding_dim=pretrained_config.embedding_sharding_dim) - deepseek.load(weights) - return deepseek diff --git a/tensorrt_llm/models/dit/__init__.py b/tensorrt_llm/models/dit/__init__.py deleted file mode 100644 index 71bf6d298d42..000000000000 --- a/tensorrt_llm/models/dit/__init__.py +++ /dev/null @@ -1,14 +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. diff --git a/tensorrt_llm/models/dit/model.py b/tensorrt_llm/models/dit/model.py deleted file mode 100644 index 227262f99405..000000000000 --- a/tensorrt_llm/models/dit/model.py +++ /dev/null @@ -1,382 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import math -from collections import OrderedDict - -import numpy as np -import tensorrt as trt - -from ..._utils import str_dtype_to_trt, trt_dtype_to_str -from ...functional import (Tensor, allgather, arange, chunk, concat, constant, - cos, exp, expand, shape, silu, sin, slice, split, - unsqueeze) -from ...layers import MLP, BertAttention, Conv2d, Embedding, LayerNorm, Linear -from ...mapping import Mapping -from ...module import Module, ModuleList -from ...parameter import Parameter -from ...plugin import current_all_reduce_helper -from ...quantization import QuantMode -from ..modeling_utils import PretrainedConfig, PretrainedModel - - -def modulate(x, shift, scale, dtype): - ones = 1.0 - if dtype is not None: - ones = constant(np.ones(1, dtype=np.float32)).cast(dtype) - return x * (ones + unsqueeze(scale, 1)) + unsqueeze(shift, 1) - - -class TimestepEmbedder(Module): - - def __init__(self, hidden_size, frequency_embedding_size=256, dtype=None): - super().__init__() - self.dtype = dtype - self.mlp1 = Linear(frequency_embedding_size, - hidden_size, - bias=True, - dtype=dtype) - self.mlp2 = Linear(hidden_size, hidden_size, bias=True, dtype=dtype) - self.frequency_embedding_size = frequency_embedding_size - - def timestep_embedding(self, t, dim, max_period=10000): - half = dim // 2 - freqs = exp( - -math.log(max_period) * - arange(start=0, end=half, dtype=trt_dtype_to_str(trt.float32)) / - constant(np.array([half], dtype=np.float32))) - args = unsqueeze(t, -1).cast(trt.float32) * unsqueeze(freqs, 0) - embedding = concat([cos(args), sin(args)], dim=-1) - if self.dtype is not None: embedding = embedding.cast(self.dtype) - assert dim % 2 == 0 - return embedding - - def forward(self, t): - t_freq = self.timestep_embedding(t, self.frequency_embedding_size) - t_emb = self.mlp2(silu(self.mlp1(t_freq))) - - return t_emb - - -class LabelEmbedder(Module): - - def __init__(self, num_classes, hidden_size, dropout_prob, dtype=None): - super().__init__() - use_cfg_embedding = dropout_prob > 0 - self.embedding_table = Embedding(num_classes + use_cfg_embedding, - hidden_size, - dtype=dtype) - self.num_classes = num_classes - self.dropout_prob = dropout_prob - - def forward(self, labels, force_drop_ids=None): - assert force_drop_ids is None - embeddings = self.embedding_table(labels) - return embeddings - - -class PatchEmbed(Module): - - def __init__(self, - img_size: int, - patch_size: int, - input_c: int, - output_c: int, - bias: bool = True, - dtype: trt.DataType = None): - super().__init__() - self.img_size = img_size - self.patch_size = patch_size - self.num_patches = (img_size // patch_size)**2 - self.proj = Conv2d(input_c, - output_c, - kernel_size=(patch_size, patch_size), - stride=(patch_size, patch_size), - bias=bias, - dtype=dtype) - - def forward(self, x): - assert x.shape[2] == self.img_size - assert x.shape[3] == self.img_size - x = self.proj(x) - x = x.flatten(2).transpose(1, 2) # NCHW -> NLC - return x - - -class DiTBlock(Module): - - def __init__(self, - hidden_size, - num_heads, - mapping=Mapping(), - mlp_ratio=4.0, - dtype=None, - quant_mode=QuantMode(0)): - super().__init__() - self.dtype = dtype - self.norm1 = LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) - self.attn = BertAttention(hidden_size, - num_heads, - tp_group=mapping.tp_group, - tp_size=mapping.tp_size, - tp_rank=mapping.tp_rank, - cp_group=mapping.cp_group, - cp_size=mapping.cp_size, - cp_rank=mapping.cp_rank, - dtype=dtype, - quant_mode=quant_mode) - self.norm2 = LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) - self.mlp = MLP(hidden_size=hidden_size, - ffn_hidden_size=int(hidden_size * mlp_ratio), - hidden_act='gelu', - tp_group=mapping.tp_group, - tp_size=mapping.tp_size, - dtype=dtype, - quant_mode=quant_mode) - self.adaLN_modulation = Linear(hidden_size, - 6 * hidden_size, - tp_group=mapping.tp_group, - tp_size=mapping.tp_size, - bias=True, - dtype=dtype) - - def forward(self, x, c, input_lengths): - c = self.adaLN_modulation(silu(c)) - shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = chunk( - c, 6, dim=1) - - x = x + unsqueeze(gate_msa, 1) * self.attn(modulate( - self.norm1(x), shift_msa, scale_msa, self.dtype), - input_lengths=input_lengths) - x = x + unsqueeze(gate_mlp, 1) * self.mlp( - modulate(self.norm2(x), shift_mlp, scale_mlp, self.dtype)) - return x - - -class FinalLayer(Module): - - def __init__(self, - hidden_size, - patch_size, - out_channels, - mapping=Mapping(), - dtype=None): - super().__init__() - self.dtype = dtype - self.norm_final = LayerNorm(hidden_size, - elementwise_affine=False, - eps=1e-6) - self.linear = Linear(hidden_size, - patch_size * patch_size * out_channels, - bias=True, - dtype=dtype) - self.adaLN_modulation = Linear(hidden_size, - 2 * hidden_size, - tp_group=mapping.tp_group, - tp_size=mapping.tp_size, - bias=True, - dtype=dtype) - - def forward(self, x, c): - shift, scale = chunk(self.adaLN_modulation(silu(c)), 2, dim=1) - - x = modulate(self.norm_final(x), shift, scale, self.dtype) - x = self.linear(x) - - return x - - -class DiT(PretrainedModel): - - def __init__(self, config: PretrainedConfig): - self.check_config(config) - super().__init__(config) - self.learn_sigma = config.learn_sigma - self.in_channels = config.in_channels - self.out_channels = config.in_channels * 2 if config.learn_sigma else config.in_channels - self.input_size = config.input_size - self.patch_size = config.patch_size - self.num_heads = config.num_attention_heads - self.dtype = str_dtype_to_trt(config.dtype) - self.cfg_scale = config.cfg_scale - self.mapping = config.mapping - - self.x_embedder = PatchEmbed(config.input_size, - config.patch_size, - config.in_channels, - config.hidden_size, - bias=True, - dtype=self.dtype) - self.t_embedder = TimestepEmbedder(config.hidden_size, dtype=self.dtype) - self.y_embedder = LabelEmbedder(config.num_classes, - config.hidden_size, - config.class_dropout_prob, - dtype=self.dtype) - num_patches = self.x_embedder.num_patches - - self.pos_embed = Parameter(shape=(1, num_patches, config.hidden_size), - dtype=self.dtype) - self.blocks = ModuleList([ - DiTBlock(config.hidden_size, - config.num_attention_heads, - mlp_ratio=config.mlp_ratio, - mapping=config.mapping, - dtype=self.dtype, - quant_mode=config.quant_mode) - for _ in range(config.num_hidden_layers) - ]) - self.final_layer = FinalLayer(config.hidden_size, - config.patch_size, - self.out_channels, - mapping=config.mapping, - dtype=self.dtype) - - # We need to invoke default `__post_init__()` for quantized layers. - # def __post_init__(self): - # return - - def check_config(self, config: PretrainedConfig): - config.set_if_not_exist('input_size', 32) - config.set_if_not_exist('patch_size', 2) - config.set_if_not_exist('in_channels', 4) - config.set_if_not_exist('mlp_ratio', 4.0) - config.set_if_not_exist('class_dropout_prob', 0.1) - config.set_if_not_exist('num_classes', 1000) - config.set_if_not_exist('learn_sigma', True) - config.set_if_not_exist('dtype', None) - config.set_if_not_exist('cfg_scale', None) - - def unpatchify(self, x: Tensor): - c = self.out_channels - p = self.x_embedder.patch_size - h = w = int(x.shape[1]**0.5) - assert h * w == x.shape[1] - - x = x.view(shape=(x.shape[0], h, w, p, p, c)) - x = x.permute((0, 5, 1, 3, 2, 4)) - imgs = x.view(shape=(x.shape[0], c, h * p, h * p)) - return imgs - - def forward(self, latent, timestep, label): - """ - Forward pass of DiT. - latent: (N, C, H, W) - timestep: (N,) - label: (N,) - """ - if self.cfg_scale is not None: - output = self.forward_with_cfg(latent, timestep, label) - else: - output = self.forward_without_cfg(latent, timestep, label) - output.mark_output('output', self.dtype) - return output - - def forward_without_cfg(self, x, t, y): - """ - Forward pass without classifier-free guidance. - """ - x = self.x_embedder(x) + self.pos_embed.value - t = self.t_embedder(t) - y = self.y_embedder(y) - self.register_network_output('t_embedder', t) - self.register_network_output('x_embedder', x) - self.register_network_output('y_embedder', y) - c = t + y - input_length = constant(np.array([x.shape[1]], dtype=np.int32)) - input_lengths = expand(input_length, unsqueeze(shape(x, 0), 0)) - # Split squeence for CP here - if self.mapping.cp_size > 1: - assert x.shape[1] % self.mapping.cp_size == 0 - x = chunk(x, self.mapping.cp_size, dim=1)[self.mapping.cp_rank] - input_lengths = input_lengths // self.mapping.cp_size - for block in self.blocks: - x = block(x, c, input_lengths) # (N, T, D) - self.register_network_output('before_final_layer', x) - x = self.final_layer(x, c) # (N, T, patch_size ** 2 * out_channels) - self.register_network_output('final_layer', x) - - # All gather after CP - if self.mapping.cp_size > 1: - x = allgather(x, self.mapping.cp_group, gather_dim=1) - x = self.unpatchify(x) # (N, out_channels, H, W) - self.register_network_output('unpatchify', x) - return x - - def forward_with_cfg(self, x, t, y): - """ - Forward pass with classifier-free guidance. - """ - batch_size = shape(x, 0) - half = slice( - x, [0, 0, 0, 0], - concat([batch_size / 2, x.shape[1], x.shape[2], x.shape[3]])) - combined = concat([half, half], dim=0) - self.register_network_output('combined', combined) - model_out = self.forward_without_cfg(combined, t, y) - - _, d, h, w = model_out.shape - eps, rest = split(model_out, [3, d - 3], dim=1) - cond_eps = slice(eps, [0, 0, 0, 0], concat([batch_size / 2, 3, h, w])) - uncond_eps = slice(eps, concat([batch_size / 2, 0, 0, 0]), - concat([batch_size / 2, 3, h, w])) - self.register_network_output('cond_eps', cond_eps) - self.register_network_output('uncond_eps', uncond_eps) - - half_eps = uncond_eps + self.cfg_scale * (cond_eps - uncond_eps) - eps = concat([half_eps, half_eps], dim=0) - self.register_network_output('eps', eps) - - return concat([eps, rest], dim=1) - - def prepare_inputs(self, max_batch_size, **kwargs): - '''@brief: Prepare inputs Tensors for the model, the given sizes are used to determine the - ranges of the dimensions of when using TRT dynamic shapes. - - @return: a list contains values which can be fed into the self.forward() - ''' - mapping = self.config.mapping - if mapping.tp_size > 1: - current_all_reduce_helper().set_workspace_tensor(mapping, 1) - - def dit_default_range(max_batch_size): - return [2, max(2, (max_batch_size + 1) // 2), max_batch_size] - - default_range = dit_default_range - if self.cfg_scale is not None: - max_batch_size *= 2 - - latent = Tensor( - name='latent', - dtype=self.dtype, - shape=[-1, self.in_channels, self.input_size, self.input_size], - dim_range=OrderedDict([ - ('batch_size', [default_range(max_batch_size)]), - ('in_channels', [[self.in_channels] * 3]), - ('latent_height', [[self.input_size] * 3]), - ('latent_width', [[self.input_size] * 3]), - ])) - timestep = Tensor(name='timestep', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([ - ('batch_size', [default_range(max_batch_size)]), - ])) - label = Tensor(name='label', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([ - ('batch_size', [default_range(max_batch_size)]), - ])) - return {'latent': latent, 'timestep': timestep, 'label': label} diff --git a/tensorrt_llm/models/eagle/__init__.py b/tensorrt_llm/models/eagle/__init__.py deleted file mode 100644 index a08b2c2049a3..000000000000 --- a/tensorrt_llm/models/eagle/__init__.py +++ /dev/null @@ -1,14 +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. diff --git a/tensorrt_llm/models/eagle/config.py b/tensorrt_llm/models/eagle/config.py deleted file mode 100644 index e7a559f34690..000000000000 --- a/tensorrt_llm/models/eagle/config.py +++ /dev/null @@ -1,222 +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 json -from typing import Optional, Union - -from transformers import LlamaConfig - -from ..._utils import get_hf_rope_theta -from ...mapping import Mapping -from ..convert_utils import infer_dtype -from ..llama.config import LLaMAConfig -from ..modeling_utils import QuantAlgo, QuantConfig - - -class EagleConfig(LLaMAConfig): - - def __init__(self, - *, - num_eagle_layers: int = 1, - max_draft_len: int = 63, - max_non_leaves_per_layer: int = 10, - **kwargs): - self.num_eagle_layers = num_eagle_layers - self.max_non_leaves_per_layer = max_non_leaves_per_layer - self.max_draft_len = max_draft_len - self.eagle_net_config = LLaMAConfig.from_dict( - kwargs["eagle_net_config"]) - del kwargs["eagle_net_config"] - super().__init__(**kwargs) - - def to_dict(self): - output = super().to_dict() - # Serialize the fields added in EagleConfig - output['num_eagle_layers'] = self.num_eagle_layers - output['max_non_leaves_per_layer'] = self.max_non_leaves_per_layer - output['max_draft_len'] = self.max_draft_len - output['eagle_net_config'] = self.eagle_net_config.to_dict() - return output - - @classmethod - def from_hugging_face( - cls, - hf_config_or_dir: Union[str, 'transformers.PretrainedConfig'], - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - import transformers - trust_remote_code = kwargs.pop('trust_remote_code', True) - speculative_config_or_dir = kwargs.pop('speculative_model_dir', None) - - if isinstance(hf_config_or_dir, transformers.PretrainedConfig): - hf_config = hf_config_or_dir - else: - hf_config_dir = str(hf_config_or_dir) - - hf_config = transformers.AutoConfig.from_pretrained( - hf_config_dir, trust_remote_code=trust_remote_code) - dtype = infer_dtype(dtype, getattr(hf_config, 'torch_dtype', None)) - - hf_config = None - hf_config_or_dir if speculative_config_or_dir is None else speculative_config_or_dir - if hf_config_or_dir is not None: - hf_config = LlamaConfig.from_pretrained(hf_config_or_dir) - - hf_config.model_type - n_head = hf_config.num_attention_heads - inter_size = hf_config.intermediate_size - n_layer = hf_config.num_hidden_layers - n_embd = hf_config.hidden_size - n_kv_head = hf_config.num_key_value_heads - rms_norm_eps = hf_config.rms_norm_eps - vocab_size = hf_config.vocab_size - rotary_scaling = hf_config.rope_scaling - rotary_base = get_hf_rope_theta(hf_config, 10000.0) - n_positions = hf_config.max_position_embeddings - hidden_act = hf_config.hidden_act - dtype = str(hf_config.torch_dtype)[6:] if dtype == 'auto' else dtype - if hasattr(hf_config, 'head_dim'): - head_dim = hf_config.head_dim - else: - head_dim = hf_config.n_embd // hf_config.n_head - if hasattr(hf_config, 'head_size'): - head_size = hf_config.head_size - else: - head_size = head_dim - - if speculative_config_or_dir is None: - hf_config_eagle = hf_config.eagle - n_head_eagle = hf_config_eagle['num_attention_heads'] - inter_size_eagle = hf_config_eagle['intermediate_size'] - n_layer_eagle = hf_config_eagle['num_hidden_layers'] - n_embd_eagle = hf_config_eagle['hidden_size'] - n_kv_head_eagle = hf_config_eagle['num_key_value_heads'] - rms_norm_eps_eagle = hf_config_eagle['rms_norm_eps'] - n_positions_eagle = hf_config_eagle['max_position_embeddings'] - else: - hf_config_eagle = LlamaConfig.from_pretrained( - speculative_config_or_dir) - n_head_eagle = hf_config_eagle.num_attention_heads - inter_size_eagle = hf_config_eagle.intermediate_size - n_layer_eagle = hf_config_eagle.num_hidden_layers - n_embd_eagle = hf_config_eagle.hidden_size - n_kv_head_eagle = hf_config_eagle.num_key_value_heads - rms_norm_eps_eagle = hf_config_eagle.rms_norm_eps - n_positions_eagle = hf_config_eagle.max_position_embeddings - - if rotary_scaling is not None: - # assert use_gpt_attention_plugin, "RoPE scaling is only supported through GPT attention plugin." - rotary_scaling = { - "type": rotary_scaling["rope_type"], - } - rotary_scaling = rotary_scaling - - eagle_net_config = { - 'architecture': "LlamaForCausalLM", - 'dtype': dtype, - 'logits_dtype': 'float32', - 'num_hidden_layers': n_layer_eagle, - 'num_attention_heads': n_head_eagle, - 'hidden_size': n_embd_eagle, - 'intermediate_size': inter_size_eagle, - 'num_key_value_heads': n_kv_head_eagle, - 'vocab_size': vocab_size, - 'position_embedding_type': 'rope_gpt_neox', - 'max_position_embeddings': n_positions_eagle, - 'hidden_act': hidden_act, - 'rotary_base': rotary_base, - 'rotary_scaling': rotary_scaling, - 'norm_epsilon': rms_norm_eps_eagle, - 'quantization': { - 'quant_algo': None, - 'kv_cache_quant_algo': None, - }, - 'mapping': { - 'world_size': mapping.world_size, - 'tp_size': mapping.tp_size, - 'pp_size': mapping.pp_size, - }, - 'use_parallel_embedding': kwargs['use_parallel_embedding'], - 'embedding_sharding_dim': kwargs['embedding_sharding_dim'], - 'head_dim': head_dim, - 'head_size': head_size - } - - config = { - 'architecture': 'EagleForCausalLM', - 'dtype': dtype, - 'logits_dtype': 'float32', - 'num_hidden_layers': n_layer, - 'num_attention_heads': n_head, - 'hidden_size': n_embd, - 'intermediate_size': inter_size, - 'num_key_value_heads': n_kv_head, - 'vocab_size': vocab_size, - 'position_embedding_type': 'rope_gpt_neox', - 'max_position_embeddings': n_positions, - 'hidden_act': hidden_act, - 'rotary_base': rotary_base, - 'rotary_scaling': rotary_scaling, - 'norm_epsilon': rms_norm_eps, - 'quantization': { - 'quant_algo': None, - 'kv_cache_quant_algo': None, - }, - 'mapping': { - 'world_size': mapping.world_size, - 'tp_size': mapping.tp_size, - 'pp_size': mapping.pp_size, - }, - 'use_parallel_embedding': kwargs['use_parallel_embedding'], - 'embedding_sharding_dim': kwargs['embedding_sharding_dim'], - 'num_eagle_layers': kwargs['speculative_config'].num_eagle_layers, - 'max_non_leaves_per_layer': - kwargs['speculative_config'].max_non_leaves_per_layer, - 'eagle_net_config': eagle_net_config - } - if quant_config: - config['quantization']['quant_algo'] = quant_config.quant_algo - config['quantization'][ - 'kv_cache_quant_algo'] = quant_config.kv_cache_quant_algo - - if quant_config.quant_algo == QuantAlgo.W4A16_GPTQ: - config['quantization'].update({ - "group_size": quant_config.group_size, - "has_zero_point": True, - "pre_quant_scale": False, - 'quant_algo': QuantAlgo.W4A16_GPTQ - }) - eagle_quant_config = {} - try: - with open( - str(speculative_config_or_dir) + '/' + - 'hf_quant_config.json') as f: - eagle_quant_config = json.load(f) - if "lm_head" in eagle_quant_config['quantization'][ - 'exclude_modules']: - eagle_quant_config['quantization']['exclude_modules'] += [ - f"eagle_nets.{i}.lm_head" for i in range( - kwargs['speculative_config'].num_eagle_layers) - ] - config['quantization'].update( - eagle_quant_config['quantization']) - config['eagle_net_config']['quantization'].update( - eagle_quant_config['quantization']) - except IOError: - pass - - return cls.from_dict(config) diff --git a/tensorrt_llm/models/eagle/model.py b/tensorrt_llm/models/eagle/model.py deleted file mode 100644 index 779f397d70ed..000000000000 --- a/tensorrt_llm/models/eagle/model.py +++ /dev/null @@ -1,1327 +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. - -from collections import OrderedDict -from typing import Optional, Union - -import numpy as np -import tensorrt as trt -from tqdm import tqdm - -from tensorrt_llm.mapping import Mapping -from tensorrt_llm.models.generation_mixin import GenerationMixin -from tensorrt_llm.models.llama.model import LLaMAForCausalLM, LLaMAModel -from tensorrt_llm.models.model_weights_loader import ModelWeightsLoader - -from ..._common import default_net, default_trtnet -from ..._utils import pad_vocab_size -from ...functional import (Tensor, _create_tensor, cast, concat, - gather_last_token_logits, index_select, shape) -from ...layers import AttentionParams, ColumnLinear, SpecDecodingParams -from ...llmapi.kv_cache_type import KVCacheType -from ...module import Module, ModuleList -from ...plugin import TRT_LLM_PLUGIN_NAMESPACE -from ..modeling_utils import QuantConfig -from .config import EagleConfig - - -class TreeParams(object): - - def __init__(self, paths: Tensor = None): - self.paths = paths # on GPU - - -def eagle_sample_and_accept_draft_plugin(lm_logits: Tensor = None, - draft_tokens: Tensor = None, - draft_lens: Tensor = None, - eagle_temperature: Tensor = None, - rand_data_validation: Tensor = None, - posterior_alpha: Tensor = None, - posterior_threshold: Tensor = None, - tree_params: TreeParams = None, - greedy_sampling: Tensor = None, - use_dynamic_tree: Tensor = None): - ''' - Takes input logits and samples golden token + predictions from draft tokens. - Runs acceptance algorithm to accept draft tokens. - When greedy_sampling is True, all decoding is done using Top1 and token equality is used - for acceptance. Otherwise, typical acceptance and multinomial samplings are used. - - Visit tests/model/eagle/test_sample_accept_draft_tokens.py for input/output examples. - - Parameters: - lm_logits : Tensor - [num_tokens, vocab_size] - Logits produced by the base model. - - draft_tokens : Tensor - [batch_size, max_decoding_draft_tokens] - Input draft tokens. Only the first draft_lens[bi] tokens are relevant for bi'th row. - - draft_lens : Tensor - [batch_size] - Lengths of the draft_tokens. 0 for context request. Actual draft length for generation requests. - - eagle_temperature : Tensor - [batch_size] - Temperature of the decoding. - - rand_data_validation : Tensor - [batch_size, max_decoding_tokens] - Random data for multinomial sampling. - - posterior_alpha : Tensor - [batch_size] - Delta in typical acceptance in https://arxiv.org/pdf/2401.10774. - - posterior_threshold : Tensor - [batch_size] - Minimum probability threshold. - Epsilon in typical acceptance in https://arxiv.org/pdf/2401.10774. - - tree_params : TreeParams - Tree params of the input draft tokens. - - greedy_sampling : Tensor - Whether to do greedy or non-greedy sampling. - - use_dynamic_tree: Tensor - Whether to use dynamic tree (i.e., Eagle-2) - - - Return: - accepted_tokens : Tensor - [batch_size, max_path_len] - Accepted token ids. Only the first num_accepted_tokens[bi] tokens are relevant for bi'th row, - - num_accepted_tokens : Tensor - [batch_size] - Number of accepted tokens per request. Each entry is >= 1. - - accepted_paths : Tensor - [batch_size] - Indices of the accepted path per request of the input paths in tree_params.paths. - - next_draft_tokens : Tensor - [batch_size, max_decoding_draft_tokens] - Empty tensor used to allocate space for the next draft tokens. - - next_draft_lens : Tensor - [batch_size] - Empty tensor used to allocate space for lens of the next draft tokens. - - next_draft_paths : Tensor - [batch_size, max_decoding_len, max_path_len] - For EAGLE-1 just a copy of input path. - - hidden_size_batch_level_starts : Tensor - [max_draft_path_len * batch_size + 1] - Empty tensor used to allocate space for eagle_prepare_drafter_inputs_plugin. - ''' - - plg_creator = trt.get_plugin_registry().get_plugin_creator( - 'EagleSampleAndAcceptDraftTokens', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert plg_creator is not None - - pf_type = trt.PluginField("type_id", - np.array([int(lm_logits.dtype)], np.int32), - trt.PluginFieldType.INT32) - - pfc = trt.PluginFieldCollection([pf_type]) - plugin = plg_creator.create_plugin("eagle_sample_and_accept_draft_plugin", - pfc) - - plug_inputs = [ - lm_logits, draft_tokens, draft_lens, eagle_temperature, - rand_data_validation, posterior_alpha, posterior_threshold, - tree_params.paths, greedy_sampling, use_dynamic_tree - ] - - plug_inputs = [i.trt_tensor for i in plug_inputs] - layer = default_trtnet().add_plugin_v2(plug_inputs, plugin) - - accepted_tokens = _create_tensor(layer.get_output(0), layer) - num_accepted_tokens = _create_tensor(layer.get_output(1), layer) - accepted_paths = _create_tensor(layer.get_output(2), layer) - next_draft_tokens = _create_tensor(layer.get_output(3), layer) - next_draft_lens = _create_tensor(layer.get_output(4), layer) - next_draft_paths = _create_tensor(layer.get_output(5), layer) - hidden_size_batch_level_starts = _create_tensor(layer.get_output(6), layer) - return tuple([ - accepted_tokens, num_accepted_tokens, accepted_paths, next_draft_tokens, - next_draft_lens, next_draft_paths, hidden_size_batch_level_starts - ]) - - -def eagle_draft_decoder_plugin( - layer_idx: int, num_eagle_layers: int, top_k_sampling: bool, - logits: Tensor, num_last_token_indices: Tensor, input_paths: Tensor, - use_dynamic_tree: Tensor, dynamic_tree_max_topK: Tensor, - input_draft_token_ids: Tensor, input_draft_lens: Tensor, - input_prev_scores: Tensor, input_current_expand_indices: Tensor, - input_all_layers_scores: Tensor, - input_all_layers_draft_token_ids: Tensor, - input_all_layers_draft_token_ids_predecessor: Tensor): - ''' - Parameters: - layer_idx : int - The index of the EagleNet. - - num_eagle_layers: int - The total number of eagle layers. - - top_k_sampling: bool - Whether to use top K sampling. Otherwise, use multinomial sampling. - - logits : Tensor - [num_logits, vocab_size] - Input logits. - - num_last_token_indices : Tensor - [1] - Number of valid logits in logits. - - input_paths: Tensor - [batch_size, max_decoding_tokens, max_path_len] - Input paths - - use_dynamic_tree: Tensor - [1] - Whether use dynamic tree (i.e., Eagle-2) - - dynamic_tree_max_topK: Tensor - [1] - Number of draft tokens expand in Eagle-2. - - input_draft_token_ids: Tensor - [batch_size, max_decoding_draft_tokens] - Draft tokens generated by previous EagleNets. - - input_draft_lens: Tensor - [batch_size] - Number of draft tokens for each request. - - input_prev_scores: Tensor - [batch_size, max_decoding_draft_tokens] - Last layer's scores - - input_current_expand_indices: Tensor - [batch_size, max_decoding_draft_tokens] - The indices of the nodes that expand in this layer. - The index is related to the final output tree, which has max_decoding_draft_tokens draft tokens. - - input_all_layers_scores: Tensor - [batch_size, num_eagle_layers, max_decoding_draft_tokens x max_decoding_draft_tokens] - For Eagle-2, record scores from all EagleNets - - input_all_layers_draft_token_ids: Tensor - [batch_size, num_eagle_layers, max_decoding_draft_tokens x max_decoding_draft_tokens] - For Eagle-2, record all draft tokens from all EagleNets - - input_all_layers_draft_token_ids_predecessor: Tensor - [batch_size, num_eagle_layers, max_decoding_draft_tokens x max_decoding_draft_tokens] - For Eagle-2, record all draft tokens' predecessor - - Return: - output_draft_token_ids: Tensor - [batch_size, max_decoding_draft_tokens] - Draft tokens generated by this EagleNets, also include the previous draft tokens. - - output_draft_draft_lens: Tensor - [batch_size] - Number of draft tokens for each request. - - output_paths: Tensor - [batch_size, max_decoding_draft_tokens, max_path_len] - The latest path. - - output_current_scores: Tensor - [batch_size, max_decoding_draft_tokens] - This layer's scores, which will be used in next layer. - - output_next_expand_indices: - [batch_size, max_decoding_draft_tokens] - The indices of the nodes that expand in next layer. - The index is related to the final output tree, which has max_decoding_draft_tokens draft tokens. - - output_all_layers_scores: - [batch_size, num_eagle_layers, max_decoding_draft_tokens x max_decoding_draft_tokens] - For Eagle-2, record scores from all EagleNets - - output_all_layers_draft_token_ids: Tensor - [batch_size, num_eagle_layers, max_decoding_draft_tokens x max_decoding_draft_tokens] - For Eagle-2, record all draft tokens from all EagleNets - - output_all_layers_draft_token_ids_predecessor: Tensor - [batch_size, num_eagle_layers, max_decoding_draft_tokens x max_decoding_draft_tokens] - For Eagle-2, record all draft tokens' predecessor - - ''' - - plg_creator = trt.get_plugin_registry().get_plugin_creator( - 'EagleDecodeDraftTokens', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert plg_creator is not None - - pf_type = trt.PluginField("type_id", np.array([int(logits.dtype)], - np.int32), - trt.PluginFieldType.INT32) - - layer_idx_t = trt.PluginField("layer_idx", - np.array(layer_idx, dtype=np.int32), - trt.PluginFieldType.INT32) - - num_eagle_layers_t = trt.PluginField( - "num_eagle_layers", np.array(num_eagle_layers, dtype=np.int32), - trt.PluginFieldType.INT32) - - top_k_sampling_t = 1 if top_k_sampling else 0 - top_k_sampling_t = trt.PluginField( - "top_k_sampling", np.array(top_k_sampling_t, dtype=np.int32), - trt.PluginFieldType.INT32) - - pfc = trt.PluginFieldCollection( - [pf_type, layer_idx_t, num_eagle_layers_t, top_k_sampling_t]) - plugin = plg_creator.create_plugin("eagle_draft_decoder_plugin", pfc) - - plug_inputs = [ - logits, input_paths, num_last_token_indices, use_dynamic_tree, - dynamic_tree_max_topK, input_draft_token_ids, input_draft_lens, - input_prev_scores, input_current_expand_indices, - input_all_layers_scores, input_all_layers_draft_token_ids, - input_all_layers_draft_token_ids_predecessor - ] - - plug_inputs = [i.trt_tensor for i in plug_inputs] - layer = default_trtnet().add_plugin_v2(plug_inputs, plugin) - - output_draft_token_ids = _create_tensor(layer.get_output(0), layer) - output_draft_lens = _create_tensor(layer.get_output(1), layer) - output_paths = _create_tensor(layer.get_output(2), layer) - output_current_scores = _create_tensor(layer.get_output(3), layer) - output_next_expand_indices = _create_tensor(layer.get_output(4), layer) - output_all_layers_scores = _create_tensor(layer.get_output(5), layer) - output_all_layers_draft_token_ids = _create_tensor(layer.get_output(6), - layer) - output_all_layers_draft_token_ids_predecessor = _create_tensor( - layer.get_output(7), layer) - return tuple([ - output_draft_token_ids, output_draft_lens, output_paths, - output_current_scores, output_next_expand_indices, - output_all_layers_scores, output_all_layers_draft_token_ids, - output_all_layers_draft_token_ids_predecessor - ]) - - -def eagle_prepare_drafter_inputs_plugin( - layer_idx: int, num_layers: int, max_non_leaves_per_layer: int, - attention_params: AttentionParams, input_ids: Tensor, - chunked_context_next_tokens: Tensor, accepted_token_ids: Tensor, - accepted_lens: Tensor, accepted_path_ids: Tensor, - next_draft_tokens: Tensor, next_draft_lens: Tensor, - next_draft_paths: Tensor, prev_draft_lens: Tensor, - prev_draft_paths: Tensor, hidden_size_batch_level_starts: Tensor, - input_gen_tokens: Tensor, - input_spec_decoding_generation_lengths: Tensor): - ''' - Prepares inputs for the EagleNet inference. - - Visit tests/model/eagle/test_prepare_drafter_inputs.py for input/output examples. - - Parameters: - layer_idx : int - Index of the EagleNet. 0 means context phase EagleNet or EagleNet0, - > 0 means EagleNetX or generation phase of EagleNet - - num_layers : int - Number of Eagle layers. - - max_non_leaves_per_layer : int - Number of nodes that can be non leaf in the tree at each level of the tree. - - attention_params : AttentionParams - - input_ids : Tensor - [num_tokens] - Tokens ids, inputs to the base model. - - chunked_context_next_tokens : Tensor - [batch_size] - The first token of the next chunk in chunked context. - -1 if current chunk is the last chunk or requests is in the gen phase. - - accepted_token_ids : Tensor - [batch_size, max_path_len] - Accepted tokens ids. - - accepted_lens : Tensor - [batch_size] - Number of accepted tokens. - - accepted_path_ids : Tensor - [batch_size] - Idx of the accepted path in prev_draft_paths. - - next_draft_tokens : Tensor - [batch_size, max_decoding_draft_tokens] - Tokens ids of the draft tokens being drafted by EagleNet - - next_draft_lens : Tensor - [batch_size] - Number of drafted tokens in next_draft_tokens - - next_draft_paths : Tensor - [batch_size, max_decoding_tokens, max_path_len] - Paths of the draft tokens for the next iteration. In EAGLE-1 is the same as prev_draft_paths - - prev_draft_lens : Tensor - [batch_size] - Number of draft tokens, inputs to the base model. - 0 for ctx requests and actual draft len for gen requests. - - prev_draft_paths : Tensor - [batch_size, max_decoding_tokens, max_path_len] - Paths of the draft tokens inputs to the base model. - - hidden_size_batch_level_starts : Tensor - [max_draft_path_len * batch_size + 1] - Exclusive sum of the starts of the segments of the hidden states in the concatenated array. - Hidden states shape is (flattened and w/o padding) - [max_draft_path_len, batch_size, num_output_tokens_i_j], where num_output_tokens_i_j - depends on the path of request j at level i. - - input_gen_tokens : Tensor - [num_gen_tokens] - Only needed to infer number of generation tokens from its shape. The content is irrelevant - - input_spec_decoding_generation_lengths : Tensor - [num_gen_requests] - Number of tokens for the base model. Only used to infer num_gen_requests from its shape, the content is irrelevant. - - Return: - sequence_length : Tensor - [batch_size] - Sequence length of the next EagleNet iteration. - For EagleNet0 equals to the (prompt_len + num_generated_tokens + accepted_lens). - For EagleNetX (X > 0) (seq_len_eagle_net_0 + spec_decoding_generation_lengths). - - context_length : Tensor - [batch_size] - Context length of the next EagleNet iteration. - For EagleNet0 it is either the actual context length of the request (for ctx requests) - or the number of accepted tokens in this iteration. EagleNet0's attn does chunked context attn. - For EagleNetX (X > 0), context length equals to the sequence length of the EagleNet0. - - spec_decoding_generation_lengths : Tensor - [batch_size] - Only relevant for EagleNetX (X > 0). - Number of draft tokens. - - spec_decoding_position_offsets : Tensor - [batch_size, max_decoding_tokens] - Only relevant for EagleNetX (X > 0). - Position offsets of the selected tokens from output_ids. - - spec_decoding_packed_mask : Tensor - [batch_size, max_decoding_tokens, ceil(max_decoding_tokens / 32)] - Only relevant for EagleNetX (X > 0). - uint32_t packed masks. - - output_ids : Tensor - [batch_size * max_non_leaves_per_layer * layer_idx] for layer_idx > 0 - [num_tokens - num_gen_tokens + num_gen_requests * (num_layers + 1)] for layer_idx == 0 - Token ids selected for the EagleNet iteration. - Tensor's actual size is larger than num_output_tokens. - - position_ids : Tensor - [batch_size] for layer_idx > 0 - [num_tokens - num_gen_tokens + num_gen_requests * (num_layers + 1)] for layer_idx == 0 - Position ids of the tokens selected for the EagleNet iteration. - Tensor's actual size is larger than num_output_tokens. - - hidden_states_indices : Tensor - [batch_size * max_non_leaves_per_layer * layer_idx] for layer_idx > 0 - [num_tokens - num_gen_tokens + num_gen_requests * (num_layers + 1)] for layer_idx == 0 - Indices of the hidden states to be selected from aggregated hidden states for the next iteration. - Tensor's actual size is larger than num_output_tokens. - - last_token_indices : Tensor - [batch_size * max_non_leaves_per_layer] - Indices of the hidden states to be converted to logits after the next EagleNet iteration. - Tensor's actual size is larger than num_output_tokens. - - num_last_token_indices : Tensor - [] - Number of logits selected after the next EagleNet iteration. - Tensors containing size of the outputs of V3 plugins. 0-D tensor. - - out_hidden_size_batch_level_starts : Tensor - [max_draft_path_len * batch_size + 1] - Same as hidden_size_batch_level_starts, but with updated path lens for the next level. - ''' - - plg_creator = trt.get_plugin_registry().get_creator( - 'EaglePrepareDrafterInputs', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert plg_creator is not None - - layer_idx = trt.PluginField("layer_idx", np.array(layer_idx, - dtype=np.int32), - trt.PluginFieldType.INT32) - - num_layers = trt.PluginField("num_layers", - np.array(num_layers, dtype=np.int32), - trt.PluginFieldType.INT32) - - max_non_leaves_per_layer = trt.PluginField( - "max_non_leaves_per_layer", - np.array(max_non_leaves_per_layer, dtype=np.int32), - trt.PluginFieldType.INT32) - - pfc = trt.PluginFieldCollection( - [layer_idx, num_layers, max_non_leaves_per_layer]) - plugin = plg_creator.create_plugin("eagle_prepare_drafter_inputs_plugin", - pfc, trt.TensorRTPhase.BUILD) - - plug_inputs = [ - attention_params.sequence_length, attention_params.context_lengths, - input_ids, chunked_context_next_tokens, accepted_token_ids, - accepted_lens, accepted_path_ids, next_draft_tokens, next_draft_lens, - next_draft_paths, prev_draft_lens, prev_draft_paths, - hidden_size_batch_level_starts, input_gen_tokens, - input_spec_decoding_generation_lengths - ] - - plug_inputs = [i.trt_tensor for i in plug_inputs] - shape_inputs = [] - layer = default_trtnet().add_plugin_v3(plug_inputs, shape_inputs, plugin) - - sequence_length = _create_tensor(layer.get_output(0), layer) - context_length = _create_tensor(layer.get_output(1), layer) - spec_decoding_generation_lengths = _create_tensor(layer.get_output(2), - layer) - spec_decoding_position_offsets = _create_tensor(layer.get_output(3), layer) - spec_decoding_packed_mask = _create_tensor(layer.get_output(4), layer) - output_ids = _create_tensor(layer.get_output(5), layer) - position_ids = _create_tensor(layer.get_output(6), layer) - hidden_states_indices = _create_tensor(layer.get_output(7), layer) - last_token_indices = _create_tensor(layer.get_output(8), layer) - num_last_token_indices = _create_tensor(layer.get_output(9), layer) - out_hidden_size_batch_level_starts = _create_tensor(layer.get_output(10), - layer) - return tuple([ - sequence_length, context_length, spec_decoding_generation_lengths, - spec_decoding_position_offsets, spec_decoding_packed_mask, output_ids, - position_ids, hidden_states_indices, last_token_indices, - num_last_token_indices, out_hidden_size_batch_level_starts - ]) - - -class EagleNet(Module): - - def __init__(self, config, logits_dtype): - super().__init__() - self.drafter = LLaMAModel(config) - self.config = config - self.logits_dtype = logits_dtype - - vocab_size_padded = pad_vocab_size(config.vocab_size, - config.mapping.tp_size) - if config.mapping.is_last_pp_rank(): - self.lm_head = ColumnLinear(config.hidden_size, - vocab_size_padded, - bias=False, - dtype=config.dtype, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size, - gather_output=True) - else: - self.lm_head = None - - def forward(self, - input_ids, - position_ids=None, - hidden_states=None, - last_token_indices=None, - spec_decoding_params=None, - kv_cache_params=None, - attention_params=None): - hidden_states, cache = self.drafter( - input_ids, - position_ids=position_ids, - use_cache=True, - spec_decoding_params=spec_decoding_params, - kv_cache_params=kv_cache_params, - attention_params=attention_params, - hidden_states_for_embed=hidden_states) - - if self.config.mapping.is_last_pp_rank(): - hidden_states = gather_last_token_logits( - hidden_states, last_token_indices, - default_net().plugin_config.remove_input_padding) - return cast(self.lm_head(hidden_states), - dtype=self.logits_dtype), hidden_states, cache - - return None, hidden_states, cache - - -class EagleForCausalLM(LLaMAForCausalLM): - config_class = EagleConfig - - def __init__(self, config: EagleConfig): - - super().__init__(config) - - self.num_eagle_layers = config.num_eagle_layers - self.max_non_leaves_per_layer = config.max_non_leaves_per_layer - self.hidden_size = config.hidden_size - self.vocab_size = config.vocab_size - vocab_size_padded = pad_vocab_size(self.vocab_size, - config.mapping.tp_size) - eagle_net_config = config.eagle_net_config - eagle_net_config.mapping = Mapping(world_size=config.mapping.world_size, - rank=config.mapping.rank, - cp_size=1, - tp_size=config.mapping.world_size, - pp_size=1) - - eagle_net_config.fc_after_embed = True - eagle_net_config.use_input_layernorm_in_first_layer = False - eagle_net_config.use_last_layernorm = False - eagle_net_config.layer_idx_offset = config.num_hidden_layers - if self.mapping.is_last_pp_rank(): - self.eagle_nets = ModuleList([ - EagleNet(config=eagle_net_config, - logits_dtype=config.logits_dtype) - for _ in range(self.num_eagle_layers) - ]) - self.max_draft_len = config.max_draft_len - - def _prepare_drafter_inputs( - self, layer_idx, input_ids, chunked_context_next_tokens, - accepted_token_ids, accepted_lens, accepted_path_ids, - next_draft_tokens, next_draft_lens, next_draft_paths, - prev_draft_lens, prev_draft_paths, input_attention_params, - input_kv_cache_params, hidden_states, - host_ctx_eagle_net_request_types, - host_ctx_eagle_net_context_lengths, - host_ctx_eagle_net_past_key_value_lengths, - host_gen_eagle_net_request_types, - host_gen_eagle_net_context_lengths, - host_gen_eagle_net_past_key_value_lengths, - hidden_size_batch_level_starts, input_gen_tokens, - input_spec_decoding_generation_lengths, spec_decoding_use): - - drafter_inputs = eagle_prepare_drafter_inputs_plugin( - layer_idx, self.num_eagle_layers, self.max_non_leaves_per_layer, - input_attention_params, input_ids, chunked_context_next_tokens, - accepted_token_ids, accepted_lens, accepted_path_ids, - next_draft_tokens, next_draft_lens, next_draft_paths, - prev_draft_lens, prev_draft_paths, hidden_size_batch_level_starts, - input_gen_tokens, input_spec_decoding_generation_lengths) - - sequence_length, context_lengths, \ - spec_decoding_generation_lengths, spec_decoding_position_offsets, \ - spec_decoding_packed_mask, output_ids, position_ids, hidden_states_indices, \ - last_token_indices, num_last_token_indices, out_hidden_size_batch_level_starts \ - = drafter_inputs - - attention_params = input_attention_params - kv_cache_params = input_kv_cache_params - attention_params.sequence_length = sequence_length - attention_params.context_lengths = context_lengths - - if layer_idx == 0: - attention_params.host_request_types = host_ctx_eagle_net_request_types - attention_params.host_context_lengths = host_ctx_eagle_net_context_lengths - kv_cache_params.host_past_key_value_lengths = host_ctx_eagle_net_past_key_value_lengths - else: - attention_params.host_request_types = host_gen_eagle_net_request_types - attention_params.host_context_lengths = host_gen_eagle_net_context_lengths - kv_cache_params.host_past_key_value_lengths = host_gen_eagle_net_past_key_value_lengths - - spec_decoding_params = None - if layer_idx > 0: - spec_decoding_params = SpecDecodingParams( - True, self.max_draft_len, spec_decoding_generation_lengths, - spec_decoding_position_offsets, spec_decoding_packed_mask, - spec_decoding_use) - - # Get hidden states for accepted ids - hidden_states = self._slice_hidden_states(hidden_states, - hidden_states_indices) - - eagle_net_inputs = {} - eagle_net_inputs["input_ids"] = output_ids - eagle_net_inputs["position_ids"] = position_ids - eagle_net_inputs["last_token_indices"] = last_token_indices - eagle_net_inputs["attention_params"] = attention_params - eagle_net_inputs["kv_cache_params"] = kv_cache_params - eagle_net_inputs["spec_decoding_params"] = spec_decoding_params - eagle_net_inputs["hidden_states"] = hidden_states - return eagle_net_inputs, out_hidden_size_batch_level_starts, num_last_token_indices - - def _slice_hidden_states(self, hidden_states, indices): - hidden_states = index_select(hidden_states, 0, indices) - - hidden_states = hidden_states.view(concat( - [shape(indices, 0), shape(hidden_states, 1)]), - zero_is_placeholder=False) - return hidden_states - - def _eagle_fwd_helper(self, lm_logits, hidden_states, *args, **kwargs): - ''' - EAGLE inference can be viewed as - TRT_Engine(Target -> Draft0 -> Draft1 -> .. -> DraftK-1) -> Runtime -> TRT_Engine(..) -> .. - Target is Base model and Draft is EagleNet. - - Each EagleNet call can be viewed as call to Draft LLM in TensorRT-LLM. - We have to - 1. prepare input tensors before the EagleNet call (like in the the runtime), - 2. call EagleNet, - 3. decode draft tokens after the EagleNet. - The only difference with normal execution of the Draft model is that in EAGLE, - all these 3 things happen inside of the TensorRT engine execution. - We do 1 and 3 inside of the plugins. - For 1. We call eagle_prepare_drafter_inputs_plugin and for 3. eagle_draft_decoder_plugin. - - The first call to the EagleNet (Draft0 == EagleNet0) is the context phase. - For context request we populate the KV cache of the EagleNet. - For generation request that have accepted tokens we emulate KV cache reuse by doing chunked attention, - where chunk is the newly accepted tokens -- all previous tokens are already in the KV cache. - - The following calls to the EagleNet (EagleNetX (X > 0)) are generation phase. - For each EagleNetX we select tokens based on the current path which are going to be used for the generation. - - Let's consider an example: prompt ABCD. EAGLE-1, i.e tree is fixed for the iteration. - Tree: - ┌───┐ - │ 0 │ - └─┬─┘ - ┌─────┴─────┐ - ┌─┴─┐ ┌─┴─┐ ┌─┴─┐ - │ 1 │ │ 2 │ │ 3 │ - └─┬─┘ └─┬─┘ └───┘ - ┌─┴─┐ ┌─┴─┐ - │ 4 │ │ 5 │ - └─┬─┘ └─┬─┘ - ┌─┴─┐ ┌─┴─┐ - │ 6 │ │ 7 │ - └───┘ └───┘ - - First iteration of the TRT engine. Request is context request: - 1. Base model is called for [ABCD] tokens produces token E. - 2. Draft0 is called for tokens [BCDE] and produces - three possibilities F, G and H for positions 1, 2 and 3 respectively. - 3. Since H (position 3) is a leaf, it is not chosen as the input to Draft1 inference. - 4. Draft1 is called for tokens [FG] with appropriate mask of: - |F|G - F|1|0 - G|0|1 - It produces tokens I and J for positions 4 and 5. - 6. Draft2 is called for inputs [FGIJ] with mask of - |F|G|I|J - F|1|0|0|0 - G|0|1|0|0 - I|1|0|1|0 - J|0|1|0|1 - Note that we could've stored FG in KV cache and provide only IJ tokens here - with mask for past KV cache, but it is not supported in TensorRT LLM attention at the moment. - - Draft2 produces tokens K and L at positions 6 and 7. - 7. Resulting outputs are: - 7.1 accepted_ids [E] - 7.2 next_draft_tokens [FGHIJKL] - - Second iteration of the TRT engine. Request is the generation request. - 1. Base model is called for [EFGHIJKL]. Let's assume that it accepts [FIKM], i.e. the left-most path in the tree. - 2. Draft0 is called as context phase for [FIKM] -- to append to kv cache of the existing [BCDE]. - It produces tokens N, O and P for positions 1, 2 and 3. - 3. Draft1 is called as generation phase for [NO] tokens. - etc. - - ''' - input_tree_params = kwargs["tree_params"] - - draft_tokens = kwargs['draft_tokens'] - draft_lens = kwargs['draft_lens'] - eagle_temperature = kwargs['eagle_temperature'] - rand_data_validation = kwargs['rand_data_validation'] - posterior_alpha = kwargs['posterior_alpha'] - posterior_threshold = kwargs['posterior_threshold'] - input_ids = kwargs['input_ids'] - chunked_context_next_tokens = kwargs['chunked_context_next_tokens'] - host_ctx_eagle_net_request_types = kwargs[ - 'host_ctx_eagle_net_request_types'] - host_ctx_eagle_net_context_lengths = kwargs[ - 'host_ctx_eagle_net_context_lengths'] - host_ctx_eagle_net_past_key_value_lengths = kwargs[ - 'host_ctx_eagle_net_past_key_value_lengths'] - host_gen_eagle_net_request_types = kwargs[ - 'host_gen_eagle_net_request_types'] - host_gen_eagle_net_context_lengths = kwargs[ - 'host_gen_eagle_net_context_lengths'] - host_gen_eagle_net_past_key_value_lengths = kwargs[ - 'host_gen_eagle_net_past_key_value_lengths'] - input_gen_tokens = kwargs["input_gen_tokens"] - greedy_sampling = kwargs["greedy_sampling"] - - # Eagle-2 - use_dynamic_tree = kwargs['use_dynamic_tree'] - dynamic_tree_max_topK = kwargs['dynamic_tree_max_topK'] - prev_scores = kwargs['prev_scores'] - current_expand_indices = kwargs['current_expand_indices'] - all_layers_scores = kwargs['all_layers_scores'] - all_layers_draft_token_ids = kwargs['all_layers_draft_token_ids'] - all_layers_draft_token_ids_predecessor = kwargs[ - 'all_layers_draft_token_ids_predecessor'] - - # Sample target tokens and accept them - # next_draft_tokens, next_draft_lens, hidden_size_batch_level_starts are outputted here just to - # reserve the tensor with max size, which eagle_draft_decoder_plugin and - # eagle_prepare_drafter_inputs_plugin are going to directly write to - output = eagle_sample_and_accept_draft_plugin( - lm_logits, draft_tokens, draft_lens, eagle_temperature, - rand_data_validation, posterior_alpha, posterior_threshold, - input_tree_params, greedy_sampling, use_dynamic_tree) - accepted_tokens, num_accepted_tokens, accepted_paths, next_draft_tokens, \ - next_draft_lens, next_draft_paths, hidden_size_batch_level_starts = output - - attention_params = kwargs["attention_params"] - kv_cache_params = kwargs["kv_cache_params"] - spec_decoding_params = kwargs["spec_decoding_params"] - - input_hidden_states = hidden_states - - # Run EAGLE nets - for li in range(self.num_eagle_layers): - # Prepare EAGLE Net inputs. - eagle_net_inputs, hidden_size_batch_level_starts, num_last_token_indices = self._prepare_drafter_inputs( - layer_idx=li, - input_ids=input_ids, - chunked_context_next_tokens=chunked_context_next_tokens, - accepted_token_ids=accepted_tokens, - accepted_lens=num_accepted_tokens, - accepted_path_ids=accepted_paths, - next_draft_tokens=next_draft_tokens, - next_draft_lens=next_draft_lens, - next_draft_paths=next_draft_paths, - prev_draft_lens=draft_lens, - prev_draft_paths=input_tree_params.paths, - input_attention_params=attention_params, - input_kv_cache_params=kv_cache_params, - hidden_states=input_hidden_states, - host_ctx_eagle_net_request_types= - host_ctx_eagle_net_request_types, - host_ctx_eagle_net_context_lengths= - host_ctx_eagle_net_context_lengths, - host_ctx_eagle_net_past_key_value_lengths= - host_ctx_eagle_net_past_key_value_lengths, - host_gen_eagle_net_request_types= - host_gen_eagle_net_request_types, - host_gen_eagle_net_context_lengths= - host_gen_eagle_net_context_lengths, - host_gen_eagle_net_past_key_value_lengths= - host_gen_eagle_net_past_key_value_lengths, - hidden_size_batch_level_starts=hidden_size_batch_level_starts, - input_gen_tokens=input_gen_tokens, - input_spec_decoding_generation_lengths=spec_decoding_params. - spec_decoding_generation_lengths, - spec_decoding_use=spec_decoding_params.spec_decoding_use) - - def single_eagle_net_iter(next_draft_tokens, next_draft_lens, - next_draft_paths, prev_scores, - current_expand_indices, all_layers_scores, - all_layers_draft_token_ids, - all_layers_draft_token_ids_predecessor): - # Run EAGLE Net - # NOTE: handle base net kv cache and eagle net kv cache are in the same tensor. - # EagleNet's kv cache is located starting at numBaseNetHiddenLayers in the kv tensor. - logits, hidden_states, _ = self.eagle_nets[li]( - **eagle_net_inputs) - - # FIXME We need to take top_k_sampling as an input - top_k_sampling = True - - # Decode draft tokens - next_draft_tokens, next_draft_lens, next_draft_paths, prev_scores, current_expand_indices, all_layers_scores, all_layers_draft_token_ids, all_layers_draft_token_ids_predecessor = eagle_draft_decoder_plugin( - li, self.num_eagle_layers, top_k_sampling, logits, - num_last_token_indices, next_draft_paths, use_dynamic_tree, - dynamic_tree_max_topK, next_draft_tokens, next_draft_lens, - prev_scores, current_expand_indices, all_layers_scores, - all_layers_draft_token_ids, - all_layers_draft_token_ids_predecessor) - - return next_draft_tokens, next_draft_lens, hidden_states, next_draft_paths, prev_scores, current_expand_indices, all_layers_scores, all_layers_draft_token_ids, all_layers_draft_token_ids_predecessor - - - next_draft_tokens, next_draft_lens, hidden_states, next_draft_paths, prev_scores, \ - current_expand_indices, all_layers_scores, all_layers_draft_token_ids, all_layers_draft_token_ids_predecessor \ - = single_eagle_net_iter(next_draft_tokens, next_draft_lens, next_draft_paths, \ - prev_scores, current_expand_indices, all_layers_scores, \ - all_layers_draft_token_ids, all_layers_draft_token_ids_predecessor) - - # Update params - if li == 0: - eagle_net_0_sequence_length = eagle_net_inputs[ - "attention_params"].sequence_length - input_hidden_states = hidden_states - else: - input_hidden_states = concat( - [input_hidden_states, hidden_states]) - - kv_cache_params = eagle_net_inputs["kv_cache_params"] - attention_params = eagle_net_inputs["attention_params"] - attention_params.context_lengths = eagle_net_0_sequence_length - attention_params.sequence_length = eagle_net_0_sequence_length - - # Mark tensors as output - accepted_tokens.mark_output('accepted_tokens') - num_accepted_tokens.mark_output('num_accepted_tokens') - accepted_paths.mark_output('accepted_paths') - next_draft_tokens.mark_output('next_draft_tokens') - next_draft_lens.mark_output('next_draft_lens') - next_draft_paths.mark_output('next_draft_paths') - - return next_draft_tokens - - def forward(self, *args, **kwargs): - extra_args = [ - "draft_tokens", "draft_lens", "eagle_temperature", - "rand_data_validation", "tree_params", - "host_ctx_eagle_net_request_types", - "host_ctx_eagle_net_context_lengths", - "host_ctx_eagle_net_past_key_value_lengths", - "host_gen_eagle_net_request_types", - "host_gen_eagle_net_context_lengths", - "host_gen_eagle_net_past_key_value_lengths", "input_gen_tokens", - "chunked_context_next_tokens", "posterior_alpha", - "posterior_threshold", "greedy_sampling", "use_dynamic_tree", - "dynamic_tree_max_topK", "prev_scores", "current_expand_indices", - "all_layers_scores", "all_layers_draft_token_ids", - "all_layers_draft_token_ids_predecessor" - ] - - base_kwargs = {k: v for k, v in kwargs.items() if k not in extra_args} - - # Base model forward - hidden_states = super().forward(*args, **base_kwargs) - - if self.mapping.is_last_pp_rank(): - extra_args = ["hidden_states"] - kwargs = {k: v for k, v in kwargs.items() if k not in extra_args} - - assert kwargs['use_cache'] and default_net( - ).plugin_config.paged_kv_cache - - if self.mapping.is_last_pp_rank(): - lm_logits, hidden_states, all_hidden_states = hidden_states - lm_logits = cast(lm_logits, self.config.logits_dtype) - # Call eagle logic to accept prev draft tokens and predict next draft tokens - next_draft_tokens = self._eagle_fwd_helper(lm_logits, - all_hidden_states, *args, - **kwargs) - else: - hidden_states.mark_output('hidden_states_output', self.config.dtype) - - if self.mapping.is_last_pp_rank(): - return next_draft_tokens - return hidden_states - - def prepare_inputs(self, *args, **kwargs): - """ - Inputs needed: - device_request_types: [bs] - draft_tokens: [bs, max_draft_len] - draft_lens: [bs] - spec_decoding_generation_lengths: [bs] - spec_decoding_position_offsets: [bs, max_gen_tokens] - spec_decoding_packed_mask: [bs, max_draft_len, packed_length] ** - eagle_temperature: [bs] - rand_data_validation: [bs, max_draft_len] - - ** The mask is tricky since the boolean mask will need to be - packed in runtime. So, the last dim will be: - packed_length = ceil((max_draft_len+1)/32) - """ - default_range = GenerationMixin.default_range - remove_input_padding = default_net().plugin_config.remove_input_padding - use_gpt_attention_plugin = default_net( - ).plugin_config.gpt_attention_plugin - use_gemm_plugin = default_net().plugin_config.gemm_plugin - paged_kv_cache = default_net().plugin_config.paged_kv_cache - multiple_profiles = default_net().plugin_config.multiple_profiles - max_batch_size = kwargs['max_batch_size'] - assert max_batch_size is not None - gt_range = default_range(max_batch_size * (self.max_draft_len + 1), - min_range=0, - opt_offset=1) - - kwargs['speculative_decoding_draft_tokens_external'] = False - kwargs['max_draft_len'] = self.max_draft_len - kwargs['spec_decoding_is_generation_length_variable'] = True - kwargs[ - 'num_hidden_layers'] = self.config.num_hidden_layers + self.config.eagle_net_config.num_hidden_layers - - # Call base class prepare inputs - inputs = super().prepare_inputs(*args, **kwargs) - - assert inputs['spec_decoding_params'] is not None - - kv_cache_type = KVCacheType.PAGED if paged_kv_cache else KVCacheType.CONTINUOUS - enable_ctx_gen_opt_profiles = GenerationMixin.has_ctx_gen_opt_profiles( - use_gpt_attention_plugin=use_gpt_attention_plugin, - use_gemm_plugin=use_gemm_plugin, - remove_input_padding=remove_input_padding, - kv_cache_type=kv_cache_type) - - num_profiles, ranges = GenerationMixin.get_profiles_ranges( - max_batch_size=max_batch_size, - max_beam_width=kwargs['max_beam_width'], - max_input_len=kwargs['max_input_len'], - max_num_tokens=kwargs['max_num_tokens'], - max_draft_len=self.max_draft_len, - opt_batch_size=None - if 'opt_batch_size' not in kwargs else kwargs['opt_batch_size'], - opt_num_tokens=None - if 'opt_num_tokens' not in kwargs else kwargs['opt_num_tokens'], - enable_ctx_gen_opt_profiles=enable_ctx_gen_opt_profiles, - multiple_profiles=multiple_profiles, - kv_cache_type=kv_cache_type) - - bb_range = ranges['bb_range'] - - draft_len_range = [self.max_draft_len for _ in range(len(bb_range))] - decoding_len_range = [(self.max_draft_len + 1) - for _ in range(len(bb_range))] - path_len_range = [(self.num_eagle_layers + 1) - for _ in range(len(bb_range))] - gen_tokens_range = [gt_range for _ in range(len(bb_range))] - num_eagle_layers_range = [ - self.num_eagle_layers for _ in range(len(bb_range)) - ] - draft_len_square_range = [(self.max_draft_len * self.max_draft_len) - for _ in range(len(bb_range))] - - draft_tokens = Tensor(name='draft_tokens', - dtype=trt.int32, - shape=[-1, self.max_draft_len], - dim_range=OrderedDict([ - ('batch_size', bb_range), - ('draft_len', draft_len_range), - ])) - draft_lens = Tensor(name='draft_lens', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([ - ('batch_size', bb_range), - ])) - eagle_temperature = Tensor(name='eagle_temperature', - dtype=trt.float32, - shape=[-1], - dim_range=OrderedDict([ - ("batch_size", bb_range), - ])) - rand_data_validation = Tensor(name='rand_data_validation', - dtype=trt.float32, - shape=[-1, self.max_draft_len + 1], - dim_range=OrderedDict([ - ('batch_size', bb_range), - ('decoding_len', decoding_len_range), - ])) - posterior_alpha = Tensor(name='posterior_alpha', - dtype=trt.float32, - shape=[-1], - dim_range=OrderedDict([ - ("batch_size", bb_range), - ])) - posterior_threshold = Tensor(name='posterior_threshold', - dtype=trt.float32, - shape=[-1], - dim_range=OrderedDict([ - ("batch_size", bb_range), - ])) - draft_paths = Tensor( - name='draft_paths', - dtype=trt.int32, - shape=[-1, self.max_draft_len + 1, self.num_eagle_layers + 1], - dim_range=OrderedDict([ - ('batch_size', bb_range), - ('decoding_len', decoding_len_range), - ('path_len', path_len_range), - ])) - - host_ctx_eagle_net_request_types = Tensor( - name='host_ctx_eagle_net_request_types', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([ - ('batch_size', bb_range), - ])) - host_ctx_eagle_net_context_lengths = Tensor( - name='host_ctx_eagle_net_context_lengths', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([ - ('batch_size', bb_range), - ])) - host_ctx_eagle_net_past_key_value_lengths = Tensor( - name='host_ctx_eagle_net_past_key_value_lengths', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([ - ('batch_size', bb_range), - ])) - host_gen_eagle_net_request_types = Tensor( - name='host_gen_eagle_net_request_types', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([ - ('batch_size', bb_range), - ])) - host_gen_eagle_net_context_lengths = Tensor( - name='host_gen_eagle_net_context_lengths', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([ - ('batch_size', bb_range), - ])) - host_gen_eagle_net_past_key_value_lengths = Tensor( - name='host_gen_eagle_net_past_key_value_lengths', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([ - ('batch_size', bb_range), - ])) - - input_gen_tokens = Tensor(name='input_gen_tokens', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([ - ('gen_tokens', gen_tokens_range), - ])) - chunked_context_next_tokens = Tensor(name='chunked_context_next_tokens', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([ - ('batch_size', bb_range), - ])) - greedy_sampling = Tensor(name='greedy_sampling', - dtype=trt.int32, - shape=[1]) - - use_dynamic_tree = Tensor(name='use_dynamic_tree', - dtype=trt.int32, - shape=[1]) - - dynamic_tree_max_topK = Tensor(name='dynamic_tree_max_topK', - dtype=trt.int32, - shape=[1]) - - prev_scores = Tensor(name='prev_scores', - dtype=trt.float32, - shape=[-1, self.max_draft_len], - dim_range=OrderedDict([ - ('batch_size', bb_range), - ('draft_len', draft_len_range), - ])) - - current_expand_indices = Tensor(name='current_expand_indices', - dtype=trt.int32, - shape=[-1, self.max_draft_len], - dim_range=OrderedDict([ - ('batch_size', bb_range), - ('draft_len', draft_len_range), - ])) - - all_layers_scores = Tensor(name='all_layers_scores', - dtype=trt.float32, - shape=[ - -1, self.num_eagle_layers, - self.max_draft_len * self.max_draft_len - ], - dim_range=OrderedDict([ - ('batch_size', bb_range), - ('num_eagle_layers', - num_eagle_layers_range), - ('draft_len_square', - draft_len_square_range), - ])) - - all_layers_draft_token_ids = Tensor( - name='all_layers_draft_token_ids', - dtype=trt.int32, - shape=[ - -1, self.num_eagle_layers, - self.max_draft_len * self.max_draft_len - ], - dim_range=OrderedDict([ - ('batch_size', bb_range), - ('num_eagle_layers', num_eagle_layers_range), - ('draft_len_square', draft_len_square_range), - ])) - - all_layers_draft_token_ids_predecessor = Tensor( - name='all_layers_draft_token_ids_predecessor', - dtype=trt.int32, - shape=[ - -1, self.num_eagle_layers, - self.max_draft_len * self.max_draft_len - ], - dim_range=OrderedDict([ - ('batch_size', bb_range), - ('num_eagle_layers', num_eagle_layers_range), - ('draft_len_square', draft_len_square_range), - ])) - - tree_params = TreeParams(paths=draft_paths) - - inputs['draft_tokens'] = draft_tokens - inputs['draft_lens'] = draft_lens - inputs['eagle_temperature'] = eagle_temperature - inputs['posterior_alpha'] = posterior_alpha - inputs['posterior_threshold'] = posterior_threshold - inputs['rand_data_validation'] = rand_data_validation - inputs['tree_params'] = tree_params - inputs[ - 'host_ctx_eagle_net_request_types'] = host_ctx_eagle_net_request_types - inputs[ - 'host_ctx_eagle_net_context_lengths'] = host_ctx_eagle_net_context_lengths - inputs[ - 'host_ctx_eagle_net_past_key_value_lengths'] = host_ctx_eagle_net_past_key_value_lengths - inputs[ - 'host_gen_eagle_net_request_types'] = host_gen_eagle_net_request_types - inputs[ - 'host_gen_eagle_net_context_lengths'] = host_gen_eagle_net_context_lengths - inputs[ - 'host_gen_eagle_net_past_key_value_lengths'] = host_gen_eagle_net_past_key_value_lengths - inputs['input_gen_tokens'] = input_gen_tokens - inputs['chunked_context_next_tokens'] = chunked_context_next_tokens - inputs['greedy_sampling'] = greedy_sampling - inputs['use_dynamic_tree'] = use_dynamic_tree - inputs['dynamic_tree_max_topK'] = dynamic_tree_max_topK - inputs['prev_scores'] = prev_scores - inputs['current_expand_indices'] = current_expand_indices - inputs['all_layers_scores'] = all_layers_scores - inputs['all_layers_draft_token_ids'] = all_layers_draft_token_ids - inputs[ - 'all_layers_draft_token_ids_predecessor'] = all_layers_draft_token_ids_predecessor - - return inputs - - @classmethod - def from_hugging_face( - cls, - hf_model_or_dir: Union[str, 'transformers.PreTrainedModel'], - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - assert hf_model_or_dir is not None - speculative_model_dir = kwargs.get('speculative_model_dir', None) - tllm_config = EagleConfig.from_hugging_face(hf_model_or_dir, - dtype=dtype, - mapping=mapping, - quant_config=quant_config, - **kwargs) - - # for rank in range(mapping.world_size): - tllm_config.mapping = Mapping(world_size=mapping.world_size, - rank=mapping.rank, - cp_size=1, - tp_size=mapping.tp_size, - pp_size=mapping.pp_size) - - model = EagleForCausalLM(tllm_config) - - def check_and_update(module, dict): - if hasattr(module, 'tllm_to_externel_key_dict'): - module.tllm_to_externel_key_dict.update(dict) - else: - module.tllm_to_externel_key_dict = dict - - def copy(tensors): - if isinstance(tensors, list): - if None in tensors: - return tensors - else: - return [tensor.clone() for tensor in tensors] - elif tensors is None: - return tensors - else: - return tensors.clone() - - shared_weight_prefixs = [] - tllm_weights = {} - customized_dict = {"drafter": ""} - if speculative_model_dir is None: - # Single checkpoint for ModelOpt - for idx, eagle_net in enumerate(model.eagle_nets): - check_and_update(eagle_net.drafter.fc, {"fc": "fc"}) - check_and_update(eagle_net.drafter.vocab_embedding, - {f"eagle_nets.{idx}": "model"}) - check_and_update(eagle_net.lm_head, {f"eagle_nets.{idx}": ""}) - shared_weight_prefixs.append(f"eagle_nets.{idx}") - customized_dict[f'eagle_nets.{idx}'] = 'eagle_module' - loader = ModelWeightsLoader(speculative_model_dir, customized_dict) - loader.update_key_mapping(model) - for tllm_key, _ in tqdm(model.named_parameters()): - if any([ - tllm_key.startswith(prefix) - for prefix in shared_weight_prefixs - ]): - tllm_weights.update(loader.load(tllm_key, preprocess=copy)) - else: - tllm_weights.update(loader.load(tllm_key)) - loader.fill(tllm_weights) - else: - # Double checkpoint for HF - for idx, eagle_net in enumerate(model.eagle_nets): - check_and_update(eagle_net.drafter.fc, {"fc": "fc"}) - check_and_update(eagle_net.drafter.vocab_embedding, - {f"eagle_nets.{idx}": ""}) - check_and_update(eagle_net.lm_head, {f"eagle_nets.{idx}": ""}) - shared_weight_prefixs.append(f"eagle_nets.{idx}") - customized_dict[f'eagle_nets.{idx}'] = '' - - # Load base model - base_loader = ModelWeightsLoader(hf_model_or_dir) - base_loader.update_key_mapping(model) - for tllm_key, _ in tqdm(model.transformer.named_parameters()): - tllm_weights.update(base_loader.load("transformer." + tllm_key)) - tllm_weights.update(base_loader.load("lm_head.weight")) - # for idx in range(args.num_eagle_layers): - for idx in range(4): - tllm_weights.update( - base_loader.load(f"eagle_nets.{idx}.lm_head.weight", - preprocess=copy)) - - # Load eagle model - eagle_loader = ModelWeightsLoader(str(speculative_model_dir), - customized_dict) - eagle_loader.update_key_mapping(model) - for tllm_key, _ in tqdm(model.eagle_nets.named_parameters()): - if not tllm_key.endswith("lm_head.weight"): - if any([ - tllm_key.startswith(prefix) - for prefix in shared_weight_prefixs - ]): - tllm_weights.update( - eagle_loader.load("eagle_nets." + tllm_key, - preprocess=copy)) - else: - tllm_weights.update( - eagle_loader.load("eagle_nets." + tllm_key)) - base_loader.fill(tllm_weights) - - return model diff --git a/tensorrt_llm/models/enc_dec/__init__.py b/tensorrt_llm/models/enc_dec/__init__.py deleted file mode 100644 index 71bf6d298d42..000000000000 --- a/tensorrt_llm/models/enc_dec/__init__.py +++ /dev/null @@ -1,14 +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. diff --git a/tensorrt_llm/models/enc_dec/model.py b/tensorrt_llm/models/enc_dec/model.py deleted file mode 100644 index be3c5afc49f6..000000000000 --- a/tensorrt_llm/models/enc_dec/model.py +++ /dev/null @@ -1,2195 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -import math -from collections import OrderedDict -from typing import List, Optional - -import tensorrt as trt -import torch - -from tensorrt_llm._common import default_net -from tensorrt_llm._utils import (numpy_to_torch, pad_vocab_size, - str_dtype_to_torch) -from tensorrt_llm.functional import (LayerNormPositionType, LayerNormType, - MLPType, PositionEmbeddingType, Tensor, - assertion, cast, gather_last_token_logits, - gelu, maximum, minimum, recv, send, shape, - transpose, unsqueeze) -# yapf: disable -from tensorrt_llm.layers import (MLP, Attention, AttentionMaskParams, - AttentionMaskType, AttentionParams, - BertAttention, ColumnLinear, Conv1d, Embedding, - FusedGatedMLP, GatedMLP, GroupNorm, - KeyValueCacheParams, LanguageAdapter, - LanguageAdapterConfig, LayerNorm, LoraParams, - PromptTuningEmbedding, RmsNorm) -# yapf: enable -from tensorrt_llm.lora_helper import (LoraConfig, - get_default_trtllm_modules_to_hf_modules, - use_lora) -from tensorrt_llm.mapping import Mapping -from tensorrt_llm.models.modeling_utils import PretrainedConfig, PretrainedModel -from tensorrt_llm.module import Module, ModuleList -from tensorrt_llm.parameter import Parameter -from tensorrt_llm.plugin.plugin import current_all_reduce_helper -from tensorrt_llm.quantization import QuantMode - -layernorm_map = { - LayerNormType.LayerNorm: LayerNorm, - LayerNormType.RmsNorm: RmsNorm, - LayerNormType.GroupNorm: GroupNorm, -} - -mlp_map = { - MLPType.MLP: MLP, - MLPType.GatedMLP: GatedMLP, - MLPType.FusedGatedMLP: FusedGatedMLP, -} - - -class EncDecTransformer(Module): - - def __init__(self, - vocab_size, - hidden_size, - has_vocab_embeddings=True, - max_position_embeddings=None, - has_position_embedding=False, - type_vocab_size=None, - has_embedding_layernorm=False, - has_embedding_scale=False, - layernorm_eps=1e-5, - layernorm_type=LayerNormType.LayerNorm, - dtype=None, - use_parallel_embedding=False, - embedding_sharding_dim=0, - mapping=Mapping(), - has_model_final_layernorm=False, - norm_epsilon=1e-05, - is_decoder=False): - super().__init__() - - self.layernorm_type = layernorm_type - ln_type = layernorm_map[layernorm_type] - self.mapping = mapping - - if self.mapping.is_first_pp_rank(): - if has_vocab_embeddings: - self.vocab_embedding = Embedding( - vocab_size, - hidden_size, - dtype=dtype, - tp_size=mapping.tp_size if use_parallel_embedding else 1, - tp_group=mapping.tp_group - if use_parallel_embedding else None, - sharding_dim=embedding_sharding_dim, - tp_rank=mapping.tp_rank) - - self.position_embedding = None - self.max_position_embeddings = max_position_embeddings - if has_position_embedding: - self.position_embedding = Embedding( - max_position_embeddings, - hidden_size, - dtype=dtype, - tp_size=mapping.tp_size if use_parallel_embedding else 1, - tp_group=mapping.tp_group - if use_parallel_embedding else None, - sharding_dim=embedding_sharding_dim, - tp_rank=mapping.tp_rank) - - self.token_type_embedding = None - if type_vocab_size: - self.token_type_embedding = Embedding( - type_vocab_size, - hidden_size, - dtype=dtype, - tp_size=mapping.tp_size if use_parallel_embedding else 1, - tp_group=mapping.tp_group - if use_parallel_embedding else None, - sharding_dim=embedding_sharding_dim, - tp_rank=mapping.tp_rank) - - # e.g. BART true, T5 false - self.ln_embed = None - if has_embedding_layernorm: - self.ln_embed = ln_type(normalized_shape=hidden_size, - eps=layernorm_eps, - dtype=dtype) - - # e.g. BART true, T5 false - self.embedding_scale = 1.0 - if has_embedding_scale: - self.embedding_scale = math.sqrt(hidden_size) - - # Note: embedding offset in BART is not considered as a standard. For the specific case, - # we just need to shrink its position embedding table by [offset:] during weight loading - - self.layers = None - self.is_decoder = is_decoder - self.has_model_final_layernorm = has_model_final_layernorm - if self.mapping.is_last_pp_rank(): - if self.has_model_final_layernorm: - self.ln_f = ln_type(normalized_shape=hidden_size, - eps=norm_epsilon, - dtype=dtype) - - def assign_module(self, module, module_name): - setattr(self, module_name, module) - - def embedding(self, - input_ids, - position_ids=None, - token_type_ids=None, - prompt_embedding_table=None, - prompt_tasks=None, - prompt_vocab_size=None): - # position_ids and token_type_ids are provided inputs - # and should not be formulated deterministically - - args = [prompt_embedding_table, prompt_tasks, prompt_vocab_size - ] if prompt_embedding_table is not None else [] - - x = self.vocab_embedding(input_ids, *args) * self.embedding_scale - self.register_network_output('word_embeddings', x) - - if self.position_embedding: - pos_emb = self.position_embedding(position_ids) - self.register_network_output('position_embeddings', pos_emb) - x = x + pos_emb - if self.token_type_embedding: - x = x + self.token_type_embedding(token_type_ids) - - if self.ln_embed: - x = self.ln_embed(x) - - return x - - -class EncoderLayer(Module): - - def __init__(self, - hidden_size, - ffn_hidden_size, - num_attention_heads, - num_kv_heads, - head_size, - max_position_embeddings=None, - q_scaling=1.0, - has_attention_qkvo_bias=False, - has_mlp_bias=False, - layernorm_position=LayerNormPositionType.pre_layernorm, - layernorm_type=LayerNormType.LayerNorm, - layernorm_eps=1e-5, - hidden_act="relu", - mlp_type=MLPType.MLP, - mapping=Mapping(), - dtype=None, - residual_scaling=1.0, - relative_attention=False, - max_distance=0, - num_buckets=0, - fp16_clamping=False, - quant_mode=QuantMode(0), - language_adapter_config: LanguageAdapterConfig = None): - super().__init__() - - # e.g. BART regular, T5 RMS - self.layernorm_type = layernorm_type - ln_type = layernorm_map[layernorm_type] - - # e.g. BART post, T5 pre - self.layernorm_position = layernorm_position - - # e.g. BART q_scaling = 1.f, T5 q_scaling = 1.f/sqrt(head_size) - self.attention = BertAttention( - hidden_size, - num_attention_heads, - attention_head_size=head_size, - num_kv_heads=num_kv_heads, - max_position_embeddings=max_position_embeddings, - q_scaling=q_scaling, - bias=has_attention_qkvo_bias, - tp_group=mapping.tp_group, - tp_size=mapping.tp_size, - tp_rank=mapping.tp_rank, - dtype=dtype, - relative_attention=relative_attention, - max_distance=max_distance, - num_buckets=num_buckets, - quant_mode=quant_mode) - - self.attention_layernorm = ln_type(normalized_shape=hidden_size, - eps=layernorm_eps, - dtype=dtype) - - # T5/BART MLP, Flan-T5 GatedMLP - self.mlp_type = mlp_type - mlp_f = mlp_map[mlp_type] - self.mlp = mlp_f( - hidden_size=hidden_size, - ffn_hidden_size=ffn_hidden_size, - hidden_act=hidden_act, - bias=has_mlp_bias, - tp_group=mapping.tp_group, - tp_size=mapping.tp_size, - dtype=dtype, - quant_mode=quant_mode, - ) - - self.mlp_layernorm = ln_type(normalized_shape=hidden_size, - eps=layernorm_eps, - dtype=dtype) - - self.residual_scaling = residual_scaling - - # T5-series model(e.g. t5-large, t5-3b, flan-t5-small) has accuracy issue due to fp16 overflow - # after residual add. We add workaround for clamping fp16 range [-64000, 64000] after every - # residual add to avoid accuracy drop. - self.fp16_clamping = fp16_clamping - - self.adapter = None - self.adapter_layer_norm = None - - if language_adapter_config: - self.adapter_layer_norm = ln_type(normalized_shape=hidden_size, - eps=layernorm_eps, - dtype=dtype) - self.adapter = LanguageAdapter( - language_adapter_config=language_adapter_config, - hidden_size=hidden_size, - hidden_act=hidden_act, - mapping=mapping, - has_mlp_bias=has_mlp_bias, - dtype=dtype, - quant_mode=quant_mode) - - def forward(self, - hidden_states: Tensor, - attention_mask=None, - input_lengths=None, - max_input_length=None, - lora_layer_params=None, - language_adapter_routings: Optional[Tensor] = None): - assert isinstance(hidden_states, Tensor) - - # self attention - residual = hidden_states * self.residual_scaling - - if self.layernorm_position == LayerNormPositionType.pre_layernorm: - hidden_states = self.attention_layernorm(hidden_states) - - attention_output = self.attention(hidden_states, - attention_mask=attention_mask, - input_lengths=input_lengths, - max_input_length=max_input_length, - lora_layer_params=lora_layer_params) - - self.register_network_output('attention_output', attention_output) - - hidden_states = residual + attention_output - - if self.fp16_clamping: - hidden_states = maximum(-64000.0, hidden_states) - hidden_states = minimum(64000.0, hidden_states) - - if self.layernorm_position == LayerNormPositionType.post_layernorm: - hidden_states = self.attention_layernorm(hidden_states) - - # MLP - residual = hidden_states * self.residual_scaling - - if self.layernorm_position == LayerNormPositionType.pre_layernorm: - hidden_states = self.mlp_layernorm(hidden_states) - - hidden_states = self.mlp(hidden_states, - lora_layer_params=lora_layer_params) - - self.register_network_output('mlp_output', hidden_states) - - hidden_states = residual + hidden_states - - if self.fp16_clamping: - hidden_states = maximum(-64000.0, hidden_states) - hidden_states = minimum(64000.0, hidden_states) - - if self.layernorm_position == LayerNormPositionType.post_layernorm: - hidden_states = self.mlp_layernorm(hidden_states) - - # MT Specific: adapters - if self.adapter: - residual = hidden_states - hidden_states = self.adapter_layer_norm(hidden_states) - hidden_states = self.adapter.layers( - hidden_states, static_routing_input=language_adapter_routings) - hidden_states = residual + hidden_states - if self.fp16_clamping: - hidden_states = maximum(-64000.0, hidden_states) - hidden_states = minimum(64000.0, hidden_states) - - return hidden_states - - -class DecoderLayer(Module): - - def __init__(self, - *, - local_layer_idx, - hidden_size, - ffn_hidden_size, - num_attention_heads, - num_kv_heads, - head_size, - max_position_embeddings=None, - q_scaling=1.0, - has_attention_qkvo_bias=False, - has_mlp_bias=False, - layernorm_position=LayerNormPositionType.pre_layernorm, - layernorm_type=LayerNormType.LayerNorm, - layernorm_eps=1e-5, - hidden_act="relu", - mlp_type=MLPType.MLP, - mapping=Mapping(), - dtype=None, - residual_scaling=1.0, - relative_attention=False, - max_distance=0, - num_buckets=0, - fp16_clamping=False, - skip_cross_kv=False, - use_implicit_relative_attention=False, - quant_mode=QuantMode(0), - language_adapter_config: LanguageAdapterConfig = None): - super().__init__() - - # e.g. BART regular, T5 RMS - self.layernorm_type = layernorm_type - ln_type = layernorm_map[layernorm_type] - - # e.g. BART post, T5 pre - self.layernorm_position = layernorm_position - - # e.g. BART q_scaling = 1.f, T5 q_scaling = 1.f/sqrt(head_size) - self.self_attention = Attention( - local_layer_idx=local_layer_idx, - hidden_size=hidden_size, - num_attention_heads=num_attention_heads, - attention_head_size=head_size, - num_kv_heads=num_kv_heads, - max_position_embeddings=max_position_embeddings, - q_scaling=q_scaling, - bias=has_attention_qkvo_bias, - attention_mask_type=AttentionMaskType.causal, - tp_group=mapping.tp_group, - tp_size=mapping.tp_size, - tp_rank=mapping.tp_rank, - dtype=dtype, - cross_attention=False, - relative_attention=relative_attention, - max_distance=max_distance if use_implicit_relative_attention else 0, - num_buckets=num_buckets, - position_embedding_type=PositionEmbeddingType.relative - if relative_attention else PositionEmbeddingType.learned_absolute, - use_implicit_relative_attention=use_implicit_relative_attention, - quant_mode=quant_mode) - - self.self_attention_layernorm = ln_type(normalized_shape=hidden_size, - eps=layernorm_eps, - dtype=dtype) - - # Note: self attn uses MMHA, mask is always causal triangular - # cross attn has two scenarios: - # - in context phase, all ones mask, same as padding type - # - in generation phase, same causal triangular mask as MMHA - # - context phase special handling is done in plugin by resetting mask type - # - # e.g. BART q_scaling = 1.f, T5 q_scaling = 1.f/sqrt(head_size) - self.cross_attention = Attention( - local_layer_idx=local_layer_idx, - hidden_size=hidden_size, - num_attention_heads=num_attention_heads, - attention_head_size=head_size, - num_kv_heads=num_kv_heads, - max_position_embeddings=max_position_embeddings, - q_scaling=q_scaling, - bias=has_attention_qkvo_bias, - attention_mask_type=AttentionMaskType.causal, - tp_group=mapping.tp_group, - tp_size=mapping.tp_size, - tp_rank=mapping.tp_rank, - dtype=dtype, - cross_attention=True, - relative_attention= - False, # Cross attention has no relative attention bias - max_distance=max_distance, - num_buckets=num_buckets, - position_embedding_type=PositionEmbeddingType.learned_absolute, - skip_cross_kv=skip_cross_kv, - quant_mode=quant_mode) - - self.cross_attention_layernorm = ln_type(normalized_shape=hidden_size, - eps=layernorm_eps, - dtype=dtype) - - # T5/BART MLP, Flan-T5 GatedMLP - self.mlp_type = mlp_type - mlp_f = mlp_map[mlp_type] - self.mlp = mlp_f( - hidden_size=hidden_size, - ffn_hidden_size=ffn_hidden_size, - hidden_act=hidden_act, - bias=has_mlp_bias, - tp_group=mapping.tp_group, - tp_size=mapping.tp_size, - dtype=dtype, - quant_mode=quant_mode, - ) - - self.mlp_layernorm = ln_type(normalized_shape=hidden_size, - eps=layernorm_eps, - dtype=dtype) - - self.residual_scaling = residual_scaling - - # T5-series model(e.g. t5-large, t5-3b, flan-t5-small) has accuracy issue due to fp16 overflow - # after residual add. We add workaround for clamping fp16 range [-64000, 64000] after every - # residual add to avoid accuracy drop. - self.fp16_clamping = fp16_clamping - - self.adapter = None - self.adapter_layer_norm = None - - if language_adapter_config: - self.adapter_layer_norm = ln_type(normalized_shape=hidden_size, - eps=layernorm_eps, - dtype=dtype) - self.adapter = LanguageAdapter( - language_adapter_config=language_adapter_config, - hidden_size=hidden_size, - hidden_act=hidden_act, - mapping=mapping, - has_mlp_bias=has_mlp_bias, - dtype=dtype, - quant_mode=quant_mode) - - def forward(self, - hidden_states: Tensor, - encoder_output: Optional[Tensor] = None, - attention_mask_params=None, - use_cache=False, - kv_cache_params=None, - attention_params=None, - lora_layer_params=None, - cross_kv_cache_gen: Optional[Tensor] = None, - cross_kv_reuse: Optional[Tensor] = None, - language_adapter_routings: Optional[Tensor] = None): - assert isinstance(hidden_states, Tensor) - - if encoder_output: - assert isinstance(encoder_output, Tensor) - - # self-attention - residual = hidden_states * self.residual_scaling - - if self.layernorm_position == LayerNormPositionType.pre_layernorm: - hidden_states = self.self_attention_layernorm(hidden_states) - - attention_output = self.self_attention( - hidden_states=hidden_states, - attention_mask=attention_mask_params.self_attention_mask, - use_cache=use_cache, - kv_cache_params=kv_cache_params, - attention_params=attention_params, - lora_layer_params=lora_layer_params) - - if use_cache: - attention_output, presents_self = attention_output - - self.register_network_output('self_attention_output', attention_output) - - hidden_states = residual + attention_output - - if self.fp16_clamping: - hidden_states = maximum(-64000.0, hidden_states) - hidden_states = minimum(64000.0, hidden_states) - - if self.layernorm_position == LayerNormPositionType.post_layernorm: - hidden_states = self.self_attention_layernorm(hidden_states) - - # cross attention - residual = hidden_states * self.residual_scaling - - if self.layernorm_position == LayerNormPositionType.pre_layernorm: - hidden_states = self.cross_attention_layernorm(hidden_states) - - attention_output = self.cross_attention( - hidden_states=hidden_states, - attention_mask=attention_mask_params.cross_attention_mask, - attention_packed_mask=attention_mask_params. - cross_attention_packed_mask, - encoder_output=encoder_output, - use_cache=use_cache, - kv_cache_params=kv_cache_params, - attention_params=attention_params, - lora_layer_params=lora_layer_params, - cross_kv_cache_gen=cross_kv_cache_gen, - cross_kv_reuse=cross_kv_reuse) - - if use_cache: - attention_output, presents_cross = attention_output - - self.register_network_output('cross_attention_output', attention_output) - - hidden_states = residual + attention_output - - if self.fp16_clamping: - hidden_states = maximum(-64000.0, hidden_states) - hidden_states = minimum(64000.0, hidden_states) - - if self.layernorm_position == LayerNormPositionType.post_layernorm: - hidden_states = self.cross_attention_layernorm(hidden_states) - - # MLP - residual = hidden_states * self.residual_scaling - - if self.layernorm_position == LayerNormPositionType.pre_layernorm: - hidden_states = self.mlp_layernorm(hidden_states) - - hidden_states = self.mlp(hidden_states, - lora_layer_params=lora_layer_params) - self.register_network_output('mlp_output', hidden_states) - - hidden_states = residual + hidden_states - - if self.fp16_clamping: - hidden_states = maximum(-64000.0, hidden_states) - hidden_states = minimum(64000.0, hidden_states) - - if self.layernorm_position == LayerNormPositionType.post_layernorm: - hidden_states = self.mlp_layernorm(hidden_states) - - # MT Specific: adapters - if self.adapter: - residual = hidden_states - hidden_states = self.adapter_layer_norm(hidden_states) - hidden_states = self.adapter.layers( - hidden_states, static_routing_input=language_adapter_routings) - hidden_states = residual + hidden_states - if self.fp16_clamping: - hidden_states = maximum(-64000.0, hidden_states) - hidden_states = minimum(64000.0, hidden_states) - - if use_cache: - return (hidden_states, presents_self, presents_cross) - return hidden_states - - -class EncoderModel(PretrainedModel): - - def __init__(self, config: PretrainedConfig): - self.check_config(config) - super().__init__(config) - self.mapping = self.config.mapping - - self.has_position_embedding = self.config.has_position_embedding - type_vocab_size = self.config.type_vocab_size - self.has_token_type_embedding = False if type_vocab_size is None else True - - # e.g. BART regular, T5 RMS - self.layernorm_type = self.config.layernorm_type - - # e.g. BART true, T5 false - self.has_attention_qkvo_bias = self.config.has_attention_qkvo_bias - self.has_mlp_bias = self.config.has_mlp_bias - - # e.g. BART false, T5 true - self.has_model_final_layernorm = self.config.has_model_final_layernorm - - self._dtype = self.config.dtype - - self.total_num_layers = self.config.num_hidden_layers - self.num_layers = self.config.num_hidden_layers // self.mapping.pp_size - - self.hidden_size = self.config.hidden_size - self.num_heads = self.config.num_attention_heads - num_kv_heads = self.num_heads - if num_kv_heads is None or num_kv_heads <= 0: - num_kv_heads = self.config.num_attention_heads - self.num_kv_heads = num_kv_heads - self.head_size = self.hidden_size // self.num_heads if self.config.head_size is None else self.config.head_size - - self.fp16_clamping = (self.config.dtype - == 'float16') and (self.config.model_type == 't5') - self.mlp_type = MLPType.MLP if not hasattr( - self.config, "mlp_type") else self.config.mlp_type - self.language_adapter_config = None if not hasattr( - self.config, - 'language_adapter_config') else LanguageAdapterConfig.from_dict( - self.config.language_adapter_config) - - self.transformer = EncDecTransformer( - self.config.vocab_size, - self.config.hidden_size, - max_position_embeddings=self.config.max_position_embeddings, - has_position_embedding=self.has_position_embedding, - type_vocab_size=type_vocab_size, - has_embedding_layernorm=self.config.has_embedding_layernorm, - has_embedding_scale=self.config.has_embedding_scale, - layernorm_eps=self.config.norm_epsilon, - layernorm_type=self.layernorm_type, - dtype=self.config.dtype, - use_parallel_embedding=self.config.use_parallel_embedding, - embedding_sharding_dim=self.config.embedding_sharding_dim, - mapping=self.mapping, - has_model_final_layernorm=self.has_model_final_layernorm, - norm_epsilon=self.config.norm_epsilon, - is_decoder=False) - - encoder_layers = ModuleList([ - EncoderLayer( - hidden_size=self.hidden_size, - ffn_hidden_size=self.config.intermediate_size, - num_attention_heads=self.num_heads, - num_kv_heads=num_kv_heads, - head_size=self.head_size, - max_position_embeddings=self.config.max_position_embeddings, - q_scaling=self.config.q_scaling, - has_attention_qkvo_bias=self.has_attention_qkvo_bias, - has_mlp_bias=self.has_mlp_bias, - layernorm_position=self.config.layernorm_position, - layernorm_eps=self.config.norm_epsilon, - layernorm_type=self.layernorm_type, - hidden_act=self.config.hidden_act, - mlp_type=self.mlp_type, - mapping=self.mapping, - dtype=self.config.dtype, - residual_scaling=1.0 - if not hasattr(self.config, "residual_scaling") else - self.config.residual_scaling, - relative_attention=self.config.relative_attention, - max_distance=self.config.max_distance, - num_buckets=self.config.num_buckets, - fp16_clamping=self.fp16_clamping, - quant_mode=self.config.quant_mode, - language_adapter_config=self.language_adapter_config) - for _ in self.mapping.pp_layers(self.total_num_layers) - ]) - self.transformer.assign_module(encoder_layers, "layers") - - def check_config(self, config: PretrainedConfig): - config.set_if_not_exist('has_position_embedding', False) - config.set_if_not_exist('type_vocab_size', None) - config.set_if_not_exist('rescale_before_lm_head', False) - config.set_if_not_exist('layernorm_type', LayerNormType.LayerNorm) - config.set_if_not_exist('layernorm_position', - LayerNormPositionType.pre_layernorm) - config.set_if_not_exist('has_attention_qkvo_bias', False) - config.set_if_not_exist('has_mlp_bias', False) - config.set_if_not_exist('has_model_final_layernorm', False) - config.set_if_not_exist('encoder_hidden_size', None) - config.set_if_not_exist('encoder_num_heads', None) - config.set_if_not_exist('encoder_num_kv_heads', None) - config.set_if_not_exist('encoder_head_size', None) - config.set_if_not_exist('model_type', 't5') - config.set_if_not_exist('skip_cross_kv', False) - config.set_if_not_exist('mlp_type', MLPType.MLP) - config.set_if_not_exist('has_embedding_scale', False) - config.set_if_not_exist('residual_scaling', 1.0) - config.set_if_not_exist('has_lm_head_bias', False) - config.set_if_not_exist('num_buckets', None) - config.set_if_not_exist('max_distance', None) - config.set_if_not_exist('relative_attention', False) - config.set_if_not_exist('residual_scaling', 1.0) - - def forward(self, - input_ids: Tensor, - input_lengths=None, - position_ids=None, - token_type_ids=None, - hidden_states=None, - max_input_length=None, - prompt_embedding_table=None, - prompt_tasks=None, - prompt_vocab_size=None, - attention_mask=None, - lora_params: LoraParams = None, - language_adapter_routings: Optional[Tensor] = None): - # In PP, layer 0 has ids as inputs, all other layers have hidden_states as inputs - if self.mapping.is_first_pp_rank(): - ptuning_args = [ - prompt_embedding_table, prompt_tasks, prompt_vocab_size - ] if prompt_embedding_table is not None else [] - - hidden_states = self.transformer.embedding(input_ids, position_ids, - token_type_ids, - *ptuning_args) - self.register_network_output('embedding_layer_output', - hidden_states) - else: - hidden_states = recv(hidden_states, self.mapping.prev_pp_rank()) - - for layer_idx, encoder_layer in enumerate(self.transformer.layers): - lora_layer_params = None - if lora_params is not None and lora_params.lora_ranks is not None: - lora_layer_params = lora_params.get_layer_params(layer_idx) - hidden_states = encoder_layer( - hidden_states=hidden_states, - attention_mask=attention_mask, - input_lengths=input_lengths, - max_input_length=max_input_length, - lora_layer_params=lora_layer_params, - language_adapter_routings=language_adapter_routings) - - if self.mapping.is_last_pp_rank(): - if self.has_model_final_layernorm: - hidden_states = self.transformer.ln_f(hidden_states) - hidden_states.mark_output('encoder_output', self._dtype) - else: - hidden_states = send(hidden_states, self.mapping.next_pp_rank()) - hidden_states.mark_output('hidden_states_output', self._dtype) - - return hidden_states - - def prepare_inputs(self, - max_batch_size, - max_input_len, - prompt_embedding_table_size: int = 0, - lora_target_modules: List[str] = None, - *args, - **kwargs): - '''@brief: Prepare inputs Tensors for the model, the given sizes are used to determine the - ranges of the dimensions of when using TRT dynamic shapes. - - @return: a list contains values which can be fed into the self.forward() - ''' - - hidden_size = self.hidden_size - - bs_range = [1, (max_batch_size + 1) // 2, max_batch_size] - inlen_range = [1, (max_input_len + 1) // 2, max_input_len] - num_tokens_range = [ - 1, - (max_input_len * max_batch_size + 1) // 2, - max_input_len * max_batch_size, - ] - - input_ids, position_ids, token_type_ids, hidden_states = None, None, None, None - remove_input_padding = default_net().plugin_config.remove_input_padding - use_lora_plugin = default_net().plugin_config.lora_plugin - - attention_mask = None - if remove_input_padding: - if self.mapping.is_first_pp_rank(): - input_ids = Tensor( - name="input_ids", - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([("num_tokens", [num_tokens_range])]), - ) - if self.has_position_embedding: - position_ids = Tensor( - name='position_ids', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([('num_tokens', - [num_tokens_range])]), - ) - if self.has_token_type_embedding: - token_type_ids = Tensor( - name='token_type_ids', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([('num_tokens', - [num_tokens_range])]), - ) - else: - hidden_states = Tensor(name='hidden_states_input', - dtype=self._dtype, - shape=[-1, hidden_size], - dim_range=OrderedDict([ - ('num_tokens', [num_tokens_range]), - ('hidden_size', [hidden_size]), - ])) - else: - if self.mapping.is_first_pp_rank(): - input_ids = Tensor( - name="input_ids", - dtype=trt.int32, - shape=[-1, -1], - dim_range=OrderedDict([("batch_size", [bs_range]), - ("input_len", [inlen_range])]), - ) - if self.has_position_embedding: - position_ids = Tensor( - name='position_ids', - dtype=trt.int32, - shape=[-1, -1], - dim_range=OrderedDict([('batch_size', [bs_range]), - ('input_len', [inlen_range])]), - ) - if self.has_token_type_embedding: - token_type_ids = Tensor( - name='token_type_ids', - dtype=trt.int32, - shape=[-1, -1], - dim_range=OrderedDict([('batch_size', [bs_range]), - ('input_len', [inlen_range])]), - ) - else: - hidden_states = Tensor(name='hidden_states_input', - dtype=self._dtype, - shape=[-1, -1, hidden_size], - dim_range=OrderedDict([ - ('batch_size', [bs_range]), - ('input_len', [inlen_range]), - ('hidden_size', [hidden_size]), - ])) - - if not default_net().plugin_config.bert_attention_plugin: - attention_mask = Tensor( - name='attention_mask', - dtype=trt.int32, - shape=[-1, -1], - dim_range=OrderedDict([ - ('batch_size', [bs_range]), - ('input_len', [inlen_range]), - ]), - ) - - # if self.mapping.tp_size > 1: - # current_all_reduce_helper().set_workspace_tensor(self.mapping, 1) - # FIXME(TRTLLM-996): Support custom allreduce for encoder models on C++ runtime - - input_lengths = Tensor( - name="input_lengths", - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([("batch_size", [bs_range])]), - ) - max_input_length = Tensor( - name="max_input_length", - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([("max_input_length", [inlen_range])]), - ) - - prompt_embedding_table = None - tasks = None - prompt_vocab_size = None - - if self.mapping.is_first_pp_rank() and prompt_embedding_table_size > 0: - p_embedding_range = [[ - 1, prompt_embedding_table_size // 2, prompt_embedding_table_size - ]] - - prompt_embedding_table = Tensor(name='prompt_embedding_table', - dtype=self._dtype, - shape=[-1, hidden_size], - dim_range=OrderedDict([ - ('prompt_embedding_table_size', - p_embedding_range), - ('hidden_size', [hidden_size]), - ])) - if remove_input_padding: - tasks = Tensor(name='tasks', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([('input_len_task', - [num_tokens_range])])) - else: - tasks = Tensor(name='tasks', - dtype=trt.int32, - shape=[-1, 1], - dim_range=OrderedDict([ - ('batch_size', bs_range), - ('broadcast_dim', [1]), - ])) - prompt_vocab_size = Tensor(name='prompt_vocab_size', - dtype=trt.int32, - shape=[1], - dim_range=OrderedDict([('size', [1])])) - ''' - LoRA plugin related inputs: - lora_target_modules for BART-encoder: - ['attn_q', 'attn_v'] - For BART-decoder, the lora_target_modules is different. - See comments in the DecoderModel.prepare_inputs() for more details. - ''' - lora_weights_pointers = None - lora_ranks = None - lora_params = None - if use_lora_plugin: - lora_weights_pointers = [] - lora_ranks = [] - # In current design, q_lora_params, k_lora_params and v_lora_params should be all enabled or all disabled at the same time. - # However, BART lora modules only contain two of them, so we use zero tensor to fill the missing ones. - missing_qkv_modules = [] - if any(x in lora_target_modules - for x in ["attn_q", "attn_k", "attn_v"]): - for lora_module in ["attn_q", "attn_k", "attn_v"]: - if lora_module not in lora_target_modules: - missing_qkv_modules.append(lora_module) - - layers_range = self.mapping.pp_layers(self.total_num_layers) - for i in layers_range: - lora_weight_pointer_dict = {} - lora_rank_dict = {} - for lora_module in (lora_target_modules + missing_qkv_modules): - lora_weight_pointer = Tensor( - name=f'{lora_module}_lora_weights_pointers_{i}', - dtype=trt.int64, - shape=[-1, 3], - dim_range=OrderedDict([('batch_size', [bs_range]), - ('in_out_scales', [3])])) - lora_weight_pointer_dict.update({ - f'{lora_module}_lora_weights_pointers': - lora_weight_pointer - }) - - lora_rank = Tensor(name=f'{lora_module}_lora_ranks_{i}', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([('batch_size', - [bs_range])])) - lora_rank_dict.update( - {f'{lora_module}_lora_ranks': lora_rank}) - - lora_weights_pointers.append(lora_weight_pointer_dict) - lora_ranks.append(lora_rank_dict) - - host_request_types = Tensor(name='host_request_types', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([('batch_size', - [bs_range])])) - - host_context_lengths = None - if remove_input_padding: - host_context_lengths = Tensor(name='host_context_lengths', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([ - ('batch_size', [bs_range]) - ])) - - lora_params = LoraParams( - lora_ranks=lora_ranks, - lora_weights_pointers=lora_weights_pointers, - host_request_types=host_request_types, - host_context_lengths=host_context_lengths, - ) - - language_adapter_routings = None - if self.language_adapter_config: - language_adapter_routings = Tensor( - name="language_adapter_routings", - dtype=trt.int32, - shape=[-1, 1], - dim_range=OrderedDict([('num_tokens', [num_tokens_range]), - ('routing_dim', [1])]), - ) - - result = { - 'input_ids': input_ids, - 'input_lengths': input_lengths, - 'position_ids': position_ids, - 'token_type_ids': token_type_ids, - 'hidden_states': hidden_states, - 'max_input_length': max_input_length, - 'prompt_embedding_table': prompt_embedding_table, - 'prompt_tasks': tasks, - 'prompt_vocab_size': prompt_vocab_size, - 'attention_mask': attention_mask, - 'lora_params': lora_params, - 'language_adapter_routings': language_adapter_routings, - } - - return result - - def use_lora(self, lora_config: LoraConfig): - use_lora(self, lora_config) - - def use_prompt_tuning(self): - embedding = self.transformer.vocab_embedding - self.transformer.vocab_embedding = PromptTuningEmbedding( - num_embeddings=embedding.num_embeddings, - embedding_dim=embedding.embedding_dim, - dtype=embedding.dtype, - tp_size=embedding.tp_size, - tp_group=embedding.tp_group, - sharding_dim=embedding.sharding_dim, - tp_rank=embedding.tp_rank) - - self.transformer.vocab_embedding.weight.value = embedding.weight.raw_value - - def precompute_relative_attention_bias(self, build_config): - pass - - -class DecoderModel(PretrainedModel): - - def __init__(self, config: PretrainedConfig): - self.check_config(config) - super().__init__(config) - - self.mapping = self.config.mapping - - self.has_position_embedding = self.config.has_position_embedding - type_vocab_size = self.config.type_vocab_size - self.has_token_type_embedding = (type_vocab_size is not None) - self.rescale_before_lm_head = self.config.rescale_before_lm_head - - # e.g. BART regular, T5 RMS - self.layernorm_type = self.config.layernorm_type - - # e.g. BART true, T5 false - self.has_attention_qkvo_bias = self.config.has_attention_qkvo_bias - self.has_mlp_bias = self.config.has_mlp_bias - - # e.g. BART false, T5 true - self.has_model_final_layernorm = self.config.has_model_final_layernorm - self._dtype = self.config.dtype - # no quantization considered for now - self._kv_dtype = self._dtype - self._logits_dtype = self.config.logits_dtype - - self.total_num_layers = self.config.num_hidden_layers - self.num_layers = self.config.num_hidden_layers // self.mapping.pp_size - - self.hidden_size = self.config.hidden_size - self.num_heads = self.config.num_attention_heads - num_kv_heads = self.num_heads - if num_kv_heads is None or num_kv_heads <= 0: - num_kv_heads = self.num_heads - self.num_kv_heads = num_kv_heads - self.head_size = self.hidden_size // self.num_heads if self.config.head_size is None else self.config.head_size - - self.encoder_hidden_size = self.config.encoder_hidden_size - self.encoder_num_heads = self.config.encoder_num_heads - encoder_num_kv_heads = None if not hasattr( - self.config, - "encoder_num_kv_heads") else self.config.encoder_num_kv_heads - if encoder_num_kv_heads is None or encoder_num_kv_heads <= 0: - encoder_num_kv_heads = self.encoder_num_heads - self.encoder_num_kv_heads = encoder_num_kv_heads - self.encoder_head_size = self.encoder_hidden_size // self.num_heads if self.config.encoder_head_size is None else self.config.encoder_head_size - - self.has_position_embedding = self.config.has_position_embedding - self.has_token_type_embedding = type_vocab_size is not None - - self.fp16_clamping = (self.config.dtype - == 'float16') and (self.config.model_type - in ['t5', 'pix2struct']) - - self.skip_cross_kv = self.config.skip_cross_kv - self.mlp_type = MLPType.MLP if not hasattr( - self.config, "mlp_type") else self.config.mlp_type - self.use_implicit_relative_attention = self.config.use_implicit_relative_attention if hasattr( - self.config, "use_implicit_relative_attention") else False - - self.language_adapter_config = None if not hasattr( - self.config, - 'language_adapter_config') else LanguageAdapterConfig.from_dict( - self.config.language_adapter_config) - - self.transformer = EncDecTransformer( - self.config.vocab_size, - self.config.hidden_size, - max_position_embeddings=self.config.max_position_embeddings, - has_position_embedding=self.config.has_position_embedding, - type_vocab_size=type_vocab_size, - has_embedding_layernorm=self.config.has_embedding_layernorm, - has_embedding_scale=self.config.has_embedding_scale, - layernorm_eps=self.config.norm_epsilon, - layernorm_type=self.config.layernorm_type, - dtype=self._dtype, - use_parallel_embedding=self.config.use_parallel_embedding, - embedding_sharding_dim=self.config.embedding_sharding_dim, - mapping=self.mapping, - has_model_final_layernorm=self.has_model_final_layernorm, - norm_epsilon=self.config.norm_epsilon, - is_decoder=True) - - layers_range = self.mapping.pp_layers(self.total_num_layers) - decoder_layers = ModuleList([ - DecoderLayer( - local_layer_idx=layer_idx - layers_range[0], - hidden_size=self.config.hidden_size, - ffn_hidden_size=self.config.intermediate_size, - num_attention_heads=self.num_heads, - num_kv_heads=self.num_kv_heads, - head_size=self.head_size, - max_position_embeddings=self.config.max_position_embeddings, - q_scaling=self.config.q_scaling, - has_attention_qkvo_bias=self.config.has_attention_qkvo_bias, - has_mlp_bias=self.config.has_mlp_bias, - layernorm_position=self.config.layernorm_position, - layernorm_eps=self.config.norm_epsilon, - layernorm_type=self.config.layernorm_type, - hidden_act=self.config.hidden_act, - mlp_type=self.mlp_type, - mapping=self.mapping, - dtype=self._dtype, - residual_scaling=self.config.residual_scaling, - relative_attention=self.config.relative_attention, - max_distance=self.config.max_distance, - num_buckets=self.config.num_buckets, - fp16_clamping=self.fp16_clamping, - skip_cross_kv=self.skip_cross_kv, - use_implicit_relative_attention=self. - use_implicit_relative_attention, - quant_mode=self.config.quant_mode, - language_adapter_config=self.language_adapter_config) - for layer_idx in layers_range - ]) - self.transformer.assign_module(decoder_layers, "layers") - - if self.mapping.is_last_pp_rank(): - vocab_size_padded = pad_vocab_size(self.config.vocab_size, - self.config.mapping.tp_size) - self.lm_head = ColumnLinear( - self.config.hidden_size, - vocab_size_padded, - bias=False if not hasattr(self.config, "has_lm_head_bias") else - self.config.has_lm_head_bias, - dtype=self.config.dtype, - tp_group=self.config.mapping.tp_group, - tp_size=self.config.mapping.tp_size, - gather_output=True, - ) - - self.trtllm_modules_to_hf_modules = { - **get_default_trtllm_modules_to_hf_modules(), - "attn_q": "self_attn.q_proj", - "attn_k": "self_attn.k_proj", - "attn_v": "self_attn.v_proj", - "attn_dense": "self_attn.o_proj", - "cross_attn_q": "encoder_attn.q_proj", - "cross_attn_k": "encoder_attn.k_proj", - "cross_attn_v": "encoder_attn.v_proj", - "cross_attn_dense": "encoder_attn.o_proj", - } - - if self.config.relative_attention and not self.use_implicit_relative_attention: - self.rel_attn_table = Parameter( - shape=(self.config.num_attention_heads // self.mapping.tp_size, - self.config.num_buckets), - dtype=self._dtype) - - def check_config(self, config: PretrainedConfig): - config.set_if_not_exist('has_position_embedding', False) - config.set_if_not_exist('type_vocab_size', None) - config.set_if_not_exist('rescale_before_lm_head', False) - config.set_if_not_exist('layernorm_type', LayerNormType.LayerNorm) - config.set_if_not_exist('layernorm_position', - LayerNormPositionType.pre_layernorm) - config.set_if_not_exist('has_attention_qkvo_bias', False) - config.set_if_not_exist('has_mlp_bias', False) - config.set_if_not_exist('has_model_final_layernorm', False) - config.set_if_not_exist('encoder_hidden_size', None) - config.set_if_not_exist('encoder_num_heads', None) - config.set_if_not_exist('encoder_num_kv_heads', None) - config.set_if_not_exist('encoder_head_size', None) - config.set_if_not_exist('model_type', 't5') - config.set_if_not_exist('skip_cross_kv', False) - config.set_if_not_exist('mlp_type', MLPType.MLP) - config.set_if_not_exist('has_embedding_scale', False) - config.set_if_not_exist('residual_scaling', 1.0) - config.set_if_not_exist('has_lm_head_bias', False) - config.set_if_not_exist('num_buckets', None) - config.set_if_not_exist('max_distance', None) - config.set_if_not_exist('relative_attention', False) - - def forward(self, - decoder_input_ids: Tensor, - encoder_output: Tensor, - position_ids=None, - token_type_ids=None, - use_cache=False, - attention_mask_params=None, - last_token_ids=None, - kv_cache_params=None, - attention_params=None, - hidden_states=None, - lora_params: LoraParams = None, - cross_kv_cache_gen: Optional[Tensor] = None, - cross_kv_reuse: Optional[Tensor] = None, - language_adapter_routings: Optional[Tensor] = None): - if self.mapping.is_first_pp_rank(): - assert isinstance(decoder_input_ids, Tensor) - else: - assert isinstance(hidden_states, Tensor) - - # In PP, layer 0 has ids as inputs, all other layers have hidden_states as inputs - if self.mapping.is_first_pp_rank(): - hidden_states = self.transformer.embedding(decoder_input_ids, - position_ids, - token_type_ids) - self.register_network_output('embedding_layer_output', - hidden_states) - else: - hidden_states = recv(hidden_states, self.mapping.prev_pp_rank()) - - kv_cache_params.fill_none_tensor_list(len(self.transformer.layers)) - - if use_cache: - presents = [] - - for i, (decoder_layer, past) in enumerate( - zip(self.transformer.layers, kv_cache_params.past_key_value)): - - lora_layer_params = None - if lora_params is not None and lora_params.lora_ranks is not None: - lora_layer_params = lora_params.get_layer_params(i) - - hidden_states = decoder_layer( - hidden_states, - encoder_output=encoder_output, - attention_mask_params=attention_mask_params, - use_cache=use_cache, - kv_cache_params=KeyValueCacheParams( - past_key_value=past, - host_past_key_value_lengths=kv_cache_params. - host_past_key_value_lengths, - host_max_attention_window_sizes=kv_cache_params. - host_max_attention_window_sizes, - host_sink_token_length=kv_cache_params. - host_sink_token_length, - cache_indirection=kv_cache_params.cache_indirection, - kv_cache_block_offsets=kv_cache_params. - kv_cache_block_offsets, - host_kv_cache_block_offsets=kv_cache_params. - host_cross_kv_cache_block_offsets, - host_kv_cache_pool_pointers=kv_cache_params. - host_kv_cache_pool_pointers, - host_kv_cache_pool_mapping=kv_cache_params. - host_kv_cache_pool_mapping, - cross_kv_cache_block_offsets=kv_cache_params. - cross_kv_cache_block_offsets, - host_cross_kv_cache_block_offsets=kv_cache_params. - host_cross_kv_cache_block_offsets, - host_cross_kv_cache_pool_pointers=kv_cache_params. - host_cross_kv_cache_pool_pointers, - host_cross_kv_cache_pool_mapping=kv_cache_params. - host_cross_kv_cache_pool_mapping), - attention_params=attention_params, - lora_layer_params=lora_layer_params, - cross_kv_cache_gen=cross_kv_cache_gen, - cross_kv_reuse=cross_kv_reuse, - language_adapter_routings=language_adapter_routings) - - if use_cache: - presents_self, presents_cross = hidden_states[1], hidden_states[ - 2] - presents.append((presents_self, presents_cross)) - hidden_states = hidden_states[0] - self.register_network_output(f'decoder_layer_{i}_output', - hidden_states) - - if self.mapping.is_last_pp_rank(): - if self.has_model_final_layernorm: - hidden_states = self.transformer.ln_f(hidden_states) - - # [bs, seq, hidden_size] or [num_tokens, hidden_size] -> [bs, hidden_size] - hidden_states = gather_last_token_logits( - hidden_states, last_token_ids, - default_net().plugin_config.remove_input_padding) - self.register_network_output('logits_before_lmhead', hidden_states) - - # Rescale output before projecting on vocab (for T5) - # See https://github.com/huggingface/transformers/blob/0b192de1f353b0e04dad4813e02e2c672de077be/src/transformers/models/t5/modeling_t5.py#L1769-L1772 - # Note: this is specific for T5, to make it more generic, one can pass in a config: - # self.config.tie_word_embeddings - default to be True for T5 - # openai whisper model didn't use this rescale - if self.rescale_before_lm_head: - hidden_states = hidden_states * (self.hidden_size**-0.5) - - # [bs, hidden_size] -> [bs, vocab_size] - lm_logits = self.lm_head(hidden_states) - lm_logits.mark_output('logits', self._logits_dtype) - else: - hidden_states = send(hidden_states, self.mapping.next_pp_rank()) - hidden_states.mark_output('hidden_states_output', self._dtype) - - if use_cache and default_net().plugin_config.paged_kv_cache == False: - for i, present in zip(self.mapping.pp_layers(self.total_num_layers), - presents): - present[0].mark_output(f'present_key_value_{i}', self._kv_dtype) - if default_net().plugin_config.gpt_attention_plugin: - present[1].mark_output(f'cross_present_key_value_{i}', - self._kv_dtype) - if self.mapping.is_last_pp_rank(): - return (lm_logits, tuple(presents)) - return (hidden_states, tuple(presents)) - else: - if self.mapping.is_last_pp_rank(): - return lm_logits - return hidden_states - - def prepare_inputs(self, - max_batch_size, - max_beam_width, - max_decoder_input_len, - max_seq_len, - max_encoder_input_len, - gather_context_logits: bool = False, - lora_target_modules: List[str] = None, - use_cache=True, - *args, - **kwargs): - '''@brief: Prepare inputs Tensors for the model, the given sizes are used to determine the - ranges of the dimensions of when using TRT dynamic shapes. - - @return: a list contains values which can be fed into the self.forward() - ''' - - # Prepare inputs - max_output_len = max_decoder_input_len + max_seq_len - - head_size = self.head_size - num_kv_heads = (self.num_kv_heads + self.mapping.tp_size - - 1) // self.mapping.tp_size - - encoder_head_size = self.encoder_head_size - encoder_num_kv_heads = (self.encoder_num_kv_heads + self.mapping.tp_size - - 1) // self.mapping.tp_size - - bb_range = [ - 1, (max_batch_size * max_beam_width + 1) // 2, - max_batch_size * max_beam_width - ] - bs_range = [1, (max_batch_size + 1) // 2, max_batch_size] - beam_width_range = [1, (max_beam_width + 1) // 2, max_beam_width] - inlen_range = [ - 1, 1, max_decoder_input_len - ] # context phase >= 1 (if forced_input_ids), generation phase = 1 - encoder_inlen_range = [ - 1, (max_encoder_input_len + 1) // 2, max_encoder_input_len - ] - mask_len_range = [1, (max_output_len + 1) // 2 + 1, max_output_len + 1] - max_output_len_range = [0, (max_output_len + 1) // 2, max_output_len] - - encoder_num_tokens_range = [ - 0, # 0 for generation phase, >0 for context phase - (max_encoder_input_len * max_batch_size + 1) // 2, - max_encoder_input_len * max_batch_size, - ] - decoder_num_tokens_range = [ - 1, - max_batch_size * max_beam_width, - max(max_decoder_input_len * max_batch_size, - max_beam_width * max_batch_size), - ] - - # No enable_two_optimization_profiles support yet - - encoder_input_len_range = [ - 0, # 0 for generation phase, >0 for context phase - (max_encoder_input_len + 1) // 2, - max_encoder_input_len - ] - max_cross_packed_mask_dim0 = max_batch_size * ( - (max_decoder_input_len + 128 - 1) // 128) * 128 - max_cross_packed_mask_dim1 = ( - (max_encoder_input_len + 256 - 1) // 256) * 256 // 32 - cross_packed_mask_dim0_range = [ - 1, (max_cross_packed_mask_dim0 + 1) // 2, max_cross_packed_mask_dim0 - ] - cross_packed_mask_dim1_range = [ - 0, # 0 for generation phase, >0 for context phase - (max_cross_packed_mask_dim1 + 1) // 2, - max_cross_packed_mask_dim1 - ] - - past_key_value = [] - sequence_length = None - host_past_key_value_lengths = None - runtime_perf_knobs = None - context_progress = None - attention_mask = None - cross_attention_mask = None - cross_attention_packed_mask = None - attention_mask_params = AttentionMaskParams() - use_gpt_attention_plugin = default_net( - ).plugin_config.gpt_attention_plugin - remove_input_padding = default_net().plugin_config.remove_input_padding - paged_kv_cache = default_net().plugin_config.paged_kv_cache - tokens_per_block = default_net().plugin_config.tokens_per_block - use_lora_plugin = default_net().plugin_config.lora_plugin - - input_ids, position_ids, token_type_ids, hidden_states = None, None, None, None - if remove_input_padding: - if self.mapping.is_first_pp_rank(): - input_ids = Tensor(name='input_ids', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([ - ('decoder_num_tokens', - [decoder_num_tokens_range]), - ])) - if self.has_position_embedding: - position_ids = Tensor(name='position_ids', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([ - ('decoder_num_tokens', - [decoder_num_tokens_range]), - ])) - if self.has_token_type_embedding: - token_type_ids = Tensor( - name='token_type_ids', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([('decoder_num_tokens', - [decoder_num_tokens_range])]), - ) - else: - hidden_states = Tensor(name='hidden_states_input', - dtype=self._dtype, - shape=[-1, self.hidden_size], - dim_range=OrderedDict([ - ('decoder_num_tokens', - [decoder_num_tokens_range]), - ('hidden_size', [self.hidden_size]), - ])) - else: - if self.mapping.is_first_pp_rank(): - input_ids = Tensor(name='input_ids', - dtype=trt.int32, - shape=[-1, -1], - dim_range=OrderedDict([ - ('batch_size_beam_width', [bb_range]), - ('input_len', [inlen_range]), - ])) - if self.has_position_embedding: - position_ids = Tensor(name='position_ids', - dtype=trt.int32, - shape=[-1, -1], - dim_range=OrderedDict([ - ('batch_size_beam_width', - [bb_range]), - ('input_len', [inlen_range]), - ])) - if self.has_token_type_embedding: - token_type_ids = Tensor( - name='token_type_ids', - dtype=trt.int32, - shape=[-1, -1], - dim_range=OrderedDict([('batch_size_beam_width', - [bb_range]), - ('input_len', [inlen_range])]), - ) - else: - hidden_states = Tensor(name='hidden_states_input', - dtype=self._dtype, - shape=[-1, -1, self.hidden_size], - dim_range=OrderedDict([ - ('batch_size_beam_width', [bb_range - ]), - ('input_len', [inlen_range]), - ('hidden_size', [self.hidden_size]), - ])) - - encoder_input_lengths = Tensor( - name="encoder_input_lengths", - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([("batch_size_beam_width", [bb_range])]), - ) - encoder_max_input_length = Tensor( - name="encoder_max_input_length", - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([("encoder_max_input_length", - [encoder_inlen_range])]), - ) - encoder_output = None - if remove_input_padding: - encoder_output = Tensor( - name="encoder_output", - dtype=self._dtype, - shape=[-1, self.encoder_hidden_size], - dim_range=OrderedDict([ - ("encoder_num_tokens", [encoder_num_tokens_range]), - ("encoder_hidden_size", [self.encoder_hidden_size]), - ]), - ) - else: - encoder_output = Tensor( - name="encoder_output", - dtype=self._dtype, - shape=[-1, -1, self.encoder_hidden_size], - dim_range=OrderedDict([ - ("batch_size_beam_width_encoder", [bb_range]), - ("encoder_input_len", [encoder_input_len_range]), - ("encoder_hidden_size", [self.encoder_hidden_size]), - ]), - ) - - if use_gpt_attention_plugin: - host_past_key_value_lengths = Tensor( - name='host_past_key_value_lengths', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([('batch_size_beam_width', [bb_range])]), - ) - - context_lengths = None - host_context_lengths = None - host_request_types = None - if use_gpt_attention_plugin and remove_input_padding: - host_context_lengths = Tensor(name='host_context_lengths', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([ - ('batch_size_beam_width', - [bb_range]) - ])) - - if use_gpt_attention_plugin: - sequence_length = Tensor( - name='sequence_length', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([('batch_size_beam_width', [bb_range])]), - ) - - context_lengths = Tensor(name='context_lengths', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([ - ('batch_size_beam_width', [bb_range]) - ])) - host_request_types = Tensor(name='host_request_types', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([ - ('batch_size_beam_width', - [bb_range]) - ])) - runtime_perf_knobs = Tensor(name='host_runtime_perf_knobs', - dtype=trt.int64, - shape=[16], - dim_range=OrderedDict([ - ('perf_knob_size', [16]) - ])) - context_progress = Tensor(name='host_context_progress', - dtype=trt.int64, - shape=[1], - dim_range=OrderedDict([ - ('context_progress_size', [1]) - ])) - - last_token_ids = None - if self.mapping.is_last_pp_rank() and not gather_context_logits: - last_token_ids = Tensor( - name="last_token_ids", - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([("batch_size_last_token_ids", [bb_range]) - ]), - ) - - if not use_gpt_attention_plugin: - attention_mask = Tensor( - name='attention_mask', - dtype=trt.int32, - shape=[-1, -1], - dim_range=OrderedDict([ - ('batch_size_beam_width', [bb_range]), - ('mask_len', [mask_len_range]), - ]), - ) - - cross_attention_mask = Tensor( - name='cross_attention_mask', - dtype=trt.int32, - shape=[-1, -1, -1], - dim_range=OrderedDict([ - ('batch_size_beam_width', [bb_range]), - ('query_len', [inlen_range]), - ('encoder_input_len_2', [encoder_input_len_range]), - ]), - ) - else: - cross_attention_mask = Tensor( - name='cross_attention_mask', - dtype=trt.bool, - shape=[-1, -1], - dim_range=OrderedDict([ - ('decoder_num_tokens_2', - [decoder_num_tokens_range - ]), # TODO should use same name as input_ids - ('encoder_input_len_2', [encoder_input_len_range]), - ]), - ) - - cross_attention_packed_mask = Tensor( - name='cross_attention_packed_mask', - dtype=trt.int32, - shape=[-1, -1], - dim_range=OrderedDict([ - ('cross_packed_mask_dim0', [cross_packed_mask_dim0_range]), - ('cross_packed_mask_dim1', [cross_packed_mask_dim1_range]), - ]), - ) - - # create the attention_mask_params. - attention_mask_params = AttentionMaskParams( - attention_mask, None, cross_attention_mask, - cross_attention_packed_mask) - - cache_indirection = Tensor( - name='cache_indirection', - dtype=trt.int32, - shape=[-1, -1, -1], - dim_range=OrderedDict([ - ('batch_size_cache', [bs_range]), - ('beam_width', [beam_width_range]), - ('max_seq_len', [max_output_len_range]), - ]), - ) - - if self.mapping.tp_size > 1: - current_all_reduce_helper().set_workspace_tensor(self.mapping, 1) - - layers_range = self.mapping.pp_layers(self.total_num_layers) - num_pp_layers = len(layers_range) - - host_max_attention_window_sizes = None - host_sink_token_length = None - if use_gpt_attention_plugin: - host_max_attention_window_sizes = Tensor( - name=f'host_max_attention_window_sizes', - dtype=trt.int32, - shape=[num_pp_layers], - dim_range=OrderedDict([('num_layers', [num_pp_layers])])) - host_sink_token_length = Tensor(name='host_sink_token_length', - dtype=trt.int32, - shape=[1], - dim_range=OrderedDict([('scalar', - [1])])) - ''' - LoRA plugin related inputs: - lora_target_modules for BART-decoder: - ['attn_q', 'cross_attn_q', - 'attn_v', 'cross_attn_v'] - This is NOT directly loaded from the adapter-config file - We make it this way because BART has LoRA weights for both self-attention and cross-attention in decoder - ''' - lora_weights_pointers = None - lora_ranks = None - lora_params = None - if use_lora_plugin: - lora_weights_pointers = [] - lora_ranks = [] - # In current design, q_lora_params, k_lora_params and v_lora_params should be all enabled or all disabled at the same time. - # However, BART lora modules only contain two of them, so we use zero tensor to fill the missing ones. - missing_qkv_modules = [] - if any(x in lora_target_modules - for x in ["attn_q", "attn_k", "attn_v"]): - for lora_module in [ - "attn_q", - "attn_k", - "attn_v", - ]: - if lora_module not in lora_target_modules: - missing_qkv_modules.append(lora_module) - if any(x in lora_target_modules - for x in ["cross_attn_q", "cross_attn_k", "cross_attn_v"]): - for lora_module in [ - "cross_attn_q", "cross_attn_k", "cross_attn_v" - ]: - if lora_module not in lora_target_modules: - missing_qkv_modules.append(lora_module) - - for i in layers_range: - lora_weight_pointer_dict = {} - lora_rank_dict = {} - for lora_module in (lora_target_modules + missing_qkv_modules): - lora_weight_pointer = Tensor( - name=f'{lora_module}_lora_weights_pointers_{i}', - dtype=trt.int64, - shape=[-1, 3], - dim_range=OrderedDict([('batch_size_beam_width', - [bb_range]), - ('in_out_scales', [3])])) - lora_weight_pointer_dict.update({ - f'{lora_module}_lora_weights_pointers': - lora_weight_pointer - }) - - lora_rank = Tensor(name=f'{lora_module}_lora_ranks_{i}', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([ - ('batch_size_beam_width', [bb_range]) - ])) - lora_rank_dict.update( - {f'{lora_module}_lora_ranks': lora_rank}) - - lora_weights_pointers.append(lora_weight_pointer_dict) - lora_ranks.append(lora_rank_dict) - - # For cross attention, we need to use encoder_input_lengths (in CPU) to pass - # as the host_context_lengths to the lora_plugin. But for self attention, we - # should keep using the original host_context_lengths. Therefore, we keep both - # of them in the lora_params. - host_encoder_input_lengths = None - if remove_input_padding: - host_encoder_input_lengths = Tensor( - name="host_encoder_input_lengths", - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([("batch_size_beam_width", [bb_range]) - ]), - ) - - lora_params = LoraParams( - lora_ranks=lora_ranks, - lora_weights_pointers=lora_weights_pointers, - host_context_lengths=host_context_lengths, - max_encoder_context_length=max_encoder_input_len, - host_request_types=host_request_types, - host_encoder_input_lengths=host_encoder_input_lengths, - ) - - kv_cache_block_offsets = None - host_kv_cache_block_offsets = None - host_kv_cache_pool_pointers = None - host_kv_cache_pool_mapping = None - - cross_kv_cache_block_offsets = None - host_cross_kv_cache_block_offsets = None - host_cross_kv_cache_pool_pointers = None - host_cross_kv_cache_pool_mapping = None - - if use_cache: - if not paged_kv_cache: - for i in layers_range: - kv_dim_range = OrderedDict([ - ('batch_size_beam_width', [bb_range]), - ('kv', [2]), - ('num_heads', [num_kv_heads]), - ('past_key_len', [max_output_len_range]), - ('head_size', [head_size]), - ]) - kv = Tensor(name=f'past_key_value_{i}', - dtype=self._kv_dtype, - shape=[-1, 2, num_kv_heads, -1, head_size], - dim_range=kv_dim_range) - - if use_gpt_attention_plugin: - cross_kv_dim_range = OrderedDict([ - ('batch_size_beam_width', [bb_range]), - ('kv', [2]), - ('cross_num_heads', [encoder_num_kv_heads]), - ('cross_past_key_len', [encoder_input_len_range]), - ('cross_head_size', [encoder_head_size]), - ]) - cross_kv = Tensor(name=f'cross_past_key_value_{i}', - dtype=self._kv_dtype, - shape=[ - -1, 2, encoder_num_kv_heads, -1, - encoder_head_size - ], - dim_range=cross_kv_dim_range) - past_key_value.append((kv, cross_kv)) - else: - # use encoder_output directly, no need to save cross_past_key_value - past_key_value.append((kv, )) - - # TODO: Remove this when TRT fix the named dimension - if not remove_input_padding: - assertion( - shape( - input_ids if self.mapping.is_first_pp_rank() else - hidden_states, 0) == shape(kv, 0), 'batch size') - - else: # paged_kv_cache == True - # PagedKV setup for KV cache of self-attention - max_blocks_per_seq_range = [[ - math.ceil(max_output_len_range[0] / tokens_per_block), - math.ceil(max_output_len_range[1] / tokens_per_block), - math.ceil(max_output_len_range[2] / tokens_per_block) - ]] - max_blocks_per_seq_range = [[ - x for x in max_blocks_per_seq_range[0] - ]] - - # PagedKV setup for KV cache of cross-attention - max_cross_blocks_per_seq_range = [[ - math.ceil(encoder_input_len_range[0] / tokens_per_block), - math.ceil(encoder_input_len_range[1] / tokens_per_block), - math.ceil(encoder_input_len_range[2] / tokens_per_block) - ]] - max_cross_blocks_per_seq_range = [[ - x for x in max_cross_blocks_per_seq_range[0] - ]] - - # TODO(oargov): add support for vgqa + vwindow, meanwhile assume a single kv cache pool - num_kv_cache_pools = 1 - - kv_cache_block_offsets = Tensor( - name=f'kv_cache_block_offsets', - dtype=trt.int32, - shape=[num_kv_cache_pools, -1, 2, -1], - dim_range=OrderedDict([ - ('num_kv_cache_pools', [num_kv_cache_pools]), - ('batch_size_beam_width', [bb_range]), - ('kv', [2]), - ('max_blocks_per_seq', max_blocks_per_seq_range), - ])) - host_kv_cache_block_offsets = Tensor( - name=f'host_kv_cache_block_offsets', - dtype=trt.int32, - shape=[num_kv_cache_pools, -1, 2, -1], - dim_range=OrderedDict([ - ('num_kv_cache_pools', [num_kv_cache_pools]), - ('batch_size_beam_width', [bb_range]), - ('kv', [2]), - ('max_blocks_per_seq', max_blocks_per_seq_range), - ])) - host_kv_cache_pool_pointers = Tensor( - name=f'host_kv_cache_pool_pointers', - dtype=trt.int64, - shape=[num_kv_cache_pools, 2], - dim_range=OrderedDict([ - ('num_pools_layers', [num_kv_cache_pools]), - ('num_pools_kv', [2]), - ])) - host_kv_cache_pool_mapping = Tensor( - name=f"host_kv_cache_pool_mapping", - dtype=trt.int32, - # 2: (Index of pool, Index of layer within pool) - shape=[num_pp_layers, 2], - dim_range=OrderedDict([ - ('pools_mapping', [num_pp_layers]), - ('layer_cache_pool_locator', [2]), - ])) - - # paged blocks for cross kv - cross_kv_cache_block_offsets = Tensor( - name=f'cross_kv_cache_block_offsets', - dtype=trt.int32, - shape=[num_kv_cache_pools, -1, 2, -1], - dim_range=OrderedDict([ - ('num_kv_cache_pools', [num_kv_cache_pools]), - ('batch_size_beam_width', [bb_range]), - ('kv', [2]), - ('max_cross_blocks_per_seq', - max_cross_blocks_per_seq_range), - ])) - host_cross_kv_cache_block_offsets = Tensor( - name=f'host_cross_kv_cache_block_offsets', - dtype=trt.int32, - shape=[num_kv_cache_pools, -1, 2, -1], - dim_range=OrderedDict([ - ('num_kv_cache_pools', [num_kv_cache_pools]), - ('batch_size_beam_width', [bb_range]), - ('kv', [2]), - ('max_cross_blocks_per_seq', - max_cross_blocks_per_seq_range), - ])) - host_cross_kv_cache_pool_pointers = Tensor( - name=f'host_cross_kv_cache_pool_pointers', - dtype=trt.int64, - shape=[num_kv_cache_pools, 2], - dim_range=OrderedDict([ - ('num_kv_cache_pools', [num_kv_cache_pools]), - ('num_pools', [2]), - ])) - host_cross_kv_cache_pool_mapping = Tensor( - name=f"host_cross_kv_cache_pool_mapping", - dtype=trt.int32, - # 2: (Index of pool, Index of layer within pool) - shape=[num_pp_layers, 2], - dim_range=OrderedDict([ - ('pools_mapping', [num_pp_layers]), - ('layer_cache_pool_locator', [2]), - ])) - - for i in layers_range: - past_key_value.append(None) - - kv_cache_params = KeyValueCacheParams( - past_key_value=past_key_value, - host_past_key_value_lengths=host_past_key_value_lengths, - host_max_attention_window_sizes=host_max_attention_window_sizes, - host_sink_token_length=host_sink_token_length, - cache_indirection=cache_indirection, - kv_cache_block_offsets=kv_cache_block_offsets, - host_kv_cache_block_offsets=host_kv_cache_block_offsets, - host_kv_cache_pool_pointers=host_kv_cache_pool_pointers, - host_kv_cache_pool_mapping=host_kv_cache_pool_mapping, - cross_kv_cache_block_offsets=cross_kv_cache_block_offsets, - host_cross_kv_cache_block_offsets= - host_cross_kv_cache_block_offsets, - host_cross_kv_cache_pool_pointers= - host_cross_kv_cache_pool_pointers, - host_cross_kv_cache_pool_mapping= - host_cross_kv_cache_pool_mapping, - ) - - attention_params = AttentionParams( - sequence_length=sequence_length, - context_lengths=context_lengths, - host_context_lengths=host_context_lengths, - max_context_length=max_decoder_input_len, - host_request_types=host_request_types, - encoder_input_lengths=encoder_input_lengths, - encoder_max_input_length=encoder_max_input_length, - host_runtime_perf_knobs=runtime_perf_knobs, - host_context_progress=context_progress) - - cross_kv_cache_gen = Tensor(name='cross_kv_cache_gen', - dtype=trt.bool, - shape=[1], - dim_range=OrderedDict([ - ('boolean', [1]), - ])) - cross_kv_reuse = None - num_heads = (self.num_heads + self.mapping.tp_size - - 1) // self.mapping.tp_size - cross_kv_out_dim = 2 * num_kv_heads * self.head_size - if self.skip_cross_kv: - if remove_input_padding: - cross_kv_reuse = Tensor( - name="cross_kv_reuse", - dtype=self._dtype, - shape=[-1, cross_kv_out_dim], - dim_range=OrderedDict([ - ("encoder_num_tokens", [encoder_num_tokens_range]), - ("encoder_kv_size", [cross_kv_out_dim]), - ]), - ) - else: - cross_kv_reuse = Tensor( - name="cross_kv_reuse", - dtype=self._dtype, - shape=[-1, -1, cross_kv_out_dim], - dim_range=OrderedDict([ - ("batch_size_beam_width_encoder", [bb_range]), - ("encoder_input_len", [encoder_input_len_range]), - ("encoder_kv_size", [cross_kv_out_dim]), - ]), - ) - - language_adapter_routings = None - if self.language_adapter_config: - language_adapter_routings = Tensor( - name="language_adapter_routings", - dtype=trt.int32, - shape=[-1, 1], - dim_range=OrderedDict([('decoder_num_tokens_range', - [decoder_num_tokens_range]), - ('routing_dim', [1])]), - ) - - result = { - 'decoder_input_ids': input_ids, - 'encoder_output': encoder_output, - 'position_ids': position_ids, - 'token_type_ids': token_type_ids, - 'use_cache': True, - 'attention_mask_params': attention_mask_params, - 'last_token_ids': last_token_ids, - 'kv_cache_params': kv_cache_params, - 'attention_params': attention_params, - 'hidden_states': hidden_states, - 'lora_params': lora_params, - 'cross_kv_cache_gen': cross_kv_cache_gen, - 'cross_kv_reuse': cross_kv_reuse, - 'language_adapter_routings': language_adapter_routings, - } - - return result - - def use_lora(self, lora_config: LoraConfig): - use_lora(self, lora_config, self.trtllm_modules_to_hf_modules) - - def precompute_relative_attention_bias(self, build_config): - if self.config.relative_attention and not self.use_implicit_relative_attention: - relative_attention_bias_builder = torch.ops.tensorrt_llm.relative_attention_bias - rel_attn_precomputed = torch.zeros( - (self.config.num_attention_heads // self.mapping.tp_size, - build_config.max_seq_len + 1, build_config.max_seq_len + 1), - dtype=str_dtype_to_torch(self.config.dtype), - device='cuda') - rel_attn_table = numpy_to_torch( - self.rel_attn_table.raw_value).to('cuda') - relative_attention_bias_builder( - rel_attn_precomputed, - rel_attn_table, - self.config.num_attention_heads // self.mapping.tp_size, - build_config.max_seq_len, - self.config.num_buckets, - False, - self.config.max_distance, - ) - for layer_idx in range(self.num_layers): - self.transformer.layers[ - layer_idx].self_attention.set_rel_attn_table( - build_config.max_seq_len, rel_attn_precomputed) - - -class WhisperEncoder(PretrainedModel): - - def __init__(self, config: PretrainedConfig): - super().__init__(config) - self._dtype = self.config.dtype - conv1 = Conv1d(config.n_mels, - config.hidden_size, - kernel_size=3, - padding=1, - dtype=self._dtype) - conv2 = Conv1d(config.hidden_size, - config.hidden_size, - kernel_size=3, - stride=2, - padding=1, - dtype=self._dtype) - self.transformer = EncDecTransformer( - 0, - self.config.hidden_size, - has_vocab_embeddings=False, - max_position_embeddings=self.config.max_position_embeddings, - has_position_embedding=self.config.has_position_embedding, - layernorm_type=LayerNormType.LayerNorm, - dtype=self.config.dtype, - has_model_final_layernorm=True, - is_decoder=False) - encoder_layers = ModuleList([ - EncoderLayer( - hidden_size=config.hidden_size, - ffn_hidden_size=config.hidden_size * 4, - num_attention_heads=config.num_attention_heads, - num_kv_heads=config.num_attention_heads, - head_size=config.hidden_size // config.num_attention_heads, - q_scaling=1.0, - has_attention_qkvo_bias=True, - has_mlp_bias=True, - hidden_act='gelu', - dtype=self._dtype) for _ in range(config.num_hidden_layers) - ]) - self.transformer.assign_module(encoder_layers, "layers") - self.transformer.assign_module(conv1, "conv1") - self.transformer.assign_module(conv2, "conv2") - self.max_audio_feature_seq_len = 3000 - self.downsample_factor = 2 - - def forward(self, - input_features: Tensor, - input_lengths=None, - position_ids=None): - if default_net().plugin_config.remove_input_padding: - # BXT,D -> 1,BxT,D -> 1,D,BxT - input_features = unsqueeze(input_features, 0) - input_features = transpose(input_features, 1, 2) - x_type = input_features.dtype - input_features = cast(input_features, self._dtype) - x = self.transformer.conv1(input_features) - x = gelu(x) - x = self.transformer.conv2(x) - x = cast(x, x_type) - x = gelu(x) - x = transpose(x, 2, 1) - x = x + cast(self.transformer.position_embedding(position_ids), x.dtype) - - if default_net().plugin_config.remove_input_padding: - #B,T,D -> BxT,D - x = x.view([-1, self.config.hidden_size]) - hidden_states = x - input_lengths = input_lengths // self.downsample_factor - for encoder_layer in self.transformer.layers: - hidden_states = encoder_layer(hidden_states, - input_lengths=input_lengths) - - x = hidden_states - x = self.transformer.ln_f(x) - x.mark_output('encoder_output', self._dtype) - return x - - def prepare_inputs(self, max_batch_size=16): - - bs_range = [1, (max_batch_size + 1) // 2, max_batch_size] - min_feat_len, optimal_feat_len = 10, 1000 # 100ms, 10s - inlen_range = [ - min_feat_len, optimal_feat_len, self.max_audio_feature_seq_len - ] - inlen_range_after_downsample = [ - min_feat_len // self.downsample_factor, - optimal_feat_len // self.downsample_factor, - self.max_audio_feature_seq_len // self.downsample_factor - ] - if not default_net().plugin_config.remove_input_padding: - x = Tensor(name="input_features", - dtype=self._dtype, - shape=[-1, self.config.n_mels, -1], - dim_range=OrderedDict([ - ("batch_size", [bs_range]), - ("feature_dim", [self.config.n_mels]), - ("feature_len_range", [inlen_range]), - ])) - position_ids = Tensor( - name='position_ids', - dtype=trt.int32, - shape=[-1, -1], - dim_range=OrderedDict([('batch_size', [bs_range]), - ('feature_len_downsample_range', - [inlen_range_after_downsample])]), - ) - else: - batch_seqlen_range = [ - 1, - (self.max_audio_feature_seq_len * max_batch_size + 1) // 2, - self.max_audio_feature_seq_len * max_batch_size, - ] - batch_seqlen_downsample_range = [ - 1, - (self.max_audio_feature_seq_len // self.downsample_factor * - max_batch_size + 1) // 2, - self.max_audio_feature_seq_len // self.downsample_factor * - max_batch_size, - ] - x = Tensor(name="input_features", - dtype=self._dtype, - shape=[-1, self.config.n_mels], - dim_range=OrderedDict([ - ("batch_seqlen_range", [batch_seqlen_range]), - ("feature_dim", [self.config.n_mels]), - ])) - position_ids = Tensor( - name='position_ids', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([('batch_seqlen_downsample_range', - [batch_seqlen_downsample_range])]), - ) - input_lengths = Tensor( - name="input_lengths", - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([("batch_size", [bs_range])]), - ) - - return { - 'input_features': x, - 'input_lengths': input_lengths, - 'position_ids': position_ids - } - - def precompute_relative_attention_bias(self, build_config): - pass diff --git a/tensorrt_llm/models/falcon/__init__.py b/tensorrt_llm/models/falcon/__init__.py deleted file mode 100644 index 71bf6d298d42..000000000000 --- a/tensorrt_llm/models/falcon/__init__.py +++ /dev/null @@ -1,14 +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. diff --git a/tensorrt_llm/models/falcon/config.py b/tensorrt_llm/models/falcon/config.py deleted file mode 100644 index 1ff2ff0391c0..000000000000 --- a/tensorrt_llm/models/falcon/config.py +++ /dev/null @@ -1,118 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -from typing import Optional, Union - -from ..._utils import get_hf_rope_theta -from ...mapping import Mapping -from ..convert_utils import infer_dtype -from ..modeling_utils import PretrainedConfig, QuantConfig - - -class FalconConfig(PretrainedConfig): - - def __init__(self, - *, - bias: bool = False, - parallel_attention: bool = False, - num_ln_in_parallel_attn: int | None = None, - new_decoder_architecture: bool = False, - rotary_base: float = 10000.0, - **kwargs): - self.bias = bias - self.parallel_attention = parallel_attention - self.num_ln_in_parallel_attn = num_ln_in_parallel_attn - self.new_decoder_architecture = new_decoder_architecture - self.rotary_base = rotary_base - super().__init__(**kwargs) - - def to_dict(self): - output = super().to_dict() - # Serialize the fields added in LLaMAConfig - output['bias'] = self.bias - output['parallel_attention'] = self.parallel_attention - output['new_decoder_architecture'] = self.new_decoder_architecture - return output - - @classmethod - def from_hugging_face( - cls, - hf_config_or_dir: Union[str, 'transformers.PretrainedConfig'], - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - import transformers - trust_remote_code = kwargs.pop('trust_remote_code', True) - - if isinstance(hf_config_or_dir, transformers.PretrainedConfig): - hf_config = hf_config_or_dir - else: - hf_config_dir = str(hf_config_or_dir) - hf_config = transformers.AutoConfig.from_pretrained( - hf_config_dir, trust_remote_code=trust_remote_code) - - # Falcon-7B config may not have num_kv_heads or n_head_kv. - # Although Falcon-180B uses GQA (num_kv_heads=8), its config - # has multi_query=True. - if getattr(hf_config, 'multi_query', False) and not getattr( - hf_config, 'new_decoder_architecture', False): - hf_config.num_kv_heads = 1 - - if hf_config.model_type == 'RefinedWeb': - # Case 1. Falcon-40B / Falcon-40B-instruct - # https://huggingface.co/tiiuae/falcon-40b/blob/main/config.json - hf_config.num_hidden_layers = hf_config.n_layer - hf_config.num_attention_heads = hf_config.n_head - hf_config.num_kv_heads = hf_config.n_head_kv - hf_config.new_decoder_architecture = True - elif hf_config.model_type == 'RefinedWebModel': - # Case 2. Falcon-7B / Falcon-7B-instruct - # https://huggingface.co/tiiuae/falcon-7b/blob/main/config.json - hf_config.num_hidden_layers = hf_config.n_layer - hf_config.num_attention_heads = hf_config.n_head - hf_config.num_kv_heads = 1 if hf_config.multi_query else hf_config.n_head - hf_config.new_decoder_architecture = False - elif hf_config.model_type != 'falcon': - raise ValueError("Shouldn't reach here.") - hf_config.model_type = 'falcon' - - dtype = infer_dtype(dtype, getattr(hf_config, 'torch_dtype', None)) - - return cls(architecture='FalconForCausalLM', - dtype=dtype, - num_hidden_layers=hf_config.num_hidden_layers, - num_attention_heads=hf_config.num_attention_heads, - num_key_value_heads=hf_config.num_kv_heads, - hidden_size=hf_config.hidden_size, - norm_epsilon=hf_config.layer_norm_epsilon, - vocab_size=hf_config.vocab_size, - position_embedding_type='alibi_with_scale' - if hf_config.alibi else 'rope_gpt_neox', - hidden_act='gelu', - bias=hf_config.bias, - parallel_attention=hf_config.parallel_attn, - num_ln_in_parallel_attn=getattr(hf_config, - 'num_ln_in_parallel_attn', - None), - new_decoder_architecture=hf_config.new_decoder_architecture, - max_position_embeddings=getattr(hf_config, - 'max_position_embeddings', - 2048), - rotary_base=get_hf_rope_theta(hf_config, 10000.0), - intermediate_size=getattr(hf_config, 'ffn_hidden_size', - None), - mapping=mapping, - quantization=quant_config, - **kwargs) diff --git a/tensorrt_llm/models/falcon/convert.py b/tensorrt_llm/models/falcon/convert.py deleted file mode 100644 index ef36e9e7383f..000000000000 --- a/tensorrt_llm/models/falcon/convert.py +++ /dev/null @@ -1,499 +0,0 @@ -import time -from typing import Dict, Optional - -import torch - -from ...quantization import QuantAlgo -from ..convert_utils import (get_weight, get_weight_and_bias, - iterate_shard_files, load_state_dict, - retrieved_layer_index_from_name) -from .config import FalconConfig - - -def split(weight: torch.Tensor, - tp_size: int, - rank: int = 0, - dim: int = 0) -> torch.Tensor: - if tp_size == 1: - return weight - elif weight.ndim == 1: - return torch.chunk(weight, tp_size)[rank].clone() - else: - return torch.chunk(weight, tp_size, dim=dim)[rank].clone() - - -def reorder_qkv_weight_or_bias(weight: torch.Tensor, - head_dim: int, - num_heads: int, - num_kv_heads: Optional[int] = None, - tp_size: int = 1, - is_bias: bool = False) -> torch.Tensor: - """ Reorder the qkv weight for TRT-LLM use. - - The shape of the fused QKV weights in HF is different from the shape that - TRT-LLM requires. In particular, the weight of HF consists of interleaved - q, k, v head weights, while that of TRT-LLM is contiguous. - HF : [q1, k1, v1, ..., qh, kh, vh] - TRT-LLM: [q1, ..., qh, k1, ..., kh, v1, vh] - where qi, vi, ki are weight vectors corresponding to attention head i. - It's similar to multi/grouped query attention cases. - - We reorder and split the weight of an attention layer to fit into TRT-LLM. - The reordered weight and bias will be - weight: (T, Qh * D + 2 * KVh * D, H) - bias : (T, Qh * D + 2 * KVh * D) - where T=tp_size, Qh=local_num_q_heads, KVh=local_num_kv_heads, D=head_dim, - H=hidden_dim. In the multi/grouped query attention, the number of K/V - attention heads are less than that of Q attention, so that K/V attention - heads may be shared across different ranks if necessary. - - For tensor parallelism, we use the first dimension to select the - corresponding weights. - """ - - # Query types and expected kv heads. - # - Conventional MHA: num_heads = num_kv_heads - # - Multi-Query Attention: num_kv_heads = 1 - # - Grouped-Query Attention: num_heads % num_kv_heads = 0 - num_kv_heads = num_kv_heads if num_kv_heads is not None else num_heads - assert num_heads % num_kv_heads == 0, \ - f'num_heads({num_heads}) must be divisible by ' \ - f'num_kv_heads({num_kv_heads})).' - - # The number of attention heads per group: N q head + 1 k head + 1 v head. - num_group_heads = num_heads // num_kv_heads + 2 - assert weight.shape[0] == num_kv_heads * num_group_heads * head_dim, \ - f'{weight.shape[0]} != {num_kv_heads} * {num_group_heads} * {head_dim}' - - qkv_in = num_heads * head_dim if not is_bias else 1 - - # Split Q/K/V weights - weight = weight.reshape(num_kv_heads, num_heads // num_kv_heads + 2, - head_dim, qkv_in) - q_w = weight[:, :-2, ...] # (nKV, num_heads // nKV, head_dim, qkv_in) - k_w = weight[:, -2:-1, ...] # (nKV, 1, head_dim, qkv_in) - v_w = weight[:, -1:, ...] # (nKV, 1, head_dim, qkv_in) - - if num_kv_heads < num_heads and num_kv_heads < tp_size: - # Duplicate K/V heads to make sure that each rank has at least one - # K/V heads. For instance, num_heads=8, num_kv_heads=2, tp_size=4, - # we will make the qkv weight as below. - # Orig: [q0 q1 q2 q3 k0 v0 q4 q5 q6 q7 k1 v0 v1] - # >>>> [[q0 q1 k0 v0], [q2 q3 k0 v0], [q4 q5 k1 v1], [q6 q7 k1 v1]] - assert tp_size % num_kv_heads == 0 - num_dups = tp_size // num_kv_heads - - # k_w and v_w have the same shape. - new_shape = (num_kv_heads, num_dups) + k_w.shape[2:] - k_w = torch.broadcast_to(k_w, size=new_shape) - v_w = torch.broadcast_to(v_w, size=new_shape) - - # Update the number of kv heads. - num_kv_heads = tp_size - - reordered = torch.concat( - [ - q_w.reshape(tp_size, num_heads // tp_size, head_dim, qkv_in), - k_w.reshape(tp_size, num_kv_heads // tp_size, head_dim, qkv_in), - v_w.reshape(tp_size, num_kv_heads // tp_size, head_dim, qkv_in), - ], - dim=1, - ) - - qkv_out = (num_heads + 2 * num_kv_heads) // tp_size * head_dim - return reordered.reshape((tp_size, qkv_out, -1)) - - -def split_qkv_weight(weight: torch.Tensor, - hidden_size: int, - num_heads: int, - tp_size: int, - rank: int, - is_bias: bool, - num_kv_heads: Optional[int] = None) -> torch.Tensor: - """ Splits the QKV matrix according to tensor parallelism """ - head_dim = hidden_size // num_heads - weight = reorder_qkv_weight_or_bias(weight, - head_dim=head_dim, - num_heads=num_heads, - num_kv_heads=num_kv_heads, - tp_size=tp_size, - is_bias=is_bias) - - # Copy a sliced tensor to prevent memory leak. A sliced tensor shares the - # memory buffer of the original tensor. So, returning without copying makes - # the buffer of a loaded "qkv" be referenced, resulting GC can't release - # those weights until the whole process ends. - if not is_bias: - return weight[rank, ...].clone() - else: - return weight[rank, ...].ravel().clone() - - -def split_matrix(weight: torch.Tensor, tp_size: int, rank: int, - dim: int) -> torch.Tensor: - return split(weight, tp_size, rank, dim=dim) - - -def get_tllm_linear_weight( - weight: torch.Tensor, - prefix: str, - bias: Optional[torch.Tensor] = None, - use_weight_only: bool = False, - plugin_weight_only_quant_type: torch.dtype = torch.int8 -) -> Dict[str, torch.Tensor]: - results = {} - if use_weight_only: - v = weight.t().contiguous() - processed_torch_weights, torch_weight_scales = \ - torch.ops.trtllm.symmetric_quantize_last_axis_of_batched_matrix( - v, plugin_weight_only_quant_type) - results[f'{prefix}.weight'] = processed_torch_weights - results[f'{prefix}.per_channel_scale'] = torch_weight_scales - else: - results[f'{prefix}.weight'] = weight - - if bias is not None: - results[f'{prefix}.bias'] = bias - - return results - - -def get_tllm_param( - param: torch.Tensor, - name: str, - use_weight_only: bool = False, - plugin_weight_only_quant_type: torch.dtype = torch.int8 -) -> Dict[str, torch.Tensor]: - results = {} - if name.endswith('.weight') and use_weight_only: - v = param.t().contiguous() - processed_torch_weights, torch_weight_scales = \ - torch.ops.trtllm.symmetric_quantize_last_axis_of_batched_matrix( - v, plugin_weight_only_quant_type) - results[name] = processed_torch_weights - results[name.replace('weight', - 'per_channel_scale')] = torch_weight_scales - else: - results[name] = param - - return results - - -def load_weights_from_hf_model(model, config: FalconConfig): - weights = {} - tik = time.time() - - model_params = dict(model.named_parameters()) - dtype = getattr(torch, config.dtype) - mapping = config.mapping - num_attention_heads = config.num_attention_heads - hidden_size = config.hidden_size - vocab_size = config.vocab_size - num_kv_heads = getattr(config, 'num_key_value_heads', num_attention_heads) - num_hidden_layers = config.num_hidden_layers - parallel_attention = config.parallel_attention - new_decoder_architecture = config.new_decoder_architecture - use_parallel_embedding = config.use_parallel_embedding - sharding_dim = config.embedding_sharding_dim - quant_algo = config.quantization.quant_algo - use_weight_only = quant_algo in [QuantAlgo.W8A16, QuantAlgo.W4A16] - if quant_algo == QuantAlgo.W8A16: - plugin_weight_only_quant_type = torch.int8 - elif quant_algo == QuantAlgo.W4A16: - plugin_weight_only_quant_type = torch.quint4x2 - else: - plugin_weight_only_quant_type = None - - layers_range = mapping.pp_layers(num_hidden_layers) - for l in layers_range: - prefix = f'transformer.h.{l}' - tllm_prex = f'transformer.layers.{l - layers_range[0]}' - qkv_weight, qkv_bias = get_weight_and_bias( - model_params, f'{prefix}.self_attention.query_key_value', dtype) - qkv_w = split_qkv_weight(qkv_weight, - hidden_size, - num_attention_heads, - mapping.tp_size, - mapping.tp_rank, - is_bias=False, - num_kv_heads=num_kv_heads) - if qkv_bias is None: - qkv_b = None - else: - qkv_b = split_qkv_weight(qkv_bias, - hidden_size, - num_attention_heads, - mapping.tp_size, - mapping.tp_rank, - is_bias=True, - num_kv_heads=num_kv_heads) - weights.update( - get_tllm_linear_weight(qkv_w, f'{tllm_prex}.attention.qkv', qkv_b, - use_weight_only, - plugin_weight_only_quant_type)) - - attn_dense_weight, attn_dense_bias = get_weight_and_bias( - model_params, f'{prefix}.self_attention.dense', dtype) - attn_dense_w = split_matrix(attn_dense_weight, - mapping.tp_size, - mapping.tp_rank, - dim=1) - weights.update( - get_tllm_linear_weight(attn_dense_w, f'{tllm_prex}.attention.dense', - attn_dense_bias, use_weight_only, - plugin_weight_only_quant_type)) - - mlp_fc_weight, mlp_fc_bias = get_weight_and_bias( - model_params, f'{prefix}.mlp.dense_h_to_4h', dtype) - mlp_fc_w = split_matrix(mlp_fc_weight, - mapping.tp_size, - mapping.tp_rank, - dim=0) - if mlp_fc_bias is None: - mlp_fc_b = None - else: - mlp_fc_b = split_matrix(mlp_fc_bias, - mapping.tp_size, - mapping.tp_rank, - dim=0) - weights.update( - get_tllm_linear_weight(mlp_fc_w, f'{tllm_prex}.mlp.fc', mlp_fc_b, - use_weight_only, - plugin_weight_only_quant_type)) - - mlp_proj_weight, mlp_proj_bias = get_weight_and_bias( - model_params, f'{prefix}.mlp.dense_4h_to_h', dtype) - mlp_proj_w = split_matrix(mlp_proj_weight, - mapping.tp_size, - mapping.tp_rank, - dim=1) - weights.update( - get_tllm_linear_weight(mlp_proj_w, f'{tllm_prex}.mlp.proj', - mlp_proj_bias, use_weight_only, - plugin_weight_only_quant_type)) - - if new_decoder_architecture: - input_ln_weight, input_ln_bias = get_weight_and_bias( - model_params, f'{prefix}.ln_attn', dtype) - if input_ln_weight is None: - input_ln_weight, input_ln_bias = get_weight_and_bias( - model_params, f'{prefix}.input_layernorm', dtype) - weights[f'{tllm_prex}.input_layernorm.weight'] = input_ln_weight - if input_ln_bias is not None: - weights[f'{tllm_prex}.input_layernorm.bias'] = input_ln_bias - - mlp_ln_weight, mlp_ln_bias = get_weight_and_bias( - model_params, f'{prefix}.ln_mlp', dtype) - if mlp_ln_weight is not None: - weights[f'{tllm_prex}.mlp_layernorm.weight'] = mlp_ln_weight - if mlp_ln_bias is not None: - weights[f'{tllm_prex}.mlp_layernorm.bias'] = mlp_ln_bias - else: - input_ln_weight, input_ln_bias = get_weight_and_bias( - model_params, f'{prefix}.input_layernorm', dtype) - weights[f'{tllm_prex}.input_layernorm.weight'] = input_ln_weight - if input_ln_bias is not None: - weights[f'{tllm_prex}.input_layernorm.bias'] = input_ln_bias - - if not parallel_attention: - post_ln_weight, post_ln_bias = get_weight_and_bias( - model_params, f'{prefix}.post_attention_layernorm', dtype) - if post_ln_weight is not None: - weights[ - f'{tllm_prex}.post_layernorm.weight'] = post_ln_weight - if post_ln_bias is not None: - weights[f'{tllm_prex}.post_layernorm.bias'] = post_ln_bias - - embed_w = get_weight(model_params, 'transformer.word_embeddings', dtype) - if mapping.is_first_pp_rank(): - if not use_parallel_embedding: - weights['transformer.vocab_embedding.weight'] = embed_w - else: - if sharding_dim == 0: - assert vocab_size % mapping.tp_size == 0 - else: - assert hidden_size % mapping.tp_size == 0 - weights['transformer.vocab_embedding.weight'] = split_matrix( - embed_w, mapping.tp_size, mapping.tp_rank, sharding_dim) - - if mapping.is_last_pp_rank(): - lm_head = get_weight(model_params, 'lm_head', dtype) - if lm_head is None: - # No lm_head in the checkpoint, cloning word_embedding. - lm_head = embed_w.clone() - weights['lm_head.weight'] = split_matrix(lm_head, - mapping.tp_size, - mapping.tp_rank, - dim=0) - ln_f_w, ln_f_b = get_weight_and_bias(model_params, 'transformer.ln_f', - dtype) - weights['transformer.ln_f.weight'] = ln_f_w - if ln_f_b is not None: - weights['transformer.ln_f.bias'] = ln_f_b - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Weights loaded. Total time: {t}') - return weights - - -def load_weights_from_hf_by_shard(model_dir: str, config: FalconConfig): - weights = {} - tik = time.time() - - dtype = getattr(torch, config.dtype) - mapping = config.mapping - num_attention_heads = config.num_attention_heads - hidden_size = config.hidden_size - vocab_size = config.vocab_size - num_kv_heads = getattr(config, 'num_key_value_heads', num_attention_heads) - num_hidden_layers = config.num_hidden_layers - use_weight_only = config.quantization.quant_algo in [ - QuantAlgo.W8A16, QuantAlgo.W4A16 - ] - if config.quantization.quant_algo == QuantAlgo.W8A16: - plugin_weight_only_quant_type = torch.int8 - elif config.quantization.quant_algo == QuantAlgo.W4A16: - plugin_weight_only_quant_type = torch.quint4x2 - else: - plugin_weight_only_quant_type = None - - layers_range = mapping.pp_layers(num_hidden_layers) - for model_file in iterate_shard_files(model_dir, mapping.tp_rank): - state_dict = load_state_dict(model_file, dtype) - for name, param in state_dict.items(): - l = retrieved_layer_index_from_name(name) - if l is not None: - if l not in layers_range: - continue - prefix = f'transformer.layers.{l-layers_range[0]}' - - if 'self_attention.query_key_value' in name: - if name.endswith('weight'): - qkv_w = split_qkv_weight(param, - hidden_size, - num_attention_heads, - mapping.tp_size, - mapping.tp_rank, - is_bias=False, - num_kv_heads=num_kv_heads) - weights.update( - get_tllm_param(qkv_w, - f'{prefix}.attention.qkv.weight', - use_weight_only, - plugin_weight_only_quant_type)) - else: - qkv_b = split_qkv_weight(param, - hidden_size, - num_attention_heads, - mapping.tp_size, - mapping.tp_rank, - is_bias=True, - num_kv_heads=num_kv_heads) - weights.update( - get_tllm_param(qkv_b, - f'{prefix}.attention.qkv.bias', - use_weight_only, - plugin_weight_only_quant_type)) - - elif 'self_attention.dense' in name: - if name.endswith('weight'): - attn_dense_w = split_matrix(param, - mapping.tp_size, - mapping.tp_rank, - dim=1) - weights.update( - get_tllm_param(attn_dense_w, - f'{prefix}.attention.dense.weight', - use_weight_only, - plugin_weight_only_quant_type)) - else: - weights.update( - get_tllm_param(param, - f'{prefix}.attention.dense.bias', - use_weight_only, - plugin_weight_only_quant_type)) - - elif 'mlp.dense_h_to_4h' in name: - if name.endswith('weight'): - mlp_fc_w = split_matrix(param, - mapping.tp_size, - mapping.tp_rank, - dim=0) - weights.update( - get_tllm_param(mlp_fc_w, f'{prefix}.mlp.fc.weight', - use_weight_only, - plugin_weight_only_quant_type)) - else: - mlp_fc_b = split_matrix(param, - mapping.tp_size, - mapping.tp_rank, - dim=0) - weights.update( - get_tllm_param(mlp_fc_b, f'{prefix}.mlp.fc.bias', - use_weight_only, - plugin_weight_only_quant_type)) - - elif 'mlp.dense_4h_to_h' in name: - if name.endswith('weight'): - mlp_proj_w = split_matrix(param, - mapping.tp_size, - mapping.tp_rank, - dim=1) - weights.update( - get_tllm_param(mlp_proj_w, - f'{prefix}.mlp.proj.weight', - use_weight_only, - plugin_weight_only_quant_type)) - else: - weights.update( - get_tllm_param(param, f'{prefix}.mlp.proj.bias', - use_weight_only, - plugin_weight_only_quant_type)) - - elif 'ln_attn' in name or 'input_layernorm' in name: - if name.endswith('weight'): - weights[f'{prefix}.input_layernorm.weight'] = param - else: - weights[f'{prefix}.input_layernorm.bias'] = param - elif 'ln_mlp' in name: - if name.endswith('weight'): - weights[f'{prefix}.mlp_layernorm.weight'] = param - else: - weights[f'{prefix}.mlp_layernorm.bias'] = param - elif 'post_attention_layernorm' in name: - if name.endswith('weight'): - weights[f'{prefix}.post_layernorm.weight'] = param - else: - weights[f'{prefix}.post_layernorm.bias'] = param - elif 'word_embeddings' in name: - if mapping.is_first_pp_rank(): - if not config.use_parallel_embedding: - weights['transformer.vocab_embedding.weight'] = param - else: - if config.embedding_sharding_dim == 0: - assert vocab_size % mapping.tp_size == 0 - else: - assert hidden_size % mapping.tp_size == 0 - weights[ - 'transformer.vocab_embedding.weight'] = split_matrix( - param, mapping.tp_size, mapping.tp_rank, - config.embedding_sharding_dim) - if mapping.is_last_pp_rank(): - weights['lm_head.weight'] = split_matrix(param, - mapping.tp_size, - mapping.tp_rank, - dim=0) - elif 'ln_f' in name: - if mapping.is_last_pp_rank(): - if name.endswith('weight'): - weights['transformer.ln_f.weight'] = param - else: - weights['transformer.ln_f.bias'] = param - del state_dict - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Weights loaded. Total time: {t}') - return weights diff --git a/tensorrt_llm/models/falcon/model.py b/tensorrt_llm/models/falcon/model.py deleted file mode 100644 index 01e523fa66c9..000000000000 --- a/tensorrt_llm/models/falcon/model.py +++ /dev/null @@ -1,273 +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. -from typing import Optional, Union - -from ..._utils import pad_vocab_size -from ...functional import Tensor, allreduce, recv, send -from ...layers import (MLP, Attention, AttentionMaskType, ColumnLinear, - Embedding, LayerNorm) -from ...mapping import Mapping -from ...module import Module -from ..modeling_utils import (DecoderLayerList, DecoderModelForCausalLM, - QuantConfig) -from .config import FalconConfig -from .convert import load_weights_from_hf_by_shard, load_weights_from_hf_model - - -class FalconDecoderLayer(Module): - - def __init__(self, config: FalconConfig, layer_idx: int): - super().__init__() - self.layer_idx = layer_idx - self.config = config - - hidden_size = config.hidden_size - dtype = config.dtype - tp_group = config.mapping.tp_group - tp_size = config.mapping.tp_size - tp_rank = config.mapping.tp_rank - layernorm_epsilon = config.norm_epsilon - - self.input_layernorm = LayerNorm(normalized_shape=hidden_size, - eps=layernorm_epsilon, - dtype=dtype) - - self.new_decoder_architecture = config.new_decoder_architecture - self.parallel_attn = config.parallel_attention - self.num_ln_in_parallel_attn = config.num_ln_in_parallel_attn - if self.num_ln_in_parallel_attn is None and self.new_decoder_architecture: - self.num_ln_in_parallel_attn = 2 - if self.is_parallel_attention: - # Not to apply allreduce inside the Attention/MLP layers. - # allreduce applies after those layer. - tp_group = None - layers_range = config.mapping.pp_layers(config.num_hidden_layers) - local_layer_idx = layer_idx - layers_range[0] - self.attention = Attention( - local_layer_idx=local_layer_idx, - hidden_size=hidden_size, - num_attention_heads=config.num_attention_heads, - num_kv_heads=config.num_key_value_heads, - max_position_embeddings=config.max_position_embeddings, - attention_mask_type=AttentionMaskType.causal, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - tp_rank=tp_rank, - bias=config.bias, - position_embedding_type=config.position_embedding_type, - rotary_embedding_base=config.rotary_base, - quant_mode=config.quantization.quant_mode, - ) - - mlp_hidden_size = hidden_size * 4 if config.intermediate_size is None else config.intermediate_size - - if self.new_decoder_architecture and self.num_ln_in_parallel_attn == 2: - # Layernorm before MLP. - self.mlp_layernorm = LayerNorm(normalized_shape=hidden_size, - eps=layernorm_epsilon, - dtype=dtype) - else: - self.mlp_layernorm = None - self.mlp = MLP( - hidden_size=hidden_size, - ffn_hidden_size=mlp_hidden_size, - hidden_act=config.hidden_act, - dtype=dtype, - bias=config.bias, - tp_group=tp_group, - tp_size=tp_size, - quant_mode=config.quantization.quant_mode, - ) - if self.is_parallel_attention: - self.post_layernorm = None - else: - self.post_layernorm = LayerNorm(normalized_shape=hidden_size, - dtype=dtype) - - @property - def is_parallel_attention(self): - return self.new_decoder_architecture or self.parallel_attn - - def forward(self, - hidden_states: Tensor, - attention_mask=None, - use_cache=False, - kv_cache_params=None, - attention_params=None): - assert isinstance(hidden_states, Tensor) - - residual = hidden_states - - if self.new_decoder_architecture and self.num_ln_in_parallel_attn == 2: - mlp_ln_output = self.mlp_layernorm(hidden_states) - hidden_states = self.input_layernorm(hidden_states) - input_ln_output = hidden_states - attention_output = self.attention(hidden_states, - attention_mask=attention_mask, - use_cache=use_cache, - kv_cache_params=kv_cache_params, - attention_params=attention_params) - - if use_cache: - attention_output, presents = attention_output - - if not self.new_decoder_architecture: - if self.parallel_attn: - hidden_states = input_ln_output - else: - hidden_states = residual + attention_output - residual = hidden_states - hidden_states = self.post_layernorm(hidden_states) - elif self.num_ln_in_parallel_attn == 2: - hidden_states = mlp_ln_output - - if (self.new_decoder_architecture and self.parallel_attn - and self.num_ln_in_parallel_attn == 1): - hidden_states = input_ln_output - - hidden_states = self.mlp(hidden_states) - - if self.is_parallel_attention: - hidden_states = hidden_states + attention_output - if self.config.mapping.tp_size > 1: - hidden_states = allreduce(hidden_states, - self.config.mapping.tp_group) - - hidden_states = residual + hidden_states - if use_cache: - return hidden_states, presents - return hidden_states - - -class FalconModel(Module): - - def __init__(self, config: FalconConfig): - super().__init__() - self.config = config - if config.mapping.is_first_pp_rank(): - self.vocab_embedding = Embedding(config.vocab_size, - config.hidden_size, - dtype=config.dtype) - - self.layers = DecoderLayerList(FalconDecoderLayer, config) - if config.mapping.is_last_pp_rank(): - self.ln_f = LayerNorm(normalized_shape=config.hidden_size, - dtype=config.dtype) - - def forward(self, - input_ids: Tensor, - position_ids=None, - use_cache=False, - attention_mask=None, - kv_cache_params=None, - attention_params=None, - hidden_states=None): - if self.config.mapping.is_first_pp_rank(): - hidden_states = self.vocab_embedding(input_ids) - else: - hidden_states = recv(hidden_states, - self.config.mapping.prev_pp_rank()) - - hidden_states = self.layers(hidden_states, - use_cache=use_cache, - attention_mask=attention_mask, - kv_cache_params=kv_cache_params, - attention_params=attention_params) - - if use_cache: - hidden_states, presents = hidden_states - - if self.config.mapping.is_last_pp_rank(): - hidden_states = self.ln_f(hidden_states) - else: - hidden_states = send(hidden_states, - self.config.mapping.next_pp_rank()) - - if use_cache: - return (hidden_states, tuple(presents)) - return hidden_states - - -class FalconForCausalLM(DecoderModelForCausalLM): - config_class = FalconConfig - - def __init__(self, config: FalconConfig): - self.check_config(config) - transformer = FalconModel(config) - - if config.mapping.is_last_pp_rank(): - vocab_size_padded = pad_vocab_size(config.vocab_size, - config.mapping.tp_size) - lm_head = ColumnLinear(config.hidden_size, - vocab_size_padded, - bias=False, - dtype=config.dtype, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size, - gather_output=True) - else: - lm_head = None - super().__init__(config, transformer, lm_head) - - def check_config(self, config): - config.set_if_not_exist('bias', True) - config.set_if_not_exist('new_decoder_architecture', False) - config.set_if_not_exist('parallel_attention', False) - - @classmethod - def from_hugging_face( - cls, - hf_model_or_dir: Union[str, 'transformers.PreTrainedModel'], - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - ''' Create a FalconForCausalLM object from give parameters - ''' - import transformers - - load_by_shard = kwargs.pop('load_by_shard', False) - # load_model_on_cpu is ignored here, since specify target device_map will fail when workers > 1. - - assert hf_model_or_dir is not None - use_preloading = isinstance(hf_model_or_dir, - transformers.PreTrainedModel) - if use_preloading: - hf_model = hf_model_or_dir - hf_config_or_dir = hf_model.config - else: - hf_model_dir = hf_model_or_dir - hf_config_or_dir = hf_model_or_dir - - config = FalconConfig.from_hugging_face(hf_config_or_dir, - dtype=dtype, - mapping=mapping, - quant_config=quant_config, - **kwargs) - - if use_preloading: - assert not load_by_shard - weights = load_weights_from_hf_model(hf_model, config) - elif load_by_shard: - weights = load_weights_from_hf_by_shard(hf_model_dir, config) - else: - hf_model = transformers.AutoModelForCausalLM.from_pretrained( - hf_model_dir, dtype='auto') - weights = load_weights_from_hf_model(hf_model, config) - - model = cls(config) - model.load(weights) - return model diff --git a/tensorrt_llm/models/gemma/__init__.py b/tensorrt_llm/models/gemma/__init__.py deleted file mode 100644 index 71bf6d298d42..000000000000 --- a/tensorrt_llm/models/gemma/__init__.py +++ /dev/null @@ -1,14 +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. diff --git a/tensorrt_llm/models/gemma/config.py b/tensorrt_llm/models/gemma/config.py deleted file mode 100644 index 5755bce17aee..000000000000 --- a/tensorrt_llm/models/gemma/config.py +++ /dev/null @@ -1,224 +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. -from json import loads -from pathlib import Path -from typing import TYPE_CHECKING, Optional, Union - -from tensorrt_llm._utils import get_hf_rope_theta -from tensorrt_llm.functional import PositionEmbeddingType -from tensorrt_llm.logger import logger -from tensorrt_llm.mapping import Mapping -from tensorrt_llm.models.convert_utils import infer_dtype -from tensorrt_llm.models.modeling_utils import (Gemma2ConfigGroup, - Gemma3ConfigGroup, - PretrainedConfig, QuantConfig) - -if TYPE_CHECKING: - from os import PathLike - - import transformers - - HfConfigOrDir = Union[str, PathLike, transformers.PretrainedConfig] - -GEMMA_ARCHITECTURE = "GemmaForCausalLM" -GEMMA2_ARCHITECTURE = "Gemma2ForCausalLM" -GEMMA3_ARCHITECTURE = "Gemma3ForCausalLM" - - -class GemmaConfig(PretrainedConfig): - - def __init__( - self, - *, - architecture: str, - rotary_base: float = 10000.0, - rotary_scaling: Optional[dict] = None, - attn_bias: bool = False, - mlp_bias: bool = False, - position_embedding_type: PositionEmbeddingType = PositionEmbeddingType. - rope_gpt_neox, - query_pre_attn_scalar: Optional[int] = None, - final_logit_softcapping: Optional[float] = None, - attn_logit_softcapping: Optional[float] = None, - mapping: Optional[Union[Mapping, dict]] = None, - _sliding_window_pattern: int = None, - rope_local_base_freq: int = None, - sliding_window: int = None, - **kwargs, - ): - use_parallel_embedding = False - if mapping: - use_parallel_embedding = mapping.tp_size > 1 if isinstance( - mapping, Mapping) else mapping["tp_size"] > 1 - if use_parallel_embedding != kwargs.pop("use_parallel_embedding", None): - """ - We always pass `bool(mapping.tp_size > 1)` - the passed value is `False` by default, and ignored either way. - We can't just raise an exception here, because this will force the user to explicitly pass `LLM(use_parallel_embedding=True)`. - """ - logger.debug( - f"Using `use_parallel_embedding={use_parallel_embedding}` for Gemma" - ) - - super().__init__( - architecture=architecture, - use_parallel_embedding=use_parallel_embedding, - rotary_base=rotary_base, - attn_bias=attn_bias, - mlp_bias=mlp_bias, - position_embedding_type=position_embedding_type, - mapping=mapping, - **kwargs, - ) - self.rotary_base = rotary_base - self.rotary_scaling = rotary_scaling - self.attn_bias = attn_bias - self.mlp_bias = mlp_bias - - self.inter_layernorms = False - if self.is_gemma_2 or self.is_gemma_3: - self.inter_layernorms = True - assert query_pre_attn_scalar is not None, "Gemma2 / Gemma3 models must configure `query_pre_attn_scalar`" - self.query_pre_attn_scalar = query_pre_attn_scalar - self.final_logit_softcapping = final_logit_softcapping - if self.is_gemma_2: - self.attn_logit_softcapping = attn_logit_softcapping - if self.is_gemma_3: - self._sliding_window_pattern = _sliding_window_pattern - self.rope_local_base_freq = rope_local_base_freq - self.sliding_window = sliding_window - - GEMMA_ADDED_FIELDS = { - "rotary_base", "rotary_scaling", "attn_bias", "mlp_bias", - "inter_layernorms" - } - GEMMA2_ADDED_FIELDS = Gemma2ConfigGroup.keys() - GEMMA3_ADDED_FIELDS = Gemma3ConfigGroup.keys() - VERBATIM = { - "num_hidden_layers", "num_attention_heads", "hidden_size", - "intermediate_size", "vocab_size", "max_position_embeddings", - "hidden_act", "use_parallel_embedding" - } | GEMMA2_ADDED_FIELDS | GEMMA3_ADDED_FIELDS - - @property - def is_gemma_2(self) -> bool: - return self.architecture == GEMMA2_ARCHITECTURE - - def gemma2_config(self): - if self.is_gemma_2: - return self.get_config_group(Gemma2ConfigGroup) - return None - - @property - def is_gemma_3(self) -> bool: - return self.architecture == GEMMA3_ARCHITECTURE - - def gemma3_config(self): - if self.is_gemma_3: - return self.get_config_group(Gemma3ConfigGroup) - return None - - def to_dict(self): - """Serialize the fields added in GemmaConfig""" - - return { - **super().to_dict(), - **{ - f: getattr(self, f) - for f in self.GEMMA_ADDED_FIELDS - }, - **({ - f: getattr(self, f) - for f in self.GEMMA2_ADDED_FIELDS - } if self.is_gemma_2 else {}), - **({ - f: getattr(self, f) - for f in self.GEMMA3_ADDED_FIELDS - } if self.is_gemma_3 else {}) - } - - @staticmethod - def get_hf_config(config_dir: "Union[str, PathLike]"): - import transformers - model_type = loads( - (Path(config_dir) / "config.json").read_text())["model_type"] - HFConfigClass = { - "gemma": transformers.GemmaConfig, - "gemma2": transformers.Gemma2Config, - "gemma3_text": transformers.Gemma3TextConfig, - }[model_type] - hf_config = HFConfigClass.from_pretrained(config_dir) - return hf_config - - @classmethod - def from_hugging_face( - cls, - hf_config_or_dir: "HfConfigOrDir", - dtype: str = "auto", - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs, - ) -> "GemmaConfig": - import transformers - if isinstance(hf_config_or_dir, transformers.PretrainedConfig): - hf_config = hf_config_or_dir - else: - hf_config = cls.get_hf_config(hf_config_or_dir) - - dtype = infer_dtype(dtype, getattr(hf_config, 'torch_dtype', None)) - - assert isinstance(quant_config, QuantConfig) or quant_config is None - assert isinstance(mapping, Mapping) or mapping is None - - # transformers 5.x moved Gemma 3 RoPE fields into a nested - # rope_parameters dict. Resolve the global and sliding-window thetas - # from either the legacy top-level attributes or the nested form. - rope_parameters = getattr(hf_config, "rope_parameters", None) or {} - full_rope = rope_parameters.get("full_attention") or {} - sliding_rope = rope_parameters.get("sliding_attention") or {} - - rotary_base = getattr(hf_config, "rope_theta", None) - if rotary_base is None: - rotary_base = full_rope.get("rope_theta") - if rotary_base is None: - rotary_base = get_hf_rope_theta(hf_config, 10000.0) - - verbatim_fields = { - k: v - for k, v in hf_config.to_dict().items() if k in cls.VERBATIM - } - # Gemma 3's local RoPE base is no longer a top-level attribute in - # transformers 5.x; pull it from rope_parameters.sliding_attention. - if ("rope_local_base_freq" in cls.VERBATIM - and verbatim_fields.get("rope_local_base_freq") is None): - local_freq = getattr(hf_config, "rope_local_base_freq", None) - if local_freq is None: - local_freq = sliding_rope.get("rope_theta") - if local_freq is not None: - verbatim_fields["rope_local_base_freq"] = local_freq - - return cls( - architecture=hf_config.architectures[0], - dtype=dtype, - head_size=hf_config.head_dim, - norm_epsilon=hf_config.rms_norm_eps, - num_key_value_heads=getattr(hf_config, "num_key_value_heads", - hf_config.num_attention_heads), - rotary_base=rotary_base, - rotary_scaling=getattr(hf_config, "rotary_scaling", None), - quantization=quant_config, - mapping=mapping, - **verbatim_fields, - **kwargs, - ) diff --git a/tensorrt_llm/models/gemma/convert.py b/tensorrt_llm/models/gemma/convert.py deleted file mode 100644 index 5aa40d9fd9c5..000000000000 --- a/tensorrt_llm/models/gemma/convert.py +++ /dev/null @@ -1,926 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -import json -import math -import os -import re -import time -from argparse import Namespace -from dataclasses import dataclass -from pathlib import Path -from types import SimpleNamespace -from typing import TYPE_CHECKING, Dict, Optional, Union - -import h5py -import numpy -import numpy as np -import sentencepiece as sp -import torch - -import tensorrt_llm -import tensorrt_llm.models.modeling_utils as tllm_utils -from tensorrt_llm._utils import (np_bfloat16, numpy_to_torch, - str_dtype_to_torch, torch_to_numpy) -from tensorrt_llm.logger import logger -from tensorrt_llm.models.convert_utils import load_calib_dataset -from tensorrt_llm.models.gemma.config import GemmaConfig -from tensorrt_llm.models.gemma.smoothquant import (capture_activation_range, - convert_hf_model, - create_model_from_config, - smooth_model) -from tensorrt_llm.models.gemma.utils import params as gemma_params -from tensorrt_llm.models.gemma.weight import (dummy_weights_awq, - load_from_fp8_gemma, - quantize_fp8_weights) -from tensorrt_llm.quantization.mode import QuantAlgo - -Weights = Dict[str, torch.Tensor] -Flattened = Dict[str, np.ndarray] - -if TYPE_CHECKING: - import transformers - - -class JAXParser: - - def load_parameters(self, - checkpoint_path: Path, - load_model_on_cpu: bool = False) -> Weights: - checkpoint_path = checkpoint_path.absolute() - os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false" - return gemma_params.nest_params( - gemma_params.param_remapper( - gemma_params.load_params(checkpoint_path))) - - def embedding_weights(self, ckpt_params) -> np.ndarray: - return ckpt_params["transformer"]["embedder"]["input_embedding"] - - def get_config(self, checkpoint_path, ckpt_params, - num_embed) -> SimpleNamespace: - from tensorrt_llm.models.gemma.utils.transformer import \ - TransformerConfig - - return TransformerConfig.from_params(ckpt_params, num_embed=num_embed) - - def rename_to_trt_llm(self, name: str) -> Optional[str]: - """Rename a gemma parameter name by the corresponding TRT-LLM style name.""" - prefix, name = name.split(".", maxsplit=1) - assert prefix == "transformer" - sub_patterns = ( - (r"embedder.input_embedding", r"vocab_embedding.weight"), - (r"layer_(\d+).pre_attention_norm.scale", - r"layers.\1.input_layernorm.weight"), - (r"layer_(\d+).attn.q_einsum.w", r"layers.\1.attention.qkv.weight"), - (r"layer_(\d+).attn.kv_einsum.w", - None), # drop as kv will be concatenated with q - (r"layer_(\d+).attn.qkv_einsum.w", - r"layers.\1.attention.qkv.weight"), - (r"layer_(\d+).attn.attn_vec_einsum.w", - r"layers.\1.attention.dense.weight"), - (r"layer_(\d+).mlp.gating_einsum", r"layers.\1.mlp.fc.weight"), - (r"layer_(\d+).mlp.linear", r"layers.\1.mlp.proj.weight"), - (r"layer_(\d+).pre_ffw_norm.scale", - r"layers.\1.post_layernorm.weight"), - (r"final_norm.scale", r"ln_f.weight"), - ) - - for source, target in sub_patterns: - if re.match(source, name): - if target is None: - return target - else: - name = re.sub(source, target, name) - return ".".join((prefix, name)) - else: - raise ValueError(f"Don't know how to rename {prefix}.{name}") - - def flatten_params(self, params) -> Flattened: - import flax.traverse_util - - new_params = flax.traverse_util.flatten_dict(params, sep=".") - # if the dtype is bfloat16, cast to float32 - for k in new_params: - if new_params[k].dtype != np.float32 and new_params[ - k].dtype != np.float16: - new_params[k] = new_params[k].astype(np.float32) - return new_params - - -class KerasParser: - - def load_parameters(self, - checkpoint_path: Path, - load_model_on_cpu: bool = False) -> h5py.File: - checkpoint_path = checkpoint_path.absolute() - config_file = "config.json" - weights_file = json.load(open(checkpoint_path / config_file))["weights"] - h5_path = checkpoint_path / weights_file - return h5py.File(h5_path, "r") - - def embedding_weights(self, ckpt_params) -> np.ndarray: - return np.array(ckpt_params["layers/reversible_embedding/vars/0"]) - - def get_config(self, checkpoint_path, ckpt_params, - num_embed) -> SimpleNamespace: - checkpoint_path = checkpoint_path.absolute() - config_file = "config.json" - config_old = json.load(open(checkpoint_path / config_file))["config"] - config_new = SimpleNamespace() - config_new.num_layers = config_old["num_layers"] - config_new.num_embed = config_old["vocabulary_size"] - config_new.embed_dim = config_old["hidden_dim"] - config_new.hidden_dim = config_old["intermediate_dim"] // 2 - config_new.num_heads = config_old["num_query_heads"] - config_new.head_dim = config_old["head_dim"] - config_new.num_kv_heads = config_old["num_key_value_heads"] - return config_new - - def rename_to_trt_llm(self, name: str) -> Optional[str]: - """Rename a gemma parameter name by the corresponding TRT-LLM style name.""" - prefix = "transformer" - name = name.replace("/gemma_decoder_block/", "/gemma_decoder_block_0/") - sub_patterns = ( - (r"layers/reversible_embedding/vars/0", r"vocab_embedding.weight"), - (r"layers/gemma_decoder_block_(\d+)/pre_attention_norm/vars/0", - r"layers.\1.input_layernorm.weight"), - (r"layers/gemma_decoder_block_(\d+)/attention/query_dense/vars/0", - r"layers.\1.attention.qkv.weight"), - (r"layers/gemma_decoder_block_(\d+)/attention/key_dense/vars/0", - None), # drop as k will be concatenated with q - (r"layers/gemma_decoder_block_(\d+)/attention/value_dense/vars/0", - None), # drop as v will be concatenated with q - (r"layers/gemma_decoder_block_(\d+)/attention/output_dense/vars/0", - r"layers.\1.attention.dense.weight"), - (r"layers/gemma_decoder_block_(\d+)/gating_ffw/vars/0", - r"layers.\1.mlp.fc.weight"), - (r"layers/gemma_decoder_block_(\d+)/gating_ffw_2/vars/0", - None), # merged with above - (r"layers/gemma_decoder_block_(\d+)/ffw_linear/vars/0", - r"layers.\1.mlp.proj.weight"), - (r"layers/gemma_decoder_block_(\d+)/pre_ffw_norm/vars/0", - r"layers.\1.post_layernorm.weight"), - (r"layers/rms_normalization/vars/0", r"ln_f.weight"), - (r"optimizer/vars/(\d+)", None), # Not used - ) - - for source, target in sub_patterns: - if re.match(source, name): - if target is None: - return target - else: - name = re.sub(source, target, name) - return ".".join((prefix, name)) - else: - raise ValueError(f"Don't know how to rename {prefix}.{name}") - - def flatten_params(self, params) -> Flattened: - f_params = {} - - def walk(name, obj): - if isinstance(obj, h5py.Dataset): - if obj.dtype == "|V2": - # bfloat16 case - f_params[name] = torch_to_numpy( - numpy_to_torch(np.array(obj).astype(np_bfloat16)).to( - torch.float32)) - else: - f_params[name] = np.array(obj) - - params.visititems(walk) - return f_params - - -class TorchParser: - - def load_parameters(self, - checkpoint_path: Path, - load_model_on_cpu: bool = False) -> Weights: - ckpt_path = list(checkpoint_path.glob('*.ckpt'))[0] - model_params = torch.load(ckpt_path)['model_state_dict'] - model_params.pop('freqs_cis') - return model_params - - def embedding_weights(self, ckpt_params) -> np.ndarray: - return ckpt_params["embedder.weight"] - - def get_config(self, checkpoint_path, ckpt_params, - num_embed) -> SimpleNamespace: - checkpoint_path = checkpoint_path.absolute() - config_file = "config.json" - with open(checkpoint_path / config_file, 'r') as f: - json_str = f.read() - json_str = json_str.replace("'", "\"") - json_str = json_str.replace(",\n}", "\n}") - config_old = json.loads(json_str) - config_new = SimpleNamespace() - config_new.num_layers = config_old["num_hidden_layers"] - config_new.num_embed = config_old["vocab_size"] - config_new.embed_dim = config_old["hidden_size"] - config_new.hidden_dim = config_old["intermediate_size"] - config_new.num_heads = config_old["num_attention_heads"] - config_new.head_dim = config_old["head_dim"] - config_new.num_kv_heads = config_old["num_key_value_heads"] - return config_new - - def rename_to_trt_llm(self, name: str) -> Optional[str]: - """Rename a gemma parameter name by the corresponding TRT-LLM style name.""" - prefix = "transformer" - sub_patterns = ( - (r"embedder.weight", r"vocab_embedding.weight"), - (r"model.layers.(\d+).input_layernorm.weight", - r"layers.\1.input_layernorm.weight"), - (r"model.layers.(\d+).self_attn.qkv_proj.weight", - r"layers.\1.attention.qkv.weight"), - (r"model.layers.(\d+).self_attn.o_proj.weight", - r"layers.\1.attention.dense.weight"), - (r"model.layers.(\d+).mlp.gate_proj.weight", - r"layers.\1.mlp.fc.weight"), - (r"model.layers.(\d+).mlp.up_proj.weight", - None), # merged with above - (r"model.layers.(\d+).mlp.down_proj.weight", - r"layers.\1.mlp.proj.weight"), - (r"model.layers.(\d+).post_attention_layernorm.weight", - r"layers.\1.post_layernorm.weight"), - (r"model.norm.weight", r"ln_f.weight"), - ) - - for source, target in sub_patterns: - if re.match(source, name): - if target is None: - return target - else: - name = re.sub(source, target, name) - return ".".join((prefix, name)) - else: - raise ValueError(f"Don't know how to rename {name}") - - def flatten_params(self, params) -> Flattened: - f_params = {} - for k, v in params.items(): - if v.dtype == torch.bfloat16: - v = v.float() - f_params[k] = torch_to_numpy(v) - return f_params - - -class HfParser: - - def load_parameters(self, - checkpoint_path: Path, - load_model_on_cpu: bool = False) -> Weights: - """`AutoModelForCausalLM.from_pretrained` will parse the correct gemma, whether Gemma or Gemma2 or future versions.""" - from transformers import AutoModelForCausalLM - hf_model = AutoModelForCausalLM.from_pretrained( - checkpoint_path, - device_map="cpu" if load_model_on_cpu else "auto", - dtype='auto', - trust_remote_code=True, - ) - model_params = dict(hf_model.named_parameters()) - return model_params - - def embedding_weights(self, ckpt_params) -> np.ndarray: - raise RuntimeError( - "This method shouldn't be called - `GemmaConfig.from_hugging_face` takes care of this." - ) - - def get_config(self, checkpoint_path, ckpt_params, - num_embed) -> SimpleNamespace: - raise RuntimeError( - "This method shouldn't be called - `GemmaConfig.from_hugging_face` takes care of this." - ) - - def rename_to_trt_llm(self, name: str) -> Optional[str]: - """Rename a gemma parameter name by the corresponding TRT-LLM style name.""" - prefix = "transformer" - sub_patterns = ( - (r"model.embed_tokens.weight", r"vocab_embedding.weight"), - (r"model.layers.(\d+).input_layernorm.weight", - r"layers.\1.input_layernorm.weight"), - (r"model.layers.(\d+).self_attn.q_proj.weight", - r"layers.\1.attention.qkv.weight"), - (r"model.layers.(\d+).self_attn.k_proj.weight", - None), # merged with above - (r"model.layers.(\d+).self_attn.v_proj.weight", - None), # merged with above - (r"model.layers.(\d+).self_attn.o_proj.weight", - r"layers.\1.attention.dense.weight"), - (r"model.layers.(\d+).self_attn.q_norm.weight", - r"layers.\1.attention.q_layernorm.weight"), - (r"model.layers.(\d+).self_attn.k_norm.weight", - r"layers.\1.attention.k_layernorm.weight"), - (r"model.layers.(\d+).mlp.gate_proj.weight", - r"layers.\1.mlp.fc.weight"), - (r"model.layers.(\d+).mlp.up_proj.weight", - None), # merged with above - (r"model.layers.(\d+).mlp.down_proj.weight", - r"layers.\1.mlp.proj.weight"), - (r"model.layers.(\d+).post_attention_layernorm.weight", - r"layers.\1.post_layernorm.weight"), - (r"model.layers.(\d+).pre_feedforward_layernorm.weight", - r"layers.\1.pre_feedforward_layernorm.weight"), - (r"model.layers.(\d+).post_feedforward_layernorm.weight", - r"layers.\1.post_feedforward_layernorm.weight"), - (r"model.norm.weight", r"ln_f.weight"), - ) - - for source, target in sub_patterns: - if re.match(source, name): - if target is None: - return target - else: - name = re.sub(source, target, name) - return ".".join((prefix, name)) - else: - raise ValueError(f"Don't know how to rename {prefix}.{name}") - - def flatten_params(self, params: Weights) -> Dict[str, numpy.ndarray]: - f_params = {} - for k, v in params.items(): - if v.dtype == torch.bfloat16: - v = v.float() - f_params[k] = torch_to_numpy(v) - return f_params - - -def split(v: np.ndarray, tp_size: int, idx: int, dim: int = 0) -> np.ndarray: - if tp_size == 1: - return v - return np.split(v, tp_size, axis=dim)[idx] - - -def split_matrix_tp(v: np.ndarray, tensor_parallel: int, rank: int, - dim: int) -> np.ndarray: - return split(v, tensor_parallel, rank, dim=dim) - - -def add_trt_llm_weight(weights: Flattened, - name: str, - param: np.ndarray, - dtype: Optional[np.dtype] = None): - assert name not in weights, f"{name} is already added." - param = numpy_to_torch(param) - if dtype is not None: - assert isinstance(dtype, - str), f"dtype must be str, but get type {type(dtype)}" - param = param.to(str_dtype_to_torch(dtype)) - weights[name] = param.contiguous() - - -def quantize(param: np.ndarray, - quant_mode: tensorrt_llm.quantization.QuantMode): - if quant_mode.is_int8_weight_only(): - quant_dtype = torch.int8 - elif quant_mode.is_int4_weight_only(): - quant_dtype = torch.quint4x2 - else: - raise ValueError(f"Invalid configuration got quant_mode={quant_mode}") - - param = numpy_to_torch(param) - param = param.t().contiguous() - - # previously this fn was available in torch.ops.fastertransformer namespace - ( - quantized_weights, - scales, - ) = torch.ops.trtllm.symmetric_quantize_last_axis_of_batched_matrix( - param, quant_dtype) - - if scales.dtype == torch.bfloat16: - scales = scales.to(torch.float32).numpy().astype("bfloat16") - else: - scales = scales.numpy() - return quantized_weights.numpy(), scales - - -Parsers = Union[JAXParser, KerasParser, TorchParser, HfParser] - - -def load_gemma_weights( - *, - trt_llm_config: tllm_utils.PretrainedConfig, - parameters_or_model_dir: "Union[Weights, Path]", - ckpt_parser: Parsers, - load_model_on_cpu: bool = True, -): - print("Loading weights...") - tik = time.time() - - tp_rank = trt_llm_config.mapping.tp_rank - tp_size = trt_llm_config.mapping.tp_size - hidden_size = trt_llm_config.hidden_size - head_dim = trt_llm_config.head_size - - weights = {} - - if isinstance(parameters_or_model_dir, Path): - logger.debug(f"Loading directory {str(parameters_or_model_dir)}...") - model_params = ckpt_parser.load_parameters( - parameters_or_model_dir, - load_model_on_cpu=load_model_on_cpu, - ) - else: - model_params = parameters_or_model_dir - - model_params = ckpt_parser.flatten_params(model_params) - - for name, param in model_params.items(): - logger.debug(f"Converting weight {name}...") - trt_llm_name = ckpt_parser.rename_to_trt_llm(name) - if trt_llm_name is None: # omit as used with other params - continue - - if "attn.q_einsum" in name: - gqa_mode = trt_llm_config.num_attention_heads != trt_llm_config.num_key_value_heads - assert gqa_mode - - # initial shape: (num_q_heads, hidden_size, head_dim) - q_param = param.transpose(1, 0, 2) - q_param = split_matrix_tp(q_param, tp_size, tp_rank, dim=1) - - # initial shape: (2, num_kv_heads, hidden_size, head_dim) - kv_name = name.replace("q_einsum", "kv_einsum") - kv_param = model_params[kv_name] - kv_param = kv_param.reshape( - trt_llm_config.num_key_value_heads * 2, - hidden_size, - head_dim, - ).transpose(1, 0, 2) - - # -> (hidden_size, num_q_heads / tp_size + 2, head_dim) - qkv_param = np.concatenate([q_param, kv_param], axis=1) - qkv_param = qkv_param.reshape(qkv_param.shape[0], -1) - qkv_param = qkv_param.transpose(1, 0) - - # If int8 kv enabled, weight-only quantization will be done later. - if trt_llm_config.quant_mode.is_weight_only() and not trt_llm_config.quant_mode.has_per_group_scaling() and \ - not trt_llm_config.quant_mode.has_int8_kv_cache(): - qkv_param_quantized, qkv_param_scales = quantize( - qkv_param, trt_llm_config.quant_mode) - add_trt_llm_weight(weights, trt_llm_name, qkv_param_quantized) - add_trt_llm_weight( - weights, - trt_llm_name.replace(".weight", ".per_channel_scale"), - qkv_param_scales, - trt_llm_config.dtype, - ) - else: - add_trt_llm_weight(weights, trt_llm_name, qkv_param, - trt_llm_config.dtype) - elif "self_attn.qkv_proj" in name: - q_param, k_param, v_param = np.split(param, [ - trt_llm_config.num_attention_heads * trt_llm_config.head_size, - trt_llm_config.num_attention_heads * trt_llm_config.head_size + - trt_llm_config.num_key_value_heads * trt_llm_config.head_size - ], - axis=0) - gqa_mode = trt_llm_config.num_attention_heads != trt_llm_config.num_key_value_heads - - q_param = split_matrix_tp(q_param, tp_size, tp_rank, dim=0) - if not gqa_mode: - k_param = split_matrix_tp(k_param, tp_size, tp_rank, dim=0) - v_param = split_matrix_tp(v_param, tp_size, tp_rank, dim=0) - - qkv_param = np.concatenate([q_param, k_param, v_param], axis=0) - if trt_llm_config.quant_mode.is_weight_only( - ) and not trt_llm_config.quant_mode.has_per_group_scaling(): - qkv_param_quantized, qkv_param_scales = quantize( - qkv_param, trt_llm_config.quant_mode) - add_trt_llm_weight(weights, trt_llm_name, qkv_param_quantized) - add_trt_llm_weight( - weights, - trt_llm_name.replace(".weight", ".per_channel_scale"), - qkv_param_scales, - trt_llm_config.dtype, - ) - else: - add_trt_llm_weight(weights, trt_llm_name, qkv_param, - trt_llm_config.dtype) - elif "attn.qkv_einsum" in name: - gqa_mode = trt_llm_config.num_attention_heads != trt_llm_config.num_key_value_heads - assert not gqa_mode - # initial shape: [3, num_heads, hidden_size, head_dim] -> [3, num_heads, head_dim, hidden_size] - qkv_param = param.transpose(0, 1, 3, 2) - qkv_param = qkv_param.reshape(qkv_param.shape[0], -1, - qkv_param.shape[3]) - qkv_param = split_matrix_tp(qkv_param, tp_size, tp_rank, dim=1) - qkv_param = qkv_param.reshape(-1, qkv_param.shape[2]) - if trt_llm_config.quant_mode.is_weight_only() and not trt_llm_config.quant_mode.has_per_group_scaling() \ - and not trt_llm_config.quant_mode.has_int8_kv_cache(): - qkv_param_quantized, qkv_param_scales = quantize( - qkv_param, trt_llm_config.quant_mode) - add_trt_llm_weight(weights, trt_llm_name, qkv_param_quantized) - add_trt_llm_weight( - weights, - trt_llm_name.replace(".weight", ".per_channel_scale"), - qkv_param_scales, - trt_llm_config.dtype, - ) - else: - add_trt_llm_weight(weights, trt_llm_name, qkv_param, - trt_llm_config.dtype) - elif "attention/query_dense" in name: - # Keras specific KQV convert - gqa_mode = trt_llm_config.num_attention_heads != trt_llm_config.num_key_value_heads - if gqa_mode: - - # initial shape: (num_q_heads, hidden_size, head_dim) - q_param = param.transpose(1, 0, 2) - q_param = split_matrix_tp(q_param, tp_size, tp_rank, dim=1) - - # initial shape: (2, num_kv_heads, hidden_size, head_dim) - k_name = name.replace("query", "key") - k_param = model_params[k_name] - v_name = name.replace("query", "value") - v_param = model_params[v_name] - kv_param = np.stack((k_param, v_param), axis=0) - - kv_param = kv_param.reshape( - trt_llm_config.num_key_value_heads * 2, - hidden_size, - head_dim, - ).transpose(1, 0, 2) - - # -> (hidden_size, num_q_heads / tp_size + 2, head_dim) - qkv_param = np.concatenate([q_param, kv_param], axis=1) - qkv_param = qkv_param.reshape(qkv_param.shape[0], -1) - qkv_param = qkv_param.transpose(1, 0) - - if trt_llm_config.quant_mode.is_weight_only( - ) and not trt_llm_config.quant_mode.has_int8_kv_cache(): - qkv_param_quantized, qkv_param_scales = quantize( - qkv_param, trt_llm_config.quant_mode) - add_trt_llm_weight(weights, trt_llm_name, - qkv_param_quantized) - add_trt_llm_weight( - weights, - trt_llm_name.replace(".weight", ".per_channel_scale"), - qkv_param_scales, - trt_llm_config.dtype, - ) - else: - add_trt_llm_weight(weights, trt_llm_name, qkv_param, - trt_llm_config.dtype) - else: - q_param = param - k_name = name.replace("query", "key") - k_param = model_params[k_name] - v_name = name.replace("query", "value") - v_param = model_params[v_name] - # initial shape: [3, num_heads, hidden_size, head_dim] -> [3, num_heads, head_dim, hidden_size] - qkv_param = np.stack((q_param, k_param, v_param), axis=0) - qkv_param = qkv_param.transpose(0, 1, 3, 2) - qkv_param = qkv_param.reshape(qkv_param.shape[0], -1, - qkv_param.shape[3]) - qkv_param = split_matrix_tp(qkv_param, tp_size, tp_rank, dim=1) - qkv_param = qkv_param.reshape(-1, qkv_param.shape[2]) - if trt_llm_config.quant_mode.is_weight_only( - ) and not trt_llm_config.quant_mode.has_int8_kv_cache(): - qkv_param_quantized, qkv_param_scales = quantize( - qkv_param, trt_llm_config.quant_mode) - add_trt_llm_weight(weights, trt_llm_name, - qkv_param_quantized) - add_trt_llm_weight( - weights, - trt_llm_name.replace(".weight", ".per_channel_scale"), - qkv_param_scales, - trt_llm_config.dtype, - ) - else: - add_trt_llm_weight(weights, trt_llm_name, qkv_param, - trt_llm_config.dtype) - elif "q_proj" in name: - mqa_mode = trt_llm_config.num_key_value_heads == 1 - - if mqa_mode: - # initial shape: (num_heads * head_dim, hidden_size) - q_param = param - q_param = split_matrix_tp(q_param, tp_size, tp_rank, dim=0) - - k_name = name.replace("q_proj", "k_proj") - k_param = model_params[k_name] - - v_name = name.replace("q_proj", "v_proj") - v_param = model_params[v_name] - else: - # initial shape: (num_heads * head_dim, hidden_size) - q_param = param - q_param = split_matrix_tp(q_param, tp_size, tp_rank, dim=0) - - k_name = name.replace("q_proj", "k_proj") - k_param = model_params[k_name] - k_param = split_matrix_tp(k_param, tp_size, tp_rank, dim=0) - - v_name = name.replace("q_proj", "v_proj") - v_param = model_params[v_name] - v_param = split_matrix_tp(v_param, tp_size, tp_rank, dim=0) - - qkv_param = np.concatenate([q_param, k_param, v_param], axis=0) - qkv_param = qkv_param.reshape(qkv_param.shape[0], -1) - - # If int8 kv enabled, weight-only quantization will be done later. - if trt_llm_config.quant_mode.is_weight_only() and not trt_llm_config.quant_mode.has_per_group_scaling() and \ - not trt_llm_config.quant_mode.has_int8_kv_cache(): - qkv_param_quantized, qkv_param_scales = quantize( - qkv_param, trt_llm_config.quant_mode) - add_trt_llm_weight(weights, trt_llm_name, qkv_param_quantized) - add_trt_llm_weight( - weights, - trt_llm_name.replace(".weight", ".per_channel_scale"), - qkv_param_scales, - trt_llm_config.dtype, - ) - else: - add_trt_llm_weight(weights, trt_llm_name, qkv_param, - trt_llm_config.dtype) - elif "attention.dense.weight" in trt_llm_name: - # initial shape: (num_heads, head_dim, hidden_size) - if len(param.shape) == 3: - param = param.reshape(-1, param.shape[2]) - param = param.transpose( - 1, 0) # (hidden_size, num_heads * head_dum) - param = split_matrix_tp(param, tp_size, tp_rank, dim=1) - if trt_llm_config.quant_mode.is_weight_only( - ) and not trt_llm_config.quant_mode.has_int8_kv_cache(): - param_quantized, param_scales = quantize( - param, trt_llm_config.quant_mode) - add_trt_llm_weight(weights, trt_llm_name, param_quantized) - add_trt_llm_weight( - weights, - trt_llm_name.replace(".weight", ".per_channel_scale"), - param_scales, - trt_llm_config.dtype, - ) - else: - add_trt_llm_weight(weights, trt_llm_name, param, - trt_llm_config.dtype) - elif "mlp.fc.weight" in trt_llm_name: - if isinstance(ckpt_parser, KerasParser): - # initial shape: (hidden_size, intermediate_size) - fc_param, gate_param = param, model_params[name.replace( - "gating_ffw", "gating_ffw_2")] - elif isinstance(ckpt_parser, TorchParser): - # initial shape: (intermediate_size, hidden_size) - fc_param, gate_param = param, model_params[name.replace( - "mlp.gate_proj", "mlp.up_proj")] - fc_param = fc_param.transpose(1, 0) - gate_param = gate_param.transpose(1, 0) - elif isinstance(ckpt_parser, HfParser): - # initial shape: (intermediate_size, hidden_size) - fc_param, gate_param = param, model_params[name.replace( - "mlp.gate_proj", "mlp.up_proj")] - fc_param = fc_param.transpose(1, 0) - gate_param = gate_param.transpose(1, 0) - else: - # initial shape: (2, hidden_size, intermediate_size) - fc_param, gate_param = param[0], param[1] - fc_param = fc_param.transpose(1, 0) - fc_param = split_matrix_tp(fc_param, tp_size, tp_rank, dim=0) - if trt_llm_config.quant_mode.is_weight_only() and not trt_llm_config.quant_mode.has_per_group_scaling() and \ - not trt_llm_config.quant_mode.has_int8_kv_cache(): - fc_param_quantized, fc_param_scales = quantize( - fc_param, trt_llm_config.quant_mode) - add_trt_llm_weight(weights, trt_llm_name, fc_param_quantized) - add_trt_llm_weight( - weights, - trt_llm_name.replace(".weight", ".per_channel_scale"), - fc_param_scales, - trt_llm_config.dtype, - ) - else: - add_trt_llm_weight(weights, trt_llm_name, fc_param, - trt_llm_config.dtype) - - gate_param = gate_param.transpose(1, 0) - gate_param = split_matrix_tp(gate_param, tp_size, tp_rank, dim=0) - trt_llm_name = trt_llm_name.replace("mlp.fc.weight", - "mlp.gate.weight") - if trt_llm_config.quant_mode.is_weight_only() and not trt_llm_config.quant_mode.has_per_group_scaling() and \ - not trt_llm_config.quant_mode.has_int8_kv_cache(): - gate_param_quantized, gate_param_scales = quantize( - gate_param, trt_llm_config.quant_mode) - add_trt_llm_weight(weights, trt_llm_name, gate_param_quantized) - add_trt_llm_weight( - weights, - trt_llm_name.replace(".weight", ".per_channel_scale"), - gate_param_scales, - trt_llm_config.dtype, - ) - else: - add_trt_llm_weight(weights, trt_llm_name, gate_param, - trt_llm_config.dtype) - elif "mlp.proj.weight" in trt_llm_name: - if not isinstance(ckpt_parser, TorchParser) and not isinstance( - ckpt_parser, HfParser): - # initial shape: (intermediate_size, hidden_size) - param = param.transpose(1, 0) - param = split_matrix_tp(param, tp_size, tp_rank, dim=1) - if trt_llm_config.quant_mode.is_weight_only() and not trt_llm_config.quant_mode.has_per_group_scaling() and \ - not trt_llm_config.quant_mode.has_int8_kv_cache(): - param_quantized, param_scales = quantize( - param, trt_llm_config.quant_mode) - add_trt_llm_weight(weights, trt_llm_name, param_quantized) - add_trt_llm_weight( - weights, - trt_llm_name.replace(".weight", ".per_channel_scale"), - param_scales, - trt_llm_config.dtype, - ) - else: - add_trt_llm_weight(weights, trt_llm_name, param, - trt_llm_config.dtype) - elif "embedder.input_embedding" in name or "reversible_embedding" in name or "embedder.weight" in name \ - or "embed_tokens.weight" in name: - # TODO: safetensor doesn't allow to save a shared tensor. - # Currently, we clone the weight but to save the disk, it - # would be better to skip saving lm_head weights and - # handle it at the loading phase. - lm_head = split_matrix_tp(param, tp_size, tp_rank, dim=0) - add_trt_llm_weight(weights, "lm_head.weight", np.copy(lm_head), - trt_llm_config.dtype) - - if trt_llm_config.use_parallel_embedding: - assert trt_llm_config.vocab_size % tp_size == 0 - param = split_matrix_tp( - param, - tp_size, - tp_rank, - dim=trt_llm_config.embedding_sharding_dim, - ) - if trt_llm_config.quant_mode.is_int8_weight_only() and not trt_llm_config.quant_mode.has_per_group_scaling() and \ - not trt_llm_config.quant_mode.has_int8_kv_cache() and trt_llm_config.quantization.exclude_modules is not None: - - # shape of embedding table: [V, K], V: vocab size, K: embedding dim - - # quantize will do following work: - # 1. transpose the input from [V, K] to [K, V] - # 2. compute V scales across dimension K - # 3. quantize the input to int8_weight - # 4. transpose the int8_weight to [V, K], but there is a bug in 'quantize' that the claimed shape is [K, V] - param_quantized, param_scales = quantize( - param, trt_llm_config.quant_mode) - - # Reshape the [K, V] to [V, K] to match the real data layout. - param_quantized = np.ascontiguousarray( - param_quantized.reshape( - [param_quantized.shape[1], param_quantized.shape[0]])) - - add_trt_llm_weight(weights, trt_llm_name, param_quantized) - add_trt_llm_weight( - weights, - trt_llm_name.replace(".weight", ".per_token_scale"), - param_scales, - trt_llm_config.dtype, - ) - else: - add_trt_llm_weight(weights, trt_llm_name, param, - trt_llm_config.dtype) - elif any(keyword in name for keyword in ( - "pre_attention_norm.scale", - "pre_ffw_norm.scale", - "final_norm.scale", - "pre_attention_norm/vars/0", - "pre_ffw_norm/vars/0", - "rms_normalization/vars/0", - "input_layernorm", - "post_attention_layernorm", - "pre_feedforward_layernorm", - "post_feedforward_layernorm", - "model.norm.weight", - "q_norm.weight", - "k_norm.weight", - )): - param = param + 1.0 # upcasted to float32 in case of bfloat16 - add_trt_llm_weight(weights, trt_llm_name, param, - trt_llm_config.dtype) - else: - raise RuntimeError(f"Unhandled {name} module weights") - del model_params - - print( - f"Weights loaded. Total time: {time.strftime('%H:%M:%S', time.gmtime(time.time() - tik))}" - ) - return weights - - -@dataclass(frozen=True, kw_only=True) -class QuantizeModifiers: - """ - Bag of additional conversion parameters, sourced from argparse or from defaults. - We may use this class to easily create wrappers of `convert_gemma` - (Otherwise, we would need to repeat these params many times or spread `**kwargs` in an untypesafe manner) - """ - - use_weight_only_with_precision: Optional[str] = None - per_channel: bool = False - per_token: bool = False - calib_dataset: str = "ccdv/cnn_dailymail" - tokenizer_dir: Optional[str] = None - enable_fp8: bool = False - fp8_kv_cache: bool = False - quant_ckpt_path: Optional[str] = None - - @classmethod - def from_args(cls, args: Namespace) -> "QuantizeModifiers": - return cls(**{ - k: v - for k, v in vars(args).items() if k in cls.__annotations__ - }) - - -def non_modelopt_quantize_if_needed( - weights: Weights, *, model_dir: Path, - quantize_modifiers: QuantizeModifiers, - trt_llm_config: Union[GemmaConfig, - tllm_utils.PretrainedConfig]) -> Weights: - tokenizer_dir = quantize_modifiers.tokenizer_dir or model_dir - - quant_cfg = trt_llm_config.quantization - - kv_cache_is_int8 = quant_cfg.kv_cache_quant_algo == QuantAlgo.INT8 - use_smooth_quant = quant_cfg._use_plugin_sq - - if use_smooth_quant or kv_cache_is_int8: - qkv_para = {} - smoother = {} - dataset = load_calib_dataset(quantize_modifiers.calib_dataset) - tokenizer = sp.SentencePieceProcessor(model_file=tokenizer_dir) - if "transformer.vocab_embedding.weight" in weights: - # To use the HF to do SmoothQuant, we need to scale the embedding. - weights["transformer.vocab_embedding.weight"] = torch.multiply( - weights["transformer.vocab_embedding.weight"].to(torch.float32), - math.sqrt(trt_llm_config.hidden_size), - ) - hf_model = create_model_from_config(trt_llm_config, weights) - act_range = capture_activation_range(hf_model, tokenizer, dataset) - if use_smooth_quant: - smooth_model(hf_model, act_range, quant_cfg.smoothquant_val, - qkv_para, smoother) - weights = convert_hf_model( - hf_model=hf_model, - mapping=trt_llm_config.mapping, - vocab_size=trt_llm_config.vocab_size or 32000, - dtype=trt_llm_config.dtype, - use_parallel_embedding=trt_llm_config.use_parallel_embedding, - sharding_dim=0, - use_weight_only=quantize_modifiers.use_weight_only_with_precision - is not None, - plugin_weight_only_quant_type=torch.int8 - if quantize_modifiers.use_weight_only_with_precision == 'int8' else - torch.quint4x2, - use_smooth_quant=use_smooth_quant, - per_channel=quantize_modifiers.per_channel, - per_token=quantize_modifiers.per_token, - int8_kv_cache=kv_cache_is_int8, - act_range=act_range, - qkv_para=qkv_para, - smoother=smoother) - if "transformer.vocab_embedding.weight" in weights: - # Revert the scaling of embedding - weights["transformer.vocab_embedding.weight"] = torch.divide( - weights["transformer.vocab_embedding.weight"].to(torch.float32), - math.sqrt(trt_llm_config.hidden_size), - ).to(str_dtype_to_torch(trt_llm_config.dtype)) - if "lm_head.weight" in weights: # Remove lm_head before saving it in unified checkpoint. - del weights["lm_head.weight"] - return weights - if quantize_modifiers.use_weight_only_with_precision and quantize_modifiers.use_weight_only_with_precision.endswith( - "awq"): - weights = dummy_weights_awq( - weights=weights, - precision=quantize_modifiers.use_weight_only_with_precision, - trt_llm_config=trt_llm_config, - group_size=128) - elif quantize_modifiers.enable_fp8 or quantize_modifiers.fp8_kv_cache: - weight_scales = quantize_fp8_weights(weights, - trt_llm_config.num_hidden_layers, - trt_llm_config.mapping) - scales = load_from_fp8_gemma(quantize_modifiers.quant_ckpt_path, - trt_llm_config.num_hidden_layers, - trt_llm_config.mapping, - quantize_modifiers.fp8_kv_cache, - weight_scales) - weights.update(scales) - - return weights - - -def load_gemma_weights_from_hf_model( - hf_model: "transformers.AutoModelForCausalLM", - config: GemmaConfig) -> Weights: - return load_gemma_weights(parameters_or_model_dir=dict( - hf_model.named_parameters()), - trt_llm_config=config, - ckpt_parser=HfParser()) diff --git a/tensorrt_llm/models/gemma/model.py b/tensorrt_llm/models/gemma/model.py deleted file mode 100644 index c8c5307c5d4c..000000000000 --- a/tensorrt_llm/models/gemma/model.py +++ /dev/null @@ -1,396 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -import math -from typing import TYPE_CHECKING, Any, Dict, Optional - -from tensorrt_llm.models.gemma.convert import (QuantizeModifiers, Weights, - load_gemma_weights_from_hf_model, - non_modelopt_quantize_if_needed) -from tensorrt_llm.quantization.mode import (MODELOPT_FLOW_QUANTIZATIONS, - QuantAlgo) - -from ..._common import default_net -from ..._utils import pad_vocab_size -from ...functional import (AllReduceFusionOp, AllReduceParams, LayerNormType, - Tensor, cast, recv, send) -from ...layers import (Attention, AttentionMaskType, AttentionParams, - ColumnLinear, Embedding, GatedMLP, KeyValueCacheParams, - LoraParams, PositionEmbeddingType, RmsNorm) -from ...lora_helper import LoraConfig, use_lora -from ...mapping import Mapping -from ...module import Module -from ..modeling_utils import (DecoderLayerList, DecoderModelForCausalLM, - QuantConfig, save_checkpoint, save_config) -from .config import GemmaConfig - -if TYPE_CHECKING: - - from .config import HfConfigOrDir - - -class GemmaDecoderLayer(Module): - - def __init__(self, config: GemmaConfig, layer_idx: int): - super().__init__() - self.layer_idx = layer_idx - self.config = config - - self.input_layernorm = RmsNorm(normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - dtype=config.dtype) - - layers_range = config.mapping.pp_layers(config.num_hidden_layers) - self.local_layer_idx = layer_idx - layers_range[0] - - q_scaling = 1.0 - max_attn_value = 0.0 - qk_layernorm = False - is_sliding = False - rotary_base = config.rotary_base - rotary_base_local = None - - gemma2_config = config.gemma2_config() - gemma3_config = config.gemma3_config() - if gemma2_config: - q_scaling = math.sqrt( - gemma2_config.query_pre_attn_scalar) / math.sqrt( - config.head_size) - max_attn_value = config.attn_logit_softcapping or 0.0 - elif gemma3_config: - qk_layernorm = True - q_scaling = math.sqrt( - gemma3_config.query_pre_attn_scalar) / math.sqrt( - config.head_size) - is_sliding = bool( - (layer_idx + 1) % gemma3_config._sliding_window_pattern) - rotary_base_local = config.rope_local_base_freq - - self.attention = Attention( - local_layer_idx=self.local_layer_idx, - hidden_size=config.hidden_size, - num_attention_heads=config.num_attention_heads, - num_kv_heads=config.num_key_value_heads, - attention_head_size=config.head_size, - qk_layernorm=qk_layernorm, - layernorm_type=LayerNormType.RmsNorm, - max_position_embeddings=config.max_position_embeddings, - dtype=config.dtype, - attention_mask_type=AttentionMaskType.causal, - bias=config.attn_bias, - position_embedding_type=PositionEmbeddingType.rope_gpt_neox, - rotary_embedding_base=rotary_base, - rotary_embedding_base_local=rotary_base_local, - rotary_embedding_scaling=config.rotary_scaling, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size, - quant_mode=config.quant_mode, - q_scaling=q_scaling, - max_attn_value=max_attn_value, - is_local=is_sliding, - ) - - mlp_hidden_size = config.hidden_size * 4 if config.intermediate_size is None else config.intermediate_size - - self.mlp = GatedMLP(hidden_size=config.hidden_size, - ffn_hidden_size=mlp_hidden_size, - hidden_act=config.hidden_act, - dtype=config.dtype, - bias=config.mlp_bias, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size, - quant_mode=config.quant_mode) - - if self.config.inter_layernorms: - self.pre_feedforward_layernorm = RmsNorm( - normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - dtype=config.dtype) - self.post_feedforward_layernorm = RmsNorm( - normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - dtype=config.dtype) - - self.post_layernorm = RmsNorm(normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - dtype=config.dtype) - - def forward(self, - hidden_states: Tensor, - attention_mask: Optional[Tensor] = None, - use_cache: bool = False, - kv_cache_params: Optional[KeyValueCacheParams] = None, - attention_params: Optional[AttentionParams] = None, - lora_layer_params: Optional[LoraParams] = None, - next_layer_input_layernorm_args=None): - # assert not ( - # default_net().plugin_config.reduce_fusion and self.has_residual_mlp - # ), "Custom all reduce and residual mlp can't be enabled at the same time." - if default_net( - ).plugin_config.reduce_fusion and self.local_layer_idx > 0: - hidden_states, residual = hidden_states #FIXME:AN need to check if appropriate residual value is hidden state is pulled out. - else: - residual = hidden_states - hidden_states = self.input_layernorm(hidden_states) - - attention_output = self.attention( - hidden_states, - attention_mask=attention_mask, - use_cache=use_cache, - kv_cache_params=kv_cache_params, - attention_params=attention_params, - norm_before_bmm1=True, - lora_layer_params=lora_layer_params, - all_reduce_params=AllReduceParams( - fusion_op=AllReduceFusionOp.RESIDUAL_RMS_PREPOST_NORM - if default_net().plugin_config.reduce_fusion else - AllReduceFusionOp.NONE, - residual=residual, - norm_weight=self.pre_feedforward_layernorm.weight.value - if self.config.inter_layernorms else None, - norm_pre_residual_weight=self.post_layernorm.weight.value, - eps=self.pre_feedforward_layernorm.eps - if self.config.inter_layernorms else 1e-06, - ), - ) - - if use_cache: - attention_output, presents = attention_output - - if default_net().plugin_config.reduce_fusion: - hidden_states, residual = attention_output - else: - if self.config.inter_layernorms: - attention_output = self.post_layernorm(attention_output) - hidden_states = residual + attention_output - residual = hidden_states - if self.config.inter_layernorms: - hidden_states = self.pre_feedforward_layernorm(hidden_states) - else: - hidden_states = self.post_layernorm(hidden_states) - - if next_layer_input_layernorm_args is not None: - hidden_states = self.mlp( - hidden_states, - lora_layer_params=lora_layer_params, - all_reduce_params=AllReduceParams( - fusion_op=AllReduceFusionOp.RESIDUAL_RMS_PREPOST_NORM - if default_net().plugin_config.reduce_fusion else - AllReduceFusionOp.NONE, - residual=residual, - norm_weight=next_layer_input_layernorm_args[0], - norm_pre_residual_weight=self.post_feedforward_layernorm. - weight.value, - eps=next_layer_input_layernorm_args[1])) - else: - hidden_states = self.mlp(hidden_states, - lora_layer_params=lora_layer_params) - - if self.config.inter_layernorms: - hidden_states = self.post_feedforward_layernorm(hidden_states) - hidden_states = residual + hidden_states - if use_cache: - return (hidden_states, presents) - return hidden_states - - -class GemmaModel(Module): - - def __init__(self, config: GemmaConfig) -> None: - super().__init__() - - self.mapping = config.mapping - if self.mapping.is_first_pp_rank(): - self.vocab_embedding = Embedding(config.vocab_size, - config.hidden_size, - dtype=config.dtype) - - self.layers = DecoderLayerList(GemmaDecoderLayer, config) - - if self.mapping.is_last_pp_rank(): - self.ln_f = RmsNorm(normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - dtype=config.dtype) - self.hidden_size = config.hidden_size - - def forward(self, - input_ids, - position_ids=None, - use_cache=False, - attention_mask=None, - kv_cache_params=None, - attention_params=None, - hidden_states=None, - prompt_embedding_table: Optional[Tensor] = None, - prompt_tasks: Optional[Tensor] = None, - prompt_vocab_size: Optional[Tensor] = None, - lora_params=None): - - ptuning_args = [ - prompt_embedding_table, prompt_tasks, prompt_vocab_size - ] if prompt_embedding_table is not None else [] - - if self.mapping.is_first_pp_rank(): - hidden_states = self.vocab_embedding(input_ids, *ptuning_args) - hidden_states = cast(hidden_states * math.sqrt(self.hidden_size), - hidden_states.dtype) - else: - hidden_states = recv(hidden_states, self.mapping.prev_pp_rank()) - hidden_states = self.layers.forward( - hidden_states, - use_cache=use_cache, - attention_mask=attention_mask, - kv_cache_params=kv_cache_params, - attention_params=attention_params, - lora_params=lora_params, - ) - - if use_cache: - hidden_states, presents = hidden_states - - if self.mapping.is_last_pp_rank(): - hidden_states = self.ln_f(hidden_states) - else: - hidden_states = send(hidden_states, self.mapping.next_pp_rank()) - - if use_cache: - return (hidden_states, tuple(presents)) - return hidden_states - - -class GemmaForCausalLM(DecoderModelForCausalLM): - config_class = GemmaConfig - - def __init__(self, config: GemmaConfig): - transformer = GemmaModel(config) - - vocab_size_padded = pad_vocab_size(config.vocab_size, - config.mapping.tp_size) - - if config.mapping.is_last_pp_rank(): - lm_head = ColumnLinear(config.hidden_size, - vocab_size_padded, - bias=False, - dtype=config.dtype, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size, - gather_output=True) - else: - lm_head = None - self.quant_mode = config.quant_mode - self.mapping = config.mapping - - super().__init__(config, transformer, lm_head) - - @staticmethod - def _load_gemma_weights_from_hf(hf_model_dir: "HfConfigOrDir", - trt_llm_config: GemmaConfig, *, - load_model_on_cpu: bool) -> Weights: - """`AutoModelForCausalLM.from_pretrained` will parse the correct gemma, whether Gemma or Gemma2 or future versions.""" - import transformers - hf_gemma = transformers.AutoModelForCausalLM.from_pretrained( - hf_model_dir, - device_map="cpu" if load_model_on_cpu else "auto", - dtype='auto', - ) - weights = load_gemma_weights_from_hf_model(hf_gemma, trt_llm_config) - del hf_gemma - return weights - - @classmethod - def from_hugging_face(cls, - hf_model_dir: "HfConfigOrDir", - dtype='float16', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - load_model_on_cpu: bool = True, - **kwargs): - config = GemmaConfig.from_hugging_face(hf_config_or_dir=hf_model_dir, - dtype=dtype, - mapping=mapping, - quant_config=quant_config, - **kwargs) - model = GemmaForCausalLM(config) - weights = cls._load_gemma_weights_from_hf( - hf_model_dir, config, load_model_on_cpu=load_model_on_cpu) - model.load(weights) - return model - - NATIVE_QUANT_FLOW = { - QuantAlgo.W8A16, QuantAlgo.W4A16, - QuantAlgo.W8A8_SQ_PER_CHANNEL_PER_TOKEN_PLUGIN, - QuantAlgo.W8A8_SQ_PER_TENSOR_PLUGIN, - QuantAlgo.W8A8_SQ_PER_CHANNEL_PER_TENSOR_PLUGIN, - QuantAlgo.W8A8_SQ_PER_TENSOR_PER_TOKEN_PLUGIN - } - - @classmethod - def assert_valid_quant_algo(cls, quant_algo: Optional[QuantAlgo]): - allowed_quant_values = { - None - } | cls.NATIVE_QUANT_FLOW | MODELOPT_FLOW_QUANTIZATIONS - assert quant_algo in allowed_quant_values, f"{quant_algo} isn't in the allowed `QuantAlgo` values for this model: {allowed_quant_values}" - - @classmethod - def quantize( - cls, - hf_model_dir: str, - output_dir: str, - dtype: str = 'float16', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - *, - gemma_config_kwargs: Dict[str, Any] = None, - **quantize_kwargs: Dict[str, Any], - ): - config = GemmaConfig.from_hugging_face(hf_model_dir, - dtype=dtype, - mapping=mapping, - quant_config=quant_config, - **(gemma_config_kwargs or {})) - - quant_algo = config.quantization.quant_algo - if quant_algo is None and config.quantization.kv_cache_quant_algo is None: - raise ValueError( - "There is no point in calling `quantize()` if both `quant_algo` and `kv_cache_quant_algo` are `None`" - ) - elif quant_algo in MODELOPT_FLOW_QUANTIZATIONS: - super().quantize(hf_model_dir, - output_dir, - dtype=config.dtype, - mapping=config.mapping, - quant_config=config.quantization, - **quantize_kwargs) - elif quant_algo in cls.NATIVE_QUANT_FLOW: - save_config(config, output_dir=output_dir, log=True) - for config in config.for_each_rank(): - hf_weights = cls._load_gemma_weights_from_hf( - hf_model_dir, config) - ranked_weights = non_modelopt_quantize_if_needed( - hf_weights, - model_dir=hf_model_dir, - quantize_modifiers=QuantizeModifiers(), - trt_llm_config=config) - save_checkpoint( - output_dir=output_dir, - weights=ranked_weights, - rank=config.mapping.rank, - ) - del hf_weights - else: - cls.assert_valid_quant_algo(quant_algo) - - def use_lora(self, lora_config: LoraConfig) -> None: - return use_lora( - self, lora_config) # Use the default trtllm->hf module mapping diff --git a/tensorrt_llm/models/gemma/smoothquant.py b/tensorrt_llm/models/gemma/smoothquant.py deleted file mode 100644 index 63311a8eb11d..000000000000 --- a/tensorrt_llm/models/gemma/smoothquant.py +++ /dev/null @@ -1,848 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -import copy -import functools -import math -import time -from collections import defaultdict -from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple - -import numpy as np -import torch -import torch.nn as nn -import torch.nn.functional as F -from tqdm import tqdm -from transformers import LlamaConfig, LlamaForCausalLM -from transformers.models.llama.modeling_llama import (LlamaAttention, - LlamaDecoderLayer, - LlamaRotaryEmbedding, - apply_rotary_pos_emb, - repeat_kv) -from transformers.pytorch_utils import Conv1D - -from tensorrt_llm._utils import pad_vocab_size -from tensorrt_llm.mapping import Mapping -from tensorrt_llm.models.convert_utils import (dup_kv_weight, generate_int8, - get_weight, smooth_gemm, - smooth_gemm_fc1_gate, split, - split_matrix_tp) - -if TYPE_CHECKING: - from transformers import AutoModelForCausalLM, Cache - - # transformers included ⬆️ `Cache` in https://github.com/huggingface/transformers/commit/633215ba58fe5114d8c8d32e415a04600e010701 - transformers 4.33, which is installed in the tests, is before this. - - -@torch.no_grad() -def capture_activation_range(model, - tokenizer, - dataset, - num_samples=1, - seq_len=512): - model.cuda().eval() - device = next(model.parameters()).device - act_scales = defaultdict(lambda: {"x": None, "y": None, "w": None}) - - # tokenizer.pad_token = tokenizer.eos_token - - def stat_tensor(name, tensor, act_scales, key): - hidden_dim = tensor.shape[-1] - tensor = tensor.view(-1, hidden_dim).abs().detach() - comming_max = torch.max(tensor, dim=0)[0].float() - - if act_scales[name][key] is None: - act_scales[name][key] = comming_max - else: - act_scales[name][key] = torch.max(act_scales[name][key], - comming_max) - - def stat_input_hook(m, x, y, name): - if isinstance(x, tuple): - x = x[0] - stat_tensor(name, x, act_scales, "x") - stat_tensor(name, y, act_scales, "y") - - if act_scales[name]["w"] is None: - act_scales[name]["w"] = m.weight.abs().clip( - 1e-8, None).max(dim=1)[0].float() - - hooks = [] - for name, m in model.named_modules(): - if isinstance(m, nn.Linear) or isinstance(m, Conv1D): - hooks.append( - m.register_forward_hook( - functools.partial(stat_input_hook, name=name))) - - for i in tqdm(range(num_samples), desc="calibrating model"): - datapoint = dataset[i:i + 1] - line = copy.copy(datapoint) - line[0] = line[0] + ' TL;DR: ' - line[0] = line[0].strip() - line[0] = line[0].replace(" n't", "n't") - # input_ids = tokenizer(line, - # return_tensors="pt", - # max_length=seq_len, - # padding=True, - # truncation=True).input_ids.to(device) - inputs = tokenizer.EncodeAsIds(line[0]) - inputs = np.array([[tokenizer.bos_id()] + inputs], dtype=np.int32) - input_ids = torch.tensor(inputs, dtype=torch.int32).to(device) - model(input_ids) - - for h in hooks: - h.remove() - - return act_scales - - -@torch.no_grad() -def smooth_model(model, scales, alpha: Optional[float], qkv_para, - smoother_dict): - # Smooth the activation and weights with smoother = $\diag{s}$ - for name, module in model.named_modules(): - if not isinstance(module, LlamaDecoderLayer): - continue - # qkv_proj - layer_name_q = name + ".self_attn.q_proj" - layer_name_k = name + ".self_attn.k_proj" - layer_name_v = name + ".self_attn.v_proj" - layer_name_qkv = name + ".self_attn.qkv_proj" - - weight = torch.cat([ - module.self_attn.q_proj.weight, module.self_attn.k_proj.weight, - module.self_attn.v_proj.weight - ], - dim=0) - - smoother = smooth_gemm(weight, scales[layer_name_q]["x"], - module.input_layernorm.weight, None, alpha) - - scales[layer_name_qkv]["x"] = scales[layer_name_q]["x"] / smoother - scales[layer_name_qkv]["w"] = weight.abs().max(dim=1)[0] - scales[layer_name_qkv]["y"] = torch.cat([ - scales[layer_name_q]["y"], scales[layer_name_k]["y"], - scales[layer_name_v]["y"] - ], - dim=0) - - # see transpose_weights function - qkv_para[layer_name_qkv] = weight.transpose(0, 1) - - # ================================================================= - layer_name = name + ".self_attn.o_proj" - smoother = smooth_gemm(module.self_attn.o_proj.weight, - scales[layer_name]["x"], None, None, alpha) - smoother_dict[layer_name] = smoother.float() - - scales[layer_name]["x"] = scales[layer_name]["x"] / smoother - scales[layer_name]["w"] = module.self_attn.o_proj.weight.abs().max( - dim=1)[0] - - # ================================================================== - fc1_layer_name = name + ".mlp.gate_proj" - gate_layer_name = name + ".mlp.up_proj" - - smoother = smooth_gemm_fc1_gate(module.mlp.gate_proj.weight, - module.mlp.up_proj.weight, - scales[fc1_layer_name]["x"], - module.post_attention_layernorm.weight, - None, alpha) - - scales[fc1_layer_name]["x"] = scales[fc1_layer_name]["x"] / smoother - scales[fc1_layer_name]["w"] = module.mlp.gate_proj.weight.abs().max( - dim=1)[0] - - scales[gate_layer_name]["x"] = scales[gate_layer_name]["x"] / smoother - scales[gate_layer_name]["w"] = module.mlp.up_proj.weight.abs().max( - dim=1)[0] - - # ================================================================== - layer_name = name + ".mlp.down_proj" - smoother = smooth_gemm(module.mlp.down_proj.weight, - scales[layer_name]["x"], None, None, alpha) - smoother_dict[layer_name] = smoother.float() - scales[layer_name]["x"] = scales[layer_name]["x"] / smoother - scales[layer_name]["w"] = module.mlp.down_proj.weight.abs().max( - dim=1)[0] - - -def get_tllm_linear_sq_weight(vals, - prefix, - shape, - tensor_parallel, - is_qkv=False, - per_token=False, - per_channel=False, - last_prefix=None, - bias=None, - smoother_value=None, - smoother_shape=None, - rank=0, - cat_dim=0, - multi_query_mode=False): - results = {} - - def multi_query_split(data, local_dim, head_size, tp_size, cur_rank): - q, k, v = torch.split(data, [local_dim, head_size, head_size], dim=-1) - q_split = torch.split(q, q.shape[-1] // tp_size, dim=-1) - k_split = torch.split(k, q.shape[-1] // tp_size, dim=-1) - v_split = torch.split(v, q.shape[-1] // tp_size, dim=-1) - return [ - torch.concat((q_split[ii], k_split[ii], v_split[ii]), dim=-1) - for ii in range(tp_size) - ][cur_rank] - - col_shape = shape if (is_qkv or per_channel) else [1, 1] - - if per_token: - if per_channel: - original_weights = vals["weight.int8.col"] - else: - original_weights = vals["weight.int8"] - local_dim = original_weights.shape[0] - head_size = (original_weights.shape[1] - local_dim) // 2 - - if multi_query_mode: - cur_weights = multi_query_split(original_weights, local_dim, - head_size, tensor_parallel, rank) - else: - cur_weights = torch.split(original_weights, - original_weights.shape[-1] // - tensor_parallel, - dim=-1)[rank] - if is_qkv: - hidden_dim = cur_weights.shape[0] - cur_weights = cur_weights.reshape(hidden_dim, -1) - results[prefix + 'weight'] = cur_weights.t().contiguous() - if smoother_value is None: - results[last_prefix] = torch.from_numpy( - np.array([1.0], dtype=np.float32)) - - if per_channel: - cur_per_channel_value = vals["scale_w_quant_orig.col"] - if smoother_value is None: - if multi_query_mode: - - cur_per_channel_value = multi_query_split( - vals["scale_w_quant_orig.col"], local_dim, head_size, - tensor_parallel, rank) - else: - cur_per_channel_value = np.split( - vals["scale_w_quant_orig.col"], - tensor_parallel, - axis=-1)[rank] - else: - cur_per_channel_value = vals["scale_w_quant_orig"] - if is_qkv: - if multi_query_mode: - cur_per_channel_value = multi_query_split( - vals["scale_w_quant_orig"], local_dim, head_size, - tensor_parallel, rank) - else: - cur_per_channel_value = torch.split( - vals["scale_w_quant_orig"], - vals["scale_w_quant_orig"].shape[-1] // tensor_parallel, - dim=-1)[rank] - - results[prefix + - 'per_channel_scale'] = cur_per_channel_value.reshape(col_shape) - else: - if per_channel: - original_weights = vals["weight.int8.col"] - else: - original_weights = vals["weight.int8"] - local_dim = original_weights.shape[0] - head_size = (original_weights.shape[1] - local_dim) // 2 - - if multi_query_mode: - cur_weights = multi_query_split(original_weights, local_dim, - head_size, tensor_parallel, rank) - else: - cur_weights = torch.split(original_weights, - original_weights.shape[-1] // - tensor_parallel, - dim=-1)[rank] - if is_qkv: - hidden_dim = cur_weights.shape[0] - cur_weights = cur_weights.reshape(hidden_dim, -1) - results[prefix + 'weight'] = cur_weights.t().contiguous() - - if per_channel: - cur_per_channel_value = vals["scale_y_accum_quant.col"] - if smoother_value is None: - if multi_query_mode: - cur_per_channel_value = multi_query_split( - vals["scale_y_accum_quant.col"], local_dim, head_size, - tensor_parallel, rank) - else: - cur_per_channel_value = np.split( - vals["scale_y_accum_quant.col"], - tensor_parallel, - axis=cat_dim)[rank] - else: - cur_per_channel_value = vals["scale_y_accum_quant"] - # QKV is always per_channel - if is_qkv: - if multi_query_mode: - cur_per_channel_value = multi_query_split( - vals["scale_y_accum_quant"], local_dim, head_size, - tensor_parallel, rank) - else: - cur_per_channel_value = np.split( - vals["scale_y_accum_quant"], - tensor_parallel, - axis=cat_dim)[rank] - - results[prefix + 'per_channel_scale'] = cur_per_channel_value.reshape( - col_shape).contiguous() - - results[last_prefix] = vals['scale_x_orig_quant'].contiguous() - - results[prefix + 'act_scale'] = vals["scale_y_quant_orig"].contiguous() - - if smoother_value is not None: - cur_smoother_value = torch.split(smoother_value, - smoother_value.shape[-1] // - tensor_parallel, - dim=cat_dim)[rank] - - results[prefix + 'smoother'] = cur_smoother_value.reshape( - smoother_shape).contiguous().to(torch.float32) - - if bias is not None: - results[prefix + 'bias'] = bias - - return results - - -def split_qkv_tp(qkv, n_head, n_kv_heads, head_size, tensor_parallel, rank): - """ - Splits the QKV matrix according to tensor parallelism - """ - kv_head_size = n_kv_heads * head_size - q, k, v = torch.split(qkv, [n_head * head_size, kv_head_size, kv_head_size], - dim=0) - q = split(q, tensor_parallel, rank, dim=0) - k = split(k, tensor_parallel, rank, dim=0) - v = split(v, tensor_parallel, rank, dim=0) - return torch.concatenate([q, k, v], dim=0).contiguous() - - -def get_tllm_linear_weight( - weight: torch.Tensor, - prefix: str, - bias: Optional[torch.Tensor] = None, - use_weight_only: bool = False, - plugin_weight_only_quant_type: torch.dtype = torch.int8 -) -> Dict[str, torch.Tensor]: - results = {} - if use_weight_only: - v = weight.t().contiguous() - processed_torch_weights, torch_weight_scales = \ - torch.ops.trtllm.symmetric_quantize_last_axis_of_batched_matrix( - v, plugin_weight_only_quant_type) - results[f'{prefix}weight'] = processed_torch_weights - results[f'{prefix}per_channel_scale'] = torch_weight_scales - else: - results[f'{prefix}weight'] = weight.contiguous() - - if bias is not None: - results[f'{prefix}bias'] = bias - - return results - - -class LlamaAttentionExtend(LlamaAttention): - - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - self.head_dim = self.config.head_size - self.q_proj = nn.Linear(self.config.hidden_size, - self.config.num_attention_heads * self.head_dim, - bias=False) - self.k_proj = nn.Linear(self.config.hidden_size, - self.config.num_key_value_heads * self.head_dim, - bias=False) - self.v_proj = nn.Linear(self.config.hidden_size, - self.config.num_key_value_heads * self.head_dim, - bias=False) - self.o_proj = nn.Linear(self.config.num_attention_heads * self.head_dim, - self.config.hidden_size, - bias=False) - self.config.head_dim = self.head_dim - self.rotary_emb = LlamaRotaryEmbedding(config=self.config) - - def forward( - self, - hidden_states: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - position_ids: Optional[torch.IntTensor] = None, - past_key_value: "Optional[Cache]" = None, - output_attentions: bool = False, - use_cache: bool = False, - cache_position: Optional[torch.IntTensor] = None, - **kwargs, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor], - Optional[Tuple[torch.Tensor]]]: - bsz, q_len, _ = hidden_states.size() - - if self.config.pretraining_tp > 1: - key_value_slicing = (self.config.num_key_value_heads * - self.head_dim) // self.config.pretraining_tp - query_slices = self.q_proj.weight.split( - (self.config.num_attention_heads * self.head_dim) // - self.config.pretraining_tp, - dim=0) - key_slices = self.k_proj.weight.split(key_value_slicing, dim=0) - value_slices = self.v_proj.weight.split(key_value_slicing, dim=0) - - query_states = [ - F.linear(hidden_states, query_slices[i]) - for i in range(self.config.pretraining_tp) - ] - query_states = torch.cat(query_states, dim=-1) - - key_states = [ - F.linear(hidden_states, key_slices[i]) - for i in range(self.config.pretraining_tp) - ] - key_states = torch.cat(key_states, dim=-1) - - value_states = [ - F.linear(hidden_states, value_slices[i]) - for i in range(self.config.pretraining_tp) - ] - value_states = torch.cat(value_states, dim=-1) - - else: - query_states = self.q_proj(hidden_states) - key_states = self.k_proj(hidden_states) - value_states = self.v_proj(hidden_states) - - query_states = query_states.view(bsz, q_len, - self.config.num_attention_heads, - self.head_dim).transpose(1, 2) - key_states = key_states.view(bsz, q_len, - self.config.num_key_value_heads, - self.head_dim).transpose(1, 2) - value_states = value_states.view(bsz, q_len, - self.config.num_key_value_heads, - self.head_dim).transpose(1, 2) - - past_key_value = getattr(self, "past_key_value", past_key_value) - cos, sin = self.rotary_emb(value_states, position_ids) - query_states, key_states = apply_rotary_pos_emb(query_states, - key_states, cos, sin) - - if past_key_value is not None: - # sin and cos are specific to RoPE models; position_ids needed for the static cache - cache_kwargs = { - "sin": sin, - "cos": cos, - "cache_position": cache_position - } - key_states, value_states = past_key_value.update( - key_states, value_states, self.layer_idx, cache_kwargs) - - key_states = repeat_kv(key_states, self.num_key_value_groups) - value_states = repeat_kv(value_states, self.num_key_value_groups) - - attn_weights = torch.matmul(query_states, key_states.transpose( - 2, 3)) / math.sqrt(self.head_dim) - - if attention_mask is not None: # no matter the length, we just slice it - if cache_position is not None: - causal_mask = attention_mask[:, :, cache_position, :key_states. - shape[-2]] - attn_weights = attn_weights + causal_mask - - # upcast attention to fp32 - attn_weights = nn.functional.softmax(attn_weights, - dim=-1, - dtype=torch.float32).to( - query_states.dtype) - attn_weights = nn.functional.dropout(attn_weights, - p=self.attention_dropout, - training=self.training) - attn_output = torch.matmul(attn_weights, value_states) - - if attn_output.size() != (bsz, self.config.num_attention_heads, q_len, - self.head_dim): - raise ValueError( - f"`attn_output` should be of size {(bsz, self.config.num_attention_heads, q_len, self.head_dim)}, but is" - f" {attn_output.size()}") - - attn_output = attn_output.transpose(1, 2).contiguous() - - # Here is what we extend. - attn_output = attn_output.reshape( - bsz, q_len, self.config.num_attention_heads * self.head_dim) - - if self.config.pretraining_tp > 1: - attn_output = attn_output.split(self.hidden_size // - self.config.pretraining_tp, - dim=2) - o_proj_slices = self.o_proj.weight.split(self.hidden_size // - self.config.pretraining_tp, - dim=1) - attn_output = sum([ - F.linear(attn_output[i], o_proj_slices[i]) - for i in range(self.config.pretraining_tp) - ]) - else: - attn_output = self.o_proj(attn_output) - - if not output_attentions: - attn_weights = None - - return attn_output, attn_weights - - -def create_model_from_config(trt_llm_config, weights): - model_config = LlamaConfig() - model_config.vocab_size = trt_llm_config.vocab_size - model_config.dtype = trt_llm_config.dtype - model_config.max_position_embeddings = trt_llm_config.max_position_embeddings - model_config.hidden_size = trt_llm_config.hidden_size - model_config.num_hidden_layers = trt_llm_config.num_hidden_layers - model_config.num_attention_heads = trt_llm_config.num_attention_heads - model_config.num_key_value_heads = trt_llm_config.num_key_value_heads - model_config.hidden_act = trt_llm_config.hidden_act - model_config.head_size = trt_llm_config.head_size - model_config.intermediate_size = trt_llm_config.intermediate_size - model = LlamaForCausalLM(model_config) - # Hack attention module since head_dim * num_heads > hidden_size for 7B. - for i in range(model_config.num_hidden_layers): - module = model.model.layers[i].self_attn - model.model.layers[i].self_attn = LlamaAttentionExtend( - module.config, module.layer_idx) - # Copy wegiht to LLAMA model. - replace_name_dict = { - 'attention.dense': 'self_attn.o_proj', - 'mlp.proj': 'mlp.down_proj', - 'mlp.gate': 'mlp.up_proj', - 'mlp.fc': 'mlp.gate_proj', - 'ln_f': 'norm', - 'post_layernorm': 'post_attention_layernorm', - 'vocab_embedding': 'embed_tokens', - } - for name in list(weights): - param = weights[name] - weights.pop(name) - new_name = name.replace('transformer', 'model') - for _name in replace_name_dict: - if _name in new_name: - new_name = new_name.replace(_name, replace_name_dict[_name]) - if 'attention.qkv' in name: - qw, kw, vw = torch.split(param, [ - model_config.num_attention_heads * model_config.head_size, - model_config.num_key_value_heads * model_config.head_size, - model_config.num_key_value_heads * model_config.head_size, - ], - dim=0) - weights[new_name.replace('attention.qkv', 'self_attn.q_proj')] = qw - weights[new_name.replace('attention.qkv', 'self_attn.k_proj')] = kw - weights[new_name.replace('attention.qkv', 'self_attn.v_proj')] = vw - else: - weights[new_name] = param - - if "lm_head.weight" not in weights: - weights["lm_head.weight"] = weights["model.embed_tokens.weight"].clone() - model.load_state_dict(weights) - return model - - -def convert_hf_model(*, hf_model: "AutoModelForCausalLM", mapping: Mapping, - vocab_size: int, dtype: str, use_parallel_embedding: bool, - sharding_dim: int, use_weight_only: bool, - plugin_weight_only_quant_type: torch.dtype, - use_smooth_quant: bool, per_channel: bool, per_token: bool, - int8_kv_cache: bool, - act_range: "defaultdict[Any, dict[str, None]]", - qkv_para: Dict, smoother: Dict): - - weights = {} - tik = time.time() - tensor_parallel = mapping.tp_size - model_params = dict(hf_model.named_parameters()) - dtype = getattr(torch, dtype) - num_attention_heads = hf_model.config.num_attention_heads - hidden_size = hf_model.config.hidden_size - intermediate_size = hf_model.config.intermediate_size - head_size = hf_model.config.head_size - num_key_value_heads = hf_model.config.num_key_value_heads - mha_mode = (num_key_value_heads == num_attention_heads) - - num_hidden_layers = hf_model.config.num_hidden_layers - layers_range = mapping.pp_layers(num_hidden_layers) - for l in layers_range: - print("Processing layer", l) - prefix = f'model.layers.{l}.' - layer_idx = int(l) - layers_range[0] - tllm_prex = f'transformer.layers.{layer_idx}.' - - if use_smooth_quant: - qkv_weight = qkv_para[prefix + 'self_attn.qkv_proj'] - qkv_out_dim = qkv_weight.shape[1] - - if not mha_mode: - hidden_size = qkv_weight.shape[0] - local_dim = hidden_size - head_size = (qkv_weight.shape[-1] - local_dim) // 2 - qkv_weight = qkv_weight.reshape(hidden_size, - local_dim + 2 * head_size) - else: - qkv_weight = qkv_weight.reshape(hidden_size, 3, - head_size * num_attention_heads) - - int8_weights = generate_int8(qkv_weight, - act_range.get(prefix + - 'self_attn.qkv_proj'), - is_qkv=True, - multi_query_mode=bool(not mha_mode)) - weights.update( - get_tllm_linear_sq_weight(int8_weights, - tllm_prex + 'attention.qkv.', - [1, qkv_out_dim // tensor_parallel], - tensor_parallel, - is_qkv=True, - per_token=per_token, - per_channel=per_channel, - last_prefix=tllm_prex + - 'input_layernorm.scale_to_int', - smoother_value=None, - smoother_shape=None, - rank=mapping.tp_rank, - cat_dim=-1, - multi_query_mode=bool(not mha_mode))) - else: - q_weight = get_weight(model_params, prefix + 'self_attn.q_proj', - dtype) - k_weight = get_weight(model_params, prefix + 'self_attn.k_proj', - dtype) - v_weight = get_weight(model_params, prefix + 'self_attn.v_proj', - dtype) - if not mha_mode: - if num_key_value_heads < tensor_parallel: - # duplicate the KV heads up to tensor_parallel - k_weight = dup_kv_weight(k_weight, num_key_value_heads, - tensor_parallel) - v_weight = dup_kv_weight(v_weight, num_key_value_heads, - tensor_parallel) - assert (k_weight.shape[0] % (mapping.tp_size * head_size)) == 0 - assert (v_weight.shape[0] % (mapping.tp_size * head_size)) == 0 - - wq = split(q_weight, mapping.tp_size, mapping.tp_rank) - wk = split(k_weight, mapping.tp_size, mapping.tp_rank) - wv = split(v_weight, mapping.tp_size, mapping.tp_rank) - - split_v = torch.concat((wq, wk, wv)) - - else: - qkv_weight = torch.cat([q_weight, k_weight, v_weight], dim=0) - - split_v = split_qkv_tp(qkv_weight, num_attention_heads, - num_key_value_heads, head_size, - tensor_parallel, mapping.tp_rank) - weights.update( - get_tllm_linear_weight(split_v, tllm_prex + 'attention.qkv.', - None, use_weight_only, - plugin_weight_only_quant_type)) - - if int8_kv_cache: - qkv_y = torch.cat([ - act_range.get(prefix + 'self_attn.q_proj')["y"], - act_range.get(prefix + 'self_attn.k_proj')["y"], - act_range.get(prefix + 'self_attn.v_proj')["y"] - ], - dim=0) - int8_kv_scales = qkv_y.max() / 127. - kv_cache_weights = {} - kv_cache_weights[ - tllm_prex + - 'attention.kv_cache_scaling_factor'] = int8_kv_scales.reshape( - [1]) - - weights.update(kv_cache_weights) - - # Attention dense. - attn_dense_weight = get_weight(model_params, - prefix + 'self_attn.o_proj', dtype) - if use_smooth_quant: - attn_dense_weight = attn_dense_weight.t() - int8_weights = generate_int8( - attn_dense_weight, act_range.get(prefix + 'self_attn.o_proj')) - weights.update( - get_tllm_linear_sq_weight( - int8_weights, - tllm_prex + 'attention.dense.', [1, hidden_size], - tensor_parallel, - is_qkv=False, - per_token=per_token, - per_channel=per_channel, - last_prefix=tllm_prex + - 'attention.quantization_scaling_factor', - smoother_value=smoother[(prefix + 'self_attn.o_proj')], - smoother_shape=[ - 1, head_size * num_attention_heads // tensor_parallel - ], - rank=mapping.tp_rank, - cat_dim=0)) - else: - attn_dense_weight = split_matrix_tp(attn_dense_weight, - tensor_parallel, - mapping.tp_rank, - dim=1) - weights.update( - get_tllm_linear_weight(attn_dense_weight, - tllm_prex + 'attention.dense.', None, - use_weight_only, - plugin_weight_only_quant_type)) - # MLP hf up to trt gate - mlp_up_weight = get_weight(model_params, prefix + 'mlp.up_proj', dtype) - if use_smooth_quant: - mlp_up_weight = mlp_up_weight.t() - int8_weights = generate_int8(mlp_up_weight, - act_range.get(prefix + 'mlp.up_proj')) - weights.update( - get_tllm_linear_sq_weight( - int8_weights, - tllm_prex + 'mlp.gate.', - [1, intermediate_size // tensor_parallel], - tensor_parallel, - is_qkv=False, - per_token=per_token, - per_channel=per_channel, - last_prefix=tllm_prex + 'post_layernorm.scale_to_int', - smoother_value=None, - smoother_shape=None, - rank=mapping.tp_rank, - cat_dim=-1)) - else: - mlp_up_weight = split_matrix_tp(mlp_up_weight, - tensor_parallel, - mapping.tp_rank, - dim=0) - weights.update( - get_tllm_linear_weight(mlp_up_weight, tllm_prex + 'mlp.gate.', - None, use_weight_only, - plugin_weight_only_quant_type)) - - # MLP trt Gate to mlp fc - mlp_gate_weight = get_weight(model_params, prefix + 'mlp.gate_proj', - dtype) - if use_smooth_quant: - mlp_gate_weight = mlp_gate_weight.t() - int8_weights = generate_int8( - mlp_gate_weight, act_range.get(prefix + 'mlp.gate_proj')) - weights.update( - get_tllm_linear_sq_weight( - int8_weights, - tllm_prex + 'mlp.fc.', - [1, intermediate_size // tensor_parallel], - tensor_parallel, - is_qkv=False, - per_token=per_token, - per_channel=per_channel, - last_prefix=tllm_prex + 'post_layernorm.scale_to_int', - smoother_value=None, - smoother_shape=None, - rank=mapping.tp_rank, - cat_dim=-1)) - else: - mlp_gate_weight = split_matrix_tp(mlp_gate_weight, - tensor_parallel, - mapping.tp_rank, - dim=0) - weights.update( - get_tllm_linear_weight(mlp_gate_weight, tllm_prex + 'mlp.fc.', - None, use_weight_only, - plugin_weight_only_quant_type)) - - # MLP down - mlp_proj_weight = get_weight(model_params, prefix + 'mlp.down_proj', - dtype) - if use_smooth_quant: - mlp_proj_weight = mlp_proj_weight.t() - int8_weights = generate_int8( - mlp_proj_weight, act_range.get(prefix + 'mlp.down_proj')) - weights.update( - get_tllm_linear_sq_weight( - int8_weights, - tllm_prex + 'mlp.proj.', [1, hidden_size], - tensor_parallel, - is_qkv=False, - per_token=per_token, - per_channel=per_channel, - last_prefix=tllm_prex + 'mlp.quantization_scaling_factor', - smoother_value=smoother[prefix + 'mlp.down_proj'], - smoother_shape=[1, intermediate_size // tensor_parallel], - rank=mapping.tp_rank, - cat_dim=-1)) - else: - mlp_proj_weight = split_matrix_tp(mlp_proj_weight, - tensor_parallel, - mapping.tp_rank, - dim=1) - weights.update( - get_tllm_linear_weight(mlp_proj_weight, tllm_prex + 'mlp.proj.', - None, use_weight_only, - plugin_weight_only_quant_type)) - - # Layer norms do not use tensor parallelism - input_ln_weight = get_weight(model_params, prefix + 'input_layernorm', - dtype) - weights[tllm_prex + 'input_layernorm.weight'] = input_ln_weight - - post_ln_weight = get_weight(model_params, - prefix + 'post_attention_layernorm', dtype) - weights[tllm_prex + 'post_layernorm.weight'] = post_ln_weight - - v = get_weight(model_params, 'model.embed_tokens', dtype) - - if use_parallel_embedding: - v = split_matrix_tp(v, - mapping.tp_size, - mapping.tp_rank, - dim=sharding_dim) - - if mapping.is_first_pp_rank(): - weights['transformer.vocab_embedding.weight'] = v - - lm_head_weights = get_weight(model_params, 'lm_head', dtype) - - if mapping.is_last_pp_rank(): - - if vocab_size % mapping.tp_size != 0: - # padding - vocab_size_padded = pad_vocab_size(vocab_size, mapping.tp_size) - pad_width = vocab_size_padded - vocab_size - - lm_head_weights = torch.from_numpy( - np.pad(lm_head_weights.detach().cpu().numpy(), - ((0, pad_width), (0, 0)), - 'constant', - constant_values=0)) - weights['lm_head.weight'] = split_matrix_tp(lm_head_weights, - tensor_parallel, - mapping.tp_rank, - dim=0) - ln_f_w = get_weight(model_params, 'model.norm', dtype) - weights['transformer.ln_f.weight'] = ln_f_w - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Weights loaded. Total time: {t}') - return weights diff --git a/tensorrt_llm/models/gemma/utils/__init__.py b/tensorrt_llm/models/gemma/utils/__init__.py deleted file mode 100644 index c506269a5304..000000000000 --- a/tensorrt_llm/models/gemma/utils/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -# Copyright 2024 DeepMind Technologies Limited. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ============================================================================ diff --git a/tensorrt_llm/models/gemma/utils/layers.py b/tensorrt_llm/models/gemma/utils/layers.py deleted file mode 100644 index 0c2f471298a8..000000000000 --- a/tensorrt_llm/models/gemma/utils/layers.py +++ /dev/null @@ -1,39 +0,0 @@ -# Copyright 2024 DeepMind Technologies Limited. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ============================================================================ -"""Base layers.""" - -import jax -import jax.numpy as jnp -from flax import linen as nn - - -class Einsum(nn.Module): - shape: tuple[int, ...] - - @nn.compact - def __call__(self, eqn: str, x: jax.Array) -> jax.Array: - w = self.param('w', nn.initializers.zeros_init(), self.shape) - return jnp.einsum(eqn, x, w) - - -class RMSNorm(nn.Module): - - @nn.compact - def __call__(self, x): - scale = self.param('scale', nn.initializers.zeros_init(), (x.shape[-1])) - var = jnp.mean(jnp.square(x), axis=-1, keepdims=True) - normed_inputs = jnp.asarray(x * jnp.reciprocal(jnp.sqrt(var + 1e-06))) - normed_inputs = normed_inputs * (1 + scale) - return normed_inputs diff --git a/tensorrt_llm/models/gemma/utils/modules.py b/tensorrt_llm/models/gemma/utils/modules.py deleted file mode 100644 index 2ec10855e496..000000000000 --- a/tensorrt_llm/models/gemma/utils/modules.py +++ /dev/null @@ -1,206 +0,0 @@ -# Copyright 2024 DeepMind Technologies Limited. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ============================================================================ -"""Transformer sub-modules. -""" - -import jax -import jax.numpy as jnp -from flax import linen as nn - -from . import layers, positional_embeddings - -K_MASK = -2.3819763e38 # Set to a large negative number. -LayerCache = dict[str, jax.Array] - - -def init_layer_cache(cache_size: int, num_heads: int, head_dim: int, - batch_size: int) -> LayerCache: - return { - 'v': - jnp.zeros((batch_size, cache_size, num_heads, head_dim), - dtype=jnp.float32), - 'k': - jnp.zeros((batch_size, cache_size, num_heads, head_dim), - dtype=jnp.float32), - } - - -class Embedder(nn.Module): - """Embedder module.""" - - vocab_size: int - embed_dim: int - - def setup(self): - self.input_embedding_table = self.param( - 'input_embedding', - nn.initializers.zeros_init(), - (self.vocab_size, self.embed_dim), - ) - - def encode(self, x: jax.Array) -> jax.Array: - x = self.input_embedding_table[(x, )] - x *= jnp.sqrt(self.embed_dim).astype(x.dtype) - return x - - def decode(self, x: jax.Array) -> jax.Array: - return jnp.dot(x, self.input_embedding_table.T) - - -class Attention(nn.Module): - """Attention module.""" - - num_heads: int - num_kv_heads: int - features: int - head_dim: int - - @property - def use_qkv_einsum(self): - return self.num_kv_heads == self.num_heads - - def setup(self): - self.attn_vec_einsum = layers.Einsum(shape=(self.num_heads, - self.head_dim, - self.features), ) - - if self.use_qkv_einsum: - self.qkv_einsum = layers.Einsum(shape=(3, self.num_heads, - self.features, - self.head_dim), ) - else: - self.q_einsum = layers.Einsum(shape=(self.num_heads, self.features, - self.head_dim), ) - self.kv_einsum = layers.Einsum(shape=(2, self.num_kv_heads, - self.features, - self.head_dim), ) - - def __call__( - self, - x: jax.Array, - segment_pos: int, - cache: LayerCache, - attn_mask: jax.Array, - time_step: int, - ) -> tuple[LayerCache, jax.Array]: - - bsz = x.shape[0] - - if self.use_qkv_einsum: - query_proj, key_proj, value_proj = self.qkv_einsum( - 'BTD,SNDH->SBTNH', x) - else: - query_proj = self.q_einsum('BTD,NDH->BTNH', x) - key_proj, value_proj = self.kv_einsum('BSD,CKDH->CBSKH', x) - - query_proj = positional_embeddings.apply_rope( - query_proj, - segment_pos, - head_dim=self.head_dim, - ) - query_scaled = query_proj * self.head_dim**-0.5 - - key_proj = positional_embeddings.apply_rope( - key_proj, - segment_pos, - head_dim=self.head_dim, - ) - - # Cache is left aligned. - cache['v'] = (cache['v'].at[:bsz, [time_step], :, :].set(value_proj) - ) # values - cache['k'] = (cache['k'].at[:bsz, [time_step], :, :].set(key_proj) - ) # rotated_keys - - logits = jnp.einsum('BTNH,BSNH->BTNS', query_scaled, cache['k']) - logits = logits.astype(jnp.float32) - - padded_logits = jnp.where( - (jnp.expand_dims(attn_mask, -2) >= K_MASK * 0.5), logits, K_MASK) - probs = jax.nn.softmax(padded_logits, axis=-1).astype(cache['k'].dtype) - - encoded = jnp.einsum('BTNS,BSNH->BTNH', probs, cache['v']) - attn_output = self.attn_vec_einsum('BTNH,NHD->BTD', encoded) - - return cache, attn_output - - -class FeedForward(nn.Module): - """Feed forward module.""" - - features: int - hidden_dim: int - - @nn.compact - def __call__(self, x): - w_gating = self.param( - 'gating_einsum', - nn.initializers.zeros_init(), - ((2, self.features, self.hidden_dim)), - ) - ff_gate = jnp.dot(x, w_gating[0]) - gate_value = nn.gelu(ff_gate) - - ff1 = jnp.dot(x, w_gating[1]) - activations = gate_value * ff1 - - w_linear = self.param( - 'linear', - nn.initializers.zeros_init(), - (self.hidden_dim, self.features), - ) - outputs = jnp.dot(activations, w_linear) - - return outputs - - -class Block(nn.Module): - """Transformer block.""" - - num_heads: int - num_kv_heads: int - embed_dim: int - head_dim: int - hidden_dim: int - - def setup(self): - self.pre_attention_norm = layers.RMSNorm() - self.attn = Attention( - num_heads=self.num_heads, - features=self.embed_dim, - head_dim=self.head_dim, - num_kv_heads=self.num_kv_heads, - ) - self.pre_ffw_norm = layers.RMSNorm() - self.mlp = FeedForward(features=self.embed_dim, - hidden_dim=self.hidden_dim) - - def __call__( - self, - x: jax.Array, - segment_pos: int, - cache: LayerCache, - attn_mask: jax.Array, - time_step: int, - ): - inputs_normalized = self.pre_attention_norm(x) - cache, attn_output = self.attn(inputs_normalized, segment_pos, cache, - attn_mask, time_step) - attn_output += x - residual = attn_output - attn_output = self.pre_ffw_norm(attn_output) - outputs = self.mlp(attn_output) - outputs = residual + outputs - return cache, outputs diff --git a/tensorrt_llm/models/gemma/utils/params.py b/tensorrt_llm/models/gemma/utils/params.py deleted file mode 100644 index f0445f918e2e..000000000000 --- a/tensorrt_llm/models/gemma/utils/params.py +++ /dev/null @@ -1,72 +0,0 @@ -# Copyright 2024 DeepMind Technologies Limited. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ============================================================================ -"""Utils for loading Gemma params. - -These utilities are just helpers for current development. They will not be -needed once Gemma switches to Orbax and changes checkpoint formats ahead of -open sourcing. -""" - -import functools -from typing import Any - -Params = dict[str, Any] - - -@functools.cache -def load_params(path: str) -> Params: - import orbax.checkpoint - """Loads parameters from a checkpoint path.""" - checkpointer = orbax.checkpoint.PyTreeCheckpointer() - params = checkpointer.restore(path) - return params - - -def param_remapper(orig_params: Params) -> Params: - """Remaps params to new module layout. - - This is needed here because the model definition does not have a separate - `mlp` module. For the real code release, we will just save the params in a - different format and this will not be needed. - - Args: - orig_params: original dict of parameters in Gemma format. - - Returns: - dict of params with different names. - """ - new_params = {} - for k, v in orig_params.items(): - if 'mlp/' in k: - layer_name, param = k.rsplit('/', maxsplit=1) - if layer_name not in new_params: - new_params[layer_name] = {} - if 'w' in v: - new_params[layer_name][param] = v['w'] - else: - new_params[k] = v - return new_params - - -def nest_params(params: Params) -> Params: - """Nests params as a dict of dicts rather than a flat dict.""" - nested_params = {} - for path, param in params.items(): - *path, leaf = path.split('/') - subdict = nested_params - for key in path: - subdict = subdict.setdefault(key, {}) - subdict[leaf] = param - return nested_params diff --git a/tensorrt_llm/models/gemma/utils/positional_embeddings.py b/tensorrt_llm/models/gemma/utils/positional_embeddings.py deleted file mode 100644 index 25707692d758..000000000000 --- a/tensorrt_llm/models/gemma/utils/positional_embeddings.py +++ /dev/null @@ -1,92 +0,0 @@ -# Copyright 2024 DeepMind Technologies Limited. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ============================================================================ -"""Utils for positional embeddings (including RoPE). -""" - -import jax -import jax.numpy as jnp - -_MAX_WAVELENGTH = 10_000 - - -def add_positional_embedding( - input_embedding: jax.Array, - position: int, - max_wavelength: int = _MAX_WAVELENGTH, -) -> jax.Array: - """Adds positional embeddings to input embeddings.""" - embed_dim = input_embedding.shape[-1] - num_timescales = embed_dim // 2 - log_timescale_increment = jnp.log(float(max_wavelength)) / jnp.maximum( - jnp.asarray(num_timescales, dtype=jnp.float32) - 1, 1) - inv_timescales = jnp.exp( - jnp.arange(num_timescales, dtype=jnp.float32) * - -log_timescale_increment) - scaled_time = position * inv_timescales - signal = jnp.concatenate([jnp.sin(scaled_time), jnp.cos(scaled_time)]) - signal = jnp.pad(signal, [[0, jnp.mod(embed_dim, 2)]]) - position_embedding = signal.astype(jnp.float32) - - return input_embedding + position_embedding - - -def _rotary_embed( - inputs: jax.Array, # [B, 1, H, D] - position: jax.Array, # [B,] - head_dim: int, - max_wavelength: int = _MAX_WAVELENGTH, -) -> jax.Array: - """Helper for RoPE.""" - fraction = 2 * jnp.arange(0, head_dim // 2) / head_dim - timescale = max_wavelength**fraction - timescale = timescale[jnp.newaxis, jnp.newaxis, jnp.newaxis, :] - - sinusoid_inp = position[:, jnp.newaxis, jnp.newaxis, - jnp.newaxis] / timescale - sin = jnp.sin(sinusoid_inp) - cos = jnp.cos(sinusoid_inp) - - first_half, second_half = jnp.split(inputs, 2, axis=-1) - first_part = first_half * cos - second_half * sin - second_part = second_half * cos + first_half * sin - - return jnp.concatenate([first_part, second_part], axis=-1) - - -def apply_rope( - inputs: jax.Array, - position: int, - head_dim: int, - max_wavelength: int = _MAX_WAVELENGTH, -) -> jax.Array: - """Applies RoPE.""" - batch_size, seq_length = inputs.shape[0:2] - - position = jnp.broadcast_to(position, [batch_size])[:, jnp.newaxis] - prefix_position = jnp.arange(seq_length, dtype=jnp.int32) - prefix_position = (position - jnp.flip(prefix_position)[jnp.newaxis, :] - ) # [B, seq_len] - prefix_position = jnp.where(prefix_position < 0, - jnp.zeros_like(prefix_position), - prefix_position).reshape((batch_size, )) - - output = _rotary_embed( - inputs, - position=prefix_position, - head_dim=head_dim, - max_wavelength=max_wavelength, - ) - - return output diff --git a/tensorrt_llm/models/gemma/utils/sampler.py b/tensorrt_llm/models/gemma/utils/sampler.py deleted file mode 100644 index 895a1557d475..000000000000 --- a/tensorrt_llm/models/gemma/utils/sampler.py +++ /dev/null @@ -1,191 +0,0 @@ -# Copyright 2024 DeepMind Technologies Limited. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ============================================================================ -"""Sampler for Gemma transformer. - -An example of a sampling class for a Gemma model. -""" - -import chex -import jax -import jax.numpy as jnp -import sentencepiece as spm - -from . import modules -from . import params as params_lib -from . import transformer as transformer_lib - - -def _compute_attention_masks(time_step: jax.Array, seq_len: int, - input_mask: jax.Array) -> jax.Array: - """Computes causal attention mask.""" - bsz = input_mask.shape[0] - batch_time_step = jnp.full((bsz, 1), time_step, dtype=jnp.uint32) - causal_padding = jnp.greater(jnp.expand_dims(jnp.arange(seq_len), 0), - batch_time_step) - causal_padding = causal_padding * jnp.expand_dims(input_mask, axis=-1) - attention_mask = ( - causal_padding[:, jnp.newaxis, jnp.newaxis, :].astype(jnp.float32) * - modules.K_MASK) - attention_mask = jnp.squeeze(attention_mask, axis=1) - return attention_mask - - -@chex.dataclass -class _SamplingState: - - # Number of tokens in the prompt. - num_input_tokens: jnp.int32 # [B] - - # Fixed-size buffer for accumulating the output tokens. - token_buffer: jnp.ndarray # [B, L] - - # Model state for conditioning the model on autoregressively. - cache: dict[str, modules.LayerCache] - - -class Sampler: - """Sampler for Gemma transformer.""" - - def __init__( - self, - transformer_config: transformer_lib.TransformerConfig, - vocab: spm.SentencePieceProcessor, - params: params_lib.Params, - cache_size: int, - buffer_size: int, - max_decode_steps: int, - ): - self.transformer = transformer_lib.Transformer( - config=transformer_config) - self.vocab = vocab - self.params = params - self.cache_size = cache_size - self.buffer_size = buffer_size - self.max_decode_steps = max_decode_steps - self._compiled_sample_fn = jax.jit(self._sample_fn) - - def _sample_step(self, params, time_step, - sampler_state: _SamplingState) -> _SamplingState: - """Performs a single sampling step.""" - time_step = jnp.asarray(time_step, dtype=jnp.int32) - last_token = sampler_state.token_buffer[:, time_step] - input_mask = last_token != self.vocab.pad_id() - attention_mask = _compute_attention_masks( - time_step, self.cache_size, input_mask).astype(jnp.float32) - - logits, cache = self.transformer.apply( - {'params': params}, - last_token, - time_step, - sampler_state.cache, - attention_mask, - time_step, - ) - - next_token_candidate = jnp.argmax(logits, axis=-1) # [B, 1] - next_token_candidate = next_token_candidate[:, 0] # [B,] - - next_token_candidate = jnp.where( - time_step < sampler_state.num_input_tokens - 1, - sampler_state.token_buffer[:, time_step + 1], - next_token_candidate, - ) - - token_buffer = sampler_state.token_buffer.at[:, time_step + 1].set( - next_token_candidate) - - return _SamplingState( - num_input_tokens=sampler_state.num_input_tokens, - token_buffer=token_buffer, - cache=cache, - ) - - def init_cache(self, bsz) -> dict[str, modules.LayerCache]: - """Initializes the attention cache for each layer.""" - return { - f'layer_{i}': - modules.init_layer_cache( - self.cache_size, - self.transformer.config.num_heads, - self.transformer.config.head_dim, - bsz, - ) - for i in range(self.transformer.config.num_layers) - } - - def init_sample_state(self, - all_input_ids: list[jax.Array]) -> _SamplingState: - """Initializes the sampling state given input prompts.""" - bsz = len(all_input_ids) - num_input_tokens = [len(input_ids) for input_ids in all_input_ids] - - token_buffer = jnp.full( - ( - bsz, - self.buffer_size, - ), - self.vocab.pad_id(), - dtype=jnp.int32, - ) - for i, (input_ids, - num_tokens) in enumerate(zip(all_input_ids, num_input_tokens)): - token_buffer = token_buffer.at[i, :num_tokens].set(input_ids) - - return _SamplingState( - num_input_tokens=jnp.array(num_input_tokens, dtype=jnp.int32), - token_buffer=token_buffer, - cache=self.init_cache(bsz), - ) - - def tokenize(self, input_string: str) -> jax.Array: - """Tokenizes the input string.""" - input_ids = self.vocab.EncodeAsIds(input_string) - input_ids = jnp.array([self.vocab.bos_id()] + - jnp.array(input_ids).tolist(), - dtype=jnp.int32) - return input_ids - - def _sample_fn( - self, - params: params_lib.Params, - initial_sampling_state: _SamplingState, - ) -> _SamplingState: - - def sample_with_params(time_step: int, sampler_state: _SamplingState): - return self._sample_step(params, time_step, sampler_state) - - return jax.lax.fori_loop(0, self.max_decode_steps, sample_with_params, - initial_sampling_state) - - def __call__(self, input_strings: list[str] | str) -> list[str]: - """Samples a completion of the input string.""" - if isinstance(input_strings, str): - input_strings = [input_strings] - all_input_ids = [self.tokenize(x) for x in input_strings] - initial_sampling_state = self.init_sample_state(all_input_ids) - - sampling_state = self._compiled_sample_fn(self.params, - initial_sampling_state) - - out_tokens = [ - buffer[num_tokens:num_tokens + self.max_decode_steps] - for buffer, num_tokens in zip(sampling_state.token_buffer, - sampling_state.num_input_tokens) - ] - decoded_outputs = [ - self.vocab.DecodeIds(out_tokens.tolist()) - for out_tokens in out_tokens - ] - return decoded_outputs diff --git a/tensorrt_llm/models/gemma/utils/transformer.py b/tensorrt_llm/models/gemma/utils/transformer.py deleted file mode 100644 index de3a0da61b7b..000000000000 --- a/tensorrt_llm/models/gemma/utils/transformer.py +++ /dev/null @@ -1,114 +0,0 @@ -"""Gemma transformer.""" - -import dataclasses - -import jax -import jax.numpy as jnp -from flax import linen as nn - -from . import layers, modules -from . import params as params_lib - -Cache = dict[str, modules.LayerCache] - - -@dataclasses.dataclass -class TransformerConfig: - """Configuration for the Gemma transformer.""" - - num_layers: int - num_embed: int - embed_dim: int - hidden_dim: int - num_heads: int - head_dim: int - num_kv_heads: int - - @classmethod - def from_params(cls, params: params_lib.Params, - num_embed: int) -> 'TransformerConfig': - """Creates a TransformerConfig from loaded parameters.""" - num_layers = (max([ - int(k.split('_')[1]) - for k in params['transformer'].keys() if 'layer_' in k - ]) + 1) - hidden_dim, embed_dim = ( - params['transformer']['layer_0']['mlp']['linear'].shape) - num_heads, head_dim, _ = (params['transformer']['layer_0']['attn'] - ['attn_vec_einsum']['w'].shape) - use_qkv_einsum = 'qkv_einsum' in params['transformer']['layer_0'][ - 'attn'] - if use_qkv_einsum: - num_kv_heads = num_heads - else: - num_kv_heads = params['transformer']['layer_0']['attn'][ - 'kv_einsum']['w'].shape[1] - return cls( - num_layers=num_layers, - num_embed=num_embed, - embed_dim=embed_dim, - hidden_dim=hidden_dim, - num_heads=num_heads, - head_dim=head_dim, - num_kv_heads=num_kv_heads, - ) - - -def init_cache(config: TransformerConfig, cache_size: int, - batch_size: int) -> Cache: - """Initializes a new Transformer cache.""" - return { - f'layer_{i}': - modules.init_layer_cache(cache_size, config.num_heads, config.head_dim, - batch_size) - for i in range(config.num_layers) - } - - -class Transformer(nn.Module): - """Gemma transformer.""" - - config: TransformerConfig - - def setup(self): - self.embedder = modules.Embedder( - vocab_size=self.config.num_embed, - embed_dim=self.config.embed_dim, - ) - self.blocks = [ - modules.Block( - name=f'layer_{i}', - num_heads=self.config.num_heads, - num_kv_heads=self.config.num_kv_heads, - embed_dim=self.config.embed_dim, - head_dim=self.config.head_dim, - hidden_dim=self.config.hidden_dim, - ) for i in range(self.config.num_layers) - ] - self.final_norm = layers.RMSNorm() - - def __call__( - self, - last_tokens: jax.Array, # [B,] - current_token_position: int, - cache: Cache, - attention_mask: jax.Array, # [B, 1, L] - time_step: int, - ) -> tuple[jax.Array, Cache]: - input_emb = self.embedder.encode(last_tokens) - x = jnp.expand_dims(input_emb, axis=1) # adding temporal dimension - - for i, block in enumerate(self.blocks): - layer_name = f'layer_{i}' - cache[layer_name], x = block( - x, - current_token_position, - cache[layer_name], - attention_mask, - time_step, - ) - - x = self.final_norm(x) - logits = self.embedder.decode(x) - - return logits, cache diff --git a/tensorrt_llm/models/gemma/weight.py b/tensorrt_llm/models/gemma/weight.py deleted file mode 100644 index 56b4a91e4048..000000000000 --- a/tensorrt_llm/models/gemma/weight.py +++ /dev/null @@ -1,181 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -import os -from typing import Dict, Optional - -import numpy as np -import torch - -from tensorrt_llm.models.modeling_utils import PretrainedConfig - -from ..._utils import numpy_to_torch -from ...logger import logger -from ...mapping import Mapping - -Weights = Dict[str, torch.Tensor] - - -def quantize_fp8_weights(weights: Weights, num_layers: int, - mapping: Mapping) -> Weights: - - def get_scaling_factor(weight): - amax = weight.max() - scale = 448.0 / amax - return scale - - layers_range = mapping.pp_layers(num_layers) - scaling_factors = {} - scaled_weights = {} - trt_llm_prefix = "transformer.layers" - for l in layers_range: - # attention.qkv.weight - for name in [ - "attention.qkv", "attention.dense", "mlp.fc", "mlp.gate", - "mlp.proj" - ]: - trt_llm_name = ".".join((trt_llm_prefix, str(l), name, "weight")) - scale_name = ".".join( - (trt_llm_prefix, str(l), name, "weights_scaling_factor")) - weight = weights[trt_llm_name].float() - dtype = weights[trt_llm_name].dtype - scale = get_scaling_factor(weight) - scaled_weights[trt_llm_name] = (weight * - scale).to(dtype).contiguous() - scaling_factors[scale_name] = numpy_to_torch( - np.asarray([1 / scale]).astype(np.float32)) - return scaling_factors - - -def load_from_fp8_gemma(quant_ckpt_path: Optional[str], num_layers: int, - mapping: Mapping, fp8_kv_cache: bool, - weight_scales: Weights): - """ - Get the fp8 scaling factors. - """ - fake_fp8_sf_dt = torch.float32 - - if quant_ckpt_path is not None and os.path.isfile(quant_ckpt_path): - fp8_gemma = np.load(quant_ckpt_path) - else: - fp8_gemma = None - logger.info( - f"There is not quantized checkpoint, use dummy fp8 scaling factors instead." - ) - weights = {} - - def get_fp8_gemma(name: str) -> np.ndarray: - if fp8_gemma is not None: - return fp8_gemma[name] - else: - return torch.tensor([1.0], dtype=fake_fp8_sf_dt).numpy() - - layers_range = mapping.pp_layers(num_layers) - for l in layers_range: - prefix = f'_np:layers:{l}' - tllm_prex = f'transformer.layers.{l-layers_range[0]}' - - weights[f'{tllm_prex}.attention.qkv.activation_scaling_factor'] = max( - get_fp8_gemma( - f'{prefix}:attention:qkv:q:activation_scaling_factor'), - get_fp8_gemma( - f'{prefix}:attention:qkv:k:activation_scaling_factor'), - get_fp8_gemma( - f'{prefix}:attention:qkv:v:activation_scaling_factor')) - weights[f'{tllm_prex}.attention.qkv.weights_scaling_factor'] = max( - get_fp8_gemma(f'{prefix}:attention:qkv:q:weights_scaling_factor'), - get_fp8_gemma(f'{prefix}:attention:qkv:k:weights_scaling_factor'), - get_fp8_gemma(f'{prefix}:attention:qkv:v:weights_scaling_factor')) - weights[ - f'{tllm_prex}.attention.dense.activation_scaling_factor'] = get_fp8_gemma( - f'{prefix}:attention:dense:activation_scaling_factor') - weights[ - f'{tllm_prex}.attention.dense.weights_scaling_factor'] = get_fp8_gemma( - f'{prefix}:attention:dense:weights_scaling_factor') - - weights[ - f'{tllm_prex}.mlp.fc.activation_scaling_factor'] = get_fp8_gemma( - f'{prefix}:mlp:fc:activation_scaling_factor') - weights[f'{tllm_prex}.mlp.fc.weights_scaling_factor'] = get_fp8_gemma( - f'{prefix}:mlp:fc:weights_scaling_factor') - - weights[ - f'{tllm_prex}.mlp.gate.activation_scaling_factor'] = get_fp8_gemma( - f'{prefix}:mlp:gate:activation_scaling_factor') - weights[f'{tllm_prex}.mlp.gate.weights_scaling_factor'] = get_fp8_gemma( - f'{prefix}:mlp:gate:weights_scaling_factor') - - weights[ - f'{tllm_prex}.mlp.proj.activation_scaling_factor'] = get_fp8_gemma( - f'{prefix}:mlp:proj:activation_scaling_factor') - weights[f'{tllm_prex}.mlp.proj.weights_scaling_factor'] = get_fp8_gemma( - f'{prefix}:mlp:proj:weights_scaling_factor') - - if fp8_kv_cache: - # Not calibrating KV cache. - scaling_factor = 1.0 - weights[ - f'{tllm_prex}.attention.kv_cache_scaling_factor'] = torch.tensor( - [scaling_factor], dtype=fake_fp8_sf_dt).numpy() - if fp8_gemma is None: - weights.update(weight_scales) - - for key in weights: - if isinstance(weights[key], np.ndarray): - weights[key] = numpy_to_torch(weights[key]) - return weights - - -def dummy_weights_awq(weights: Weights, precision: str, - trt_llm_config: PretrainedConfig, group_size: int): - packer = torch.ops.trtllm.pack_int8_tensor_to_packed_int4 - use_fp8_kv_cache = trt_llm_config.quant_mode.has_fp8_kv_cache() - use_int8_kv_cache = trt_llm_config.quant_mode.has_int8_kv_cache() - num_layers = trt_llm_config.num_hidden_layers - for name in list(weights): - if any([ - _name in name for _name in [ - 'mlp.proj.weight', 'mlp.gate.weight', 'mlp.fc.weight', - 'attention.qkv.weight', 'attention.dense.weight' - ] - ]): - print("Processing:", name) - weight = np.ascontiguousarray(weights[name].T) - in_dim, out_dim = weight.shape - scale = np.amax(weight) / 7 - weights_scaling_factor = np.ones([out_dim, in_dim // group_size - ]) * scale.astype(np.float32) - weight_smoothed = (weight.astype(np.float32) / scale).astype( - np.int8) - weight_smoothed[weight_smoothed < -8] = -8 - weight_smoothed[weight_smoothed > 7] = 7 - prequant_scaling_factor = np.ones([in_dim], dtype=weight.dtype) - weights[name] = packer( - torch.from_numpy(weight_smoothed)).T.contiguous().numpy() - weights[name.replace( - 'weight', 'prequant_scaling_factor')] = prequant_scaling_factor - weights[name.replace( - 'weight', - 'weights_scaling_factor')] = weights_scaling_factor.astype( - weight.dtype) - if precision == "w4a8_awq": - alpha = np.array([1], dtype=np.float32) - weights[name.replace('weight', 'alpha')] = alpha - if use_fp8_kv_cache or use_int8_kv_cache: - for l in range(num_layers): - t = np.array([1], dtype=np.float32) - weights[ - f"transformer.layers.{l}.attention.kv_cache_scaling_factor"] = t - - return weights diff --git a/tensorrt_llm/models/generation_mixin.py b/tensorrt_llm/models/generation_mixin.py deleted file mode 100644 index 5f590f1ee4bb..000000000000 --- a/tensorrt_llm/models/generation_mixin.py +++ /dev/null @@ -1,889 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -import math -from collections import OrderedDict -from typing import List, Optional - -import tensorrt as trt - -from ..functional import Tensor -from ..layers import MropeParams, SpecDecodingParams -from ..llmapi.kv_cache_type import KVCacheType -from ..mapping import Mapping -from ..plugin import current_all_reduce_helper - - -class GenerationMixin: - - @staticmethod - def has_ctx_gen_opt_profiles( - use_gpt_attention_plugin: bool = False, - use_gemm_plugin: bool = False, - use_mamba_conv1d_plugin: bool = False, - remove_input_padding: bool = False, - paged_state: bool = False, - kv_cache_type: KVCacheType = KVCacheType.CONTINUOUS) -> bool: - res = False - - if not use_gpt_attention_plugin or not use_gemm_plugin: - use_in_flight_batching = False - # Refer to modelConfig.h: supportsInflightBatching(), this should be consistent its implementation. - # We skip check transformer or rnn arch for simplification. - if remove_input_padding and use_gpt_attention_plugin: - use_in_flight_batching = kv_cache_type in [ - KVCacheType.PAGED, KVCacheType.DISABLED - ] - elif remove_input_padding and use_mamba_conv1d_plugin: - use_in_flight_batching = paged_state == True - - res = not use_in_flight_batching - - return res - - @staticmethod - def default_range(max_range, offset=0, min_range=1, opt_offset=0): - result = [ - min_range, (max_range + min_range + opt_offset) // 2, max_range - ] - return [elem + offset for elem in result] - - @staticmethod - def split_num_tokens_range(max_num_tokens): - split_point = [64, 128, 256, 512, 1024] - num_tokens_ranges = [] - for i, p in enumerate(split_point): - if i == 0 and max_num_tokens <= p: - return [1, max_num_tokens, max_num_tokens] - elif max_num_tokens <= p: - num_tokens_ranges.append( - [split_point[i - 1], max_num_tokens, max_num_tokens]) - return num_tokens_ranges - elif i == 0 and max_num_tokens > p: - num_tokens_ranges = [[1, 64, 64]] - else: - num_tokens_ranges.append( - [split_point[i - 1], split_point[i], split_point[i]]) - num_tokens_ranges.append( - [split_point[-1], max_num_tokens, max_num_tokens]) - return num_tokens_ranges - - @staticmethod - def get_profiles_ranges( - *, - max_batch_size, - max_beam_width, - max_input_len, - max_num_tokens, - max_draft_len, - opt_batch_size, - opt_num_tokens, - enable_ctx_gen_opt_profiles, - multiple_profiles, - kv_cache_type: KVCacheType = KVCacheType.CONTINUOUS): - default_range = GenerationMixin.default_range - if opt_batch_size: - bb_range_cxt = [1, opt_batch_size, max_batch_size] - bb_range_gen = [ - 1, opt_batch_size * max_beam_width, - max_batch_size * max_beam_width - ] - else: - bb_range_cxt = default_range(max_batch_size) - bb_range_gen = default_range(max_batch_size * max_beam_width) - tokens_per_engine_step = max_draft_len + 1 - tokens_per_engine_step_range = [ - 1, tokens_per_engine_step, tokens_per_engine_step - ] - bbd_range_ctx = [ - bb_range_cxt[i] * (tokens_per_engine_step if i != 0 else 1) - for i in range(len(bb_range_cxt)) - ] - bbd_range_gen = [ - bb_range_gen[i] * (tokens_per_engine_step if i != 0 else 1) - for i in range(len(bb_range_gen)) - ] - inlen_range_cxt = default_range(max_input_len) - inlen_range_gen = [1, 1, tokens_per_engine_step] - if enable_ctx_gen_opt_profiles: - num_profiles = 2 - bb_range = [bb_range_cxt, bb_range_gen] - bbd_range = [bbd_range_ctx, bbd_range_gen] - inlen_range = [inlen_range_cxt, inlen_range_gen] - position_ids_inlen_range = [inlen_range_cxt, [1, 1, 1]] - num_tokens_range_ctx = default_range(max_batch_size * max_input_len) - # Draft tokens cannot be combined with beam search - num_tokens_range_gen = default_range( - max_batch_size * max(tokens_per_engine_step, max_beam_width)) - num_tokens_range = [num_tokens_range_ctx, num_tokens_range_gen] - - # Only keep context range when kv cache is disabled. - if kv_cache_type == KVCacheType.DISABLED: - num_profiles = 1 - bb_range = [bb_range[0]] - bbd_range = [bbd_range[0]] - inlen_range = [inlen_range[0]] - position_ids_inlen_range = [position_ids_inlen_range[0]] - num_tokens_range_ctx = [num_tokens_range_ctx[0]] - # Draft tokens cannot be combined with beam search - num_tokens_range_gen = [num_tokens_range_gen[0]] - num_tokens_range = [num_tokens_range[0]] - - else: - if multiple_profiles: - num_tokens_range = GenerationMixin.split_num_tokens_range( - max_num_tokens) - else: - if opt_num_tokens is None: - opt_num_tokens = min(max_num_tokens, - max_batch_size * max_beam_width) - num_tokens_range = [[1, opt_num_tokens, max_num_tokens]] - num_profiles = len(num_tokens_range) - bb_range = [bb_range_gen] * num_profiles - bbd_range = [bbd_range_gen] * num_profiles - inlen_range = [[1, 1, max_input_len]] * num_profiles - position_ids_inlen_range = [[1, 1, max_input_len]] * num_profiles - tokens_per_engine_step_range = [tokens_per_engine_step_range - ] * num_profiles - - position_ids_num_tokens_range = num_tokens_range - # If max_draft_len != 0, the input_ids may include draft tokens. And the length of position_ids may be not the same as input_ids. - # In extreme cases, input_ids may contain (max_draft_token + 1) * N, and the actual position_ids value is only 1 * N. - # Therefore, we need to adjust the min value in the ranges of position_ids. - if max_draft_len != 0: - position_ids_num_tokens_range = list( - map( - lambda x: - [math.ceil(x[0] / (max_draft_len + 1)), x[1], x[2]], - num_tokens_range)) - - ranges = { - 'bb_range': bb_range, - 'bbd_range': bbd_range, - 'inlen_range': inlen_range, - 'position_ids_inlen_range': position_ids_inlen_range, - 'num_tokens_range': num_tokens_range, - 'tokens_per_engine_step_range': tokens_per_engine_step_range, - 'position_ids_num_tokens_range': position_ids_num_tokens_range, - } - return num_profiles, ranges - - def prepare_attention_inputs( - self, - *, - max_batch_size, - max_beam_width, - max_input_len, - max_seq_len, - num_kv_heads, - head_size, - num_layers, - kv_dtype, - kv_cache_type: KVCacheType, - num_profiles=1, - enable_ctx_gen_opt_profiles=False, - remove_input_padding=False, - use_gpt_attention_plugin=False, - tokens_per_block=32, - mapping=Mapping(), - streamingllm=False, - attn_layer_idx=None, - opt_batch_size=None, - num_kv_heads_per_layer: Optional[List[int]] = None): - - if attn_layer_idx is not None and num_kv_heads_per_layer is not None: - assert len(attn_layer_idx) == len(num_kv_heads_per_layer), ( - f"Expected len(attn_layer_idx) ({len(attn_layer_idx)})" - f" == len(num_kv_heads_per_layer) ({len(num_kv_heads_per_layer)})" - ) - default_range = GenerationMixin.default_range - - if opt_batch_size: - bb_range_cxt = [1, opt_batch_size, max_batch_size] - bb_range_gen = [ - 1, opt_batch_size * max_beam_width, - max_batch_size * max_beam_width - ] - else: - bb_range_cxt = default_range(max_batch_size) - bb_range_gen = default_range(max_batch_size * max_beam_width) - - _bs_range = default_range(max_batch_size) - _beam_width_range = default_range(max_beam_width) - _max_len_range = default_range(max_seq_len) - _mask_len_ctx = default_range(max_input_len) - _kv_cache_range_ctx = [0, 0, 0] - _kv_cache_range_gen = default_range(max_seq_len, -1) - if kv_cache_type == KVCacheType.DISABLED: - _kv_cache_range = default_range(max_seq_len) - else: - kv_max_seq_len = max_seq_len - if streamingllm: - # add the max bubble length - kv_max_seq_len += tokens_per_block - 1 - if max_beam_width > 1: - # support cyclic kv cache cases that use one more block - kv_max_seq_len += tokens_per_block - _kv_cache_range = default_range(kv_max_seq_len) - - if enable_ctx_gen_opt_profiles: - if kv_cache_type != KVCacheType.DISABLED: - assert num_profiles == 2 - bb_range = [bb_range_cxt, bb_range_gen] - mask_len_range = [_mask_len_ctx, _max_len_range] - if use_gpt_attention_plugin: - kv_cache_range = [_kv_cache_range, _kv_cache_range] - else: - kv_cache_range = [_kv_cache_range_ctx, _kv_cache_range_gen] - else: - assert num_profiles == 1 - bb_range = [bb_range_cxt] - mask_len_range = [_mask_len_ctx] - if use_gpt_attention_plugin: - kv_cache_range = [_kv_cache_range] - else: - kv_cache_range = [_kv_cache_range_ctx] - - else: - bb_range = [bb_range_gen] * num_profiles - mask_len_range = [_max_len_range] * num_profiles - kv_cache_range = [_kv_cache_range] * num_profiles - bs_range = [_bs_range] * num_profiles - beam_width_range = [_beam_width_range] * num_profiles - max_len_range = [_max_len_range] * num_profiles - - num_kv_heads = (num_kv_heads + mapping.tp_size - 1) // mapping.tp_size - if num_kv_heads_per_layer is not None: - num_kv_heads_per_layer = [ - (nheads + mapping.tp_size - 1) // mapping.tp_size - for nheads in num_kv_heads_per_layer - ] - - layers_range = mapping.pp_layers(num_layers) - if attn_layer_idx is None: - attn_layer_idx = [i for i in range(num_layers)] - # layer indices of attention layers local to the current pp rank - local_attn_layers = [i for i in layers_range if i in attn_layer_idx] - # number of attention layers local to previous pp ranks - num_attn_layers_lower_ranks = attn_layer_idx.index(local_attn_layers[0]) - num_attn_layers = len(local_attn_layers) - num_layers_prev_rank = layers_range[ - 0] // mapping.pp_rank if mapping.pp_rank != 0 else len(layers_range) - past_key_value = [] - kv_cache_block_offsets = None - host_kv_cache_block_offsets = None - host_kv_cache_pool_pointers = None - host_kv_cache_pool_mapping = None - if kv_cache_type == KVCacheType.DISABLED: - past_key_value = [None] * num_layers_prev_rank - else: - if kv_cache_type != KVCacheType.PAGED: - for layer_idx in layers_range: - if layer_idx not in local_attn_layers: - # not an attention layer ==> give it None pkv input - past_key_value.append(None) - continue - - attn_idx = local_attn_layers.index(layer_idx) - if num_kv_heads_per_layer is not None: - heads_dim_name = f"num_heads_{layer_idx}" - kv_heads = num_kv_heads_per_layer[ - num_attn_layers_lower_ranks + attn_idx] - else: - heads_dim_name = "num_heads" - kv_heads = num_kv_heads - - kv_dim_range = OrderedDict([ - ('batch_size_beam_width', bb_range), - ('kv', [2] * num_profiles), - (heads_dim_name, [kv_heads] * num_profiles), - ('past_key_len', kv_cache_range), - ('head_size', [head_size] * num_profiles), - ]) - - kv = Tensor(name=f'past_key_value_{layer_idx}', - dtype=kv_dtype, - shape=[-1, 2, kv_heads, -1, head_size], - dim_range=kv_dim_range) - past_key_value.append(kv) - else: - if enable_ctx_gen_opt_profiles: - max_blocks_per_seq_range = [ - [ - math.ceil(kv_cache_range[0][0] / tokens_per_block), - math.ceil(kv_cache_range[0][1] / tokens_per_block), - math.ceil(kv_cache_range[0][2] / tokens_per_block) - ], - [ - math.ceil(kv_cache_range[1][0] / tokens_per_block), - math.ceil(kv_cache_range[1][1] / tokens_per_block), - math.ceil(kv_cache_range[1][2] / tokens_per_block) - ] - ] - else: - max_blocks_per_seq_range = [[ - math.ceil(kv_cache_range[0][0] / tokens_per_block), - math.ceil(kv_cache_range[0][1] / tokens_per_block), - math.ceil(kv_cache_range[0][2] / tokens_per_block) - ]] * num_profiles - - NUM_KV_CACHE_POOLS = -1 # the number of unique variable window sizes, which is only known at runtime, affects the number of pools. - # dim_range=(min=1, opt=1 (this is the usual case - non vgqa, non vsliding_window), max=num_layers, - # TODO(nhaber): Benchmark if making NUM_KV_CACHE_POOLS dynamic has a significant performance hit? - - kv_pools_range = [[1, 1, len(local_attn_layers)]] * num_profiles - kv_cache_block_offsets = Tensor( - name=f'kv_cache_block_offsets', - dtype=trt.int32, - shape=[NUM_KV_CACHE_POOLS, -1, 2, -1], - dim_range=OrderedDict([ - ('num_kv_cache_pools', kv_pools_range), - ('batch_size_beam_width', bb_range), - ('kv', [2] * num_profiles), - ('max_blocks_per_seq', max_blocks_per_seq_range), - ])) - host_kv_cache_block_offsets = Tensor( - name=f'host_kv_cache_block_offsets', - dtype=trt.int32, - shape=[NUM_KV_CACHE_POOLS, -1, 2, -1], - dim_range=OrderedDict([ - ('num_kv_cache_pools', kv_pools_range), - ('batch_size_beam_width', bb_range), - ('kv', [2] * num_profiles), - ('max_blocks_per_seq', max_blocks_per_seq_range), - ])) - host_kv_cache_pool_pointers = Tensor( - name=f'host_kv_cache_pool_pointers', - dtype=trt.int64, - shape=[NUM_KV_CACHE_POOLS, 2], - dim_range=OrderedDict([ - ('num_pools_layers', kv_pools_range), - ('num_pools_kv', [2] * num_profiles), - ])) - - host_kv_cache_pool_mapping = Tensor( - name=f'host_kv_cache_pool_mapping', - dtype=trt.int32, - shape=[num_attn_layers, - 2], # 2: (Index of pool, Index of layer within pool) - dim_range=OrderedDict([ - ('pools_mapping', [num_attn_layers] * num_profiles), - ('layer_cache_pool_locator', [2] * num_profiles) - ])) - - past_key_value = [None] * num_layers_prev_rank - - assert len(past_key_value) == num_layers_prev_rank - sequence_length = None - context_lengths = None - host_context_lengths = None - host_past_key_value_lengths = None - host_max_attention_window_sizes = None - host_sink_token_length = None - attention_mask = None - cache_indirection = None - host_request_types = None - runtime_perf_knobs = None - context_progress = None - - if use_gpt_attention_plugin: - if kv_cache_type != KVCacheType.DISABLED: - sequence_length = Tensor( - name='sequence_length', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([('batch_size_beam_width', bb_range) - ]), - ) - - host_request_types = Tensor( - name='host_request_types', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([('batch_size_beam_width', bb_range)]), - ) - if kv_cache_type != KVCacheType.DISABLED: - host_past_key_value_lengths = Tensor( - name='host_past_key_value_lengths', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([('batch_size_beam_width', bb_range) - ]), - ) - context_lengths = Tensor( - name='context_lengths', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([('batch_size_beam_width', bb_range)]), - ) - runtime_perf_knobs = Tensor(name='host_runtime_perf_knobs', - dtype=trt.int64, - shape=[16], - dim_range=OrderedDict([ - ('perf_knob_size', - [16] * num_profiles) - ])) - context_progress = Tensor(name='host_context_progress', - dtype=trt.int64, - shape=[1], - dim_range=OrderedDict([ - ('context_progress_size', - [1] * num_profiles) - ])) - else: - attention_mask = Tensor( - name='attention_mask', - dtype=trt.int32, - shape=[-1, -1], - dim_range=OrderedDict([ - ('batch_size_beam_width', bb_range), - ('mask_len', mask_len_range), - ]), - ) - - if use_gpt_attention_plugin and remove_input_padding: - host_context_lengths = Tensor( - name='host_context_lengths', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([('batch_size_beam_width', bb_range)]), - ) - - if use_gpt_attention_plugin: - # TODO: change shape to [1] - host_max_attention_window_sizes = Tensor( - name=f'host_max_attention_window_sizes', - dtype=trt.int32, - shape=[num_attn_layers], - dim_range=OrderedDict([('num_layers', - [num_attn_layers] * num_profiles)])) - - host_sink_token_length = Tensor(name='host_sink_token_length', - dtype=trt.int32, - shape=[1], - dim_range=OrderedDict([ - ('scalar', [1] * num_profiles) - ])) - - if kv_cache_type != KVCacheType.DISABLED: - cache_indirection = Tensor( - name='cache_indirection', - dtype=trt.int32, - shape=[-1, -1, -1], - dim_range=OrderedDict([ - ('batch_size_cache', bs_range), - ('beam_width', beam_width_range), - ('max_seq_len', max_len_range), - ]), - ) - - return { - 'attention_mask': attention_mask, - 'sequence_length': sequence_length, - 'host_past_key_value_lengths': host_past_key_value_lengths, - 'host_max_attention_window_sizes': host_max_attention_window_sizes, - 'host_sink_token_length': host_sink_token_length, - 'past_key_value': past_key_value, - 'cache_indirection': cache_indirection, - 'kv_cache_block_offsets': kv_cache_block_offsets, - 'host_kv_cache_block_offsets': host_kv_cache_block_offsets, - 'host_kv_cache_pool_pointers': host_kv_cache_pool_pointers, - 'host_kv_cache_pool_mapping': host_kv_cache_pool_mapping, - 'context_lengths': context_lengths, - 'host_context_lengths': host_context_lengths, - 'host_request_types': host_request_types, - 'host_runtime_perf_knobs': runtime_perf_knobs, - 'host_context_progress': context_progress, - } - - def prepare_basic_inputs( - self, - *, - max_batch_size, - max_beam_width, - max_input_len, - max_seq_len, - max_num_tokens, - hidden_size, - num_kv_heads, - head_size, - num_layers, - kv_dtype, - kv_cache_type: KVCacheType, - remove_input_padding=False, - use_gpt_attention_plugin=False, - use_gemm_plugin=False, - tokens_per_block=32, - gather_context_logits=False, - dtype=None, - num_heads=None, - mapping=Mapping(), - opt_num_tokens=None, - prompt_embedding_table_size: int = 0, - position_encoding_2d=False, - use_lora_plugin: bool = False, - lora_target_modules: List[str] = None, - speculative_decoding_draft_tokens_external: bool = False, - spec_decoding_is_generation_length_variable: bool = False, - max_draft_len=0, - multiple_profiles: bool = False, - streamingllm: bool = False, - opt_batch_size=None, - pp_reduce_scatter: bool = False, - mrope_rotary_cos_sin_size: int = None, - ): - - enable_ctx_gen_opt_profiles = GenerationMixin.has_ctx_gen_opt_profiles( - use_gpt_attention_plugin=use_gpt_attention_plugin, - use_gemm_plugin=use_gemm_plugin, - remove_input_padding=remove_input_padding, - kv_cache_type=kv_cache_type) - - num_profiles, ranges = GenerationMixin.get_profiles_ranges( - max_batch_size=max_batch_size, - max_beam_width=max_beam_width, - max_input_len=max_input_len, - max_num_tokens=max_num_tokens, - max_draft_len=max_draft_len, - opt_batch_size=opt_batch_size, - opt_num_tokens=opt_num_tokens, - enable_ctx_gen_opt_profiles=enable_ctx_gen_opt_profiles, - multiple_profiles=multiple_profiles, - kv_cache_type=kv_cache_type) - bb_range = ranges['bb_range'] - bbd_range = ranges['bbd_range'] - inlen_range = ranges['inlen_range'] - num_tokens_range = ranges['num_tokens_range'] - position_ids_inlen_range = ranges['position_ids_inlen_range'] - tokens_per_engine_step_range = ranges['tokens_per_engine_step_range'] - position_ids_num_tokens_range = ranges['position_ids_num_tokens_range'] - - input_ids = None - position_ids = None - hidden_states = None - if remove_input_padding: - if mapping.is_first_pp_rank(): - input_ids = Tensor(name='input_ids', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([ - ('num_tokens', num_tokens_range), - ])) - if position_encoding_2d: - position_ids = Tensor( - name='position_ids', - dtype=trt.int32, - shape=[2, -1], - dim_range=OrderedDict([ - ('2', [2] * num_profiles), - ('position_ids_num_tokens_range', - position_ids_num_tokens_range), - ]), - ) - else: - position_ids = Tensor( - name='position_ids', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([ - ('position_ids_num_tokens_range', - position_ids_num_tokens_range), - ]), - ) - else: - assert dtype is not None - assert num_heads is not None - pp_hidden_size = hidden_size // mapping.tp_size if pp_reduce_scatter else hidden_size - hidden_states = Tensor( - name='hidden_states_input', - dtype=dtype, - shape=[-1, pp_hidden_size], - dim_range=OrderedDict([ - ('num_tokens', num_tokens_range), - ('hidden_size', [pp_hidden_size] * num_profiles), - ]), - ) - - else: - if mapping.is_first_pp_rank(): - input_ids = Tensor(name='input_ids', - dtype=trt.int32, - shape=[-1, -1], - dim_range=OrderedDict([ - ('batch_size_beam_width', bb_range), - ('input_len', inlen_range), - ])) - if position_encoding_2d: - position_ids = Tensor( - name='position_ids', - dtype=trt.int32, - shape=[-1, 2, -1], - dim_range=OrderedDict([ - ('batch_size_beam_width', bb_range), - ('2', [2] * num_profiles), - ('position_ids_inlen_range', - position_ids_inlen_range), - ]), - ) - else: - position_ids = Tensor( - name='position_ids', - dtype=trt.int32, - shape=[-1, -1], - dim_range=OrderedDict([ - ('batch_size_beam_width', bb_range), - ('position_ids_inlen_range', - position_ids_inlen_range), - ]), - ) - else: - assert dtype is not None - assert num_heads is not None - hidden_states = Tensor( - name='hidden_states_input', - dtype=dtype, - shape=[-1, -1, hidden_size], - dim_range=OrderedDict([ - ('batch_size_beam_width', bb_range), - ('input_len', inlen_range), - ('hidden_size', [hidden_size] * num_profiles), - ]), - ) - - if mapping.tp_size > 1: - current_all_reduce_helper().set_workspace_tensor( - mapping, num_profiles) - - prompt_embedding_table = None - tasks = None - prompt_vocab_size = None - if prompt_embedding_table_size > 0: - _p_embedding_range = [ - 1, prompt_embedding_table_size // 2, prompt_embedding_table_size - ] - p_embedding_range = [_p_embedding_range] * num_profiles - - prompt_embedding_table = Tensor(name='prompt_embedding_table', - dtype=dtype, - shape=[-1, hidden_size], - dim_range=OrderedDict([ - ('prompt_embedding_table_size', - p_embedding_range), - ('hidden_size', - [hidden_size] * num_profiles), - ])) - if remove_input_padding: - tasks = Tensor(name='tasks', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([ - ('input_len_task', num_tokens_range), - ])) - else: - tasks = Tensor(name='tasks', - dtype=trt.int32, - shape=[-1, 1], - dim_range=OrderedDict([ - ('batch_size_beam_width', bb_range), - ('broadcast_dim', [1] * num_profiles), - ])) - prompt_vocab_size = Tensor(name='prompt_vocab_size', - dtype=trt.int32, - shape=[1], - dim_range=OrderedDict([ - ('size', [1] * num_profiles) - ])) - - lora_weights_pointers = None - lora_ranks = None - if use_lora_plugin: - lora_weights_pointers = [] - lora_ranks = [] - layers_range = mapping.pp_layers(num_layers) - for i in layers_range: - lora_weight_pointer_dict = {} - lora_rank_dict = {} - for lora_module in lora_target_modules: - - lora_weight_pointer = Tensor( - name=f'{lora_module}_lora_weights_pointers_{i}', - dtype=trt.int64, - shape=[-1, 3], - dim_range=OrderedDict([ - ('batch_size_beam_width', bb_range), - ('in_out_scales', [3] * num_profiles), - ])) - lora_weight_pointer_dict.update({ - f"{lora_module}_lora_weights_pointers": - lora_weight_pointer - }) - - lora_rank = Tensor( - name=f'{lora_module}_lora_ranks_{i}', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([('batch_size_beam_width', - bb_range)]), - ) - lora_rank_dict.update( - {f"{lora_module}_lora_ranks": lora_rank}) - - lora_weights_pointers.append(lora_weight_pointer_dict) - lora_ranks.append(lora_rank_dict) - - last_token_ids = None - if mapping.is_last_pp_rank() and not gather_context_logits: - if not remove_input_padding and max_draft_len > 0: - last_token_ids = Tensor( - name='last_token_ids', - dtype=trt.int32, - shape=[-1, -1], - dim_range=OrderedDict([ - ('batch_size_beam_width', bb_range), - ('last_token_ids', tokens_per_engine_step_range), - ]), - ) - else: - last_token_ids = Tensor( - name='last_token_ids', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([ - ('batch_size_last_token_ids', bbd_range), - ]), - ) - - spec_decoding_params = None - # Use positional offsets and packed mask only when not in SpS spec decoding - if speculative_decoding_draft_tokens_external == False and max_draft_len > 0: - tokens_per_engine_step = max_draft_len + 1 - # 32 bits packed mask aligned. - num_packed_masks = (tokens_per_engine_step + 32 - 1) // 32 - packed_mask_len_range = [[0, 1, num_packed_masks]] * num_profiles - # total number of spec decoding tokens for all sequences (sequence length can be variable). - num_gen_tokens_range = [ - GenerationMixin.default_range( - max_batch_size * max_beam_width * tokens_per_engine_step, - min_range=0) - ] * num_profiles - bb_range_0 = [[0] + bbr[1:] for bbr in bb_range] - - # support variable sequence lengths for medusa. - spec_decoding_generation_lengths = Tensor( - name='spec_decoding_generation_lengths', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([('batch_size_beam_width_0', bb_range_0) - ]), - ) - - # position offsets that are fixed during the whole session. - # it will be shared among all sequences. - spec_decoding_position_offsets = Tensor( - name='spec_decoding_position_offsets', - dtype=trt.int32, - shape=[-1, -1], - dim_range=OrderedDict([ - ('batch_size_beam_width_0', bb_range_0), - ('spec_decoding_position_ids_dim0', - tokens_per_engine_step_range), - ]), - ) - - spec_decoding_packed_mask = Tensor( - name='spec_decoding_packed_mask', - dtype=trt.int32, - shape=[-1, -1], - dim_range=OrderedDict([ - ('spec_decoding_packed_mask_dim0', num_gen_tokens_range), - ('spec_decoding_packed_mask_dim1', packed_mask_len_range), - ]), - ) - spec_decoding_use = Tensor(name='spec_decoding_use', - dtype=trt.int32, - shape=[1], - dim_range=OrderedDict([ - ('spec_decoding_use_dim', - [1] * num_profiles), - ])) - spec_decoding_params = SpecDecodingParams( - spec_decoding_is_generation_length_variable= - spec_decoding_is_generation_length_variable, - spec_decoding_max_generation_length=tokens_per_engine_step, - spec_decoding_generation_lengths= - spec_decoding_generation_lengths, - spec_decoding_position_offsets=spec_decoding_position_offsets, - spec_decoding_packed_mask=spec_decoding_packed_mask, - spec_decoding_use=spec_decoding_use) - - mrope_params = None - if mrope_rotary_cos_sin_size is not None: - mrope_rotary_cos_sin = Tensor( - name='mrope_rotary_cos_sin', - dtype=trt.float32, - shape=[-1, mrope_rotary_cos_sin_size], - dim_range=OrderedDict([ - ('batch_size_beam_width', bb_range), - ('mult_dim', [mrope_rotary_cos_sin_size] * num_profiles), - ]), - ) - - mrope_position_deltas = Tensor( - name='mrope_position_deltas', - dtype=trt.int32, - shape=[-1, 1], - dim_range=OrderedDict([('batch_size_beam_width', bb_range), - ('mult_dim_delta', [1] * num_profiles)]), - ) - mrope_params = MropeParams( - mrope_rotary_cos_sin=mrope_rotary_cos_sin, - mrope_position_deltas=mrope_position_deltas, - ) - - basic_inputs = { - 'input_ids': input_ids, - 'hidden_states_input': hidden_states, - 'position_ids': position_ids, - 'last_token_ids': last_token_ids, - 'prompt_embedding_table': prompt_embedding_table, - 'tasks': tasks, - 'prompt_vocab_size': prompt_vocab_size, - 'lora_ranks': lora_ranks, - 'lora_weights_pointers': lora_weights_pointers, - 'spec_decoding_params': spec_decoding_params, - 'mrope_params': mrope_params, - } - - attention_inputs = self.prepare_attention_inputs( - max_batch_size=max_batch_size, - max_beam_width=max_beam_width, - max_input_len=max_input_len, - max_seq_len=max_seq_len, - num_kv_heads=num_kv_heads, - head_size=head_size, - num_layers=num_layers, - kv_dtype=kv_dtype, - num_profiles=num_profiles, - enable_ctx_gen_opt_profiles=enable_ctx_gen_opt_profiles, - remove_input_padding=remove_input_padding, - use_gpt_attention_plugin=use_gpt_attention_plugin, - kv_cache_type=kv_cache_type, - tokens_per_block=tokens_per_block, - mapping=mapping, - streamingllm=streamingllm, - opt_batch_size=opt_batch_size) - for key, value in attention_inputs.items(): - basic_inputs[key] = value - - return basic_inputs diff --git a/tensorrt_llm/models/gpt/__init__.py b/tensorrt_llm/models/gpt/__init__.py deleted file mode 100644 index 71bf6d298d42..000000000000 --- a/tensorrt_llm/models/gpt/__init__.py +++ /dev/null @@ -1,14 +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. diff --git a/tensorrt_llm/models/gpt/config.py b/tensorrt_llm/models/gpt/config.py deleted file mode 100644 index ba09d1f86942..000000000000 --- a/tensorrt_llm/models/gpt/config.py +++ /dev/null @@ -1,324 +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. - -from typing import Optional, Union - -import torch - -from ..._utils import get_hf_rope_theta -from ...layers import MoeConfig -from ...logger import logger -from ...mapping import Mapping -from ..convert_utils import infer_dtype -from ..modeling_utils import PretrainedConfig, QuantConfig - - -class GPTConfig(PretrainedConfig): - - def __init__(self, - *, - gpt_variant: str = 'gpt2', - bias: bool = True, - q_scaling: float = 1.0, - embedding_scale: Optional[float] = None, - apply_query_key_layer_scaling: bool = False, - rotary_pct: float = 1.0, - rotary_base: float = 10000.0, - rotary_scaling: Optional[dict] = None, - inner_layernorm: bool = False, - norm_before_bmm1: bool = False, - moe: Optional[Union[MoeConfig, dict]] = None, - **kwargs): - self.gpt_variant = gpt_variant - self.bias = bias - self.q_scaling = q_scaling - self.embedding_scale = embedding_scale - self.apply_query_key_layer_scaling = apply_query_key_layer_scaling - self.rotary_pct = rotary_pct - self.rotary_base = rotary_base - self.rotary_scaling = rotary_scaling - self.inner_layernorm = inner_layernorm - self.norm_before_bmm1 = norm_before_bmm1 - if moe is None: - # Legacy MOE config fields - moe = MoeConfig( - num_experts=kwargs.pop('moe_num_experts', 0), - top_k=kwargs.pop('moe_top_k', 0), - normalization_mode=kwargs.pop( - 'moe_normalization_mode', - MoeConfig.ExpertScaleNormalizationMode.RENORMALIZE)) - elif isinstance(moe, dict): - moe = MoeConfig.from_dict(moe) - assert isinstance(moe, MoeConfig) - self.moe = moe.validate() - - super().__init__(**kwargs) - - def to_dict(self): - output = super().to_dict() - # Serialize the fields added in GPTConfig - output['gpt_variant'] = self.gpt_variant - output['bias'] = self.bias - output['q_scaling'] = self.q_scaling - output['embedding_scale'] = self.embedding_scale - output[ - 'apply_query_key_layer_scaling'] = self.apply_query_key_layer_scaling - output['rotary_pct'] = self.rotary_pct - output['rotary_base'] = self.rotary_base - output['rotary_scaling'] = self.rotary_scaling - output['inner_layernorm'] = self.inner_layernorm - output['norm_before_bmm1'] = self.norm_before_bmm1 - output['moe'] = self.moe.to_dict() - return output - - @classmethod - def from_hugging_face( - cls, - hf_config_or_dir: Union[str, 'transformers.PretrainedConfig'], - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - import transformers - trust_remote_code = kwargs.pop('trust_remote_code', True) - - from .convert import get_needed_padding - - if isinstance(hf_config_or_dir, transformers.PretrainedConfig): - hf_config = hf_config_or_dir - else: - hf_config = transformers.AutoConfig.from_pretrained( - hf_config_or_dir, trust_remote_code=trust_remote_code) - - gpt_variant = kwargs.pop('gpt_variant', None) - if gpt_variant is None: - logger.info("Inferring gpt variant from path...") - for v in [ - 'starcoder2', 'starcoder', 'santacoder', 'gpt2', - 'persimmon', 'fuyu', 'kosmos-2', 'jais', 'nemotron' - ]: - if v in hf_config._name_or_path or v == hf_config.model_type: - gpt_variant = v - break - if gpt_variant == 'fuyu': - gpt_variant = 'persimmon' - - assert gpt_variant in [ - 'gpt2', 'santacoder', 'starcoder', 'starcoder2', 'persimmon', - 'kosmos-2', 'jais', 'nemotron' - ] - logger.info(f"Gpt variant: {gpt_variant}") - - if gpt_variant in ['starcoder2', 'nemotron', 'persimmon']: - hf_config.n_embd = hf_config.hidden_size - hf_config.n_inner = hf_config.intermediate_size - hf_config.n_head = hf_config.num_attention_heads - hf_config.n_kv_head = hf_config.num_key_value_heads if hasattr( - hf_config, 'num_key_value_heads') else hf_config.n_head - hf_config.n_layer = hf_config.num_hidden_layers - hf_config.n_positions = hf_config.max_position_embeddings - hf_config.activation_function = 'gelu' if gpt_variant == 'starcoder2' else 'squared-relu' - if gpt_variant == "nemotron": - hf_config.layer_norm_eps = hf_config.norm_eps - hf_config.layer_norm_epsilon = hf_config.norm_epsilon if gpt_variant == 'starcoder2' else hf_config.layer_norm_eps - hf_config.bias = hf_config.use_bias if gpt_variant == 'starcoder2' else gpt_variant != 'nemotron' - hf_config.position_embedding_type = 'rope_gpt_neox' - hf_config.rotary_base = get_hf_rope_theta(hf_config, 10000.0) - hf_config.rotary_pct = getattr( - hf_config, 'partial_rotary_factor', - getattr(hf_config, 'rope_percent', 1.0)) - try: - # only for persimmon, not starcoder2 - hf_config.vocab_size = hf_config.text_config.vocab_size - except AttributeError: - pass - elif gpt_variant == "kosmos-2": - hf_config.n_embd = hf_config.text_config.embed_dim - hf_config.n_inner = hf_config.text_config.ffn_dim - hf_config.n_head = hf_config.text_config.attention_heads - hf_config.n_kv_head = hf_config.n_head - hf_config.n_layer = hf_config.text_config.layers - hf_config.n_positions = hf_config.text_config.max_position_embeddings - hf_config.activation_function = hf_config.text_config.activation_function - hf_config.layer_norm_epsilon = hf_config.text_config.layer_norm_eps - hf_config.bias = True - hf_config.vocab_size = hf_config.text_config.vocab_size - else: - if hf_config.n_inner is None: - hf_config.n_inner = hf_config.n_embd * 4 - if gpt_variant in ['santacoder', 'starcoder']: - hf_config.n_kv_head = 1 - else: - hf_config.n_kv_head = hf_config.n_head - - if gpt_variant == 'jais': - hf_config.q_scaling = (hf_config.n_embd // hf_config.n_head)**0.5 - if hasattr(hf_config, 'width_scale'): - hf_config.logits_scale = hf_config.width_scale - else: - hf_config.logits_scale = hf_config.mup_output_alpha * hf_config.mup_width_scale - - if hasattr(hf_config, 'mup_embeddings_scale'): - hf_config.embeddings_scale = hf_config.mup_embeddings_scale - else: - assert hasattr(hf_config, 'embeddings_scale') - - hf_config.n_inner += get_needed_padding(hf_config.n_inner, - mapping.tp_size) - - if gpt_variant == 'kosmos-2': - if hf_config.text_config.scale_embedding: - hf_config.embeddings_scale = hf_config.n_embd**0.5 - - dtype = infer_dtype(dtype, getattr(hf_config, 'torch_dtype', None)) - - return cls(architecture=hf_config.architectures[0], - dtype=dtype, - num_hidden_layers=hf_config.n_layer, - num_attention_heads=hf_config.n_head, - num_key_value_heads=hf_config.n_kv_head, - hidden_size=hf_config.n_embd, - intermediate_size=hf_config.n_inner, - norm_epsilon=hf_config.layer_norm_epsilon, - vocab_size=hf_config.vocab_size, - position_embedding_type=getattr(hf_config, - 'position_embedding_type', - 'learned_absolute'), - max_position_embeddings=hf_config.n_positions, - hidden_act=hf_config.activation_function, - gpt_variant=gpt_variant, - bias=getattr(hf_config, 'bias', True), - apply_query_key_layer_scaling=getattr( - hf_config, 'apply_query_key_layer_scaling', False), - rotary_pct=getattr(hf_config, 'rotary_pct', 1.0), - rotary_base=getattr(hf_config, 'rotary_base', 10000.0), - rotary_scaling=getattr(hf_config, 'rotary_scaling', None), - qk_layernorm=gpt_variant == 'persimmon', - inner_layernorm=gpt_variant == 'kosmos-2', - norm_before_bmm1=gpt_variant == 'kosmos-2', - q_scaling=getattr(hf_config, 'q_scaling', 1), - embedding_scale=getattr(hf_config, 'embeddings_scale', None), - mapping=mapping, - quantization=quant_config, - **kwargs) - - @classmethod - def from_nemo(cls, - nemo_ckpt_dir: str, - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - import transformers - - from .convert import (UnpackedNemoCheckpointDir, cpu_map_location, - gpu_map_location, rename_keys) - - load_model_on_cpu = kwargs.pop('load_model_on_cpu', False) - nemo_rename_key = kwargs.pop('nemo_rename_key', []) - layer_rename_config = { - pattern.split(':')[0]: pattern.split(':')[1] - for pattern in nemo_rename_key - } - - unpacked_checkpoints_dir = UnpackedNemoCheckpointDir( - nemo_ckpt_dir, load_checkpoints_to_cpu=load_model_on_cpu) - nemo_model_config = unpacked_checkpoints_dir.model_config - - training_tp_size = nemo_model_config.get("tensor_model_parallel_size", - 1) - training_pp_size = nemo_model_config.get("pipeline_model_parallel_size", - 1) - - checkpoints_paths = unpacked_checkpoints_dir.get_checkpoints_paths( - training_tp_size, - training_pp_size, - ) - if unpacked_checkpoints_dir._load_checkpoints_to_cpu: - map_location_fn = cpu_map_location - else: - map_location_fn = gpu_map_location - model_00 = torch.load(checkpoints_paths[0][0], - map_location=map_location_fn) - model_00 = rename_keys(model_00, layer_rename_config) - vocab_size = model_00[ - "model.language_model.embedding.word_embeddings.weight"].shape[ - 0] * training_tp_size - del model_00 - - hf_config = transformers.GPT2Config( - vocab_size=vocab_size, - n_positions=nemo_model_config['max_position_embeddings'], - n_embd=nemo_model_config['hidden_size'], - n_layer=nemo_model_config['num_layers'], - n_head=nemo_model_config['num_attention_heads'], - n_inner=nemo_model_config['ffn_hidden_size'], - activation_function=nemo_model_config['activation'], - layer_norm_epsilon=nemo_model_config['layernorm_epsilon'], - ) - hf_config.n_kv_head = hf_config.n_head - hf_config.bias = nemo_model_config['bias'] - hf_config.apply_query_key_layer_scaling = False - - hf_config.position_embedding_type = nemo_model_config.get( - 'position_embedding_type', 'learned_absolute') - if hf_config.position_embedding_type == 'rope': - hf_config.position_embedding_type = 'rope_gpt_neox' - hf_config.rotary_base = nemo_model_config.get('rotary_base', 10000.0) - hf_config.rotary_pct = nemo_model_config.get('rotary_percentage', 1.0) - assert hf_config.rotary_pct >= 0 and hf_config.rotary_pct <= 1 - - rotary_scaling_factor = nemo_model_config.get( - 'seq_len_interpolation_factor', None) - if rotary_scaling_factor is None: - hf_config.rotary_scaling = None - else: - assert rotary_scaling_factor > 1 - hf_config.rotary_scaling = { - 'type': 'linear', - 'factor': rotary_scaling_factor - } - - if dtype == 'auto': - dtype = nemo_model_config.get('precision', None) - if dtype is None: - dtype = 'float16' - elif 'bf16' in dtype or 'bfloat16' in dtype: - dtype = 'bfloat16' - else: - dtype = 'float16' - logger.info(f"Specified dtype 'auto'; inferred dtype {dtype!r}.") - - return cls(architecture='GPTForCausalLM', - dtype=dtype, - num_hidden_layers=hf_config.n_layer, - num_attention_heads=hf_config.n_head, - num_key_value_heads=hf_config.n_kv_head, - hidden_size=hf_config.n_embd, - intermediate_size=hf_config.n_inner, - norm_epsilon=hf_config.layer_norm_epsilon, - vocab_size=hf_config.vocab_size, - position_embedding_type=hf_config.position_embedding_type, - max_position_embeddings=hf_config.n_positions, - hidden_act=hf_config.activation_function, - bias=hf_config.bias, - apply_query_key_layer_scaling=hf_config. - apply_query_key_layer_scaling, - rotary_pct=hf_config.rotary_pct, - rotary_base=hf_config.rotary_base, - rotary_scaling=hf_config.rotary_scaling, - mapping=mapping, - quantization=quant_config, - **kwargs) diff --git a/tensorrt_llm/models/gpt/convert.py b/tensorrt_llm/models/gpt/convert.py deleted file mode 100644 index 799fac22d40f..000000000000 --- a/tensorrt_llm/models/gpt/convert.py +++ /dev/null @@ -1,1455 +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. - -import copy -import functools -import os -import shutil -import tarfile -import time -from collections import defaultdict -from pathlib import Path -from typing import Dict, Optional, Union - -import numpy as np -import safetensors -import torch -import torch.nn as nn -import yaml -from tqdm import tqdm -from transformers import AutoModelForCausalLM, AutoTokenizer -from transformers.models.gpt2.modeling_gpt2 import GPT2Block -from transformers.pytorch_utils import Conv1D - -from ..._utils import pad_vocab_size, str_dtype_to_torch -from ...logger import logger -from ...quantization import QuantAlgo -from ..convert_utils import (generate_int8, get_weight, get_weight_and_bias, - load_calib_dataset, - retrieved_layer_index_from_name, smooth_gemm) -from .config import GPTConfig - - -def rename_keys(model_state, layer_rename_config: Dict[str, str]): - if not layer_rename_config: - return model_state - - new_state_dict = {} - for key, value in model_state.items(): - for old, new in layer_rename_config.items(): - key = key.replace(old, new) - assert key not in new_state_dict, f"Key already exists: {key}" - new_state_dict[key] = value - - return new_state_dict - - -def get_needed_padding(value: int, multiple: int) -> int: - return (multiple - value % multiple) % multiple - - -def pad_array_up_to(v: torch.Tensor, axis: int, multiple: int) -> torch.Tensor: - a = [0 for i in range(len(v.shape) * 2)] - a[axis * 2 - 1] = get_needed_padding(v.shape[axis], multiple) - return torch.nn.functional.pad(v, a) - - -def split(param: torch.Tensor, - tp_rank: int, - tp_size: int, - is_column: bool = True) -> torch.Tensor: - """Split linear layer's weight, bias or scaling factors for tensor parallelism.""" - if param is None: - return None - assert param.ndim in [1, 2] - if tp_size == 1: - return param - if param.numel() == 1: - return param - if param.ndim == 1 and not is_column: - return param - split_dim = 0 if (param.ndim == 1 or is_column) else 1 - return torch.chunk(param, tp_size, dim=split_dim)[tp_rank].contiguous() - - -def split_qkv( - param: torch.Tensor, - tp_rank: int, - tp_size: int, - hidden_size: int, - num_heads: int, - num_kv_heads: Optional[int] = None, -) -> torch.Tensor: - """Split qkv layer's weight, bias or scaling factors for tensor parallelism. - - param: (num_heads*head_dim + 2*num_kv_heads*head_dim, in_dim) - """ - if param is None: - return None - assert hidden_size % num_heads == 0 - head_dim = hidden_size // num_heads - num_kv_heads = num_kv_heads if num_kv_heads is not None else num_heads - assert num_heads % num_kv_heads == 0 - assert num_heads % tp_size == 0 - - q_param, k_param, v_param = torch.split( - param, [hidden_size, num_kv_heads * head_dim, num_kv_heads * head_dim], - dim=0) - - if num_kv_heads < tp_size: - assert tp_size % num_kv_heads == 0 - num_dups = tp_size // num_kv_heads - remain_shape = k_param.shape[1:] - k_param = k_param.view( - num_kv_heads, head_dim, - *remain_shape).repeat_interleave(num_dups, dim=0).view( - num_kv_heads * head_dim * num_dups, *remain_shape) - v_param = v_param.view( - num_kv_heads, head_dim, - *remain_shape).repeat_interleave(num_dups, dim=0).view( - num_kv_heads * head_dim * num_dups, *remain_shape) - else: - assert num_kv_heads % tp_size == 0 - - q_param = split(q_param, tp_rank, tp_size, is_column=True) - k_param = split(k_param, tp_rank, tp_size, is_column=True) - v_param = split(v_param, tp_rank, tp_size, is_column=True) - return torch.cat([q_param, k_param, v_param], dim=0) - - -def split_embedding( - param: torch.Tensor, - tp_rank: int, - tp_size: int, - use_parallel_embedding: bool = False, - sharding_dim: int = 0, -) -> torch.Tensor: - if param is None: - return None - if not use_parallel_embedding: - return param - - vocab_size, hidden_size = param.size() - if sharding_dim == 0: - if vocab_size % tp_size != 0: - vocab_size_padded = pad_vocab_size(vocab_size, tp_size) - pad_width = vocab_size_padded - vocab_size - param = torch.nn.functional.pad(param, (0, 0, 0, pad_width), - value=0) - else: - assert hidden_size % tp_size == 0 - return split(param, tp_rank, tp_size, is_column=(sharding_dim == 0)) - - -def get_tllm_linear_weight( - weight: torch.Tensor, - prefix: str, - bias: Optional[torch.Tensor] = None, - use_weight_only: bool = False, - plugin_weight_only_quant_type: torch.dtype = torch.int8 -) -> Dict[str, torch.Tensor]: - results = {} - if use_weight_only: - v = weight.t().contiguous() - processed_torch_weights, torch_weight_scales = \ - torch.ops.trtllm.symmetric_quantize_last_axis_of_batched_matrix( - v, plugin_weight_only_quant_type) - results[f'{prefix}.weight'] = processed_torch_weights - results[f'{prefix}.per_channel_scale'] = torch_weight_scales - else: - results[f'{prefix}.weight'] = weight - - if bias is not None: - results[f'{prefix}.bias'] = bias - - return results - - -@torch.no_grad() -def capture_activation_range(model, - tokenizer, - dataset, - num_samples=512, - seq_len=512): - model.eval() - device = next(model.parameters()).device - act_scales = defaultdict(lambda: {"x": None, "y": None, "w": None}) - - def stat_tensor(name, tensor, act_scales, key): - hidden_dim = tensor.shape[-1] - tensor = tensor.view(-1, hidden_dim).abs().detach() - comming_max = torch.max(tensor, dim=0)[0].float() - - if act_scales[name][key] is None: - act_scales[name][key] = comming_max - else: - act_scales[name][key] = torch.max(act_scales[name][key], - comming_max) - - def stat_input_hook(m, x, y, name): - if isinstance(x, tuple): - x = x[0] - stat_tensor(name, x, act_scales, "x") - stat_tensor(name, y, act_scales, "y") - - if act_scales[name]["w"] is None: - act_scales[name]["w"] = m.weight.abs().clip(1e-8, - None).max(dim=0)[0] - - hooks = [] - for name, m in model.named_modules(): - if isinstance(m, nn.Linear) or isinstance(m, Conv1D): - hooks.append( - m.register_forward_hook( - functools.partial(stat_input_hook, name=name))) - - for i in tqdm(range(num_samples), desc="calibrating model"): - input_ids = tokenizer(dataset[i], - return_tensors="pt", - max_length=seq_len, - truncation=True).input_ids.to(device) - model(input_ids) - - for h in hooks: - h.remove() - - return act_scales - - -@torch.no_grad() -def smooth_gpt_model(model, scales, alpha): - # Smooth the activation and weights with smoother = $\diag{s}$ - for name, module in model.named_modules(): - if not isinstance(module, GPT2Block): - continue - - # qkv_proj - layer_name = name + ".attn.c_attn" - smoother = smooth_gemm(module.attn.c_attn.weight.T, - scales[layer_name]["x"], module.ln_1.weight, - module.ln_1.bias, alpha) - scales[layer_name]["x"] = scales[layer_name]["x"] / smoother - scales[layer_name]["w"] = module.attn.c_attn.weight.abs().max(dim=0)[0] - - # fc1 - layer_name = name + ".mlp.c_fc" - smoother = smooth_gemm(module.mlp.c_fc.weight.T, - scales[layer_name]["x"], module.ln_2.weight, - module.ln_2.bias, alpha) - scales[layer_name]["x"] = scales[layer_name]["x"] / smoother - scales[layer_name]["w"] = module.mlp.c_fc.weight.abs().max(dim=0)[0] - - -def get_tllm_linear_sq_weight(vals, - prefix, - shape, - tensor_parallel, - is_qkv=False, - per_token=False, - per_channel=False, - last_prefix=None, - bias=None, - smoother_value=None, - smoother_shape=None, - rank=0, - cat_dim=0, - multi_query_mode=False): - results = {} - - def multi_query_split(data, local_dim, head_size, tp_size, cur_rank): - q, k, v = torch.split(data, [local_dim, head_size, head_size], dim=-1) - q_split = torch.split(q, q.shape[-1] // tp_size, dim=-1) - k_split = torch.split(k, q.shape[-1] // tp_size, dim=-1) - v_split = torch.split(v, q.shape[-1] // tp_size, dim=-1) - return [ - torch.concat((q_split[ii], k_split[ii], v_split[ii]), dim=-1) - for ii in range(tp_size) - ][cur_rank] - - col_shape = shape if (is_qkv or per_channel) else [1, 1] - - if per_token: - if per_channel: - original_weights = vals["weight.int8.col"] - else: - original_weights = vals["weight.int8"] - local_dim = original_weights.shape[0] - head_size = (original_weights.shape[1] - local_dim) // 2 - - if multi_query_mode: - cur_weights = multi_query_split(original_weights, local_dim, - head_size, tensor_parallel, rank) - else: - cur_weights = np.split(original_weights, - tensor_parallel, - axis=cat_dim)[rank] - if is_qkv: - hidden_dim = cur_weights.shape[0] - cur_weights = cur_weights.reshape(hidden_dim, -1) - results[prefix + 'weight'] = cur_weights.t().contiguous() - if smoother_value is None: - results[last_prefix] = torch.from_numpy( - np.array([1.0], dtype=np.float32)) - - if per_channel: - cur_per_channel_value = vals["scale_w_quant_orig.col"] - if smoother_value is None: - if multi_query_mode: - cur_per_channel_value = multi_query_split( - vals["scale_w_quant_orig.col"], local_dim, head_size, - tensor_parallel, rank) - else: - cur_per_channel_value = np.split( - vals["scale_w_quant_orig.col"], - tensor_parallel, - axis=cat_dim)[rank] - else: - cur_per_channel_value = vals["scale_w_quant_orig"] - if is_qkv: - if multi_query_mode: - cur_per_channel_value = multi_query_split( - vals["scale_w_quant_orig"], local_dim, head_size, - tensor_parallel, rank) - else: - cur_per_channel_value = np.split(vals["scale_w_quant_orig"], - tensor_parallel, - axis=cat_dim)[rank] - - results[prefix + 'per_channel_scale'] = cur_per_channel_value.reshape( - col_shape).contiguous() - else: - if per_channel: - original_weights = np.array(vals["weight.int8.col"]) - else: - original_weights = np.array(vals["weight.int8"]) - local_dim = original_weights.shape[0] - head_size = (original_weights.shape[1] - local_dim) // 2 - - if multi_query_mode: - cur_weights = multi_query_split(original_weights, local_dim, - head_size, tensor_parallel, rank) - else: - cur_weights = np.split(original_weights, - tensor_parallel, - axis=cat_dim)[rank] - if is_qkv: - hidden_dim = cur_weights.shape[0] - cur_weights = cur_weights.reshape(hidden_dim, -1) - results[prefix + - 'weight'] = torch.from_numpy(cur_weights).t().contiguous() - - if per_channel: - cur_per_channel_value = vals["scale_y_accum_quant.col"] - if smoother_value is None: - if multi_query_mode: - cur_per_channel_value = multi_query_split( - vals["scale_y_accum_quant.col"], local_dim, head_size, - tensor_parallel, rank) - else: - cur_per_channel_value = np.split( - vals["scale_y_accum_quant.col"], - tensor_parallel, - axis=cat_dim)[rank] - else: - cur_per_channel_value = vals["scale_y_accum_quant"] - # QKV is always per_channel - if is_qkv: - if multi_query_mode: - cur_per_channel_value = multi_query_split( - vals["scale_y_accum_quant"], local_dim, head_size, - tensor_parallel, rank) - else: - cur_per_channel_value = np.split( - vals["scale_y_accum_quant"], - tensor_parallel, - axis=cat_dim)[rank] - - results[prefix + 'per_channel_scale'] = cur_per_channel_value.reshape( - col_shape).contiguous() - - results[last_prefix] = vals['scale_x_orig_quant'].contiguous() - - results[prefix + 'act_scale'] = vals["scale_y_quant_orig"].contiguous() - - if smoother_value is not None: - cur_smoother_value = np.split(smoother_value, - tensor_parallel, - axis=cat_dim)[rank] - results[prefix + 'smoother'] = cur_smoother_value.reshape( - smoother_shape).contiguous().to(torch.float32) - - if bias is not None: - results[prefix + 'bias'] = bias - - return results - - -def load_weights_from_hf_model(hf_model, - config: GPTConfig, - act_range: Optional[dict] = None): - quant_algo = config.quantization.quant_algo - use_weight_only = quant_algo in [QuantAlgo.W8A16, QuantAlgo.W4A16] - if quant_algo == QuantAlgo.W8A16: - plugin_weight_only_quant_type = torch.int8 - elif quant_algo == QuantAlgo.W4A16: - plugin_weight_only_quant_type = torch.quint4x2 - else: - plugin_weight_only_quant_type = None - - use_smooth_quant = config.quantization._use_plugin_sq - per_channel = use_smooth_quant and 'PER_CHANNEL' in quant_algo - per_token = use_smooth_quant and 'PER_TOKEN' in quant_algo - int8_kv_cache = config.quantization.kv_cache_quant_algo == QuantAlgo.INT8 - if use_smooth_quant or int8_kv_cache: - assert act_range is not None - - weights = {} - tik = time.time() - - hf_config = hf_model.config - model_params = dict(hf_model.named_parameters()) - dtype = getattr(torch, config.dtype) - gpt_variant = config.gpt_variant - num_attention_heads = config.num_attention_heads - hidden_size = config.hidden_size - vocab_size = config.vocab_size - num_kv_heads = config.num_key_value_heads - num_hidden_layers = config.num_hidden_layers - multi_query_mode = (num_kv_heads != num_attention_heads) - - mapping = config.mapping - layers_range = mapping.pp_layers(num_hidden_layers) - for l in layers_range: - if gpt_variant in ['starcoder2', 'nemotron']: - prefix = f'model.layers.{l}' - elif gpt_variant == 'persimmon': - is_fuyu = f'language_model.model.embed_tokens.weight' in model_params - prefix = f'language_model.model.layers.{l}' if is_fuyu else f'model.layers.{l}' - elif gpt_variant == 'kosmos-2': - prefix = f'text_model.model.layers.{l}' - else: - prefix = f'transformer.h.{l}' - tllm_prex = f'transformer.layers.{l-layers_range[0]}' - - # (1) Attention QKV Linear - if gpt_variant == 'santacoder': - q_w, q_b = get_weight_and_bias(model_params, - f'{prefix}.attn.q_attn', dtype) - kv_w, kv_b = get_weight_and_bias(model_params, - f'{prefix}.attn.kv_attn', dtype) - qkv_w = torch.cat([q_w, kv_w], dim=-1) - qkv_b = torch.cat([q_b, kv_b], dim=-1) - elif gpt_variant in ['starcoder2', 'nemotron', 'kosmos-2']: - q_w, q_b = get_weight_and_bias(model_params, - f'{prefix}.self_attn.q_proj', dtype) - k_w, k_b = get_weight_and_bias(model_params, - f'{prefix}.self_attn.k_proj', dtype) - v_w, v_b = get_weight_and_bias(model_params, - f'{prefix}.self_attn.v_proj', dtype) - qkv_w = torch.cat([q_w.cuda(), k_w.cuda(), v_w.cuda()], dim=0) - qkv_b = torch.cat([q_b.cuda(), k_b.cuda(), - v_b.cuda()], dim=0) if q_b is not None else None - elif gpt_variant == 'persimmon': - qkv_w, qkv_b = get_weight_and_bias( - model_params, f'{prefix}.self_attn.query_key_value', dtype) - else: - qkv_w, qkv_b = get_weight_and_bias(model_params, - f'{prefix}.attn.c_attn', dtype) - if gpt_variant in ['gpt2', 'santacoder', 'jais']: - qkv_w = qkv_w.t().contiguous() # transpose for Conv1D - - if use_smooth_quant: - qkv_out_dim = qkv_w.shape[0] - qkv_w_t = qkv_w.t() - if not multi_query_mode: - qkv_w_t = qkv_w_t.reshape(hidden_size, 3, hidden_size) - int8_weights = generate_int8(qkv_w_t, - act_range.get(f'{prefix}.attn.c_attn'), - is_qkv=True, - multi_query_mode=multi_query_mode) - qkv_b = split_qkv(qkv_b, mapping.tp_rank, mapping.tp_size, - hidden_size, num_attention_heads, num_kv_heads) - weights.update( - get_tllm_linear_sq_weight( - int8_weights, - f'{tllm_prex}.attention.qkv.', - [1, qkv_out_dim // mapping.tp_size], - mapping.tp_size, - is_qkv=True, - per_token=per_token, - per_channel=per_channel, - last_prefix=f'{tllm_prex}.input_layernorm.scale_to_int', - bias=qkv_b, - smoother_value=None, - smoother_shape=None, - rank=mapping.tp_rank, - cat_dim=-1, - multi_query_mode=multi_query_mode)) - else: - if gpt_variant == 'persimmon': - qkv_w = split(qkv_w, - mapping.tp_rank, - mapping.tp_size, - is_column=True) - - qkv_b = split(qkv_b, - mapping.tp_rank, - mapping.tp_size, - is_column=True) - else: - qkv_w = split_qkv(qkv_w, mapping.tp_rank, mapping.tp_size, - hidden_size, num_attention_heads, - num_kv_heads) - qkv_b = split_qkv(qkv_b, mapping.tp_rank, mapping.tp_size, - hidden_size, num_attention_heads, - num_kv_heads) - weights.update( - get_tllm_linear_weight(qkv_w, f'{tllm_prex}.attention.qkv', - qkv_b, use_weight_only, - plugin_weight_only_quant_type)) - - if int8_kv_cache: - qkv_w_t = qkv_w.t() - if not multi_query_mode: - qkv_w_t = qkv_w_t.reshape(hidden_size, 3, hidden_size) - int8_weights = generate_int8(qkv_w_t, - act_range.get(f'{prefix}.attn.c_attn'), - is_qkv=True, - multi_query_mode=multi_query_mode) - weights[ - f'{tllm_prex}.attention.kv_cache_scaling_factor'] = int8_weights[ - 'scale_y_quant_orig'].contiguous() - - # (2) Attention Dense Linear - if gpt_variant in ['starcoder2', 'nemotron']: - attn_dense_w, attn_dense_b = get_weight_and_bias( - model_params, f'{prefix}.self_attn.o_proj', dtype) - elif gpt_variant == 'persimmon': - attn_dense_w, attn_dense_b = get_weight_and_bias( - model_params, f'{prefix}.self_attn.dense', dtype) - elif gpt_variant == 'kosmos-2': - attn_dense_w, attn_dense_b = get_weight_and_bias( - model_params, f'{prefix}.self_attn.out_proj', dtype) - else: - attn_dense_w, attn_dense_b = get_weight_and_bias( - model_params, f'{prefix}.attn.c_proj', dtype) - if gpt_variant in ['gpt2', 'santacoder', 'jais']: - attn_dense_w = attn_dense_w.t().contiguous() # transpose for Conv1D - - if use_smooth_quant: - attn_dense_w_t = attn_dense_w.t() - int8_weights = generate_int8(attn_dense_w_t, - act_range.get(f'{prefix}.attn.c_proj')) - # change it to the real smoother if dense layer is applied smooth quant - fake_smoother_value = torch.ones([1, hidden_size], - dtype=torch.float32) - weights.update( - get_tllm_linear_sq_weight( - int8_weights, - f'{tllm_prex}.attention.dense.', [1, hidden_size], - mapping.tp_size, - is_qkv=False, - per_token=per_token, - per_channel=per_channel, - last_prefix= - f'{tllm_prex}.attention.quantization_scaling_factor', - bias=attn_dense_b, - smoother_value=fake_smoother_value, - smoother_shape=[1, hidden_size // mapping.tp_size], - rank=mapping.tp_rank, - cat_dim=0)) - else: - attn_dense_w = split(attn_dense_w, - mapping.tp_rank, - mapping.tp_size, - is_column=False) - weights.update( - get_tllm_linear_weight(attn_dense_w, - f'{tllm_prex}.attention.dense', - attn_dense_b, use_weight_only, - plugin_weight_only_quant_type)) - - # (3) MLP FC Linear - if gpt_variant == 'persimmon': - suffix = "mlp.dense_h_to_4h" - elif gpt_variant == 'kosmos-2': - suffix = "ffn.fc1" - elif gpt_variant == 'nemotron': - suffix = "mlp.up_proj" - else: - suffix = "mlp.c_fc" - mlp_fc_w, mlp_fc_b = get_weight_and_bias(model_params, - f'{prefix}.{suffix}', dtype) - if gpt_variant in ['gpt2', 'santacoder', 'jais']: - mlp_fc_w = mlp_fc_w.t().contiguous() # transpose for Conv1D - if gpt_variant in ['jais']: - mlp_fc_w = pad_array_up_to(mlp_fc_w, 0, mapping.tp_size) - mlp_fc_b = pad_array_up_to(mlp_fc_b, 0, mapping.tp_size) - if use_smooth_quant: - mlp_fc_w_t = mlp_fc_w.t() - int8_weights = generate_int8(mlp_fc_w_t, - act_range.get(f'{prefix}.mlp.c_fc')) - mlp_fc_b = split(mlp_fc_b, - mapping.tp_rank, - mapping.tp_size, - is_column=True) - weights.update( - get_tllm_linear_sq_weight( - int8_weights, - f'{tllm_prex}.mlp.fc.', - [1, 4 * hidden_size // mapping.tp_size], - mapping.tp_size, - is_qkv=False, - per_token=per_token, - per_channel=per_channel, - last_prefix=f'{tllm_prex}.post_layernorm.scale_to_int', - bias=mlp_fc_b, - smoother_value=None, - smoother_shape=None, - rank=mapping.tp_rank, - cat_dim=-1)) - else: - mlp_fc_w = split(mlp_fc_w, - mapping.tp_rank, - mapping.tp_size, - is_column=True) - mlp_fc_b = split(mlp_fc_b, - mapping.tp_rank, - mapping.tp_size, - is_column=True) - if gpt_variant in ['jais']: - weights.update( - get_tllm_linear_weight(mlp_fc_w, f'{tllm_prex}.mlp.gate', - mlp_fc_b, use_weight_only, - plugin_weight_only_quant_type)) - else: - weights.update( - get_tllm_linear_weight(mlp_fc_w, f'{tllm_prex}.mlp.fc', - mlp_fc_b, use_weight_only, - plugin_weight_only_quant_type)) - if gpt_variant in ['jais']: - mlp_fc2_w, mlp_fc2_b = get_weight_and_bias( - model_params, f'{prefix}.mlp.c_fc2', dtype) - mlp_fc2_w = mlp_fc2_w.t().contiguous() - mlp_fc2_w = pad_array_up_to(mlp_fc2_w, 0, mapping.tp_size) - mlp_fc2_b = pad_array_up_to(mlp_fc2_b, 0, mapping.tp_size) - mlp_fc2_w = split(mlp_fc2_w, - mapping.tp_rank, - mapping.tp_size, - is_column=True) - mlp_fc2_b = split(mlp_fc2_b, - mapping.tp_rank, - mapping.tp_size, - is_column=True) - weights.update( - get_tllm_linear_weight(mlp_fc2_w, f'{tllm_prex}.mlp.fc', - mlp_fc2_b, use_weight_only, - plugin_weight_only_quant_type)) - - # (4) MLP Proj Layer - if gpt_variant == 'persimmon': - suffix = "mlp.dense_4h_to_h" - elif gpt_variant == 'kosmos-2': - suffix = "ffn.fc2" - elif gpt_variant == 'nemotron': - suffix = "mlp.down_proj" - else: - suffix = "mlp.c_proj" - mlp_proj_w, mlp_proj_b = get_weight_and_bias(model_params, - f'{prefix}.{suffix}', - dtype) - if gpt_variant in ['gpt2', 'santacoder', 'jais']: - mlp_proj_w = mlp_proj_w.t().contiguous() # transpose for Conv1D - if gpt_variant in ['jais']: - mlp_proj_w = pad_array_up_to(mlp_proj_w, 1, mapping.tp_size) - if use_smooth_quant: - mlp_proj_w_t = mlp_proj_w.t() - int8_weights = generate_int8(mlp_proj_w_t, - act_range.get(f'{prefix}.mlp.c_proj')) - # change it to the real smoother if proj layer is applied smooth quant - fake_smoother_value = torch.ones([1, 4 * hidden_size], - dtype=torch.float32) - weights.update( - get_tllm_linear_sq_weight( - int8_weights, - f'{tllm_prex}.mlp.proj.', [1, hidden_size], - mapping.tp_size, - is_qkv=False, - per_token=per_token, - per_channel=per_channel, - last_prefix=f'{tllm_prex}.mlp.quantization_scaling_factor', - bias=mlp_proj_b, - smoother_value=fake_smoother_value, - smoother_shape=[1, 4 * hidden_size // mapping.tp_size], - rank=mapping.tp_rank, - cat_dim=0)) - else: - mlp_proj_w = split(mlp_proj_w, - mapping.tp_rank, - mapping.tp_size, - is_column=False) - weights.update( - get_tllm_linear_weight(mlp_proj_w, f'{tllm_prex}.mlp.proj', - mlp_proj_b, use_weight_only, - plugin_weight_only_quant_type)) - - # (5) Input layernorm - apply_layernorm_1p = gpt_variant == 'nemotron' - if gpt_variant in ['starcoder2', 'nemotron', 'persimmon']: - input_ln_w, input_ln_b = get_weight_and_bias( - model_params, f'{prefix}.input_layernorm', dtype) - elif gpt_variant == 'kosmos-2': - input_ln_w, input_ln_b = get_weight_and_bias( - model_params, f'{prefix}.self_attn_layer_norm', dtype) - else: - input_ln_w, input_ln_b = get_weight_and_bias( - model_params, f'{prefix}.ln_1', dtype) - if apply_layernorm_1p: - input_ln_w += 1.0 - weights[f'{tllm_prex}.input_layernorm.weight'] = input_ln_w - if input_ln_b is not None: - weights[f'{tllm_prex}.input_layernorm.bias'] = input_ln_b - - # (6) Post layernorm - if gpt_variant in ['starcoder2', 'nemotron', 'persimmon']: - post_ln_w, post_ln_b = get_weight_and_bias( - model_params, f'{prefix}.post_attention_layernorm', dtype) - elif gpt_variant == 'kosmos-2': - post_ln_w, post_ln_b = get_weight_and_bias( - model_params, f'{prefix}.final_layer_norm', dtype) - else: - post_ln_w, post_ln_b = get_weight_and_bias(model_params, - f'{prefix}.ln_2', dtype) - if apply_layernorm_1p: - post_ln_w += 1.0 - weights[f'{tllm_prex}.post_layernorm.weight'] = post_ln_w - if post_ln_b is not None: - weights[f'{tllm_prex}.post_layernorm.bias'] = post_ln_b - - if gpt_variant == 'persimmon': - q_layernorm_w, q_layernorm_b = get_weight_and_bias( - model_params, f'{prefix}.self_attn.q_layernorm', dtype) - - weights[f'{tllm_prex}.attention.q_layernorm.weight'] = q_layernorm_w - weights[f'{tllm_prex}.attention.q_layernorm.bias'] = q_layernorm_b - - k_layernorm_w, k_layernorm_b = get_weight_and_bias( - model_params, f'{prefix}.self_attn.k_layernorm', dtype) - - weights[f'{tllm_prex}.attention.k_layernorm.weight'] = k_layernorm_w - weights[f'{tllm_prex}.attention.k_layernorm.bias'] = k_layernorm_b - - if gpt_variant == 'kosmos-2': - q_layernorm_w, q_layernorm_b = get_weight_and_bias( - model_params, f'{prefix}.self_attn.inner_attn_ln', dtype) - - weights[ - f'{tllm_prex}.attention.inner_layernorm.weight'] = q_layernorm_w - weights[ - f'{tllm_prex}.attention.inner_layernorm.bias'] = q_layernorm_b - - k_layernorm_w, k_layernorm_b = get_weight_and_bias( - model_params, f'{prefix}.ffn.ffn_layernorm', dtype) - - weights[f'{tllm_prex}.mlp.inner_layernorm.weight'] = k_layernorm_w - weights[f'{tllm_prex}.mlp.inner_layernorm.bias'] = k_layernorm_b - - if mapping.is_first_pp_rank(): - if gpt_variant in ['starcoder2', 'nemotron']: - embed_w = get_weight(model_params, 'model.embed_tokens', dtype) - elif gpt_variant == 'kosmos-2': - embed_w = get_weight(model_params, 'text_model.model.embed_tokens', - dtype) - elif gpt_variant == 'persimmon': - embed_w = get_weight(model_params, - ('language_model.' if is_fuyu else '') + - 'model.embed_tokens', dtype) - else: - embed_w = get_weight(model_params, 'transformer.wte', dtype) - weights['transformer.vocab_embedding.weight'] = split_embedding( - embed_w, - mapping.tp_rank, - mapping.tp_size, - use_parallel_embedding=config.use_parallel_embedding, - sharding_dim=config.embedding_sharding_dim) - - if gpt_variant == 'kosmos-2': - padding_idx = hf_config.text_config.pad_token_id - sin_pos_embedding = hf_model.text_model.model.embed_positions.get_embedding( - padding_idx + 1 + hf_config.text_config.max_position_embeddings, - hf_config.text_config.embed_dim, - padding_idx=padding_idx) # [2 + num_embeddings, embed_dim] - pos_embed_w = sin_pos_embedding[2:].to(dtype).detach().cpu() - else: - pos_embed_w = get_weight(model_params, 'transformer.wpe', dtype) - if pos_embed_w is not None: - weights['transformer.position_embedding.weight'] = split_embedding( - pos_embed_w, - mapping.tp_rank, - mapping.tp_size, - use_parallel_embedding=config.use_parallel_embedding, - sharding_dim=config.embedding_sharding_dim) - - if mapping.is_last_pp_rank(): - if gpt_variant in ['starcoder2', 'nemotron']: - embed_w = get_weight(model_params, 'lm_head', dtype) - if embed_w is None: - embed_w = get_weight(model_params, 'model.embed_tokens', dtype) - elif gpt_variant == 'persimmon': - embed_w = get_weight(model_params, - ('language_model.' if is_fuyu else '') + - 'lm_head', dtype) - elif gpt_variant == 'kosmos-2': - embed_w = get_weight(model_params, 'text_model.model.embed_tokens', - dtype) - else: - embed_w = get_weight(model_params, 'transformer.wte', dtype) - - if vocab_size % mapping.tp_size != 0: - vocab_size_padded = pad_vocab_size(vocab_size, mapping.tp_size) - pad_width = vocab_size_padded - vocab_size - embed_w = torch.nn.functional.pad(embed_w, (0, 0, 0, pad_width), - value=0) - if hasattr(hf_config, 'logits_scale'): - embed_w *= hf_config.logits_scale - weights['lm_head.weight'] = split(embed_w.clone(), - mapping.tp_rank, - mapping.tp_size, - is_column=True) - - if gpt_variant in ['starcoder2', 'nemotron']: - ln_f_w, ln_f_b = get_weight_and_bias(model_params, 'model.norm', - dtype) - elif gpt_variant == 'persimmon': - ln_f_w, ln_f_b = get_weight_and_bias( - model_params, ('language_model.' if is_fuyu else '') + - 'model.final_layernorm', dtype) - elif gpt_variant == 'kosmos-2': - ln_f_w, ln_f_b = get_weight_and_bias(model_params, - 'text_model.model.layer_norm', - dtype) - else: - ln_f_w, ln_f_b = get_weight_and_bias(model_params, - 'transformer.ln_f', dtype) - if apply_layernorm_1p: - ln_f_w += 1.0 - weights['transformer.ln_f.weight'] = ln_f_w - if ln_f_b is not None: - weights['transformer.ln_f.bias'] = ln_f_b - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Weights loaded. Total time: {t}') - return weights - - -def quantize(hf_model_dir: str, - output_dir: str, - config: GPTConfig, - device: str = 'cuda', - calib_dataset: str = 'cnn_dailymail', - trust_remote_code: bool = True): - os.makedirs(output_dir, exist_ok=True) - config.to_json_file(os.path.join(output_dir, 'config.json')) - - mapping = config.mapping - assert mapping.rank == 0, "quantize should be called at rank 0 only" - - quant_config = config.quantization - use_smooth_quant = quant_config._use_plugin_sq - int8_kv_cache = quant_config.kv_cache_quant_algo == QuantAlgo.INT8 - - assert use_smooth_quant or int8_kv_cache, "Call from_hugging_face when there is no quantization" - assert hf_model_dir is not None - ## only load and call smooth quant routine once for all ranks - hf_model = AutoModelForCausalLM.from_pretrained( - hf_model_dir, - device_map='auto' if device != 'cpu' else 'cpu', - dtype='auto' if not use_smooth_quant else torch.float16, - trust_remote_code=trust_remote_code) - - os.environ["TOKENIZERS_PARALLELISM"] = os.environ.get( - "TOKENIZERS_PARALLELISM", "false") - tokenizer = AutoTokenizer.from_pretrained( - hf_model_dir, - trust_remote_code=trust_remote_code, - use_fast=False, - padding_side='left') - - dataset = load_calib_dataset(calib_dataset) - act_range = capture_activation_range(hf_model, tokenizer, dataset) - if use_smooth_quant: - smooth_gpt_model(hf_model, act_range, quant_config.smoothquant_val) - - for rank in range(mapping.world_size): - # To avoid changing the mapping arg in-place, also the given mapping from caller is rank agnostic, since quantize is called from only one rank - config = copy.deepcopy(config) - config.set_rank(rank) - weights = load_weights_from_hf_model( - hf_model, - config=config, - act_range=act_range, - ) - safetensors.torch.save_file( - weights, os.path.join(output_dir, f'rank{rank}.safetensors')) - del weights - - -def load_hf_gpt(model_dir: str, load_model_on_cpu: bool = False): - if 'kosmos-2' in model_dir: - # AutoModelForVision2Seq was removed in transformers 5.x; import lazily - # so `import tensorrt_llm` works under newer transformers versions even - # though kosmos-2 conversion itself still requires transformers<5. - from transformers import AutoModelForVision2Seq - hf_model = AutoModelForVision2Seq.from_pretrained( - model_dir, trust_remote_code=True) - else: - hf_model = AutoModelForCausalLM.from_pretrained( - model_dir, - device_map='auto' if not load_model_on_cpu else 'cpu', - dtype='auto', - trust_remote_code=True, - ) - return hf_model - - -def cpu_map_location(storage, loc): - return storage.cpu() - - -def gpu_map_location(storage, loc): - if loc.startswith("cuda"): - training_gpu_idx = int(loc.split(":")[1]) - inference_gpu_idx = training_gpu_idx % torch.cuda.device_count() - return storage.cuda(inference_gpu_idx) - elif loc.startswith("cpu"): - return storage.cpu() - else: - raise ValueError(f"Not handled {loc}") - - -def copy_tokenizer_files(config, out_dir): - basenames = { - "model": "tokenizer", - "vocab_file": "vocab", - "merge_file": "merges", - } - - for key in basenames.keys(): - if config[key] is None: - continue - path = Path(config[key]) - if not path.exists(): - logger.debug(f"Tokenizer {key}: {path} file not found") - continue - - dst_path = out_dir / f"{basenames[key]}{path.suffix}" - logger.debug(f"Copy tokenizer {key}: {path}->{dst_path}") - shutil.copy(path.as_posix(), dst_path.as_posix()) - - -def update_tokenizer_paths(tokenizer_config: Dict, - tokenizer_file_paths: Dict[str, Optional[str]]): - for key, new_path in tokenizer_file_paths.items(): - old_path = tokenizer_config[key] - if old_path is None: - continue - old_path = Path(old_path) - if new_path: - logger.debug(f"Update tokenizer {key} {old_path} -> {new_path}") - tokenizer_config[key] = new_path.as_posix() - elif not old_path.exists(): - logger.warning( - f"Tokenizer {key}'s path {old_path} does not exists: set it to None" - ) - tokenizer_config[key] = None - return tokenizer_config - - -def unpack_nemo_ckpt(nemo_archive_path: Union[str, Path], - out_dir_path: Union[str, Path]): - nemo_archive_path = Path(nemo_archive_path) - if not nemo_archive_path.exists(): - raise FileNotFoundError(f"{nemo_archive_path} does not exist") - - for tar_mode in ["r:", "r:gz"]: - try: - with tarfile.open(nemo_archive_path, mode=tar_mode) as tar_file: - - def is_within_directory(directory, target): - - abs_directory = os.path.abspath(directory) - abs_target = os.path.abspath(target) - - prefix = os.path.commonprefix([abs_directory, abs_target]) - - return prefix == abs_directory - - def safe_members(tar_file): - members = [] - for member in tar_file.getmembers(): - member_path = os.path.join(out_dir_path, member.name) - if not is_within_directory(out_dir_path, member_path): - raise Exception( - "Attempted Path Traversal in Tar File") - members.append(member) - return members - - for member in safe_members(tar_file): - tar_file.extract(member, - path=out_dir_path, - numeric_owner=False, - filter=tarfile.data_filter) - - return out_dir_path - except tarfile.ReadError: - pass - - raise RuntimeError(f"Could not unpack {nemo_archive_path}") - - -def extract_layers_with_prefix(model_, prefix): - length_to_trim = len(prefix) - model_state = model_.get("state_dict", model_) - return { - key[length_to_trim:]: model_state[key] - for key in model_state.keys() if prefix in key - } - - -class UnpackedNemoCheckpointDir: - - def __init__(self, - checkpoints_dir: Union[str, Path], - load_checkpoints_to_cpu: bool = False): - self._checkpoints_dir = Path(checkpoints_dir) - self._load_checkpoints_to_cpu = load_checkpoints_to_cpu - - @property - @functools.lru_cache - def model_config(self): - model_config = None - - model_config_filename = "model_config.yaml" - model_configs_paths = list( - self._checkpoints_dir.rglob(model_config_filename)) - if model_configs_paths: - if len(model_configs_paths) > 1: - raise RuntimeError( - f"There are more than single {model_config_filename} " - f"in {self._checkpoints_dir}: {', '.join(map(lambda p: p.as_posix(), model_configs_paths))}" - ) - model_config_path = model_configs_paths[0] - logger.debug(f"Loading model config from {model_config_path}") - with model_config_path.open("r") as model_config_file: - model_config = yaml.load(model_config_file, - Loader=yaml.SafeLoader) - else: - logger.debug("Searching model config in checkpoints") - # try to obtain from checkpoint - checkpoint_name = self.checkpoint_name - checkpoints_paths = sorted( - self._checkpoints_dir.rglob(checkpoint_name)) - if checkpoints_paths: - # assume that parallel ranks 0 checkpoint should have model config embedded - checkpoint_path = checkpoints_paths[0] - - map_location_fn = cpu_map_location if self._load_checkpoints_to_cpu else gpu_map_location - - model_00 = torch.load(checkpoint_path, - map_location=map_location_fn) - if "hyper_parameters" in model_00 and "cfg" in model_00[ - "hyper_parameters"]: - model_config = model_00["hyper_parameters"]["cfg"] - logger.debug( - f"Loaded model config from checkpoint {checkpoint_path}" - ) - else: - logger.debug( - f"Could not find model config in checkpoint {checkpoint_path}" - ) - del model_00 - - if model_config is None: - logger.warning( - f"Could not find checkpoint with NeMo model config in {self._checkpoints_dir}" - ) - - logger.debug(f"Loaded model config {model_config}") - - return model_config - - @property - def checkpoints_dir(self): - return self._checkpoints_dir - - def get_checkpoints_paths(self, - tensor_model_parallel_size=1, - pipeline_model_parallel_size=1): - """ - Injects tensor/pipeline model parallel ranks into the filepath. - Does nothing if not using model parallelism. - """ - - checkpoint_path_without_rank = self.checkpoints_dir / self.checkpoint_name - - def _inject_parallel_ranks(tp_rank, pp_rank): - if tensor_model_parallel_size > 1 or pipeline_model_parallel_size > 1: - if pipeline_model_parallel_size is None or pipeline_model_parallel_size == 1: - checkpoint_path = (checkpoint_path_without_rank.parent / - f"mp_rank_{tp_rank:02d}" / - checkpoint_path_without_rank.name) - else: - checkpoint_path = ( - checkpoint_path_without_rank.parent / - f"tp_rank_{tp_rank:02d}_pp_rank_{pp_rank:03d}" / - checkpoint_path_without_rank.name) - return checkpoint_path - else: - return checkpoint_path_without_rank - - return [[ - _inject_parallel_ranks(tp_rank=tp_rank, pp_rank=pp_rank) - for pp_rank in range(pipeline_model_parallel_size) - ] for tp_rank in range(tensor_model_parallel_size)] - - @property - @functools.lru_cache - def checkpoint_name(self): - patterns = [ - "model_weights.ckpt", # older megatron checkpoints - "*last.ckpt", # newer format of checkpoints - ] - for pattern in patterns: - model_files = sorted(list(self._checkpoints_dir.rglob(pattern))) - if model_files: - return model_files[0].name - - raise ValueError( - f"Could not find checkpoint files in {self._checkpoints_dir}") - - @functools.lru_cache - def get_tokenizer_file_path(self, tokenizer_key, file_key, - default_filename_pattern): - model_config = self.model_config - file_property = None - if tokenizer_key in model_config and file_key in model_config[ - tokenizer_key]: - file_property = model_config[tokenizer_key][file_key] - elif file_key in model_config: - file_property = model_config[file_key] - - logger.debug( - f"model_config[{tokenizer_key}][{file_key}]={file_property}") - - if file_property and file_property.startswith("nemo:"): - filename = file_property.split("nemo:")[1] - filename_pattern = f"*{filename}" - elif file_property and file_property.startswith("/artifacts/"): - filename = Path(file_property).name - filename_pattern = f"*{filename}" - elif file_property is None or file_property == "None": - filename_pattern = None - else: - filename_pattern = default_filename_pattern - logger.warning( - f"Tokenizer file from config: {tokenizer_key}.{file_key}={file_property} " - f"looks like unsupported path. Pattern {filename_pattern} will be used." - ) - - file_path = None - if filename_pattern is not None: - files_paths = list(self._checkpoints_dir.glob(filename_pattern)) - if files_paths: - assert len(files_paths) == 1 - file_path = files_paths[0] - - return file_path - - @functools.lru_cache - def get_all_tokenizer_file_paths(self): - return { - "model": - self.get_tokenizer_file_path("tokenizer", "model", "*.model"), - "vocab_file": - self.get_tokenizer_file_path("tokenizer", "vocab_file", "*vocab*"), - "merge_file": - self.get_tokenizer_file_path("tokenizer", "merge_file", - "*merge*.txt"), - } - - -@torch.no_grad() -def load_torch_checkpoints(checkpoints_paths, - merge_factor, - tp_rank, - pp_rank, - map_location_fn, - handle_model_level_weights, - layer_rename_config: Dict[str, str] = {}): - models = [] - for k in range(merge_factor): - rank_weights = checkpoints_paths[tp_rank * merge_factor + k][pp_rank] - model = torch.load(rank_weights, map_location=map_location_fn) - model = rename_keys(model, layer_rename_config) - handle_model_level_weights(model, tp_rank * merge_factor + k, pp_rank) - layers = extract_layers_with_prefix(model, - "model.language_model.encoder.") - models.append(layers) - return models - - -@torch.no_grad() -def load_weights_from_nemo(nemo_ckpt_dir: str, config: GPTConfig, **kwargs): - assert config.mapping.pp_size == 1, \ - "Pipeline parallelism is not supported." - assert not config.quantization.quant_mode.has_any_quant(), \ - "Quantization is not supported." - - load_model_on_cpu = kwargs.pop('load_model_on_cpu', False) - nemo_rename_key = kwargs.pop('nemo_rename_key', []) - layer_rename_config = { - pattern.split(':')[0]: pattern.split(':')[1] - for pattern in nemo_rename_key - } - - unpacked_checkpoints_dir = UnpackedNemoCheckpointDir( - nemo_ckpt_dir, load_checkpoints_to_cpu=load_model_on_cpu) - nemo_model_config = unpacked_checkpoints_dir.model_config - - checkpoints_paths = unpacked_checkpoints_dir.get_checkpoints_paths( - nemo_model_config.get("tensor_model_parallel_size", 1), - nemo_model_config.get("pipeline_model_parallel_size", 1), - ) - - if unpacked_checkpoints_dir._load_checkpoints_to_cpu: - map_location_fn = cpu_map_location - else: - map_location_fn = gpu_map_location - dtype = str_dtype_to_torch(config.dtype) - - # load position_embedding from rank 0 - model_00 = torch.load(checkpoints_paths[0][0], map_location=map_location_fn) - model_00 = model_00.get("state_dict", model_00) - model_00 = rename_keys(model_00, layer_rename_config) - has_position_embedding = "model.language_model.embedding.position_embeddings.weight" in model_00 - has_lm_head = "model.language_model.output_layer.weight" in model_00 - del model_00 - - num_layers = nemo_model_config["num_layers"] - training_tp_size = nemo_model_config.get("tensor_model_parallel_size", 1) - training_pp_size = nemo_model_config.get("pipeline_model_parallel_size", 1) - inference_tp_size = config.mapping.tp_size - inference_tp_rank = config.mapping.tp_rank - - apply_layernorm_1p = (nemo_model_config.get('normalization', - '') == "layernorm1p") - split_gated_activation = ("swiglu" - in nemo_model_config.get('activation', "gelu")) - num_attention_heads = nemo_model_config["num_attention_heads"] - # use_attention_nemo_shape = True - transpose_weights = True - # multi_query_mode = False - local_dim = None - - # merge_factor: how many TP training nodes are merged into an inference TP node - # split_factor: in how many parts a TP training node is split - gcd = np.gcd(training_tp_size, inference_tp_size) - merge_factor = training_tp_size // gcd - split_factor = inference_tp_size // gcd - - model_level_weights = defaultdict(list) - - def handle_model_level_weights(model, tp_idx: int, pp_idx: int): - if tp_idx == 0 and pp_idx == 0: - if has_position_embedding: - val = model[ - "model.language_model.embedding.position_embeddings.weight"].detach( - ).cpu() - model_level_weights[ - "transformer.position_embedding.weight"].append(val) - if pp_idx == 0: - val = model.get( - "state_dict", model - )["model.language_model.embedding.word_embeddings.weight"].detach( - ).cpu() - model_level_weights["transformer.vocab_embedding.weight"].append( - val) - if has_lm_head and pp_idx == training_pp_size - 1: - val = model.get( - "state_dict", - model)["model.language_model.output_layer.weight"].detach().cpu( - ) - model_level_weights["lm_head.weight"].append(val) - - weights = {} - tik = time.time() - tp_rank = inference_tp_rank // split_factor - # for tp_rank in range(training_tp_size // merge_factor): - for pp_rank in range(training_pp_size): - models = load_torch_checkpoints(checkpoints_paths, merge_factor, - tp_rank, pp_rank, map_location_fn, - handle_model_level_weights, - layer_rename_config) - for name in list(models[0].keys()): - params = [model[name].detach().cpu() for model in models] - if transpose_weights and params[0].ndim == 2: - params = [p.T for p in params] - if "layernorm.weight" in name and apply_layernorm_1p: - params = [p + 1.0 for p in params] - - l = retrieved_layer_index_from_name(name) - if l is not None: - new_l = l + pp_rank * num_layers // training_pp_size - prefix = f'transformer.layers.{new_l}' - - if 'attention.query_key_value' in name: - if name.endswith('weight'): - hidden_dim = params[0].shape[0] - if local_dim is None: - local_dim = params[0].shape[-1] // 3 - - # multi_query_mode = False; use_attention_nemo_shape = True - head_num = num_attention_heads // training_tp_size - size_per_head = hidden_dim // num_attention_heads - params = [ - param.reshape(hidden_dim, head_num, 3, - size_per_head) for param in params - ] - params = [param.permute(0, 2, 1, 3) for param in params] - params = [ - param.reshape(hidden_dim, 3, local_dim) - for param in params - ] - cat_dim = -1 - param = torch.concat(params, dim=cat_dim) - param = torch.chunk(param, split_factor, - dim=cat_dim)[inference_tp_rank % - split_factor] - weights[ - f'{prefix}.attention.qkv.weight'] = param.reshape( - hidden_dim, -1).t() - else: - if local_dim is None: - local_dim = params[0].shape[-1] // 3 - - # multi_query_mode = False; use_attention_nemo_shape = True - head_num = num_attention_heads // training_tp_size - size_per_head = local_dim // head_num - params = [ - param.reshape(head_num, 3, size_per_head) - for param in params - ] - params = [param.permute(1, 0, 2) for param in params] - params = [ - param.reshape(3, local_dim) for param in params - ] - cat_dim = -1 - param = torch.concat(params, dim=cat_dim) - param = torch.chunk(param, split_factor, - dim=cat_dim)[inference_tp_rank % - split_factor] - weights[f'{prefix}.attention.qkv.bias'] = param.reshape( - -1) - - elif 'attention.dense' in name: - if name.endswith('weight'): - cat_dim = 0 - param = torch.concat(params, dim=cat_dim) - param = torch.chunk(param, split_factor, - dim=cat_dim)[inference_tp_rank % - split_factor] - weights[f'{prefix}.attention.dense.weight'] = param.t() - else: - weights[f'{prefix}.attention.dense.bias'] = params[0] - - elif 'mlp.dense_h_to_4h' in name: - if name.endswith('weight'): - if split_gated_activation: - params = [torch.chunk(p, 2, dim=-1) for p in params] - params, gate_params = list(zip(*params)) - cat_dim = -1 - param = torch.concat(params, dim=cat_dim) - param = torch.chunk(param, split_factor, - dim=cat_dim)[inference_tp_rank % - split_factor] - weights[f'{prefix}.mlp.fc.weight'] = param.t() - if split_gated_activation: - gate_param = torch.concat(gate_params, dim=cat_dim) - gate_param = torch.chunk( - gate_param, split_factor, - dim=cat_dim)[inference_tp_rank % split_factor] - weights[f'{prefix}.mlp.gate.weight'] = gate_param.t( - ) - else: - if split_gated_activation: - params = [torch.chunk(p, 2, dim=-1) for p in params] - params, gate_params = list(zip(*params)) - cat_dim = -1 - param = torch.concat(params, dim=cat_dim) - param = torch.chunk(param, split_factor, - dim=cat_dim)[inference_tp_rank % - split_factor] - weights[f'{prefix}.mlp.fc.bias'] = param - if split_gated_activation: - gate_param = torch.concat(gate_params, dim=cat_dim) - gate_param = torch.chunk( - gate_param, split_factor, - dim=cat_dim)[inference_tp_rank % split_factor] - weights[f'{prefix}.mlp.gate.bias'] = gate_param - - elif 'mlp.dense_4h_to_h' in name: - if name.endswith('weight'): - cat_dim = 0 - param = torch.concat(params, dim=cat_dim) - param = torch.chunk(param, split_factor, - dim=cat_dim)[inference_tp_rank % - split_factor] - weights[f'{prefix}.mlp.proj.weight'] = param.t() - else: - weights[f'{prefix}.mlp.proj.bias'] = params[0] - - elif 'input_layernorm' in name: - if name.endswith('weight'): - weights[f'{prefix}.input_layernorm.weight'] = params[0] - else: - weights[f'{prefix}.input_layernorm.bias'] = params[0] - elif 'post_attention_layernorm' in name: - if name.endswith('weight'): - weights[f'{prefix}.post_layernorm.weight'] = params[0] - else: - weights[f'{prefix}.post_layernorm.bias'] = params[0] - - elif 'final_layernorm' in name: - if name.endswith('weight'): - weights['transformer.ln_f.weight'] = params[0] - else: - weights['transformer.ln_f.bias'] = params[0] - for model in models: - del model[name] - del models - - for key in list(model_level_weights.keys()): - weights[key] = torch.concat(model_level_weights[key], dim=0) - weights[key] = torch.chunk(weights[key], split_factor, - dim=0)[inference_tp_rank % split_factor] - del model_level_weights[key] - for key, param in weights.items(): - weights[key] = weights[key].to(dtype).contiguous() - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Weights loaded. Total time: {t}') - return weights diff --git a/tensorrt_llm/models/gpt/model.py b/tensorrt_llm/models/gpt/model.py deleted file mode 100644 index 89267a90f71e..000000000000 --- a/tensorrt_llm/models/gpt/model.py +++ /dev/null @@ -1,417 +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. - -from typing import Optional, Union - -from ..._utils import pad_vocab_size -from ...functional import (Tensor, is_gated_activation, non_gated_version, recv, - send) -from ...layers import (MLP, MOE, Attention, AttentionMaskType, ColumnLinear, - Embedding, GatedMLP, LayerNorm, MoeConfig, - PositionEmbeddingType) -from ...lora_helper import LoraConfig, use_lora -from ...mapping import Mapping -from ...module import Module -from ...quantization import QuantMode -from ..model_weights_loader import ModelWeightsLoader -from ..modeling_utils import (DecoderLayerList, DecoderModelForCausalLM, - QuantConfig) -from .config import GPTConfig -from .convert import (load_hf_gpt, load_weights_from_hf_model, - load_weights_from_nemo) - - -def MLPFactory(hidden_size, - ffn_hidden_size, - hidden_act, - bias=True, - dtype=None, - moe_config: MoeConfig = MoeConfig(), - tp_group=None, - tp_size=1, - mapping=Mapping(), - quant_mode=QuantMode(0), - inner_layernorm=False, - eps=1e-05): - if moe_config.has_moe(): - return MOE(moe_config, - hidden_size, - ffn_hidden_size, - hidden_act, - mapping=mapping, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - quant_mode=quant_mode) - MLPClass = GatedMLP if is_gated_activation(hidden_act) else MLP - hidden_act = non_gated_version(hidden_act) - return MLPClass( - hidden_size, - ffn_hidden_size, - hidden_act, - bias, - dtype, - tp_group, - tp_size, - quant_mode, - inner_layernorm=inner_layernorm, - eps=eps, - ) - - -class GPTDecoderLayer(Module): - - def __init__(self, config: GPTConfig, layer_idx: int): - super().__init__() - self.layer_idx = layer_idx - self.config = config - - tp_group = config.mapping.tp_group - tp_size = config.mapping.tp_size - tp_rank = config.mapping.tp_rank - - self.input_layernorm = LayerNorm(normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - dtype=config.dtype) - - layers_range = config.mapping.pp_layers(config.num_hidden_layers) - local_layer_idx = layer_idx - layers_range[0] - inner_layernorm = config.inner_layernorm if hasattr( - config, "inner_layernorm") else False - attention_head_size = config.head_size if hasattr(config, - "head_size") else None - self.attention = Attention( - local_layer_idx=local_layer_idx, - hidden_size=config.hidden_size, - num_attention_heads=config.num_attention_heads, - num_kv_heads=config.num_key_value_heads, - max_position_embeddings=config.max_position_embeddings, - num_layers=config.num_hidden_layers, - q_scaling=config.q_scaling, - apply_query_key_layer_scaling=config.apply_query_key_layer_scaling, - dtype=config.dtype, - attention_mask_type=AttentionMaskType.causal, - attention_head_size=attention_head_size, - position_embedding_type=config.position_embedding_type, - rotary_embedding_percentage=config.rotary_pct, - rotary_embedding_base=config.rotary_base, - rotary_embedding_scaling=config.rotary_scaling, - bias=config.bias, - tp_group=tp_group, - tp_size=tp_size, - tp_rank=tp_rank, - quant_mode=config.quant_mode, - qk_layernorm=config.qk_layernorm, - inner_layernorm=inner_layernorm, - eps=config.norm_epsilon) - - mlp_hidden_size = config.hidden_size * 4 if config.intermediate_size is None else config.intermediate_size - self.norm_before_bmm1 = config.norm_before_bmm1 if hasattr( - config, "norm_before_bmm1") else False - - self.mlp = MLPFactory(hidden_size=config.hidden_size, - ffn_hidden_size=mlp_hidden_size, - hidden_act=config.hidden_act, - dtype=config.dtype, - bias=config.bias, - moe_config=config.moe, - tp_group=tp_group, - tp_size=tp_size, - mapping=config.mapping, - quant_mode=config.quant_mode, - inner_layernorm=inner_layernorm, - eps=config.norm_epsilon) - - self.post_layernorm = LayerNorm(normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - dtype=config.dtype) - - def forward(self, - hidden_states: Tensor, - attention_mask=None, - use_cache=False, - kv_cache_params=None, - attention_params=None, - lora_layer_params=None, - spec_decoding_params=None): - - assert isinstance(hidden_states, Tensor) - - residual = hidden_states - - hidden_states = self.input_layernorm(hidden_states) - - attention_output = self.attention( - hidden_states, - attention_mask=attention_mask, - use_cache=use_cache, - spec_decoding_params=spec_decoding_params, - kv_cache_params=kv_cache_params, - attention_params=attention_params, - lora_layer_params=lora_layer_params, - norm_before_bmm1=self.norm_before_bmm1) - - if use_cache: - attention_output, presents = attention_output - - hidden_states = residual + attention_output - - residual = hidden_states - hidden_states = self.post_layernorm(hidden_states) - - hidden_states = self.mlp(hidden_states, - lora_layer_params=lora_layer_params) - - hidden_states = residual + hidden_states - - if use_cache: - return (hidden_states, presents) - return hidden_states - - -class GPTModel(Module): - - def __init__(self, config: GPTConfig): - super().__init__() - self.mapping = config.mapping - self.position_embedding_type = config.position_embedding_type - if config.mapping.is_first_pp_rank(): - self.vocab_embedding = Embedding(config.vocab_size, - config.hidden_size, - dtype=config.dtype) - - self.embedding_scale = config.embedding_scale - - if config.position_embedding_type == PositionEmbeddingType.learned_absolute: - self.position_embedding = Embedding( - num_embeddings=config.max_position_embeddings, - embedding_dim=config.hidden_size, - dtype=config.dtype) - - self.layers = DecoderLayerList(GPTDecoderLayer, config) - - if config.mapping.is_last_pp_rank(): - self.ln_f = LayerNorm(normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - dtype=config.dtype) - - def forward(self, - input_ids, - position_ids, - use_cache=False, - attention_mask=None, - kv_cache_params=None, - attention_params=None, - hidden_states=None, - prompt_embedding_table=None, - prompt_tasks=None, - prompt_vocab_size=None, - lora_params=None, - spec_decoding_params=None): - if self.mapping.is_first_pp_rank(): - ptuning_args = [ - prompt_embedding_table, prompt_tasks, prompt_vocab_size - ] if prompt_embedding_table is not None else [] - hidden_states = self.vocab_embedding(input_ids, *ptuning_args) - if self.embedding_scale is not None: - hidden_states *= self.embedding_scale - if self.position_embedding_type == PositionEmbeddingType.learned_absolute: - hidden_states = hidden_states + self.position_embedding( - position_ids) - else: - hidden_states = recv(hidden_states, self.mapping.prev_pp_rank()) - - hidden_states = self.layers(hidden_states, - use_cache=use_cache, - attention_mask=attention_mask, - kv_cache_params=kv_cache_params, - attention_params=attention_params, - lora_params=lora_params, - spec_decoding_params=spec_decoding_params) - if use_cache: - hidden_states, presents = hidden_states - - if self.mapping.is_last_pp_rank(): - hidden_states = self.ln_f(hidden_states) - else: - hidden_states = send(hidden_states, self.mapping.next_pp_rank()) - - if use_cache: - return (hidden_states, tuple(presents)) - return hidden_states - - -class GPTForCausalLM(DecoderModelForCausalLM): - config_class = GPTConfig - - def __init__(self, config: GPTConfig): - transformer = GPTModel(config) - - if config.mapping.is_last_pp_rank(): - vocab_size_padded = pad_vocab_size(config.vocab_size, - config.mapping.tp_size) - lm_head = ColumnLinear(config.hidden_size, - vocab_size_padded, - bias=False, - dtype=config.dtype, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size, - gather_output=True) - else: - lm_head = None - self.trtllm_modules_to_hf_modules = { - "attn_q": "q_proj", - "attn_k": "k_proj", - "attn_v": "v_proj", - "attn_dense": "o_proj", - "mlp_h_to_4h": "c_fc", - "mlp_4h_to_h": "c_proj", - } - super().__init__(config, transformer, lm_head) - - @classmethod - def from_hugging_face( - cls, - hf_model_or_dir: Union[str, 'transformers.PreTrainedModel'], - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - ''' Create a LLaMAForCausalLM object from give parameters - ''' - import transformers - - load_model_on_cpu = kwargs.pop('load_model_on_cpu', False) - is_prequantized_to_fp8 = kwargs.pop('is_prequantized_to_fp8', False) - - assert hf_model_or_dir is not None - use_preloading = isinstance(hf_model_or_dir, - transformers.PreTrainedModel) - if use_preloading: - hf_model = hf_model_or_dir - hf_config_or_dir = hf_model.config - else: - hf_model_dir = hf_model_or_dir - hf_config_or_dir = hf_model_or_dir - - config = GPTConfig.from_hugging_face(hf_config_or_dir, - dtype=dtype, - mapping=mapping, - quant_config=quant_config, - **kwargs) - if is_prequantized_to_fp8: - custom_dict = {'fc': 'up_proj'} - loader = ModelWeightsLoader(hf_model_dir, custom_dict) - model = cls(config) - - # This is to account for all layernorms in nemotron variants being NemotronLayerNorm-1P. - def apply_layernorm_1p(weights): - return weights + 1.0 - - loader.update_key_mapping(model) - tllm_weights = {} - for tllm_key, _ in model.named_parameters(): - if config.gpt_variant == "nemotron" and ( - 'layernorm.weight' in tllm_key - or 'ln_f.weight' in tllm_key): - tllm_weights.update( - loader.load(tllm_key, - preprocess=apply_layernorm_1p, - custom_postprocess_kwargs={})) - else: - tllm_weights.update( - loader.load(tllm_key, custom_postprocess_kwargs={})) - loader.fill(tllm_weights) - else: - if not use_preloading: - hf_model = load_hf_gpt(hf_model_dir, load_model_on_cpu) - weights = load_weights_from_hf_model(hf_model, config) - model = cls(config) - model.load(weights) - return model - - @classmethod - def quantize( - cls, - hf_model_dir: str, - output_dir: str, - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - *, - device: str = 'cuda', - calib_dataset: str = 'cnn_dailymail', - calib_batches: int = 512, - calib_batch_size: int = 1, - calib_max_seq_length: int = 512, - random_seed: int = 1234, - tokenizer_max_seq_length: int = 2048, - **kwargs, - ): - if quant_config._requires_modelopt_quantization: - # modelopt quantization flow - super().quantize(hf_model_dir, - output_dir, - dtype=dtype, - mapping=mapping, - quant_config=quant_config, - device=device, - calib_dataset=calib_dataset, - calib_batches=calib_batches, - calib_batch_size=calib_batch_size, - calib_max_seq_length=calib_max_seq_length, - random_seed=random_seed, - tokenizer_max_seq_length=tokenizer_max_seq_length) - elif quant_config._requires_calibration: - # non-modelopt quantization flow - from . import convert - - config = GPTConfig.from_hugging_face(hf_model_dir, - dtype=dtype, - mapping=mapping, - quant_config=quant_config, - **kwargs) - convert.quantize(hf_model_dir, - output_dir, - config=config, - device=device, - calib_dataset=calib_dataset) - else: - raise ValueError( - f"The quant_config ({quant_config}) does not require calibration, try {cls.__name__}.from_hugging_face instead." - ) - - @classmethod - def from_nemo(cls, - nemo_ckpt_dir: str, - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - config = GPTConfig.from_nemo(nemo_ckpt_dir, - dtype=dtype, - mapping=mapping, - quant_config=quant_config, - **kwargs) - - weights = load_weights_from_nemo(nemo_ckpt_dir, config, **kwargs) - - model = cls(config) - model.load(weights) - return model - - def use_lora(self, lora_config: LoraConfig): - use_lora(self, lora_config, self.trtllm_modules_to_hf_modules) diff --git a/tensorrt_llm/models/gptj/__init__.py b/tensorrt_llm/models/gptj/__init__.py deleted file mode 100644 index 71bf6d298d42..000000000000 --- a/tensorrt_llm/models/gptj/__init__.py +++ /dev/null @@ -1,14 +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. diff --git a/tensorrt_llm/models/gptj/config.py b/tensorrt_llm/models/gptj/config.py deleted file mode 100644 index e4a7d3d99e9b..000000000000 --- a/tensorrt_llm/models/gptj/config.py +++ /dev/null @@ -1,55 +0,0 @@ -from typing import Mapping, Optional, Union - -from ...mapping import Mapping -from ..convert_utils import infer_dtype -from ..modeling_utils import PretrainedConfig, QuantConfig - - -class GPTJConfig(PretrainedConfig): - """ - This is the configuration class to store the configuration of GPTJ model. - """ - - def __init__(self, *, rotary_dim: int = 64, **kwargs): - self.rotary_dim = rotary_dim - super().__init__(**kwargs) - - def to_dict(self): - output = super().to_dict() - output.update(rotary_dim=self.rotary_dim) - return output - - @classmethod - def from_hugging_face( - cls, - hf_config_or_dir: Union[str, 'transformers.PretrainedConfig'], - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - import transformers - trust_remote_code = kwargs.pop('trust_remote_code', True) - - if isinstance(hf_config_or_dir, transformers.PretrainedConfig): - hf_config = hf_config_or_dir - else: - hf_config_dir = str(hf_config_or_dir) - hf_config = transformers.AutoConfig.from_pretrained( - hf_config_dir, trust_remote_code=trust_remote_code) - - dtype = infer_dtype(dtype, getattr(hf_config, 'torch_dtype', None)) - - return cls(architecture=hf_config.architectures[0], - dtype=dtype, - num_hidden_layers=hf_config.num_hidden_layers, - num_attention_heads=hf_config.num_attention_heads, - hidden_size=hf_config.hidden_size, - norm_epsilon=hf_config.layer_norm_epsilon, - vocab_size=hf_config.vocab_size, - position_embedding_type='rope_gptj', - max_position_embeddings=hf_config.max_position_embeddings, - hidden_act='gelu', - rotary_dim=hf_config.rotary_dim, - mapping=mapping, - quantization=quant_config, - **kwargs) diff --git a/tensorrt_llm/models/gptj/convert.py b/tensorrt_llm/models/gptj/convert.py deleted file mode 100644 index 1f10b3b260b0..000000000000 --- a/tensorrt_llm/models/gptj/convert.py +++ /dev/null @@ -1,170 +0,0 @@ -import time -from typing import Dict, Optional - -import torch - -from tensorrt_llm.quantization import QuantAlgo - -from ..convert_utils import get_weight, get_weight_and_bias, split_matrix_tp -from .config import GPTJConfig - - -def get_tllm_linear_weight( - weight: torch.Tensor, - prefix: str, - bias: Optional[torch.Tensor] = None, - use_weight_only: bool = False, - plugin_weight_only_quant_type: torch.dtype = torch.int8 -) -> Dict[str, torch.Tensor]: - results = {} - if use_weight_only: - v = weight.t().contiguous() - processed_torch_weights, torch_weight_scales = \ - torch.ops.trtllm.symmetric_quantize_last_axis_of_batched_matrix( - v, plugin_weight_only_quant_type) - results[f'{prefix}.weight'] = processed_torch_weights - results[f'{prefix}.per_channel_scale'] = torch_weight_scales - else: - results[f'{prefix}.weight'] = weight.contiguous() - - if bias is not None: - results[f'{prefix}.bias'] = bias - - return results - - -def get_tllm_param( - param: torch.Tensor, - name: str, - use_weight_only: bool = False, - plugin_weight_only_quant_type: torch.dtype = torch.int8 -) -> Dict[str, torch.Tensor]: - results = {} - if name.endswith('.weight') and use_weight_only: - v = param.t().contiguous() - processed_torch_weights, torch_weight_scales = \ - torch.ops.trtllm.symmetric_quantize_last_axis_of_batched_matrix( - v, plugin_weight_only_quant_type) - results[name] = processed_torch_weights - results[name.replace('weight', - 'per_channel_scale')] = torch_weight_scales - else: - results[name] = param - - return results - - -def load_weights_from_hf_model(hf_model, config: GPTJConfig): - quant_algo = config.quantization.quant_algo - use_weight_only = quant_algo in [QuantAlgo.W8A16, QuantAlgo.W4A16] - if quant_algo == QuantAlgo.W8A16: - plugin_weight_only_quant_type = torch.int8 - elif quant_algo == QuantAlgo.W4A16: - plugin_weight_only_quant_type = torch.quint4x2 - else: - plugin_weight_only_quant_type = None - - weights = {} - tik = time.time() - - model_params = dict(hf_model.named_parameters()) - dtype = getattr(torch, config.dtype) - num_hidden_layers = config.num_hidden_layers - mapping = config.mapping - - layers_range = mapping.pp_layers(num_hidden_layers) - for l in layers_range: - prefix = f'transformer.h.{l}' - tllm_prex = f'transformer.layers.{l-layers_range[0]}' - # Attention QKV (no bias) - q_weight = get_weight(model_params, f'{prefix}.attn.q_proj', dtype) - k_weight = get_weight(model_params, f'{prefix}.attn.k_proj', dtype) - v_weight = get_weight(model_params, f'{prefix}.attn.v_proj', dtype) - q_w = split_matrix_tp(q_weight, mapping.tp_size, mapping.tp_rank, dim=0) - k_w = split_matrix_tp(k_weight, mapping.tp_size, mapping.tp_rank, dim=0) - v_w = split_matrix_tp(v_weight, mapping.tp_size, mapping.tp_rank, dim=0) - qkv_w = torch.concatenate([q_w, k_w, v_w], dim=0) - weights.update( - get_tllm_linear_weight(qkv_w, f'{tllm_prex}.attention.qkv', None, - use_weight_only, - plugin_weight_only_quant_type)) - # Attention dense (not bias) - attn_dense_weight = get_weight(model_params, f'{prefix}.attn.out_proj', - dtype) - attn_dense_w = split_matrix_tp(attn_dense_weight, - mapping.tp_size, - mapping.tp_rank, - dim=1) - weights.update( - get_tllm_linear_weight(attn_dense_w, f'{tllm_prex}.attention.dense', - None, use_weight_only, - plugin_weight_only_quant_type)) - # MLP fc_in (with bias) - mlp_fc_weight, mlp_fc_bias = get_weight_and_bias( - model_params, f'{prefix}.mlp.fc_in', dtype) - mlp_fc_w = split_matrix_tp(mlp_fc_weight, - mapping.tp_size, - mapping.tp_rank, - dim=0) - mlp_fc_b = split_matrix_tp(mlp_fc_bias, - mapping.tp_size, - mapping.tp_rank, - dim=0) - weights.update( - get_tllm_linear_weight(mlp_fc_w, f'{tllm_prex}.mlp.fc', mlp_fc_b, - use_weight_only, - plugin_weight_only_quant_type)) - # MLP fc_out (with bias) - mlp_proj_weight, mlp_proj_bias = get_weight_and_bias( - model_params, f'{prefix}.mlp.fc_out', dtype) - mlp_proj_w = split_matrix_tp(mlp_proj_weight, - mapping.tp_size, - mapping.tp_rank, - dim=1) - # Only rank0 will get bias - if mapping.tp_size > 1 and mapping.tp_rank > 0: - mlp_proj_bias = torch.zeros(mlp_proj_weight.shape[0], - dtype=mlp_proj_weight.dtype) - weights.update( - get_tllm_linear_weight(mlp_proj_w, f'{tllm_prex}.mlp.proj', - mlp_proj_bias, use_weight_only, - plugin_weight_only_quant_type)) - - input_ln_weight, input_ln_bias = get_weight_and_bias( - model_params, f'{prefix}.ln_1', dtype) - weights[f'{tllm_prex}.input_layernorm.weight'] = input_ln_weight - weights[f'{tllm_prex}.input_layernorm.bias'] = input_ln_bias - - if mapping.is_first_pp_rank(): - # Embedding - embed_w = get_weight(model_params, 'transformer.wte', dtype) - if config.use_parallel_embedding: - embed_w = split_matrix_tp(embed_w, - mapping.tp_size, - mapping.tp_rank, - dim=0) - weights['transformer.vocab_embedding.weight'] = embed_w - - if mapping.is_last_pp_rank(): - # lm_head weight and bias - lm_head_w, ln_head_bias = get_weight_and_bias(model_params, 'lm_head', - dtype) - weights['lm_head.weight'] = split_matrix_tp(lm_head_w, - mapping.tp_size, - mapping.tp_rank, - dim=0) - weights['lm_head.bias'] = split_matrix_tp(ln_head_bias, - mapping.tp_size, - mapping.tp_rank, - dim=0) - ln_f_w, ln_f_b = get_weight_and_bias(model_params, 'transformer.ln_f', - dtype) - # ln_f weight and bias - weights['transformer.ln_f.weight'] = ln_f_w - if ln_f_b is not None: - weights['transformer.ln_f.bias'] = ln_f_b - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Weights loaded. Total time: {t}') - return weights diff --git a/tensorrt_llm/models/gptj/model.py b/tensorrt_llm/models/gptj/model.py deleted file mode 100644 index ff2d1deab24c..000000000000 --- a/tensorrt_llm/models/gptj/model.py +++ /dev/null @@ -1,202 +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. - -from typing import Optional, Union - -from ..._utils import pad_vocab_size -from ...functional import PositionEmbeddingType, Tensor, allreduce -from ...layers import (MLP, Attention, AttentionMaskType, ColumnLinear, - Embedding, LayerNorm) -from ...mapping import Mapping -from ...module import Module -from ..modeling_utils import DecoderLayerList, DecoderModelForCausalLM -from .config import GPTJConfig -from .convert import load_weights_from_hf_model - - -class GPTJDecoderLayer(Module): - - def __init__(self, config: GPTJConfig, layer_idx: int): - super().__init__() - self.layer_idx = layer_idx - self.config = config - - hidden_size = config.hidden_size - num_attention_heads = config.num_attention_heads - rotary_dim = config.rotary_dim - dtype = config.dtype - tp_size = config.mapping.tp_size - tp_rank = config.mapping.tp_rank - layernorm_epsilon = config.norm_epsilon - - self.input_layernorm = LayerNorm(normalized_shape=hidden_size, - eps=layernorm_epsilon, - dtype=dtype) - - layers_range = config.mapping.pp_layers(config.num_hidden_layers) - local_layer_idx = layer_idx - layers_range[0] - self.attention = Attention( - local_layer_idx=local_layer_idx, - hidden_size=hidden_size, - num_attention_heads=num_attention_heads, - rotary_embedding_percentage=rotary_dim / - (hidden_size // num_attention_heads), - max_position_embeddings=config.max_position_embeddings, - attention_mask_type=AttentionMaskType.causal, - dtype=dtype, - tp_group=None, - tp_size=tp_size, - tp_rank=tp_rank, - bias=False, - position_embedding_type=PositionEmbeddingType.rope_gptj, - quant_mode=config.quant_mode) - - self.mlp = MLP(hidden_size=hidden_size, - ffn_hidden_size=hidden_size * 4, - hidden_act=config.hidden_act, - dtype=dtype, - bias=True, - tp_group=None, - tp_size=tp_size, - quant_mode=config.quant_mode) - - def forward(self, - hidden_states: Tensor, - attention_mask=None, - use_cache=False, - kv_cache_params=None, - attention_params=None): - assert isinstance(hidden_states, Tensor) - - residual = hidden_states - - hidden_states = self.input_layernorm(hidden_states) - - attention_output = self.attention(hidden_states, - attention_mask=attention_mask, - use_cache=use_cache, - kv_cache_params=kv_cache_params, - attention_params=attention_params) - - if use_cache: - attention_output, presents = attention_output - attention_output = attention_output - - feed_forward_hidden_states = self.mlp(hidden_states) - hidden_states = attention_output + feed_forward_hidden_states - if self.config.mapping.tp_size > 1: - hidden_states = allreduce(hidden_states, - self.config.mapping.tp_group) - hidden_states = hidden_states + residual - - if use_cache: - return (hidden_states, presents) - return hidden_states - - -class GPTJModel(Module): - - def __init__(self, config: GPTJConfig): - super().__init__() - self.config = config - - if config.mapping.is_first_pp_rank(): - self.vocab_embedding = Embedding(config.vocab_size, - config.hidden_size, - dtype=config.dtype) - self.layers = DecoderLayerList(GPTJDecoderLayer, config) - if config.mapping.is_last_pp_rank(): - self.ln_f = LayerNorm(normalized_shape=config.hidden_size, - dtype=config.dtype) - - def forward(self, - input_ids: Tensor, - position_ids=None, - use_cache=False, - attention_mask=None, - kv_cache_params=None, - attention_params=None): - - hidden_states = self.vocab_embedding(input_ids) - - hidden_states = self.layers(hidden_states, - use_cache=use_cache, - attention_mask=attention_mask, - kv_cache_params=kv_cache_params, - attention_params=attention_params) - - if use_cache: - hidden_states, presents = hidden_states - - hidden_states = self.ln_f(hidden_states) - - if use_cache: - return (hidden_states, tuple(presents)) - return hidden_states - - -class GPTJForCausalLM(DecoderModelForCausalLM): - config_class = GPTJConfig - - def __init__(self, config: GPTJConfig): - transformer = GPTJModel(config) - vocab_size_padded = pad_vocab_size(config.vocab_size, - config.mapping.tp_size) - if config.mapping.is_last_pp_rank(): - lm_head = ColumnLinear(config.hidden_size, - vocab_size_padded, - bias=True, - dtype=config.dtype, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size, - gather_output=True) - else: - lm_head = None - super().__init__(config, transformer, lm_head) - - @classmethod - def from_hugging_face( - cls, - hf_model_or_dir: Union[str, 'transformers.PreTrainedModel'], - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config=None, - **kwargs): - import transformers - use_preloading = isinstance(hf_model_or_dir, - transformers.PreTrainedModel) - if use_preloading: - hf_model = hf_model_or_dir - hf_config_or_dir = hf_model.config - else: - hf_model_dir = hf_model_or_dir - hf_config_or_dir = hf_model_or_dir - - config = GPTJConfig.from_hugging_face(hf_config_or_dir, - dtype=dtype, - mapping=mapping, - quant_config=quant_config, - **kwargs) - - if not use_preloading: - trust_remote_code = kwargs.pop('trust_remote_code', True) - - hf_model = transformers.AutoModelForCausalLM.from_pretrained( - hf_model_dir, dtype='auto', trust_remote_code=trust_remote_code) - weights = load_weights_from_hf_model(hf_model, config) - - model = GPTJForCausalLM(config) - model.load(weights) - return model diff --git a/tensorrt_llm/models/gptneox/__init__.py b/tensorrt_llm/models/gptneox/__init__.py deleted file mode 100644 index 71bf6d298d42..000000000000 --- a/tensorrt_llm/models/gptneox/__init__.py +++ /dev/null @@ -1,14 +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. diff --git a/tensorrt_llm/models/gptneox/model.py b/tensorrt_llm/models/gptneox/model.py deleted file mode 100644 index 0ac5d356336c..000000000000 --- a/tensorrt_llm/models/gptneox/model.py +++ /dev/null @@ -1,147 +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. - -from ..._utils import pad_vocab_size -from ...functional import PositionEmbeddingType, Tensor -from ...layers import (MLP, Attention, AttentionMaskType, ColumnLinear, - Embedding, LayerNorm) -from ...module import Module -from ..modeling_utils import (DecoderLayerList, DecoderModelForCausalLM, - PretrainedConfig) - - -class GPTNeoXDecoderLayer(Module): - - def __init__(self, config: PretrainedConfig, layer_idx: int): - super().__init__() - self.layer_idx = layer_idx - self.config = config - hidden_size = config.hidden_size - dtype = config.dtype - tp_group = config.mapping.tp_group - tp_size = config.mapping.tp_size - - self.input_layernorm = LayerNorm(normalized_shape=hidden_size, - dtype=dtype) - - self.post_attention_layernorm = LayerNorm(normalized_shape=hidden_size, - dtype=dtype) - - layers_range = config.mapping.pp_layers(config.num_hidden_layers) - local_layer_idx = layer_idx - layers_range[0] - self.attention = Attention( - local_layer_idx=local_layer_idx, - hidden_size=hidden_size, - num_attention_heads=config.num_attention_heads, - rotary_embedding_percentage=config.rotary_pct, - rotary_embedding_base=config.rotary_emb_base, - position_embedding_type=PositionEmbeddingType.rope_gpt_neox, - max_position_embeddings=config.max_position_embeddings, - dtype=dtype, - attention_mask_type=AttentionMaskType.causal, - bias=True, - tp_group=tp_group, - tp_size=tp_size, - quant_mode=config.quant_mode) - - self.mlp = MLP(hidden_size=hidden_size, - ffn_hidden_size=hidden_size * 4, - hidden_act=config.hidden_act, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - quant_mode=config.quant_mode) - - def forward(self, - hidden_states: Tensor, - attention_mask=None, - use_cache=False, - kv_cache_params=None, - attention_params=None): - residual = hidden_states - - input_layernorm_output = self.input_layernorm(hidden_states) - post_attention_layernorm_output = self.post_attention_layernorm( - hidden_states) - - attention_output = self.attention(input_layernorm_output, - attention_mask=attention_mask, - use_cache=use_cache, - kv_cache_params=kv_cache_params, - attention_params=attention_params, - norm_before_bmm1=True) - - if use_cache: - attention_output, presents = attention_output - - feed_forward_hidden_states = self.mlp(post_attention_layernorm_output) - hidden_states = attention_output + feed_forward_hidden_states + residual - if use_cache: - return (hidden_states, presents) - return hidden_states - - -class GPTNeoXModel(Module): - - def __init__(self, config: PretrainedConfig): - super().__init__() - self.vocab_embedding = Embedding(num_embeddings=config.vocab_size, - embedding_dim=config.hidden_size, - dtype=config.dtype) - - self.layers = DecoderLayerList(GPTNeoXDecoderLayer, config) - - self.ln_f = LayerNorm(normalized_shape=config.hidden_size, - dtype=config.dtype) - - def forward(self, - input_ids: Tensor, - position_ids=None, - use_cache=False, - attention_mask=None, - kv_cache_params=None, - attention_params=None): - hidden_states = self.vocab_embedding(input_ids) - - hidden_states = self.layers(hidden_states, - use_cache=use_cache, - attention_mask=attention_mask, - kv_cache_params=kv_cache_params, - attention_params=attention_params) - if use_cache: - hidden_states, presents = hidden_states - - hidden_states = self.ln_f(hidden_states) - - if use_cache: - return (hidden_states, tuple(presents)) - return hidden_states - - -class GPTNeoXForCausalLM(DecoderModelForCausalLM): - - def __init__(self, config: PretrainedConfig): - transformer = GPTNeoXModel(config) - vocab_size_padded = pad_vocab_size(config.vocab_size, - config.mapping.tp_size) - - lm_head = ColumnLinear(config.hidden_size, - vocab_size_padded, - bias=False, - dtype=config.dtype, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size, - gather_output=True) - super().__init__(config, transformer, lm_head) diff --git a/tensorrt_llm/models/grok/__init__.py b/tensorrt_llm/models/grok/__init__.py deleted file mode 100644 index 71bf6d298d42..000000000000 --- a/tensorrt_llm/models/grok/__init__.py +++ /dev/null @@ -1,14 +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. diff --git a/tensorrt_llm/models/grok/convert.py b/tensorrt_llm/models/grok/convert.py deleted file mode 100644 index c0ef4b630b9f..000000000000 --- a/tensorrt_llm/models/grok/convert.py +++ /dev/null @@ -1,518 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -import time -from pathlib import Path -from typing import Optional - -import jax -import numpy as np -import torch -from jax import dlpack as jax_dlpack -from torch.utils import dlpack as torch_dlpack - -from ..._utils import pad_vocab_size, release_gc -from ...layers import MoeConfig -from ...quantization import QuantAlgo -from ..convert_utils import split -from ..modeling_utils import PretrainedConfig, QuantConfig, optimize_model - - -def get_jax_weight(config, prefix, dtype, postfix='.weight', key_name='scale'): - - return torch.as_tensor((config[prefix + postfix][key_name])._value, - dtype=dtype).T - - -def get_jax_weight_scale_tp(params, key, rank): - jax_obj = params[key]['w'] - jax_scales = jax.device_put(jax_obj.scales, device=jax.devices('gpu')[rank]) - torch_scales = torch_dlpack.from_dlpack(jax_dlpack.to_dlpack(jax_scales)) - return torch.as_tensor( - np.asarray(jax_obj.weight.addressable_shards[rank].data)), torch_scales - - -def get_jax_weight_scale(params, key): - jax_obj = params[key]['w'] - - jax_scales = jax.device_put(jax_obj.scales, device=jax.devices('cpu')[0]) - - torch_scales = torch_dlpack.from_dlpack( - jax_dlpack.to_dlpack(jax_scales, copy=False)) - return torch.as_tensor(np.asarray(jax_obj.weight), - dtype=torch.int8), torch_scales - - -def get_tllm_linear_weight( - weight, - torch_weight_scales, - prefix, - plugin_weight_only_quant_type=torch.int8, - postfix='weight', -): - results = {} - processed_weight = torch.ops.trtllm.preprocess_weights_for_mixed_gemm( - weight if weight.is_contiguous() else weight.contiguous(), - plugin_weight_only_quant_type, torch.bfloat16) - results[prefix + postfix] = processed_weight - - results[prefix + 'per_channel_scale'] = torch_weight_scales.contiguous() - - return results - - -def convert_grok(hf_model, - config, - mapping, - vocab_size=32000, - dtype='float32', - use_parallel_embedding=False, - sharding_dim=0, - use_weight_only=False, - use_gemm_woq_plugin=False, - plugin_weight_only_quant_type=torch.int8, - moe_config=None): - - weights = {} - tik = time.time() - tensor_parallel = mapping.tp_size - model_params = hf_model - dtype = getattr(torch, dtype) - - config['num_attention_heads'] - config['hidden_size'] - - layers_range = mapping.pp_layers(config['num_hidden_layers']) - - def convert_layer(l): - prefix = f'transformer/decoder_layer_{l}/' - print(prefix) - tllm_prex = f'transformer.layers.{l - layers_range[0]}.' - - wq, q_scale = get_jax_weight_scale_tp( - model_params, prefix + 'multi_head_attention/query', - mapping.tp_rank) - wk, k_scale = get_jax_weight_scale_tp( - model_params, prefix + 'multi_head_attention/key', mapping.tp_rank) - wv, v_scale = get_jax_weight_scale_tp( - model_params, prefix + 'multi_head_attention/value', - mapping.tp_rank) - - qs = split(q_scale, mapping.tp_size, mapping.tp_rank, dim=1) - ks = split(k_scale, mapping.tp_size, mapping.tp_rank, dim=1) - vs = split(v_scale, mapping.tp_size, mapping.tp_rank, dim=1) - split_v = torch.concat((wq, wk, wv), dim=1) - scale_v = torch.concat((qs, ks, vs), dim=1) - - weights.update( - get_tllm_linear_weight(split_v, scale_v.squeeze(), - tllm_prex + 'attention.qkv.', - plugin_weight_only_quant_type)) - - attn_dense_weight, attn_dense_scales = get_jax_weight_scale_tp( - model_params, prefix + 'multi_head_attention/linear', - mapping.tp_rank) - - split_scales = split(attn_dense_scales, - tensor_parallel, - mapping.tp_rank, - dim=0) - - weights.update( - get_tllm_linear_weight(attn_dense_weight, split_scales.squeeze(), - tllm_prex + 'attention.dense.', - plugin_weight_only_quant_type)) - if mapping.moe_ep_size > 1: - w3, s3 = get_jax_weight_scale( - model_params, f'transformer/decoder_layer_{l}/moe/linear_v') - - w2, s2 = get_jax_weight_scale( - model_params, f'transformer/decoder_layer_{l}/moe/linear_1') - - w1, s1 = get_jax_weight_scale( - model_params, f'transformer/decoder_layer_{l}/moe/linear') - - # moe expert parallel - w3_split = split(w3, - mapping.moe_ep_size, - mapping.moe_ep_rank, - dim=0) - w2_split = split(w2, - mapping.moe_ep_size, - mapping.moe_ep_rank, - dim=0) - w1_split = split(w1, - mapping.moe_ep_size, - mapping.moe_ep_rank, - dim=0) - - s3_split = split(s3, - mapping.moe_ep_size, - mapping.moe_ep_rank, - dim=0) - s2_split = split(s2, - mapping.moe_ep_size, - mapping.moe_ep_rank, - dim=0) - s1_split = split(s1, - mapping.moe_ep_size, - mapping.moe_ep_rank, - dim=0) - - # moe tensor parallel - w3_split = split(w3_split, - mapping.moe_tp_size, - mapping.moe_tp_rank, - dim=2) - w2_split = split(w2_split, - mapping.moe_tp_size, - mapping.moe_tp_rank, - dim=1) - w1_split = split(w1_split, - mapping.moe_tp_size, - mapping.moe_tp_rank, - dim=2) - - s3_split = split(s3_split, - mapping.moe_tp_size, - mapping.moe_tp_rank, - dim=2) - s2_split = split(s2_split, - mapping.moe_tp_size, - mapping.moe_tp_rank, - dim=1) - s1_split = split(s1_split, - mapping.moe_tp_size, - mapping.moe_tp_rank, - dim=2) - else: - w3_split, s3 = get_jax_weight_scale_tp( - model_params, f'transformer/decoder_layer_{l}/moe/linear_v', - mapping.tp_rank) - - w2_split, s2 = get_jax_weight_scale_tp( - model_params, f'transformer/decoder_layer_{l}/moe/linear_1', - mapping.tp_rank) - - w1_split, s1 = get_jax_weight_scale_tp( - model_params, f'transformer/decoder_layer_{l}/moe/linear', - mapping.tp_rank) - - s3_split = split(s3, - mapping.moe_tp_size, - mapping.moe_tp_rank, - dim=2) - s2_split = split(s2, - mapping.moe_tp_size, - mapping.moe_tp_rank, - dim=1) - s1_split = split(s1, - mapping.moe_tp_size, - mapping.moe_tp_rank, - dim=2) - weights.update( - get_tllm_linear_weight(w2_split, - s2_split.reshape(moe_config.num_experts, -1), - tllm_prex + 'mlp.proj.', - plugin_weight_only_quant_type)) - - weights.update( - get_tllm_linear_weight( - torch.concat([w3_split, w1_split], dim=-1), - torch.concat([s3_split, s1_split], - dim=-1).reshape(moe_config.num_experts, -1), - tllm_prex + 'mlp.fc.', - plugin_weight_only_quant_type, - )) - - moe_experts_gate_weights = get_jax_weight(model_params, - prefix + 'router', - torch.float32, - postfix='', - key_name='w').contiguous() - - weights[tllm_prex + 'mlp.router.weight'] = moe_experts_gate_weights - - # Layer norms do not use tensor parallelism - input_ln_weight = get_jax_weight(model_params, - prefix + 'rms_norm', - dtype, - postfix='') - weights[tllm_prex + 'input_layernorm.weight'] = input_ln_weight - - post_attn_weight = get_jax_weight(model_params, - prefix + 'rms_norm_1', - dtype, - postfix='') - weights[tllm_prex + 'post_attn_layernorm.weight'] = post_attn_weight - - post_ln_weight = get_jax_weight(model_params, - prefix + 'rms_norm_2', - dtype, - postfix='') - weights[tllm_prex + 'post_layernorm.weight'] = post_ln_weight - - post_mlp_weight = get_jax_weight(model_params, - prefix + 'rms_norm_3', - dtype, - postfix='') - weights[tllm_prex + 'post_mlp_layernorm.weight'] = post_mlp_weight - - for l in layers_range: - convert_layer(l) - release_gc() - - v = get_jax_weight(model_params, - 'language_model/in_out_embed', - dtype, - postfix='', - key_name='embeddings').T - tie_word_embeddings = config['tie_word_embeddings'] - if tie_word_embeddings: - # lm_head.weight has the same weights as embedding - if mapping.is_last_pp_rank(): - if vocab_size % mapping.tp_size != 0: - # padding - vocab_size_padded = pad_vocab_size(vocab_size, mapping.tp_size) - pad_width = vocab_size_padded - vocab_size - - v = torch.nn.functional.pad(v, (0, pad_width, 0, 0), 'constant', - 0) - weights['lm_head.weight'] = split(v, mapping.tp_size, - mapping.tp_rank) - - if use_parallel_embedding: - v = split(v, mapping.tp_size, mapping.tp_rank, dim=sharding_dim) - - if mapping.is_first_pp_rank(): - weights['transformer.vocab_embedding.weight'] = v - - ln_f_w = get_jax_weight(model_params, - 'language_model/rms_norm', - dtype, - postfix='') - weights['transformer.ln_f.weight'] = ln_f_w - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Weights loaded. Total time: {t}') - return weights - - -def create_config_from_xai(dtype, - mapping, - quantization: QuantConfig = None, - override_fields: dict = {}): - config = {} - hf_config = { - "architectures": ["Grok1ModelForCausalLM"], - "vocab_size": 131072, - "hidden_size": 6144, - "intermediate_size": 32768, - "num_hidden_layers": 64, - "num_attention_heads": 48, - "num_key_value_heads": 8, - "attn_output_multiplier": 0.08838834764831845, - "embedding_multiplier_scale": 78.38367176906169, - "output_multiplier_scale": 0.5773502691896257, - "max_attn_value": 30.0, - "max_position_embeddings": 8192, - "rms_norm_eps": 1e-5, - "use_cache": True, - "pad_token_id": 0, - "bos_token_id": 1, - "eos_token_id": 2, - "tie_word_embeddings": True, - "num_experts_per_tok": 2, - "num_experts": 8, - "output_router_logits": False, - "router_aux_loss_coef": 0.001, - "torch_dtype": "bfloat16", - "transformers_version": "4.35.0" - } - # same for from_meta and from_cli_args - n_head = hf_config['num_attention_heads'] - inter_size = hf_config['intermediate_size'] - n_layer = hf_config['num_hidden_layers'] - n_embd = hf_config['hidden_size'] - n_kv_head = hf_config['num_key_value_heads'] - rms_norm_eps = hf_config['rms_norm_eps'] - vocab_size = hf_config['vocab_size'] - n_positions = hf_config['max_position_embeddings'] - hidden_act = 'geglu' - config['rotary_scaling'] = None - rotary_base = 10000.0 - - config[ - 'moe_normalization_mode'] = MoeConfig.ExpertScaleNormalizationMode.NONE - - moe_num_experts = hf_config['num_experts'] - - moe_top_k = hf_config['num_experts_per_tok'] - - attn_output_multiplier = hf_config['attn_output_multiplier'] - embedding_multiplier_scale = hf_config['embedding_multiplier_scale'] - - output_multiplier_scale = hf_config['output_multiplier_scale'] - max_attn_value = hf_config['max_attn_value'] - - architecture = hf_config['architectures'][0] - - attn_bias = False - - config.update({ - 'architecture': - architecture, - 'dtype': - dtype, - 'logits_dtype': - 'float32', - 'num_hidden_layers': - n_layer, - 'num_attention_heads': - n_head, - 'hidden_size': - n_embd, - 'intermediate_size': - inter_size, - 'num_key_value_heads': - n_kv_head, - 'vocab_size': - vocab_size, - 'position_embedding_type': - 'rope_gpt_neox', - 'max_position_embeddings': - n_positions, - 'hidden_act': - hidden_act, - 'rotary_base': - rotary_base, - 'norm_epsilon': - rms_norm_eps, - 'moe_num_experts': - moe_num_experts, - 'moe_top_k': - moe_top_k, - 'moe_normalization_mode': - MoeConfig.ExpertScaleNormalizationMode.NONE, - #TODO: should have directly map from the Mapping object to the TRT-LLM checkpoint fields - 'mapping': { - 'world_size': mapping.tp_size * mapping.pp_size, - 'tp_size': mapping.tp_size, - 'pp_size': mapping.pp_size, - 'moe_tp_size': mapping.moe_tp_size, - 'moe_ep_size': mapping.moe_ep_size, - }, - 'attn_bias': - attn_bias, - "attn_output_multiplier": - attn_output_multiplier, - "embedding_multiplier_scale": - embedding_multiplier_scale, - "output_multiplier_scale": - output_multiplier_scale, - "max_attn_value": - max_attn_value, - "tie_word_embeddings": - True, - }) - - config['quantization'] = quantization.model_dump() - config.update(override_fields) - - return config - - -def from_hugging_face(cls, - model_dir, - dtype, - *, - mapping, - quantization: QuantConfig = None, - override_fields={}, - skip_loading_weights=False, - preloaded_model=None): - ''' Create a LLaMAForCausalLM object from give parameters - ''' - assert model_dir is not None - if isinstance(model_dir, Path): # some code relies on this as string - model_dir = str(model_dir) - - config = create_config_from_xai(dtype, - mapping, - quantization, - override_fields=override_fields) - - pretrained_config = PretrainedConfig.from_dict(config) - pretrained_config.set_rank(mapping.rank) # TODO:remove this hack - - grok = cls.from_config(pretrained_config) - grok = optimize_model( - grok, - use_parallel_embedding=pretrained_config.use_parallel_embedding, - ) - - if skip_loading_weights: - return grok - - weights = load_weights_from_xai(config=config, - mapping=mapping, - model=preloaded_model) - - grok.load(weights) - return grok - - -def quantize(dtype, - model_dir, - output_dir, - mapping, - quantization: QuantConfig, - *, - override_fields, - dataset_cache_dir: Optional[str] = None): - ''' - Quantize the save the model as TRT-LLM checkpoint to output_dir - ''' - pass #The official grok-1 model is published under int8 wo format, we don't need to quantize again. - - -def load_weights_from_xai(*, config, mapping, model): - assert model is not None - plugin_weight_only_quant_type = None # the value does not matter when use_weight_only is False - quant_algo = config['quantization']['quant_algo'] - assert quant_algo == QuantAlgo.W8A16 - plugin_weight_only_quant_type = torch.int8 - - moe_config = MoeConfig( - num_experts=config['moe_num_experts'], - top_k=config['moe_top_k'], - normalization_mode=config['moe_normalization_mode']).validate() - - use_weight_only = quant_algo in [QuantAlgo.W8A16] - - weights = convert_grok( - model, - config, - mapping, - vocab_size=config['vocab_size'], - dtype=config['dtype'], - use_weight_only=use_weight_only, - use_gemm_woq_plugin=not config.get('disable_weight_only_quant_plugin', - False), - plugin_weight_only_quant_type=plugin_weight_only_quant_type, - use_parallel_embedding=config.get('use_parallel_embedding', False), - sharding_dim=config.get('embedding_sharding_dim', 0), - moe_config=moe_config) - return weights diff --git a/tensorrt_llm/models/grok/model.py b/tensorrt_llm/models/grok/model.py deleted file mode 100644 index 9ff22cd71c71..000000000000 --- a/tensorrt_llm/models/grok/model.py +++ /dev/null @@ -1,273 +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. -from typing import Optional - -from ..._utils import pad_vocab_size -from ...functional import Tensor, recv, send -from ...layers import (MOE, Attention, AttentionMaskType, ColumnLinear, - Embedding, MoeConfig, PositionEmbeddingType, RmsNorm) -from ...lora_helper import LoraConfig, use_lora -from ...mapping import Mapping -from ...module import Module -from ..modeling_utils import (DecoderLayerList, DecoderModelForCausalLM, - PretrainedConfig, QuantConfig) - - -class GrokDecoderLayer(Module): - - def __init__(self, config: PretrainedConfig, layer_idx: int): - super().__init__() - self.layer_idx = layer_idx - self.config = config - - self.input_layernorm = RmsNorm(normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - dtype=config.dtype) - - layers_range = config.mapping.pp_layers(config.num_hidden_layers) - local_layer_idx = layer_idx - layers_range[0] - self.attention = Attention( - local_layer_idx=local_layer_idx, - hidden_size=config.hidden_size, - attention_head_size=config.head_size, - num_attention_heads=config.num_attention_heads, - num_kv_heads=config.num_key_value_heads, - max_position_embeddings=config.max_position_embeddings, - dtype=config.dtype, - attention_mask_type=AttentionMaskType.causal, - bias=config.attn_bias, - position_embedding_type=PositionEmbeddingType.rope_gpt_neox, - rotary_embedding_base=config.rotary_base, - rotary_embedding_scaling=config.rotary_scaling, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size, - tp_rank=config.mapping.tp_rank, - quant_mode=config.quant_mode, - max_attn_value=config.max_attn_value) - - mlp_hidden_size = config.hidden_size * 4 if config.intermediate_size is None else config.intermediate_size - self.post_attn_layernorm = RmsNorm(normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - dtype=config.dtype) - - self.post_mlp_layernorm = RmsNorm(normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - dtype=config.dtype) - mlp_kwargs = {} - assert config.moe_num_experts > 1, "Grok model is a MoE model." - ClsMLP = MOE - moe_config = MoeConfig( - num_experts=config.moe_num_experts, - top_k=config.moe_top_k, - normalization_mode=config.moe_normalization_mode).validate() - mlp_kwargs = { - "moe_config": moe_config, - "mapping": config.mapping, - } - self.mlp = ClsMLP(hidden_size=config.hidden_size, - ffn_hidden_size=mlp_hidden_size, - hidden_act=config.hidden_act, - dtype=config.dtype, - bias=config.mlp_bias, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size, - quant_mode=config.quant_mode, - **mlp_kwargs) - self.post_layernorm = RmsNorm(normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - dtype=config.dtype) - - def forward(self, - hidden_states, - attention_mask=None, - use_cache=False, - spec_decoding_params=None, - kv_cache_params=None, - attention_params=None, - lora_layer_params=None): - residual = hidden_states - hidden_states = self.input_layernorm(hidden_states) - - attention_output = self.attention( - hidden_states, - attention_mask=attention_mask, - use_cache=use_cache, - spec_decoding_params=spec_decoding_params, - kv_cache_params=kv_cache_params, - attention_params=attention_params, - lora_layer_params=lora_layer_params) - - if use_cache: - attention_output, presents = attention_output - - attention_output = self.post_attn_layernorm(attention_output) - hidden_states = residual + attention_output - - residual_attn = hidden_states - - # regular llama/mixtral layers - hidden_states = self.post_layernorm(hidden_states) - hidden_states = self.mlp(hidden_states, - lora_layer_params=lora_layer_params) - hidden_states = self.post_mlp_layernorm(hidden_states) - hidden_states = residual_attn + hidden_states - - if use_cache: - return (hidden_states, presents) - return hidden_states - - -class GrokModel(Module): - - def __init__(self, config: PretrainedConfig) -> None: - super().__init__() - - self.mapping = config.mapping - if self.mapping.is_first_pp_rank(): - self.vocab_embedding = Embedding(config.vocab_size, - config.hidden_size, - dtype=config.dtype) - - self.layers = DecoderLayerList(GrokDecoderLayer, config) - - self.embedding_multiplier_scale = config.embedding_multiplier_scale - - if self.mapping.is_last_pp_rank(): - self.ln_f = RmsNorm(normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - dtype=config.dtype) - - def forward(self, - input_ids, - position_ids=None, - use_cache=False, - attention_mask=None, - spec_decoding_params=None, - kv_cache_params=None, - attention_params=None, - hidden_states=None, - prompt_embedding_table: Optional[Tensor] = None, - prompt_tasks: Optional[Tensor] = None, - prompt_vocab_size: Optional[Tensor] = None, - lora_params=None): - - ptuning_args = [ - prompt_embedding_table, prompt_tasks, prompt_vocab_size - ] if prompt_embedding_table is not None else [] - - if self.mapping.is_first_pp_rank(): - hidden_states = self.vocab_embedding(input_ids, *ptuning_args) - hidden_states *= 78.38367176906169 - else: - hidden_states = recv(hidden_states, self.mapping.prev_pp_rank()) - - hidden_states = self.layers.forward( - hidden_states, - use_cache=use_cache, - attention_mask=attention_mask, - kv_cache_params=kv_cache_params, - attention_params=attention_params, - lora_params=lora_params, - spec_decoding_params=spec_decoding_params) - - if use_cache: - hidden_states, presents = hidden_states - - if self.mapping.is_last_pp_rank(): - hidden_states = self.ln_f(hidden_states) - else: - hidden_states = send(hidden_states, self.mapping.next_pp_rank()) - - if use_cache: - return (hidden_states, tuple(presents)) - return hidden_states - - -class GrokForCausalLM(DecoderModelForCausalLM): - - def __init__(self, config: PretrainedConfig): - self.check_config(config) - transformer = GrokModel(config) - vocab_size_padded = pad_vocab_size(config.vocab_size, - config.mapping.tp_size) - if config.mapping.is_last_pp_rank(): - lm_head = ColumnLinear(config.hidden_size, - vocab_size_padded, - bias=False, - dtype=config.dtype, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size, - gather_output=True) - else: - lm_head = None - self.quant_mode = config.quant_mode - self.mapping = config.mapping - super().__init__(config, transformer, lm_head) - - def check_config(self, config): - config.set_if_not_exist('mlp_bias', False) - config.set_if_not_exist('attn_bias', False) - config.set_if_not_exist('rotary_base', 10000.0) - config.set_if_not_exist('rotary_scaling', None) - config.set_if_not_exist('moe_num_experts', 0) - config.set_if_not_exist('moe_top_k', 0) - config.set_if_not_exist('moe_normalization_mode', - MoeConfig.ExpertScaleNormalizationMode.NONE) - - @classmethod - def from_hugging_face(cls, - hf_model_dir, - dtype='float16', - mapping: Optional[Mapping] = None, - **kwargs): - from . import convert - if mapping is None: - mapping = Mapping() - grok = convert.from_hugging_face( - cls, - hf_model_dir, - dtype, - mapping=mapping, - quantization=kwargs.get('quantization', QuantConfig()), - override_fields=kwargs.get('override_fields', {}), - skip_loading_weights=kwargs.get('skip_loading_weights', False), - preloaded_model=kwargs.get('preloaded_model', None)) - return grok - - def default_plugin_config(self, **kwargs): - plugin_config = super().default_plugin_config(**kwargs) - if self.quant_mode.is_int4_weight_only_per_group(): - plugin_config.set_weight_only_groupwise_quant_matmul_plugin() - return plugin_config - - @classmethod - def quantize( - cls, - hf_model_dir, - output_dir, - quant_config: QuantConfig, - *, - dtype='float16', - mapping: Optional[Mapping] = None, - calib_batches=512, - calib_batch_size=1, - random_seed=1234, - tokenizer_max_seq_length=2048, - **kwargs, - ): - pass - - def use_lora(self, lora_config: LoraConfig): - use_lora(self, lora_config) diff --git a/tensorrt_llm/models/grok/weight.py b/tensorrt_llm/models/grok/weight.py deleted file mode 100644 index 0c412305a388..000000000000 --- a/tensorrt_llm/models/grok/weight.py +++ /dev/null @@ -1,35 +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. -from typing import Union - -import numpy as np -import torch - - -def split(v: Union[np.ndarray, torch.Tensor], - tp_size: int, - tp_rank: int, - dim=0): - if tp_size == 1: - return v - assert len(v.shape) > 1 or dim == 0 - if isinstance(v, np.ndarray): - return np.ascontiguousarray( - np.split(v, tp_size, axis=dim)[tp_rank].copy()) - else: - assert v.shape[dim] % tp_size == 0, \ - 'Unable to split: shape={v.shape} (dim={dim}) tp_size={tp_size}.' - split_size = v.shape[dim] // tp_size - return v.split(split_size, dim=dim)[tp_rank].clone().detach() diff --git a/tensorrt_llm/models/llama/__init__.py b/tensorrt_llm/models/llama/__init__.py deleted file mode 100644 index 71bf6d298d42..000000000000 --- a/tensorrt_llm/models/llama/__init__.py +++ /dev/null @@ -1,14 +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. diff --git a/tensorrt_llm/models/llama/config.py b/tensorrt_llm/models/llama/config.py deleted file mode 100644 index 54038e32c4ff..000000000000 --- a/tensorrt_llm/models/llama/config.py +++ /dev/null @@ -1,280 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -import json -import math -import sys -from pathlib import Path -from typing import Optional, Union - -from ..._utils import get_hf_rope_theta -from ...layers import MoeConfig -from ...mapping import Mapping -from ..convert_utils import infer_dtype -from ..modeling_utils import PretrainedConfig, QuantConfig - - -class LLaMAConfig(PretrainedConfig): - - def __init__(self, - *, - mlp_bias: bool = False, - attn_bias: bool = False, - rotary_base: float = 10000.0, - rotary_scaling: Optional[dict] = None, - residual_mlp: bool = False, - disable_weight_only_quant_plugin: bool = False, - moe: Optional[Union[MoeConfig, dict]] = None, - remove_duplicated_kv_heads: bool = False, - embedding_multiplier: float = 1.0, - attention_multiplier: float = 1.0, - residual_multiplier: float = 1.0, - output_multiplier_scale: float = 1.0, - **kwargs): - self.mlp_bias = mlp_bias - self.attn_bias = attn_bias - self.rotary_base = rotary_base - self.rotary_scaling = rotary_scaling - self.residual_mlp = residual_mlp - self.disable_weight_only_quant_plugin = disable_weight_only_quant_plugin - if moe is None: - # Legacy MOE config fields - moe = MoeConfig( - num_experts=kwargs.pop('moe_num_experts', 0), - top_k=kwargs.pop('moe_top_k', 0), - normalization_mode=kwargs.pop( - 'moe_normalization_mode', - MoeConfig.ExpertScaleNormalizationMode.RENORMALIZE)) - elif isinstance(moe, dict): - moe = MoeConfig.from_dict(moe) - assert isinstance(moe, MoeConfig) - self.moe = moe.validate() - self.remove_duplicated_kv_heads = remove_duplicated_kv_heads - self.fc_after_embed = False - self.use_input_layernorm_in_first_layer = True - self.use_last_layernorm = True - self.layer_idx_offset = 0 - self.embedding_multiplier = embedding_multiplier - self.attention_multiplier = attention_multiplier - self.residual_multiplier = residual_multiplier - self.output_multiplier_scale = output_multiplier_scale - self.has_partial_lora_mask = False - - super().__init__(**kwargs) - - def to_dict(self): - output = super().to_dict() - # Serialize the fields added in LLaMAConfig - output['mlp_bias'] = self.mlp_bias - output['attn_bias'] = self.attn_bias - output['rotary_base'] = self.rotary_base - output['rotary_scaling'] = self.rotary_scaling - output['residual_mlp'] = self.residual_mlp - output[ - 'disable_weight_only_quant_plugin'] = self.disable_weight_only_quant_plugin - output['fc_after_embed'] = self.fc_after_embed - output[ - 'use_input_layernorm_in_first_layer'] = self.use_input_layernorm_in_first_layer - output['use_last_layernorm'] = self.use_last_layernorm - output['layer_idx_offset'] = self.layer_idx_offset - output['moe'] = self.moe.to_dict() - return output - - @classmethod - def from_hugging_face( - cls, - hf_config_or_dir: Union[str, 'transformers.PretrainedConfig'], - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - import transformers - - trust_remote_code = kwargs.pop('trust_remote_code', True) - has_partial_lora_mask = False - - if isinstance(hf_config_or_dir, transformers.PretrainedConfig): - hf_config = hf_config_or_dir - else: - hf_config_dir = str(hf_config_or_dir) - if "vila" in hf_config_dir.lower(): - sys.path.append(hf_config_dir + "/../VILA") - from llava.model import LlavaLlamaConfig # noqa - from llava.model import LlavaLlamaModel - transformers.AutoConfig.register("llava_llama", - LlavaLlamaConfig, - exist_ok=True) - transformers.AutoModelForCausalLM.register( - LlavaLlamaConfig, LlavaLlamaModel) - - hf_config = transformers.AutoConfig.from_pretrained( - hf_config_dir, trust_remote_code=trust_remote_code) - if hf_config.model_type == "llava": - # LLaVA = Vision model + Llama LLM - # We load a llava config and use its' text config as llama config - from transformers import LlavaConfig - hf_config = LlavaConfig.from_pretrained( - hf_config_dir).text_config - if hf_config.model_type == "llava_next": - from transformers import LlavaNextConfig - hf_config = LlavaNextConfig.from_pretrained( - hf_config_dir).text_config - if hf_config.model_type == "llava_llama": - hf_config.llm_cfg["architecture"] = hf_config.llm_cfg[ - "architectures"][0] - hf_config.llm_cfg["dtype"] = hf_config.llm_cfg["torch_dtype"] - hf_config = PretrainedConfig.from_dict(hf_config.llm_cfg) - if hf_config.model_type == 'internlmxcomposer2': - # InternLM-XComposer2 has a mask for partial lora - # Therefore we need an additional flag for this mask - has_partial_lora_mask = True - if hf_config.model_type == 'mistral3': - from transformers import Mistral3Config - hf_config = Mistral3Config.from_pretrained( - hf_config_dir).text_config - hf_config.architectures = ["MistralForCausalLM"] - - num_key_value_heads = getattr(hf_config, "num_key_value_heads", - hf_config.num_attention_heads) - if hf_config.model_type == "exaone": - hidden_act = hf_config.activation_function - # NOTE - # EXAONE also uses RMS norm but they represent as layer_norm_epsilon. - norm_epsilon = getattr(hf_config, "layer_norm_epsilon", 1e-5) - else: - hidden_act = hf_config.hidden_act - norm_epsilon = hf_config.rms_norm_eps - head_dim = getattr( - hf_config, "head_dim", - hf_config.hidden_size // hf_config.num_attention_heads) - head_size = getattr(hf_config, "kv_channels", head_dim) - attn_bias = getattr(hf_config, 'bias', False) or getattr( - hf_config, 'attention_bias', False) - rotary_scaling = getattr(hf_config, "rope_scaling", None) - rotary_base = get_hf_rope_theta(hf_config, 10000.0) - residual_mlp = getattr(hf_config, "parallel_attn_mlp_res", False) - disable_weight_only_quant_plugin = kwargs.pop( - 'disable_weight_only_quant_plugin', False) - remove_duplicated_kv_heads = kwargs.pop('remove_duplicated_kv_heads', - False) - embedding_multiplier = getattr(hf_config, "embedding_multiplier", 1.0) - attention_multiplier = getattr(hf_config, "attention_multiplier", 1.0) - if attention_multiplier != 1.0: - attention_multiplier *= math.sqrt(head_size) - residual_multiplier = getattr(hf_config, "residual_multiplier", 1.0) - output_multiplier_scale = 1.0 / getattr(hf_config, "logits_scaling", - 1.0) - if hf_config.model_type in ["mixtral", "arctic", "granitemoe"]: - # HF LLaMA-type models are implicitly using gated activation. - # With our MoE implementation, we must make it explicit - hidden_act = "swiglu" - moe_normalization_mode = MoeConfig.ExpertScaleNormalizationMode.RENORMALIZE - else: - moe_normalization_mode = None - moe_num_experts = getattr(hf_config, "num_local_experts", 0) - moe_top_k = getattr(hf_config, "num_experts_per_tok", 0) - moe_config = MoeConfig(num_experts=moe_num_experts, - top_k=moe_top_k, - normalization_mode=moe_normalization_mode) - moe_config.validate() - - dtype = infer_dtype(dtype, getattr(hf_config, 'torch_dtype', None)) - tie_word_embeddings = getattr(hf_config, 'tie_word_embeddings', False) - - return cls( - architecture=hf_config.architectures[0], - dtype=dtype, - num_hidden_layers=hf_config.num_hidden_layers, - num_attention_heads=hf_config.num_attention_heads, - hidden_size=hf_config.hidden_size, - intermediate_size=hf_config.intermediate_size, - num_key_value_heads=num_key_value_heads, - head_size=head_size, - vocab_size=hf_config.vocab_size, - position_embedding_type='rope_gpt_neox', - max_position_embeddings=hf_config.max_position_embeddings, - hidden_act=hidden_act, - norm_epsilon=norm_epsilon, - attn_bias=attn_bias, - rotary_base=rotary_base, - rotary_scaling=rotary_scaling, - residual_mlp=residual_mlp, - disable_weight_only_quant_plugin=disable_weight_only_quant_plugin, - moe=moe_config, - mapping=mapping, - quantization=quant_config, - has_partial_lora_mask=has_partial_lora_mask, - remove_duplicated_kv_heads=remove_duplicated_kv_heads, - tie_word_embeddings=tie_word_embeddings, - embedding_multiplier=embedding_multiplier, - attention_multiplier=attention_multiplier, - residual_multiplier=residual_multiplier, - output_multiplier_scale=output_multiplier_scale, - **kwargs) - - @classmethod - def from_meta_ckpt(cls, - meta_ckpt_dir: str, - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - - with open(Path(meta_ckpt_dir, "params.json")) as fp: - meta_config: dict = json.load(fp) - - n_embd = meta_config["dim"] - n_head = meta_config["n_heads"] - n_kv_head = meta_config.get("n_kv_heads", n_head) - vocab_size = meta_config.get("vocab_size", 32000) - - # Reset vocab_size to 32000 for LLama v2 checkpoint. - if vocab_size == -1: - vocab_size = 32000 - - if "hidden_dim" in meta_config: - inter_size = meta_config["hidden_dim"] - else: - multiple_of = meta_config.get("multiple_of", 1) - n_embd_ = int(4 * n_embd * 2 / 3) - ffn_dim_multiplier = meta_config.get("ffn_dim_multiplier", 1) - inter_size = multiple_of * ( - (int(n_embd_ * ffn_dim_multiplier) + multiple_of - 1) // - multiple_of) - - dtype = infer_dtype(dtype, 'bfloat16') - - if meta_config.get('use_scaled_rope'): - rotary_scaling = {"type": "llama3"} - else: - rotary_scaling = meta_config.get("rope_scaling") - - # meta checkpoint don't have vocab_size|hidden_act|rotary_base specified, use same default value as HF - return cls(architecture="LlamaForCausalLM", - dtype=dtype, - num_hidden_layers=meta_config["n_layers"], - num_attention_heads=n_head, - hidden_size=n_embd, - intermediate_size=inter_size, - num_key_value_heads=n_kv_head, - vocab_size=vocab_size, - position_embedding_type='rope_gpt_neox', - max_position_embeddings=2048, - hidden_act='silu', - rotary_scaling=rotary_scaling, - rotary_base=meta_config.get('rope_theta', 10000), - norm_epsilon=meta_config["norm_eps"], - mapping=mapping, - quantization=quant_config, - **kwargs) diff --git a/tensorrt_llm/models/llama/convert.py b/tensorrt_llm/models/llama/convert.py deleted file mode 100644 index c36526f8283c..000000000000 --- a/tensorrt_llm/models/llama/convert.py +++ /dev/null @@ -1,2449 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -import copy -import functools -import os -import sys -import time -from collections import defaultdict -from pathlib import Path -from typing import List, Optional - -import numpy as np -import safetensors -import torch -import torch.nn as nn -from tqdm import tqdm -from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer -from transformers.models.llama.modeling_llama import LlamaDecoderLayer -from transformers.pytorch_utils import Conv1D - -from ..._utils import pad_vocab_size, release_gc, str_dtype_to_torch -from ...logger import logger -from ...quantization import QuantAlgo -from ...quantization.quantize import (qserve_pack_reorder_per_channel, - qserve_pack_reorder_per_group, - qserve_quantize_weight_per_channel, - qserve_quantize_weight_per_group) -from ..convert_utils import (dup_kv_weight, generate_int8, - get_tllm_linear_weight, iterate_shard_files, - load_calib_dataset, load_state_dict, - retrieved_layer_index_from_name, smooth_gemm, - smooth_gemm_fc1_gate, split, split_matrix_tp, - split_qkv_bias_tp, split_qkv_tp) -from ..modeling_utils import PretrainedConfig -from .config import LLaMAConfig - - -@torch.no_grad() -def smooth_llama_model(model, scales, alpha, llama_qkv_para, llama_smoother): - # Smooth the activation and weights with smoother = $\diag{s}$ - for name, module in model.named_modules(): - if not isinstance( - module, - LlamaDecoderLayer) and module.__class__.__name__ not in [ - "InternLMDecoderLayer", "MistralDecoderLayer" - ]: - continue - # qkv_proj - layer_name_q = name + ".self_attn.q_proj" - layer_name_k = name + ".self_attn.k_proj" - layer_name_v = name + ".self_attn.v_proj" - layer_name_qkv = name + ".self_attn.qkv_proj" - - weight = torch.cat([ - module.self_attn.q_proj.weight, module.self_attn.k_proj.weight, - module.self_attn.v_proj.weight - ], - dim=0) - - smoother = smooth_gemm(weight, scales[layer_name_q]["x"], - module.input_layernorm.weight, None, alpha) - - scales[layer_name_qkv]["x"] = scales[layer_name_q]["x"] / smoother - scales[layer_name_qkv]["w"] = weight.abs().max(dim=1)[0] - scales[layer_name_qkv]["y"] = torch.cat([ - scales[layer_name_q]["y"], scales[layer_name_k]["y"], - scales[layer_name_v]["y"] - ], - dim=0) - - # see transpose_weights function - llama_qkv_para[layer_name_qkv] = weight.transpose(0, 1) - - # ================================================================= - layer_name = name + ".self_attn.o_proj" - smoother = smooth_gemm(module.self_attn.o_proj.weight, - scales[layer_name]["x"], None, None, alpha) - llama_smoother[layer_name] = smoother.float() - - scales[layer_name]["x"] = scales[layer_name]["x"] / smoother - scales[layer_name]["w"] = module.self_attn.o_proj.weight.abs().max( - dim=1)[0] - - # ================================================================== - fc1_layer_name = name + ".mlp.gate_proj" - gate_layer_name = name + ".mlp.up_proj" - - smoother = smooth_gemm_fc1_gate(module.mlp.gate_proj.weight, - module.mlp.up_proj.weight, - scales[fc1_layer_name]["x"], - module.post_attention_layernorm.weight, - None, alpha) - - scales[fc1_layer_name]["x"] = scales[fc1_layer_name]["x"] / smoother - scales[fc1_layer_name]["w"] = module.mlp.gate_proj.weight.abs().max( - dim=1)[0] - - scales[gate_layer_name]["x"] = scales[gate_layer_name]["x"] / smoother - scales[gate_layer_name]["w"] = module.mlp.up_proj.weight.abs().max( - dim=1)[0] - - # ================================================================== - layer_name = name + ".mlp.down_proj" - smoother = smooth_gemm(module.mlp.down_proj.weight, - scales[layer_name]["x"], None, None, alpha) - llama_smoother[layer_name] = smoother.float() - scales[layer_name]["x"] = scales[layer_name]["x"] / smoother - scales[layer_name]["w"] = module.mlp.down_proj.weight.abs().max( - dim=1)[0] - - # ================================================================== - if hasattr(module, 'residual_mlp'): - fc1_layer_name = name + ".residual_mlp.w1" - gate_layer_name = name + ".residual_mlp.w3" - - smoother = smooth_gemm_fc1_gate(module.residual_mlp.w1.weight, - module.residual_mlp.w3.weight, - scales[fc1_layer_name]["x"], - module.residual_layernorm.weight, - None, alpha) - - scales[fc1_layer_name]["x"] = scales[fc1_layer_name]["x"] / smoother - scales[fc1_layer_name]["w"] = module.residual_mlp.w1.weight.abs( - ).max(dim=1)[0] - - scales[gate_layer_name][ - "x"] = scales[gate_layer_name]["x"] / smoother - scales[gate_layer_name]["w"] = module.residual_mlp.w3.weight.abs( - ).max(dim=1)[0] - - # ================================================================== - layer_name = name + ".residual_mlp.w2" - smoother = smooth_gemm(module.residual_mlp.w2.weight, - scales[layer_name]["x"], None, None, alpha) - llama_smoother[layer_name] = smoother.float() - scales[layer_name]["x"] = scales[layer_name]["x"] / smoother - scales[layer_name]["w"] = module.residual_mlp.w2.weight.abs().max( - dim=1)[0] - - -@torch.no_grad() -def capture_activation_range(model, - tokenizer, - dataset, - num_samples=512, - seq_len=512): - model.eval() - device = next(model.parameters()).device - act_scales = defaultdict(lambda: {"x": None, "y": None, "w": None}) - - tokenizer.pad_token = tokenizer.eos_token - - def stat_tensor(name, tensor, act_scales, key): - hidden_dim = tensor.shape[-1] - tensor = tensor.view(-1, hidden_dim).abs().detach() - comming_max = torch.max(tensor, dim=0)[0].float() - - if act_scales[name][key] is None: - act_scales[name][key] = comming_max - else: - act_scales[name][key] = torch.max(act_scales[name][key], - comming_max) - - def stat_input_hook(m, x, y, name): - if isinstance(x, tuple): - x = x[0] - stat_tensor(name, x, act_scales, "x") - stat_tensor(name, y, act_scales, "y") - - if act_scales[name]["w"] is None: - act_scales[name]["w"] = m.weight.abs().clip(1e-8, - None).max(dim=1)[0] - - hooks = [] - for name, m in model.named_modules(): - if isinstance(m, nn.Linear) or isinstance(m, Conv1D): - hooks.append( - m.register_forward_hook( - functools.partial(stat_input_hook, name=name))) - - for i in tqdm(range(num_samples), desc="calibrating model"): - datapoint = dataset[i:i + 1] - line = copy.copy(datapoint) - line[0] = line[0] + ' TL;DR: ' - line[0] = line[0].strip() - line[0] = line[0].replace(" n't", "n't") - input_ids = tokenizer(line, - return_tensors="pt", - max_length=seq_len, - padding=True, - truncation=True).input_ids.to(device) - model(input_ids) - for h in hooks: - h.remove() - return act_scales - - -def get_weight(named_params, prefix, dtype): - if named_params[prefix + '.weight'].dtype != dtype: - named_params[prefix + - '.weight'].data = named_params[prefix + - '.weight'].to(dtype) - return named_params[prefix + '.weight'].detach() - - -def get_weight_and_scale(named_params, - prefix, - dtype, - mapping=None, - split_scale=False): - if prefix + '.weight_scale' not in named_params: - return get_weight(named_params, prefix, dtype), None - else: - assert named_params[prefix + '.weight'].dtype == torch.float8_e4m3fn - assert named_params[prefix + '.weight_scale'].dtype == torch.float32 - weight_scale = named_params[prefix + '.weight_scale'].detach() - if split_scale: - weight_scale = split(weight_scale, - mapping.tp_size, - mapping.tp_rank, - dim=0) - return named_params[prefix + - '.weight'].detach(), weight_scale.reshape(-1) - - -def get_bias(named_params, prefix, dtype): - if named_params[prefix + '.bias'].dtype != dtype: - named_params[prefix + '.bias'].data = named_params[prefix + - '.bias'].to(dtype) - return named_params[prefix + '.bias'].detach() - - -def get_weight_and_bias(named_params, prefix, dtype): - return get_weight(named_params, prefix, - dtype), get_bias(named_params, prefix, dtype) - - -def fp8_per_channel_quant_weight_gpu(weight, clamp_val, rank=0): - weight = weight.to("cuda:" + str(rank)) - # activation range bound. - x = weight.to(torch.float32).clamp(clamp_val[0], clamp_val[1]) - xmax = x.abs().max(-1, keepdim=True).values - # minimum scaling factor. - torch_weight_scales = (xmax / 448.0).clamp(min=1.0 / (448.0 * 512.0)) - out = x / torch_weight_scales - torch_weight_scales = torch_weight_scales.reshape(-1) - out = torch.clamp(out, -448, 448) - processed_torch_weights = out.to(torch.float8_e4m3fn) - - processed_torch_weights = processed_torch_weights.to( - torch.float8_e4m3fn).cpu() - torch_weight_scales = torch_weight_scales.cpu() - - return processed_torch_weights, torch_weight_scales - - -def get_tllm_linear_sq_weight(vals, - prefix, - shape, - tensor_parallel, - is_qkv=False, - per_token=False, - per_channel=False, - last_prefix=None, - bias=None, - smoother_value=None, - smoother_shape=None, - rank=0, - cat_dim=0, - multi_query_mode=False): - results = {} - - def multi_query_split(data, local_dim, head_size, tp_size, cur_rank): - q, k, v = torch.split(data, [local_dim, head_size, head_size], dim=-1) - q_split = torch.chunk(q, tp_size, dim=-1) - k_split = torch.chunk(k, tp_size, dim=-1) - v_split = torch.chunk(v, tp_size, dim=-1) - return [ - torch.concat((q_split[ii], k_split[ii], v_split[ii]), axis=-1) - for ii in range(tp_size) - ][cur_rank] - - col_shape = shape if (is_qkv or per_channel) else [1, 1] - - if per_token: - if per_channel: - original_weights = torch.Tensor(vals["weight.int8.col"]) - else: - original_weights = torch.Tensor(vals["weight.int8"]) - local_dim = original_weights.shape[0] - head_size = (original_weights.shape[1] - local_dim) // 2 - - if multi_query_mode: - cur_weights = multi_query_split(original_weights, local_dim, - head_size, tensor_parallel, rank) - else: - cur_weights = torch.chunk(original_weights, - tensor_parallel, - dim=cat_dim)[rank] - if is_qkv: - hidden_dim = cur_weights.shape[0] - cur_weights = cur_weights.reshape(hidden_dim, -1) - results[prefix + 'weight'] = cur_weights.t().contiguous() - if smoother_value is None: - results[last_prefix] = torch.Tensor([1.0]).to(torch.float32) - - if per_channel: - cur_per_channel_value = vals["scale_w_quant_orig.col"] - if smoother_value is None: - if multi_query_mode: - cur_per_channel_value = multi_query_split( - vals["scale_w_quant_orig.col"], local_dim, head_size, - tensor_parallel, rank) - else: - cur_per_channel_value = torch.chunk( - vals["scale_w_quant_orig.col"], - tensor_parallel, - dim=cat_dim)[rank] - else: - cur_per_channel_value = vals["scale_w_quant_orig"] - if is_qkv: - if multi_query_mode: - cur_per_channel_value = multi_query_split( - vals["scale_w_quant_orig"], local_dim, head_size, - tensor_parallel, rank) - else: - cur_per_channel_value = torch.chunk( - vals["scale_w_quant_orig"], - tensor_parallel, - dim=cat_dim)[rank] - - results[prefix + 'per_channel_scale'] = cur_per_channel_value.reshape( - col_shape).contiguous() - else: - if per_channel: - original_weights = torch.Tensor(vals["weight.int8.col"]) - else: - original_weights = torch.Tensor(vals["weight.int8"]) - local_dim = original_weights.shape[0] - head_size = (original_weights.shape[1] - local_dim) // 2 - - if multi_query_mode: - cur_weights = multi_query_split(original_weights, local_dim, - head_size, tensor_parallel, rank) - else: - cur_weights = torch.chunk(original_weights, - tensor_parallel, - dim=cat_dim)[rank] - if is_qkv: - hidden_dim = cur_weights.shape[0] - cur_weights = cur_weights.reshape(hidden_dim, -1) - results[prefix + 'weight'] = cur_weights.t().contiguous() - - if per_channel: - cur_per_channel_value = vals["scale_y_accum_quant.col"] - if smoother_value is None: - if multi_query_mode: - cur_per_channel_value = multi_query_split( - vals["scale_y_accum_quant.col"], local_dim, head_size, - tensor_parallel, rank) - else: - cur_per_channel_value = torch.chunk( - vals["scale_y_accum_quant.col"], - tensor_parallel, - dim=cat_dim)[rank] - else: - cur_per_channel_value = vals["scale_y_accum_quant"] - # QKV is always per_channel - if is_qkv: - if multi_query_mode: - cur_per_channel_value = multi_query_split( - vals["scale_y_accum_quant"], local_dim, head_size, - tensor_parallel, rank) - else: - cur_per_channel_value = torch.chunk( - vals["scale_y_accum_quant"], - tensor_parallel, - dim=cat_dim)[rank] - - results[prefix + - 'per_channel_scale'] = torch.Tensor(cur_per_channel_value).to( - torch.float32).reshape(col_shape).contiguous() - results[prefix + 'act_scale'] = torch.Tensor( - [[vals['scale_y_quant_orig']]]).to(torch.float32).contiguous() - results[last_prefix] = torch.Tensor([vals['scale_x_orig_quant'] - ]).to(torch.float32).contiguous() - - if smoother_value is not None: - cur_smoother_value = torch.chunk(smoother_value, - tensor_parallel, - dim=cat_dim)[rank] - results[prefix + 'smoother'] = cur_smoother_value.reshape( - smoother_shape).contiguous().to(torch.float32) - - if bias is not None: - results[prefix + 'bias'] = bias - - return results - - -def get_prefix_and_param_name_map(architecture, use_safetensors=False): - - key_postfix = "" - if use_safetensors: - key_postfix = ".weight" - - architecture = architecture.lower() - if "exaone" in architecture: - model_prefix = "transformer" - param_name_map = { - "vocab_embedding": "wte" + key_postfix, # vocab_embedding - "lm_head": "lm_head" + key_postfix, # lm_head - "ln_f": "ln_f" + key_postfix, # ln_f - "attention.qkv": "attn.attention", # attention.qkv - "qkv_suffix": "_proj" + key_postfix, # qkv_suffix - "attention.dense": - "attn.attention.out_proj" + key_postfix, # attention.dense - "mlp.gate": "mlp.c_fc_1" + key_postfix, # mlp.gate - "mlp.proj": "mlp.c_proj" + key_postfix, # mlp.proj - "mlp.fc": "mlp.c_fc_0" + key_postfix, # mlp.fc - "input_layernorm": "ln_1" + key_postfix, # input_layernorm - "post_layernorm": "ln_2" + key_postfix, # post_layernorm - } - layer_prefix = 'h' - else: # LLaMA - model_prefix = "model" - param_name_map = { - "vocab_embedding": "embed_tokens" + key_postfix, # vocab_embedding - "lm_head": "lm_head" + key_postfix, # lm_head - "ln_f": "norm" + key_postfix, # ln_f - "attention.qkv": "self_attn", # attention.qkv - "qkv_suffix": "_proj" + key_postfix, # qkv suffix - "attention.dense": - "self_attn.o_proj" + key_postfix, # attention.dense - "mlp.gate": "mlp.up_proj" + key_postfix, # mlp.gate - "mlp.proj": "mlp.down_proj" + key_postfix, # mlp.proj - "mlp.fc": "mlp.gate_proj" + key_postfix, # mlp.fc - "input_layernorm": - "input_layernorm" + key_postfix, # input_layernorm - "post_layernorm": - "post_attention_layernorm" + key_postfix, # post_layernorm - } - layer_prefix = 'layers' - - return model_prefix, layer_prefix, param_name_map - - -def load_hf_llama(model_dir: str, load_model_on_cpu: bool = False): - if "vila" in model_dir: - sys.path.append(model_dir + "/../VILA") - from llava.model import LlavaLlamaConfig, LlavaLlamaModel # noqa - from transformers import AutoModel - model = AutoModel.from_pretrained( - model_dir, - device_map='auto', - trust_remote_code=True, - ) - return model.llm - - hf_config = AutoConfig.from_pretrained(model_dir, trust_remote_code=True) - model_cls = AutoModelForCausalLM - if hf_config.model_type == "llava": - from transformers import LlavaForConditionalGeneration - model_cls = LlavaForConditionalGeneration - if hf_config.model_type == "llava_next": - from transformers import LlavaNextForConditionalGeneration - model_cls = LlavaNextForConditionalGeneration - model = model_cls.from_pretrained( - model_dir, - device_map='auto' if not load_model_on_cpu else 'cpu', - dtype='auto', - trust_remote_code=True, - ) - if hf_config.model_type in ["llava", "llava_next"]: - model = model.language_model - return model - - -def load_weights_from_hf_model(hf_model, - config: LLaMAConfig, - act_range: Optional[dict] = None, - qkv_para: Optional[dict] = None, - smoother: Optional[dict] = None): - quant_algo = config.quantization.quant_algo - use_weight_only = quant_algo in [QuantAlgo.W8A16, QuantAlgo.W4A16] - if quant_algo == QuantAlgo.W8A16: - plugin_weight_only_quant_type = torch.int8 - elif quant_algo == QuantAlgo.W4A16: - plugin_weight_only_quant_type = torch.quint4x2 - else: - plugin_weight_only_quant_type = None - use_gemm_woq_plugin = (not config.disable_weight_only_quant_plugin) - use_fp8_rowwise = quant_algo in [QuantAlgo.FP8_PER_CHANNEL_PER_TOKEN] - - use_smooth_quant = config.quantization._use_plugin_sq - per_channel = use_smooth_quant and 'PER_CHANNEL' in quant_algo - per_token = use_smooth_quant and 'PER_TOKEN' in quant_algo - int8_kv_cache = config.quantization.kv_cache_quant_algo == QuantAlgo.INT8 - fp8_kv_cache = config.quantization.kv_cache_quant_algo == QuantAlgo.FP8 - if use_smooth_quant or int8_kv_cache: - assert act_range is not None - assert qkv_para is not None - assert smoother is not None - - weights = {} - tik = time.time() - model_params = dict(hf_model.named_parameters()) - dtype = getattr(torch, config.dtype) - - mapping = config.mapping - moe_config = config.moe - mha_mode = (config.num_key_value_heads == config.num_attention_heads) - layers_range = config.mapping.pp_layers(config.num_hidden_layers) - exclude_layers_id = [0, config.num_hidden_layers - 1] - - model_prefix, layer_prefix, param_name_map = get_prefix_and_param_name_map( - config.architecture) - - def convert_layer(l): - prefix = f'{model_prefix}.{layer_prefix}.{l}.' - tllm_prex = f'transformer.layers.{l - layers_range[0]}.' - q_weight = get_weight( - model_params, prefix + f'{param_name_map["attention.qkv"]}.q_proj', - dtype) - k_weight = get_weight( - model_params, prefix + f'{param_name_map["attention.qkv"]}.k_proj', - dtype) - v_weight = get_weight( - model_params, prefix + f'{param_name_map["attention.qkv"]}.v_proj', - dtype) - - # Meta's recipe of not using fp8 rowwise for the first and last layer. - use_fp8_rowwise_in_layer = use_fp8_rowwise and ( - l not in exclude_layers_id) - - if not mha_mode: - if config.num_key_value_heads < mapping.tp_size: - # duplicate the KV heads up to tensor_parallel - k_weight = dup_kv_weight(k_weight, config.num_key_value_heads, - mapping.tp_size) - v_weight = dup_kv_weight(v_weight, config.num_key_value_heads, - mapping.tp_size) - assert (k_weight.shape[0] % - (mapping.tp_size * config.head_size)) == 0 - assert (v_weight.shape[0] % - (mapping.tp_size * config.head_size)) == 0 - - wq = split(q_weight, mapping.tp_size, mapping.tp_rank) - wk = split(k_weight, mapping.tp_size, mapping.tp_rank) - wv = split(v_weight, mapping.tp_size, mapping.tp_rank) - - split_v = torch.concat((wq, wk, wv)) - - else: - qkv_weight = torch.cat([q_weight, k_weight, v_weight], dim=0) - - split_v = split_qkv_tp(qkv_weight, config.num_attention_heads, - config.hidden_size, mapping.tp_size, - mapping.tp_rank) - - if prefix + f'{param_name_map["attention.qkv"]}.q_proj.bias' in model_params: - # only used in Internlm 7B models - q_bias = get_bias( - model_params, - prefix + f'{param_name_map["attention.qkv"]}.q_proj', dtype) - k_bias = get_bias( - model_params, - prefix + f'{param_name_map["attention.qkv"]}.k_proj', dtype) - v_bias = get_bias( - model_params, - prefix + f'{param_name_map["attention.qkv"]}.v_proj', dtype) - qkv_bias = torch.cat((q_bias, k_bias, v_bias)) - split_bias_v = split_qkv_bias_tp(qkv_bias, - config.num_attention_heads, - config.hidden_size, - mapping.tp_size, mapping.tp_rank) - else: - split_bias_v = None - - if use_smooth_quant: - qkv_weight = qkv_para[prefix + - f'{param_name_map["attention.qkv"]}.qkv_proj'] - qkv_out_dim = qkv_weight.shape[1] - - if not mha_mode: - local_dim = qkv_weight.shape[0] - kv_hidden_size = (qkv_weight.shape[-1] - local_dim) // 2 - qkv_weight = qkv_weight.reshape(local_dim, - local_dim + 2 * kv_hidden_size) - else: - qkv_weight = qkv_weight.reshape(config.hidden_size, 3, - config.hidden_size) - - int8_weights = generate_int8( - qkv_weight, - act_range.get(prefix + - f'{param_name_map["attention.qkv"]}.qkv_proj'), - is_qkv=True, - multi_query_mode=bool(not mha_mode)) - - weights.update( - get_tllm_linear_sq_weight(int8_weights, - tllm_prex + 'attention.qkv.', - [1, qkv_out_dim // mapping.tp_size], - mapping.tp_size, - is_qkv=True, - bias=split_bias_v, - per_token=per_token, - per_channel=per_channel, - last_prefix=tllm_prex + - 'input_layernorm.scale_to_int', - smoother_value=None, - smoother_shape=None, - rank=mapping.tp_rank, - cat_dim=-1, - multi_query_mode=bool(not mha_mode))) - else: - weights.update( - get_tllm_linear_weight(split_v, - tllm_prex + 'attention.qkv.', - split_bias_v, - use_weight_only, - plugin_weight_only_quant_type, - dtype, - use_gemm_woq_plugin, - use_fp8_rowwise=False)) - - if int8_kv_cache: - qkv_y = torch.cat([ - act_range.get(prefix + - f'{param_name_map["attention.qkv"]}.q_proj')["y"], - act_range.get(prefix + - f'{param_name_map["attention.qkv"]}.k_proj')["y"], - act_range.get(prefix + - f'{param_name_map["attention.qkv"]}.v_proj')["y"] - ], - dim=0) - - int8_kv_scales = qkv_y.max() / 127. - - kv_cache_weights = {} - - kv_cache_weights[ - tllm_prex + - 'attention.kv_cache_scaling_factor'] = int8_kv_scales.reshape( - [1]) - - weights.update(kv_cache_weights) - elif fp8_kv_cache: - # FIXME: set it to 1.0f for fp8 kv cache. - weights[tllm_prex + - 'attention.kv_cache_scaling_factor'] = torch.tensor( - [1.0], dtype=torch.float32) - - attn_dense_weight = get_weight( - model_params, prefix + param_name_map["attention.dense"], dtype) - split_v = split_matrix_tp(attn_dense_weight, - mapping.tp_size, - mapping.tp_rank, - dim=1) - - if prefix + f'{param_name_map["attention.dense"]}.bias' in model_params: - attn_dense_bias = get_bias( - model_params, prefix + param_name_map["attention.dense"], dtype) - else: - attn_dense_bias = None - if use_smooth_quant: - attn_dense_weight = attn_dense_weight.t() - proj_out_dim = attn_dense_weight.shape[0] - - int8_weights = generate_int8( - attn_dense_weight, - act_range.get(prefix + param_name_map["attention.dense"])) - weights.update( - get_tllm_linear_sq_weight( - int8_weights, - tllm_prex + 'attention.dense.', [1, config.hidden_size], - mapping.tp_size, - is_qkv=False, - bias=attn_dense_bias, - per_token=per_token, - per_channel=per_channel, - last_prefix=tllm_prex + - 'attention.quantization_scaling_factor', - smoother_value=smoother[( - prefix + param_name_map["attention.dense"])], - smoother_shape=[1, proj_out_dim // mapping.tp_size], - rank=mapping.tp_rank, - cat_dim=0)) - else: - weights.update( - get_tllm_linear_weight(split_v, - tllm_prex + 'attention.dense.', - attn_dense_bias, - use_weight_only, - plugin_weight_only_quant_type, - dtype, - use_gemm_woq_plugin, - use_fp8_rowwise=False)) - - if moe_config.has_moe(): - rank_experts = list(range(moe_config.num_experts)) - if mapping.has_moe_ep(): - rank_experts = mapping.ep_experts(moe_config.num_experts) - architecture = config.architecture.lower() - if "granite" not in architecture: - for suffix in ["w1", "w2", "w3"]: - model_params[f'model.layers.{l}.block_sparse_moe.experts.{suffix}.weight'] = \ - torch.stack([model_params[f'model.layers.{l}.block_sparse_moe.experts.{expert}.{suffix}.weight'].detach() - for expert in rank_experts]) - w3 = model_params[ - f'model.layers.{l}.block_sparse_moe.experts.w3.weight'] - w2 = model_params[ - f'model.layers.{l}.block_sparse_moe.experts.w2.weight'] - w1 = model_params[ - f'model.layers.{l}.block_sparse_moe.experts.w1.weight'] - else: - w2 = model_params[ - f'model.layers.{l}.block_sparse_moe.output_linear.weight'] - half_size = model_params[ - f'model.layers.{l}.block_sparse_moe.input_linear.weight'].shape[ - -2] // 2 - w1, w3 = model_params[ - f'model.layers.{l}.block_sparse_moe.input_linear.weight']\ - .split(half_size, dim=-2) - w1 = w1[rank_experts[0]:rank_experts[-1] + 1] - w2 = w2[rank_experts[0]:rank_experts[-1] + 1] - w3 = w3[rank_experts[0]:rank_experts[-1] + 1] - - if mapping.has_moe_tp(): - w3 = split(w3, mapping.moe_tp_size, mapping.moe_tp_rank, dim=1) - w2 = split(w2, mapping.moe_tp_size, mapping.moe_tp_rank, dim=2) - w1 = split(w1, mapping.moe_tp_size, mapping.moe_tp_rank, dim=1) - - model_params[ - f'model.layers.{l}.block_sparse_moe.experts.w3w1.weight'] = torch.concat( - [w3, w1], dim=-2) - - model_params[ - f'model.layers.{l}.block_sparse_moe.experts.w2.weight'] = w2 - - ## block_sparse_moe.experts.w2.weight - moe_experts_w2_weights = get_weight( - model_params, prefix + 'block_sparse_moe.experts.w2', dtype) - weights.update( - get_tllm_linear_weight(moe_experts_w2_weights, - tllm_prex + 'mlp.proj.', None, - use_weight_only, - plugin_weight_only_quant_type, dtype, - use_gemm_woq_plugin)) - ##block_sparse_moe.experts.w3w1.weight - moe_experts_w3w1_weights = get_weight( - model_params, prefix + 'block_sparse_moe.experts.w3w1', dtype) - weights.update( - get_tllm_linear_weight(moe_experts_w3w1_weights, - tllm_prex + 'mlp.fc.', None, - use_weight_only, - plugin_weight_only_quant_type, dtype, - use_gemm_woq_plugin)) - - if config.residual_mlp: - residual_mlp_gate_weights = get_weight( - model_params, prefix + 'residual_mlp.w3', dtype) - if use_smooth_quant: - residual_mlp_gate_weights = residual_mlp_gate_weights.t() - int8_weights = generate_int8( - residual_mlp_gate_weights, - act_range.get(prefix + 'residual_mlp.w3')) - weights.update( - get_tllm_linear_sq_weight( - int8_weights, - tllm_prex + 'residual_mlp.gate.', - [1, config.hidden_size // mapping.tp_size], - mapping.tp_size, - is_qkv=False, - per_token=per_token, - per_channel=per_channel, - last_prefix=tllm_prex + - 'post_layernorm.scale_to_int', - smoother_value=None, - smoother_shape=None, - rank=mapping.tp_rank, - cat_dim=-1)) - else: - split_v = split_matrix_tp(residual_mlp_gate_weights, - mapping.tp_size, - mapping.tp_rank, - dim=0) - weights.update( - get_tllm_linear_weight(split_v, - tllm_prex + 'residual_mlp.gate.', - None, use_weight_only, - plugin_weight_only_quant_type, - dtype, use_gemm_woq_plugin)) - - residual_mlp_fc_weight = get_weight(model_params, - prefix + 'residual_mlp.w1', - dtype) - if use_smooth_quant: - residual_mlp_fc_weight = residual_mlp_fc_weight.t( - ) #verified - int8_weights = generate_int8( - residual_mlp_fc_weight, - act_range.get(prefix + 'residual_mlp.w1')) - weights.update( - get_tllm_linear_sq_weight( - int8_weights, - tllm_prex + 'residual_mlp.fc.', - [1, config.hidden_size // mapping.tp_size], - mapping.tp_size, - is_qkv=False, - per_token=per_token, - per_channel=per_channel, - last_prefix=tllm_prex + - 'post_layernorm.scale_to_int', - smoother_value=None, - smoother_shape=None, - rank=mapping.tp_rank, - cat_dim=-1)) - else: - split_v = split_matrix_tp(residual_mlp_fc_weight, - mapping.tp_size, - mapping.tp_rank, - dim=0) - weights.update( - get_tllm_linear_weight(split_v, - tllm_prex + 'residual_mlp.fc.', - None, use_weight_only, - plugin_weight_only_quant_type, - dtype, use_gemm_woq_plugin)) - - residual_mlp_proj_weight = get_weight( - model_params, prefix + 'residual_mlp.w2', dtype) - - if use_smooth_quant: - residual_mlp_proj_weight = residual_mlp_proj_weight.t() - int8_weights = generate_int8( - residual_mlp_proj_weight, - act_range.get(prefix + 'residual_mlp.w2')) - weights.update( - get_tllm_linear_sq_weight( - int8_weights, - tllm_prex + 'residual_mlp.proj.', - [1, config.hidden_size], - mapping.tp_size, - is_qkv=False, - per_token=per_token, - per_channel=per_channel, - last_prefix=tllm_prex + - 'residual_mlp.quantization_scaling_factor', - smoother_value=smoother[prefix + 'residual_mlp.w2'], - smoother_shape=[ - 1, config.hidden_size // mapping.tp_size - ], - rank=mapping.tp_rank, - cat_dim=0)) - else: - split_v = split_matrix_tp(residual_mlp_proj_weight, - mapping.tp_size, - mapping.tp_rank, - dim=1) - weights.update( - get_tllm_linear_weight(split_v, - tllm_prex + 'residual_mlp.proj.', - None, use_weight_only, - plugin_weight_only_quant_type, - dtype, use_gemm_woq_plugin)) - - architecture = config.architecture.lower() - if "granite" not in architecture: - moe_experts_gate_weights = get_weight( - model_params, prefix + 'block_sparse_moe.gate', - torch.float32) - else: - moe_experts_gate_weights = get_weight( - model_params, prefix + 'block_sparse_moe.router.layer', - torch.float32) - weights.update( - get_tllm_linear_weight( - moe_experts_gate_weights, - tllm_prex + 'mlp.router.', - None, - False, # Router should never be quantized - plugin_weight_only_quant_type, - dtype, - use_gemm_woq_plugin)) - else: - mlp_gate_weight, mlp_gate_weight_scale = get_weight_and_scale( - model_params, prefix + param_name_map["mlp.gate"], dtype, - mapping, True) - split_v = split_matrix_tp(mlp_gate_weight, - mapping.tp_size, - mapping.tp_rank, - dim=0) - if use_smooth_quant: - - mlp_gate_weight = mlp_gate_weight.t() - int8_weights = generate_int8( - # mlp_gate_weight, act_range.get(prefix + 'mlp.up_proj')) - mlp_gate_weight, - act_range.get(prefix + param_name_map["mlp.gate"])) - - weights.update( - get_tllm_linear_sq_weight( - int8_weights, - tllm_prex + 'mlp.gate.', - [1, config.intermediate_size // mapping.tp_size], - mapping.tp_size, - is_qkv=False, - per_token=per_token, - per_channel=per_channel, - last_prefix=tllm_prex + 'post_layernorm.scale_to_int', - smoother_value=None, - smoother_shape=None, - rank=mapping.tp_rank, - cat_dim=-1)) - else: - weights.update( - get_tllm_linear_weight( - split_v, - tllm_prex + 'mlp.gate.', - None, - use_weight_only, - plugin_weight_only_quant_type, - dtype, - use_gemm_woq_plugin, - use_fp8_rowwise_in_layer, - weight_scale=mlp_gate_weight_scale, - clamp_value=config.quantization.clamp_val)) - - mlp_fc_weight, mlp_fc_weight_scale = get_weight_and_scale( - model_params, prefix + param_name_map["mlp.fc"], dtype, mapping, - True) - split_v = split_matrix_tp(mlp_fc_weight, - mapping.tp_size, - mapping.tp_rank, - dim=0) - - if use_smooth_quant: - mlp_fc_weight = mlp_fc_weight.t() #verified - int8_weights = generate_int8( - mlp_fc_weight, - act_range.get(prefix + param_name_map["mlp.fc"])) - weights.update( - get_tllm_linear_sq_weight( - int8_weights, - tllm_prex + 'mlp.fc.', - [1, config.intermediate_size // mapping.tp_size], - mapping.tp_size, - is_qkv=False, - per_token=per_token, - per_channel=per_channel, - last_prefix=tllm_prex + 'post_layernorm.scale_to_int', - smoother_value=None, - smoother_shape=None, - rank=mapping.tp_rank, - cat_dim=-1)) - else: - weights.update( - get_tllm_linear_weight( - split_v, - tllm_prex + 'mlp.fc.', - None, - use_weight_only, - plugin_weight_only_quant_type, - dtype, - use_gemm_woq_plugin, - use_fp8_rowwise_in_layer, - weight_scale=mlp_fc_weight_scale, - clamp_value=config.quantization.clamp_val)) - - mlp_proj_weight, mlp_proj_weight_scale = get_weight_and_scale( - model_params, prefix + param_name_map["mlp.proj"], dtype) - split_v = split_matrix_tp(mlp_proj_weight, - mapping.tp_size, - mapping.tp_rank, - dim=1) - - if use_smooth_quant: - mlp_proj_weight = mlp_proj_weight.t() - int8_weights = generate_int8( - mlp_proj_weight, - act_range.get(prefix + param_name_map["mlp.proj"])) - weights.update( - get_tllm_linear_sq_weight( - int8_weights, - tllm_prex + 'mlp.proj.', [1, config.hidden_size], - mapping.tp_size, - is_qkv=False, - per_token=per_token, - per_channel=per_channel, - last_prefix=tllm_prex + - 'mlp.quantization_scaling_factor', - smoother_value=smoother[prefix + - param_name_map["mlp.proj"]], - smoother_shape=[ - 1, config.intermediate_size // mapping.tp_size - ], - rank=mapping.tp_rank, - cat_dim=0)) - else: - weights.update( - get_tllm_linear_weight( - split_v, - tllm_prex + 'mlp.proj.', - None, - use_weight_only, - plugin_weight_only_quant_type, - dtype, - use_gemm_woq_plugin, - use_fp8_rowwise_in_layer, - weight_scale=mlp_proj_weight_scale, - clamp_value=config.quantization.clamp_val)) - - # Layer norms do not use tensor parallelism - input_ln_weight = get_weight(model_params, - prefix + param_name_map["input_layernorm"], - dtype) - weights[tllm_prex + 'input_layernorm.weight'] = input_ln_weight - - post_ln_weight = get_weight(model_params, - prefix + param_name_map["post_layernorm"], - dtype) - weights[tllm_prex + 'post_layernorm.weight'] = post_ln_weight - - if config.residual_mlp: - residual_ln_weight = get_weight(model_params, - prefix + 'residual_layernorm', - dtype) - weights[tllm_prex + - 'residual_layernorm.weight'] = residual_ln_weight - - cur_block_weights = [ - weight_name for weight_name in model_params - if weight_name.find(prefix) != -1 - ] - for weight_name in cur_block_weights: - model_params[weight_name] = None - - for l in layers_range: - convert_layer(l) - release_gc() - - vocab_embedding = get_weight( - model_params, f'{model_prefix}.{param_name_map["vocab_embedding"]}', - dtype) - - if mapping.is_first_pp_rank(): - if config.use_parallel_embedding: - weights['transformer.vocab_embedding.weight'] = split_matrix_tp( - vocab_embedding, - mapping.tp_size, - mapping.tp_rank, - dim=config.embedding_sharding_dim) - else: - weights['transformer.vocab_embedding.weight'] = vocab_embedding - - if mapping.is_last_pp_rank(): - if hf_model.config.tie_word_embeddings: - # lm_head.weight has the same weights as embedding - lm_head_weights = vocab_embedding.clone() - else: - lm_head_weights = get_weight(model_params, - param_name_map["lm_head"], dtype) - - if config.vocab_size % mapping.tp_size != 0: - # padding - vocab_size_padded = pad_vocab_size(config.vocab_size, - mapping.tp_size) - pad_width = vocab_size_padded - config.vocab_size - - lm_head_weights = torch.nn.functional.pad(lm_head_weights, - (0, 0, 0, pad_width), - 'constant', - value=0) - weights['lm_head.weight'] = split_matrix_tp(lm_head_weights, - mapping.tp_size, - mapping.tp_rank, - dim=0) - ln_f_w = get_weight(model_params, - f'{model_prefix}.{param_name_map["ln_f"]}', dtype) - weights['transformer.ln_f.weight'] = ln_f_w - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Weights loaded. Total time: {t}') - return weights - - -def quantize(hf_model_dir: str, - output_dir: str, - config: LLaMAConfig, - device: str = 'cuda', - calib_dataset: str = 'cnn_dailymail', - trust_remote_code: bool = True, - calib_batches: int = 512, - calib_max_seq_length: int = 512): - ''' - Quantize the save the model as TRT-LLM checkpoint to output_dir - ''' - os.makedirs(output_dir, exist_ok=True) - config.to_json_file(os.path.join(output_dir, 'config.json')) - - mapping = config.mapping - assert mapping.rank == 0, "quantize should be called at rank 0 only" - - quant_config = config.quantization - use_smooth_quant = quant_config._use_plugin_sq - int8_kv_cache = quant_config.kv_cache_quant_algo == QuantAlgo.INT8 - - assert use_smooth_quant or int8_kv_cache, "Call from_hugging_face when there is no quantization" - assert hf_model_dir is not None - ## only load and call smooth quant routine once for all ranks - hf_config = AutoConfig.from_pretrained(hf_model_dir, - trust_remote_code=trust_remote_code) - assert "llava" not in hf_config.model_type, "Smooth quant llava/vila/llava_next is not supported yet" - hf_model = AutoModelForCausalLM.from_pretrained( - hf_model_dir, - device_map='auto' if device != 'cpu' else 'cpu', - dtype='auto' if not use_smooth_quant else torch.float16, - trust_remote_code=trust_remote_code) - - os.environ["TOKENIZERS_PARALLELISM"] = os.environ.get( - "TOKENIZERS_PARALLELISM", "false") - tokenizer = AutoTokenizer.from_pretrained( - hf_model_dir, - trust_remote_code=trust_remote_code, - use_fast=False, - padding_side='left') - - dataset = load_calib_dataset(calib_dataset) - - if calib_batches == -1: # use the whole dataset if calib_batches is -1 - calib_batches = len(dataset) - - act_range = capture_activation_range(hf_model, - tokenizer, - dataset, - num_samples=calib_batches, - seq_len=calib_max_seq_length) - qkv_para, smoother = {}, {} - if use_smooth_quant: - smooth_llama_model(hf_model, act_range, quant_config.smoothquant_val, - qkv_para, smoother) - - for rank in range(mapping.world_size): - # To avoid changing the mapping arg in-place, also the given mapping from caller is rank agnostic, since quantize is called from only one rank - config = copy.deepcopy(config) - config.set_rank(rank) - weights = load_weights_from_hf_model( - hf_model, - config=config, - act_range=act_range, - qkv_para=qkv_para, - smoother=smoother, - ) - safetensors.torch.save_file( - weights, os.path.join(output_dir, f'rank{rank}.safetensors')) - del weights - - -class QkvWeightHelper: - """ A helper utility for loading QKV weights from sharded files. """ - - def __init__(self, config: PretrainedConfig): - self.hidden_size = config.hidden_size - self.num_heads = config.num_attention_heads - self.num_kv_heads = config.num_key_value_heads - self.tp_size = config.mapping.tp_size - self.tp_rank = config.mapping.tp_rank - self.is_mha = self.num_heads == self.num_kv_heads - self.head_size = None if not hasattr(config, - "head_size") else config.head_size - self._qkv_weights = {} - self.remove_duplicated_kv_heads = getattr(config, - 'remove_duplicated_kv_heads', - False) - - @staticmethod - def is_qkv_weight(name): - for k in ['q_proj', 'k_proj', 'v_proj']: - if 'self_attn' in name and k in name: - return True - return False - - def add_weight(self, i: int, name: str, weight: torch.Tensor): - if 'q_proj' in name: - tag = 'q' - elif 'k_proj' in name: - tag = 'k' - elif 'v_proj' in name: - tag = 'v' - else: - raise ValueError(f'Got an unexpected parameter of name {name}') - if i not in self._qkv_weights: - self._qkv_weights[i] = {} - self._qkv_weights[i][tag] = weight - - def is_qkv_prepared(self, layer_idx): - if layer_idx not in self._qkv_weights: - return False - weights = self._qkv_weights[layer_idx] - return 'q' in weights and 'k' in weights and 'v' in weights - - def split_qkv_weights(self, layer_idx): - if not self.is_qkv_prepared(layer_idx): - return None - weights = self._qkv_weights.pop(layer_idx) # to prevent memory leak. - q, k, v = (torch.tensor(weights[t]) for t in ['q', 'k', 'v']) - - if self.remove_duplicated_kv_heads: - head_size = self.hidden_size // self.num_heads if self.head_size is None else self.head_size - k = k.reshape( - [k.shape[0] // head_size // 2, 2, head_size, self.hidden_size]) - v = v.reshape( - [v.shape[0] // head_size // 2, 2, head_size, self.hidden_size]) - assert (k[:, 0] == k[:, 1]).all() - assert (v[:, 0] == v[:, 1]).all() - k = k[:, 0].reshape([-1, self.hidden_size]) - v = v[:, 0].reshape([-1, self.hidden_size]) - - if not self.is_mha: - head_size = self.hidden_size // self.num_heads if self.head_size is None else self.head_size - if self.num_kv_heads < self.tp_size: - # duplicate the KV heads up to tensor_parallel - k = dup_kv_weight(k, self.num_kv_heads, self.tp_size) - v = dup_kv_weight(v, self.num_kv_heads, self.tp_size) - assert k.shape[0] % (self.tp_size * head_size) == 0 - assert v.shape[0] % (self.tp_size * head_size) == 0 - wq = split(q, self.tp_size, self.tp_rank) - wk = split(k, self.tp_size, self.tp_rank) - wv = split(v, self.tp_size, self.tp_rank) - fused_qkv = torch.cat((wq, wk, wv), dim=0) - else: - qkv = torch.cat([q, k, v], dim=0) - qkv = qkv.reshape(3, q.shape[0], q.shape[1]) - fused_qkv = split(qkv, self.tp_size, self.tp_rank, dim=1) - fused_qkv = fused_qkv.reshape(3 * (q.shape[0] // self.tp_size), - q.shape[1]) - return fused_qkv - - -def load_weights_from_hf_by_shard(model_dir: str, config: LLaMAConfig): - '''Weights-only quantization is the only supported quantization recipe here.''' - logger.info('Loading weights from HF LLaMA...') - quant_algo = config.quantization.quant_algo - use_weight_only = quant_algo in [QuantAlgo.W8A16, QuantAlgo.W4A16] - if quant_algo == QuantAlgo.W8A16: - plugin_weight_only_quant_type = torch.int8 - elif quant_algo == QuantAlgo.W4A16: - plugin_weight_only_quant_type = torch.quint4x2 - else: - plugin_weight_only_quant_type = None - - weights = {} - tik = time.time() - dtype = getattr(torch, config.dtype) - - mapping = config.mapping - moe_config = config.moe - assert not moe_config.has_moe(), "MoE does not support sharded load" - assert "Exaone" not in config.architecture, "EXAONE model currently not support sharded load" - - from transformers import AutoConfig - hf_config = AutoConfig.from_pretrained(model_dir) - - quant_mode = config.quant_mode - if quant_mode.is_int8_weight_only(): - plugin_weight_only_quant_type = torch.int8 - elif quant_mode.is_int4_weight_only(): - plugin_weight_only_quant_type = torch.quint4x2 - elif config.quant_mode.has_fp8_rowwise(): - plugin_weight_only_quant_type = torch.float8_e4m3fn - else: - plugin_weight_only_quant_type = None - use_weight_only = quant_mode.is_weight_only() - use_fp8_rowwise = quant_mode.has_fp8_rowwise() - # Meta's recipe of not using fp8 rowwise for the first and last layer. - exclude_layers_id = [0, config.num_hidden_layers - 1] - - layers_range = mapping.pp_layers(config.num_hidden_layers) - - qkv_weight_helper = QkvWeightHelper(config) - - def convert_to_dtype(name, param, model_params, dtype): - # fp8 rowwise weights will only load fp8 weights and scales for the mlp layer. - if ('weight_scale' in name or name.replace('weight', 'weight_scale') in model_params) \ - and use_fp8_rowwise: - assert 'mlp' in name, "only MLP layers support fp8 rowwise currently." - return param - else: - return param.to(dtype) - - def fp8_rowwise_quantization(name, - param, - model_params, - clamp_value, - split_scale=False): - # check if weights are already quantized. - loaded_weight_scale = model_params.get( - name.replace('weight', 'weight_scale')) - if loaded_weight_scale is not None: - assert param.dtype == torch.float8_e4m3fn, "weight data type must be torch.float8_e4m3fn" - if split_scale: - assert mapping is not None - loaded_weight_scale = split(loaded_weight_scale, - mapping.tp_size, - mapping.tp_rank, - dim=0) - - return param, loaded_weight_scale.reshape(-1) - else: - return fp8_per_channel_quant_weight_gpu(param, clamp_value) - - for model_file in iterate_shard_files(model_dir, - rank=mapping.tp_rank, - progress_bar=False): - logger.debug(f'Loading file {str(model_file)}...') - model_params = load_state_dict(model_file) - for name, param in model_params.items(): - logger.debug(f'Converting weight {name}...') - layer_idx = retrieved_layer_index_from_name(name) - tllm_prex = f'transformer.layers.{layer_idx}.' - - param = convert_to_dtype(name, param, model_params, dtype) - - if layer_idx is None: - layer = None - else: - if layer_idx not in layers_range: - continue - else: - tllm_prex = f'transformer.layers.{layer_idx - layers_range[0]}.' - - # Meta's recipe of not using fp8 rowwise for the first and last layer. - use_fp8_rowwise_in_layer = use_fp8_rowwise and ( - layer_idx not in exclude_layers_id) - - if 'model.embed_tokens.weight' in name: - if hf_config.tie_word_embeddings: - # lm_head.weight has the same weights as embedding - if mapping.is_last_pp_rank(): - - if config.vocab_size % mapping.tp_size != 0: - # padding - vocab_size_padded = pad_vocab_size( - config.vocab_size, mapping.tp_size) - pad_width = vocab_size_padded - config.vocab_size - param = torch.from_numpy( - np.pad(param.detach().cpu().numpy(), - ((0, pad_width), (0, 0)), - 'constant', - constant_values=0)) - weights['lm_head.weight'] = split( - param, mapping.tp_size, mapping.tp_rank) - if config.use_parallel_embedding: - param = split(param, mapping.tp_size, mapping.tp_rank, - config.embedding_sharding_dim) - if mapping.is_first_pp_rank(): - weights['transformer.vocab_embedding.weight'] = param - elif 'model.norm.weight' in name: - if mapping.is_last_pp_rank(): - weights['transformer.ln_f.weight'] = param - elif 'lm_head.weight' in name: - if mapping.is_last_pp_rank(): - if config.vocab_size % mapping.tp_size != 0: - # padding - vocab_size_padded = pad_vocab_size( - config.vocab_size, mapping.tp_size) - pad_width = vocab_size_padded - config.vocab_size - param = torch.from_numpy( - np.pad(param.detach().cpu().numpy(), - ((0, pad_width), (0, 0)), - 'constant', - constant_values=0)) - weights['lm_head.weight'] = split(param, mapping.tp_size, - mapping.tp_rank) - elif 'input_layernorm.weight' in name: - weights[tllm_prex + 'input_layernorm.weight'] = param - elif 'post_attention_layernorm.weight' in name: - weights[tllm_prex + 'post_layernorm.weight'] = param - elif qkv_weight_helper.is_qkv_weight(name): - qkv_weight_helper.add_weight(layer_idx, name, param) - if not qkv_weight_helper.is_qkv_prepared(layer_idx): - continue - split_v = qkv_weight_helper.split_qkv_weights(layer_idx) - if use_weight_only: - param = split_v.transpose() - processed_torch_weights, torch_weight_scales = \ - torch.ops.trtllm.symmetric_quantize_last_axis_of_batched_matrix( - param, plugin_weight_only_quant_type) - weights[tllm_prex + - 'attention.qkv.weight'] = processed_torch_weights - weights[ - tllm_prex + - 'attention.qkv.per_channel_scale'] = torch_weight_scales - else: - weights[tllm_prex + 'attention.qkv.weight'] = split_v - elif 'self_attn.o_proj.weight' in name: - split_v = split(param, mapping.tp_size, mapping.tp_rank, dim=1) - if use_weight_only: - processed_torch_weights, torch_weight_scales = \ - torch.ops.trtllm.symmetric_quantize_last_axis_of_batched_matrix( - split_v.transpose(), plugin_weight_only_quant_type) - weights[tllm_prex + - 'attention.dense.weight'] = processed_torch_weights - weights[ - tllm_prex + - 'attention.dense.per_channel_scale'] = torch_weight_scales - else: - weights[tllm_prex + 'attention.dense.weight'] = split_v - elif name.endswith('mlp.up_proj.weight'): - split_v = split(param, mapping.tp_size, mapping.tp_rank, dim=0) - if use_weight_only: - processed_torch_weights, torch_weight_scales = \ - torch.ops.trtllm.symmetric_quantize_last_axis_of_batched_matrix( - split_v.transpose(), plugin_weight_only_quant_type) - weights[tllm_prex + - 'mlp.gate.weight'] = processed_torch_weights - weights[tllm_prex + - 'mlp.gate.per_channel_scale'] = torch_weight_scales - elif use_fp8_rowwise_in_layer: - processed_torch_weights, torch_weight_scales = fp8_rowwise_quantization( - name, split_v, model_params, - config.quantization.clamp_val, True) - weights[tllm_prex + - 'mlp.gate.weight'] = processed_torch_weights.view( - plugin_weight_only_quant_type) - weights[ - tllm_prex + - 'mlp.gate.per_channel_scale'] = torch_weight_scales.to( - torch.float32) - else: - weights[tllm_prex + 'mlp.gate.weight'] = split_v - elif name.endswith('mlp.down_proj.weight'): - split_v = split(param, mapping.tp_size, mapping.tp_rank, dim=1) - if use_weight_only: - processed_torch_weights, torch_weight_scales = \ - torch.ops.trtllm.symmetric_quantize_last_axis_of_batched_matrix( - split_v.transpose(), plugin_weight_only_quant_type) - weights[tllm_prex + - 'mlp.proj.weight'] = processed_torch_weights - weights[tllm_prex + - 'mlp.proj.per_channel_scale'] = torch_weight_scales - elif use_fp8_rowwise_in_layer: - processed_torch_weights, torch_weight_scales = fp8_rowwise_quantization( - name, split_v, model_params, - config.quantization.clamp_val) - weights[tllm_prex + - 'mlp.proj.weight'] = processed_torch_weights.view( - plugin_weight_only_quant_type) - weights[ - tllm_prex + - 'mlp.proj.per_channel_scale'] = torch_weight_scales.to( - torch.float32) - else: - weights[tllm_prex + 'mlp.proj.weight'] = split_v - - elif name.endswith('mlp.gate_proj.weight'): - split_v = split(param, mapping.tp_size, mapping.tp_rank, dim=0) - if use_weight_only: - processed_torch_weights, torch_weight_scales = \ - torch.ops.trtllm.symmetric_quantize_last_axis_of_batched_matrix( - split_v.transpose(), plugin_weight_only_quant_type) - layer.mlp.fc.weight.value = processed_torch_weights - layer.mlp.fc.per_channel_scale.value = torch_weight_scales - weights[tllm_prex + - 'mlp.fc.weight'] = processed_torch_weights - weights[tllm_prex + - 'mlp.fc.per_channel_scale'] = torch_weight_scales - elif use_fp8_rowwise_in_layer: - processed_torch_weights, torch_weight_scales = fp8_rowwise_quantization( - name, split_v, model_params, - config.quantization.clamp_val, True) - weights[tllm_prex + - 'mlp.fc.weight'] = processed_torch_weights.view( - plugin_weight_only_quant_type) - weights[ - tllm_prex + - 'mlp.fc.per_channel_scale'] = torch_weight_scales.to( - torch.float32) - else: - weights[tllm_prex + 'mlp.fc.weight'] = split_v - - del model_params - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - logger.info(f'Weights loaded. Total time: {t}') - return weights - - -def load_weights_from_hf_safetensors(model_dir: str, config: LLaMAConfig): - logger.info('Loading weights from Huggingface {} safetensors...'.format( - config.architecture.split('ForCausalLM')[0])) - tik = time.time() - import json - import os - - import safetensors - weights = {} - - model_dir = model_dir if model_dir.endswith("/") else model_dir + "/" - safetensors_map = {} - has_safetensor_index_json = True - try: - with open(model_dir + "model.safetensors.index.json", 'r') as fr: - sharding_map = json.load(fr) - for k, v in sharding_map['weight_map'].items(): - safetensors_map[k] = int(v[6:11]) - 1 - except FileNotFoundError: - has_safetensor_index_json = False - - shard_files = [] - for name in os.listdir(model_dir): - if name.endswith(".safetensors"): - if has_safetensor_index_json and name not in sharding_map[ - 'weight_map'].values(): - continue - shard_files.append(name) - shard_files.sort() - safetensors_ptrs = [ - safetensors.safe_open(model_dir + shard_file, - framework="pt", - device="cpu") for shard_file in shard_files - ] - - mapping = config.mapping - num_hidden_layers = config.num_hidden_layers - vocab_size = config.vocab_size - pad_vocab = vocab_size % mapping.tp_size != 0 - vocab_size_padded = pad_vocab_size(config.vocab_size, mapping.tp_size) - dtype = config.dtype - - moe_config = config.moe - - kv_tp_size = None - kv_tp_rank = None - if config.num_key_value_heads < mapping.tp_size: - kv_tp_size = config.num_key_value_heads - kv_tp_rank = mapping.tp_rank * kv_tp_size // mapping.tp_size - - model_prefix, layer_prefix, param_name_map = get_prefix_and_param_name_map( - config.architecture, use_safetensors=True) - - torch_dtype = str_dtype_to_torch(dtype) - - def load(key, - tp_dim=-1, - no_prefix=0, - is_expert_weights=False, - tp_size=None, - tp_rank=None): - if not no_prefix: - key = f'{model_prefix}.' + key - ptr_idx = safetensors_map[key] if key in safetensors_map else 0 - - if key not in safetensors_ptrs[ptr_idx].keys(): - return None - - tensor_slice = safetensors_ptrs[ptr_idx].get_slice(key) - tensor_shape = tensor_slice.get_shape() - if tp_dim == -1: - res = tensor_slice[:] - elif tp_dim >= 0 and tp_dim < len(tensor_shape): - if is_expert_weights: - if tp_size is None: - tp_size = mapping.moe_tp_size - if tp_rank is None: - tp_rank = mapping.moe_tp_rank - else: - if tp_size is None: - tp_size = mapping.tp_size - if tp_rank is None: - tp_rank = mapping.tp_rank - dim_size = tensor_shape[tp_dim] - if dim_size % tp_size != 0: - logger.error( - f"Current weight {key}'s shape {tensor_shape} is invalid at dimension {tp_dim} for TP size {tp_size}" - ) - indices = [slice(None)] * len(tensor_shape) - indices[tp_dim] = slice(dim_size * tp_rank // tp_size, - dim_size * (tp_rank + 1) // tp_size) - res = tensor_slice[indices] - else: - raise ValueError( - f"Invalid TP dim {tp_dim} for weight {key}'s shape {tensor_shape}" - ) - return res.to(torch_dtype).contiguous( - ) if "block_sparse_moe.gate" not in key and "block_sparse_moe.router" not in key else res.to( - torch.float32) - - def load_and_set(target, - key, - tp_dim=-1, - no_prefix=0, - is_expert_weights=False): - res = load(key, tp_dim, no_prefix, is_expert_weights) - weights[target] = res - if "weight" in key: - bias = load(key.replace("weight", "bias"), -1, no_prefix, - is_expert_weights) - if bias is not None: - weights[target.replace("weight", "bias")] = bias - - if mapping.is_first_pp_rank(): - weights['transformer.vocab_embedding.weight'] = load( - param_name_map["vocab_embedding"], config.embedding_sharding_dim - if config.use_parallel_embedding else -1) # vocab_embedding - - if mapping.is_last_pp_rank(): - v = load(param_name_map["lm_head"], -1, 1) if pad_vocab else load( - param_name_map["lm_head"], 0, 1) # lm_head - if v is None: - v = load(param_name_map["vocab_embedding"], - -1 if pad_vocab else 0).clone().detach() - - if pad_vocab: - v = torch.nn.functional.pad( - v, (0, 0, 0, vocab_size_padded - vocab_size), 'constant', 0) - v = split(v, mapping.tp_size, mapping.tp_rank) - weights['lm_head.weight'] = v - weights['transformer.ln_f.weight'] = load( - param_name_map["ln_f"]) # ln_f - - layers_range = mapping.pp_layers(num_hidden_layers) - for l in layers_range: - layer_idx = l - layers_range[0] - prefix = f'{layer_prefix}.{l}.' - tllm_prex = f'transformer.layers.{layer_idx}' - - # Attention - qkv_list = [] - for comp in ["q", "k", "v"]: - tp_size = kv_tp_size if comp != "q" else None - tp_rank = kv_tp_rank if comp != "q" else None - weight_part = load(prefix + f'{param_name_map["attention.qkv"]}.' + - comp + param_name_map["qkv_suffix"], - 0, - tp_size=tp_size, - tp_rank=tp_rank) - qkv_list.append(weight_part) - bias_part = load( - (prefix + f'{param_name_map["attention.qkv"]}.' + comp + - param_name_map["qkv_suffix"]).replace("weight", "bias"), - 0, - tp_size=tp_size, - tp_rank=tp_rank) - if bias_part is not None: - qkv_list.append(bias_part) - if len(qkv_list) == 3: - # No bias - weights[f'{tllm_prex}.attention.qkv.weight'] = torch.cat( - qkv_list, 0) - else: - weights[f'{tllm_prex}.attention.qkv.weight'] = torch.cat( - qkv_list[::2], 0) - weights[f'{tllm_prex}.attention.qkv.bias'] = torch.cat( - qkv_list[1::2], 0) - load_and_set(f'{tllm_prex}.attention.dense.weight', - prefix + param_name_map["attention.dense"], - 1) # attention.dense - - # MLP - if not moe_config.has_moe(): - load_and_set(f'{tllm_prex}.mlp.gate.weight', - prefix + param_name_map["mlp.gate"], 0) # mlp.gate - load_and_set(f'{tllm_prex}.mlp.proj.weight', - prefix + param_name_map["mlp.proj"], 1) # mlp.proj - load_and_set(f'{tllm_prex}.mlp.fc.weight', - prefix + param_name_map["mlp.fc"], 0) # mlp.fc - - else: - architecture = config.architecture.lower() - if "granite" not in architecture: - weights[f'{tllm_prex}.mlp.router.weight'] = load( - prefix + 'block_sparse_moe.gate.weight') - else: - weights[f'{tllm_prex}.mlp.router.weight'] = load( - prefix + 'block_sparse_moe.router.layer.weight') - rank_experts = list(range(moe_config.num_experts)) - if mapping.has_moe_ep(): - rank_experts = mapping.ep_experts(moe_config.num_experts) - - if "granite" not in architecture: - expert_weight_list = [] - for suffix in range(3): - tp_dim = -1 - if mapping.has_moe_tp(): - tp_dim = 1 if suffix == 1 else 0 - expert_weight_list.append( - torch.stack( - list( - load( - prefix + - f'block_sparse_moe.experts.{expert}.w{suffix + 1}.weight', - tp_dim=tp_dim, - is_expert_weights=True) - for expert in rank_experts))) - - w1 = expert_weight_list[0] - w2 = expert_weight_list[1] - w3 = expert_weight_list[2] - else: - w2 = load(prefix + f'block_sparse_moe.output_linear.weight', - is_expert_weights=True) #TODO: correct this - w13 = load(prefix + f'block_sparse_moe.input_linear.weight', - is_expert_weights=True) - - half_size = w13.shape[-2] // 2 - w1, w3 = w13.split(half_size, dim=-2) - w1 = w1[rank_experts[0]:rank_experts[-1] + 1] - w2 = w2[rank_experts[0]:rank_experts[-1] + 1] - w3 = w3[rank_experts[0]:rank_experts[-1] + 1] - - weights[f'{tllm_prex}.mlp.fc.weight'] = \ - torch.concat([w3, w1], dim=-2).contiguous() - weights[f'{tllm_prex}.mlp.proj.weight'] = w2.contiguous() - - load_and_set(f'{tllm_prex}.input_layernorm.weight', prefix + - param_name_map["input_layernorm"]) # input_layernorm - load_and_set(f'{tllm_prex}.post_layernorm.weight', prefix + - param_name_map["post_layernorm"]) # post_layernorm - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - logger.info(f'Weights loaded. Total time: {t}') - - return weights - - -def load_weights_from_gptq(quant_ckpt_path: str, config: LLaMAConfig): - logger.info('Loading weights from groupwise GPTQ LLaMA safetensors...') - weights = {} - tik = time.time() - - num_hidden_layers = config.num_hidden_layers - vocab_size = config.vocab_size - dtype = config.dtype - mapping = config.mapping - quant_algo = config.quantization.quant_algo - - gptq_llama = safetensors.safe_open(quant_ckpt_path, - framework="pt", - device=0) - gptq_prefix = "model." - gptq_suffix_list = [".qweight", ".qzeros", ".scales"] - gptq_key_list = [ - "embed_tokens.weight", # vocab_embedding - "lm_head.weight", # lm_head - "norm.weight", # ln_f - "self_attn.", # attention.qkv - "_proj", # qkv suffix - "self_attn.o_proj", # attention.dense - "mlp.up_proj", # mlp.gate - "mlp.down_proj", # mlp.proj - "mlp.gate_proj", # mlp.fc - "input_layernorm.weight", # input_layernorm - "post_attention_layernorm.weight", # post_layernorm - ] - split_sym = "." - - packer = torch.ops.trtllm.pack_int8_tensor_to_packed_int4 - preprocessor = torch.ops.trtllm.preprocess_weights_for_mixed_gemm - torch_dtype = str_dtype_to_torch(dtype) - - def load(key, no_prefix=0): - if no_prefix: - return gptq_llama.get_tensor(key) - else: - return gptq_llama.get_tensor(gptq_prefix + key) - - def torch_split(v, dim): - if v.shape[dim] % mapping.tp_size != 0: - logger.error( - "Current weight shape is invalid for mapping.tp_size=" + - str(mapping.tp_size)) - assert False, "Invalid TP size" - return v.split(v.shape[dim] // mapping.tp_size, - dim=dim)[mapping.tp_rank] - - def unpack_int32_into_int8(w_packed): - # Unpack inputs packed in int32/float32 into uint4 and store them in int8 format - w_packed_int4x2 = w_packed.contiguous().view(torch.uint8) - w_unpacked = torch.zeros(w_packed_int4x2.shape[0], - w_packed_int4x2.shape[1] * 2, - dtype=torch.int8) - w_unpacked[:, ::2] = w_packed_int4x2 % 16 - w_unpacked[:, 1::2] = w_packed_int4x2 // 16 - return w_unpacked.contiguous() - - def process_and_assign_weight(v: List[torch.Tensor], - tllm_prex: str, - quant_algo: QuantAlgo, - tp_dim: int = -1): - if tp_dim == -1: - qweight_int32, qzeros_int32, scales_fp16 = [ - item.cpu() for item in v - ] - else: - qweight_int32, qzeros_int32, scales_fp16 = [ - torch_split(item, tp_dim).cpu() for item in v - ] - - USE_UINT4_INPUT = 1 # Set to true if checkpoint store UINT4 weights - USE_UINT8_INPUT = 1 # Set to true if checkpoint store UINT8 weights - USE_GPTQ_FOR_LLAMA = 1 # GPTQ-for-LLaMA added 1 to zeros - - if quant_algo == QuantAlgo.W4A16_GPTQ: - # unpack inputs packed in int32 into int4 and store them in int8 format - qweight_unpacked_int8 = unpack_int32_into_int8( - qweight_int32.T).T.contiguous() - 8 - qweight_interleaved = preprocessor( - packer(qweight_unpacked_int8), torch.quint4x2, - torch.float16).view(torch.float16) - # zeros = zeros * scales - qzeros_unpacked_int32 = unpack_int32_into_int8(qzeros_int32) - if not USE_UINT4_INPUT: - # Correcting UINT4 values back to INT4 order - mask_negative = qzeros_unpacked_int32[qzeros_unpacked_int32 < 0] - mask_positive = qzeros_unpacked_int32[qzeros_unpacked_int32 >= - 0] - qzeros_unpacked_int32 = qzeros_unpacked_int32 + 16 * mask_negative - 16 * mask_positive - zeros_x_scales_fp16 = (-qzeros_unpacked_int32 + 8 * USE_UINT4_INPUT - - USE_GPTQ_FOR_LLAMA) * scales_fp16 - else: - # unpack inputs packed in int32 into int8 - qweight_unpacked_int8 = ( - qweight_int32.T.contiguous().view(torch.uint8).T.contiguous() - - 128).to(torch.int8) - qweight_interleaved = preprocessor(qweight_unpacked_int8, - torch.int8, torch.float16).view( - torch.float16) - qzeros_unpacked_int32 = qzeros_int32.view(torch.uint8) - zeros_x_scales_fp16 = (-qzeros_unpacked_int32 + - 128 * USE_UINT8_INPUT - - USE_GPTQ_FOR_LLAMA) * scales_fp16 - - zeros_x_scales_fp16 = zeros_x_scales_fp16.half() - - results = { - f'{tllm_prex}.weight': qweight_interleaved, - f'{tllm_prex}.weights_scaling_factor': scales_fp16, - f'{tllm_prex}.zero': zeros_x_scales_fp16, - } - return results - - # Load weights from GPTQ checkpoint into TRT-LLM module - # 1. vocab_embedding - v = load(gptq_key_list[0]) - if mapping.is_first_pp_rank(): - # tensorrt_llm_llama.vocab_embedding.weight.value = v.to( - # torch_dtype).cpu().numpy() - weights['transformer.vocab_embedding.weight'] = v.to(torch_dtype) - # 2. lm_head - v = load(gptq_key_list[1], "no_prefix") - if mapping.is_last_pp_rank(): - # tensorrt_llm_llama.lm_head.weight.value = torch_split( - # v, 0).to(torch_dtype).cpu().numpy() - if vocab_size % mapping.tp_size != 0: - # padding - vocab_size_padded = pad_vocab_size(vocab_size, mapping.tp_size) - pad_width = vocab_size_padded - vocab_size - v = torch.from_numpy( - np.pad(v.detach().cpu().numpy(), ((0, pad_width), (0, 0)), - 'constant', - constant_values=0)) - weights['lm_head.weight'] = torch_split(v, 0).to(torch_dtype) - - # 3. ln_f - v = load(gptq_key_list[2]) - if mapping.is_last_pp_rank(): - # tensorrt_llm_llama.ln_f.weight.value = v.to(torch_dtype).cpu().numpy() - weights['transformer.ln_f.weight'] = v.to(torch_dtype) - # 4. Weights inside each layer - layers_range = mapping.pp_layers(num_hidden_layers) - for l in layers_range: - layer_idx = l - layers_range[0] - prefix = "layers" + split_sym + str(layer_idx) + split_sym - logger.info(f'Process weights in layer: {layer_idx}') - # layer = tensorrt_llm_llama.layers[layer_idx] - tllm_prex = f'transformer.layers.{l-layers_range[0]}' - # 4.1 attention.qkv - qkv_weight_list = [] - for suf in gptq_suffix_list: - qkv_list = [] - for comp in ["q", "k", "v"]: - comp_part = load(prefix + gptq_key_list[3] + comp + - gptq_key_list[4] + suf) - comp_part = torch_split(comp_part, 1) - qkv_list.append(comp_part) - qkv_weight_list.append(torch.cat(qkv_list, dim=1)) - - # process_and_assign_weight(layer.attention.qkv, qkv_weight_list) - weights.update( - process_and_assign_weight(qkv_weight_list, - f'{tllm_prex}.attention.qkv', quant_algo)) - # 4.2 attention.dense - v = [load(prefix + gptq_key_list[5] + suf) for suf in gptq_suffix_list] - # process_and_assign_weight(layer.attention.dense, v, 0) - weights.update( - process_and_assign_weight(v, - f'{tllm_prex}.attention.dense', - quant_algo, - tp_dim=0)) - # 4.3 mlp.gate - v = [load(prefix + gptq_key_list[6] + suf) for suf in gptq_suffix_list] - # process_and_assign_weight(layer.mlp.gate, v, 1) - weights.update( - process_and_assign_weight(v, - f'{tllm_prex}.mlp.gate', - quant_algo, - tp_dim=1)) - # 4.4 mlp.proj - v = [load(prefix + gptq_key_list[7] + suf) for suf in gptq_suffix_list] - # process_and_assign_weight(layer.mlp.proj, v, 0) - weights.update( - process_and_assign_weight(v, - f'{tllm_prex}.mlp.proj', - quant_algo, - tp_dim=0)) - # 4.5 mlp.fc - v = [load(prefix + gptq_key_list[8] + suf) for suf in gptq_suffix_list] - # process_and_assign_weight(layer.mlp.fc, v, 1) - weights.update( - process_and_assign_weight(v, - f'{tllm_prex}.mlp.fc', - quant_algo, - tp_dim=1)) - # 4.6 input_layernorm - v = load(prefix + gptq_key_list[9]) - # layer.input_layernorm.weight.value = v.to(torch_dtype).cpu().numpy() - weights[f'{tllm_prex}.input_layernorm.weight'] = v.to(torch_dtype) - - # 4.7 post_layernorm - v = load(prefix + gptq_key_list[10]) - # layer.post_layernorm.weight.value = v.to(torch_dtype).cpu().numpy() - weights[f'{tllm_prex}.post_layernorm.weight'] = v.to(torch_dtype) - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - logger.info(f'Weights loaded. Total time: {t}') - - return weights - - -def load_weights_from_deepcompressor(quant_ckpt_path: str, config: LLaMAConfig): - logger.info( - 'Loading weights from DeepCompressor torch checkpoint for QServe W4A8 inference...' - ) - weights = {} - tik = time.time() - - per_group = config.quant_mode.has_per_group_scaling() - group_size = 128 if per_group else -1 - - num_hidden_layers = config.num_hidden_layers - vocab_size = config.vocab_size - dtype = config.dtype - mapping = config.mapping - torch_dtype = str_dtype_to_torch(dtype) - assert torch_dtype == torch.float16, "Currently QServe only supports float16" - - # weight - fake_quant_weights = torch.load(quant_ckpt_path + '/model.pt', - map_location='cpu', - weights_only=True) - # scale.0, scale.1, zero - quant_params = torch.load(quant_ckpt_path + '/scale.pt', - map_location='cpu', - weights_only=True) - - def load(key): - if 'zero' in key: - v = quant_params[key] - # https://github.com/mit-han-lab/qserve/blob/64ee627dfd747510809998d3592439f05a71ba31/scripts/ckpt_converter/checkpoint_converter.py#L99 - if v.min() < 0: - v = v + 8 - return v - if 'scale' in key: - return quant_params[key] - return fake_quant_weights[key] - - if per_group: - deepcompressor_suffix = [ - 'weight', 'weight.scale.0', 'weight.scale.1', 'weight.scaled_zero' - ] - qserve_suffix = ['weight', 's1_scales', 's2_scales', 's2_szeros'] - else: - deepcompressor_suffix = [ - 'weight', 'weight.scale.0', 'weight.scaled_zero' - ] - qserve_suffix = ['weight', 's1_scales', 's1_szeros'] - - def tp_split_tensor(v: torch.Tensor, dim): - if v.shape[dim] % mapping.tp_size != 0: - logger.error( - f"Current weight shape is invalid for mapping.tp_size={mapping.tp_size}" - ) - assert False, "Invalid TP size" - return v.split(v.shape[dim] // mapping.tp_size, - dim=dim)[mapping.tp_rank].contiguous() - - def tp_split_weight_and_params(v: List[torch.Tensor], column_linear: bool): - if per_group: - weight, s1_scales, s2_scales, s2_szeros = v - # weight (out_features, in_features) - # weight.scale.0 (out_features, 1, 1, 1) - # weight.scale.1 (out_features, 1, in_features/group_size, 1) - # weight.scaled_zero (out_features, 1, in_features/group_size, 1) - if column_linear: - weight = tp_split_tensor(weight, 0) - s1_scales = tp_split_tensor(s1_scales, 0) - s2_scales = tp_split_tensor(s2_scales, 0) - s2_szeros = tp_split_tensor(s2_szeros, 0) - else: - weight = tp_split_tensor(weight, 1) - s1_scales = s1_scales - s2_scales = tp_split_tensor(s2_scales, 2) - s2_szeros = tp_split_tensor(s2_szeros, 2) - return [weight, s1_scales, s2_scales, s2_szeros] - else: - weight, s1_scales, s1_szeros = v - # weight (out_features, in_features) - # weight.scale.0 (out_features, 1, 1, 1) - # weight.zero (out_features, 1, 1, 1) - if column_linear: - weight = tp_split_tensor(weight, 0) - s1_scales = tp_split_tensor(s1_scales, 0) - s1_szeros = tp_split_tensor(s1_szeros, 0) - else: - weight = tp_split_tensor(weight, 1) - s1_scales = s1_scales - s1_szeros = s1_szeros - return [weight, s1_scales, s1_szeros] - - def process_weight_and_params(v: List[torch.Tensor], tllm_prex: str): - if per_group: - weight, s1_scales, s2_scales, s2_szeros = v - qweight = qserve_quantize_weight_per_group(weight, s1_scales, - s2_scales, s2_szeros, - group_size) - qweight, s1_scales, s2_scales, s2_zeros = qserve_pack_reorder_per_group( - qweight, s1_scales, s2_scales, s2_szeros, group_size) - - return { - # Note: Linear modules in TRTLLM do not use the name 'qweight' - f'{tllm_prex}.{qserve_suffix[0]}': qweight, - f'{tllm_prex}.{qserve_suffix[1]}': s1_scales, - f'{tllm_prex}.{qserve_suffix[2]}': s2_scales, - f'{tllm_prex}.{qserve_suffix[3]}': s2_zeros, - } - else: - weight, s1_scales, s1_szeros = v - qweight = qserve_quantize_weight_per_channel( - weight, s1_scales, s1_szeros) - qweight, s1_scales, s1_szeros = qserve_pack_reorder_per_channel( - qweight, s1_scales, s1_szeros) - - return { - # Note: Linear modules in TRTLLM use the name 'weight' instead of 'qweight' - f'{tllm_prex}.{qserve_suffix[0]}': qweight, - f'{tllm_prex}.{qserve_suffix[1]}': s1_scales, - f'{tllm_prex}.{qserve_suffix[2]}': s1_szeros, - } - - # Load weights - # 1. vocab_embedding - v = load('model.embed_tokens.weight') - if mapping.is_first_pp_rank(): - weights['transformer.vocab_embedding.weight'] = v.to(torch_dtype) - - # 2. lm_head - v = load('lm_head.weight') - if mapping.is_last_pp_rank(): - if vocab_size % mapping.tp_size != 0: - # padding - vocab_size_padded = pad_vocab_size(vocab_size, mapping.tp_size) - pad_width = vocab_size_padded - vocab_size - v = torch.nn.functional.pad(v, (0, 0, 0, pad_width)) - weights['lm_head.weight'] = tp_split_tensor(v, 0).to(torch_dtype) - - # 3. ln_f - v = load('model.norm.weight') - if mapping.is_last_pp_rank(): - weights['transformer.ln_f.weight'] = v.to(torch_dtype) - - # 4. Weights inside each layer - layers_range = mapping.pp_layers(num_hidden_layers) - for layer_idx in layers_range: - prefix = f'model.layers.{layer_idx}' - logger.info(f'Processing weights in layer: {layer_idx}') - tllm_prex = f'transformer.layers.{layer_idx - layers_range[0]}' - - # 4.1 attention.qkv - qkv_list = [] - for comp in ["q", "k", "v"]: - v = [ - load(f'{prefix}.self_attn.{comp}_proj.{suffix}') - for suffix in deepcompressor_suffix - ] - v = tp_split_weight_and_params(v, column_linear=True) - qkv_list.append(v) - # Concat qkv - q, k, v = qkv_list - qkv = [ - torch.concat((q[i], k[i], v[i]), dim=0) - for i in range(len(deepcompressor_suffix)) - ] - weights.update( - process_weight_and_params(qkv, f'{tllm_prex}.attention.qkv')) - - # 4.2 attention.dense - v = [ - load(f'{prefix}.self_attn.o_proj.{suffix}') - for suffix in deepcompressor_suffix - ] - v = tp_split_weight_and_params(v, column_linear=False) - weights.update( - process_weight_and_params(v, f'{tllm_prex}.attention.dense')) - - # TODO: The naming here is tricky. - # The implementation of GatedMLP is act(fc(x)) * gate(x). - # However, the common convention is act(gate_proj(x)) * up_proj(x). - - # 4.3 mlp.gate - v = [ - load(f'{prefix}.mlp.up_proj.{suffix}') - for suffix in deepcompressor_suffix - ] - v = tp_split_weight_and_params(v, column_linear=True) - weights.update(process_weight_and_params(v, f'{tllm_prex}.mlp.gate')) - - # 4.4 mlp.fc - v = [ - load(f'{prefix}.mlp.gate_proj.{suffix}') - for suffix in deepcompressor_suffix - ] - v = tp_split_weight_and_params(v, column_linear=True) - weights.update(process_weight_and_params(v, f'{tllm_prex}.mlp.fc')) - - # 4.5 mlp.proj - v = [ - load(f'{prefix}.mlp.down_proj.{suffix}') - for suffix in deepcompressor_suffix - ] - v = tp_split_weight_and_params(v, column_linear=False) - weights.update(process_weight_and_params(v, f'{tllm_prex}.mlp.proj')) - - # 4.6 input_layernorm - v = load(f'{prefix}.input_layernorm.weight') - weights[f'{tllm_prex}.input_layernorm.weight'] = v.to(torch_dtype) - - # 4.7 post_layernorm - v = load(f'{prefix}.post_attention_layernorm.weight') - weights[f'{tllm_prex}.post_layernorm.weight'] = v.to(torch_dtype) - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - logger.info(f'Weights loaded. Total time: {t}') - - # TODO: All the RMSNorm weight, including ln_f, input_layernorm, post_layernorm, are actually all 1s - # Could implement a simplified module without weight - return weights - - -def load_torch_meta_ckpt(meta_ckpt_path: Path): - ''' - meta_ckpt_path: The format of meta_ckpt_path is like /consolidated.xx There are two possible cases: - 1. A file like /consolidated.xx.pth, loading it by torch.load directly - 2. A folder like /consolidated.xx/, need to load all weights in the folder. - ''' - file_path = meta_ckpt_path.parent / (meta_ckpt_path.name + ".pth") - if file_path.exists() and file_path.is_file(): - return torch.load(file_path, map_location="cpu") - else: - folder_path = meta_ckpt_path - assert folder_path.exists() and folder_path.is_dir() - - ckpts = list(Path(folder_path).glob("consolidated-*.pth")) - - all_weights = {} - for ckpt in ckpts: - _weight = torch.load(ckpt, map_location="cpu") - all_weights = all_weights | _weight - del _weight - return all_weights - - -def load_weights_from_meta_ckpt(meta_ckpt_dir: str, config: LLaMAConfig): - torch_dtype = str_dtype_to_torch(config.dtype) - mapping = config.mapping - use_fp8_rowwise = config.quant_mode.has_fp8_rowwise() - if config.quant_mode.has_any_quant() and not use_fp8_rowwise: - logger.error( - "Meta ckpts only support fp8_rowwise quantization currently.") - weights = {} - # Meta's recipe of not using fp8 rowwise for the first and last layer. - exclude_layers_id = [0, config.num_hidden_layers - 1] - - def gather_ckpts(ckpts): - gathered = {} - for k in ckpts[0]: - d = 0 - # TODO not sure should we consider tok here. - if any([n in k for n in ["wo", "w2"]]): - d = 1 - if "norm" in k or "rope" in k: # no TP - gathered[k] = ckpts[0][k].clone() - else: - gathered[k] = torch.cat([pt[k] for pt in ckpts], dim=d).clone() - return gathered - - def split_ckpt(ckpt, ranks_per_ckpt, ckpt_rank): - split_ckpt = {} - for k, v in ckpt.items(): - d = 0 - if any(n in k for n in - ["wo", "feed_forward.w2", "tok", "feed_forward.gate"]): - d = 1 - if "norm" in k or "rope" in k: # no TP - split_ckpt[k] = v.clone() - elif config.num_key_value_heads < mapping.tp_size and any( - n in k for n in ["wk", "wv"]): - assert mapping.tp_size % config.num_key_value_heads == 0 - # special case: we need to duplicate KV head - tmp = dup_kv_weight(v, config.num_key_value_heads, - mapping.tp_size) - split_ckpt[k] = torch.split(tmp, - tmp.shape[d] // ranks_per_ckpt, - dim=d)[ckpt_rank].clone() - else: - split_ckpt[k] = torch.split(v, - v.shape[d] // ranks_per_ckpt, - dim=d)[ckpt_rank].clone() - return split_ckpt - - def get_current_weights(num_ckpts): - if num_ckpts > mapping.tp_size: - # combine ckpts - assert (num_ckpts % mapping.tp_size) == 0 - nf = num_ckpts // mapping.tp_size - fs = nf * mapping.tp_rank - file_ids = list(range(fs, fs + nf)) - ckpts = [] - for f in file_ids: - ckpt = load_torch_meta_ckpt( - Path(meta_ckpt_dir, f"consolidated.{f:02d}")) - ckpts.append(ckpt) - return gather_ckpts(ckpts) - elif num_ckpts < mapping.tp_size: - # split ckpt - assert (mapping.tp_size % num_ckpts) == 0 - ranks_per_ckpt = mapping.tp_size // num_ckpts - ckpt_fid = mapping.tp_rank // ranks_per_ckpt - ckpt_rank = mapping.tp_rank % ranks_per_ckpt - nH_per_ckpt = config.num_attention_heads // num_ckpts - assert (nH_per_ckpt % ranks_per_ckpt) == 0 - ckpt = load_torch_meta_ckpt( - Path(meta_ckpt_dir, f"consolidated.{ckpt_fid:02d}")) - return split_ckpt(ckpt, ranks_per_ckpt, ckpt_rank) - - # num_ckpts == tensor_parallel, 1:1 mapping from files to TP - return load_torch_meta_ckpt( - Path(meta_ckpt_dir, f"consolidated.{mapping.tp_rank:02d}")) - - def permute(w, nH, d, dH): - # due to MQA's wk, nH*dH != d could be true - return w.view(nH, dH // 2, 2, d).transpose(1, 2).reshape(nH * dH, d) - - def extract_layer_idx(name): - ss = name.split('.') - for s in ss: - if s.isdigit(): - return s - return None - - if not hasattr(load_weights_from_meta_ckpt, "saved_embed"): - load_weights_from_meta_ckpt.saved_embed = None - - def combine_embeddings(embeds, num_ckpts): - if len(embeds) == 1: - return embeds[0] - assert [ - embeds[i].shape == embeds[i + 1].shape - for i in range(len(embeds) - 1) - ] - if embeds[0].shape[0] == config.vocab_size // num_ckpts: - merge_dim = 0 - elif embeds[0].shape[1] == config.hidden_size // num_ckpts: - merge_dim = 1 - else: - logger.error("Unable to infer embedding split dimension") - assert False, "Unable to infer embedding split dimension" - return torch.cat(embeds, dim=merge_dim) - - def gather_embedding(cur_embed, name: str, num_ckpts): - if mapping.tp_size == 1: - # even if num_ckpts > 1, get_current_weights will already have it gathered - return cur_embed - if load_weights_from_meta_ckpt.saved_embed is None: - embeds = [None] * num_ckpts - for i in range(num_ckpts): - ckpt = load_torch_meta_ckpt( - Path(meta_ckpt_dir, f"consolidated.{i:02d}")) - embeds[i] = ckpt[name] - embed = combine_embeddings(embeds, num_ckpts).to(torch_dtype) - load_weights_from_meta_ckpt.saved_embed = embed - - return load_weights_from_meta_ckpt.saved_embed - - logger.info('Loading weights from Meta LLaMA checkpoints ...') - tik = time.time() - - num_kv_heads = config.num_key_value_heads - mha_mode = (num_kv_heads == config.num_attention_heads) - - ckpts = list(Path(meta_ckpt_dir).glob("consolidated.*")) - num_ckpts = len(ckpts) - # llama/llama2 doesn't have MQA. So, simplifying loader logic by not worrying about it. - assert num_kv_heads > 1 or num_kv_heads >= num_ckpts, \ - f"We don't know how the {num_kv_heads} KV heads are distributed among {num_ckpts} checkpoints." - - tik = time.time() - ckpt = get_current_weights(num_ckpts) - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - logger.info(f'[{mapping.rank}] get_current_weights. Total time: {t}') - - head_size = config.hidden_size // config.num_attention_heads - layers_range = mapping.pp_layers(config.num_hidden_layers) - - for l in layers_range: - prefix = f'layers.{l}.attention.' - q_weight = permute(ckpt[prefix + 'wq.weight'].clone(), - nH=(config.num_attention_heads // mapping.tp_size), - d=config.hidden_size, - dH=head_size) - if num_kv_heads < mapping.tp_size and num_ckpts >= mapping.tp_size: - assert mapping.tp_size % num_kv_heads == 0 - assert False, "Not supported yet" - k_weight = permute(ckpt[prefix + 'wk.weight'].clone(), - nH=((num_kv_heads + mapping.tp_size - 1) // - mapping.tp_size), - d=config.hidden_size, - dH=head_size) - v_weight = ckpt[prefix + 'wv.weight'].clone() - - qkv_weight = torch.cat([q_weight, k_weight, v_weight], dim=0) - ckpt[prefix + 'qkv.weight'] = qkv_weight - - for k, v in tqdm(ckpt.items()): - dtype = torch_dtype if 'feed_forward.gate' not in k else torch.float32 - - v = v.to(dtype) - if "tok_embeddings" in k: - if not config.use_parallel_embedding: - v = gather_embedding(v, k, num_ckpts) - elif config.embedding_sharding_dim == 0: - # this needs a gather and then resplit along different dims - v = gather_embedding(v, k, num_ckpts) - v = split(v, mapping.tp_size, mapping.tp_rank, 0) - if mapping.is_first_pp_rank(): - weights['transformer.vocab_embedding.weight'] = v - elif "output" in k: - if mapping.is_last_pp_rank(): - if config.vocab_size % mapping.tp_size != 0: - # padding - vocab_size_padded = pad_vocab_size(config.vocab_size, - mapping.tp_size) - pad_width = vocab_size_padded - config.vocab_size - v = torch.from_numpy( - np.pad(v.detach().cpu().numpy(), - ((0, pad_width), (0, 0)), - 'constant', - constant_values=0)) - weights['lm_head.weight'] = v.detach().clone() - elif k == "norm.weight": - if mapping.is_last_pp_rank(): - weights['transformer.ln_f.weight'] = v - else: - # layer specific weights - layer_idx = extract_layer_idx(k) - if layer_idx is None or int(layer_idx) not in layers_range: - continue - - # Meta's recipe of not using fp8 rowwise for the first and last layer. - use_fp8_rowwise_in_layer = use_fp8_rowwise and ( - int(layer_idx) not in exclude_layers_id) - idx = int(layer_idx) - layers_range[0] - tllm_prex = f'transformer.layers.{idx}.' - - if 'attention_norm.weight' in k: - weights[tllm_prex + 'input_layernorm.weight'] = v - elif 'ffn_norm.weight' in k: - weights[tllm_prex + 'post_layernorm.weight'] = v - elif 'feed_forward.w3.weight' in k: - if use_fp8_rowwise_in_layer: - processed_torch_weights, torch_weight_scales = fp8_per_channel_quant_weight_gpu( - v, config.quantization.clamp_val) - weights[tllm_prex + - 'mlp.gate.weight'] = processed_torch_weights - weights[tllm_prex + - 'mlp.gate.per_channel_scale'] = torch_weight_scales - else: - weights[tllm_prex + 'mlp.gate.weight'] = v - elif 'feed_forward.w2.weight' in k: - if use_fp8_rowwise_in_layer: - processed_torch_weights, torch_weight_scales = fp8_per_channel_quant_weight_gpu( - v, config.quantization.clamp_val) - weights[tllm_prex + - 'mlp.proj.weight'] = processed_torch_weights - weights[tllm_prex + - 'mlp.proj.per_channel_scale'] = torch_weight_scales - else: - weights[tllm_prex + 'mlp.proj.weight'] = v - elif 'feed_forward.w1.weight' in k: - if use_fp8_rowwise_in_layer: - processed_torch_weights, torch_weight_scales = fp8_per_channel_quant_weight_gpu( - v, config.quantization.clamp_val) - weights[tllm_prex + - 'mlp.fc.weight'] = processed_torch_weights - weights[tllm_prex + - 'mlp.fc.per_channel_scale'] = torch_weight_scales - else: - weights[tllm_prex + 'mlp.fc.weight'] = v - elif 'attention.wo.weight' in k: - weights[tllm_prex + 'attention.dense.weight'] = v - elif 'attention.qkv.weight' in k: - weights[tllm_prex + 'attention.qkv.weight'] = v - elif 'feed_forward.gate' in k: - weights[tllm_prex + 'mlp.router.weight'] = v - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - logger.info(f'Weights loaded. Total time: {t}') - return weights diff --git a/tensorrt_llm/models/llama/model.py b/tensorrt_llm/models/llama/model.py deleted file mode 100644 index 2e272772adb0..000000000000 --- a/tensorrt_llm/models/llama/model.py +++ /dev/null @@ -1,614 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -import os -from typing import Optional, Union - -import transformers - -from ..._common import default_net -from ..._utils import pad_vocab_size -from ...functional import (AllReduceFusionOp, AllReduceParams, Tensor, - allgather, concat, constant, div, non_gated_version, - recv, send, unsqueeze) -from ...layers import (MOE, Attention, AttentionMaskType, ColumnLinear, - Embedding, FusedGatedMLP, GatedMLP, - PositionEmbeddingType, RmsNorm) -from ...lora_helper import LoraConfig, use_lora -from ...mapping import Mapping -from ...module import Module -from ...quantization.functional import fused_layernorm -from ..convert_utils import has_safetensors -from ..model_weights_loader import ModelWeightsLoader -from ..modeling_utils import (DecoderLayerList, DecoderModelForCausalLM, - QuantConfig) -from .config import LLaMAConfig -from .convert import (load_hf_llama, load_weights_from_deepcompressor, - load_weights_from_gptq, load_weights_from_hf_by_shard, - load_weights_from_hf_model, - load_weights_from_hf_safetensors, - load_weights_from_meta_ckpt) - - -class LLaMADecoderLayer(Module): - - def __init__(self, config: LLaMAConfig, layer_idx: int): - super().__init__() - self.layer_idx = layer_idx - layer_idx += config.layer_idx_offset - self.config = config - self.mapping = config.mapping - - if (self.config.use_input_layernorm_in_first_layer - and self.layer_idx == 0) or self.layer_idx > 0: - self.input_layernorm = RmsNorm(normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - dtype=config.dtype) - - layers_range = config.mapping.pp_layers(config.num_hidden_layers) - self.local_layer_idx = layer_idx - layers_range[0] - self.is_last_local_layer = layer_idx == layers_range[-1] - self.attention = Attention( - local_layer_idx=self.local_layer_idx, - hidden_size=config.hidden_size, - attention_head_size=config.head_size, - num_attention_heads=config.num_attention_heads, - num_kv_heads=config.num_key_value_heads, - max_position_embeddings=config.max_position_embeddings, - dtype=config.dtype, - attention_mask_type=AttentionMaskType.causal, - bias=config.attn_bias, - position_embedding_type=PositionEmbeddingType.rope_gpt_neox, - rotary_embedding_base=config.rotary_base, - rotary_embedding_scaling=config.rotary_scaling, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size, - tp_rank=config.mapping.tp_rank, - q_scaling=1.0 / config.attention_multiplier, - quant_mode=config.quant_mode, - cp_group=config.mapping.cp_group, - cp_size=config.mapping.cp_size, - cp_rank=config.mapping.cp_rank) - - mlp_hidden_size = config.hidden_size * 4 if config.intermediate_size is None else config.intermediate_size - - ClsMLP = GatedMLP - mlp_kwargs = {} - if config.moe.has_moe(): - ClsMLP = MOE - mlp_kwargs = { - "moe_config": config.moe, - "mapping": config.mapping, - } - self.mlp = ClsMLP(hidden_size=config.hidden_size, - ffn_hidden_size=mlp_hidden_size, - hidden_act=config.hidden_act, - dtype=config.dtype, - bias=config.mlp_bias, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size, - quant_mode=config.quant_mode, - **mlp_kwargs) - - self.post_layernorm = RmsNorm(normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - dtype=config.dtype) - - # Residual MLP that applies on pre-attention input - # TODO: change to self.has_residual_mlp = self.config.residual_mlp after ModelOpt quantize config is updated - self.has_residual_mlp = False - if hasattr(self.config, - "residual_mlp") and self.config.residual_mlp is True: - self.has_residual_mlp = True - - if self.has_residual_mlp: - self.residual_layernorm = RmsNorm( - normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - dtype=config.dtype) - ClsMLP = GatedMLP # TODO: may use FusedGatedMLP to further speedup - self.residual_mlp = ClsMLP( - hidden_size=config.hidden_size, - ffn_hidden_size=config. - hidden_size, # residual mlp uses hidden_size - hidden_act=non_gated_version( - config.hidden_act), # back to non-gated - dtype=config.dtype, - bias=config.mlp_bias, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size, - quant_mode=config.quant_mode) - - def forward(self, - hidden_states, - attention_mask=None, - use_cache=False, - spec_decoding_params=None, - kv_cache_params=None, - attention_params=None, - lora_layer_params=None, - next_layer_input_layernorm_args=None): - assert not ( - default_net().plugin_config.reduce_fusion and self.has_residual_mlp - ), "Custom all reduce and residual mlp can't be enabled at the same time." - assert not ( - default_net().plugin_config.reduce_fusion - and default_net().plugin_config.user_buffer - and default_net().plugin_config.pp_reduce_scatter - ), "User buffer reduce fusion enabled with PP reduce scatter is not supported now." - assert not ( - default_net().plugin_config.reduce_fusion - and default_net().plugin_config.norm_quant_fusion - ), "Reduce fusion and quant fusion can't be enabled at the same time." - if default_net( - ).plugin_config.reduce_fusion and self.local_layer_idx > 0: - hidden_states, residual = hidden_states - elif default_net( - ).plugin_config.norm_quant_fusion and self.local_layer_idx > 0: - hidden_states, residual = hidden_states - else: - residual = hidden_states - if (self.config.use_input_layernorm_in_first_layer - and self.layer_idx == 0) or self.layer_idx > 0: - hidden_states = self.input_layernorm(hidden_states) - - reduce_fusion_op = AllReduceFusionOp.NONE - if default_net().plugin_config.reduce_fusion: - if default_net().plugin_config.user_buffer: - if self.config.quant_mode.has_fp8_qdq(): - reduce_fusion_op = AllReduceFusionOp.RESIDUAL_RMS_NORM_QUANT_FP8 - elif self.config.quant_mode.has_nvfp4(): - assert default_net( - ).plugin_config.gemm_plugin == "nvfp4", "UB with nvfp4 model must use nvfp4 gemm plugin" - reduce_fusion_op = AllReduceFusionOp.RESIDUAL_RMS_NORM_QUANT_NVFP4 - else: - assert False, "UB must enabled with fp8 or nvfp4 model" - else: - reduce_fusion_op = AllReduceFusionOp.RESIDUAL_RMS_NORM - - reduce_fusion_scale = None - if default_net().plugin_config.reduce_fusion and default_net( - ).plugin_config.user_buffer: - if isinstance(self.mlp, FusedGatedMLP): - if self.config.quant_mode.has_fp8_qdq(): - reduce_fusion_scale = constant( - self.mlp.fused_fc.activation_scaling_factor.raw_value. - copy()) - elif self.config.quant_mode.has_nvfp4(): - reduce_fusion_scale = constant( - [1.0] / self.mlp.fused_fc. - activation_global_scaling_factor.raw_value) - else: - if self.config.quant_mode.has_fp8_qdq(): - reduce_fusion_scale = constant( - self.mlp.fc.activation_scaling_factor.raw_value.copy()) - elif self.config.quant_mode.has_nvfp4(): - reduce_fusion_scale = constant( - [1.0] / - self.mlp.fc.activation_global_scaling_factor.raw_value) - attention_output = self.attention( - hidden_states, - attention_mask=attention_mask, - use_cache=use_cache, - spec_decoding_params=spec_decoding_params, - kv_cache_params=kv_cache_params, - attention_params=attention_params, - lora_layer_params=lora_layer_params, - all_reduce_params=AllReduceParams( - fusion_op=reduce_fusion_op, - residual=residual, - norm_weight=self.post_layernorm.weight.value, - scale=reduce_fusion_scale, - eps=self.post_layernorm.eps)) - if use_cache: - attention_output, presents = attention_output - - if self.has_residual_mlp: - hidden_states = residual + attention_output - residual_attn = hidden_states - # arctic layer w/ residual mlp - - # residual mlp - hidden_states = self.residual_layernorm(hidden_states) - hidden_states = self.residual_mlp(hidden_states) - residual_mlp = residual_attn + hidden_states - - # parallel moe - # parallel moe layers applies on PRE-ATTENTION input residual, therefore achieving pre-fetching and better parallelism - hidden_states = self.post_layernorm(residual) - hidden_states = self.mlp(hidden_states, - lora_layer_params=lora_layer_params) - hidden_states = residual_mlp + hidden_states - else: - if default_net().plugin_config.reduce_fusion: - hidden_states, residual = attention_output - elif default_net().plugin_config.norm_quant_fusion: - hidden_states, residual_attn, act_per_block_scale = fused_layernorm( - input=attention_output, - normalized_shape=self.config.hidden_size, - residual=residual, - weight=self.post_layernorm.weight.value, - scale=div( - 1, self.mlp.fc.activation_global_scaling_factor.value) - if self.mlp.fc.activation_global_scaling_factor.value else - None, - eps=self.post_layernorm.eps, - p_dtype=self.config.dtype) - - hidden_states, residual_attn = ( - hidden_states, act_per_block_scale), residual_attn - assert isinstance(hidden_states, tuple) - else: - hidden_states = residual + attention_output * self.config.residual_multiplier - residual = hidden_states - hidden_states = self.post_layernorm(hidden_states) - if next_layer_input_layernorm_args is not None: - #this is middle layer - hidden_states = self.mlp( - hidden_states, - lora_layer_params=lora_layer_params, - all_reduce_params=AllReduceParams( - fusion_op=reduce_fusion_op, - residual=residual_attn - if default_net().plugin_config.norm_quant_fusion else - residual, - norm_weight=next_layer_input_layernorm_args[0], - scale=next_layer_input_layernorm_args[2], - eps=next_layer_input_layernorm_args[1])) - if default_net().plugin_config.norm_quant_fusion: - hidden_states, residual, act_per_block_scale = fused_layernorm( - input=hidden_states, - normalized_shape=self.config.hidden_size, - residual=residual_attn, - weight=next_layer_input_layernorm_args[0], - scale=div(1, next_layer_input_layernorm_args[2]) - if next_layer_input_layernorm_args[2] else None, - eps=next_layer_input_layernorm_args[1], - p_dtype=self.config.dtype) - hidden_states = (hidden_states, - act_per_block_scale), residual - else: - if default_net( - ).plugin_config.pp_reduce_scatter and self.is_last_local_layer and not self.mapping.is_last_pp_rank( - ): - hidden_states = self.mlp( - hidden_states, - lora_layer_params=lora_layer_params, - last_local_layer_residual=residual) - else: - if (default_net().plugin_config.reduce_fusion - and default_net().plugin_config.user_buffer): - hidden_states, residual = self.mlp( - hidden_states, - lora_layer_params=lora_layer_params, - all_reduce_params=AllReduceParams( - fusion_op=AllReduceFusionOp.LAST_PROCESS_FOR_UB, - residual=residual)) - else: - hidden_states = self.mlp( - hidden_states, lora_layer_params=lora_layer_params) - hidden_states = residual + hidden_states * self.config.residual_multiplier - if use_cache: - return (hidden_states, presents) - return hidden_states - - -class LLaMAModel(Module): - - def __init__(self, config: LLaMAConfig) -> None: - super().__init__() - - self.mapping = config.mapping - self.vocab_size = config.vocab_size - self.has_partial_lora_mask = config.has_partial_lora_mask - self.hidden_size = config.hidden_size - if self.mapping.is_first_pp_rank(): - self.vocab_embedding = Embedding(config.vocab_size, - config.hidden_size, - dtype=config.dtype) - self.embedding_multiplier = config.embedding_multiplier - - self.layers = DecoderLayerList(LLaMADecoderLayer, config) - - if config.fc_after_embed: - self.fc = ColumnLinear(2 * config.hidden_size, - config.hidden_size, - bias=True, - dtype=config.dtype, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size, - gather_output=True) - - if self.mapping.is_last_pp_rank(): - self.ln_f = None - if config.use_last_layernorm: - self.ln_f = RmsNorm(normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - dtype=config.dtype) - - def forward(self, - input_ids, - position_ids=None, - use_cache=False, - attention_mask=None, - spec_decoding_params=None, - kv_cache_params=None, - attention_params=None, - hidden_states=None, - hidden_states_for_embed=None, - prompt_embedding_table: Optional[Tensor] = None, - prompt_tasks: Optional[Tensor] = None, - prompt_vocab_size: Optional[Tensor] = None, - lora_params=None): - - ptuning_args = [ - prompt_embedding_table, prompt_tasks, prompt_vocab_size - ] if prompt_embedding_table is not None else [] - - if self.mapping.is_first_pp_rank(): - hidden_states = self.vocab_embedding(input_ids, *ptuning_args) - hidden_states *= self.embedding_multiplier - else: - hidden_states = recv(hidden_states, self.mapping.prev_pp_rank()) - if default_net().plugin_config.pp_reduce_scatter: - hidden_states = allgather(hidden_states, - self.mapping.tp_group, - gather_dim=0) - # reshape to (-1, hidden_size) - hidden_states = hidden_states.view( - concat([-1, self.hidden_size])) - - if hidden_states_for_embed is not None: - hidden_states = concat([hidden_states, hidden_states_for_embed], - dim=-1) - hidden_states = self.fc(hidden_states) - - if lora_params is not None and self.has_partial_lora_mask: - partial_lora_mask = input_ids > (self.vocab_size - 1) - lora_params.partial_lora_mask = unsqueeze(partial_lora_mask, -1) - - hidden_states = self.layers.forward( - hidden_states, - use_cache=use_cache, - attention_mask=attention_mask, - kv_cache_params=kv_cache_params, - attention_params=attention_params, - lora_params=lora_params, - spec_decoding_params=spec_decoding_params) - - if use_cache: - hidden_states, presents = hidden_states - - if self.mapping.is_last_pp_rank(): - if self.ln_f: - hidden_states = self.ln_f(hidden_states) - else: - hidden_states = send(hidden_states, self.mapping.next_pp_rank()) - - if use_cache: - return (hidden_states, tuple(presents)) - return hidden_states - - -class LLaMAForCausalLM(DecoderModelForCausalLM): - config_class = LLaMAConfig - - def __init__(self, config: LLaMAConfig): - transformer = LLaMAModel(config) - vocab_size_padded = pad_vocab_size(config.vocab_size, - config.mapping.tp_size) - if config.mapping.is_last_pp_rank(): - lm_head = ColumnLinear(config.hidden_size, - vocab_size_padded, - bias=False, - dtype=config.dtype, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size, - gather_output=True) - else: - lm_head = None - self.quant_mode = config.quant_mode - self.mapping = config.mapping - super().__init__(config, transformer, lm_head) - - @classmethod - def from_hugging_face( - cls, - hf_model_or_dir: Union[str, 'transformers.PreTrainedModel'], - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - ''' Create a LLaMAForCausalLM object from give parameters - ''' - import transformers - - load_by_shard = kwargs.pop('load_by_shard', False) - load_model_on_cpu = kwargs.pop('load_model_on_cpu', False) - quant_ckpt_path = kwargs.pop('quant_ckpt_path', None) - use_autoawq = kwargs.pop('use_autoawq', None) - if os.environ.get("TRTLLM_DISABLE_UNIFIED_CONVERTER" - ) is not None and not isinstance( - hf_model_or_dir, transformers.PreTrainedModel): - if "vila" in hf_model_or_dir or "llava" in hf_model_or_dir: - hf_model_or_dir = load_hf_llama(hf_model_or_dir, - load_model_on_cpu) - elif not load_by_shard and not has_safetensors( - hf_model_or_dir) and ( - quant_config is None - or not quant_config.quant_mode.has_any_quant()): - hf_model_or_dir = load_hf_llama(hf_model_or_dir, - load_model_on_cpu) - - assert hf_model_or_dir is not None - use_preloading = isinstance(hf_model_or_dir, - transformers.PreTrainedModel) - if use_preloading: - hf_model = hf_model_or_dir - hf_config_or_dir = hf_model.config - else: - hf_model_dir = hf_model_or_dir - hf_config_or_dir = hf_model_or_dir - - config = LLaMAConfig.from_hugging_face(hf_config_or_dir, - dtype=dtype, - mapping=mapping, - quant_config=quant_config, - **kwargs) - if config.remove_duplicated_kv_heads: - config.num_key_value_heads = config.num_key_value_heads // 2 - if os.environ.get("TRTLLM_DISABLE_UNIFIED_CONVERTER") is None: - custom_dict = {} - model_name = hf_model.config.model_type if use_preloading else hf_model_or_dir - if "llava" in model_name: - custom_dict = { - "transformer": "language_model.model", - "lm_head": "language_model.lm_head" - } - elif "vila" in model_name: - hf_model_dir += "/llm" - elif "exaone" in model_name.lower(): - custom_dict = { - "transformer": "transformer", - "layers": "h", - "vocab_embedding": "wte", - "lm_head": "lm_head", - "ln_f": "ln_f", - "attention": "attn.attention", - "dense": "out_proj", - "gate": "c_fc_1", - "proj": "c_proj", - "fc": "c_fc_0", - "input_layernorm": "ln_1", - "post_layernorm": "ln_2", - } - elif config.tie_word_embeddings: - custom_dict = {"lm_head": "model.embed_tokens"} - - if quant_ckpt_path is not None: - hf_model_dir = quant_ckpt_path - arg_dict = {"use_autoawq": True} if use_autoawq else {} - - loader = ModelWeightsLoader(hf_model_dir, custom_dict) - model = cls(config) - loader.generate_tllm_weights(model, arg_dict) - else: - if use_preloading: - assert not load_by_shard - weights = load_weights_from_hf_model(hf_model, config) - elif load_by_shard: - weights = load_weights_from_hf_by_shard(hf_model_dir, config) - elif has_safetensors( - hf_model_dir) and not config.quant_mode.has_any_quant(): - weights = load_weights_from_hf_safetensors(hf_model_dir, config) - elif quant_ckpt_path is not None: - if quant_config.quant_mode.is_int4_weight_only(): - weights = load_weights_from_gptq(quant_ckpt_path, config) - elif quant_config.quant_mode.is_qserve_w4a8(): - weights = load_weights_from_deepcompressor( - quant_ckpt_path, config) - else: - raise ValueError( - "quant_ckpt_path should be specified only for GPTQ or QServe" - ) - else: - hf_model = load_hf_llama(hf_model_dir, load_model_on_cpu) - weights = load_weights_from_hf_model(hf_model, config) - model = cls(config) - model.load(weights) - return model - - def default_plugin_config(self, **kwargs): - plugin_config = super().default_plugin_config(**kwargs) - if self.quant_mode.is_int4_weight_only_per_group(): - plugin_config.weight_only_groupwise_quant_matmul_plugin = 'auto' - return plugin_config - - @classmethod - def from_meta_ckpt(cls, - meta_ckpt_dir: str, - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - config = LLaMAConfig.from_meta_ckpt(meta_ckpt_dir, - dtype=dtype, - mapping=mapping, - quant_config=quant_config, - **kwargs) - - weights = load_weights_from_meta_ckpt(meta_ckpt_dir, config) - - model = cls(config) - model.load(weights) - return model - - @classmethod - def quantize( - cls, - hf_model_dir: str, - output_dir: str, - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - *, - device: str = 'cuda', - calib_dataset: str = 'cnn_dailymail', - calib_batches: int = 512, - calib_batch_size: int = 1, - calib_max_seq_length: int = 512, - random_seed: int = 1234, - tokenizer_max_seq_length: int = 2048, - **kwargs, - ): - if quant_config._requires_modelopt_quantization: - # modelopt quantization flow - super().quantize(hf_model_dir, - output_dir, - dtype=dtype, - mapping=mapping, - quant_config=quant_config, - device=device, - calib_dataset=calib_dataset, - calib_batches=calib_batches, - calib_batch_size=calib_batch_size, - calib_max_seq_length=calib_max_seq_length, - random_seed=random_seed, - tokenizer_max_seq_length=tokenizer_max_seq_length) - elif quant_config._requires_calibration: - # non-modelopt quantization flow - from . import convert - - config = LLaMAConfig.from_hugging_face(hf_model_dir, - dtype=dtype, - mapping=mapping, - quant_config=quant_config, - **kwargs) - trust_remote_code = kwargs.pop("trust_remote_code", True) - - convert.quantize(hf_model_dir, - output_dir, - config=config, - device=device, - calib_dataset=calib_dataset, - trust_remote_code=trust_remote_code, - calib_batches=calib_batches, - calib_max_seq_length=calib_max_seq_length) - else: - raise ValueError( - f"The quant_config ({quant_config}) does not require calibration, try {cls.__name__}.from_hugging_face instead." - ) - - def use_lora(self, lora_config: LoraConfig): - use_lora(self, lora_config) diff --git a/tensorrt_llm/models/mamba/__init__.py b/tensorrt_llm/models/mamba/__init__.py deleted file mode 100644 index 71bf6d298d42..000000000000 --- a/tensorrt_llm/models/mamba/__init__.py +++ /dev/null @@ -1,14 +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. diff --git a/tensorrt_llm/models/mamba/config.py b/tensorrt_llm/models/mamba/config.py deleted file mode 100644 index 84e2f189c46a..000000000000 --- a/tensorrt_llm/models/mamba/config.py +++ /dev/null @@ -1,313 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -import json -import os -from enum import Enum -from typing import List, Optional, Union - -import transformers - -from ...mapping import Mapping -from ..convert_utils import infer_dtype -from ..modeling_utils import PretrainedConfig, QuantConfig - - -class CheckpointType(str, Enum): - mistral_inference = "mistral_inference" - state_spaces = "state_spaces" - hf = "hf" - - -def get_ckpt_type(model_path): - hf_config = transformers.AutoConfig.from_pretrained(model_path, - trust_remote_code=True) - if hasattr(hf_config, "ssm_cfg") and hf_config.ssm_cfg: - return CheckpointType.state_spaces - if os.path.exists(os.path.join(model_path, "params.json")): - return CheckpointType.mistral_inference - return CheckpointType.hf - - -class MambaConfig(PretrainedConfig): - - def __init__(self, - *, - residual_in_fp32: bool = True, - pad_vocab_size_multiple: int = -1, - layer_types: List[str] = ["recurrent"], - **kwargs): - self.residual_in_fp32 = residual_in_fp32 - self.pad_vocab_size_multiple = pad_vocab_size_multiple - self.layer_types = layer_types - super().__init__(**kwargs) - - def to_dict(self): - output = super().to_dict() - # Serialize the fields added in MambaConfig - - return output - - def update(self, data_dict): - self.__dict__.update(data_dict) - - @classmethod - def from_hugging_face( - cls, - hf_config_or_dir: Union[str, 'transformers.PretrainedConfig'], - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - import transformers - - ckpt_type = get_ckpt_type(hf_config_or_dir) - - mamba_version = 'Mamba1' - if ckpt_type == CheckpointType.hf: - if isinstance(hf_config_or_dir, transformers.PretrainedConfig): - hf_config = hf_config_or_dir - else: - hf_config_dir = str(hf_config_or_dir) - - hf_config = transformers.AutoConfig.from_pretrained( - hf_config_dir, trust_remote_code=True) - - dtype = infer_dtype(dtype, getattr(hf_config, 'torch_dtype', None)) - - vocab_size = hf_config.vocab_size - pad_vocab_size_multiple = getattr(hf_config, - "pad_vocab_size_multiple", 1) - if vocab_size % pad_vocab_size_multiple != 0: - vocab_size += pad_vocab_size_multiple - ( - vocab_size % pad_vocab_size_multiple) - return cls(architecture="MambaForCausalLM", - dtype=dtype, - num_hidden_layers=hf_config.num_hidden_layers, - num_attention_heads=mapping.world_size, - hidden_size=hf_config.hidden_size, - intermediate_size=hf_config.intermediate_size, - vocab_size=vocab_size, - mamba_version=mamba_version, - hidden_act=hf_config.hidden_act, - rms_norm=hf_config.rms_norm, - residual_in_fp32=hf_config.residual_in_fp32, - pad_vocab_size_multiple=pad_vocab_size_multiple, - rnn_hidden_size=hf_config.intermediate_size, - rnn_conv_dim_size=hf_config.intermediate_size, - state_size=hf_config.state_size, - conv_kernel=hf_config.conv_kernel, - use_bias=hf_config.use_bias, - mapping=mapping, - quantization=quant_config, - **kwargs) - elif ckpt_type == CheckpointType.state_spaces: - - mamba_version = 'Mamba2' - if isinstance(hf_config_or_dir, transformers.PretrainedConfig): - hf_config = hf_config_or_dir - else: - hf_config_dir = str(hf_config_or_dir) - - hf_config = transformers.AutoConfig.from_pretrained( - hf_config_dir, trust_remote_code=True) - - dtype = infer_dtype(dtype, getattr(hf_config, 'torch_dtype', None)) - - vocab_size = hf_config.vocab_size - pad_vocab_size_multiple = getattr(hf_config, - "pad_vocab_size_multiple", 1) - if vocab_size % pad_vocab_size_multiple != 0: - vocab_size += pad_vocab_size_multiple - ( - vocab_size % pad_vocab_size_multiple) - assert hasattr(hf_config, - 'ssm_cfg') and hf_config.ssm_cfg['layer'] == 'Mamba2' - config = json.load( - open(os.path.join(hf_config_or_dir, 'config.json'))) - ssm_cfg = config.pop('ssm_cfg') - cfg_to_mamba_cfg = { - 'd_model': 'hidden_size', - 'n_layer': 'num_hidden_layers', - 'fused_add_norm': None, - 'tie_embeddings': None, - } - ssm_cfg_to_mamba_cfg = { - 'd_state': 'state_size', - 'd_conv': 'conv_kernel', - 'bias': 'use_bias', - 'headdim': 'head_dim', - 'ngroups': 'n_groups', - 'chunk_size': 'chunk_size', - 'rmsnorm': 'ssm_rmsnorm', - } - for k in cfg_to_mamba_cfg: - if k in config: - v = config.pop(k) - if cfg_to_mamba_cfg[k] is not None: - config[cfg_to_mamba_cfg[k]] = v - for k in ssm_cfg_to_mamba_cfg: - if k in ssm_cfg and ssm_cfg_to_mamba_cfg[k] is not None: - config[ssm_cfg_to_mamba_cfg[k]] = ssm_cfg[k] - - if 'expand' in config: - expand = config['expand'] - hf_config.intermediate_size = expand * config['hidden_size'] - else: - hf_config.intermediate_size = 2 * config['hidden_size'] - mamba2_default_cfg = { - 'n_groups': 1, - 'hidden_size': hf_config.d_model, - 'head_dim': 64, - 'chunk_size': 256, - 'state_size': 128, - } - hf_config.update(mamba2_default_cfg) - - conv_dim = hf_config.intermediate_size + 2 * hf_config.n_groups * hf_config.state_size - ssm_rmsnorm = getattr(hf_config, "ssm_rmsnorm", hf_config.rms_norm) - mamba2_cfg = { - 'rnn_head_size': hf_config.head_dim, - 'rnn_conv_dim_size': conv_dim, - 'ngroups': hf_config.n_groups, - 'chunk_size': hf_config.chunk_size, - 'ssm_rmsnorm': ssm_rmsnorm, - } - hf_config.update(mamba2_cfg) - - return cls(architecture="MambaForCausalLM", - dtype=dtype, - num_hidden_layers=hf_config.n_layer, - num_attention_heads=mapping.world_size - if mapping is not None else 1, - hidden_size=hf_config.hidden_size, - intermediate_size=hf_config.intermediate_size, - vocab_size=vocab_size, - mamba_version=mamba_version, - hidden_act=hf_config.hidden_act, - rms_norm=hf_config.rms_norm, - residual_in_fp32=hf_config.residual_in_fp32, - pad_vocab_size_multiple=pad_vocab_size_multiple, - rnn_hidden_size=hf_config.intermediate_size, - rnn_conv_dim_size=hf_config.rnn_conv_dim_size, - state_size=hf_config.state_size, - conv_kernel=hf_config.conv_kernel, - use_bias=hf_config.use_bias, - mapping=mapping, - quantization=quant_config, - rnn_head_size=hf_config.rnn_head_size, - ngroups=hf_config.ngroups, - chunk_size=hf_config.chunk_size, - ssm_rmsnorm=hf_config.ssm_rmsnorm, - **kwargs) - elif ckpt_type == CheckpointType.mistral_inference: - mamba_version = 'Mamba2' - - config = json.load( - open(os.path.join(hf_config_or_dir, 'params.json'))) - cfg_to_mamba_cfg = { - 'dim': 'hidden_size', - 'n_layers': 'num_hidden_layers', - 'n_groups': 'n_groups', - 'fused_add_norm': None, - 'tie_embeddings': None, - 'model_type': None, - } - for k in cfg_to_mamba_cfg: - if k in config: - v = config.pop(k) - if cfg_to_mamba_cfg[k] is not None: - config[cfg_to_mamba_cfg[k]] = v - - config['architecture'] = 'MambaForCuasualLM' - config['dtype'] = dtype - config['num_attention_heads'] = mapping.world_size - - hf_config = MambaConfig(**config) - mamba2_default_cfg = { - 'n_groups': 8, - 'hidden_size': 4096, - 'head_dim': 64, - 'chunk_size': 256, - 'state_size': 128, - 'conv_kernel': 4, - 'use_bias': False - } - - hf_config.update(mamba2_default_cfg) - conv_dim = hf_config.intermediate_size + 2 * hf_config.n_groups * hf_config.state_size - ssm_rmsnorm = getattr(hf_config, "ssm_rmsnorm", hf_config.rms_norm) - mamba2_cfg = { - 'rnn_head_size': hf_config.head_dim, - 'rnn_conv_dim_size': conv_dim, - 'ngroups': hf_config.n_groups, - 'chunk_size': hf_config.chunk_size, - 'ssm_rmsnorm': ssm_rmsnorm, - } - hf_config.update(mamba2_cfg) - - if 'expand' in config: - expand = config['expand'] - hf_config.intermediate_size = expand * hf_config.hidden_size - else: - hf_config.intermediate_size = 2 * hf_config.hidden_size - vocab_size = hf_config.vocab_size - pad_vocab_size_multiple = getattr(hf_config, - "pad_vocab_size_multiple", 1) - if vocab_size % pad_vocab_size_multiple != 0: - vocab_size += pad_vocab_size_multiple - ( - vocab_size % pad_vocab_size_multiple) - - return cls( - architecture="MambaForCausalLM", - dtype=dtype, - num_hidden_layers=hf_config.num_hidden_layers, - num_attention_heads=mapping.world_size, - hidden_size=hf_config.hidden_size, - intermediate_size=hf_config.intermediate_size, - # num_key_value_heads=num_key_value_heads, - vocab_size=vocab_size, - mamba_version=mamba_version, - hidden_act=hf_config.hidden_act, - rms_norm=hf_config.rms_norm, - residual_in_fp32=hf_config.residual_in_fp32, - pad_vocab_size_multiple=pad_vocab_size_multiple, - rnn_hidden_size=hf_config.intermediate_size, - rnn_conv_dim_size=hf_config.rnn_conv_dim_size, - state_size=hf_config.state_size, - conv_kernel=hf_config.conv_kernel, - use_bias=hf_config.use_bias, - mapping=mapping, - quantization=quant_config, - rnn_head_size=hf_config.rnn_head_size, - ngroups=hf_config.n_groups, - chunk_size=hf_config.chunk_size, - ssm_rmsnorm=hf_config.ssm_rmsnorm, - **kwargs) - else: - pass - - if isinstance(hf_config_or_dir, transformers.PretrainedConfig): - hf_config = hf_config_or_dir - else: - hf_config_dir = str(hf_config_or_dir) - - hf_config = transformers.AutoConfig.from_pretrained( - hf_config_dir, trust_remote_code=True) - - vocab_size = hf_config.vocab_size - pad_vocab_size_multiple = getattr(hf_config, "pad_vocab_size_multiple", - 1) - if vocab_size % pad_vocab_size_multiple != 0: - vocab_size += pad_vocab_size_multiple - (vocab_size % - pad_vocab_size_multiple) diff --git a/tensorrt_llm/models/mamba/convert.py b/tensorrt_llm/models/mamba/convert.py deleted file mode 100644 index f55bda43c7c2..000000000000 --- a/tensorrt_llm/models/mamba/convert.py +++ /dev/null @@ -1,245 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -import copy -import re -import time -from pathlib import Path -from typing import Union - -import torch - -import tensorrt_llm -from tensorrt_llm.models.convert_utils import (iterate_shard_files, - load_state_dict) - - -def get_weight(config, prefix, dtype): - return config[prefix + '.weight'].to(dtype).detach() - - -def get_bias(config, prefix, dtype): - if (prefix + '.bias') in config: - return config[prefix + '.bias'].to(dtype).detach() - return None - - -def get_weight_and_bias(config, prefix, dtype_w, dtype_b): - return get_weight(config, prefix, - dtype_w), get_bias(config, prefix, dtype_b) - - -def split(v, tp_size, idx, dim=0): - assert v.shape[dim] % tp_size == 0 - split_size = v.shape[dim] // tp_size - if tp_size == 1: - return v - return torch.split(v, split_size, dim=dim)[idx] - - -def rename_hf_to_tllm(name: str): - """ Rename a HF parameter name by the corresponding TRT-LLM style name. """ - # remove model - if 'model.' in name: - name = name.replace('model.', '') - - # change layer name - if 'embeddings.' in name: - name = name.replace('embeddings', 'vocab_embedding') - elif 'embedding.' in name: - name = name.replace('embedding', 'vocab_embedding') - norm_pattern = r'\d\.norm\.' - if 'mixer.' in name: - name = name.replace('mixer.', 'ssm.') - elif re.search(norm_pattern, name): - name = name.replace('norm.', 'input_layernorm.') - elif 'norm_f.' in name: - name = name.replace('norm_f.', 'ln_f.') - - # Parameter names in ssm layers - if 'A_log' in name: - name = name.replace('A_log', 'A') - elif 'dt_proj.bias' in name: - name = name.replace('dt_proj.bias', 'dt_bias') - return name - - -def convert_hf_mamba(hf_mamba, dtype='float32'): - weights = {} - tik = time.time() - - model_params = dict(hf_mamba.named_parameters()) - dtype = getattr(torch, dtype) - - # Parameter names in mamba block - for l in range(hf_mamba.config.num_hidden_layers): - # ssm layer - prefix = f'backbone.layers.{l}.mixer.' - tllm_prex = f'backbone.layers.{l}.ssm.' - for layer in ['conv1d', 'x_proj', 'dt_proj', 'out_proj']: - dtype_b = torch.float32 if layer == 'dt_proj' else dtype - weight, bias = get_weight_and_bias(model_params, prefix + layer, - dtype, dtype_b) - if layer == 'conv1d': - weight = weight.unsqueeze(3) - tllm_weight_name = tllm_prex + layer + '.weight' - tllm_bias_name = tllm_prex + ('dt_bias' if layer == 'dt_proj' else - layer + '.bias') - weights[tllm_weight_name] = weight - if bias is not None: - weights[tllm_bias_name] = bias - # in_proj - weight, bias = get_weight_and_bias(model_params, prefix + 'in_proj', - dtype, dtype) - in_proj_weights = torch.split(weight, weight.size(0) // 2, dim=0) - tllm_weight_name = tllm_prex + 'in_proj.weight' - weights[tllm_weight_name.replace('proj', 'proj_x')] = in_proj_weights[0] - weights[tllm_weight_name.replace('proj', 'proj_z')] = in_proj_weights[1] - if bias is not None: - in_proj_biases = torch.split(bias, bias.size(0) // 2, dim=0) - tllm_bias_name = tllm_prex + 'in_proj.bias' - weights[tllm_bias_name.replace('proj', - 'proj_x')] = in_proj_biases[0] - weights[tllm_bias_name.replace('proj', - 'proj_x')] = in_proj_biases[1] - - # A and D - Aparam = model_params[prefix + 'A_log'].float().detach() - Aparam = Aparam.permute(1, 0).contiguous() - weights[tllm_prex + 'A'] = -torch.exp(Aparam) - weights[tllm_prex + 'D'] = model_params[prefix + 'D'].float().detach() - # norm - prefix = f'backbone.layers.{l}.norm' - tllm_prex = f'backbone.layers.{l}.input_layernorm.' - weight, bias = get_weight_and_bias(model_params, prefix, dtype, dtype) - weights[tllm_prex + 'weight'] = weight - if bias is not None: - weights[tllm_prex + 'bias'] = bias - - # others - for layer in ['backbone.embeddings', 'backbone.norm_f']: - weight, bias = get_weight_and_bias(model_params, layer, dtype, dtype) - layer = layer.replace('embeddings', 'vocab_embedding') - layer = layer.replace('norm_f', 'ln_f') - weights[layer + '.weight'] = weight - if bias is not None: - weights[layer + '.bias'] = bias - weights['lm_head.weight'], _ = get_weight_and_bias(model_params, - 'backbone.embeddings', - dtype, dtype) - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Weights loaded. Total time: {t}') - return weights - - -def convert_from_hf_checkpoint(mamba_config: dict, model_dir: Union[str, Path]): - - print('Loading weights from HF Mamba...') - tik = time.time() - - tp_rank = mamba_config.mapping.tp_rank - tp_size = mamba_config.mapping.tp_size - d_inner = mamba_config.rnn_hidden_size - d_state = mamba_config.state_size - dtype = mamba_config.dtype - mamba_version = mamba_config.mamba_version - weights = {} - if isinstance(dtype, str): - dtype = tensorrt_llm.str_dtype_to_torch(dtype) - - for model_file in iterate_shard_files(model_dir, 0): - # logger.debug(f'Loading file {str(model_file)}...') - model_params = load_state_dict(model_file, dtype=dtype) - for name, param in model_params.items(): - # logger.debug(f'Converting weight {name}...') - tllm_name = rename_hf_to_tllm(name) - param = param.detach().cpu() - if 'A_log' in name: - param = -torch.exp(param.float()) - if mamba_version == 'Mamba1': - param = param.permute(1, 0).contiguous() - elif 'D' in name: - param = param.float() - elif 'dt_proj.bias' in name: - param = param.float() - elif 'dt_bias' in name: - param = param.float() - elif 'conv1d.weight' in name: - param = param.unsqueeze(3) - - # split in_proj in Mamba1 - if 'in_proj' in name and mamba_version == 'Mamba1': - in_proj_params = torch.split(param, param.size(0) // 2, dim=0) - weights[tllm_name.replace('proj', 'proj_x')] = in_proj_params[0] - weights[tllm_name.replace('proj', 'proj_z')] = in_proj_params[1] - elif 'in_proj' in name and mamba_version == 'Mamba2': - nheads = d_inner // mamba_config.rnn_head_size - ngroups = mamba_config.ngroups - - in_proj_z, in_proj_x, in_proj_b, in_proj_c, in_proj_dt = torch.split( - param, [ - d_inner, d_inner, ngroups * d_state, ngroups * d_state, - nheads - ], - dim=0) - in_proj_z = split(in_proj_z, tp_size, tp_rank, dim=0) - in_proj_x = split(in_proj_x, tp_size, tp_rank, dim=0) - in_proj_b = split(in_proj_b, tp_size, tp_rank, dim=0) - in_proj_c = split(in_proj_c, tp_size, tp_rank, dim=0) - in_proj_dt = split(in_proj_dt, tp_size, tp_rank, dim=0) - in_proj = torch.concat( - [in_proj_z, in_proj_x, in_proj_b, in_proj_c, in_proj_dt]) - weights[tllm_name] = in_proj.contiguous() - elif 'conv1d' in name and mamba_version == 'Mamba2': - ngroups = mamba_config.ngroups - conv_x, conv_b, conv_c = torch.split( - param, [d_inner, ngroups * d_state, ngroups * d_state], - dim=0) - conv_x = split(conv_x, tp_size, tp_rank, dim=0) - conv_b = split(conv_b, tp_size, tp_rank, dim=0) - conv_c = split(conv_c, tp_size, tp_rank, dim=0) - conv = torch.concat([conv_x, conv_b, conv_c]) - weights[tllm_name] = conv.contiguous() - elif any(keyword in name for keyword in ( - 'mixer.norm.weight', - 'A_log', - 'D', - 'dt_proj.bias', - 'dt_bias', - )) and mamba_version == 'Mamba2': - weights[tllm_name] = split(param, tp_size, tp_rank, dim=0) - elif 'out_proj' in name and mamba_version == 'Mamba2': - weights[tllm_name] = split(param, tp_size, tp_rank, - dim=1).contiguous() - else: - weights[tllm_name] = param - del model_params - - # lm_head - emb = weights['backbone.vocab_embedding.weight'] - if 'lm_head.weight' not in weights or weights['lm_head.weight'].data_ptr( - ) == emb.data_ptr(): - weights['lm_head.weight'] = copy.deepcopy(emb) - if mamba_version == 'Mamba2': - weights['lm_head.weight'] = split(weights['lm_head.weight'], - tp_size, - tp_rank, - dim=0) - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Weights loaded. Total time: {t}') - return weights diff --git a/tensorrt_llm/models/mamba/model.py b/tensorrt_llm/models/mamba/model.py deleted file mode 100644 index 0ccd4abd3a38..000000000000 --- a/tensorrt_llm/models/mamba/model.py +++ /dev/null @@ -1,471 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -import os -from collections import OrderedDict -from typing import List, Optional, Union - -import tensorrt as trt -from transformers import AutoModelForCausalLM - -from ..._common import default_net -from ..._utils import str_dtype_to_trt -from ...functional import (Tensor, arange, cast, concat, expand, - gather_last_token_logits, shape, unsqueeze) -from ...layers import ColumnLinear, Embedding, LayerNorm, Mamba, Mamba2, RmsNorm -from ...mapping import Mapping -from ...module import Module, ModuleList -from ...plugin import current_all_reduce_helper -from ..generation_mixin import GenerationMixin -from ..modeling_utils import PretrainedConfig, PretrainedModel, QuantConfig -from .config import MambaConfig -from .convert import convert_from_hf_checkpoint, convert_hf_mamba - - -class MambaLayer(Module): - - def __init__(self, config: PretrainedConfig, layer_idx: int): - super().__init__() - self.dtype = config.dtype - self.residual_in_fp32 = config.residual_in_fp32 - n_layer = config.num_hidden_layers - self.last_layer = layer_idx == n_layer - 1 - - if config.mamba_version == 'Mamba1': - assert config.mapping.tp_size == 1, "Mamba1 can not support tensor parallelism." - self.ssm = Mamba(config.hidden_size, - config.rnn_hidden_size, - d_state=config.state_size, - d_conv=config.conv_kernel, - bias=config.use_bias, - dtype=config.dtype) - elif config.mamba_version == 'Mamba2': - self.ssm = Mamba2(config.hidden_size, - config.rnn_hidden_size, - d_state=config.state_size, - d_conv=config.conv_kernel, - headdim=config.rnn_head_size, - ngroups=config.ngroups, - chunk_size=config.chunk_size, - bias=config.use_bias, - rmsnorm=config.ssm_rmsnorm, - dtype=config.dtype, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size) - if config.rms_norm: - self.input_layernorm = RmsNorm(normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - dtype=config.dtype) - else: - self.input_layernorm = LayerNorm( - normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - dtype=config.dtype) - - def forward(self, - hidden_states: Tensor, - residual: Tensor, - conv_state: Tensor, - ssm_state: Tensor, - host_request_types: Tensor, - last_token_ids: Tensor, - host_context_lengths: Optional[Tensor] = None, - slot_mapping: Optional[Tensor] = None, - conv_indices: Optional[Tensor] = None): - - hidden_states = self.input_layernorm(hidden_states) - - ssm_out, present_conv, present_ssm = self.ssm( - hidden_states, - conv_state=conv_state, - ssm_state=ssm_state, - host_request_types=host_request_types, - last_token_ids=last_token_ids, - host_context_lengths=host_context_lengths, - slot_mapping=slot_mapping, - conv_indices=conv_indices) - if self.residual_in_fp32: - residual = residual + cast(ssm_out, 'float32') - hidden_states = cast(residual, self.dtype) - else: - residual = residual + ssm_out - hidden_states = residual - - if self.last_layer: - return hidden_states, None, present_conv, present_ssm - else: - return hidden_states, residual, present_conv, present_ssm - - -class MambaModel(Module): - - def __init__(self, config: PretrainedConfig): - super().__init__() - self.d_conv = config.conv_kernel - self.d_inner = config.rnn_hidden_size // config.mapping.tp_size - n_layer = config.num_hidden_layers - self.residual_in_fp32 = config.residual_in_fp32 - if config.vocab_size % config.pad_vocab_size_multiple != 0: - config.vocab_size += config.pad_vocab_size_multiple - ( - config.vocab_size % config.pad_vocab_size_multiple) - self.vocab_embedding = Embedding(config.vocab_size, - config.hidden_size, - dtype=config.dtype) - self.layers = ModuleList( - [MambaLayer(config, i) for i in range(n_layer)]) - if config.rms_norm: - self.ln_f = RmsNorm(normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - dtype=config.dtype) - else: - self.ln_f = LayerNorm(normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - dtype=config.dtype) - - def forward(self, - input_ids, - conv_states, - ssm_states, - host_request_types, - last_token_ids, - host_context_lengths, - slot_mapping: Optional[Tensor] = None): - hidden_states = self.vocab_embedding(input_ids) - - # Get conv state indices - indices = None - if not default_net().plugin_config.mamba_conv1d_plugin: - batch_size = shape(input_ids, 0) - indices = expand( - unsqueeze(arange(0, self.d_conv - 1, dtype='int32'), 0), - concat([batch_size, self.d_conv - 1])) - offsets = expand(unsqueeze(last_token_ids, 1), - concat([batch_size, self.d_conv - 1])) - indices = unsqueeze(indices + offsets, 1) - indices = expand( - indices, concat([batch_size, self.d_inner, self.d_conv - 1])) - - residual = cast(hidden_states, - 'float32') if self.residual_in_fp32 else hidden_states - hidden_values = [hidden_states, residual] - present_convs, present_ssms = [], [] - for layer, past_conv, past_ssm in zip(self.layers, conv_states, - ssm_states): - hidden_values = layer(hidden_values[0], hidden_values[1], past_conv, - past_ssm, host_request_types, last_token_ids, - host_context_lengths, slot_mapping, indices) - present_convs.append(hidden_values[2]) - present_ssms.append(hidden_values[3]) - hidden_states = hidden_values[0] - hidden_states = self.ln_f(hidden_states) - return hidden_states, tuple(present_convs), tuple(present_ssms) - - -class MambaForCausalLM(PretrainedModel): - config_class = MambaConfig - - def __init__(self, config: PretrainedConfig): - super().__init__(config) - dtype = config.dtype - logits_dtype = config.logits_dtype - if isinstance(dtype, str): - self.dtype = str_dtype_to_trt(dtype) - else: - assert isinstance(dtype, trt.DataType) - self.dtype = dtype - - self.config = config - self.mamba_version = config.mamba_version - self.d_inner = config.rnn_hidden_size // config.mapping.tp_size - self.d_conv = config.conv_kernel - self.d_state = config.state_size - self.conv_dim = config.rnn_conv_dim_size // config.mapping.tp_size - self.gather_context_logits = False - - if isinstance(logits_dtype, str): - self._logits_dtype = str_dtype_to_trt(logits_dtype) - else: - assert isinstance(logits_dtype, trt.DataType) - self._logits_dtype = logits_dtype - - self.backbone = MambaModel(config) - self.lm_head = ColumnLinear(config.hidden_size, - config.vocab_size, - bias=False, - dtype=dtype, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size, - gather_output=True) - - def __post_init__(self): - return - - def forward(self, - input_ids, - conv_states, - ssm_states, - host_request_types, - last_token_ids, - last_token_ids_for_logits, - host_context_lengths, - slot_mapping: Optional[Tensor] = None): - hidden_states, present_convs, present_ssms = self.backbone( - input_ids, conv_states, ssm_states, host_request_types, - last_token_ids, host_context_lengths, slot_mapping) - - if not self.gather_context_logits: - hidden_states = gather_last_token_logits( - hidden_states, last_token_ids_for_logits, - default_net().plugin_config.remove_input_padding) - - lm_logits = self.lm_head(hidden_states) - lm_logits.mark_output('logits', self._logits_dtype) - if not default_net().plugin_config.paged_state: - for i, present_conv in enumerate(present_convs): - present_conv.mark_output(f'present_conv_state_{i}', self.dtype) - for i, present_ssm in enumerate(present_ssms): - present_ssm.mark_output(f'present_rnn_state_{i}', self.dtype) - - return (lm_logits, present_convs, present_ssms) - - def prepare_inputs( - self, - max_batch_size, - max_input_len, - max_seq_len, - max_num_tokens, - use_cache, - max_beam_width: int = 1, - opt_num_tokens: int = None, - opt_batch_size: int = 0, - prompt_embedding_table_size: int = 0, - max_draft_len: int = 0, - gather_context_logits: bool = False, - lora_target_modules: List[str] = None, - speculative_decoding_draft_tokens_external: bool = False): - '''@brief: Prepare inputs Tensors for the model, the given sizes are used to determine the - ranges of the dimensions of when using TRT dynamic shapes. - - @return: a list contains values which can be fed into the self.forward() - ''' - assert speculative_decoding_draft_tokens_external == False, "Speculative decoding is not supported in Mamba" - assert max_beam_width == 1, "We don't support beam search for the Mamba model." - - remove_input_padding = default_net().plugin_config.remove_input_padding - use_gemm_plugin = default_net().plugin_config.gemm_plugin - paged_state = default_net().plugin_config.paged_state - multiple_profiles = default_net().plugin_config.multiple_profiles - use_mamba_conv1d_plugin = default_net( - ).plugin_config.mamba_conv1d_plugin - - self.gather_context_logits = gather_context_logits - mapping = self.config.mapping - - # basic inputs - enable_ctx_gen_opt_profiles = GenerationMixin.has_ctx_gen_opt_profiles( - use_gemm_plugin=use_gemm_plugin, - use_mamba_conv1d_plugin=use_mamba_conv1d_plugin, - remove_input_padding=remove_input_padding, - paged_state=paged_state) - - num_profiles, ranges = GenerationMixin.get_profiles_ranges( - max_batch_size=max_batch_size, - max_beam_width=max_beam_width, - max_input_len=max_input_len, - max_num_tokens=max_num_tokens, - max_draft_len=max_draft_len, - opt_batch_size=opt_batch_size, - opt_num_tokens=opt_num_tokens, - enable_ctx_gen_opt_profiles=enable_ctx_gen_opt_profiles, - multiple_profiles=multiple_profiles) - - if remove_input_padding: - assert use_mamba_conv1d_plugin, "mamba_conv1d_plugin is needed to support remove_input_padding" - input_ids = Tensor(name='input_ids', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([ - ('num_tokens', ranges['num_tokens_range']), - ])) - else: - input_ids = Tensor(name='input_ids', - dtype=trt.int32, - shape=[-1, -1], - dim_range=OrderedDict([ - ('batch_size_beam_width', - ranges['bb_range']), - ('input_len', ranges['inlen_range']), - ])) - if mapping.tp_size > 1: - current_all_reduce_helper().set_workspace_tensor( - mapping, num_profiles) - - # recurrent inputs - conv_states = [] - ssm_states = [] - if use_mamba_conv1d_plugin: - conv_state_dim_range = OrderedDict([ - ('batch_size', ranges['bb_range']), - ('kernel_size', [self.d_conv - 1] * num_profiles), - ('dim_size', [self.conv_dim] * num_profiles), - ]) - else: - conv_state_dim_range = OrderedDict([ - ('batch_size', ranges['bb_range']), - ('dim_size', [self.conv_dim] * num_profiles), - ('kernel_size', [self.d_conv - 1] * num_profiles), - ]) - - if self.mamba_version == 'Mamba2': - headdim = self.config.rnn_head_size - nheads = self.d_inner // headdim - ssm_state_dim_range = OrderedDict([ - ('batch_size', ranges['bb_range']), - ('head_size', [nheads] * num_profiles), - ('state_size', [self.d_state] * num_profiles), - ('headdim_size', [headdim] * num_profiles), - ]) - ssm_state_shape = [-1, nheads, self.d_state, headdim] - else: - ssm_state_dim_range = OrderedDict([ - ('batch_size', ranges['bb_range']), - ('state_size', [self.d_state] * num_profiles), - ('dim_size', [self.d_inner] * num_profiles), - ]) - ssm_state_shape = [-1, self.d_state, self.d_inner] - one_dim_range = OrderedDict([ - ('buffer_count', [1] * num_profiles), - ]) - - for i in range(self.config.num_hidden_layers): - if default_net().plugin_config.paged_state: - conv_state = Tensor(name=f'conv_state_ptr_{i}', - dtype=str_dtype_to_trt('int64'), - shape=[1], - dim_range=one_dim_range) - - ssm_state = Tensor(name=f'rnn_state_ptr_{i}', - dtype=str_dtype_to_trt('int64'), - shape=[1], - dim_range=one_dim_range) - else: - if use_mamba_conv1d_plugin: - conv_state = Tensor( - name=f'past_conv_state_{i}', - dtype=self.dtype, - shape=[-1, self.d_conv - 1, self.conv_dim], - dim_range=conv_state_dim_range) - else: - conv_state = Tensor( - name=f'past_conv_state_{i}', - dtype=self.dtype, - shape=[-1, self.conv_dim, self.d_conv - 1], - dim_range=conv_state_dim_range) - - ssm_state = Tensor(name=f'past_rnn_state_{i}', - dtype=self.dtype, - shape=ssm_state_shape, - dim_range=ssm_state_dim_range) - - conv_states.append(conv_state) - ssm_states.append(ssm_state) - - host_request_types = Tensor( - name='host_request_types', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([('batch_size', ranges['bb_range'])]), - ) - - if remove_input_padding: - host_context_lengths = Tensor( - name='host_context_lengths', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([('batch_size', ranges['bb_range'])]), - ) - else: - host_context_lengths = None - - last_token_ids = Tensor( - name='last_token_ids', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([ - ('batch_size', ranges['bbd_range']), - ]), - ) - last_token_ids_for_logits = None - if not gather_context_logits: - last_token_ids_for_logits = last_token_ids - - return_dict = { - 'input_ids': input_ids, - 'conv_states': conv_states, - 'ssm_states': ssm_states, - 'host_request_types': host_request_types, - 'last_token_ids': last_token_ids, - 'last_token_ids_for_logits': last_token_ids_for_logits, - 'host_context_lengths': host_context_lengths, - } - - if default_net().plugin_config.paged_state: - slot_mapping = Tensor( - name='slot_mapping', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([('batch_size', ranges['bb_range'])]), - ) - return_dict['slot_mapping'] = slot_mapping - - return return_dict - - @classmethod - def from_hugging_face( - cls, - hf_model_or_dir: Union[str, 'transformers.PreTrainedModel'], - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - import transformers - - assert hf_model_or_dir is not None - use_preloading = isinstance(hf_model_or_dir, - transformers.PreTrainedModel) - if use_preloading: - hf_model = hf_model_or_dir - hf_config_or_dir = hf_model.config - else: - hf_model_dir = hf_model_or_dir - hf_config_or_dir = hf_model_or_dir - config = MambaConfig.from_hugging_face(hf_config_or_dir, - dtype=dtype, - mapping=mapping, - quant_config=quant_config, - **kwargs) - - if not os.path.exists(hf_model_dir): - hf_model = AutoModelForCausalLM.from_pretrained( - hf_model_dir, dtype="auto", trust_remote_code=True) - - assert isinstance(hf_model, transformers.PreTrainedModel) - weights = convert_hf_mamba(hf_model, dtype) - else: - weights = convert_from_hf_checkpoint(config, hf_model_dir) - - model = cls(config) - model.load(weights) - - return model diff --git a/tensorrt_llm/models/medusa/__init__.py b/tensorrt_llm/models/medusa/__init__.py deleted file mode 100644 index 2a36ca922710..000000000000 --- a/tensorrt_llm/models/medusa/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-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. diff --git a/tensorrt_llm/models/medusa/config.py b/tensorrt_llm/models/medusa/config.py deleted file mode 100644 index 0da2777244a4..000000000000 --- a/tensorrt_llm/models/medusa/config.py +++ /dev/null @@ -1,114 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import json -from typing import Optional, Union - -from ...mapping import Mapping -from ..convert_utils import infer_dtype -from ..llama.config import LLaMAConfig -from ..modeling_utils import PretrainedConfig, QuantConfig -from ..qwen.config import QWenConfig - - -# Medusa-specific config is stored and retrieved from GenericMedusaConfig. -class MedusaConfig(PretrainedConfig): - - def __init__(self, - *, - num_medusa_heads: int = 4, - num_medusa_layers: int = 1, - max_draft_len: int = 63, - **kwargs): - - model_type = str(kwargs.get('model_type', '')).lower() - generic_medusa_config = QWenConfig if 'qwen' in model_type else LLaMAConfig - self.config = generic_medusa_config(**kwargs) - - # Add objects - self.config.num_medusa_heads = num_medusa_heads - self.config.num_medusa_layers = num_medusa_layers - self.config.max_draft_len = max_draft_len - - def to_dict(self): - output = self.config.to_dict() - output['num_medusa_heads'] = self.config.num_medusa_heads - output['num_medusa_layers'] = self.config.num_medusa_layers - output['max_draft_len'] = self.config.max_draft_len - return output - - # Specialization to redirect accesses to self.config - def __getattr__(self, name): - return getattr(self.config, name) - - def __getstate__(self): - return self.__dict__ - - def __setstate__(self, state): - self.__dict__.update(state) - - @classmethod - def from_hugging_face( - cls, - hf_config_or_dir: Union[str, 'transformers.PretrainedConfig'], - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - import transformers - - trust_remote_code = kwargs.pop('trust_remote_code', True) - speculative_config_or_dir = kwargs.pop('speculative_model_dir', None) - speculative_config = kwargs.pop("speculative_config", None) - - if isinstance(hf_config_or_dir, transformers.PretrainedConfig): - hf_config = hf_config_or_dir - else: - hf_config_dir = str(hf_config_or_dir) - - hf_config = transformers.AutoConfig.from_pretrained( - hf_config_dir, trust_remote_code=trust_remote_code) - dtype = infer_dtype(dtype, getattr(hf_config, 'torch_dtype', None)) - - if hasattr(hf_config, "medusa"): - # is modelOpt ckpt - num_medusa_heads = hf_config.medusa["num_medusa_heads"] - num_medusa_layers = hf_config.medusa["num_medusa_layers"] - else: - config_file = speculative_config_or_dir / "config.json" - with open(config_file) as fp: - config = json.load(fp) - - num_medusa_heads = speculative_config.num_medusa_heads if speculative_config is not None else config.get( - 'num_medusa_heads', None) - num_medusa_layers = config.get('medusa_num_layers', None) - - return cls(architecture="MedusaForCausalLM", - dtype=dtype, - num_hidden_layers=hf_config.num_hidden_layers, - num_attention_heads=hf_config.num_attention_heads, - hidden_size=hf_config.hidden_size, - intermediate_size=hf_config.intermediate_size, - num_key_value_heads=hf_config.num_key_value_heads, - vocab_size=hf_config.vocab_size, - position_embedding_type='rope_gpt_neox', - max_position_embeddings=hf_config.max_position_embeddings, - hidden_act=hf_config.hidden_act, - norm_epsilon=hf_config.rms_norm_eps, - mapping=mapping, - quantization=quant_config, - num_medusa_heads=num_medusa_heads, - num_medusa_layers=num_medusa_layers, - **kwargs) diff --git a/tensorrt_llm/models/medusa/model.py b/tensorrt_llm/models/medusa/model.py deleted file mode 100644 index 525c4ec391f0..000000000000 --- a/tensorrt_llm/models/medusa/model.py +++ /dev/null @@ -1,267 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-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. - -import json -from typing import Optional, Union - -from transformers import AutoModelForCausalLM - -from tensorrt_llm._utils import numpy_to_torch -from tensorrt_llm.models.llama.model import LLaMAForCausalLM -from tensorrt_llm.models.medusa.weight import load_medusa_hf -from tensorrt_llm.models.qwen.model import QWenForCausalLM - -from ..._common import default_net -from ..._utils import pad_vocab_size -from ...functional import ACT2FN, stack -from ...layers import ColumnLinear -from ...mapping import Mapping -from ...module import Module, ModuleList -from ..modeling_utils import PretrainedModel, QuantConfig -from .config import MedusaConfig -from .weight import convert_hf_llama - - -class MedusaLayer(Module): - - def __init__( - self, - hidden_size, - hidden_act="silu", - dtype=None, - mapping=Mapping(), - ): - super().__init__() - self.linear = ColumnLinear(hidden_size, - hidden_size, - dtype=dtype, - tp_group=mapping.tp_group, - tp_size=mapping.tp_size, - gather_output=True) - self.hidden_act = hidden_act - - def forward(self, x): - return x + ACT2FN[self.hidden_act](self.linear(x)) - - -class MedusaHead(Module): - - def __init__( - self, - num_layers, - hidden_size, - vocab_size, - hidden_act="silu", - dtype=None, - mapping=Mapping(), - ): - super().__init__() - self.medusa_layers = ModuleList([ - MedusaLayer(hidden_size=hidden_size, - hidden_act=hidden_act, - dtype=dtype, - mapping=mapping) for _ in range(num_layers) - ]) - self.lm_head = ColumnLinear(hidden_size, - vocab_size, - bias=False, - dtype=dtype, - tp_group=mapping.tp_group, - tp_size=mapping.tp_size, - gather_output=True) - return - - def forward(self, x): - hidden_states = x - for layer in self.medusa_layers: - hidden_states = layer(hidden_states) - return self.lm_head(hidden_states) - - -# MedusaForCausalLm is a thin wrapper that picks parent class for GenericMedusaForCausalLM. -# All medusa functionality is defined in GenericMedusaForCausalLM. -class MedusaForCausalLm(PretrainedModel): - config_class = MedusaConfig - - def __init__(self, config: MedusaConfig): - super().__init__(config) - - BaseLM = QWenForCausalLM if hasattr( - config, - "model_type") and "qwen" in config.model_type else LLaMAForCausalLM - - class GenericMedusaForCausalLM(BaseLM): - - def __init__(self, config: MedusaConfig): - super().__init__(config) - self.num_medusa_heads = config.num_medusa_heads - self.num_medusa_layers = config.num_medusa_layers - self.hidden_size = config.hidden_size - self.vocab_size = config.vocab_size - vocab_size_padded = pad_vocab_size(self.vocab_size, - config.mapping.tp_size) - self.medusa_heads = ModuleList([ - MedusaHead(num_layers=self.num_medusa_layers, - hidden_size=config.hidden_size, - vocab_size=vocab_size_padded, - hidden_act=config.hidden_act, - dtype=config.dtype, - mapping=config.mapping) - for _ in range(self.num_medusa_heads) - ]) - self.max_medusa_token_len = config.max_draft_len - - def forward(self, *args, **kwargs): - output_original = True - hidden_states = super().forward(*args, **kwargs) - - if kwargs['use_cache']: - if default_net().plugin_config.paged_kv_cache: - lm_logits, hidden_states, _ = hidden_states - else: - lm_logits, presents, hidden_states = hidden_states - - if self.mapping.is_last_pp_rank(): - medusa_logits = [] - for i in range(self.num_medusa_heads): - medusa_logits.append( - self.medusa_heads[i](hidden_states)) - # [num_medusa_heads, batch_size, num_medusa_tokens + 1, padded_vocab_size]. - # Remove padding [num_medusa_heads, batch_size * num_medusa_tokens + 1, padded_vocab_size]. - medusa_logits = stack(medusa_logits, dim=0) - medusa_logits.mark_output('medusa_logits', - self.config.logits_dtype) - else: - hidden_states.mark_output('hidden_states_output', - self.config.dtype) - - if kwargs['use_cache'] and default_net( - ).plugin_config.paged_kv_cache == False: - if self.mapping.is_last_pp_rank(): - if output_original: - return (medusa_logits, lm_logits, presents) - return (medusa_logits, presents) - return (hidden_states, presents) - else: - if self.mapping.is_last_pp_rank(): - if output_original: - return medusa_logits, lm_logits - return medusa_logits - return hidden_states - - def prepare_inputs(self, *args, **kwargs): - kwargs['speculative_decoding_draft_tokens_external'] = False - kwargs['max_draft_len'] = self.max_medusa_token_len - return super().prepare_inputs(*args, **kwargs) - - self.model = GenericMedusaForCausalLM(config) - - # Specialization to redirect accesses to self.model - def __getattribute__(self, name): - if name == 'model' or '__' in name: - return object.__getattribute__(self, name) - else: - model = object.__getattribute__(self, 'model') - return model.__getattribute__(name) - - # Override specialized __setattr__ defined in Module - def __setattr__(self, name, value) -> None: - object.__setattr__(self, name, value) - - @classmethod - def from_hugging_face( - cls, - hf_model_or_dir: Union[str, 'transformers.PreTrainedModel'], - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - import transformers - - assert hf_model_or_dir is not None - speculative_model_dir = kwargs.get('speculative_model_dir', None) - - use_preloading = isinstance(hf_model_or_dir, - transformers.PreTrainedModel) - if use_preloading: - hf_model = hf_model_or_dir - hf_config_or_dir = hf_model.config - else: - hf_model_dir = hf_model_or_dir - hf_config_or_dir = hf_model_or_dir - - config = MedusaConfig.from_hugging_face(hf_config_or_dir, - dtype=dtype, - mapping=mapping, - quant_config=quant_config, - **kwargs) - - # ModelOpt ckpt has combined base model and Medusa-head - is_modelopt_ckpt = True if not speculative_model_dir else False - - if not use_preloading: - trust_remote_code = kwargs.pop('trust_remote_code', True) - - if is_modelopt_ckpt: - hf_model = LLaMAForCausalLM.from_hugging_face( - hf_model_dir, - dtype, - mapping=mapping, - quant_config=quant_config, - **kwargs) - else: - hf_model = AutoModelForCausalLM.from_pretrained( - hf_model_dir, - dtype="auto", - trust_remote_code=trust_remote_code) - - assert isinstance(hf_model, transformers.PreTrainedModel) - - if is_modelopt_ckpt: - weights = { - name: numpy_to_torch(param.raw_value) - for name, param in hf_model.named_parameters() - } - else: - weights = convert_hf_llama( - hf_model, - config.mapping, - dtype='float16', - use_parallel_embedding=config.use_parallel_embedding) - - model = cls(config) - - if is_modelopt_ckpt: - num_medusa_heads = config.config.num_medusa_heads - num_medusa_layers = config.config.num_medusa_layers - speculative_model_dir = hf_model_or_dir - else: - config_file = speculative_model_dir / "config.json" - with open(config_file) as fp: - model_config = json.load(fp) - - num_medusa_heads = kwargs[ - 'speculative_config'].num_medusa_heads if 'speculative_config' in kwargs else model_config.get( - 'medusa_num_heads', None) - num_medusa_layers = model_config.get('medusa_num_layers', None) - medusa_weights = load_medusa_hf(medusa_path=speculative_model_dir, - num_medusa_heads=num_medusa_heads, - num_medusa_layers=num_medusa_layers, - mapping=mapping, - dtype="float16", - is_modelopt_ckpt=is_modelopt_ckpt) - weights.update(medusa_weights) - model.load(weights) - return model diff --git a/tensorrt_llm/models/medusa/weight.py b/tensorrt_llm/models/medusa/weight.py deleted file mode 100644 index 6964dbdd3ee5..000000000000 --- a/tensorrt_llm/models/medusa/weight.py +++ /dev/null @@ -1,648 +0,0 @@ -import copy -import functools -import time -from collections import defaultdict -from pathlib import Path - -import numpy as np -import torch -import torch.nn as nn -from tqdm import tqdm -from transformers.models.llama.modeling_llama import LlamaDecoderLayer -from transformers.pytorch_utils import Conv1D - -from tensorrt_llm._utils import str_dtype_to_torch -from tensorrt_llm.logger import logger -from tensorrt_llm.mapping import Mapping -from tensorrt_llm.models.convert_utils import (dup_kv_weight, generate_int8, - smooth_gemm, - smooth_gemm_fc1_gate, split, - split_matrix_tp, split_qkv_tp) - - -def get_tllm_linear_weight(weight, - prefix, - bias=None, - use_weight_only=False, - plugin_weight_only_quant_type=torch.int8, - postfix='weight'): - results = {} - if use_weight_only: - v = weight.t().contiguous().cpu() - processed_torch_weights, torch_weight_scales = \ - torch.ops.trtllm.symmetric_quantize_last_axis_of_batched_matrix( - v, plugin_weight_only_quant_type) - results[prefix + postfix] = processed_torch_weights - results[prefix + 'per_channel_scale'] = torch_weight_scales - else: - results[prefix + postfix] = weight.contiguous() - - if bias is not None: - results[prefix + 'bias'] = bias - - return results - - -def load_medusa_hf(medusa_path: str, - num_medusa_heads: int, - num_medusa_layers: int, - mapping=Mapping(), - dtype='float32', - use_weight_only=False, - plugin_weight_only_quant_type=None, - is_modelopt_ckpt=False): - logger.info("Loading Medusa heads' weights ...") - - if is_modelopt_ckpt: - from safetensors.torch import load_file - state_dict = {} - for filename in sorted(Path(medusa_path).glob("*.safetensors")): - print(f"Loading the weights of Medusa heads from {filename}") - state_dict.update(load_file(filename)) - else: - is_ckpt_safetensors = False - - ckpt_file = Path(medusa_path) / "medusa_lm_head.pt" - if not ckpt_file.exists(): - ckpt_file = Path(medusa_path) / "medusa_lm_head.safetensors" - is_ckpt_safetensors = True - - if is_ckpt_safetensors: - logger.info("Safetensors Found ...") - from safetensors.torch import load_file - state_dict = load_file(ckpt_file) - else: - state_dict = torch.load(ckpt_file, map_location="cpu") - - torch_dtype = str_dtype_to_torch(dtype) - weights = {} - - prefix = "medusa_heads." if is_modelopt_ckpt else "" - for h in range(num_medusa_heads): - for l in range(num_medusa_layers): - w = state_dict[f"{prefix}{h}.{l}.linear.weight"].clone().to( - torch_dtype) - - split_v = split(w, mapping.tp_size, mapping.tp_rank) - weights.update( - get_tllm_linear_weight( - split_v, f'medusa_heads.{h}.medusa_layers.{l}.linear.', - None, use_weight_only, plugin_weight_only_quant_type)) - - b = state_dict[f"{prefix}{h}.{l}.linear.bias"].clone().to( - torch_dtype) - - weights['medusa_heads.{}.medusa_layers.{}.linear.bias'.format( - h, l)] = split(b, mapping.tp_size, mapping.tp_rank) - - lm = state_dict[f"{prefix}{h}.{num_medusa_layers}.weight"].clone().to( - torch_dtype) # LM Head - - weights['medusa_heads.{}.lm_head.weight'.format(h)] = split( - lm, mapping.tp_size, mapping.tp_rank) - - # scaling factors - if is_modelopt_ckpt: - scaling_dtype = str_dtype_to_torch("float32") - weights[f'medusa_heads.{h}.medusa_layers.{l}.linear.activation_scaling_factor'] = \ - state_dict[f"{prefix}{h}.{l}.linear.input_scale"].clone().to(scaling_dtype) - - weights[f'medusa_heads.{h}.medusa_layers.{l}.linear.weights_scaling_factor'] = \ - state_dict[f"{prefix}{h}.{l}.linear.weight_scale"].clone().to(scaling_dtype) - - weights['medusa_heads.{}.lm_head.activation_scaling_factor'.format(h)] = \ - state_dict[f"{prefix}{h}.{num_medusa_layers}.input_scale"].clone().to(scaling_dtype) - - weights['medusa_heads.{}.lm_head.weights_scaling_factor'.format(h)] = \ - state_dict[f"{prefix}{h}.{num_medusa_layers}.weight_scale"].clone().to(scaling_dtype) - - return weights - - -@torch.no_grad() -def smooth_llama_model(model, scales, alpha, llama_qkv_para, llama_smoother): - # Smooth the activation and weights with smoother = $\diag{s}$ - for name, module in model.named_modules(): - if not isinstance(module, LlamaDecoderLayer): - continue - # qkv_proj - layer_name_q = name + ".self_attn.q_proj" - layer_name_k = name + ".self_attn.k_proj" - layer_name_v = name + ".self_attn.v_proj" - layer_name_qkv = name + ".self_attn.qkv_proj" - - weight = torch.cat([ - module.self_attn.q_proj.weight, module.self_attn.k_proj.weight, - module.self_attn.v_proj.weight - ], - dim=0) - - smoother = smooth_gemm(weight, scales[layer_name_q]["x"], - module.input_layernorm.weight, None, alpha) - - scales[layer_name_qkv]["x"] = scales[layer_name_q]["x"] / smoother - scales[layer_name_qkv]["w"] = weight.abs().max(dim=1)[0] - scales[layer_name_qkv]["y"] = torch.cat([ - scales[layer_name_q]["y"], scales[layer_name_k]["y"], - scales[layer_name_v]["y"] - ], - dim=0) - - # see transpose_weights function - llama_qkv_para[layer_name_qkv] = weight.transpose(0, 1) - - # ================================================================= - layer_name = name + ".self_attn.o_proj" - smoother = smooth_gemm(module.self_attn.o_proj.weight, - scales[layer_name]["x"], None, None, alpha) - llama_smoother[layer_name] = smoother.float() - - scales[layer_name]["x"] = scales[layer_name]["x"] / smoother - scales[layer_name]["w"] = module.self_attn.o_proj.weight.abs().max( - dim=1)[0] - - # ================================================================== - fc1_layer_name = name + ".mlp.gate_proj" - gate_layer_name = name + ".mlp.up_proj" - - smoother = smooth_gemm_fc1_gate(module.mlp.gate_proj.weight, - module.mlp.up_proj.weight, - scales[fc1_layer_name]["x"], - module.post_attention_layernorm.weight, - None, alpha) - - scales[fc1_layer_name]["x"] = scales[fc1_layer_name]["x"] / smoother - scales[fc1_layer_name]["w"] = module.mlp.gate_proj.weight.abs().max( - dim=1)[0] - - scales[gate_layer_name]["x"] = scales[gate_layer_name]["x"] / smoother - scales[gate_layer_name]["w"] = module.mlp.up_proj.weight.abs().max( - dim=1)[0] - - # ================================================================== - layer_name = name + ".mlp.down_proj" - smoother = smooth_gemm(module.mlp.down_proj.weight, - scales[layer_name]["x"], None, None, alpha) - llama_smoother[layer_name] = smoother.float() - scales[layer_name]["x"] = scales[layer_name]["x"] / smoother - scales[layer_name]["w"] = module.mlp.down_proj.weight.abs().max( - dim=1)[0] - - -@torch.no_grad() -def capture_activation_range(model, - tokenizer, - dataset, - num_samples=512, - seq_len=512): - model.eval() - device = next(model.parameters()).device - act_scales = defaultdict(lambda: {"x": None, "y": None, "w": None}) - - tokenizer.pad_token = tokenizer.eos_token - - def stat_tensor(name, tensor, act_scales, key): - hidden_dim = tensor.shape[-1] - tensor = tensor.view(-1, hidden_dim).abs().detach() - comming_max = torch.max(tensor, dim=0)[0].float() - - if act_scales[name][key] is None: - act_scales[name][key] = comming_max - else: - act_scales[name][key] = torch.max(act_scales[name][key], - comming_max) - - def stat_input_hook(m, x, y, name): - if isinstance(x, tuple): - x = x[0] - stat_tensor(name, x, act_scales, "x") - stat_tensor(name, y, act_scales, "y") - - if act_scales[name]["w"] is None: - act_scales[name]["w"] = m.weight.abs().clip(1e-8, - None).max(dim=1)[0] - - hooks = [] - for name, m in model.named_modules(): - if isinstance(m, nn.Linear) or isinstance(m, Conv1D): - hooks.append( - m.register_forward_hook( - functools.partial(stat_input_hook, name=name))) - - for i in tqdm(range(num_samples), desc="calibrating model"): - datapoint = dataset[i:i + 1] - line = copy.copy(datapoint) - line[0] = line[0] + ' TL;DR: ' - line[0] = line[0].strip() - line[0] = line[0].replace(" n't", "n't") - input_ids = tokenizer(line, - return_tensors="pt", - max_length=seq_len, - padding=True, - truncation=True).input_ids.to(device) - model(input_ids) - for h in hooks: - h.remove() - return act_scales - - -def get_weight(config, prefix, dtype): - if config[prefix + '.weight'].dtype != dtype: - config[prefix + '.weight'].data = config[prefix + '.weight'].to(dtype) - return config[prefix + '.weight'] - - -def get_bias(config, prefix, dtype): - if config[prefix + '.bias'].dtype != dtype: - config[prefix + '.bias'].data = config[prefix + '.bias'].to(dtype) - return config[prefix + '.bias'] - - -def get_weight_and_bias(config, prefix, dtype): - return get_weight(config, prefix, dtype), get_bias(config, prefix, dtype) - - -def get_tllm_linear_sq_weight(vals, - prefix, - shape, - tensor_parallel, - is_qkv=False, - per_token=False, - per_channel=False, - last_prefix=None, - bias=None, - smoother_value=None, - smoother_shape=None, - rank=0, - cat_dim=0, - multi_query_mode=False): - results = {} - - def multi_query_split(data, local_dim, head_size, tp_size, cur_rank): - q, k, v = np.split(data, [local_dim, local_dim + head_size], axis=-1) - q_split = np.split(q, tp_size, axis=-1) - k_split = np.split(k, tp_size, axis=-1) - v_split = np.split(v, tp_size, axis=-1) - return [ - np.concatenate((q_split[ii], k_split[ii], v_split[ii]), axis=-1) - for ii in range(tp_size) - ][cur_rank] - - col_shape = shape if (is_qkv or per_channel) else [1, 1] - - if per_token: - original_weights = vals["weight.int8.col"] - - local_dim = original_weights.shape[0] - head_size = (original_weights.shape[1] - local_dim) // 2 - if multi_query_mode: - cur_weights = multi_query_split(original_weights, local_dim, - head_size, tensor_parallel, rank) - else: - cur_weights = np.split(original_weights, - tensor_parallel, - axis=cat_dim)[rank] - if is_qkv: - hidden_dim = cur_weights.shape[0] - cur_weights = cur_weights.reshape(hidden_dim, -1) - results[prefix + - 'weight'] = torch.from_numpy(cur_weights).t().contiguous() - if smoother_value is None: - results[last_prefix] = torch.from_numpy( - np.array([1.0], dtype=np.float32)) - - if smoother_value is None: - if multi_query_mode: - cur_per_channel_value = multi_query_split( - vals["scale_w_quant_orig.col"], local_dim, head_size, - tensor_parallel, rank) - else: - cur_per_channel_value = np.split(vals["scale_w_quant_orig.col"], - tensor_parallel, - axis=cat_dim)[rank] - else: - cur_per_channel_value = vals["scale_w_quant_orig.col"] - results[prefix + 'per_channel_scale'] = torch.from_numpy( - np.array(cur_per_channel_value, - dtype=np.float32).reshape(col_shape)).contiguous() - else: - original_weights = np.array(vals["weight.int8"]) - cur_weights = np.split(original_weights, tensor_parallel, - axis=cat_dim)[rank] - - if is_qkv: - hidden_dim = cur_weights.shape[0] - cur_weights = cur_weights.reshape(hidden_dim, -1) - results[prefix + - 'weight'] = torch.from_numpy(cur_weights).t().contiguous() - - cur_per_channel_value = vals["scale_y_accum_quant"] - - results[prefix + 'per_channel_scale'] = torch.from_numpy( - np.array([cur_per_channel_value], - dtype=np.float32).reshape(col_shape)).contiguous() - - results[last_prefix] = torch.from_numpy( - np.array([vals['scale_x_orig_quant']], - dtype=np.float32)).contiguous() - - results[prefix + 'act_scale'] = torch.from_numpy( - np.array([[vals["scale_y_quant_orig"]]], - dtype=np.float32)).contiguous() - - if smoother_value is not None: - cur_smoother_value = np.split(smoother_value, - tensor_parallel, - axis=cat_dim)[rank] - results[prefix + 'smoother'] = cur_smoother_value.reshape( - smoother_shape).contiguous().to(torch.float32) - - if bias is not None: - results[prefix + 'bias'] = bias - - return results - - -def convert_hf_llama(hf_model, - mapping, - rank=0, - dtype='float32', - use_parallel_embedding=False, - sharding_dim=0, - use_weight_only=False, - plugin_weight_only_quant_type=torch.int8, - use_smooth_quant=False, - per_channel=False, - per_token=False, - int8_kv_cache=False, - act_range=[], - qkv_para=[], - smoother=[], - lora_config=None): - - weights = {} - tik = time.time() - tensor_parallel = mapping.tp_size - model_params = dict(hf_model.named_parameters()) - dtype = getattr(torch, dtype) - num_attention_heads = hf_model.config.num_attention_heads - hidden_size = hf_model.config.hidden_size - intermediate_size = hf_model.config.intermediate_size - num_key_value_heads = hf_model.config.num_key_value_heads - mha_mode = (num_key_value_heads == num_attention_heads) - - num_hidden_layers = hf_model.config.num_hidden_layers - layers_range = mapping.pp_layers(num_hidden_layers) - for l in layers_range: - layer_idx = l - layers_range[0] - prefix = f'model.layers.{l}.' - tllm_prex = f'transformer.layers.{layer_idx}.' - - q_weight = get_weight(model_params, prefix + 'self_attn.q_proj', dtype) - k_weight = get_weight(model_params, prefix + 'self_attn.k_proj', dtype) - v_weight = get_weight(model_params, prefix + 'self_attn.v_proj', dtype) - - if not mha_mode: - head_size = hidden_size // num_attention_heads - if num_key_value_heads < tensor_parallel: - # duplicate the KV heads up to tensor_parallel - k_weight = dup_kv_weight(k_weight, num_key_value_heads, - tensor_parallel) - v_weight = dup_kv_weight(v_weight, num_key_value_heads, - tensor_parallel) - assert (k_weight.shape[0] % (mapping.tp_size * head_size)) == 0 - assert (v_weight.shape[0] % (mapping.tp_size * head_size)) == 0 - - wq = split(q_weight, mapping.tp_size, mapping.tp_rank) - wk = split(k_weight, mapping.tp_size, mapping.tp_rank) - wv = split(v_weight, mapping.tp_size, mapping.tp_rank) - - split_v = torch.concat((wq, wk, wv)) - - else: - qkv_weight = torch.cat([q_weight, k_weight, v_weight], dim=0) - - split_v = split_qkv_tp(qkv_weight, num_attention_heads, hidden_size, - tensor_parallel, mapping.tp_rank) - if use_smooth_quant: - qkv_weight = qkv_para[prefix + 'self_attn.qkv_proj'] - - if not mha_mode: - hidden_size = qkv_weight.shape[0] - local_dim = hidden_size - head_size = (qkv_weight.shape[-1] - local_dim) // 2 - qkv_weight = qkv_weight.reshape(hidden_size, - local_dim + 2 * head_size) - else: - qkv_weight = qkv_weight.reshape(hidden_size, 3, hidden_size) - - int8_weights = generate_int8(qkv_weight, - act_range.get(prefix + - 'self_attn.qkv_proj'), - is_qkv=True, - multi_query_mode=bool(not mha_mode)) - - weights.update( - get_tllm_linear_sq_weight( - int8_weights, - tllm_prex + 'attention.qkv.', [ - 1, 3 * hidden_size // tensor_parallel - if mha_mode else hidden_size // tensor_parallel + - (hidden_size // num_key_value_heads) // - tensor_parallel * 2 - ], - tensor_parallel, - is_qkv=True, - per_token=per_token, - per_channel=per_channel, - last_prefix=tllm_prex + 'input_layernorm.scale_to_int', - smoother_value=None, - smoother_shape=None, - rank=mapping.tp_rank, - cat_dim=-1, - multi_query_mode=bool(not mha_mode))) - else: - weights.update( - get_tllm_linear_weight(split_v, tllm_prex + 'attention.qkv.', - None, use_weight_only, - plugin_weight_only_quant_type)) - - if int8_kv_cache: - qkv_y = torch.cat([ - act_range.get(prefix + 'self_attn.q_proj')["y"], - act_range.get(prefix + 'self_attn.k_proj')["y"], - act_range.get(prefix + 'self_attn.v_proj')["y"] - ], - dim=0) - - int8_kv_scales = qkv_y.max() / 127. - - kv_cache_weights = {} - - kv_cache_weights[ - tllm_prex + - 'attention.kv_cache_scaling_factor'] = int8_kv_scales.reshape( - [1]) - - attn_dense_weight = get_weight(model_params, - prefix + 'self_attn.o_proj', dtype) - split_v = split_matrix_tp(attn_dense_weight, - tensor_parallel, - mapping.tp_rank, - dim=1) - if use_smooth_quant: - attn_dense_weight = attn_dense_weight.t() - int8_weights = generate_int8( - attn_dense_weight, act_range.get(prefix + 'self_attn.o_proj')) - weights.update( - get_tllm_linear_sq_weight( - int8_weights, - tllm_prex + 'attention.dense.', [1, hidden_size], - tensor_parallel, - is_qkv=False, - per_token=per_token, - per_channel=per_channel, - last_prefix=tllm_prex + - 'attention.quantization_scaling_factor', - smoother_value=smoother[(prefix + 'self_attn.o_proj')], - smoother_shape=[1, hidden_size // tensor_parallel], - rank=mapping.tp_rank, - cat_dim=0)) - else: - weights.update( - get_tllm_linear_weight(split_v, tllm_prex + 'attention.dense.', - None, use_weight_only, - plugin_weight_only_quant_type)) - - mlp_gate_weight = get_weight(model_params, prefix + 'mlp.up_proj', - dtype) - split_v = split_matrix_tp(mlp_gate_weight, - tensor_parallel, - mapping.tp_rank, - dim=0) - if use_smooth_quant: - mlp_gate_weight = mlp_gate_weight.t() - int8_weights = generate_int8(mlp_gate_weight, - act_range.get(prefix + 'mlp.up_proj')) - - weights.update( - get_tllm_linear_sq_weight( - int8_weights, - tllm_prex + 'mlp.gate.', - [1, intermediate_size // tensor_parallel], - tensor_parallel, - is_qkv=False, - per_token=per_token, - per_channel=per_channel, - last_prefix=tllm_prex + 'post_layernorm.scale_to_int', - smoother_value=None, - smoother_shape=None, - rank=mapping.tp_rank, - cat_dim=-1)) - else: - weights.update( - get_tllm_linear_weight(split_v, tllm_prex + 'mlp.gate.', None, - use_weight_only, - plugin_weight_only_quant_type)) - - mlp_fc_weight = get_weight(model_params, prefix + 'mlp.gate_proj', - dtype) - split_v = split_matrix_tp(mlp_fc_weight, - tensor_parallel, - mapping.tp_rank, - dim=0) - - if use_smooth_quant: - mlp_fc_weight = mlp_fc_weight.t() #verified - int8_weights = generate_int8( - mlp_fc_weight, act_range.get(prefix + 'mlp.gate_proj')) - weights.update( - get_tllm_linear_sq_weight( - int8_weights, - tllm_prex + 'mlp.fc.', - [1, intermediate_size // tensor_parallel], - tensor_parallel, - is_qkv=False, - per_token=per_token, - per_channel=per_channel, - last_prefix=tllm_prex + 'post_layernorm.scale_to_int', - smoother_value=None, - smoother_shape=None, - rank=mapping.tp_rank, - cat_dim=-1)) - else: - weights.update( - get_tllm_linear_weight(split_v, tllm_prex + 'mlp.fc.', None, - use_weight_only, - plugin_weight_only_quant_type)) - - mlp_proj_weight = get_weight(model_params, prefix + 'mlp.down_proj', - dtype) - split_v = split_matrix_tp(mlp_proj_weight, - tensor_parallel, - mapping.tp_rank, - dim=1) - - if use_smooth_quant: - mlp_proj_weight = mlp_proj_weight.t() - int8_weights = generate_int8( - mlp_proj_weight, act_range.get(prefix + 'mlp.down_proj')) - weights.update( - get_tllm_linear_sq_weight( - int8_weights, - tllm_prex + 'mlp.proj.', [1, hidden_size], - tensor_parallel, - is_qkv=False, - per_token=per_token, - per_channel=per_channel, - last_prefix=tllm_prex + 'mlp.quantization_scaling_factor', - smoother_value=smoother[prefix + 'mlp.down_proj'], - smoother_shape=[1, intermediate_size // tensor_parallel], - rank=mapping.tp_rank, - cat_dim=0)) - else: - weights.update( - get_tllm_linear_weight(split_v, tllm_prex + 'mlp.proj.', None, - use_weight_only, - plugin_weight_only_quant_type)) - # Layer norms do not use tensor parallelism - input_ln_weight = get_weight(model_params, prefix + 'input_layernorm', - dtype) - weights[tllm_prex + 'input_layernorm.weight'] = input_ln_weight - - post_ln_weight = get_weight(model_params, - prefix + 'post_attention_layernorm', dtype) - weights[tllm_prex + 'post_layernorm.weight'] = post_ln_weight - - v = get_weight(model_params, 'model.embed_tokens', dtype) - - if hf_model.config.tie_word_embeddings: - # lm_head.weight has the same weights as embedding - if mapping.is_last_pp_rank(): - weights['lm_head.weight'] = split(v, mapping.tp_size, - mapping.tp_rank) - - if use_parallel_embedding: - v = split_matrix_tp(v, - mapping.tp_size, - mapping.tp_rank, - dim=sharding_dim) - - if mapping.is_first_pp_rank(): - weights['transformer.vocab_embedding.weight'] = v - - lm_head_weights = get_weight(model_params, 'lm_head', dtype) - - if mapping.is_last_pp_rank(): - weights['lm_head.weight'] = split_matrix_tp(lm_head_weights, - tensor_parallel, - mapping.tp_rank, - dim=0) - - ln_f_w = get_weight(model_params, 'model.norm', dtype) - weights['transformer.ln_f.weight'] = ln_f_w - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Weights loaded. Total time: {t}') - return weights diff --git a/tensorrt_llm/models/mllama/__init__.py b/tensorrt_llm/models/mllama/__init__.py deleted file mode 100644 index 71bf6d298d42..000000000000 --- a/tensorrt_llm/models/mllama/__init__.py +++ /dev/null @@ -1,14 +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. diff --git a/tensorrt_llm/models/mllama/config.py b/tensorrt_llm/models/mllama/config.py deleted file mode 100644 index cbd7f1b8f387..000000000000 --- a/tensorrt_llm/models/mllama/config.py +++ /dev/null @@ -1,256 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -import json -from pathlib import Path -from typing import List, Optional, Union - -from ..._utils import get_hf_rope_theta -from ...functional import LayerNormPositionType, LayerNormType, MLPType -from ...mapping import Mapping -from ..convert_utils import infer_dtype -from ..modeling_utils import PretrainedConfig, QuantConfig - - -class MLLaMAConfig(PretrainedConfig): - - def __init__(self, - *, - mlp_bias: bool = False, - attn_bias: bool = False, - rotary_base: float = 10000.0, - rotary_scaling: Optional[dict] = None, - residual_mlp: bool = False, - disable_weight_only_quant_plugin: bool = False, - cross_attention: bool = True, - cross_attention_layers: List[int] = None, - vision_output_dim: int = 0, - has_position_embedding=False, - type_vocab_size=None, - rescale_before_lm_head=False, - layernorm_type=LayerNormType.RmsNorm, - layernorm_position=LayerNormPositionType.pre_layernorm, - has_attention_qkvo_bias=False, - has_mlp_bias=False, - has_model_final_layernorm=True, - model_type='MLLaMAForCausalLM', - skip_cross_kv=False, - mlp_type=MLPType.GatedMLP, - has_embedding_scale=False, - residual_scaling=1.0, - has_lm_head_bias=False, - num_buckets=None, - max_distance=0, - relative_attention=False, - **kwargs): - self.mlp_bias = mlp_bias - self.attn_bias = attn_bias - self.rotary_base = rotary_base - self.rotary_scaling = rotary_scaling - self.residual_mlp = residual_mlp - self.disable_weight_only_quant_plugin = disable_weight_only_quant_plugin - - assert cross_attention - self.cross_attention = cross_attention - self.cross_attention_layers = cross_attention_layers - assert vision_output_dim != 0 - self.vision_output_dim = vision_output_dim - - self.has_position_embedding = has_position_embedding - self.type_vocab_size = type_vocab_size - self.rescale_before_lm_head = rescale_before_lm_head - self.layernorm_type = layernorm_type - self.layernorm_position = layernorm_position - self.has_attention_qkvo_bias = has_attention_qkvo_bias - self.has_mlp_bias = has_mlp_bias - self.has_model_final_layernorm = has_model_final_layernorm - self.model_type = model_type - self.skip_cross_kv = skip_cross_kv - self.mlp_type = mlp_type - self.has_embedding_scale = has_embedding_scale - self.residual_scaling = residual_scaling - self.has_lm_head_bias = has_lm_head_bias - self.num_buckets = num_buckets - self.max_distance = max_distance - self.relative_attention = relative_attention - self.skip_cross_attn_blocks = True - - kwargs.pop('embed_vocab_size', None) - kwargs.pop('num_kv_heads_per_layer', None) - kwargs.pop('num_kv_heads_per_cross_attn_layer', None) - super().__init__(**kwargs) - - @property - def embed_vocab_size(self): - return self.vocab_size + 8 #FIXME The vocab_size of embedding contains the special tokens for image - - @property - def num_kv_heads_per_layer(self): - num_kv_heads_per_layer = [ - self.num_key_value_heads for _ in range(self.num_hidden_layers) - ] - for layer_idx in self.cross_attention_layers: - num_kv_heads_per_layer[layer_idx] = 0 - return num_kv_heads_per_layer - - @property - def num_kv_heads_per_cross_attn_layer(self): - num_kv_heads_per_cross_attn_layer = [ - 0 for _ in range(self.num_hidden_layers) - ] - for layer_idx in self.cross_attention_layers: - num_kv_heads_per_cross_attn_layer[ - layer_idx] = self.num_key_value_heads - return num_kv_heads_per_cross_attn_layer - - def to_dict(self): - output = super().to_dict() - # Serialize the fields added in MLLaMAConfig - output['mlp_bias'] = self.mlp_bias - output['attn_bias'] = self.attn_bias - output['rotary_base'] = self.rotary_base - output['rotary_scaling'] = self.rotary_scaling - output['residual_mlp'] = self.residual_mlp - output[ - 'disable_weight_only_quant_plugin'] = self.disable_weight_only_quant_plugin - output['cross_attention'] = self.cross_attention - output['cross_attention_layers'] = self.cross_attention_layers - output['vision_output_dim'] = self.vision_output_dim - output['embed_vocab_size'] = self.embed_vocab_size - output['num_kv_heads_per_layer'] = self.num_kv_heads_per_layer - output[ - 'num_kv_heads_per_cross_attn_layer'] = self.num_kv_heads_per_cross_attn_layer - output['skip_cross_attn_blocks'] = self.skip_cross_attn_blocks - return output - - @classmethod - def from_hugging_face( - cls, - hf_config_or_dir: Union[str, 'transformers.PretrainedConfig'], - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - import transformers - - if isinstance(hf_config_or_dir, transformers.PretrainedConfig): - hf_config = hf_config_or_dir - else: - hf_config_dir = str(hf_config_or_dir) - hf_config = transformers.AutoConfig.from_pretrained( - hf_config_dir, trust_remote_code=True) - - hf_text_config = hf_config.text_config - hf_vision_config = hf_config.vision_config - num_key_value_heads = getattr(hf_text_config, "num_key_value_heads", - hf_text_config.num_attention_heads) - - hidden_act = hf_text_config.hidden_act if hasattr( - hf_text_config, "hidden_act") else hf_text_config.hidden_activation - norm_epsilon = hf_text_config.rms_norm_eps - - head_dim = getattr( - hf_text_config, "head_dim", - hf_text_config.hidden_size // hf_text_config.num_attention_heads) - head_size = getattr(hf_text_config, "kv_channels", head_dim) - attn_bias = getattr(hf_text_config, 'bias', False) or getattr( - hf_text_config, 'attention_bias', False) - rotary_scaling = getattr(hf_text_config, "rope_scaling", None) - rotary_base = get_hf_rope_theta(hf_text_config, 10000.0) - residual_mlp = getattr(hf_text_config, "parallel_attn_mlp_res", False) - disable_weight_only_quant_plugin = kwargs.pop( - 'disable_weight_only_quant_plugin', False) - - dtype = infer_dtype(dtype, getattr(hf_config, 'torch_dtype', None)) - - return cls( - architecture=hf_config.architectures[0], - dtype=dtype, - num_hidden_layers=hf_text_config.num_hidden_layers, - num_attention_heads=hf_text_config.num_attention_heads, - hidden_size=hf_text_config.hidden_size, - intermediate_size=hf_text_config.intermediate_size, - num_key_value_heads=num_key_value_heads, - head_size=head_size, - vocab_size=hf_text_config.vocab_size, - position_embedding_type='rope_gpt_neox', - max_position_embeddings=hf_text_config.max_position_embeddings, - hidden_act=hidden_act, - norm_epsilon=norm_epsilon, - attn_bias=attn_bias, - rotary_base=rotary_base, - rotary_scaling=rotary_scaling, - residual_mlp=residual_mlp, - disable_weight_only_quant_plugin=disable_weight_only_quant_plugin, - mapping=mapping, - quantization=quant_config, - cross_attention_layers=hf_text_config.cross_attention_layers, - vision_output_dim=hf_vision_config.vision_output_dim, - **kwargs) - - @classmethod - def from_meta_ckpt(cls, - meta_ckpt_dir: str, - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - - with open(Path(meta_ckpt_dir, "params.json")) as fp: - meta_config: dict = json.load(fp) - - n_embd = meta_config["dim"] - n_head = meta_config["n_heads"] - n_kv_head = meta_config.get("n_kv_heads", n_head) - vocab_size = meta_config.get("vocab_size", 32000) - - # Reset vocab_size to 32000 for LLama v2 checkpoint. - if vocab_size == -1: - vocab_size = 32000 - - if "hidden_dim" in meta_config: - inter_size = meta_config["hidden_dim"] - else: - multiple_of = meta_config.get("multiple_of", 1) - n_embd_ = int(4 * n_embd * 2 / 3) - ffn_dim_multiplier = meta_config.get("ffn_dim_multiplier", 1) - inter_size = multiple_of * ( - (int(n_embd_ * ffn_dim_multiplier) + multiple_of - 1) // - multiple_of) - - dtype = infer_dtype(dtype, 'bfloat16') - - if meta_config.get('use_scaled_rope'): - rotary_scaling = {"type": "llama3"} - else: - rotary_scaling = meta_config.get("rope_scaling") - - # meta checkpoint don't have vocab_size|hidden_act|rotary_base specified, use same default value as HF - return cls(architecture="MLLaMAForCausalLM", - dtype=dtype, - num_hidden_layers=meta_config["n_layers"], - num_attention_heads=n_head, - hidden_size=n_embd, - intermediate_size=inter_size, - num_key_value_heads=n_kv_head, - vocab_size=vocab_size, - position_embedding_type='rope_gpt_neox', - max_position_embeddings=2048, - hidden_act='silu', - rotary_scaling=rotary_scaling, - rotary_base=meta_config.get('rope_theta', 10000), - norm_epsilon=meta_config["norm_eps"], - mapping=mapping, - quantization=quant_config, - **kwargs) diff --git a/tensorrt_llm/models/mllama/model.py b/tensorrt_llm/models/mllama/model.py deleted file mode 100644 index a610a480959d..000000000000 --- a/tensorrt_llm/models/mllama/model.py +++ /dev/null @@ -1,1569 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -import math -from collections import OrderedDict -from typing import List, Optional, Union - -import tensorrt as trt -import torch - -from tensorrt_llm._common import default_net -from tensorrt_llm._utils import numpy_to_torch, str_dtype_to_torch -from tensorrt_llm.functional import (Conditional, LayerNormPositionType, - LayerNormType, MLPType, - PositionEmbeddingType, Tensor, assertion, - gather_last_token_logits, maximum, minimum, - recv, reduce, send, shape, tanh) -from tensorrt_llm.layers import (MLP, Attention, AttentionMaskParams, - AttentionMaskType, AttentionParams, - ColumnLinear, Embedding, FusedGatedMLP, - GatedMLP, GroupNorm, KeyValueCacheParams, - LayerNorm, LoraParams, RmsNorm) -from tensorrt_llm.llmapi.kv_cache_type import KVCacheType -from tensorrt_llm.lora_helper import (LoraConfig, - get_default_trtllm_modules_to_hf_modules, - use_lora) -from tensorrt_llm.mapping import Mapping -from tensorrt_llm.models.model_weights_loader import ModelWeightsLoader -from tensorrt_llm.models.modeling_utils import PretrainedModel, QuantConfig -from tensorrt_llm.module import Module, ModuleList -from tensorrt_llm.parameter import Parameter -from tensorrt_llm.quantization import QuantMode - -from .config import MLLaMAConfig - -layernorm_map = { - LayerNormType.LayerNorm: LayerNorm, - LayerNormType.RmsNorm: RmsNorm, - LayerNormType.GroupNorm: GroupNorm, -} - -mlp_map = { - MLPType.MLP: MLP, - MLPType.GatedMLP: GatedMLP, - MLPType.FusedGatedMLP: FusedGatedMLP, -} - -ADD_DEBUG_TENSOR = False - - -class CrossAttentionTransformerBlock(Module): - - def __init__( - self, - *, - local_layer_idx, - hidden_size, - ffn_hidden_size, - num_attention_heads, - num_kv_heads, - head_size, - max_position_embeddings=None, - q_scaling=1.0, - has_attention_qkvo_bias=False, - has_mlp_bias=False, - layernorm_position=LayerNormPositionType.pre_layernorm, - layernorm_type=LayerNormType.RmsNorm, - layernorm_eps=1e-5, - hidden_act="gated-silu", - mlp_type=MLPType.GatedMLP, - mapping=Mapping(), - dtype=None, - residual_scaling=1.0, - relative_attention=False, - max_distance=0, - num_buckets=0, - fp16_clamping=False, - skip_cross_kv=False, - use_implicit_relative_attention=False, - rotary_embedding_base=None, - rotary_embedding_scaling=None, - quant_mode=QuantMode(0), - ): - super().__init__() - self.local_layer_idx = local_layer_idx - self.layernorm_type = layernorm_type - ln_type = layernorm_map[layernorm_type] - - self.layernorm_position = layernorm_position - assert self.layernorm_position == LayerNormPositionType.pre_layernorm - - self.cross_attention = Attention( - local_layer_idx=local_layer_idx, - hidden_size=hidden_size, - num_attention_heads=num_attention_heads, - attention_head_size=head_size, - num_kv_heads=num_kv_heads, - max_position_embeddings=max_position_embeddings, - q_scaling=q_scaling, - bias=has_attention_qkvo_bias, - attention_mask_type=AttentionMaskType.causal, - tp_group=mapping.tp_group, - tp_size=mapping.tp_size, - tp_rank=mapping.tp_rank, - dtype=dtype, - cross_attention=True, - relative_attention= - False, # Cross attention has no relative attention bias - max_distance=max_distance, - num_buckets=num_buckets, - position_embedding_type=PositionEmbeddingType. - learned_absolute, # we don't use rope for cross attn - skip_cross_kv=skip_cross_kv, - qk_layernorm=True, - layernorm_type=layernorm_type, - quant_mode=quant_mode, - ) - - self.input_layernorm = ln_type(normalized_shape=hidden_size, - eps=layernorm_eps, - dtype=dtype) - self.gate_attn = Parameter(shape=tuple((1, )), dtype=dtype) - - self.mlp_type = mlp_type - mlp_f = mlp_map[mlp_type] - - self.mlp = mlp_f( - hidden_size=hidden_size, - ffn_hidden_size=ffn_hidden_size, - hidden_act=hidden_act, - bias=has_mlp_bias, - tp_group=mapping.tp_group, - tp_size=mapping.tp_size, - dtype=dtype, - ) - - self.post_layernorm = ln_type(normalized_shape=hidden_size, - eps=layernorm_eps, - dtype=dtype) - self.gate_ffwd = Parameter(shape=tuple((1, )), dtype=dtype) - - self.residual_scaling = residual_scaling - - self.fp16_clamping = fp16_clamping - self.no_ffn = False - - def forward(self, - hidden_states: Tensor, - encoder_output: Optional[Tensor] = None, - attention_mask_params=None, - use_cache=False, - kv_cache_params=None, - attention_params=None, - lora_layer_params=None, - cross_kv_cache_gen: Optional[Tensor] = None, - cross_kv_reuse: Optional[Tensor] = None, - full_text_row_masked_out_mask: Tensor = None, - skip_cross_attn_blocks: Tensor = None): - - assert isinstance(hidden_states, Tensor) - - if encoder_output: - assert isinstance(encoder_output, Tensor) - - if ADD_DEBUG_TENSOR: - hidden_states.mark_output( - f'{self.local_layer_idx:2d}/1.0: hidden_states', - hidden_states.dtype) - # cross attention - residual = hidden_states * self.residual_scaling - - # skip input_layernorm - if skip_cross_attn_blocks is not None: - input_ln_conditional = Conditional(skip_cross_attn_blocks) - skip_result = input_ln_conditional.add_input(hidden_states) - hidden_states = input_ln_conditional.add_input(hidden_states) - hidden_states = self.input_layernorm(hidden_states) - hidden_states = input_ln_conditional.add_output( - skip_result, hidden_states) - else: - hidden_states = self.input_layernorm(hidden_states) - - if ADD_DEBUG_TENSOR: - hidden_states.mark_output( - f'{self.local_layer_idx:2d}/2.1: normed_input', - hidden_states.dtype) - # pass full_text_row_masked_out_mask and xattn_mask - attention_output = self.cross_attention( - hidden_states=hidden_states, - attention_mask=attention_mask_params.cross_attention_mask, - attention_packed_mask=attention_mask_params. - cross_attention_packed_mask, - encoder_output=encoder_output, - use_cache=use_cache, - kv_cache_params=kv_cache_params, - attention_params=attention_params, - lora_layer_params=lora_layer_params, - cross_kv_cache_gen=cross_kv_cache_gen, - cross_kv_reuse=cross_kv_reuse, - skip_attn=skip_cross_attn_blocks, - ) - - if use_cache: - attention_output, presents_cross = attention_output - if ADD_DEBUG_TENSOR: - attention_output.mark_output( - f'{self.local_layer_idx:2d}/3.1: cross_attention_output', - attention_output.dtype) - - attn_residual_scale = tanh(self.gate_attn.value.cast(trt.float32)).cast( - attention_output.dtype) - - attention_input = hidden_states - hidden_states = residual + attn_residual_scale * attention_output - - # use to skip attention_output with residual - # Since conditional does not work for gpt_attention_plugin, we replace the - # attention_output by hidden_states (input of attention) now. - if skip_cross_attn_blocks is not None: - attn_conditional = Conditional(skip_cross_attn_blocks) - skip_result = attn_conditional.add_input(attention_input) - hidden_states = attn_conditional.add_input(hidden_states) - hidden_states = attn_conditional.add_output(skip_result, - hidden_states) - - if ADD_DEBUG_TENSOR: - hidden_states.mark_output( - f'{self.local_layer_idx:2d}/3.2: cross_attn_output_with_residual', - hidden_states.dtype) - - if self.fp16_clamping: - hidden_states = maximum(-64000.0, hidden_states) - hidden_states = minimum(64000.0, hidden_states) - - # MLP - # skip post_layernorm and mlp - if skip_cross_attn_blocks is not None: - mlp_conditional = Conditional(skip_cross_attn_blocks) - skip_case = mlp_conditional.add_input(hidden_states) - hidden_states = mlp_conditional.add_input(hidden_states) - - attention_output = attention_output * full_text_row_masked_out_mask # TODO should move this mask into attention? - - residual = hidden_states * self.residual_scaling - - hidden_states = self.post_layernorm(hidden_states) - - hidden_states = self.mlp(hidden_states, - lora_layer_params=lora_layer_params) - if ADD_DEBUG_TENSOR: - hidden_states.mark_output( - f'{self.local_layer_idx:2d}/4.1: mlp_output', - hidden_states.dtype) - - hidden_states = hidden_states * full_text_row_masked_out_mask - if ADD_DEBUG_TENSOR: - hidden_states.mark_output( - f'{self.local_layer_idx:2d}/4.2: masked_mlp_output', - hidden_states.dtype) - ffn_residual_scale = tanh(self.gate_ffwd.value.cast(trt.float32)).cast( - hidden_states.dtype) - hidden_states = residual + ffn_residual_scale * hidden_states * float( - not self.no_ffn) - - if skip_cross_attn_blocks is not None: - hidden_states = mlp_conditional.add_output(skip_case, hidden_states) - - if self.fp16_clamping: - hidden_states = maximum(-64000.0, hidden_states) - hidden_states = minimum(64000.0, hidden_states) - - if ADD_DEBUG_TENSOR: - hidden_states.mark_output( - f'{self.local_layer_idx:2d}/4.4: transformer_out', - hidden_states.dtype) - - if use_cache: - return (hidden_states, presents_cross) - return hidden_states - - -class TransformerBlock(Module): - - def __init__( - self, - *, - local_layer_idx, - hidden_size, - ffn_hidden_size, - num_attention_heads, - num_kv_heads, - head_size, - max_position_embeddings=None, - q_scaling=1.0, - has_attention_qkvo_bias=False, - has_mlp_bias=False, - layernorm_position=LayerNormPositionType.pre_layernorm, - layernorm_type=LayerNormType.RmsNorm, - layernorm_eps=1e-5, - hidden_act="gated-silu", - mlp_type=MLPType.GatedMLP, - mapping=Mapping(), - dtype=None, - residual_scaling=1.0, - relative_attention=False, - max_distance=0, - num_buckets=0, - fp16_clamping=False, - skip_cross_kv=False, - use_implicit_relative_attention=False, - rotary_embedding_base=None, - rotary_embedding_scaling=None, - quant_mode=QuantMode(0), - ): - super().__init__() - self.local_layer_idx = local_layer_idx - self.layernorm_type = layernorm_type - ln_type = layernorm_map[layernorm_type] - - self.layernorm_position = layernorm_position - assert self.layernorm_position == LayerNormPositionType.pre_layernorm - - self.self_attention = Attention( - local_layer_idx=local_layer_idx, - hidden_size=hidden_size, - num_attention_heads=num_attention_heads, - attention_head_size=head_size, - num_kv_heads=num_kv_heads, - max_position_embeddings=max_position_embeddings, - q_scaling=q_scaling, - bias=has_attention_qkvo_bias, - attention_mask_type=AttentionMaskType.causal, - tp_group=mapping.tp_group, - tp_size=mapping.tp_size, - tp_rank=mapping.tp_rank, - dtype=dtype, - cross_attention=False, - relative_attention=relative_attention, - max_distance=max_distance if use_implicit_relative_attention else 0, - num_buckets=num_buckets, - position_embedding_type=PositionEmbeddingType.relative - if relative_attention else PositionEmbeddingType.rope_gpt_neox, - use_implicit_relative_attention=use_implicit_relative_attention, - rotary_embedding_base=rotary_embedding_base, - rotary_embedding_scaling=rotary_embedding_scaling, - quant_mode=quant_mode, - ) - - self.input_layernorm = ln_type(normalized_shape=hidden_size, - eps=layernorm_eps, - dtype=dtype) - - self.mlp_type = mlp_type - mlp_f = mlp_map[mlp_type] - self.mlp = mlp_f( - hidden_size=hidden_size, - ffn_hidden_size=ffn_hidden_size, - hidden_act=hidden_act, - bias=has_mlp_bias, - tp_group=mapping.tp_group, - tp_size=mapping.tp_size, - dtype=dtype, - ) - - self.post_layernorm = ln_type(normalized_shape=hidden_size, - eps=layernorm_eps, - dtype=dtype) - - self.residual_scaling = residual_scaling - - self.fp16_clamping = fp16_clamping - - def forward( - self, - hidden_states: Tensor, - encoder_output: Optional[Tensor] = None, # not used - attention_mask_params=None, - use_cache=False, - kv_cache_params=None, - attention_params=None, - lora_layer_params=None, - cross_kv_cache_gen: Optional[Tensor] = None, - cross_kv_reuse: Optional[Tensor] = None, - full_text_row_masked_out_mask: Tensor = None, # not used - skip_cross_attn_blocks=None, - ): - assert isinstance(hidden_states, Tensor) - - # self-attention - residual = hidden_states * self.residual_scaling - if ADD_DEBUG_TENSOR: - hidden_states.mark_output( - f'{self.local_layer_idx:2d}/1.0: hidden_states', - hidden_states.dtype) - - hidden_states = self.input_layernorm(hidden_states) - if ADD_DEBUG_TENSOR: - hidden_states.mark_output( - f'{self.local_layer_idx:2d}/2.1: normed attn_input', - hidden_states.dtype) - - attention_output = self.self_attention( - hidden_states=hidden_states, - attention_mask=attention_mask_params.self_attention_mask, - use_cache=use_cache, - kv_cache_params=kv_cache_params, - attention_params=attention_params, - lora_layer_params=lora_layer_params) - - if use_cache: - attention_output, presents_self = attention_output - - if ADD_DEBUG_TENSOR: - attention_output.mark_output( - f'{self.local_layer_idx:2d}/3.1: self_attention_output', - attention_output.dtype) - - hidden_states = residual + attention_output - if ADD_DEBUG_TENSOR: - hidden_states.mark_output( - f'{self.local_layer_idx:2d}/3.1: attention_output_with_residual', - hidden_states.dtype) - - if self.fp16_clamping: - hidden_states = maximum(-64000.0, hidden_states) - hidden_states = minimum(64000.0, hidden_states) - - # MLP - residual = hidden_states * self.residual_scaling - - hidden_states = self.post_layernorm(hidden_states) - if ADD_DEBUG_TENSOR: - hidden_states.mark_output( - f'{self.local_layer_idx:2d}/3.2: normed_mlp_input', - hidden_states.dtype) - - hidden_states = self.mlp(hidden_states, - lora_layer_params=lora_layer_params) - - if ADD_DEBUG_TENSOR: - hidden_states.mark_output( - f'{self.local_layer_idx:2d}/4.1: mlp_output', - hidden_states.dtype) - - hidden_states = residual + hidden_states - if ADD_DEBUG_TENSOR: - hidden_states.mark_output( - f'{self.local_layer_idx:2d}/4.2: mlp_output_residual', - hidden_states.dtype) - - if self.fp16_clamping: - hidden_states = maximum(-64000.0, hidden_states) - hidden_states = minimum(64000.0, hidden_states) - - if use_cache: - return (hidden_states, presents_self) - return hidden_states - - -class MLLaMAModel(Module): - - def __init__(self, config: MLLaMAConfig) -> None: - super().__init__() - self.config = config - self.position_embedding_type = config.position_embedding_type - - self.mapping = self.config.mapping - - self.layernorm_type = self.config.layernorm_type - ln_type = layernorm_map[self.layernorm_type] - - self.has_attention_qkvo_bias = self.config.has_attention_qkvo_bias - self.has_mlp_bias = self.config.has_mlp_bias - - self.has_model_final_layernorm = self.config.has_model_final_layernorm - self._dtype = self.config.dtype - # no quantization considered for now - self._kv_dtype = self._dtype - self._logits_dtype = self.config.logits_dtype - - self.total_num_layers = self.config.num_hidden_layers - self.num_layers = self.config.num_hidden_layers // self.mapping.pp_size - - self.hidden_size = self.config.hidden_size - self.encoder_hidden_size = self.config.hidden_size - self.num_heads = self.config.num_attention_heads - # num_kv_heads = self.num_heads - num_kv_heads = self.config.num_key_value_heads - if num_kv_heads is None or num_kv_heads <= 0: - num_kv_heads = self.num_heads - self.num_kv_heads = num_kv_heads - self.head_size = self.hidden_size // self.num_heads if self.config.head_size is None else self.config.head_size - - self.fp16_clamping = False - - self.skip_cross_kv = self.config.skip_cross_kv - self.mlp_type = MLPType.MLP if not hasattr( - self.config, "mlp_type") else self.config.mlp_type - self.use_implicit_relative_attention = self.config.use_implicit_relative_attention if hasattr( - self.config, "use_implicit_relative_attention") else False - - self.cross_attention_layers = self.config.cross_attention_layers - - if self.mapping.is_first_pp_rank(): - self.vocab_embedding = Embedding( - self.config.embed_vocab_size, - self.config.hidden_size, - dtype=self._dtype, - tp_size=self.mapping.tp_size - if self.config.use_parallel_embedding else 1, - tp_group=self.mapping.tp_group - if self.config.use_parallel_embedding else None, - sharding_dim=self.config.embedding_sharding_dim, - tp_rank=self.mapping.tp_rank) - - layers_range = self.mapping.pp_layers(self.total_num_layers) - _layers = [] - for layer_idx in layers_range: - local_layer_idx = layer_idx - layers_range[0] - args = { - "local_layer_idx": local_layer_idx, - "hidden_size": self.config.hidden_size, - "ffn_hidden_size": self.config.intermediate_size, - "num_attention_heads": self.num_heads, - "num_kv_heads": self.num_kv_heads, - "head_size": self.head_size, - "max_position_embeddings": self.config.max_position_embeddings, - "layernorm_position": self.config.layernorm_position, - "layernorm_eps": self.config.norm_epsilon, - "layernorm_type": self.config.layernorm_type, - "hidden_act": self.config.hidden_act, - "mlp_type": self.mlp_type, - "mapping": self.mapping, - "dtype": self._dtype, - "residual_scaling": self.config.residual_scaling, - "max_distance": self.config.max_distance, - "num_buckets": self.config.num_buckets, - "fp16_clamping": self.fp16_clamping, - "skip_cross_kv": self.skip_cross_kv, - "rotary_embedding_base": self.config.rotary_base, - "rotary_embedding_scaling": self.config.rotary_scaling, - "quant_mode": self.config.quant_mode, - } - if layer_idx in self.cross_attention_layers: - assert layers_range[0] == 0, "not support PP now" - _layers.append(CrossAttentionTransformerBlock(**args)) - else: - _layers.append(TransformerBlock(**args)) - - self.layers = ModuleList(_layers) - - if self.mapping.is_last_pp_rank(): - self.ln_f = None - if self.has_model_final_layernorm: - self.ln_f = ln_type(normalized_shape=self.config.hidden_size, - eps=self.config.norm_epsilon, - dtype=self.config.dtype) - - if self.config.relative_attention and not self.use_implicit_relative_attention: - self.rel_attn_table = Parameter( - shape=(self.config.num_attention_heads // self.mapping.tp_size, - self.config.num_buckets), - dtype=self._dtype) - - def forward( - self, - decoder_input_ids: Tensor, - encoder_output: Tensor, - use_cache=False, - attention_mask_params=None, - last_token_ids=None, - kv_cache_params=None, - attention_params=None, - hidden_states=None, - lora_params: LoraParams = None, - cross_kv_cache_gen: Optional[Tensor] = None, - cross_kv_reuse: Optional[Tensor] = None, - prompt_embedding_table: Optional[Tensor] = None, - prompt_tasks: Optional[Tensor] = None, - prompt_vocab_size: Optional[Tensor] = None, - skip_cross_attn_blocks: Optional[Tensor] = None, - ): - if self.mapping.is_first_pp_rank(): - assert isinstance(decoder_input_ids, Tensor) - else: - assert isinstance(hidden_states, Tensor) - - # In PP, layer 0 has ids as inputs, all other layers have hidden_states as inputs - if self.mapping.is_first_pp_rank(): - hidden_states = self.vocab_embedding(decoder_input_ids) - self.register_network_output('embedding_layer_output', - hidden_states) - else: - hidden_states = recv(hidden_states, self.mapping.prev_pp_rank()) - - kv_cache_params.fill_none_tensor_list(len(self.layers)) - - full_text_row_masked_out_mask = reduce( - (attention_mask_params.cross_attention_mask).cast( - hidden_states.dtype), - trt.ReduceOperation.MAX, - dim=-1, - keepdim=True) - if ADD_DEBUG_TENSOR: - full_text_row_masked_out_mask.mark_output( - "full_text_row_masked_out_mask", - full_text_row_masked_out_mask.dtype) - - cross_attention_mask_type = attention_mask_params.cross_attention_mask.dtype - attention_mask_params.cross_attention_mask = ( - attention_mask_params.cross_attention_mask.cast( - full_text_row_masked_out_mask.dtype) * - full_text_row_masked_out_mask).cast(cross_attention_mask_type) - - invert_mask = 1.0 - attention_mask_params.cross_attention_mask.cast( - hidden_states.dtype) - invert_full_text_row_masked_out_mask = 1.0 - full_text_row_masked_out_mask - final_mask = invert_mask - invert_full_text_row_masked_out_mask - attention_mask_params.cross_attention_mask = final_mask.cast( - cross_attention_mask_type) - if ADD_DEBUG_TENSOR: - attention_mask_params.cross_attention_mask.mark_output( - "attention_mask_params.cross_attention_mask", - attention_mask_params.cross_attention_mask.dtype) - - if use_cache: - presents = [] - for i, (decoder_layer, past) in enumerate( - zip(self.layers, kv_cache_params.past_key_value)): - - lora_layer_params = None - if lora_params is not None and lora_params.lora_ranks is not None: - lora_layer_params = lora_params.get_layer_params(i) - hidden_states = decoder_layer( - hidden_states, - encoder_output=encoder_output, - attention_mask_params=attention_mask_params, - use_cache=use_cache, - kv_cache_params=KeyValueCacheParams( - past_key_value=past, - host_past_key_value_lengths=kv_cache_params. - host_past_key_value_lengths, - host_max_attention_window_sizes=kv_cache_params. - host_max_attention_window_sizes, - host_sink_token_length=kv_cache_params. - host_sink_token_length, - cache_indirection=kv_cache_params.cache_indirection, - kv_cache_block_offsets=kv_cache_params. - kv_cache_block_offsets, - host_kv_cache_block_offsets=kv_cache_params. - host_cross_kv_cache_block_offsets, - host_kv_cache_pool_pointers=kv_cache_params. - host_kv_cache_pool_pointers, - host_kv_cache_pool_mapping=kv_cache_params. - host_kv_cache_pool_mapping, - cross_kv_cache_block_offsets=kv_cache_params. - cross_kv_cache_block_offsets, - host_cross_kv_cache_block_offsets=kv_cache_params. - host_cross_kv_cache_block_offsets, - host_cross_kv_cache_pool_pointers=kv_cache_params. - host_cross_kv_cache_pool_pointers, - host_cross_kv_cache_pool_mapping=kv_cache_params. - host_cross_kv_cache_pool_mapping, - ), - skip_cross_attn_blocks=skip_cross_attn_blocks if isinstance( - decoder_layer, CrossAttentionTransformerBlock) else None, - attention_params=attention_params, - lora_layer_params=lora_layer_params, - cross_kv_cache_gen=cross_kv_cache_gen, - cross_kv_reuse=cross_kv_reuse, - full_text_row_masked_out_mask=full_text_row_masked_out_mask, - ) - - if use_cache: - present = hidden_states[1] - presents.append((present)) - hidden_states = hidden_states[0] - - if self.mapping.is_last_pp_rank(): - if self.ln_f: - hidden_states = self.ln_f(hidden_states) - else: - hidden_states = send(hidden_states, self.mapping.next_pp_rank()) - - if use_cache: - return (hidden_states, tuple(presents)) - return hidden_states - - def precompute_relative_attention_bias(self, build_config): - if self.config.relative_attention and not self.use_implicit_relative_attention: - relative_attention_bias_builder = torch.ops.tensorrt_llm.relative_attention_bias - rel_attn_precomputed = torch.zeros( - (self.config.num_attention_heads // self.mapping.tp_size, - build_config.max_seq_len + 1, build_config.max_seq_len + 1), - dtype=str_dtype_to_torch(self.config.dtype), - device='cuda') - rel_attn_table = numpy_to_torch( - self.rel_attn_table.raw_value).to('cuda') - relative_attention_bias_builder( - rel_attn_precomputed, - rel_attn_table, - self.config.num_attention_heads // self.mapping.tp_size, - build_config.max_seq_len, - self.config.num_buckets, - False, - self.config.max_distance, - ) - for layer_idx in range(self.num_layers): - self.layers[layer_idx].self_attention.set_rel_attn_table( - build_config.max_seq_len, rel_attn_precomputed) - - -# TODO try to inherit the DecoderModelForCausalLM -class MLLaMAForCausalLM(PretrainedModel): - config_class = MLLaMAConfig - - def __init__(self, config: MLLaMAConfig): - super().__init__(config) - Attention.create_attention_const_params(self, config) - self.position_embedding_type = config.position_embedding_type - self.transformer = MLLaMAModel(config) - - self.mapping = self.config.mapping - - self.has_model_final_layernorm = self.config.has_model_final_layernorm - self._dtype = self.config.dtype - self._kv_dtype = self._dtype - self._logits_dtype = self.config.logits_dtype - - if self.mapping.is_last_pp_rank(): - self.lm_head = ColumnLinear( - self.config.hidden_size, - self.config.vocab_size, - bias=False if not hasattr(self.config, "has_lm_head_bias") else - self.config.has_lm_head_bias, - dtype=self.config.dtype, - tp_group=self.config.mapping.tp_group, - tp_size=self.config.mapping.tp_size, - gather_output=True, - ) - - self.trtllm_modules_to_hf_modules = { - **get_default_trtllm_modules_to_hf_modules(), - "attn_q": "self_attn.q_proj", - "attn_k": "self_attn.k_proj", - "attn_v": "self_attn.v_proj", - "attn_dense": "self_attn.o_proj", - "cross_attn_q": "encoder_attn.q_proj", - "cross_attn_k": "encoder_attn.k_proj", - "cross_attn_v": "encoder_attn.v_proj", - "cross_attn_dense": "encoder_attn.o_proj", - } - - def forward( - self, - decoder_input_ids: Tensor, - encoder_output: Tensor, - use_cache=False, - attention_mask_params=None, - last_token_ids=None, - kv_cache_params=None, - attention_params=None, - hidden_states=None, - lora_params: LoraParams = None, - cross_kv_cache_gen: Optional[Tensor] = None, - cross_kv_reuse: Optional[Tensor] = None, - prompt_embedding_table: Optional[Tensor] = None, - prompt_tasks: Optional[Tensor] = None, - prompt_vocab_size: Optional[Tensor] = None, - skip_cross_attn_blocks: Optional[Tensor] = None, - ): - if self.mapping.is_first_pp_rank(): - assert isinstance(decoder_input_ids, Tensor) - else: - assert isinstance(hidden_states, Tensor) - attention_params = Attention.fill_attention_params( - self, attention_params) - hidden_states = self.transformer( - decoder_input_ids=decoder_input_ids, - encoder_output=encoder_output, - use_cache=use_cache, - attention_mask_params=attention_mask_params, - last_token_ids=last_token_ids, - kv_cache_params=kv_cache_params, - attention_params=attention_params, - hidden_states=hidden_states, - lora_params=lora_params, - cross_kv_cache_gen=cross_kv_cache_gen, - cross_kv_reuse=cross_kv_reuse, - prompt_embedding_table=prompt_embedding_table, - prompt_tasks=prompt_tasks, - prompt_vocab_size=prompt_vocab_size, - skip_cross_attn_blocks=skip_cross_attn_blocks, - ) - - if use_cache: - hidden_states, presents = hidden_states - - if self.mapping.is_last_pp_rank(): - # [bs, seq, hidden_size] or [num_tokens, hidden_size] -> [bs, hidden_size] - hidden_states = gather_last_token_logits( - hidden_states, last_token_ids, - default_net().plugin_config.remove_input_padding) - - # [bs, hidden_size] -> [bs, vocab_size] - lm_logits = self.lm_head(hidden_states) - lm_logits.mark_output(f'logits', self._logits_dtype) - else: - hidden_states = send(hidden_states, self.mapping.next_pp_rank()) - hidden_states.mark_output(f'hidden_states_output', self._dtype) - - if use_cache and default_net().plugin_config.paged_kv_cache == False: - for i, present in zip(self.mapping.pp_layers(self.total_num_layers), - presents): - present[0].mark_output(f'present_key_value_{i}', self._kv_dtype) - if default_net().plugin_config.gpt_attention_plugin: - present[1].mark_output(f'cross_present_key_value_{i}', - self._kv_dtype) - if self.mapping.is_last_pp_rank(): - return (lm_logits, tuple(presents)) - return (hidden_states, tuple(presents)) - else: - if self.mapping.is_last_pp_rank(): - return lm_logits - return hidden_states - - def prepare_inputs(self, - max_batch_size, - max_beam_width, - max_decoder_input_len, - max_seq_len, - max_encoder_input_len, - gather_context_logits: bool = False, - gather_generation_logits: bool = False, - lora_target_modules: List[str] = None, - prompt_embedding_table_size: int = 0, - use_cache=True, - *args, - **kwargs): - '''@brief: Prepare inputs Tensors for the model, the given sizes are used to determine the - ranges of the dimensions of when using TRT dynamic shapes. - - @return: a list contains values which can be fed into the self.forward() - ''' - - # Prepare inputs - max_output_len = max_decoder_input_len + max_seq_len - - head_size = self.transformer.head_size - num_kv_heads = (self.transformer.num_kv_heads + self.mapping.tp_size - - 1) // self.mapping.tp_size - - encoder_head_size = head_size - encoder_num_kv_heads = num_kv_heads - - bb_range = [ - 1, (max_batch_size * max_beam_width + 1) // 2, - max_batch_size * max_beam_width - ] - bs_range = [1, (max_batch_size + 1) // 2, max_batch_size] - beam_width_range = [1, (max_beam_width + 1) // 2, max_beam_width] - inlen_range = [ - 1, 1, max_decoder_input_len - ] # context phase >= 1 (if forced_input_ids), generation phase = 1 - encoder_inlen_range = [ - 1, (max_encoder_input_len + 1) // 2, max_encoder_input_len - ] - mask_len_range = [1, (max_output_len + 1) // 2 + 1, max_output_len + 1] - max_output_len_range = [0, (max_output_len + 1) // 2, max_output_len] - - encoder_num_tokens_range = [ - 0, # 0 for generation phase, >0 for context phase - (max_encoder_input_len * max_batch_size + 1) // 2, - max_encoder_input_len * max_batch_size, - ] - decoder_num_tokens_range = [ - 1, - max_batch_size * max_beam_width, - max(max_decoder_input_len * max_batch_size, - max_beam_width * max_batch_size), - ] - - # No enable_two_optimization_profiles support yet - encoder_input_len_range = [ - 0, # 0 for generation phase, >0 for context phase - (max_encoder_input_len + 1) // 2, - max_encoder_input_len - ] - # pack masks into bits (store as int32). - max_cross_packed_mask_dim0 = max_batch_size * ( - (max_decoder_input_len + 128 - 1) // 128) * 128 - max_cross_packed_mask_dim1 = ( - (max_encoder_input_len + 256 - 1) // 256) * 256 // 32 - cross_packed_mask_dim0_range = [ - 1, (max_cross_packed_mask_dim0 + 1) // 2, max_cross_packed_mask_dim0 - ] - cross_packed_mask_dim1_range = [ - 0, # 0 for generation phase, >0 for context phase - (max_cross_packed_mask_dim1 + 1) // 2, - max_cross_packed_mask_dim1 - ] - past_key_value = [] - sequence_length = None - host_past_key_value_lengths = None - attention_mask_params = AttentionMaskParams() - use_gpt_attention_plugin = default_net( - ).plugin_config.gpt_attention_plugin - remove_input_padding = default_net().plugin_config.remove_input_padding - paged_kv_cache = default_net().plugin_config.paged_kv_cache - tokens_per_block = default_net().plugin_config.tokens_per_block - use_lora_plugin = default_net().plugin_config.lora_plugin - kv_cache_type = None - if not use_cache: - kv_cache_type = KVCacheType.DISABLED - else: - if paged_kv_cache: - kv_cache_type = KVCacheType.PAGED - else: - kv_cache_type = KVCacheType.CONTINUOUS - - input_ids, hidden_states = None, None - if remove_input_padding: - if self.mapping.is_first_pp_rank(): - input_ids = Tensor(name='input_ids', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([ - ('decoder_num_tokens', - [decoder_num_tokens_range]), - ])) - else: - hidden_states = Tensor(name='hidden_states_input', - dtype=self._dtype, - shape=[-1, self.hidden_size], - dim_range=OrderedDict([ - ('decoder_num_tokens', - [decoder_num_tokens_range]), - ('hidden_size', [self.hidden_size]), - ])) - else: - if self.mapping.is_first_pp_rank(): - input_ids = Tensor(name='input_ids', - dtype=trt.int32, - shape=[-1, -1], - dim_range=OrderedDict([ - ('batch_size_beam_width', [bb_range]), - ('input_len', [inlen_range]), - ])) - else: - hidden_states = Tensor(name='hidden_states_input', - dtype=self._dtype, - shape=[-1, -1, self.hidden_size], - dim_range=OrderedDict([ - ('batch_size_beam_width', [bb_range - ]), - ('input_len', [inlen_range]), - ('hidden_size', [self.hidden_size]), - ])) - - encoder_input_lengths = Tensor( - name="encoder_input_lengths", - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([("batch_size_beam_width", [bb_range])]), - ) - encoder_max_input_length = Tensor( - name="encoder_max_input_length", - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([("encoder_max_input_length", - [encoder_inlen_range])]), - ) - if remove_input_padding: - encoder_output = Tensor( - name="encoder_output", - dtype=self._dtype, - shape=[-1, self.config.hidden_size], - dim_range=OrderedDict([ - ("encoder_num_tokens", [encoder_num_tokens_range]), - ("hidden_size", [self.config.hidden_size]), - ]), - ) - else: - encoder_output = Tensor( - name="encoder_output", - dtype=self._dtype, - shape=[-1, -1, self.config.hidden_size], - dim_range=OrderedDict([ - ("batch_size_beam_width_encoder", [bb_range]), - ("encoder_input_len", [encoder_input_len_range]), - ("hidden_size", [self.config.hidden_size]), - ]), - ) - - context_lengths = None - host_context_lengths = None - host_request_types = None - host_runtime_perf_knobs = None - host_context_progress = None - if use_gpt_attention_plugin and remove_input_padding: - host_context_lengths = Tensor(name='host_context_lengths', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([ - ('batch_size_beam_width', - [bb_range]) - ])) - - if use_gpt_attention_plugin: - if kv_cache_type != KVCacheType.DISABLED: - sequence_length = Tensor( - name='sequence_length', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([('batch_size_beam_width', [bb_range]) - ]), - ) - host_past_key_value_lengths = Tensor( - name='host_past_key_value_lengths', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([('batch_size_beam_width', [bb_range]) - ]), - ) - - context_lengths = Tensor(name='context_lengths', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([ - ('batch_size_beam_width', [bb_range]) - ])) - host_request_types = Tensor(name='host_request_types', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([ - ('batch_size_beam_width', - [bb_range]) - ])) - - host_runtime_perf_knobs = Tensor(name='host_runtime_perf_knobs', - dtype=trt.int64, - shape=[16], - dim_range=OrderedDict([ - ('perf_knob_size', [16]) - ])) - - host_context_progress = Tensor(name='host_context_progress', - dtype=trt.int64, - shape=[1], - dim_range=OrderedDict([ - ('context_progress_size', [1]) - ])) - - last_token_ids = None - if self.mapping.is_last_pp_rank() and not gather_context_logits: - last_token_ids = Tensor( - name="last_token_ids", - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([("batch_size_last_token_ids", [bb_range]) - ]), - ) - - attention_mask = None - if not use_gpt_attention_plugin: - attention_mask = Tensor( - name='attention_mask', - dtype=trt.int32, - shape=[-1, -1], - dim_range=OrderedDict([ - ('batch_size_beam_width', [bb_range]), - ('mask_len', [mask_len_range]), - ]), - ) - assert False, "not support non-attention-plugin case now" - - cross_attention_mask = Tensor( - name='cross_attention_mask', - dtype=trt.bool, - shape=[-1, -1], - dim_range=OrderedDict([ - ('decoder_num_tokens_2', - [decoder_num_tokens_range - ]), # TODO should use same name as input_ids - ('encoder_input_len_2', [encoder_input_len_range]), - ]), - ) - - cross_attention_packed_mask = Tensor( - name='cross_attention_packed_mask', - dtype=trt.int32, - shape=[-1, -1], - dim_range=OrderedDict([ - ('cross_packed_mask_dim0', [cross_packed_mask_dim0_range]), - ('cross_packed_mask_dim1', [cross_packed_mask_dim1_range]), - ]), - ) - - # create the attention_mask_params. - attention_mask_params = AttentionMaskParams( - attention_mask, None, cross_attention_mask, - cross_attention_packed_mask) - - cache_indirection = Tensor( - name='cache_indirection', - dtype=trt.int32, - shape=[-1, -1, -1], - dim_range=OrderedDict([ - ('batch_size_cache', [bs_range]), - ('beam_width', [beam_width_range]), - ('max_seq_len', [max_output_len_range]), - ]), - ) - - layers_range = self.mapping.pp_layers(self.transformer.total_num_layers) - num_pp_layers = len(layers_range) - - host_max_attention_window_sizes = None - host_sink_token_length = None - if use_gpt_attention_plugin: - host_max_attention_window_sizes = Tensor( - name=f'host_max_attention_window_sizes', - dtype=trt.int32, - shape=[num_pp_layers], - dim_range=OrderedDict([('num_layers', [num_pp_layers])])) - host_sink_token_length = Tensor(name='host_sink_token_length', - dtype=trt.int32, - shape=[1], - dim_range=OrderedDict([('scalar', - [1])])) - # TODO LoRA for mllama is not verified. - lora_weights_pointers = None - lora_ranks = None - lora_params = None - if use_lora_plugin: - lora_weights_pointers = [] - lora_ranks = [] - missing_qkv_modules = [] - if any(x in lora_target_modules - for x in ["attn_q", "attn_k", "attn_v"]): - for lora_module in [ - "attn_q", - "attn_k", - "attn_v", - ]: - if lora_module not in lora_target_modules: - missing_qkv_modules.append(lora_module) - if any(x in lora_target_modules - for x in ["cross_attn_q", "cross_attn_k", "cross_attn_v"]): - for lora_module in [ - "cross_attn_q", "cross_attn_k", "cross_attn_v" - ]: - if lora_module not in lora_target_modules: - missing_qkv_modules.append(lora_module) - - # For LoRA - for i in layers_range: - lora_weight_pointer_dict = {} - lora_rank_dict = {} - for lora_module in (lora_target_modules + missing_qkv_modules): - lora_weight_pointer = Tensor( - name=f'{lora_module}_lora_weights_pointers_{i}', - dtype=trt.int64, - shape=[-1, 3], - dim_range=OrderedDict([('batch_size_beam_width', - [bb_range]), - ('in_out_scales', [3])])) - lora_weight_pointer_dict.update({ - f'{lora_module}_lora_weights_pointers': - lora_weight_pointer - }) - - lora_rank = Tensor(name=f'{lora_module}_lora_ranks_{i}', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([ - ('batch_size_beam_width', [bb_range]) - ])) - lora_rank_dict.update( - {f'{lora_module}_lora_ranks': lora_rank}) - - lora_weights_pointers.append(lora_weight_pointer_dict) - lora_ranks.append(lora_rank_dict) - - # For cross attention, we need to use encoder_input_lengths (in CPU) to pass - # as the host_context_lengths to the lora_plugin. But for self attention, we - # should keep using the original host_context_lengths. Therefore, we keep both - # of them in the lora_params. - host_encoder_input_lengths = None - if remove_input_padding: - host_encoder_input_lengths = Tensor( - name="host_encoder_input_lengths", - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([("batch_size_beam_width", [bb_range]) - ]), - ) - - lora_params = LoraParams( - lora_ranks=lora_ranks, - lora_weights_pointers=lora_weights_pointers, - host_context_lengths=host_context_lengths, - max_context_length=max_decoder_input_len, - max_encoder_context_length=max_encoder_input_len, - host_request_types=host_request_types, - host_encoder_input_lengths=host_encoder_input_lengths, - ) - - kv_cache_block_offsets = None - host_kv_cache_block_offsets = None - host_kv_cache_pool_pointers = None - host_kv_cache_pool_mapping = None - - cross_kv_cache_block_offsets = None - host_cross_kv_cache_block_offsets = None - host_cross_kv_cache_pool_pointers = None - host_cross_kv_cache_pool_mapping = None - - if use_cache: - if not paged_kv_cache: - for i in layers_range: - kv_dim_range = OrderedDict([ - ('batch_size_beam_width', [bb_range]), - ('kv', [2]), - ('num_heads', [num_kv_heads]), - ('past_key_len', [max_output_len_range]), - ('head_size', [head_size]), - ]) - kv = Tensor(name=f'past_key_value_{i}', - dtype=self._kv_dtype, - shape=[-1, 2, num_kv_heads, -1, head_size], - dim_range=kv_dim_range) - - past_key_value.append(kv) - - if i in self.transformer.cross_attention_layers: - xa_layer_id = self.transformer.cross_attention_layers.index( - i) + layers_range[-1] - cross_kv_dim_range = OrderedDict([ - ('batch_size_beam_width', [bb_range]), - ('kv', [2]), - ('cross_num_heads', [encoder_num_kv_heads]), - ('cross_past_key_len', [encoder_input_len_range]), - ('cross_head_size', [encoder_head_size]), - ]) - cross_kv = Tensor( - name=f'cross_past_key_value_{xa_layer_id}', - dtype=self._kv_dtype, - shape=[ - -1, 2, encoder_num_kv_heads, -1, encoder_head_size - ], - dim_range=cross_kv_dim_range) - past_key_value.append(kv) - - # TODO: Remove this when TRT fix the named dimension - if not remove_input_padding: - assertion( - shape( - input_ids if self.mapping.is_first_pp_rank() else - hidden_states, 0) == shape(kv, 0), 'batch size') - - else: # paged_kv_cache == True - # PagedKV setup for KV cache of self-attention - max_blocks_per_seq_range = [[ - math.ceil(max_output_len_range[0] / tokens_per_block), - math.ceil(max_output_len_range[1] / tokens_per_block), - math.ceil(max_output_len_range[2] / tokens_per_block) - ]] - max_blocks_per_seq_range = [[ - x for x in max_blocks_per_seq_range[0] - ]] - - # PagedKV setup for KV cache of cross-attention - max_cross_blocks_per_seq_range = [[ - math.ceil(encoder_input_len_range[0] / tokens_per_block), - math.ceil(encoder_input_len_range[1] / tokens_per_block), - math.ceil(encoder_input_len_range[2] / tokens_per_block) - ]] - max_cross_blocks_per_seq_range = [[ - x for x in max_cross_blocks_per_seq_range[0] - ]] - - num_kv_cache_pools = 2 - - kv_cache_block_offsets = Tensor( - name=f'kv_cache_block_offsets', - dtype=trt.int32, - shape=[num_kv_cache_pools, -1, 2, -1], - dim_range=OrderedDict([ - ('num_kv_cache_pools', [num_kv_cache_pools]), - ('batch_size_beam_width', [bb_range]), - ('kv', [2]), - ('max_blocks_per_seq', max_blocks_per_seq_range), - ])) - host_kv_cache_block_offsets = Tensor( - name=f'host_kv_cache_block_offsets', - dtype=trt.int32, - shape=[num_kv_cache_pools, -1, 2, -1], - dim_range=OrderedDict([ - ('num_kv_cache_pools', [num_kv_cache_pools]), - ('batch_size_beam_width', [bb_range]), - ('kv', [2]), - ('max_blocks_per_seq', max_blocks_per_seq_range), - ])) - host_kv_cache_pool_pointers = Tensor( - name=f'host_kv_cache_pool_pointers', - dtype=trt.int64, - shape=[num_kv_cache_pools, 2], - dim_range=OrderedDict([ - ('num_kv_cache_pools', [num_kv_cache_pools]), - ('num_pools', [2]), - ])) - host_kv_cache_pool_mapping = Tensor( - name=f"host_kv_cache_pool_mapping", - dtype=trt.int32, - # 2: (Index of pool, Index of layer within pool) - shape=[num_pp_layers, 2], - dim_range=OrderedDict([ - ('pools_mapping', [num_pp_layers]), - ('layer_cache_pool_locator', [2]), - ])) - - # paged blocks for cross kv - cross_kv_cache_block_offsets = Tensor( - name=f'cross_kv_cache_block_offsets', - dtype=trt.int32, - shape=[num_kv_cache_pools, -1, 2, -1], - dim_range=OrderedDict([ - ('num_kv_cache_pools', [num_kv_cache_pools]), - ('batch_size_beam_width', [bb_range]), - ('kv', [2]), - ('max_cross_blocks_per_seq', - max_cross_blocks_per_seq_range), - ])) - host_cross_kv_cache_block_offsets = Tensor( - name=f'host_cross_kv_cache_block_offsets', - dtype=trt.int32, - shape=[num_kv_cache_pools, -1, 2, -1], - dim_range=OrderedDict([ - ('num_kv_cache_pools', [num_kv_cache_pools]), - ('batch_size_beam_width', [bb_range]), - ('kv', [2]), - ('max_cross_blocks_per_seq', - max_cross_blocks_per_seq_range), - ])) - host_cross_kv_cache_pool_pointers = Tensor( - name=f'host_cross_kv_cache_pool_pointers', - dtype=trt.int64, - shape=[num_kv_cache_pools, 2], - dim_range=OrderedDict([ - ('num_kv_cache_pools', [num_kv_cache_pools]), - ('num_pools', [2]), - ])) - host_cross_kv_cache_pool_mapping = Tensor( - name=f"host_cross_kv_cache_pool_mapping", - dtype=trt.int32, - # 2: (Index of pool, Index of layer within pool) - shape=[num_pp_layers, 2], - dim_range=OrderedDict([ - ('pools_mapping', [num_pp_layers]), - ('layer_cache_pool_locator', [2]), - ])) - - for i in layers_range: - past_key_value.append(None) - - kv_cache_params = KeyValueCacheParams( - past_key_value=past_key_value, - host_past_key_value_lengths=host_past_key_value_lengths, - host_max_attention_window_sizes=host_max_attention_window_sizes, - host_sink_token_length=host_sink_token_length, - cache_indirection=cache_indirection, - kv_cache_block_offsets=kv_cache_block_offsets, - host_kv_cache_block_offsets=host_kv_cache_block_offsets, - host_kv_cache_pool_pointers=host_kv_cache_pool_pointers, - host_kv_cache_pool_mapping=host_kv_cache_pool_mapping, - cross_kv_cache_block_offsets=cross_kv_cache_block_offsets, - host_cross_kv_cache_block_offsets= - host_cross_kv_cache_block_offsets, - host_cross_kv_cache_pool_pointers= - host_cross_kv_cache_pool_pointers, - host_cross_kv_cache_pool_mapping= - host_cross_kv_cache_pool_mapping, - ) - - attention_params = AttentionParams( - sequence_length=sequence_length, - context_lengths=context_lengths, - host_context_lengths=host_context_lengths, - max_context_length=max_decoder_input_len, - host_request_types=host_request_types, - host_runtime_perf_knobs=host_runtime_perf_knobs, - host_context_progress=host_context_progress, - encoder_input_lengths=encoder_input_lengths, - encoder_max_input_length=encoder_max_input_length, - ) - - cross_kv_cache_gen = Tensor(name='cross_kv_cache_gen', - dtype=trt.bool, - shape=[1], - dim_range=OrderedDict([ - ('boolean', [1]), - ])) - cross_kv_reuse = None - num_heads = (self.transformer.num_heads + self.mapping.tp_size - - 1) // self.mapping.tp_size - cross_kv_out_dim = 2 * num_kv_heads * self.transformer.head_size - if self.transformer.skip_cross_kv: - if remove_input_padding: - cross_kv_reuse = Tensor( - name="cross_kv_reuse", - dtype=self._dtype, - shape=[-1, cross_kv_out_dim], - dim_range=OrderedDict([ - ("encoder_num_tokens", [encoder_num_tokens_range]), - ("encoder_kv_size", [cross_kv_out_dim]), - ]), - ) - else: - cross_kv_reuse = Tensor( - name="cross_kv_reuse", - dtype=self._dtype, - shape=[-1, -1, cross_kv_out_dim], - dim_range=OrderedDict([ - ("batch_size_beam_width_encoder", [bb_range]), - ("encoder_input_len", [encoder_input_len_range]), - ("encoder_kv_size", [cross_kv_out_dim]), - ]), - ) - - skip_cross_attn_blocks = None - if self.config.skip_cross_attn_blocks: - skip_cross_attn_blocks = Tensor(name='skip_cross_attn_blocks', - dtype=trt.bool, - shape=[1], - dim_range=OrderedDict([ - ('boolean', [1]), - ])) - - prompt_embedding_table = None - tasks = None - prompt_vocab_size = None - - if self.mapping.is_first_pp_rank() and prompt_embedding_table_size > 0: - p_embedding_range = [[ - 1, prompt_embedding_table_size // 2, prompt_embedding_table_size - ]] - - prompt_embedding_table = Tensor( - name='prompt_embedding_table', - dtype=self._dtype, - shape=[-1, self.transformer.hidden_size], - dim_range=OrderedDict([ - ('prompt_embedding_table_size', p_embedding_range), - ('hidden_size', [self.transformer.hidden_size]), - ])) - if remove_input_padding: - num_tokens_range = [ - 1, - (max_decoder_input_len * max_batch_size + 1) // 2, - max_decoder_input_len * max_batch_size, - ] - tasks = Tensor(name='tasks', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([ - ('decoder_num_tokens', - [decoder_num_tokens_range]), - ])) - else: - tasks = Tensor(name='tasks', - dtype=trt.int32, - shape=[-1, 1], - dim_range=OrderedDict([ - ('batch_size', bs_range), - ('broadcast_dim', [1]), - ])) - prompt_vocab_size = Tensor(name='prompt_vocab_size', - dtype=trt.int32, - shape=[1], - dim_range=OrderedDict([('size', [1])])) - - result = { - 'decoder_input_ids': input_ids, - 'encoder_output': encoder_output, - 'use_cache': True, - 'attention_mask_params': attention_mask_params, - 'last_token_ids': last_token_ids, - 'kv_cache_params': kv_cache_params, - 'attention_params': attention_params, - 'hidden_states': hidden_states, - 'lora_params': lora_params, - 'cross_kv_cache_gen': cross_kv_cache_gen, - 'cross_kv_reuse': cross_kv_reuse, - 'prompt_embedding_table': prompt_embedding_table, - 'prompt_tasks': tasks, - 'prompt_vocab_size': prompt_vocab_size, - 'skip_cross_attn_blocks': skip_cross_attn_blocks, - } - - return result - - def use_lora(self, lora_config: LoraConfig): - use_lora(self, lora_config, self.trtllm_modules_to_hf_modules) - - @classmethod - def from_hugging_face( - cls, - hf_model_or_dir: Union[str, 'transformers.PreTrainedModel'], - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - ''' Create a MLLaMAForCausalLM object from give parameters - ''' - import transformers - - kwargs.pop('load_by_shard', False) - kwargs.pop('load_model_on_cpu', False) - quant_ckpt_path = kwargs.pop('quant_ckpt_path', None) - - assert hf_model_or_dir is not None - use_preloading = isinstance(hf_model_or_dir, - transformers.PreTrainedModel) - if use_preloading: - hf_model = hf_model_or_dir - hf_config_or_dir = hf_model.config - else: - hf_model_dir = hf_model_or_dir - hf_config_or_dir = hf_model_or_dir - - config = MLLaMAConfig.from_hugging_face(hf_config_or_dir, - dtype=dtype, - mapping=mapping, - quant_config=quant_config, - **kwargs) - - custom_dict = { - "lm_head": "language_model.lm_head", - "transformer.ln_f": "language_model.model.norm", - "transformer": "language_model.model", - "self_attention": "self_attn", - "cross_attention": "cross_attn", - "vocab_embedding": "embed_tokens", - "gate_attn": "cross_attn_attn_gate", - "gate_ffwd": "cross_attn_mlp_gate", - "q_layernorm": "q_norm", - "k_layernorm": "k_norm", - } - - if quant_ckpt_path is not None: - hf_model_dir = quant_ckpt_path - - loader = ModelWeightsLoader(hf_model_dir, custom_dict) - model = cls(config) - loader.generate_tllm_weights(model) - - return model diff --git a/tensorrt_llm/models/mmdit_sd3/__init__.py b/tensorrt_llm/models/mmdit_sd3/__init__.py deleted file mode 100644 index 0d106f45ce46..000000000000 --- a/tensorrt_llm/models/mmdit_sd3/__init__.py +++ /dev/null @@ -1,14 +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. diff --git a/tensorrt_llm/models/mmdit_sd3/config.py b/tensorrt_llm/models/mmdit_sd3/config.py deleted file mode 100644 index 6d15c4d5c961..000000000000 --- a/tensorrt_llm/models/mmdit_sd3/config.py +++ /dev/null @@ -1,132 +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. -from typing import Any, Dict, Optional, Sequence, Tuple - -from ...mapping import Mapping -from ..convert_utils import infer_dtype -from ..modeling_utils import PretrainedConfig, QuantConfig - - -class SD3Transformer2DModelConfig(PretrainedConfig): - - def __init__( - self, - *, - sample_size: int = 128, - patch_size: int = 2, - in_channels: int = 16, - num_layers: int = 24, - attention_head_dim: int = 64, - num_attention_heads: int = 24, - joint_attention_dim: int = 4096, - caption_projection_dim: int = 1536, - pooled_projection_dim: int = 2048, - out_channels: int = 16, - pos_embed_max_size: int = 384, - dual_attention_layers: - Tuple[int] = ( - 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12 - ), # () for sd3.0; (0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12) for sd3.5 - qk_norm: Optional[str] = None, - skip_layers: Optional[Sequence[int]] = None, - use_pretrained_pos_emb: bool = False, - **kwargs): - - kwargs.update({ - 'hidden_size': attention_head_dim * num_attention_heads, - 'num_hidden_layers': num_layers, - 'num_attention_heads': num_attention_heads - }) - super().__init__(**kwargs) - self.sample_size = sample_size - self.patch_size = patch_size - self.in_channels = in_channels - self.num_layers = num_layers - self.attention_head_dim = attention_head_dim - self.num_attention_heads = num_attention_heads - self.joint_attention_dim = joint_attention_dim - self.caption_projection_dim = caption_projection_dim - self.pooled_projection_dim = pooled_projection_dim - self.out_channels = out_channels - self.pos_embed_max_size = pos_embed_max_size - self.dual_attention_layers = dual_attention_layers - self.qk_norm = qk_norm - self.skip_layers = skip_layers - self.use_pretrained_pos_emb = use_pretrained_pos_emb - - def to_dict(self): - output = super().to_dict() - # Serialize the fields added in SD3Transformer2DModelConfig - output['sample_size'] = self.sample_size - output['patch_size'] = self.patch_size - output['in_channels'] = self.in_channels - output['num_layers'] = self.num_layers - output['attention_head_dim'] = self.attention_head_dim - output['num_attention_heads'] = self.num_attention_heads - output['joint_attention_dim'] = self.joint_attention_dim - output['caption_projection_dim'] = self.caption_projection_dim - output['pooled_projection_dim'] = self.pooled_projection_dim - output['out_channels'] = self.out_channels - output['pos_embed_max_size'] = self.pos_embed_max_size - output['dual_attention_layers'] = self.dual_attention_layers - output['qk_norm'] = self.qk_norm - output['skip_layers'] = self.skip_layers - output['use_pretrained_pos_emb'] = self.use_pretrained_pos_emb - return output - - @classmethod - def from_hugging_face_config(cls, - hf_config: Dict[str, Any], - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - sample_size = hf_config['sample_size'] - patch_size = hf_config['patch_size'] - in_channels = hf_config['in_channels'] - num_layers = hf_config['num_layers'] - attention_head_dim = hf_config['attention_head_dim'] - num_attention_heads = hf_config['num_attention_heads'] - joint_attention_dim = hf_config['joint_attention_dim'] - caption_projection_dim = hf_config['caption_projection_dim'] - pooled_projection_dim = hf_config['pooled_projection_dim'] - out_channels = hf_config['out_channels'] - pos_embed_max_size = hf_config['pos_embed_max_size'] - dual_attention_layers = hf_config['dual_attention_layers'] - qk_norm = hf_config['qk_norm'] - skip_layers = None - use_pretrained_pos_emb = kwargs.get('use_pretrained_pos_emb', False) - dtype = infer_dtype(dtype, hf_config.get('torch_dtype')) - - return cls(architecture='SD3Transformer2DModel', - sample_size=sample_size, - patch_size=patch_size, - in_channels=in_channels, - num_layers=num_layers, - attention_head_dim=attention_head_dim, - num_attention_heads=num_attention_heads, - joint_attention_dim=joint_attention_dim, - caption_projection_dim=caption_projection_dim, - pooled_projection_dim=pooled_projection_dim, - out_channels=out_channels, - pos_embed_max_size=pos_embed_max_size, - dual_attention_layers=dual_attention_layers, - qk_norm=qk_norm, - skip_layers=skip_layers, - use_pretrained_pos_emb=use_pretrained_pos_emb, - dtype=dtype, - mapping=mapping, - quantization=quant_config, - **kwargs) diff --git a/tensorrt_llm/models/mmdit_sd3/model.py b/tensorrt_llm/models/mmdit_sd3/model.py deleted file mode 100644 index 546abbeade2c..000000000000 --- a/tensorrt_llm/models/mmdit_sd3/model.py +++ /dev/null @@ -1,620 +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. -from collections import OrderedDict -from typing import Any, Dict, List, Optional - -from ..._utils import str_dtype_to_torch -from ...functional import (Tensor, allgather, chunk, concat, einsum, pad, shape, - unsqueeze) -from ...layers import LayerNorm, Linear -from ...layers.attention import DiffusersAttention -from ...layers.embedding import (CombinedTimestepTextProjEmbeddings, - SD3PatchEmbed) -from ...layers.mlp import (LinearActivation, LinearApproximateGELU, LinearGEGLU, - LinearGELU, LinearSwiGLU) -from ...layers.normalization import (AdaLayerNormContinuous, AdaLayerNormZero, - SD35AdaLayerNormZeroX) -from ...logger import logger -from ...mapping import Mapping -from ...module import Module, ModuleList -from ..model_weights_loader import ModelWeightsLoader -from ..modeling_utils import PretrainedModel -from .config import SD3Transformer2DModelConfig - - -class FeedForward(Module): - - def __init__( - self, - dim: int, - dim_out: Optional[int] = None, - mult: int = 4, - activation_fn: str = "geglu", - inner_dim=None, - bias: bool = True, - mapping=Mapping(), - dtype=None, - ): - super().__init__() - - self.mapping = mapping - self.dtype = dtype - - if inner_dim is None: - inner_dim = int(dim * mult) - dim_out = dim_out if dim_out is not None else dim - - if activation_fn == "gelu": - raise NotImplementedError('GELU only support tanh now.') - if activation_fn == "gelu-approximate": - act_fn = LinearGELU(dim, - inner_dim, - approximate="tanh", - bias=bias, - mapping=mapping, - dtype=dtype) - elif activation_fn == "geglu": - act_fn = LinearGEGLU(dim, - inner_dim, - approximate="tanh", - bias=bias, - mapping=mapping, - dtype=dtype) - elif activation_fn == "geglu-approximate": - act_fn = LinearApproximateGELU(dim, - inner_dim, - bias=bias, - mapping=mapping, - dtype=dtype) - elif activation_fn == "swiglu": - act_fn = LinearSwiGLU(dim, - inner_dim, - bias=bias, - mapping=mapping, - dtype=dtype) - elif activation_fn == "linear-silu": - act_fn = LinearActivation(dim, - inner_dim, - bias=bias, - activation="silu", - mapping=mapping, - dtype=dtype) - - self.net = ModuleList([ - act_fn, - Linear(inner_dim, - dim_out, - bias=bias, - tp_group=self.mapping.tp_group, - tp_size=self.mapping.tp_size, - dtype=self.dtype) - ]) - - def forward(self, hidden_states: Tensor): - for module in self.net: - hidden_states = module(hidden_states) - return hidden_states - - -class JointTransformerBlock(Module): - - def __init__(self, - dim: int, - num_attention_heads: int, - attention_head_dim: int, - context_pre_only: bool = False, - qk_norm: Optional[str] = None, - use_dual_attention: bool = False, - mapping=Mapping(), - dtype=None): - super().__init__() - - self.use_dual_attention = use_dual_attention - self.context_pre_only = context_pre_only - context_norm_type = "ada_norm_continous" if context_pre_only else "ada_norm_zero" - - if use_dual_attention: - self.norm1 = SD35AdaLayerNormZeroX(dim, - mapping=mapping, - dtype=dtype) - else: - self.norm1 = AdaLayerNormZero(dim, mapping=mapping, dtype=dtype) - - if context_norm_type == "ada_norm_continous": - self.norm1_context = AdaLayerNormContinuous( - dim, - dim, - elementwise_affine=False, - eps=1e-6, - bias=True, - norm_type="layer_norm", - dtype=dtype) - elif context_norm_type == "ada_norm_zero": - self.norm1_context = AdaLayerNormZero(dim, dtype=dtype) - else: - raise ValueError( - f"Unknown context_norm_type: {context_norm_type}, currently only support `ada_norm_continous`, `ada_norm_zero`" - ) - - self.attn = DiffusersAttention( - query_dim=dim, - cross_attention_dim=None, - added_kv_proj_dim=dim, - dim_head=attention_head_dim, - heads=num_attention_heads, - out_dim=dim, - context_pre_only=context_pre_only, - bias=True, - qk_norm=qk_norm, - eps=1e-6, - mapping=mapping, - dtype=dtype, - ) - - if use_dual_attention: - self.attn2 = DiffusersAttention( - query_dim=dim, - cross_attention_dim=None, - dim_head=attention_head_dim, - heads=num_attention_heads, - out_dim=dim, - bias=True, - qk_norm=qk_norm, - eps=1e-6, - mapping=mapping, - dtype=dtype, - ) - else: - self.attn2 = None - - self.norm2 = LayerNorm(dim, - elementwise_affine=False, - eps=1e-6, - dtype=dtype) - self.ff = FeedForward(dim=dim, - dim_out=dim, - activation_fn="gelu-approximate", - mapping=mapping, - dtype=dtype) - - if not context_pre_only: - self.norm2_context = LayerNorm(dim, - elementwise_affine=False, - eps=1e-6, - dtype=dtype) - self.ff_context = FeedForward(dim=dim, - dim_out=dim, - activation_fn="gelu-approximate", - mapping=mapping, - dtype=dtype) - else: - self.norm2_context = None - self.ff_context = None - - # let chunk size default to None - self._chunk_size = None - self._chunk_dim = 0 - - def set_chunk_feed_forward(self, - chunk_size: Optional[int] = None, - dim: int = 0): - # Sets chunk feed-forward - self._chunk_size = chunk_size - self._chunk_dim = dim - - @staticmethod - def _chunked_feed_forward(ff: Module, hidden_states: Tensor, chunk_dim: int, - chunk_size: int): - # "feed_forward_chunk_size" can be used to save memory - if hidden_states.shape[chunk_dim] % chunk_size != 0: - raise ValueError( - f"`hidden_states` dimension to be chunked: {hidden_states.shape[chunk_dim]} has to be divisible by chunk size: {chunk_size}. Make sure to set an appropriate `chunk_size` when calling `unet.enable_forward_chunking`." - ) - - num_chunks = hidden_states.shape[chunk_dim] // chunk_size - ff_output = concat( - [ - ff(hid_slice) - for hid_slice in chunk(hidden_states, num_chunks, dim=chunk_dim) - ], - dim=chunk_dim, - ) - return ff_output - - def forward(self, - hidden_states: Tensor, - encoder_hidden_states: Tensor, - temb: Tensor, - joint_attention_kwargs: Optional[Dict[str, Any]] = None, - *args, - **kwargs): - joint_attention_kwargs = joint_attention_kwargs or {} - if self.use_dual_attention: - norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp, norm_hidden_states2, gate_msa2 = self.norm1( - hidden_states, emb=temb) - else: - norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1( - hidden_states, emb=temb) - - if self.context_pre_only: - norm_encoder_hidden_states = self.norm1_context( - encoder_hidden_states, temb) - else: - norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.norm1_context( - encoder_hidden_states, emb=temb) - - # Attention. - attn_output, context_attn_output = self.attn( - hidden_states=norm_hidden_states, - encoder_hidden_states=norm_encoder_hidden_states, - **joint_attention_kwargs, - ) - - # Process attention outputs for the `hidden_states`. - attn_output = unsqueeze(gate_msa, 1) * attn_output - hidden_states = hidden_states + attn_output - - if self.use_dual_attention: - attn_output2 = self.attn2(hidden_states=norm_hidden_states2, - **joint_attention_kwargs) - attn_output2 = unsqueeze(gate_msa2, 1) * attn_output2 - hidden_states = hidden_states + attn_output2 - - norm_hidden_states = self.norm2(hidden_states) - norm_hidden_states = norm_hidden_states * ( - 1 + unsqueeze(scale_mlp, 1)) + unsqueeze(shift_mlp, 1) - - if self._chunk_size is not None: - # "feed_forward_chunk_size" can be used to save memory - ff_output = self._chunked_feed_forward(self.ff, norm_hidden_states, - self._chunk_dim, - self._chunk_size) - else: - ff_output = self.ff(norm_hidden_states) - ff_output = unsqueeze(gate_mlp, 1) * ff_output - hidden_states = hidden_states + ff_output - - # Process attention outputs for the `encoder_hidden_states`. - if self.context_pre_only: - encoder_hidden_states = None - else: - context_attn_output = unsqueeze(c_gate_msa, 1) * context_attn_output - encoder_hidden_states = encoder_hidden_states + context_attn_output - - norm_encoder_hidden_states = self.norm2_context( - encoder_hidden_states) - norm_encoder_hidden_states = norm_encoder_hidden_states * ( - 1 + unsqueeze(c_scale_mlp, 1)) + unsqueeze(c_shift_mlp, 1) - if self._chunk_size is not None: - # "feed_forward_chunk_size" can be used to save memory - context_ff_output = self._chunked_feed_forward( - self.ff_context, norm_encoder_hidden_states, - self._chunk_dim, self._chunk_size) - else: - context_ff_output = self.ff_context(norm_encoder_hidden_states) - encoder_hidden_states = encoder_hidden_states + unsqueeze( - c_gate_mlp, 1) * context_ff_output - - return encoder_hidden_states, hidden_states - - -class SD3Transformer2DModel(PretrainedModel): - config_class = SD3Transformer2DModelConfig - - def __init__(self, config: SD3Transformer2DModelConfig): - super().__init__(config) - self.quant_mode = config.quant_mode - self.mapping = config.mapping - self.dtype = config.dtype - - self.in_channels = config.in_channels - default_out_channels = config.in_channels - self.out_channels = config.out_channels if config.out_channels is not None else default_out_channels - self.inner_dim = config.num_attention_heads * config.attention_head_dim - - self.pos_embed = SD3PatchEmbed( - height=config.sample_size, - width=config.sample_size, - patch_size=config.patch_size, - in_channels=self.in_channels, - embed_dim=self.inner_dim, - pos_embed_max_size=config. - pos_embed_max_size, # hard-code as HF implementation - dtype=self.dtype) - self.time_text_embed = CombinedTimestepTextProjEmbeddings( - embedding_dim=self.inner_dim, - pooled_projection_dim=config.pooled_projection_dim, - mapping=self.mapping, - dtype=self.dtype) - self.context_embedder = Linear(config.joint_attention_dim, - config.caption_projection_dim, - tp_group=self.mapping.tp_group, - tp_size=self.mapping.tp_size, - dtype=self.dtype) - - self.transformer_blocks = ModuleList([ - JointTransformerBlock( - dim=self.inner_dim, - num_attention_heads=config.num_attention_heads, - attention_head_dim=config.attention_head_dim, - context_pre_only=(i == config.num_layers - 1), - qk_norm=config.qk_norm, - use_dual_attention=True - if i in config.dual_attention_layers else False, - mapping=self.mapping, - dtype=self.dtype) for i in range(config.num_layers) - ]) - - self.norm_out = AdaLayerNormContinuous(self.inner_dim, - self.inner_dim, - elementwise_affine=False, - eps=1e-6, - dtype=self.dtype) - self.proj_out = Linear(self.inner_dim, - config.patch_size * config.patch_size * - self.out_channels, - bias=True, - tp_group=self.mapping.tp_group, - tp_size=self.mapping.tp_size, - dtype=self.dtype) - - self.skip_layers = config.skip_layers - self.use_pretrained_pos_emb = config.use_pretrained_pos_emb - self.config = config - - def forward(self, - hidden_states: Tensor, - encoder_hidden_states: Optional[Tensor] = None, - pooled_projections: Optional[Tensor] = None, - timestep: Optional[Tensor] = None, - block_controlnet_hidden_states: List[Tensor] = None, - joint_attention_kwargs: Optional[Dict[str, Any]] = None): - height, width = hidden_states.shape[-2:] - hidden_states = self.pos_embed( - hidden_states) # takes care of adding positional embeddings too. - temb = self.time_text_embed(timestep, pooled_projections) - encoder_hidden_states = self.context_embedder(encoder_hidden_states) - - if self.mapping.cp_size > 1: - hidden_states = chunk(hidden_states, - chunks=self.mapping.cp_size, - dim=1)[self.mapping.cp_rank] - encoder_redundant = encoder_hidden_states.shape[ - 1] % self.mapping.cp_size - encoder_padding_index = tuple( - [0, 0] * (encoder_hidden_states.ndim() - 2) + - [0, self.mapping.cp_size - encoder_redundant]) - if encoder_redundant != 0: - encoder_hidden_states = pad(encoder_hidden_states, - pad=encoder_padding_index) - encoder_hidden_states = chunk(encoder_hidden_states, - chunks=self.mapping.cp_size, - dim=1)[self.mapping.cp_rank] - for index_block, block in enumerate(self.transformer_blocks): - # Skip specified layers - is_skip = True if self.skip_layers is not None and index_block in self.skip_layers else False - - if not is_skip: - encoder_hidden_states, hidden_states = block( - hidden_states=hidden_states, - encoder_hidden_states=encoder_hidden_states, - temb=temb, - joint_attention_kwargs=joint_attention_kwargs, - ) - - # controlnet residual - if block_controlnet_hidden_states is not None and block.context_pre_only is False: - interval_control = len(self.transformer_blocks) / len( - block_controlnet_hidden_states) - hidden_states = hidden_states + block_controlnet_hidden_states[ - int(index_block / interval_control)] - - hidden_states = self.norm_out(hidden_states, temb) - hidden_states = self.proj_out(hidden_states) - if self.mapping.cp_size > 1: - hidden_states = allgather(hidden_states, - group=self.mapping.cp_group, - gather_dim=1) - - # unpatchify - patch_size = self.config.patch_size - height = height // patch_size - width = width // patch_size - - hidden_states = hidden_states.view( - concat([ - shape(hidden_states, 0), height, width, patch_size, patch_size, - self.out_channels - ])) - hidden_states = einsum("nhwpqc->nchpwq", [hidden_states]) - output = hidden_states.view( - concat([ - shape(hidden_states, 0), self.out_channels, height * patch_size, - width * patch_size - ])) - - output.mark_output("output") - return output - - def prepare_inputs(self, max_batch_size, **kwargs): - - def sd3_default_range(max_batch_size): - return [1, max(1, (max_batch_size + 1) // 2), max_batch_size] - - default_range = sd3_default_range - prompt_embeds_len = 256 + 77 # [NOTE] tokenizer_max_length = 77; max_sequence_length = 256 - - hidden_states = Tensor(name='hidden_states', - dtype=self.dtype, - shape=[ - -1, self.in_channels, - self.config.sample_size, - self.config.sample_size - ], - dim_range=OrderedDict([ - ('batch_size', - [default_range(max_batch_size)]), - ('in_channels', [[self.in_channels] * 3]), - ('height', [[self.config.sample_size] * 3]), - ('width', [[self.config.sample_size] * 3]), - ])) - encoder_hidden_states = Tensor( - name='encoder_hidden_states', - dtype=self.dtype, - shape=[-1, prompt_embeds_len, self.config.joint_attention_dim], - dim_range=OrderedDict([ - ('batch_size', [default_range(max_batch_size)]), - ('txt_len', [[prompt_embeds_len] * 3]), - ('joint_attention_dim', [[self.config.joint_attention_dim] * 3 - ]), - ])) - pooled_projections = Tensor( - name='pooled_projections', - dtype=self.dtype, - shape=[-1, self.config.pooled_projection_dim], - dim_range=OrderedDict([ - ('batch_size', [default_range(max_batch_size)]), - ('pooled_projection_dim', - [[self.config.pooled_projection_dim] * 3]), - ])) - timestep = Tensor(name='timestep', - dtype=self.dtype, - shape=[-1], - dim_range=OrderedDict([ - ('batch_size', [default_range(max_batch_size)]), - ])) - return { - "hidden_states": hidden_states, - "encoder_hidden_states": encoder_hidden_states, - "pooled_projections": pooled_projections, - "timestep": timestep, - } - - @classmethod - def from_pretrained(cls, - pretrained_model_name_or_path: str, - dtype='float16', - mapping=Mapping(), - **kwargs): - quant_ckpt_path = kwargs.pop('quant_ckpt_path', None) - - from diffusers import StableDiffusion3Pipeline - - transformer = StableDiffusion3Pipeline.from_pretrained( - pretrained_model_name_or_path, - torch_dtype=str_dtype_to_torch(dtype)).transformer - - config = SD3Transformer2DModelConfig.from_hugging_face_config( - transformer.config, dtype=dtype, mapping=mapping, **kwargs) - - hf_model_dir = transformer.config._name_or_path - custom_dict = {} - if quant_ckpt_path is not None: - hf_model_dir = quant_ckpt_path - - loader = SD3ModelWeightsLoader(hf_model_dir, custom_dict) - model = cls(config) - loader.generate_tllm_weights(model) - return model - - def load(self, weights, from_pruned=False): - required_names = set() - for name, param in self.named_parameters(): - if self.use_pretrained_pos_emb and 'pos_embed' in name: - required_names.add(name) - continue - if param.is_inited(): - continue - if name not in weights: - # Exemption for embedding sharing - if name.endswith('lm_head.weight') and any( - k.endswith('vocab_embedding.weight') - for k in weights.keys()): - continue - if name.endswith('lm_head.per_channel_scale') and any( - k.endswith('vocab_embedding.per_channel_scale') - for k in weights.keys()): - continue - required_names.add(name) - - provided_names = set(weights.keys()) - if not required_names.issubset(provided_names): - raise RuntimeError( - f"Required but not provided tensors:{required_names.difference(provided_names)}" - ) - if not provided_names.issubset(required_names): - logger.warning( - f"Provided but not required tensors: {provided_names.difference(required_names)}" - ) - - for name, param in self.named_parameters(): - if name in provided_names: - if not from_pruned: - try: - param.value = weights[name] - except Exception as e: - raise RuntimeError( - f"Encounter error '{e}' for parameter '{name}'") - else: - param.set_value_or_dummy(weights[name]) - - def enable_forward_chunking(self, - chunk_size: Optional[int] = None, - dim: int = 0): - raise NotImplementedError() - - def disable_forward_chunking(self): - raise NotImplementedError() - - @property - def attn_processors(self): - return None - - def set_attn_processor(self, processor): - raise NotImplementedError() - - def fuse_qkv_projections(self): - raise NotImplementedError() - - def unfuse_qkv_projections(self): - raise NotImplementedError() - - def _set_gradient_checkpointing(self, module, value=False): - raise NotImplementedError() - - -class SD3ModelWeightsLoader(ModelWeightsLoader): - - def translate_to_external_key(self, tllm_key: str, - tllm_to_externel_key_dict: dict): - """Convert and load external checkpoint into a TensorRT LLM model. - """ - trtllm_to_hf_name = { - r"transformer_blocks.(\d+).ff(\w*).net.1.weight": - "transformer_blocks.*.ff*.net.2.weight", - r"transformer_blocks.(\d+).ff(\w*).net.1.bias": - "transformer_blocks.*.ff*.net.2.bias", - } - import re - for k, v in trtllm_to_hf_name.items(): - m = re.match(k, tllm_key) - if m is not None: - matched_pos = m.groups() - placeholders = v.count('*') - assert len(matched_pos) == placeholders - for i in range(len(matched_pos)): - v = v.replace('*', matched_pos[i], 1) - return v - return tllm_key diff --git a/tensorrt_llm/models/model_weights_loader.py b/tensorrt_llm/models/model_weights_loader.py deleted file mode 100644 index a130883e95c7..000000000000 --- a/tensorrt_llm/models/model_weights_loader.py +++ /dev/null @@ -1,402 +0,0 @@ -import glob -import math -import os -import weakref -from enum import Enum -from typing import Callable, List, Optional - -import tensorrt as trt -import torch -from safetensors import safe_open -from tqdm import tqdm -from transformers import PreTrainedModel - -from .._utils import trt_dtype_to_torch -from ..layers.moe import MOEWeightWrapper -from ..logger import logger -from ..quantization.layers import (WeightOnlyGroupwiseQuantColumnLinear, - WeightOnlyGroupwiseQuantRowLinear) - - -class ModelWeightsFormat(Enum): - IN_MEMORY = "in_mem" - SAFETENSORS = "safetensors" - BINARY = "bin" - PYTORCH = "pth" - - -class ModelWeightsLoader: - """Convert and load external checkpoint into a TensorRT LLM model. - - Attributes: - model_dir : Model directory or in-memory torch model. - format : Checkpoint file format. - shards : Shard pointer list (safetensors) or shard dict lists (other types) - shard map : Dict of external checkpoints' keys -> shard index. - tllm_to_externel_key_dict : Dict of TRT-LLM keywords -> External checkpoints' keywords, based on HF LLaMA. - customized_key_dict : Customized dict for updating the default tllm_to_externel_key_dict. - """ - - def __init__(self, model_dir, customized_key_dict: dict = {}) -> None: - - # Checkpoint file format information - self.model_dir = model_dir - self.format = None - self.model = None - self.shards = [] - self.shard_map = {} - - # Key translator vocabulary - self.tllm_to_externel_key_dict = { - "transformer": "model", - "vocab_embedding": "embed_tokens", - "layers": "layers", - "lm_head": "lm_head", - "ln_f": "norm", - "attention": "self_attn", - "qkv": ["q_proj", "k_proj", "v_proj"], - "dense": "o_proj", - "gate": "up_proj", - "proj": "down_proj", - "fc": "gate_proj", - "input_layernorm": "input_layernorm", - "post_layernorm": "post_attention_layernorm", - "kv_cache_scaling_factor": ["k_proj.k_scale", "v_proj.v_scale"], - "kv_cache_rcp_scaling_factor": ["k_proj.k_scale", "v_proj.v_scale"], - } - self.tllm_to_externel_key_dict.update(customized_key_dict) - - self.detect_format() - self.preload() - - def translate_to_external_key( - self, - tllm_key: str, - tllm_to_externel_key_dict: Optional[dict] = None - ) -> str | List[str]: - """Translate TRT-LLM key into HF key or HF key list (e.g. QKV/MoE/GPTQ) - - tllm_key will get translated into HF format section by section. - If one section is responded with multiple hf_keys in a list, \ - the translated keys will also get multiplied accordingly. - tllm_key : "transformer.layers.0.attention. qkv .weight" - | | | | | | - translated: [" model .layers.0.self_attn.q_proj.weight, - " model .layers.0.self_attn.k_proj.weight, - " model .layers.0.self_attn.v_proj.weight] - - Args: - tllm_key (str): Input TRT-LLM key. - tllm_to_externel_key_dict (dict): User specified dict with higher priority. \ - Generated from layer attributes automatically. - - Returns: - hf_keys (str | list[str]) : Translated HF key(s). - """ - tllm_keys = [tllm_key] - d = self.tllm_to_externel_key_dict.copy() - if tllm_to_externel_key_dict is not None: - d.update(tllm_to_externel_key_dict) - for k, v in d.items(): - if k in tllm_key: - # Ensure replacement happen when k covers several full sections in tllm_key - if not any([ - ('.' + k + '.') in tllm_key, - k == tllm_key, - tllm_key.startswith(k) and (k + '.') in tllm_key, - tllm_key.endswith(k) and ('.' + k) in tllm_key, - ]): - continue - if isinstance(v, list): - tllm_keys = [t for t in tllm_keys for _ in range(len(v))] - tllm_keys = [ - s.replace(k, v[idx % len(v)]) - for idx, s in enumerate(tllm_keys) - ] - else: - tllm_keys = [s.replace(k, v) for s in tllm_keys] - - for idx, k in enumerate(tllm_keys): - while ".." in k: - k = k.replace("..", ".") - if k.startswith("."): - k = k[1:] - if k.endswith("."): - k = k[:-1] - tllm_keys[idx] = k - - return tllm_keys[0] if len(tllm_keys) == 1 else tllm_keys - - def detect_format(self): - if isinstance(self.model_dir, dict) or isinstance( - self.model_dir, PreTrainedModel): - self.format = ModelWeightsFormat.IN_MEMORY - elif os.path.isfile(self.model_dir): - if self.model_dir.endswith(".safetensors"): - self.format = ModelWeightsFormat.SAFETENSORS - elif self.model_dir.endswith(".bin"): - self.format = ModelWeightsFormat.BINARY - elif self.model_dir.endswith(".pth"): - self.format = ModelWeightsFormat.PYTORCH - else: - raise NotImplementedError( - "Only safetensors/pickle/binary files are supported.") - elif os.path.isdir(self.model_dir): - file_list = os.listdir(self.model_dir) - if any([f.endswith(".safetensors") for f in file_list]): - self.format = ModelWeightsFormat.SAFETENSORS - elif any([f.endswith(".bin") for f in file_list]): - self.format = ModelWeightsFormat.BINARY - elif any([f.endswith(".pth") for f in file_list]): - self.format = ModelWeightsFormat.PYTORCH - else: - raise NotImplementedError( - "Only safetensors/pickle/binary directories are supported.") - else: - raise NotImplementedError( - "args.model_dir is not a directory, a file or an in-memory module!" - ) - - def preload(self): - # Initialize shards and load_func - if isinstance(self.model_dir, PreTrainedModel): - shard_files = [dict(self.model_dir.named_parameters())] - elif isinstance(self.model_dir, dict): - shard_files = [self.model_dir] - elif os.path.isfile(self.model_dir): - shard_files = [self.model_dir] - elif os.path.isdir(self.model_dir): - shard_files = glob.glob(self.model_dir + "/*." + self.format.value) - else: - raise NotImplementedError( - "args.model_dir is not a directory, a file or an in-memory module!" - ) - shard_files.sort() - if self.format == ModelWeightsFormat.SAFETENSORS: - self.shards = [ - safe_open(f, framework="pt", device="cpu") for f in shard_files - ] - elif self.format == ModelWeightsFormat.BINARY or self.format == ModelWeightsFormat.PYTORCH: - self.shards = [ - torch.load(f, weights_only=True, map_location="cpu", mmap=True) - for f in shard_files - ] - elif self.format == ModelWeightsFormat.IN_MEMORY: - self.shards = [shard_files[0]] - else: - raise NotImplementedError( - "Only *.safetensors/*.pth/*.bin files are supported.") - for idx, shard in enumerate(self.shards): - self.shard_map.update({k: idx for k in shard.keys()}) - - def load_tensor(self, key, tp_size=1, tp_dim=-1, tp_rank=0): - # Retrieve shard index - if key in self.shard_map: - ptr_idx = self.shard_map[key] - else: - if "language_model." + key in self.shard_map: - key = "language_model." + key - ptr_idx = self.shard_map[key] - else: - return None - - if self.format == ModelWeightsFormat.SAFETENSORS: - tensor = self.shards[ptr_idx].get_slice(key) - tensor_shape = tensor.get_shape() - if tensor_shape == []: - tensor = self.shards[ptr_idx].get_tensor(key).unsqueeze(0) - tensor_shape = tensor.shape - else: - tensor = self.shards[ptr_idx][key] - tensor_shape = tensor.shape - - if tp_size <= 1 or tp_dim < 0: - return tensor[:] - else: - if len(tensor_shape) == 1 and (tp_dim > 0 or tensor_shape[0] == 1): - return tensor[:] - else: - width = tensor_shape[tp_dim] - if width == 1: - return tensor[:] - slice_width = math.ceil(width / tp_size) - slice_start = tp_rank * slice_width - slice_end = min((tp_rank + 1) * slice_width, width) - slice_obj = [slice(None)] * len(tensor_shape) - slice_obj[tp_dim] = slice(slice_start, slice_end) - res = tensor[tuple(slice_obj)] - return res - - def load(self, - tllm_key: str, - preprocess: Callable[[int], None] = None, - skip_tp: bool = False, - custom_postprocess_kwargs: dict = {}): - """Load tensor from shards - - This function contains following steps: - 1. Translate tllm_key into external key(s). - 2. Load tensor/tensors partially according to layer attributes. - 3. Call preprocess() if it is not None. - 4. Call layer's post processing function. - 5. Return the dict for updating weight dict. - - Args: - tllm_key (str): TRT-LLM key from model iterators - preprocess (function, Optional): Customized preprocess function for step 3. - skip_tp (bool): Skip TP in case of the derived TP config is inappropriate. - """ - tp_rank = self.model.config.mapping.tp_rank - - sub_module = self.model - for attr in tllm_key.split(".")[:-1]: - sub_module = getattr(sub_module, attr) - param = self.model - for attr in tllm_key.split("."): - param = getattr(param, attr) - if param.is_buffer: - return {} - assert sub_module is not None and param is not None, f"{tllm_key} got Nonetype for parameter or parent module." - - tllm_to_externel_key_dict = getattr(sub_module, - "tllm_to_externel_key_dict", None) - tp_dim = getattr(sub_module, "tp_dim", -1) - require_weight_transpose = ( - isinstance(sub_module, WeightOnlyGroupwiseQuantColumnLinear) - or isinstance(sub_module, WeightOnlyGroupwiseQuantRowLinear)) - if tp_dim >= 0 and require_weight_transpose: - if sub_module.prequant_scaling_factor is not None: - if tllm_key.endswith("prequant_scaling_factor"): - tp_dim = 1 - tp_dim - elif tllm_key.endswith("weights_scaling_factor"): - tp_dim = -1 - elif tllm_key.endswith("weight"): - tp_dim = 1 - tp_dim - tp_size = getattr(sub_module, "tp_size", 1) - # Disable auto TP when num_kv_heads is invalid for split - if getattr(sub_module, "is_qkv", - False) and self.model.config.num_key_value_heads < tp_size: - tp_dim = -1 - tp_size = 1 - if skip_tp: - tp_dim = -1 - tp_size = 1 - if isinstance(sub_module, MOEWeightWrapper): - tp_rank = self.model.config.mapping.moe_tp_rank - external_key = self.translate_to_external_key( - tllm_key, tllm_to_externel_key_dict) - if isinstance(external_key, list): - v = [ - self.load_tensor(k, tp_size, tp_dim, tp_rank) - for k in external_key - ] - else: - v = self.load_tensor(external_key, tp_size, tp_dim, tp_rank) - - if preprocess is not None: - v = preprocess(v) - - if not hasattr(sub_module, "postprocess"): - if isinstance(v, list): - raise ValueError( - f"Param {tllm_key} is translated into {external_key}, post-process function is required." - ) - elif v is None: - weight_dict = {} - else: - weight_dict = {tllm_key: v.to(trt_dtype_to_torch(param.dtype))} - else: - postprocess_kwargs = {"config": self.model.config} - postprocess_kwargs.update(custom_postprocess_kwargs) - v = sub_module.postprocess(tllm_key, v, **postprocess_kwargs) - if isinstance(v, dict): - weight_dict = v - else: - weight_dict = {tllm_key: v} - - for k, v in weight_dict.items(): - if v is not None and not v.is_contiguous(): - weight_dict[k] = v.contiguous() - - return weight_dict - - def update_key_mapping(self, model): - self.model = weakref.ref(model)() - # Auto PP - config = model.config - if config.mapping.has_pp(): - pp_layers = config.mapping.pp_layers(config.num_hidden_layers) - self.tllm_to_externel_key_dict.update({ - f"layers.{tllm_local_layer_idx}": - f"{self.tllm_to_externel_key_dict['layers']}.{hf_global_layer_idx}" - for tllm_local_layer_idx, hf_global_layer_idx in enumerate( - pp_layers) - }) - if self.tllm_to_externel_key_dict['layers'] != 'layers': - del self.tllm_to_externel_key_dict['layers'] - - # Share embedding; only applies to standard structure with lm_head and transformer.vocab_embedding - if hasattr(self.model, 'lm_head') and hasattr( - self.model, 'transformer') and hasattr(self.model.transformer, - 'vocab_embedding'): - lm_head_weights = self.load_tensor( - self.translate_to_external_key('lm_head.weight')) - vocab_embed_weights = self.load_tensor( - self.translate_to_external_key( - 'transformer.vocab_embedding.weight')) - if lm_head_weights is None and vocab_embed_weights is not None: - self.tllm_to_externel_key_dict[ - 'lm_head'] = self.tllm_to_externel_key_dict[ - 'transformer'] + '.' + self.tllm_to_externel_key_dict[ - 'vocab_embedding'] - elif lm_head_weights is not None and vocab_embed_weights is None: - self.tllm_to_externel_key_dict[ - 'vocab_embedding'] = self.tllm_to_externel_key_dict[ - 'lm_head'] - self.model.transformer.vocab_embedding.tllm_to_externel_key_dict = { - 'transformer': '' - } - - def fill(self, weights): - for tllm_key, param in self.model.named_parameters(): - if param.is_buffer: - continue - if tllm_key.endswith('embed_positions_for_gpt_attention'): - continue - w_shape = weights[tllm_key].shape - # WAR for 4bit datatype shape mismatch. - if w_shape != param.shape and param.dtype != trt.fp4: - logger.warning( - f'{tllm_key} has invalid shape {w_shape}. Expected {param.shape}.' - ) - pad = torch.nn.functional.pad - pad_dim = [] - for dim in range(weights[tllm_key].dim()): - current_dim = -1 - dim - pad_dim.append(0) - pad_dim.append( - max(0, param.shape[current_dim] - w_shape[current_dim])) - try: - logger.warning( - f'{tllm_key} is going to be padded by {pad_dim}.') - weights[tllm_key] = pad(weights[tllm_key], - tuple(pad_dim), - value=0) - assert weights[tllm_key].shape == param.shape - except: - raise ValueError( - f'Parameter {tllm_key} has invalid shape {weights[tllm_key].shape} compared with expected shape {param.shape}. Auto padding failed.' - ) - param.value = weights[tllm_key] - - def generate_tllm_weights(self, - model, - custom_postprocess_kwargs: dict = {}): - # For customization, please copy this function and make changes inside the for loop. - self.update_key_mapping(model) - tllm_weights = {} - for tllm_key, _ in tqdm(model.named_parameters()): - tllm_weights.update( - self.load(tllm_key, - custom_postprocess_kwargs=custom_postprocess_kwargs)) - self.fill(tllm_weights) diff --git a/tensorrt_llm/models/modeling_utils.py b/tensorrt_llm/models/modeling_utils.py index b445c128e360..f5a9f27364f7 100644 --- a/tensorrt_llm/models/modeling_utils.py +++ b/tensorrt_llm/models/modeling_utils.py @@ -1,3 +1,18 @@ +# 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. + import argparse import copy import dataclasses @@ -7,51 +22,30 @@ import re from enum import IntFlag, auto from functools import cached_property -from pathlib import Path -from typing import (TYPE_CHECKING, Callable, Dict, Generator, List, Optional, - Union) +from typing import TYPE_CHECKING, Dict, Generator, List, Optional, Union -import numpy as np -import safetensors -import torch from pydantic import Field, PrivateAttr -from .._common import default_net -from .._utils import (QuantModeWrapper, get_init_params, numpy_to_torch, - release_gc, str_dtype_to_torch, str_dtype_to_trt, - trt_dtype_to_torch) +from .._utils import QuantModeWrapper from ..bindings.executor import RuntimeDefaults -from ..functional import (PositionEmbeddingType, Tensor, allgather, constant, - cp_split_plugin, gather_last_token_logits, - index_select, tanh, view) -from ..layers import (MLP, AttentionParams, Embedding, FusedGatedMLP, - FusedRgLru, GatedMLP, KeyValueCacheParams, LoraParams, - PromptTuningEmbedding, RgLru) -from ..layers.attention import Attention, BertAttention -from ..layers.linear import ColumnLinear, Linear, RowLinear -from ..layers.lora import Dora, Lora -from ..layers.moe import MOE, MoeOOTB -from ..llmapi.kv_cache_type import KVCacheType +from ..functional import PositionEmbeddingType from ..llmapi.utils import StrictBaseModel from ..logger import logger from ..mapping import Mapping -from ..module import Module, ModuleList -from ..parameter import Parameter -from ..plugin import init_all_reduce_helper -from ..quantization import QuantMode -from ..quantization.functional import preprocess_weights_for_mixed_gemm -from ..quantization.layers import (FP8Linear, Fp8RowwiseFusedGatedMLP, - Fp8RowwiseGatedMLP, - WeightOnlyGroupwiseQuantLinear, - WeightOnlyGroupwiseQuantRowLinear, - WeightOnlyQuantLinear, - WeightOnlyQuantRowLinear) from ..quantization.mode import (KV_CACHE_QUANT_ALGO_LIST, QUANT_ALGO_LIST, - W8A8_SQ_PLUGIN_LIST, QuantAlgo) -from ..quantization.utils import fp4_utils -from ..top_model_mixin import TopModelMixin -from .convert_utils import weight_only_quantize_dict -from .generation_mixin import GenerationMixin + W8A8_SQ_PLUGIN_LIST, QuantAlgo, QuantMode) + +# QuantConfig and LayerQuantConfig live in the (TensorRT-free) +# tensorrt_llm.quantization package; re-exported here for backward +# compatibility with existing import sites. + +__all__ = [ + 'PretrainedConfig', + 'SpeculativeDecodingMode', + 'QuantConfig', + 'LayerQuantConfig', + 'QuantAlgo', +] @dataclasses.dataclass(kw_only=True, frozen=True) @@ -578,1409 +572,3 @@ def for_each_rank(self) -> "Generator[Self, None, None]": config_copy = copy.deepcopy(self) config_copy.set_rank(rank) yield config_copy - - -class DecoderLayerList(ModuleList): - - def __init__(self, cls, config): - self.num_hidden_layers = config.num_hidden_layers - self.layer_list = config.mapping.pp_layers(config.num_hidden_layers) - self.quant_mode = config.quant_mode - super().__init__([cls(config, idx) for idx in self.layer_list]) - - def forward(self, - hidden_states, - use_cache=False, - attention_mask=None, - kv_cache_params=None, - attention_params=None, - mrope_params=None, - position_ids=None, - lora_params=None, - spec_decoding_params=None, - vision_token_mask=None): - kv_cache_params.fill_none_tensor_list(len(self.layer_list)) - - if use_cache: - presents = [] - - for layer_idx, (layer, past) in enumerate( - zip(self, kv_cache_params.past_key_value)): - - lora_layer_params = None - if lora_params is not None and lora_params.lora_ranks is not None: - lora_layer_params = lora_params.get_layer_params(layer_idx) - - kwargs = {} - if position_ids is not None: - kwargs['position_ids'] = position_ids - if vision_token_mask is not None: - kwargs['vision_token_mask'] = vision_token_mask - if lora_layer_params is not None: - kwargs['lora_layer_params'] = lora_layer_params - if spec_decoding_params is not None: - kwargs['spec_decoding_params'] = spec_decoding_params - if mrope_params is not None: - kwargs['mrope_params'] = mrope_params - - if default_net().plugin_config.reduce_fusion: - if layer_idx + self.layer_list[0] < self.layer_list[-1]: - qkv_activation_scaling_factor = None - if default_net().plugin_config.user_buffer: - qkv_linear = self[layer_idx + 1].attention.qkv - if self.quant_mode.has_fp8_qdq(): - qkv_activation_scaling_factor = constant( - qkv_linear.activation_scaling_factor.raw_value. - copy()) - elif self.quant_mode.has_nvfp4(): - qkv_activation_scaling_factor = constant( - qkv_linear.activation_global_scaling_factor. - raw_value.copy()) - kwargs['next_layer_input_layernorm_args'] = ( - self[layer_idx + 1].input_layernorm.weight.value, - self[layer_idx + 1].input_layernorm.eps, - qkv_activation_scaling_factor) - else: - kwargs['next_layer_input_layernorm_args'] = None - elif default_net().plugin_config.norm_quant_fusion: - if layer_idx < self.layer_list[-1] - self.layer_list[0]: - try: - activation_scaling_factor = constant( - self[layer_idx + 1].attention.qkv. - activation_global_scaling_factor.raw_value.copy()) - except: - activation_scaling_factor = None - kwargs['next_layer_input_layernorm_args'] = ( - self[layer_idx + 1].input_layernorm.weight.value, - self[layer_idx + 1].input_layernorm.eps, - activation_scaling_factor) - else: - kwargs['next_layer_input_layernorm_args'] = None - - hidden_states = layer( - hidden_states, - use_cache=use_cache, - attention_mask=attention_mask, - kv_cache_params=KeyValueCacheParams( - past_key_value=[past], - host_past_key_value_lengths=kv_cache_params. - host_past_key_value_lengths, - host_max_attention_window_sizes=kv_cache_params. - host_max_attention_window_sizes, - host_sink_token_length=kv_cache_params. - host_sink_token_length, - kv_cache_block_offsets=kv_cache_params. - kv_cache_block_offsets, - host_kv_cache_block_offsets=kv_cache_params. - host_kv_cache_block_offsets, - host_kv_cache_pool_pointers=kv_cache_params. - host_kv_cache_pool_pointers, - host_kv_cache_pool_mapping=kv_cache_params. - host_kv_cache_pool_mapping, - cache_indirection=kv_cache_params.cache_indirection), - attention_params=attention_params, - **kwargs) - - if use_cache: - presents.append(hidden_states[1]) - hidden_states = hidden_states[0] - - if use_cache: - return hidden_states, presents - return hidden_states - - -class PostInitCaller(type): - - def __call__(cls, *args, **kwargs): - obj = type.__call__(cls, *args, **kwargs) - obj.__post_init__() - return obj - - -class PretrainedModel(Module, - GenerationMixin, - TopModelMixin, - metaclass=PostInitCaller): - - def __init__(self, config: PretrainedConfig): - super().__init__() - init_all_reduce_helper() - self.config = config - - def __post_init__(self): - from ..quantization.quantize import quantize - quantize(self, self.config.quantization) - - # Currently, use_parallel_embedding must be enabled before weight loading; - # otherwise, the model will be inconsistent with the weights loaded from checkpoint. - optimize_model( - self, use_parallel_embedding=self.config.use_parallel_embedding) - - def release(self): - release_gc() - - def __del__(self): - self.release() - - def check_config(self, config): - raise NotImplementedError( - f"{self.__class__} is an abstract class. Only classes inheriting this class can be called." - ) - - @classmethod - def from_config(cls, config: PretrainedConfig): - return cls(config) - - @classmethod - def from_checkpoint( - cls, - ckpt_dir: str, - rank: Optional[int] = None, - config: Optional[PretrainedConfig] = None, - *, - preprocess_weights_hook: Optional[Callable[[Dict[str, Tensor]], - Dict[str, Tensor]]] = None): - if config is None: - config = PretrainedConfig.from_json_file( - os.path.join(ckpt_dir, 'config.json')) - - if rank is not None: - config.set_rank(rank) - - rank = config.mapping.rank - if config.mapping.cp_size > 1: - # cp_tp_pp rank -> tp_pp rank: because different cp ranks share the same ckpt. - cp_size = config.mapping.cp_size - # rank = pp_rank × tp_size × cp_size + tp_rank × cp_size + cp_rank. - # rank // cp_size is equivalent to pp_rank × tp_size + tp_rank. - rank = rank // cp_size - weights_path = os.path.join(ckpt_dir, f'rank{rank}.safetensors') - - assert os.path.isfile(weights_path) - weights = safetensors.torch.load_file(weights_path) - is_checkpoint_pruned = getattr(config, 'is_pruned', False) - - if preprocess_weights_hook is not None: - weights = preprocess_weights_hook(weights) - - weights = preprocess_weights(weights, - config, - from_pruned=is_checkpoint_pruned) - model = cls(config) - model.load(weights, from_pruned=is_checkpoint_pruned) - return model - - def load(self, weights, from_pruned=False): - required_names = set() - for name, param in self.named_parameters(): - if param.is_inited(): - continue - if name not in weights: - # Exemption for embedding sharing - if name.endswith('lm_head.weight') and any( - k.endswith('vocab_embedding.weight') - for k in weights.keys()): - continue - if name.endswith('lm_head.per_channel_scale') and any( - k.endswith('vocab_embedding.per_channel_scale') - for k in weights.keys()): - continue - required_names.add(name) - - provided_names = set(weights.keys()) - - if not required_names.issubset(provided_names): - raise RuntimeError( - f"Required but not provided tensors:{required_names.difference(provided_names)}" - ) - if not provided_names.issubset(required_names): - logger.warning( - f"Provided but not required tensors: {provided_names.difference(required_names)}" - ) - - for name, param in self.named_parameters(): - if name in provided_names: - if not from_pruned: - try: - param.value = weights[name] - except Exception as e: - raise RuntimeError( - f"Encounter error '{e}' for parameter '{name}'") - else: - param.set_value_or_dummy(weights[name]) - - def save_checkpoint(self, output_dir, save_config=True): - # multiple ranks could share same config.json, so adding a save_config parameter to let user avoiding writing config.json in all ranks - rank = self.config.mapping.rank - weights = { - name: numpy_to_torch(param.raw_value) - for name, param in self.named_parameters() - } - # If there are some tensors share memory, this will lead to error when we call "save_file". So, for repeated tensors, we - # clone the tensors to prevent this issue. - data_ptrs = set() - for name, param in weights.items(): - if param.data_ptr() in data_ptrs: - weights[name] = param.clone() - data_ptrs.add(weights[name].data_ptr()) - safetensors.torch.save_file( - weights, os.path.join(output_dir, f'rank{rank}.safetensors')) - if save_config: - self.config.to_json_file(os.path.join(output_dir, 'config.json')) - - def prepare_inputs( - self, - max_batch_size, - max_input_len, - max_seq_len, - max_num_tokens, - use_cache, - max_beam_width: int = 1, - opt_num_tokens: int = None, - prompt_embedding_table_size: int = 0, - position_encoding_2d: bool = False, - max_draft_len: int = 0, - speculative_decoding_draft_tokens_external: bool = False, - spec_decoding_is_generation_length_variable: bool = False, - gather_context_logits: bool = False, - lora_target_modules: List[str] = None, - opt_batch_size: int = 0, - num_hidden_layers: int = None, - mrope_rotary_cos_sin_size: int = None, - ): - '''@brief: Prepare inputs Tensors for the model, the given sizes are used to determine the - ranges of the dimensions when using TRT dynamic shapes. - - @return: a list containing values which can be fed into the self.forward() - ''' - - # Prepare inputs - remove_input_padding = default_net().plugin_config.remove_input_padding - use_gpt_attention_plugin = default_net( - ).plugin_config.gpt_attention_plugin - use_gemm_plugin = default_net().plugin_config.gemm_plugin - paged_kv_cache = default_net().plugin_config.paged_kv_cache - tokens_per_block = default_net().plugin_config.tokens_per_block - use_lora_plugin = default_net().plugin_config.lora_plugin - multiple_profiles = default_net().plugin_config.multiple_profiles - streamingllm = default_net().plugin_config.streamingllm - pp_reduce_scatter = default_net().plugin_config.pp_reduce_scatter - - kv_cache_type = None - if not use_cache: - kv_cache_type = KVCacheType.DISABLED - else: - if paged_kv_cache: - kv_cache_type = KVCacheType.PAGED - else: - kv_cache_type = KVCacheType.CONTINUOUS - - model_inputs = self.prepare_basic_inputs( - max_batch_size=max_batch_size, - max_beam_width=max_beam_width, - max_input_len=max_input_len, - max_seq_len=max_seq_len, - hidden_size=self.config.hidden_size, - num_kv_heads=self.config.num_key_value_heads, - head_size=self.config.head_size, - num_layers=num_hidden_layers - if num_hidden_layers is not None else self.config.num_hidden_layers, - kv_dtype=str_dtype_to_trt(self.config.kv_dtype), - remove_input_padding=remove_input_padding, - use_gpt_attention_plugin=use_gpt_attention_plugin, - use_gemm_plugin=use_gemm_plugin, - kv_cache_type=kv_cache_type, - tokens_per_block=tokens_per_block, - num_heads=self.config.num_attention_heads, - max_num_tokens=max_num_tokens, - opt_num_tokens=opt_num_tokens, - dtype=str_dtype_to_trt(self.config.dtype), - prompt_embedding_table_size=prompt_embedding_table_size, - position_encoding_2d=position_encoding_2d, - mapping=self.config.mapping, - gather_context_logits=gather_context_logits, - use_lora_plugin=use_lora_plugin, - max_draft_len=max_draft_len, - speculative_decoding_draft_tokens_external= - speculative_decoding_draft_tokens_external, - spec_decoding_is_generation_length_variable= - spec_decoding_is_generation_length_variable, - lora_target_modules=lora_target_modules, - multiple_profiles=multiple_profiles, - streamingllm=streamingllm, - opt_batch_size=opt_batch_size, - pp_reduce_scatter=pp_reduce_scatter, - mrope_rotary_cos_sin_size=mrope_rotary_cos_sin_size) - - result = { - 'input_ids': - model_inputs['input_ids'], - 'position_ids': - model_inputs['position_ids'], - 'use_cache': - kv_cache_type != KVCacheType.DISABLED, - 'last_token_ids': - model_inputs['last_token_ids'], - 'attention_mask': - model_inputs['attention_mask'], - 'kv_cache_params': - KeyValueCacheParams( - past_key_value=model_inputs['past_key_value'], - host_past_key_value_lengths=model_inputs[ - 'host_past_key_value_lengths'], - host_max_attention_window_sizes=model_inputs[ - 'host_max_attention_window_sizes'], - host_sink_token_length=model_inputs['host_sink_token_length'], - kv_cache_block_offsets=model_inputs['kv_cache_block_offsets'], - host_kv_cache_block_offsets=model_inputs[ - 'host_kv_cache_block_offsets'], - host_kv_cache_pool_pointers=model_inputs[ - 'host_kv_cache_pool_pointers'], - host_kv_cache_pool_mapping=model_inputs[ - 'host_kv_cache_pool_mapping'], - cache_indirection=model_inputs['cache_indirection'], - ), - 'attention_params': - AttentionParams( - sequence_length=model_inputs['sequence_length'], - context_lengths=model_inputs['context_lengths'], - host_context_lengths=model_inputs['host_context_lengths'], - max_context_length=max_input_len, - host_request_types=model_inputs['host_request_types'], - host_runtime_perf_knobs=model_inputs['host_runtime_perf_knobs'], - host_context_progress=model_inputs['host_context_progress'], - ) - } - - if prompt_embedding_table_size > 0: - result['prompt_embedding_table'] = model_inputs[ - 'prompt_embedding_table'] - result['prompt_tasks'] = model_inputs['tasks'] - result['prompt_vocab_size'] = model_inputs['prompt_vocab_size'] - if model_inputs['hidden_states_input'] is not None: - result['hidden_states'] = model_inputs['hidden_states_input'] - if use_lora_plugin: - result['lora_params'] = LoraParams( - model_inputs['lora_ranks'], - model_inputs['lora_weights_pointers'], - host_context_lengths=model_inputs['host_context_lengths'], - host_request_types=model_inputs['host_request_types']) - if model_inputs['spec_decoding_params'] is not None: - result['spec_decoding_params'] = model_inputs[ - 'spec_decoding_params'] - if model_inputs['mrope_params'] is not None: - result['mrope_params'] = model_inputs['mrope_params'] - - return result - - @classmethod - def quantize( - cls, - hf_model_dir: str, - output_dir: str, - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - *, - device: str = 'cuda', - calib_dataset: str = 'cnn_dailymail', - calib_batches: int = 512, - calib_batch_size: int = 1, - calib_max_seq_length: int = 512, - random_seed: int = 1234, - tokenizer_max_seq_length: int = 2048, - **kwargs, - ): - config_cls = getattr(cls, 'config_class', None) - if config_cls is None: - raise NotImplementedError( - f"{cls.__name__} has not implemented corresponding config class, which is needed for correct config parsing." - ) - config: PretrainedConfig = config_cls.from_hugging_face( - hf_model_dir, - dtype=dtype, - mapping=mapping, - quant_config=quant_config, - **kwargs) - if config.mapping.moe_ep_size > 1: - raise NotImplementedError( - "Quantization for expert parallelism is not supported") - if not config.quantization._requires_modelopt_quantization: - raise ValueError( - f"The quant_config ({quant_config}) should not call modelopt quantization" - ) - - from ..quantization import quantize_and_export - quantize_and_export( - model_dir=str(hf_model_dir), - device=device, - calib_dataset=calib_dataset, - dtype=config.dtype, - qformat=config.quantization._get_modelopt_qformat(), - kv_cache_dtype=config.quantization._get_modelopt_kv_cache_dtype(), - calib_size=calib_batches, - batch_size=calib_batch_size, - calib_max_seq_length=calib_max_seq_length, - awq_block_size=config.quantization.group_size, - output_dir=output_dir, - tp_size=config.mapping.tp_size, - pp_size=config.mapping.pp_size, - cp_size=config.mapping.cp_size, - seed=random_seed, - tokenizer_max_seq_length=tokenizer_max_seq_length, - ) - - -class DecoderModelForCausalLM(PretrainedModel): - - def __init__(self, config: PretrainedConfig, transformer, lm_head): - super().__init__(config) - self.transformer = transformer - self.lm_head = lm_head - self.mup_width_multiplier = getattr(config, 'mup_width_multiplier', - None) - # Create constant attention parameters to be reused by all layers. - Attention.create_attention_const_params(self, config) - self.position_embedding_type = config.position_embedding_type - - def forward(self, - input_ids: Tensor, - position_ids=None, - use_cache=False, - last_token_ids=None, - attention_mask=None, - kv_cache_params=None, - attention_params=None, - mrope_params=None, - hidden_states=None, - prompt_embedding_table: Optional[Tensor] = None, - prompt_tasks: Optional[Tensor] = None, - prompt_vocab_size: Optional[Tensor] = None, - lora_params=None, - spec_decoding_params=None): - - # fill attention params. - attention_params = Attention.fill_attention_params( - self, attention_params) - - # split the sequence for context parallelism - if self.config.mapping.cp_size > 1: - if len(input_ids.shape) == 1: - # input shape is [-1] - input_ids, cp_join_index = cp_split_plugin( - input_ids, - attention_params.host_request_types, - attention_params.host_context_lengths, - self.config.mapping.cp_size, - self.config.mapping.cp_rank, - ) - else: - assert False, "Context parallelism with non-remove-padding is not supported yet." - - is_gemma_2_cg = self.config.has_config_group(Gemma2ConfigGroup) - is_gemma_3_cg = self.config.has_config_group(Gemma3ConfigGroup) - - kwargs = { - 'input_ids': input_ids, - 'position_ids': position_ids, - 'use_cache': use_cache, - 'attention_mask': attention_mask, - 'kv_cache_params': kv_cache_params, - 'attention_params': attention_params, - } - if lora_params is not None: - kwargs['lora_params'] = lora_params - if hidden_states is not None: - kwargs['hidden_states'] = hidden_states - if prompt_embedding_table is not None: - kwargs['prompt_embedding_table'] = prompt_embedding_table - if prompt_tasks is not None: - kwargs['prompt_tasks'] = prompt_tasks - if prompt_vocab_size is not None: - kwargs['prompt_vocab_size'] = prompt_vocab_size - - if spec_decoding_params is not None: - kwargs['spec_decoding_params'] = spec_decoding_params - if mrope_params is not None: - kwargs['mrope_params'] = mrope_params - - hidden_states = self.transformer.forward(**kwargs) - - if use_cache: - hidden_states, presents = hidden_states - - # All gather and rebuild sequence after transformer layer for context parallelism - if self.config.mapping.cp_size > 1: - if len(hidden_states.shape) == 2: - hidden_states = allgather(hidden_states, - self.config.mapping.cp_group, - gather_dim=0) - hidden_states = view(hidden_states, - [-1, hidden_states.shape[-1]]) - hidden_states = index_select(hidden_states, 0, cp_join_index) - else: - assert False, "Context parallelism with non-remove-padding is not supported yet." - - if self.config.mapping.is_last_pp_rank(): - all_hidden_states = hidden_states - hidden_states = gather_last_token_logits( - hidden_states, last_token_ids, - default_net().plugin_config.remove_input_padding) - - # [batch_size, hidden_size] -> [batch_size, vocab_size] - lm_logits = self.lm_head(hidden_states) - if hasattr(self.config, 'output_multiplier_scale'): - lm_logits *= getattr(self.config, 'output_multiplier_scale', 1) - if self.mup_width_multiplier is not None: - lm_logits = lm_logits / self.mup_width_multiplier - if is_gemma_2_cg or is_gemma_3_cg: - softcap = self.config.get_config_group( - Gemma2ConfigGroup if not is_gemma_3_cg else - Gemma3ConfigGroup).final_logit_softcapping - if softcap: - lm_logits = lm_logits * float(1 / softcap) - lm_logits = tanh(lm_logits) * float(softcap) - lm_logits.mark_output('logits', self.config.logits_dtype) - else: - hidden_states.mark_output('hidden_states_output', self.config.dtype) - - if use_cache and not default_net().plugin_config.paged_kv_cache: - for i, present in zip( - self.config.mapping.pp_layers( - self.config.num_hidden_layers), presents): - present.mark_output(f'present_key_value_{i}', - self.config.kv_dtype) - if self.config.mapping.is_last_pp_rank(): - return (lm_logits, presents, hidden_states) - return (hidden_states, presents) - else: - if self.config.mapping.is_last_pp_rank(): - return lm_logits, hidden_states, all_hidden_states - return hidden_states - - -def fuse_gate_mlp( - model: PretrainedModel, - gemm_swiglu_plugin_dtype: Optional[str] = None, - low_latency_gemm_swiglu_plugin_dtype: Optional[str] = None, -) -> PretrainedModel: - from ..quantization.quantize import fp8_quantize - - for name, mlp, layer in model.named_modules_with_parent(): - if isinstance(mlp, GatedMLP): - init_params = get_init_params(mlp) - - hidden_act = init_params["hidden_act"] - if hidden_act not in ["silu", "gelu"]: - logger.warning( - f"fuse_gate_mlp cannot be done for {name} due to unsupported activation {hidden_act}. Skipping." - ) - continue - - init_params["inner_layernorm"] = mlp.inner_layernorm is not None - fused_layer = FusedGatedMLP(**init_params) - - fc_name = name + '.fc' - layer_quant_cfg = model.config._get_quant_cfg(fc_name) - layer_quant_algo = layer_quant_cfg.quant_algo - if layer_quant_algo != QuantAlgo.FP8 and layer_quant_algo is not None: - continue - - if isinstance(model.config.quantization.exclude_modules, list) \ - and fc_name in model.config.quantization.exclude_modules: - layer_quant_algo = None - - if layer_quant_algo == QuantAlgo.FP8: - fused_layer = fp8_quantize(fused_layer, layer_quant_cfg) - - if isinstance(mlp.dtype, str): - dtype = str_dtype_to_torch(mlp.dtype) - else: - dtype = trt_dtype_to_torch(mlp.dtype) - - gate_weight = numpy_to_torch(mlp.gate.weight.raw_value) - fc_weight = numpy_to_torch(mlp.fc.weight.raw_value) - assert gate_weight.dtype == fc_weight.dtype - need_qdq = gate_weight.dtype == torch.float8_e4m3fn - - gate_weight = gate_weight.to(dtype) - fc_weight = fc_weight.to(dtype) - # dequantize if needed - if need_qdq: - gate_weight = gate_weight.to(dtype) * numpy_to_torch( - mlp.gate.weights_scaling_factor.raw_value) - fc_weight = fc_weight.to(dtype) * numpy_to_torch( - mlp.fc.weights_scaling_factor.raw_value) - - # concat - fused_weight = torch.cat([gate_weight, fc_weight], dim=0) - - fused_weight_scaling_factor = numpy_to_torch( - max( - mlp.gate.weights_scaling_factor.raw_value, - mlp.fc.weights_scaling_factor.raw_value, - )) - # quantize if needed - if need_qdq: - fused_weight = (fused_weight / - fused_weight_scaling_factor).to( - torch.float8_e4m3fn) - - if gemm_swiglu_plugin_dtype == 'fp8' or low_latency_gemm_swiglu_plugin_dtype == 'fp8': - # gemm_swiglu_plugin needs (k, n) weights - # but weights should still be k-major for fp8 - fused_layer.fused_fc.weight = Parameter( - shape=(fused_layer.fused_fc.in_features, - fused_layer.fused_fc.out_features), - dtype='fp8') - fused_layer.fused_fc.weight.value = fused_weight.view( - fused_layer.fused_fc.in_features, - fused_layer.fused_fc.out_features) - else: - fused_layer.fused_fc.weight.value = fused_weight - fused_layer.fused_fc.weights_scaling_factor.value = fused_weight_scaling_factor - - fused_layer.fused_fc.activation_scaling_factor.value = max( - mlp.gate.activation_scaling_factor.raw_value, - mlp.fc.activation_scaling_factor.raw_value, - ) - - if mlp.bias: - fused_layer.fused_fc.bias.value = np.concatenate( - [mlp.gate.bias.raw_value, mlp.fc.bias.raw_value], - axis=0) - elif layer_quant_algo is None: - fused_layer.fused_fc.weight.value = np.concatenate( - [ - mlp.gate.weight.raw_value, - mlp.fc.weight.raw_value, - ], - axis=0, - ) - if mlp.bias: - fused_layer.fused_fc.bias.value = np.concatenate( - [mlp.gate.bias.raw_value, mlp.fc.bias.raw_value], - axis=0) - else: - raise ValueError(f'Unsupported quant algo: {layer_quant_algo}') - - fused_layer.proj = mlp.proj - fused_layer.inner_layernorm = mlp.inner_layernorm - - _, mlp_name = name.rsplit('.', 1) - setattr(layer, mlp_name, fused_layer) - - elif isinstance(mlp, Fp8RowwiseGatedMLP): - init_params = get_init_params(mlp) - - hidden_act = init_params["hidden_act"] - if hidden_act not in ["silu", "gelu"]: - logger.warning( - f"fuse_gate_mlp cannot be done for {name} due to unsupported activation {hidden_act}. Skipping." - ) - continue - - if mlp.clamp_val is not None: - init_params["clamp_val"] = mlp.clamp_val.raw_value.tolist() - fused_layer = Fp8RowwiseFusedGatedMLP(**init_params) - fused_layer.fused_fc.weight.value = np.concatenate( - [ - mlp.gate.weight.raw_value, - mlp.fc.weight.raw_value, - ], - axis=0, - ) - fused_layer.fused_fc.per_channel_scale.value = np.concatenate( - [ - mlp.gate.per_channel_scale.raw_value, - mlp.fc.per_channel_scale.raw_value, - ], - axis=0, - ) - if mlp.bias: - fused_layer.fused_fc.bias.value = np.concatenate( - [mlp.gate.bias.raw_value, mlp.fc.bias.raw_value], axis=0) - - fused_layer.proj = mlp.proj - _, mlp_name = name.rsplit('.', 1) - setattr(layer, mlp_name, fused_layer) - - return model - - -def unfuse_qkv_gemm(model: PretrainedModel) -> PretrainedModel: - '''Split all the models' Attention layer's QKV GEMM into 3 GEMMs layer.q layer.k, layer.v and return the changed model - ''' - from ..quantization.quantize import quantize - - for name, layer in model.named_modules(): - if isinstance(layer, Attention) and not layer.cross_attention: - assert layer.tp_size == 1, "unfuse_qkv_gemm requires tp_size == 1" - if layer.qkv is None: - continue - qkv_params = get_init_params(layer.qkv, ColumnLinear) - qkv_params["bias"] = qkv_params["bias"] is not None - qkv_params["strict_dtype"] = qkv_params.get( - "strict_dtype") is not None - q = ColumnLinear( - **{ - **qkv_params, - "out_features": - layer.tp_size * layer.num_attention_heads * - layer.attention_head_size, - }) - k = ColumnLinear( - **{ - **qkv_params, - "out_features": - layer.tp_size * layer.num_attention_kv_heads * - layer.attention_head_size, - }) - v = ColumnLinear( - **{ - **qkv_params, - "out_features": - layer.tp_size * layer.num_attention_kv_heads * - layer.attention_head_size, - }) - layer_quant_cfg = model.config._get_quant_cfg(name + '.qkv') - q = quantize(q, layer_quant_cfg) - k = quantize(k, layer_quant_cfg) - v = quantize(v, layer_quant_cfg) - out_features = q.out_features + k.out_features + v.out_features - if isinstance(layer.qkv, ( - WeightOnlyQuantLinear, - WeightOnlyQuantRowLinear, - WeightOnlyGroupwiseQuantLinear, - WeightOnlyGroupwiseQuantRowLinear, - )): - out_dim = 1 - else: - out_dim = 0 - if layer.qkv.weight.is_inited(): - qkv_weight = layer.qkv.weight.raw_value - weights = np.split(qkv_weight, [ - qkv_weight.shape[out_dim] * q.out_features // out_features, - qkv_weight.shape[out_dim] * - (q.out_features + k.out_features) // out_features, - ], - axis=out_dim) - for gemm, weight in zip([q, k, v], weights): - gemm.weight.value = weight - if layer.qkv.bias is not None and layer.qkv.bias.is_inited(): - qkv_bias = layer.qkv.bias.raw_value - biases = np.split(qkv_bias, [ - qkv_bias.shape[out_dim] * q.out_features // out_features, - qkv_bias.shape[out_dim] * - (q.out_features + k.out_features) // out_features, - ], - axis=out_dim) - for gemm, bias in zip([q, k, v], biases): - gemm.bias.value = bias - for name, parameter in layer.qkv._parameters.items(): - if name not in ["weight", "bias"]: - for gemm in [q, k, v]: - setattr(gemm, name, parameter) - layer.q = q - layer.k = k - layer.v = v - layer.qkv = None - return model - - -def fuse_rg_lru(model: PretrainedModel) -> PretrainedModel: - for name, rg_lru, parent in model.named_modules_with_parent(): - if isinstance(rg_lru, RgLru): - fused_layer = FusedRgLru(**get_init_params(rg_lru)) - fused_layer.gate.weight.value = np.concatenate( - [ - rg_lru.input_gate.weight.raw_value, - rg_lru.recurrent_gate.weight.raw_value, - ], - axis=-1, - ) - fused_layer.gate.bias.value = np.concatenate( - [ - rg_lru.input_gate.bias.raw_value, - rg_lru.recurrent_gate.bias.raw_value, - ], - axis=-1, - ) - fused_layer.recurrent_param.value = rg_lru.recurrent_param.raw_value - rg_lru_name = name.rsplit('.', 1)[-1] - setattr(parent, rg_lru_name, fused_layer) - return model - - -def set_prompt_tuning(model: PretrainedModel) -> PretrainedModel: - '''Replace the given models embedding layer with a PromptTuningEmbedding layer in-place, return the changed model - Pre-conditions: vocab_embedding exists - Post-conditions: isinstance(vocab_embedding, PromptTuningEmbedding) - - ''' - for name, embedding, parent in model.named_modules_with_parent(): - layer_name = name.rsplit('.', 1)[-1] - if layer_name == "vocab_embedding" and isinstance(embedding, Embedding): - ptuning_embedding = PromptTuningEmbedding( - **get_init_params(embedding)) - ptuning_embedding.weight.value = embedding.weight.raw_value - parent.vocab_embedding = ptuning_embedding - return model - - -def add_lora(model: PretrainedModel, - max_lora_rank: Optional[int], - with_dora: bool = False) -> PretrainedModel: - ''' Add lora layers to the Attention/BertAttention/Linear/RowLinear/FusedGatedMLP layers to the given model, return the changed model - ''' - for name, layer in model.named_modules(): - max_rank = max_lora_rank - if isinstance(layer, (Attention, BertAttention)): - if max_rank is None: - max_rank = min( - layer.hidden_size, - layer.num_attention_heads * layer.attention_head_size, - layer.num_attention_kv_heads * layer.attention_head_size) - layer.qkv_lora = Lora( - in_hidden_size=layer.hidden_size, - out_hidden_sizes=[ - layer.num_attention_heads * layer.attention_head_size, - layer.num_attention_kv_heads * layer.attention_head_size, - layer.num_attention_kv_heads * layer.attention_head_size - ], - max_low_rank=max_rank, - ) - - if with_dora: - layer.qkv_dora = Dora(out_hidden_sizes=[ - layer.num_attention_heads * layer.attention_head_size, - layer.num_attention_kv_heads * layer.attention_head_size, - layer.num_attention_kv_heads * layer.attention_head_size - ], ) - - if isinstance(layer, (Linear, RowLinear)): - if max_rank is None: - max_rank = min(layer.in_features, layer.out_features) - layer.lora = Lora( - in_hidden_size=layer.in_features, - out_hidden_sizes=[layer.out_features], - max_low_rank=max_rank, - ) - if with_dora: - layer.dora = Dora(out_hidden_sizes=[layer.out_features]) - - if isinstance(layer, (MLP, FusedGatedMLP)): - if max_rank is None: - max_rank = min(layer.hidden_size, - layer.ffn_hidden_size // layer.tp_size) - layer.lora = Lora( - in_hidden_size=layer.hidden_size, - out_hidden_sizes=[ - layer.ffn_hidden_size // layer.tp_size, - layer.ffn_hidden_size // layer.tp_size - ], - max_low_rank=max_rank, - ) - - if isinstance(layer, FusedGatedMLP): - layer.fused_gate_up_lora = Lora( - in_hidden_size=layer.hidden_size, - out_hidden_sizes=[ - layer.ffn_hidden_size * 2 // layer.tp_size - ], - max_low_rank=max_rank, - ) - - if with_dora: - layer.dora = Dora(out_hidden_sizes=[ - layer.ffn_hidden_size // layer.tp_size, - layer.ffn_hidden_size // layer.tp_size - ], ) - - if isinstance(layer, FusedGatedMLP): - layer.fused_gate_up_dora = Dora(out_hidden_sizes=[ - layer.ffn_hidden_size * 2 // layer.tp_size - ], ) - - if isinstance(layer, MOE): - if max_rank is None: - max_rank = min(layer.hidden_size, - layer.ffn_hidden_size // layer.tp_size) - layer.max_low_rank = max_rank - return model - - -def to_ootb_moe(model: PretrainedModel) -> PretrainedModel: - ''' Use OOTB MoE instead of MoE plugin, return the changed model - ''' - for name, layer, parent in model.named_modules_with_parent(): - if isinstance(layer, MOE): - layer_name = name.rsplit('.', 1)[-1] - ootb_layer = layer.to(MoeOOTB, model.config.quantization) - setattr(parent, layer_name, ootb_layer) - return model - - -def parallelize_embedding(model: PretrainedModel) -> PretrainedModel: - for name, embedding, parent in model.named_modules_with_parent(): - layer_name = name.rsplit('.', 1)[-1] - if isinstance(embedding, Embedding) and embedding.tp_group is None: - init_params = get_init_params(embedding) - init_params["tp_group"] = model.config.mapping.tp_group - init_params["tp_size"] = model.config.mapping.tp_size - init_params["tp_rank"] = model.config.mapping.tp_rank - init_params["sharding_dim"] = model.config.embedding_sharding_dim - new_embedding = embedding.__class__(**init_params) - setattr(parent, layer_name, new_embedding) - return model - - -def share_embedding(model: PretrainedModel) -> PretrainedModel: - lm_head = None - vocab_embedding = None - for name, layer in model.named_modules(): - layer_name = name.rsplit('.', 1)[-1] - if layer_name == "lm_head": - lm_head = layer - if layer_name == "vocab_embedding": - vocab_embedding = layer - if lm_head is not None and vocab_embedding is not None: - break - - # Cannot find either lm_head or vocab_embedding, e.g., pipeline parallel - if lm_head is None or vocab_embedding is None: - return model - - # lm_head and vocab_embedding have different shapes, e.g., tensor parallel without embedding parallel - if lm_head.weight.shape != vocab_embedding.weight.shape: - return model - - # lm_head can have a different type if quantized - if lm_head.weight.dtype != vocab_embedding.weight.dtype: - return model - - # Don't assume weight can be shared if vocab_embedding is not initialized, e.g., dummy weights - if not vocab_embedding.weight.is_inited(): - return model - - if lm_head.weight.is_inited(): - lm_head_weight = numpy_to_torch(lm_head.weight.raw_value) - vocab_embed_weight = numpy_to_torch(vocab_embedding.weight.raw_value) - # The lm_head and vocab_embedding have different weights - if (lm_head_weight - vocab_embed_weight).abs().max().item() > 1e-6: - return model - - lm_head.weight = vocab_embedding.weight - if getattr(lm_head, 'per_channel_scale', None) and getattr( - vocab_embedding, 'per_channel_scale', None): - lm_head.per_channel_scale = vocab_embedding.per_token_scale - return model - - -def set_fp8_context_fmha(model: PretrainedModel) -> PretrainedModel: - for name, layer in model.named_modules(): - if isinstance(layer, Attention) and hasattr( - layer.dense, 'activation_scaling_factor'): - scale = [1.0] / layer.dense.activation_scaling_factor.raw_value - layer.attention_output_orig_quant_scale = Parameter( - value=scale.astype(np.float32), dtype='float32') - elif isinstance(layer, Attention) and hasattr( - layer.dense, 'activation_global_scaling_factor'): - scale = [1.0 - ] / layer.dense.activation_global_scaling_factor.raw_value - layer.attention_output_orig_quant_scale = Parameter( - value=scale.astype(np.float32), dtype='float32') - - return model - - -def set_fuse_fp4_quant(model: PretrainedModel) -> PretrainedModel: - for name, layer in model.named_modules(): - if isinstance(layer, Attention) and hasattr( - layer.dense, 'activation_global_scaling_factor'): - scale = [1.0 - ] / layer.dense.activation_global_scaling_factor.raw_value - layer.attention_output_sf_scale = Parameter(value=scale.astype( - np.float32), - dtype='float32') - - return model - - -def optimize_model( - model: PretrainedModel, - use_parallel_embedding: bool = False, - share_embedding_table: bool = False, - use_ootb_moe: bool = False, - use_fused_mlp: bool = False, - gemm_swiglu_plugin_dtype: Optional[str] = None, - low_latency_gemm_swiglu_plugin_dtype: Optional[str] = None, - use_fused_rg_lru: bool = False, - use_unfused_qkv_gemm: bool = False, - use_prompt_tuning: bool = False, - use_lora: bool = False, - max_lora_rank: Optional[int] = None, - use_fp8_context_fmha: bool = False, - fuse_fp4_quant: bool = False, - use_optimize_cross_qkv: bool = False, - use_dora: bool = False, -) -> PretrainedModel: - """ - Run optimization passes on model. - There are dependencies between some passes, - so we always run passes in the order of arguments to guarantee the execution order. - """ - # before weight loading - if use_parallel_embedding: - model = parallelize_embedding(model) - - if share_embedding_table: - # if share_embedding_table is enabled, only one copy of the embedding table is stored in converted ckpt - # this pass is required to make lm_head.weight and vocab_embedding.weight point to the same tensor - # however even if share_embedding_table is not enabled, trt would still only keep one copy of the table if the weights are identical - model = share_embedding(model) - - # After weight loading - if use_ootb_moe: - model = to_ootb_moe(model) - if use_fused_mlp: - model = fuse_gate_mlp(model, gemm_swiglu_plugin_dtype, - low_latency_gemm_swiglu_plugin_dtype) - if use_fused_rg_lru: - model = fuse_rg_lru(model) - if use_unfused_qkv_gemm: - model = unfuse_qkv_gemm(model) - if use_prompt_tuning: - model = set_prompt_tuning(model) - if use_lora: - model = add_lora(model, max_lora_rank, with_dora=use_dora) - if use_fp8_context_fmha: - model = set_fp8_context_fmha(model) - if fuse_fp4_quant: - model = set_fuse_fp4_quant(model) - if not use_lora and use_optimize_cross_qkv is True: - # This optimization is not supported when we use lora - model = optimize_cross_qkv(model) - - return model - - -def optimize_cross_qkv(model): - """ - For cross attention layer, we can skip computing the query of encoder_output. - So, add a new attribute 'kv' in the cross_attention layer. This might lead to - additional memory cost on model size, but save the memory usage on runtime. - - Currently, this function only detects ColumnLinear and FP8Linear. It does not support - other quantization now. - """ - for name, attn, layer in model.named_modules_with_parent(): - if isinstance(attn, Attention) and attn.cross_attention and \ - (type(attn.qkv) == ColumnLinear or type(attn.qkv) == FP8Linear): - old_qkv = attn.qkv - linear_class = type(old_qkv) - new_kv = linear_class( - in_features=attn.hidden_size, - out_features=2 * attn.tp_size * attn.num_attention_kv_heads * - attn.attention_head_size, - bias=old_qkv.bias, - dtype=old_qkv.dtype, - tp_group=old_qkv.tp_group, - tp_size=old_qkv.tp_size, - gather_output=old_qkv.gather_output, - prefer_managed_weight=old_qkv.prefer_managed_weight, - is_qkv=old_qkv.is_qkv, - ) - - old_qkv_weight_value = old_qkv.weight.raw_value - if (old_qkv_weight_value.shape == np.asarray([ - (attn.num_attention_heads + 2 * attn.num_attention_kv_heads) * - attn.attention_head_size, attn.hidden_size - ])).all(): - - q_weight, kv_weight = np.array_split( - old_qkv_weight_value.reshape( - attn.num_attention_heads + - 2 * attn.num_attention_kv_heads, - attn.attention_head_size, attn.hidden_size), - [attn.num_attention_heads], - axis=0) - new_kv.weight.value = kv_weight.reshape([ - 2 * attn.num_attention_kv_heads * attn.attention_head_size, - attn.hidden_size - ]) - elif (old_qkv_weight_value.shape == np.asarray([ - attn.hidden_size, - (attn.num_attention_heads + 2 * attn.num_attention_kv_heads) * - attn.attention_head_size - ])).all(): - q_weight, kv_weight = np.array_split( - old_qkv_weight_value.reshape( - attn.hidden_size, attn.num_attention_heads + - 2 * attn.num_attention_kv_heads, - attn.attention_head_size), [attn.num_attention_heads], - axis=1) - new_kv.weight.value = kv_weight.reshape([ - attn.hidden_size, - 2 * attn.num_attention_kv_heads * attn.attention_head_size - ]) - else: - assert False - - if isinstance(attn.qkv, FP8Linear): - new_kv.activation_scaling_factor.value = old_qkv.activation_scaling_factor.raw_value - new_kv.weights_scaling_factor.value = old_qkv.weights_scaling_factor.raw_value - - if old_qkv.bias: - q_bias, kv_bias = np.array_split(old_qkv.bias.raw_value.reshape( - attn.num_attention_heads + 2 * attn.num_attention_kv_heads, - attn.attention_head_size), [attn.num_attention_heads], - axis=0) - new_kv.bias.value = kv_bias.reshape([ - 2 * attn.num_attention_kv_heads * attn.attention_head_size - ]) - setattr(attn, "kv", new_kv) - - return model - - -def preprocess_perlayer_weights(weights, - model_config, - quant_algo, - from_pruned=False): - exclude_modules = model_config.quantization.exclude_modules - - # INT4_AWQ - if quant_algo == QuantAlgo.W4A8_AWQ or quant_algo == QuantAlgo.W4A16_AWQ: - preprocessor = preprocess_weights_for_mixed_gemm - if quant_algo == QuantAlgo.W4A8_AWQ: - activation_type = torch.float8_e4m3fn - elif quant_algo == QuantAlgo.W4A16_AWQ: - activation_type = torch.float16 - for name, param in weights.items(): - if from_pruned and param.numel() == 0: - continue - if name.endswith('weight') and param.dtype == torch.int8: - dtype = torch.float16 - if model_config.dtype == "bfloat16": - dtype = torch.bfloat16 - weights[name] = preprocessor(param.transpose(-1, -2), - torch.quint4x2, - activation_type).view(dtype) - if name.endswith('weights_scaling_factor'): - weights[name] = param.transpose(-1, -2).contiguous().to( - str_dtype_to_torch(model_config.dtype)) - if name.endswith('prequant_scaling_factor'): - if len(weights[name].shape) == 2: - # MoE experts share the same scaling factor. - param = param[0, :] - weights[name] = param.reshape(1, -1) - if model_config.mapping.tp_rank > 0: - if name.endswith('attention.dense.bias') or name.endswith( - 'mlp.proj.bias'): - weights[name] = torch.zeros_like(param) - - if quant_algo == QuantAlgo.W4A8_AWQ: - for name in list(weights): - if name.endswith('weights_scaling_factor'): - activation_scaling_factor = weights.pop( - name.replace('weights_scaling_factor', - 'activation_scaling_factor')) - weights_scaling_factor_2 = weights.pop( - name.replace('weights_scaling_factor', - 'weights_scaling_factor_2')) - weights[name] /= weights_scaling_factor_2 - weights[name] = weights[name].to(torch.float16).view( - str_dtype_to_torch(model_config.dtype)) - weights[name.replace( - 'weights_scaling_factor', - 'prequant_scaling_factor')] /= activation_scaling_factor - weights[name.replace( - 'weights_scaling_factor', 'alpha' - )] = activation_scaling_factor * weights_scaling_factor_2 - weights[name.replace('weights_scaling_factor', - 'activation_scaling_factor' - )] = activation_scaling_factor - - # FP8 - elif quant_algo == QuantAlgo.FP8: - for name, param in weights.items(): - if name.endswith('weight') and param.dtype == torch.int8: - weights[name] = param.view(torch.float8_e4m3fn) - # lm_head is not always quantized to FP8 - if "lm_head.weight" in weights and weights[ - 'lm_head.weight'].dtype is not torch.float8_e4m3fn: - weights.pop('lm_head.weights_scaling_factor', None) - weights.pop('lm_head.activation_scaling_factor', None) - elif quant_algo == QuantAlgo.FP8_PER_CHANNEL_PER_TOKEN: - for name, param in weights.items(): - if name.endswith('weight') and param.dtype == torch.int8: - weights[name] = param.view(torch.float8_e4m3fn) - # lm_head is not quantized to FP8 - if "lm_head.weight" in weights: - assert weights['lm_head.weight'].dtype == str_dtype_to_torch( - model_config.dtype) - weights.pop('lm_head.weights_scaling_factor', None) - weights.pop('lm_head.activation_scaling_factor', None) - # FP4 - elif quant_algo == QuantAlgo.NVFP4: - # Interleave block scale for NVFP4 plugin. - for name in list(weights): - if name.endswith('weights_scaling_factor'): - out_features, in_features = weights[name].shape - nrows = fp4_utils.pad_up(out_features, 128) - ncols = fp4_utils.pad_up(in_features, 4) - new_name = name.replace('weights_scaling_factor', - 'weights_block_scaling_factor') - weights[new_name] = weights[name] - weights[ - new_name + - "_interleaved"] = torch.ops.trtllm.block_scale_interleave( - weights[name].view(fp4_utils.float4_sf_dtype).cpu( - ).contiguous()).reshape(nrows, ncols).view( - fp4_utils.float4_sf_dtype) - weights.pop(name) - if name.endswith('weights_scaling_factor_2'): - new_name = name.replace('weights_scaling_factor_2', - 'weights_global_scaling_factor') - weights[new_name] = weights[name] - weights.pop(name) - if name.endswith('activation_scaling_factor'): - new_name = name.replace('activation_scaling_factor', - 'activation_global_scaling_factor') - weights[new_name] = weights[name] - weights.pop(name) - for name in list(weights): - if name.endswith('weights_global_scaling_factor'): - weight_global_sf = weights[name] - act_global_sf = weights[name.replace( - 'weights_global_scaling_factor', - 'activation_global_scaling_factor')] - weights[name.replace( - 'weights_global_scaling_factor', - 'alpha')] = act_global_sf * weight_global_sf - elif quant_algo in [QuantAlgo.W4A16, QuantAlgo.W8A16]: - weights = weight_only_quantize_dict(weights=weights, - quant_algo=quant_algo, - exclude_modules=exclude_modules, - plugin=True) - - -def preprocess_weights(weights: Dict[str, torch.Tensor], - model_config: PretrainedConfig, - from_pruned=False) -> None: - """This function in-place modifies weights and model_config, making them compatible with each other. - - Note: Typically, it should be called before model creation and weight loading. For example, - preprocess_weights(weights, model_config) - model = XXXForCausalLM(model_config) - model.load(weights) - """ - quant_config = model_config.quantization - quant_algo = quant_config.quant_algo - - pattern_info = ['fc', 'gate', 'proj', 'qkv', 'dense'] - - def process_kv_scaling_factor(weights: Dict[str, torch.Tensor]): - new_entries = {} - names_to_delete = set() - - # If k, v cache scaling factors are stored separately, combine them into kv cache scaling factor. - for name, param in weights.items(): - if name.endswith('.k_cache_scaling_factor'): - v_name = name.replace('k_cache_scaling_factor', - 'v_cache_scaling_factor') - assert v_name in weights, f"{v_name} not found" - kv_name = name.replace('k_cache_scaling_factor', - 'kv_cache_scaling_factor') - new_entries[kv_name] = torch.max(weights[name], weights[v_name]) - names_to_delete.update([name, v_name]) - weights.update(new_entries) - for k in names_to_delete: - del weights[k] - - new_entries = [] - # The unified converter generate_tllm_weights() already generates these rcp weights, but legacy - # converters do not. Handle it here. - for name, param in weights.items(): - if name.endswith('.kv_cache_scaling_factor'): - rcp_name = name.replace('kv_cache_scaling_factor', - 'kv_cache_rcp_scaling_factor') - if rcp_name not in weights: - new_entries.append((rcp_name, torch.reciprocal(param))) - weights.update(new_entries) - - process_kv_scaling_factor(weights) - - per_layer_weights = {} - - for name, param in weights.items(): - in_mode = False - for info in pattern_info: - pattern = rf'(.*?{info}.*?)' - pattern_match = re.match(pattern, name) - if pattern_match: - base_name = pattern_match.group(1) - if base_name not in per_layer_weights.keys(): - per_layer_weights[base_name] = {} - per_layer_weights[base_name][name] = param - in_mode = True - break - if not in_mode: - # [lm_head.weight, ln_f.weight, vocab_embedding.weight] - base_name = name.rsplit('.', 1)[0] - if base_name not in per_layer_weights.keys(): - per_layer_weights[base_name] = {} - per_layer_weights[base_name][name] = param - - new_weights = {} - for base_name, layer_weights in per_layer_weights.items(): - if quant_algo != QuantAlgo.MIXED_PRECISION: - layer_quant_algo = quant_algo - else: - quant_cfg = quant_config._get_quant_cfg(base_name) - if not quant_cfg.quant_algo: - new_weights.update(layer_weights) - continue - - layer_quant_algo = quant_cfg.quant_algo - - preprocess_perlayer_weights(layer_weights, model_config, - layer_quant_algo, from_pruned) - new_weights.update(layer_weights) - - weights = new_weights - for name, param in weights.items(): - if model_config.architecture == 'GPTJForCausalLM': - if model_config.mapping.tp_rank > 0: - if 'attention.dense.bias' in name or 'mlp.proj.bias' in name: - weights[name] = torch.zeros_like(param) - - return weights - - -def get_kv_cache_type_from_legacy(use_cache: bool, - paged_kv_cache: bool) -> KVCacheType: - if use_cache: - if paged_kv_cache: - return KVCacheType.PAGED - else: - return KVCacheType.CONTINUOUS - else: - return KVCacheType.DISABLED - - -def save_config(config: PretrainedConfig, *, output_dir: str, - log: bool) -> None: - config_path = Path(output_dir) / "config.json" - if log: - logger.debug(f"Saving TensorRT LLM configuration to {config_path}") - config_path.parent.mkdir(exist_ok=True, parents=True) - config_path.write_text(json.dumps(config.to_dict(), indent=4)) - - -def save_checkpoint(*, output_dir: str, weights: dict, rank: int) -> None: - """ Checkpoint saver for weight loader.""" - safetensors.torch.save_file( - weights, os.path.join(output_dir, f'rank{rank}.safetensors')) diff --git a/tensorrt_llm/models/mpt/__init__.py b/tensorrt_llm/models/mpt/__init__.py deleted file mode 100644 index 71bf6d298d42..000000000000 --- a/tensorrt_llm/models/mpt/__init__.py +++ /dev/null @@ -1,14 +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. diff --git a/tensorrt_llm/models/mpt/model.py b/tensorrt_llm/models/mpt/model.py deleted file mode 100644 index ecfa343e363f..000000000000 --- a/tensorrt_llm/models/mpt/model.py +++ /dev/null @@ -1,176 +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. - -from ..._utils import pad_vocab_size -from ...functional import PositionEmbeddingType, Tensor -from ...layers import (MLP, Attention, AttentionMaskType, ColumnLinear, - Embedding, LayerNorm) -from ...module import Module -from ..modeling_utils import (DecoderLayerList, DecoderModelForCausalLM, - PretrainedConfig) - - -class MPTDecoderLayer(Module): - - def __init__(self, config: PretrainedConfig, layer_idx: int): - super().__init__() - self.layer_idx = layer_idx - self.config = config - - hidden_size = config.hidden_size - dtype = config.dtype - tp_size = config.mapping.tp_size - tp_rank = config.mapping.tp_rank - tp_group = config.mapping.tp_group - layernorm_epsilon = config.norm_epsilon - - self.input_layernorm = LayerNorm(normalized_shape=hidden_size, - eps=layernorm_epsilon, - bias=False, - dtype=dtype) - - layers_range = config.mapping.pp_layers(config.num_hidden_layers) - local_layer_idx = layer_idx - layers_range[0] - self.attention = Attention( - local_layer_idx=local_layer_idx, - hidden_size=hidden_size, - num_attention_heads=config.num_attention_heads, - num_kv_heads=config.num_key_value_heads, - attention_mask_type=AttentionMaskType.causal, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - tp_rank=tp_rank, - bias=config.bias, - position_embedding_type=PositionEmbeddingType.alibi, - quant_mode=config.quant_mode, - clip_qkv=config.clip_qkv, - alibi_bias_max=config.alibi_bias_max) - - self.mlp = MLP(hidden_size=hidden_size, - ffn_hidden_size=hidden_size * 4, - hidden_act=config.hidden_act, - dtype=dtype, - bias=config.bias, - tp_group=tp_group, - tp_size=tp_size, - quant_mode=config.quant_mode) - - self.post_layernorm = LayerNorm(normalized_shape=hidden_size, - eps=layernorm_epsilon, - bias=False, - dtype=dtype) - - def forward(self, - hidden_states: Tensor, - attention_mask=None, - use_cache=False, - kv_cache_params=None, - attention_params=None): - - assert isinstance(hidden_states, Tensor) - - residual = hidden_states - - hidden_states = self.input_layernorm(hidden_states) - - attention_output = self.attention(hidden_states, - attention_mask=attention_mask, - use_cache=use_cache, - kv_cache_params=kv_cache_params, - attention_params=attention_params) - - if use_cache: - attention_output, presents = attention_output - - hidden_states = residual + attention_output - - residual = hidden_states - hidden_states = self.post_layernorm(hidden_states) - - hidden_states = self.mlp(hidden_states) - - hidden_states = residual + hidden_states - - if use_cache: - return (hidden_states, presents) - return hidden_states - - -class MPTModel(Module): - - def __init__(self, config: PretrainedConfig): - super().__init__() - self.config = config - - if config.mapping.is_first_pp_rank(): - self.vocab_embedding = Embedding(config.vocab_size, - config.hidden_size, - dtype=config.dtype) - self.layers = DecoderLayerList(MPTDecoderLayer, config) - if config.mapping.is_last_pp_rank(): - self.ln_f = LayerNorm(normalized_shape=config.hidden_size, - bias=False, - dtype=config.dtype) - - def forward(self, - input_ids, - position_ids, - use_cache=False, - attention_mask=None, - kv_cache_params=None, - attention_params=None): - - hidden_states = self.vocab_embedding(input_ids) - - hidden_states = self.layers(hidden_states, - use_cache=use_cache, - attention_mask=attention_mask, - kv_cache_params=kv_cache_params, - attention_params=attention_params) - - if use_cache: - hidden_states, presents = hidden_states - - hidden_states = self.ln_f(hidden_states) - - if use_cache: - return (hidden_states, tuple(presents)) - return hidden_states - - -class MPTForCausalLM(DecoderModelForCausalLM): - - def __init__(self, config: PretrainedConfig): - self.check_config(config) - transformer = MPTModel(config) - vocab_size_padded = pad_vocab_size(config.vocab_size, - config.mapping.tp_size) - if config.mapping.is_last_pp_rank(): - lm_head = ColumnLinear(config.hidden_size, - vocab_size_padded, - bias=config.bias, - dtype=config.dtype, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size, - gather_output=True) - else: - lm_head = None - super().__init__(config, transformer, lm_head) - - def check_config(self, config): - config.set_if_not_exist('bias', False) - config.set_if_not_exist('clip_qkv', None) - config.set_if_not_exist('alibi_bias_max', 8) diff --git a/tensorrt_llm/models/multimodal_encoders/__init__.py b/tensorrt_llm/models/multimodal_encoders/__init__.py deleted file mode 100644 index 71bf6d298d42..000000000000 --- a/tensorrt_llm/models/multimodal_encoders/__init__.py +++ /dev/null @@ -1,14 +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. diff --git a/tensorrt_llm/models/multimodal_encoders/config.py b/tensorrt_llm/models/multimodal_encoders/config.py deleted file mode 100644 index 2ae8135caae2..000000000000 --- a/tensorrt_llm/models/multimodal_encoders/config.py +++ /dev/null @@ -1,116 +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. -from typing import Optional, Union - -import torch - -from ..._utils import torch_dtype_to_str -from ...logger import logger -from ...mapping import Mapping -from ..modeling_utils import PretrainedConfig, QuantConfig - - -class LlavaNextVisionConfig(PretrainedConfig): - - def __init__(self, - *, - image_size: int, - patch_size: int, - text_hidden_size: int, - projector_hidden_act: str = 'gelu', - num_channels: int = 3, - vision_model_type: str = 'clip_vision_model', - **kwargs): - self.image_size = image_size - self.patch_size = patch_size - self.text_hidden_size = text_hidden_size - self.num_channels = num_channels - self.projector_hidden_act = projector_hidden_act - self.vision_model_type = vision_model_type - - super().__init__(**kwargs) - - @classmethod - def from_hugging_face( - cls, - hf_config_or_dir: Union[str, 'transformers.PretrainedConfig'], - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - import transformers - - if isinstance(hf_config_or_dir, transformers.PretrainedConfig): - hf_config = hf_config_or_dir - else: - hf_config_dir = str(hf_config_or_dir) - - hf_config = transformers.AutoConfig.from_pretrained( - hf_config_dir, trust_remote_code=True) - if hf_config.model_type == "llava_next": - from transformers import LlavaNextConfig - hf_config = LlavaNextConfig.from_pretrained(hf_config_dir) - else: - logger.error("Provided model type is not llava_next.") - - text_hidden_size = hf_config.text_config.hidden_size - # Extract only the vision config - llava_next_vision_config = hf_config.vision_config - - # llava-next uses the second last layer as vision output - num_feature_layers = llava_next_vision_config.num_hidden_layers + hf_config.vision_feature_layer + 1 - - vision_model_type = getattr(llava_next_vision_config, - "vision_model_type", "clip_vision_model") - - num_key_value_heads = getattr( - llava_next_vision_config, "num_key_value_heads", - llava_next_vision_config.num_attention_heads) - - # Default configs from HF - hidden_act = 'quick_gelu' - norm_epsilon = 1e-5 - - head_size = llava_next_vision_config.hidden_size // llava_next_vision_config.num_attention_heads - - if dtype == 'auto': - dtype = getattr(hf_config, 'torch_dtype', None) - if dtype is None: - dtype = 'float16' - if isinstance(dtype, torch.dtype): - dtype = torch_dtype_to_str(dtype) - if dtype == 'float32': - dtype = 'float16' - - return cls( - image_size=llava_next_vision_config.image_size, - patch_size=llava_next_vision_config.patch_size, - text_hidden_size=text_hidden_size, - projector_hidden_act=hf_config.projector_hidden_act, - vision_model_type=vision_model_type, - architecture=hf_config.architectures[0], - dtype=dtype, - num_hidden_layers=num_feature_layers, - num_attention_heads=llava_next_vision_config.num_attention_heads, - hidden_size=llava_next_vision_config.hidden_size, - intermediate_size=llava_next_vision_config.intermediate_size, - num_key_value_heads=num_key_value_heads, - head_size=head_size, - vocab_size=llava_next_vision_config.vocab_size, - hidden_act=hidden_act, - norm_epsilon=norm_epsilon, - mapping=mapping, - quantization=quant_config, - **kwargs) diff --git a/tensorrt_llm/models/multimodal_encoders/model.py b/tensorrt_llm/models/multimodal_encoders/model.py deleted file mode 100644 index b60337b0f740..000000000000 --- a/tensorrt_llm/models/multimodal_encoders/model.py +++ /dev/null @@ -1,175 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -import os -from collections import OrderedDict -from typing import Optional - -import safetensors - -from ..._utils import numpy_to_torch -from ...functional import ACT2FN, Tensor, concat, shape, slice -from ...layers import Linear -from ...logger import logger -from ...mapping import Mapping -from ...models import CLIPVisionTransformer -from ...module import Module -from ...parameter import Parameter -from ..model_weights_loader import ModelWeightsLoader -from ..modeling_utils import PretrainedModel, QuantConfig -from .config import LlavaNextVisionConfig - - -# Adapted from https://github.com/huggingface/transformers/blob/v4.39.0/src/transformers/models/llava_next/modeling_llava_next.py#L149 -class LlavaNextMultiModalProjector(Module): - - def __init__(self, config: LlavaNextVisionConfig): - super().__init__() - - self.linear_1 = Linear(config.hidden_size, - config.text_hidden_size, - dtype=config.dtype) - self.act = ACT2FN[config.projector_hidden_act] - self.linear_2 = Linear(config.text_hidden_size, - config.text_hidden_size, - dtype=config.dtype) - - def forward(self, image_features): - hidden_states = self.linear_1(image_features) - hidden_states = self.act(hidden_states) - hidden_states = self.linear_2(hidden_states) - return hidden_states - - -class LlavaNextVisionWrapper(PretrainedModel): - - def __init__(self, config: LlavaNextVisionConfig): - super().__init__(config) - self.vision_tower = None - self.config = config - if config.vision_model_type == "clip_vision_model": - self.vision_tower = CLIPVisionTransformer( - image_size=config.image_size, - num_channels=config.num_channels, - patch_size=config.patch_size, - hidden_size=config.hidden_size, - num_attention_heads=config.num_attention_heads, - max_position_embeddings=config.max_position_embeddings, - norm_epsilon=config.norm_epsilon, - intermediate_size=config.intermediate_size, - hidden_act=config.hidden_act, - num_hidden_layers=config.num_hidden_layers, - require_ln_f=False, - mapping=config.mapping, - dtype=config.dtype) - else: - logger.error( - "Currently TRT-LLM only supports CLIP vision transformer.") - - self.multi_modal_projector = LlavaNextMultiModalProjector(config) - self.image_newline = Parameter(shape=(config.text_hidden_size, ), - dtype=config.dtype) - - def forward(self, pixel_values, position_ids=None): - image_features = self.vision_tower(pixel_values) - select_size = concat([ - shape(image_features, 0), image_features.shape[1] - 1, - shape(image_features, 2) - ]) - selected_image_feature = slice(image_features, - starts=[0, 1, 0], - sizes=select_size) # (bs, 576, c) - image_features = self.multi_modal_projector(selected_image_feature) - image_features.mark_output('image_features', self.config.dtype) - return image_features # (bs, 576, c) - - @classmethod - def from_hugging_face(cls, - hf_model_dir: str, - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - ''' Create a LlavaNextVisionWrapper object from give parameters - ''' - if os.environ.get("TRTLLM_DISABLE_UNIFIED_CONVERTER") is not None: - logger.error( - "Please enable unified converter to convert llava-next checkpoints." - ) - - config = LlavaNextVisionConfig.from_hugging_face( - hf_model_dir, - dtype=dtype, - mapping=mapping, - quant_config=quant_config, - **kwargs) - - custom_dict = {} - if "llava" in hf_model_dir: - custom_dict = { - "vision_tower": "vision_tower.vision_model", - "input_layernorm": "layer_norm1", - "post_layernorm": "layer_norm2", - "fc": "fc1", - "proj": "fc2", - "dense": "out_proj", - "pre_layernorm": "pre_layrnorm", - "ln_f": "post_layernorm", - } - loader = ModelWeightsLoader(hf_model_dir, custom_dict) - model = cls(config) - loader.generate_tllm_weights(model) - return model - - def save_checkpoint(self, output_dir, save_config=True): - rank = self.config.mapping.rank - weights = { - name: numpy_to_torch(param.raw_value) - for name, param in self.named_parameters() - } - image_newline = { - "image_newline": numpy_to_torch(self.image_newline.raw_value) - } - safetensors.torch.save_file( - weights, os.path.join(output_dir, f'rank{rank}.safetensors')) - safetensors.torch.save_file( - image_newline, - os.path.join(output_dir, f'image_newlines.safetensors')) - if save_config: - self.config.to_json_file(os.path.join(output_dir, 'config.json')) - - def prepare_inputs(self, max_batch_size, **kwargs): - '''@brief: Prepare inputs Tensors for the model, the given sizes are used to determine the - ranges of the dimensions of when using TRT dynamic shapes. - - @return: a list contains values which can be fed into the self.forward() - ''' - - batch_size_range = [ - 1, max(1, (max_batch_size + 1) // 2), max_batch_size - ] - pixel_values = Tensor( - name='pixel_values', - dtype=self.config.dtype, - shape=[ - -1, self.config.num_channels, self.config.image_size, - self.config.image_size - ], - dim_range=OrderedDict([ - ('batch_size', [batch_size_range]), - ('in_channels', [[self.config.num_channels] * 3]), - ('latent_height', [[self.config.image_size] * 3]), - ('latent_width', [[self.config.image_size] * 3]), - ])) - return {'pixel_values': pixel_values} diff --git a/tensorrt_llm/models/nemotron_nas/__init__.py b/tensorrt_llm/models/nemotron_nas/__init__.py deleted file mode 100644 index 71bf6d298d42..000000000000 --- a/tensorrt_llm/models/nemotron_nas/__init__.py +++ /dev/null @@ -1,14 +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. diff --git a/tensorrt_llm/models/nemotron_nas/config.py b/tensorrt_llm/models/nemotron_nas/config.py deleted file mode 100644 index 11d02df84b07..000000000000 --- a/tensorrt_llm/models/nemotron_nas/config.py +++ /dev/null @@ -1,208 +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. -from dataclasses import asdict -from typing import Any, Dict, List, Optional, Union - -from tensorrt_llm._utils import get_hf_rope_theta -from tensorrt_llm.functional import PositionEmbeddingType -from tensorrt_llm.mapping import Mapping -from tensorrt_llm.models.convert_utils import infer_dtype -from tensorrt_llm.models.modeling_utils import PretrainedConfig, QuantConfig -from tensorrt_llm.models.nemotron_nas.convert import \ - hf_block_configs_to_layer_configs -from tensorrt_llm.models.nemotron_nas.layer_config import ( - AttentionConfig, AttentionImplementation, DeciLayerConfig, FFNConfig) - - -class DeciConfig(PretrainedConfig): - - def __init__(self, - *, - architecture: str = 'DeciLMForCausalLM', - dtype: str, - hidden_size: int, - num_hidden_layers: int, - num_attention_heads: int, - vocab_size: int, - hidden_act: str = 'gelu', - logits_dtype: str = 'float32', - norm_epsilon: float = 0.00001, - position_embedding_type: Union[ - PositionEmbeddingType, - str] = PositionEmbeddingType.rope_gpt_neox, - rotary_base: float = 10000.0, - rotary_scaling: Optional[dict] = None, - max_position_embeddings: int, - num_key_value_heads: Optional[int] = None, - intermediate_size: Optional[int] = None, - mapping: Optional[Union[Mapping, dict]] = None, - quantization: Optional[Union[QuantConfig, dict]] = None, - use_parallel_embedding: bool = False, - embedding_sharding_dim: int = 0, - head_size: Optional[int] = None, - qk_layernorm: bool = False, - layer_configs: Optional[List[Union[DeciLayerConfig, - Dict[str, - Dict[str, - Any]]]]] = None, - block_configs: Optional[object] = None, - **kwargs): - super().__init__(architecture=architecture, - dtype=dtype, - hidden_size=hidden_size, - num_hidden_layers=num_hidden_layers, - num_attention_heads=num_attention_heads, - vocab_size=vocab_size, - hidden_act=hidden_act, - logits_dtype=logits_dtype, - norm_epsilon=norm_epsilon, - position_embedding_type=position_embedding_type, - max_position_embeddings=max_position_embeddings, - num_key_value_heads=num_key_value_heads, - intermediate_size=intermediate_size, - mapping=mapping, - quantization=quantization, - use_parallel_embedding=use_parallel_embedding, - embedding_sharding_dim=embedding_sharding_dim, - head_size=head_size, - qk_layernorm=qk_layernorm, - **kwargs) - - self.rotary_base = rotary_base - self.rotary_scaling = rotary_scaling - - if block_configs is not None: - assert layer_configs is None - self.layer_configs = hf_block_configs_to_layer_configs( - block_configs, - num_attention_heads=num_attention_heads, - hidden_size=hidden_size) - elif layer_configs is not None: - assert len( - layer_configs - ) == num_hidden_layers, f"num_hidden_layers ({num_hidden_layers}) must match len(layer_configs) ({len(layer_configs)})" - - self.layer_configs = self._ensure_layer_configs(layer_configs) - else: - self.layer_configs = None - - # HACK: this is needed for many parts of the code - self.layer_types = [ - AttentionImplementation( - self.get_layer_config(layer_idx).attention.impl).value - for layer_idx in range(self.num_hidden_layers) - ] - - # HACK: this is here since the runtime doesn't parse the layer_configs yet - self.num_kv_heads_per_layer = [] - for layer_idx in range(self.num_hidden_layers): - layer_config = self.get_layer_config(layer_idx) - if layer_config.is_attention_layer: - self.num_kv_heads_per_layer.append( - layer_config.attention.num_key_value_heads) - - def _ensure_layer_configs( - self, layer_configs: List[Union[DeciLayerConfig, Dict[str, Any]]] - ) -> List[DeciLayerConfig]: - return [ - DeciLayerConfig.from_dict(c) if isinstance(c, dict) else c - for c in layer_configs - ] - - def to_dict(self): - output = super().to_dict() - if self.layer_configs is not None: - output["layer_configs"] = [asdict(c) for c in self.layer_configs] - return output - - def get_layer_config(self, layer_idx: int) -> DeciLayerConfig: - if self.layer_configs is not None: - conf = self.layer_configs[layer_idx] - else: - conf = DeciLayerConfig() - - attention_impl = conf.attention.impl - num_key_value_heads = conf.attention.num_key_value_heads or self.num_key_value_heads - ffn_impl = conf.ffn.impl - intermediate_size = conf.ffn.intermediate_size or self.intermediate_size - - return DeciLayerConfig( - attention=AttentionConfig(impl=attention_impl, - num_key_value_heads=num_key_value_heads), - ffn=FFNConfig(impl=ffn_impl, intermediate_size=intermediate_size)) - - def get_layer_num_kv_heads(self, layer_idx) -> int: - layer_config = self.get_layer_config(layer_idx) - assert layer_config.is_attention_layer, f"Layer {layer_idx} is not an attention layer" - return layer_config.attention.num_key_value_heads or self.num_key_value_heads - - @classmethod - def from_hugging_face( - cls, - hf_config_or_dir: Union[str, 'transformers.PretrainedConfig'], - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - trust_remote_code: bool = True, - **kwargs): - import transformers - - if isinstance(hf_config_or_dir, transformers.PretrainedConfig): - hf_config = hf_config_or_dir - else: - hf_config = transformers.AutoConfig.from_pretrained( - hf_config_or_dir, trust_remote_code=trust_remote_code) - - assert hf_config.model_type in ( - "deci", - "nemotron-nas"), f"Unsupported model type: {hf_config.model_type}" - - block_configs = getattr(hf_config, "block_configs", None) - if block_configs is not None: - layer_configs = hf_block_configs_to_layer_configs( - block_configs, - num_attention_heads=hf_config.num_attention_heads, - hidden_size=hf_config.hidden_size) - else: - # older deci arch - num_key_value_heads_per_layer = getattr( - hf_config, "num_key_value_heads_per_layer", None) - if num_key_value_heads_per_layer is not None: - layer_configs = [ - DeciLayerConfig(attention=AttentionConfig( - num_key_value_heads=num_key_value_heads)) - for num_key_value_heads in num_key_value_heads_per_layer - ] - else: - layer_configs = None - - dtype = infer_dtype(dtype, getattr(hf_config, 'torch_dtype', None)) - - return cls(dtype=dtype, - hidden_size=hf_config.hidden_size, - hidden_act=hf_config.hidden_act, - intermediate_size=hf_config.intermediate_size, - num_attention_heads=hf_config.num_attention_heads, - num_hidden_layers=hf_config.num_hidden_layers, - num_key_value_heads=hf_config.num_key_value_heads, - norm_epsilon=hf_config.rms_norm_eps, - rotary_scaling=hf_config.rope_scaling, - rotary_base=get_hf_rope_theta(hf_config, 10000.0), - vocab_size=hf_config.vocab_size, - max_position_embeddings=hf_config.max_position_embeddings, - mapping=mapping, - quantization=quant_config, - layer_configs=layer_configs, - **kwargs) diff --git a/tensorrt_llm/models/nemotron_nas/convert.py b/tensorrt_llm/models/nemotron_nas/convert.py deleted file mode 100644 index ba26414ae8d4..000000000000 --- a/tensorrt_llm/models/nemotron_nas/convert.py +++ /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. -import enum -import json -import time -from abc import ABC, abstractmethod -from contextlib import contextmanager -from dataclasses import asdict -from pathlib import Path -from typing import Any, Dict, Iterator, List, Optional, TypedDict, Union - -import safetensors -import torch - -from tensorrt_llm._utils import pad_vocab_size -from tensorrt_llm.logger import logger -from tensorrt_llm.models.convert_utils import dup_kv_weight, split -from tensorrt_llm.models.nemotron_nas.layer_config import ( - AttentionConfig, AttentionImplementation, DeciLayerConfig, FFNConfig, - FFNImplementation) -from tensorrt_llm.quantization.mode import QuantAlgo - - -def _ffn_mult_to_intermediate_size(ffn_mult: float, n_embd: int) -> int: - intermediate_size = int(2 * ffn_mult * n_embd / 3) - return _find_multiple(intermediate_size, 256) - - -def _find_multiple(n: int, k: int) -> int: - if n % k == 0: - return n - return n + k - (n % k) - - -# BlockConfig is a custom class defined inside deci huggingface checkpoints, we can't import it -def hf_block_config_to_layer_config(block_config: Union["BlockConfig", dict], - num_attn_heads: int, - hidden_size: int) -> DeciLayerConfig: - """`block_config` (`Union[BlockConfig, dict]`): A `dict` when exported from `ModelOpt`; A `dataclass` at the HF phase - """ - block_config = block_config if isinstance(block_config, - dict) else asdict(block_config) - attn = block_config["attention"] - if attn["no_op"]: - attn_impl = AttentionImplementation.NO_OP - num_key_value_heads = None - elif attn["replace_with_linear"]: - attn_impl = AttentionImplementation.LINEAR - num_key_value_heads = None - elif attn.get("sparsify", None): - raise NotImplementedError("Sparsification is not supported") - else: - attn_impl = AttentionImplementation.ATTENTION - num_key_value_heads = num_attn_heads // attn["n_heads_in_group"] - - ffn = block_config["ffn"] - if ffn["no_op"]: - ffn_impl = FFNImplementation.NO_OP - intermediate_size = None - elif ffn["replace_with_linear"]: - ffn_impl = FFNImplementation.LINEAR - intermediate_size = None - elif ffn.get("sparsify", None): - raise NotImplementedError("Sparsification is not supported") - else: - ffn_impl = FFNImplementation.MLP - intermediate_size = _ffn_mult_to_intermediate_size( - ffn["ffn_mult"], hidden_size) - - return DeciLayerConfig(attention=AttentionConfig( - impl=attn_impl, num_key_value_heads=num_key_value_heads), - ffn=FFNConfig(impl=ffn_impl, - intermediate_size=intermediate_size)) - - -def hf_block_configs_to_layer_configs( - block_configs: Union["BlockConfig", dict], *, num_attention_heads: int, - hidden_size: int) -> List[DeciLayerConfig]: - return [ - hf_block_config_to_layer_config(block_config, num_attention_heads, - hidden_size) - for block_config in block_configs - ] - - -@contextmanager -def timed_loading() -> Iterator[None]: - tik = time.time() - yield - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - logger.info(f'Weights loaded. Total time: {t}') - - -class TpDim(enum.IntEnum): - NO_TP = -1 - COLWISE = 0 - ROWWISE = 1 - - -class SafetensorsIndex(TypedDict): - metadata: Dict[str, Any] - weight_map: Dict[str, str] - - -class WeightsLoader(ABC): - - @abstractmethod - def read_weight(self, name: str) -> torch.Tensor: - ... - - def get_weight(self, - name: str, - tp_dim: TpDim = TpDim.NO_TP, - tp_size: int = 1, - tp_rank: int = 0) -> torch.Tensor: - weight = self.read_weight(name) - if tp_dim != TpDim.NO_TP: - weight = split(weight, tp_size, tp_rank, dim=tp_dim) - return weight - - def get_kv_weight(self, - name: str, - num_heads: int, - tp_size: int = 1, - tp_rank: int = 0) -> torch.Tensor: - weight = self.read_weight(name) - if tp_size > num_heads: - weight = dup_kv_weight(weight, num_heads, tp_size) - if tp_size > 1: - weight = split(weight, tp_size, tp_rank, dim=0) - - return weight - - -class HFModelWeightsLoader(WeightsLoader): - - def __init__(self, *, hf_model: "transformers.PreTrainedModel", - dtype: str) -> None: - self.model_params = dict(hf_model.named_parameters()) - self.dtype = getattr(torch, dtype) - - def read_weight(self, name: str) -> torch.Tensor: - weight = self.model_params[name] - if weight.dtype != self.dtype: - weight = weight.to(self.dtype) - weight = weight.detach() - return weight - - -class SafetensorsWeightsLoader(WeightsLoader): - - def __init__(self, *, model_dir: Path, dtype: str) -> None: - self.model_dir = model_dir - self.dtype = getattr(torch, dtype) - - # the index has a weight map that maps weight names to the files they are found in - safetensor_index_json = self.model_dir / "model.safetensors.index.json" - has_safetensor_index_json = safetensor_index_json.is_file() - if has_safetensor_index_json: - with safetensor_index_json.open("r") as fr: - self.sharding_map: SafetensorsIndex = json.load(fr) - else: - self.sharding_map = SafetensorsIndex(metadata={}, weight_map={}) - - shard_files = {f.name for f in self.model_dir.glob("*.safetensors")} - if has_safetensor_index_json: - # only read the files that have weights according to the index - shard_files &= set(self.sharding_map["weight_map"].values()) - self.shard_files = sorted(list(shard_files)) - - self.safetensors_files = { - shard_file: - safetensors.safe_open(model_dir / shard_file, - framework="pt", - device="cpu") - for shard_file in shard_files - } - - def read_weight(self, name: str) -> torch.Tensor: - shard_filename = self.sharding_map['weight_map'].get( - name, self.shard_files[0]) - res = self.safetensors_files[shard_filename].get_tensor(name) - return res.to(self.dtype).contiguous() - - -def load_model_weights(loader: WeightsLoader, - config: "DeciConfig") -> Dict[str, torch.Tensor]: - mapping = config.mapping - num_hidden_layers = config.num_hidden_layers - vocab_size = config.vocab_size - pad_vocab = vocab_size % mapping.tp_size != 0 - vocab_size_padded = pad_vocab_size(vocab_size, mapping.tp_size) - - weights = {} - - def load_weight(name: str, tp_dim: TpDim = TpDim.NO_TP) -> torch.Tensor: - return loader.get_weight(name=name, - tp_dim=tp_dim, - tp_rank=mapping.tp_rank, - tp_size=mapping.tp_size) - - with timed_loading(): - if mapping.is_first_pp_rank(): - weights['transformer.vocab_embedding.weight'] = load_weight( - "model.embed_tokens.weight", - TpDim(config.embedding_sharding_dim) - if config.use_parallel_embedding else - TpDim.NO_TP) # vocab_embedding - - if mapping.is_last_pp_rank(): - v = load_weight("lm_head.weight", - TpDim.NO_TP) if pad_vocab else load_weight( - "lm_head.weight", TpDim.COLWISE) # lm_head - if pad_vocab: - v = torch.nn.functional.pad( - v, (0, 0, 0, vocab_size_padded - vocab_size), 'constant', 0) - v = split(v, mapping.tp_size, mapping.tp_rank) - weights['lm_head.weight'] = v - weights['transformer.ln_f.weight'] = load_weight( - "model.norm.weight") # ln_f - - layers_range = mapping.pp_layers(num_hidden_layers) - for l in layers_range: - layer_config = config.get_layer_config(l) - layer_idx = l - layers_range[0] - tllm_prex = f'transformer.layers.{layer_idx}' - - # Attention - if layer_config.is_attention_layer: - weights[f'{tllm_prex}.input_layernorm.weight'] = load_weight( - f"model.layers.{l}.input_layernorm.weight" - ) # input_layernorm - - q = load_weight(f"model.layers.{l}.self_attn.q_proj.weight", - TpDim.COLWISE) - k = loader.get_kv_weight( - f"model.layers.{l}.self_attn.k_proj.weight", - num_heads=layer_config.attention.num_key_value_heads, - tp_size=mapping.tp_size, - tp_rank=mapping.tp_rank) - v = loader.get_kv_weight( - f"model.layers.{l}.self_attn.v_proj.weight", - num_heads=layer_config.attention.num_key_value_heads, - tp_size=mapping.tp_size, - tp_rank=mapping.tp_rank) - weights[f'{tllm_prex}.attention.qkv.weight'] = torch.cat( - [q, k, v], 0) - weights[f'{tllm_prex}.attention.dense.weight'] = load_weight( - f"model.layers.{l}.self_attn.o_proj.weight", - TpDim.ROWWISE) # attention.dense - - elif layer_config.is_linear_attention_layer: - weights[f'{tllm_prex}.input_layernorm.weight'] = load_weight( - f"model.layers.{l}.input_layernorm.weight" - ) # input_layernorm - - weights[f'{tllm_prex}.attention.weight'] = load_weight( - f"model.layers.{l}.self_attn.linear_attn.weight", - TpDim.COLWISE) - - elif not layer_config.is_noop_attention_layer: - raise NotImplementedError( - f"Loading weights for layer with attention of type {layer_config.attention.impl} is not supported" - ) - - # MLP - if layer_config.is_mlp_layer: - weights[f'{tllm_prex}.post_layernorm.weight'] = load_weight( - f"model.layers.{l}.post_attention_layernorm.weight" - ) # post_layernorm - - weights[f'{tllm_prex}.ffn.gate.weight'] = load_weight( - f"model.layers.{l}.mlp.up_proj.weight", - TpDim.COLWISE) # mlp.gate - weights[f'{tllm_prex}.ffn.proj.weight'] = load_weight( - f"model.layers.{l}.mlp.down_proj.weight", - TpDim.ROWWISE) # mlp.proj - weights[f'{tllm_prex}.ffn.fc.weight'] = load_weight( - f"model.layers.{l}.mlp.gate_proj.weight", - TpDim.COLWISE) # mlp.fc - - elif layer_config.is_linear_ffn_layer: - weights[f'{tllm_prex}.post_layernorm.weight'] = load_weight( - f"model.layers.{l}.post_attention_layernorm.weight" - ) # post_layernorm - - weights[f'{tllm_prex}.ffn.weight'] = load_weight( - f"model.layers.{l}.mlp.linear_mlp.weight", TpDim.COLWISE) - - elif not layer_config.is_noop_ffn_layer: - raise NotImplementedError( - f"Loading weights for a layer with FFN of type {layer_config.ffn.impl} is not implemented yet" - ) - - return weights - - -def load_weights_from_hf_model( - hf_model: "transformers.PreTrainedModel", - config: "DeciConfig", - act_range: Optional[dict] = None, - qkv_para: Optional[dict] = None, - smoother: Optional[dict] = None) -> Dict[str, torch.Tensor]: - quant_algo = config.quantization.quant_algo - use_weight_only = quant_algo in [QuantAlgo.W8A16, QuantAlgo.W4A16] - if quant_algo == QuantAlgo.W8A16: - torch.int8 - elif quant_algo == QuantAlgo.W4A16: - torch.quint4x2 - else: - pass - - use_smooth_quant = config.quantization._use_plugin_sq - int8_kv_cache = config.quantization.kv_cache_quant_algo == QuantAlgo.INT8 - if use_smooth_quant or int8_kv_cache: - assert act_range is not None - assert qkv_para is not None - assert smoother is not None - - # TODO(oargov): add support for these quants - assert not use_weight_only, "WOQ is not supported yet" - assert not use_smooth_quant, "SmoothQuant is not supported yet" - assert not int8_kv_cache, "INT8 KV cache is not supported yet" - - # TODO(oargov): support moe - moe_config = getattr(config, "moe", None) - assert moe_config is None, "MoE is not supported yet" - - # TODO(oargov): implement resisdual mlp - residual_mlp = getattr(config, "residual_mlp", None) - assert not residual_mlp, "Residual MLP is not supported yet" - - loader = HFModelWeightsLoader(hf_model=hf_model, dtype=config.dtype) - logger.info('Converting weights from Huggingface model...') - return load_model_weights(loader=loader, config=config) - - -def load_weights_from_hf_safetensors( - model_dir: Union[str, Path], - config: "DeciConfig") -> Dict[str, torch.Tensor]: - - if isinstance(model_dir, str): - model_dir = Path(model_dir) - - loader = SafetensorsWeightsLoader(model_dir=model_dir, dtype=config.dtype) - logger.info('Loading weights from Huggingface safetensors...') - return load_model_weights(loader=loader, config=config) - - -def update_weights_following_modelopt_optimization( - weights: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: - # Rename MLPs to FFNs to match TRTLLM implementation expectation - weights = {k.replace('.mlp.', '.ffn.'): v for k, v in weights.items()} - - # Move all linear attentions to their expected locations - weights = { - k.replace('.attn_replacing_linear.', '.attention.'): v - for k, v in weights.items() - } - - # Move all linear MLPs to their expected locations - weights = { - k.replace('.mlp_replacing_linear.', '.ffn.'): v - for k, v in weights.items() - } - - return weights diff --git a/tensorrt_llm/models/nemotron_nas/layer_config.py b/tensorrt_llm/models/nemotron_nas/layer_config.py deleted file mode 100644 index 84ed3487da7c..000000000000 --- a/tensorrt_llm/models/nemotron_nas/layer_config.py +++ /dev/null @@ -1,86 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import enum -from dataclasses import dataclass, field -from typing import Any, Dict, Optional - - -class AttentionImplementation(str, enum.Enum): - ATTENTION = "attention" - LINEAR = "linear" - NO_OP = "no_op" - - -class FFNImplementation(str, enum.Enum): - MLP = "mlp" - LINEAR = "linear" - NO_OP = "no_op" - - -@dataclass(frozen=True, kw_only=True) -class AttentionConfig: - impl: AttentionImplementation = AttentionImplementation.ATTENTION - num_key_value_heads: Optional[int] = None - - @property - def needs_kv_cache(self) -> bool: - return self.impl == AttentionImplementation.ATTENTION - - -@dataclass(frozen=True, kw_only=True) -class FFNConfig: - impl: FFNImplementation = FFNImplementation.MLP - intermediate_size: Optional[int] = None - - -@dataclass(frozen=True, kw_only=True) -class DeciLayerConfig: - attention: AttentionConfig = field(default_factory=AttentionConfig) - ffn: FFNConfig = field(default_factory=FFNConfig) - - @classmethod - def from_dict(cls, d: Dict[str, Any]) -> "DeciLayerConfig": - assert "attention" in d, "Missing attention configuration" - assert "ffn" in d, "Missing mlp configuration" - - return cls( - attention=AttentionConfig(**d["attention"]), - ffn=FFNConfig(**d["ffn"]), - ) - - @property - def is_attention_layer(self) -> bool: - return self.attention.impl == AttentionImplementation.ATTENTION - - @property - def is_mlp_layer(self) -> bool: - return self.ffn.impl == FFNImplementation.MLP - - @property - def is_noop_attention_layer(self) -> bool: - return self.attention.impl == AttentionImplementation.NO_OP - - @property - def is_linear_attention_layer(self) -> bool: - return self.attention.impl == AttentionImplementation.LINEAR - - @property - def is_noop_ffn_layer(self) -> bool: - return self.ffn.impl == FFNImplementation.NO_OP - - @property - def is_linear_ffn_layer(self) -> bool: - return self.ffn.impl == FFNImplementation.LINEAR diff --git a/tensorrt_llm/models/nemotron_nas/model.py b/tensorrt_llm/models/nemotron_nas/model.py deleted file mode 100644 index f6562d3a7649..000000000000 --- a/tensorrt_llm/models/nemotron_nas/model.py +++ /dev/null @@ -1,816 +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. -from dataclasses import dataclass -from typing import List, Optional, Tuple, Type, Union - -from tensorrt_llm.functional import (AllReduceFusionOp, AllReduceParams, - AttentionMaskType, PositionEmbeddingType, - Tensor, gather_last_token_logits, recv, - send) -from tensorrt_llm.layers.attention import (Attention, AttentionParams, - KeyValueCacheParams, - SpecDecodingParams) -from tensorrt_llm.layers.embedding import Embedding -from tensorrt_llm.layers.linear import ColumnLinear -from tensorrt_llm.layers.lora import LoraParams -from tensorrt_llm.layers.mlp import GatedMLP -from tensorrt_llm.layers.normalization import RmsNorm -from tensorrt_llm.llmapi.kv_cache_type import KVCacheType -from tensorrt_llm.mapping import Mapping -from tensorrt_llm.models.convert_utils import has_safetensors -from tensorrt_llm.models.modeling_utils import DecoderModelForCausalLM -from tensorrt_llm.models.nemotron_nas.config import DeciConfig -from tensorrt_llm.models.nemotron_nas.convert import ( - load_weights_from_hf_model, load_weights_from_hf_safetensors, - update_weights_following_modelopt_optimization) -from tensorrt_llm.module import Module, ModuleList -from tensorrt_llm.plugin.plugin import init_all_reduce_helper - -from ..._common import default_net -from ..._utils import pad_vocab_size -from ..modeling_utils import PretrainedConfig, QuantConfig, preprocess_weights - - -@dataclass -class DeciLMLayerOutput: - hidden_states: Tensor - present_kv: Optional[Tensor] = None - - -@dataclass -class DeciLMLayerListOutput: - hidden_states: Tensor - present_kvs: List[Tensor] - - -class NoOp(Module): - - def forward(self, hidden_states: Tensor, *args, **kwargs) -> int: - return 0 - - -class NoOpAttention(NoOp): - - def forward(self, - hidden_states: Tensor, - attention_mask=None, - use_cache: bool = False, - *args, - **kwargs) -> Union[int, Tuple[int, None]]: - out = super().forward(hidden_states=hidden_states, - attention_mask=attention_mask, - use_cache=use_cache, - *args, - **kwargs) - if use_cache: - return out, None - return out - - -class LinearAttention(ColumnLinear): - - def forward(self, - hidden_states: Tensor, - attention_mask=None, - use_cache: bool = False, - *args, - **kwargs) -> Union[Tensor, Tuple[Tensor, None]]: - out = super().forward(x=hidden_states, - lora_runtime_params=None, - lora_hidden_state=None) - - if use_cache: - return out, None - return out - - -class LinearFFN(ColumnLinear): - - def forward(self, - hidden_states, - lora_layer_params=None, - all_reduce_params: Optional[AllReduceParams] = None) -> Tensor: - return super().forward(x=hidden_states, - lora_runtime_params=None, - lora_hidden_state=None) - - -NoOpFFN = NoOp -NoOpLayerNorm = NoOp - - -class DeciLMDecoderLayer(Module): - - def __init__(self, config: DeciConfig, layer_idx: int): - super().__init__() - self.layer_idx = layer_idx - self.config = config - - layers_range = config.mapping.pp_layers(config.num_hidden_layers) - self.local_layer_idx = layer_idx - layers_range[0] - - self.layer_config = self.config.get_layer_config(self.layer_idx) - - self._init_attention() - self._init_ffn() - - @property - def input_layernorm_was_fused(self) -> bool: - """ - The previous layer ran our input_layernorm for us if: - 1. The reduce_fusion plugin is enabled and - 2. We are not the first local model layer and - 3. The previous layer is an MLP layer - """ - return default_net( - ).plugin_config.reduce_fusion and self.local_layer_idx > 0 and self.config.get_layer_config( - self.layer_idx - - 1).is_mlp_layer and self.needs_input_layernorm_fusion - - @property - def needs_input_layernorm_fusion(self) -> bool: - """ - This layer needs the previous layer to perform input_layernorm fusion if: - 1. The reduce_fusion plugin is enabled and - 2. This is not a NOOP attention layer (otherwise it has no input_layernorm) - """ - return default_net( - ).plugin_config.reduce_fusion and not self.layer_config.is_noop_attention_layer - - @property - def can_fuse_post_layernorm(self) -> bool: - """ - This layer can fuse attention and post_layernorm if: - 1. The reduce_fusion plugin is enabled and - 2. It is an attention layer and - 3. It is not a NOOP FFN layer (othrewise it has no post_layernorm) - """ - return default_net( - ).plugin_config.reduce_fusion and self.layer_config.is_attention_layer and not self.layer_config.is_noop_ffn_layer - - @property - def can_fuse_input_layernorm(self) -> bool: - """ - This layer can run the next layer's input_layernorm if: - 1. The reduce_fusion plugin is enable and - 2. It is an MLP layer - """ - return default_net( - ).plugin_config.reduce_fusion and self.layer_config.is_mlp_layer - - def _init_attention(self) -> None: - """ - Initialize some attention alternative - """ - # normal attention - if self.layer_config.is_attention_layer: - # according to recurrentgemma, len(layer_types) can be less than num_hidden_layers - # in this case, the list should wrap-around - # for example, if layer_types = ["attention", "recurrent", "recurrent"], and we have 5 layers, we get: - # layer 0 ==> attention - # layer 1 ==> recurrent - # layer 2 ==> recurrent - # layer 3 ==> attention - # layer 4 ==> recurrent - # we check which layers are local to our rank - layers_range = self.config.mapping.pp_layers( - self.config.num_hidden_layers) - # then take the size of layer_types in the config - layer_type_len = len(self.config.layer_types) - # collect the layer types of all the local layers - local_layer_types = [ - self.config.layer_types[layer_id % layer_type_len] - for layer_id in layers_range - ] - # and see how many of them are attention layers to determine our local attention layer idx - local_attn_layer_idx = local_layer_types[:self. - local_layer_idx].count( - "attention") - - # Iterate over all local layer configs, getting num_kv_heads of the attention ones - num_kv_heads_per_local_layer = [ - layer_config.attention.num_key_value_heads for layer_config in - [self.config.layer_configs[idx] for idx in layers_range] - if layer_config.is_attention_layer - ] - - # adjust num heads according to tp size - num_kv_heads_per_local_layer = [ - (nheads + self.config.mapping.tp_size - 1) // - self.config.mapping.tp_size - for nheads in num_kv_heads_per_local_layer - ] - nheads_tp = (self.layer_config.attention.num_key_value_heads + - self.config.mapping.tp_size - - 1) // self.config.mapping.tp_size - - self.input_layernorm = RmsNorm( - normalized_shape=self.config.hidden_size, - eps=self.config.norm_epsilon, - dtype=self.config.dtype, - ) - - self.attention = Attention( - local_layer_idx=local_attn_layer_idx, - hidden_size=self.config.hidden_size, - attention_head_size=self.config.head_size, - num_attention_heads=self.config.num_attention_heads, - num_kv_heads=self.layer_config.attention.num_key_value_heads, - max_position_embeddings=self.config.max_position_embeddings, - dtype=self.config.dtype, - attention_mask_type=AttentionMaskType.causal, - bias=False, - position_embedding_type=PositionEmbeddingType.rope_gpt_neox, - rotary_embedding_base=self.config.rotary_base, - rotary_embedding_scaling=self.config.rotary_scaling, - tp_group=self.config.mapping.tp_group, - tp_size=self.config.mapping.tp_size, - tp_rank=self.config.mapping.tp_rank, - quant_mode=self.config.quant_mode) - - elif self.layer_config.is_noop_attention_layer: - self.input_layernorm = NoOpLayerNorm() - self.attention = NoOpAttention() - - elif self.layer_config.is_linear_attention_layer: - self.input_layernorm = RmsNorm( - normalized_shape=self.config.hidden_size, - eps=self.config.norm_epsilon, - dtype=self.config.dtype, - ) - - self.attention = LinearAttention( - in_features=self.config.hidden_size, - out_features=self.config.hidden_size, - bias=False, - dtype=self.config.dtype, - tp_group=self.config.mapping.tp_group, - tp_size=self.config.mapping.tp_size, - gather_output=True) - - else: - raise NotImplementedError( - f"Attention of type {str(self.layer_config.attention.impl)} is not implemented" - ) - - def _init_ffn(self) -> None: - """ - Initialize some ffn alternative - """ - - if self.layer_config.is_mlp_layer: - intermediate_size = self.layer_config.ffn.intermediate_size or self.config.intermediate_size - mlp_hidden_size = intermediate_size or self.config.hidden_size * 4 - - self.post_layernorm = RmsNorm( - normalized_shape=self.config.hidden_size, - eps=self.config.norm_epsilon, - dtype=self.config.dtype, - ) - - self.ffn = GatedMLP( - hidden_size=self.config.hidden_size, - ffn_hidden_size=mlp_hidden_size, - hidden_act=self.config.hidden_act, - bias=False, - dtype=self.config.dtype, - tp_group=self.config.mapping.tp_group, - tp_size=self.config.mapping.tp_size, - quant_mode=self.config.quant_mode, - ) - - elif self.layer_config.is_noop_ffn_layer: - self.post_layernorm = NoOpLayerNorm() - self.ffn = NoOpFFN() - - elif self.layer_config.is_linear_ffn_layer: - self.post_layernorm = RmsNorm( - normalized_shape=self.config.hidden_size, - eps=self.config.norm_epsilon, - dtype=self.config.dtype, - ) - - self.ffn = LinearFFN(in_features=self.config.hidden_size, - out_features=self.config.hidden_size, - bias=False, - dtype=self.config.dtype, - tp_group=self.config.mapping.tp_group, - tp_size=self.config.mapping.tp_size, - gather_output=True) - - else: - raise NotImplementedError( - f"FFN of type {str(self.layer_config.ffn.impl)} is not implemented" - ) - - def forward(self, - hidden_states: Tensor | Tuple[Tensor, Tensor], - attention_mask: Optional[Tensor] = None, - use_cache: bool = False, - spec_decoding_params=None, - kv_cache_params: Optional[KeyValueCacheParams] = None, - attention_params: Optional[AttentionParams] = None, - lora_layer_params: Optional[LoraParams] = None, - next_layer_input_layernorm_args: Optional[Tuple[Tensor, - float]] = None): - if self.input_layernorm_was_fused: - # previous layer already performed our layer norm - assert isinstance(hidden_states, tuple) - hidden_states, residual = hidden_states - else: - residual = hidden_states - hidden_states = self.input_layernorm(hidden_states) - - if self.can_fuse_post_layernorm: - all_reduce_params = AllReduceParams( - fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM, - residual=residual, - norm_weight=self.post_layernorm.weight.value, - eps=self.post_layernorm.eps) - else: - all_reduce_params = None - - attention_output = self._run_attention( - hidden_states=hidden_states, - attention_mask=attention_mask, - use_cache=use_cache, - spec_decoding_params=spec_decoding_params, - kv_cache_params=kv_cache_params, - attention_params=attention_params, - lora_layer_params=lora_layer_params, - all_reduce_params=all_reduce_params) - - if use_cache: - attention_output, present_kv = attention_output - else: - present_kv = None - - if self.can_fuse_post_layernorm: - hidden_states, residual = attention_output - else: - hidden_states = residual + attention_output - residual = hidden_states - hidden_states = self.post_layernorm(hidden_states) - - if next_layer_input_layernorm_args is not None: - assert self.can_fuse_input_layernorm - norm_weight, eps = next_layer_input_layernorm_args - all_reduce_params = AllReduceParams( - fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM, - residual=residual, - norm_weight=norm_weight, - eps=eps) - hidden_states = self._run_ffn(hidden_states, - lora_layer_params=lora_layer_params, - all_reduce_params=all_reduce_params) - - else: - hidden_states = self._run_ffn(hidden_states, - lora_layer_params=lora_layer_params) - hidden_states = residual + hidden_states - - return DeciLMLayerOutput(hidden_states=hidden_states, - present_kv=present_kv) - - def _run_attention( - self, - hidden_states: Tensor, - attention_mask: Optional[Tensor] = None, - use_cache: bool = False, - spec_decoding_params=None, - kv_cache_params: Optional[KeyValueCacheParams] = None, - attention_params: Optional[AttentionParams] = None, - lora_layer_params: Optional[LoraParams] = None, - all_reduce_params: Optional[AllReduceParams] = None - ) -> Union[Tensor, Tuple[Tensor, None]]: - """ - Ideally, this functionality would be encapsulated in a LinearAttention class, but during - FP8 and lower quantization, our linear classes get overrun by ModelOpt, thus we must - control the attention inputs at the DecoderLayer level. - """ - if self.layer_config.is_linear_attention_layer: - out = self.attention(hidden_states) - return out, None if use_cache else out - else: - if not self.layer_config.is_attention_layer: - assert all_reduce_params is None, f"Layer with attention of type {self.layer_config.attention.impl} can't do reduce_fusion" - - return self.attention(hidden_states=hidden_states, - attention_mask=attention_mask, - use_cache=use_cache, - spec_decoding_params=spec_decoding_params, - kv_cache_params=kv_cache_params, - attention_params=attention_params, - lora_layer_params=lora_layer_params, - all_reduce_params=all_reduce_params) - - def _run_ffn(self, - hidden_states, - lora_layer_params=None, - all_reduce_params: Optional[AllReduceParams] = None): - """ - Ideally, this functionality would be encapsulated in a LinearMLP class, but during - FP8 and lower quantization, our linear classes get overrun by ModelOpt, thus we must - control the MLP inputs at the DecoderLayer level. - """ - if all_reduce_params is not None: - assert self.layer_config.is_mlp_layer, f"Layer with FFN of type {self.layer_config.ffn.impl} can't do reduce_fusion" - - if self.layer_config.is_linear_ffn_layer: - return self.ffn(hidden_states) - else: - return self.ffn(hidden_states, - lora_layer_params=lora_layer_params, - all_reduce_params=all_reduce_params) - - -class DeciLMDecoderLayerList(ModuleList): - - def __init__(self, cls: Type[DeciLMDecoderLayer], config: DeciConfig): - self.num_hidden_layers = config.num_hidden_layers - # global indices of local layers - self.layer_list = config.mapping.pp_layers(config.num_hidden_layers) - super().__init__([cls(config, idx) for idx in self.layer_list]) - # global indices of local attention layers - self.attention_layer_list = [ - self.layer_list[i] for i, layer in enumerate(self) - if layer.layer_config.is_attention_layer - ] - - def forward( - self, - hidden_states: Tensor, - use_cache: bool, - attention_mask: Optional[Tensor], - kv_cache_params: KeyValueCacheParams, - attention_params: Optional[AttentionParams] = None, - position_ids: Optional[Tensor] = None, - lora_params: Optional[LoraParams] = None, - spec_decoding_params: Optional[SpecDecodingParams] = None, - ) -> DeciLMLayerListOutput: - kv_cache_params.fill_none_tensor_list(len(self.layer_list)) - - presents = [] - - # put None where we don't have attention layers - pkv_iter = iter(kv_cache_params.past_key_value) - - past_key_values = [x for x in pkv_iter] - - for layer_idx, (layer, past) in enumerate(zip(self, past_key_values)): - next_layer_input_layernorm_args = None - if default_net().plugin_config.reduce_fusion: - if layer_idx < self.layer_list[-1]: - # this is not the last layer - next_layer = self[layer_idx + 1] - if layer.can_fuse_input_layernorm and next_layer.needs_input_layernorm_fusion: - # this layer can fuse the next layer's input_layernorm - next_layer_input_layernorm_args = ( - next_layer.input_layernorm.weight.value, - next_layer.input_layernorm.eps) - - layer_out = layer( - hidden_states=hidden_states, - attention_mask=attention_mask, - attention_params=attention_params, - kv_cache_params=KeyValueCacheParams( - past_key_value=[past], - host_past_key_value_lengths=kv_cache_params. - host_past_key_value_lengths, - host_max_attention_window_sizes=kv_cache_params. - host_max_attention_window_sizes, - host_sink_token_length=kv_cache_params. - host_sink_token_length, - kv_cache_block_offsets=kv_cache_params. - kv_cache_block_offsets, - host_kv_cache_block_offsets=kv_cache_params. - host_kv_cache_block_offsets, - host_kv_cache_pool_pointers=kv_cache_params. - host_kv_cache_pool_pointers, - host_kv_cache_pool_mapping=kv_cache_params. - host_kv_cache_pool_mapping, - cache_indirection=kv_cache_params.cache_indirection, - ), - spec_decoding_params=spec_decoding_params, - use_cache=use_cache, - lora_layer_params=lora_params.get_layer_config(layer_idx) - if lora_params is not None - and lora_params.lora_ranks is not None else None, - next_layer_input_layernorm_args=next_layer_input_layernorm_args) - - hidden_states = layer_out.hidden_states - if use_cache and layer_out.present_kv is not None: - presents.append(layer_out.present_kv) - - return DeciLMLayerListOutput(hidden_states=hidden_states, - present_kvs=presents) - - -class DeciLMModel(Module): - - def __init__(self, config: DeciConfig) -> None: - super().__init__() - init_all_reduce_helper() - - self.mapping = config.mapping - if self.mapping.is_first_pp_rank(): - # first rank in pipeline-parallel handles token embedding - assert config.vocab_size is not None - self.vocab_embedding = Embedding(config.vocab_size, - config.hidden_size, - dtype=config.dtype) - - self.position_embedding_type = config.position_embedding_type - self.layers = DeciLMDecoderLayerList(DeciLMDecoderLayer, config) - - if self.mapping.is_last_pp_rank(): - # last rank in pipeline-parallel handles final norm - self.ln_f = RmsNorm( - normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - dtype=config.dtype, - ) - - def _vocab_embedding(self, - input_ids: Tensor, - prompt_embedding_table: Optional[Tensor] = None, - prompt_tasks: Optional[Tensor] = None, - prompt_vocab_size: Optional[Tensor] = None) -> Tensor: - # prompt tuning - ptuning_args = ([ - prompt_embedding_table, prompt_tasks, prompt_vocab_size - ] if prompt_embedding_table is not None else []) - - hidden_states = self.vocab_embedding(input_ids, *ptuning_args) - return hidden_states - - def forward( - self, - input_ids, - position_ids=None, - use_cache: bool = False, - attention_mask: Optional[Tensor] = None, - spec_decoding_params=None, - kv_cache_params: Optional[KeyValueCacheParams] = None, - attention_params: Optional[AttentionParams] = None, - hidden_states: Optional[Tensor] = None, - prompt_embedding_table: Optional[Tensor] = None, - prompt_tasks: Optional[Tensor] = None, - prompt_vocab_size: Optional[Tensor] = None, - lora_params: Optional[LoraParams] = None, - ) -> DeciLMLayerListOutput: - - if self.mapping.is_first_pp_rank(): - # first pipeline rank ==> do prompt embedding - hidden_states = self._vocab_embedding( - input_ids=input_ids, - prompt_embedding_table=prompt_embedding_table, - prompt_tasks=prompt_tasks, - prompt_vocab_size=prompt_vocab_size) - else: - # receive hidden states from prior rank in the pipeline - hidden_states = recv(hidden_states, self.mapping.prev_pp_rank()) - - layers_out = self.layers.forward( - hidden_states, - use_cache=use_cache, - attention_mask=attention_mask, - kv_cache_params=kv_cache_params, - attention_params=attention_params, - lora_params=lora_params, - spec_decoding_params=spec_decoding_params, - ) - - if self.mapping.is_last_pp_rank(): - # last pipeline rank ==> do final norm - hidden_states = self.ln_f(layers_out.hidden_states) - else: - # send hidden states to next rank in the pipeline - hidden_states = send(layers_out.hidden_states, - self.mapping.next_pp_rank()) - - return DeciLMLayerListOutput(hidden_states=hidden_states, - present_kvs=layers_out.present_kvs) - - -class DeciLMForCausalLM(DecoderModelForCausalLM): - config_class = DeciConfig - - def __init__(self, config: DeciConfig): - - transformer = DeciLMModel(config) - vocab_size_padded = pad_vocab_size(config.vocab_size, - config.mapping.tp_size) - - if config.mapping.is_last_pp_rank(): - # last pipeline rank needs to do calculate logits - lm_head = ColumnLinear( - config.hidden_size, - vocab_size_padded, - bias=False, - dtype=config.dtype, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size, - gather_output=True, - ) - else: - lm_head = None - super().__init__(config, transformer, lm_head) - - # Create constant attention parameters to be reused by all layers. - Attention.create_attention_const_params(self, config) - self.position_embedding_type = config.position_embedding_type - - @classmethod - def from_hugging_face(cls, - hf_model_or_dir: Union[ - str, 'transformers.PreTrainedModel'], - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - load_by_shard: bool = False, - load_model_on_cpu: bool = False, - trust_remote_code: bool = True, - **kwargs) -> "DeciLMForCausalLM": - import transformers - - # TODO(oargov): add support for these - assert not load_by_shard, "load_by_shard is not implemented yet" - - use_preloading = isinstance(hf_model_or_dir, - transformers.PreTrainedModel) - if use_preloading: - hf_config_or_dir = hf_model_or_dir.config - else: - hf_config_or_dir = hf_model_or_dir - - config = DeciConfig.from_hugging_face( - hf_config_or_dir, - dtype=dtype, - mapping=mapping, - quant_config=quant_config, - trust_remote_code=trust_remote_code, - **kwargs) - - if use_preloading: - assert not load_by_shard - weights = load_weights_from_hf_model(hf_model_or_dir, config) - elif has_safetensors( - hf_model_or_dir) and not config.quant_mode.has_any_quant(): - weights = load_weights_from_hf_safetensors(hf_model_or_dir, config) - else: - hf_model = transformers.AutoModelForCausalLM.from_pretrained( - hf_model_or_dir, - device_map='auto' if not load_model_on_cpu else 'cpu', - dtype=dtype, - trust_remote_code=trust_remote_code, - ) - weights = load_weights_from_hf_model(hf_model, config) - preprocess_weights(weights, config) - - model = DeciLMForCausalLM(config) - model.load(weights) - return model - - @classmethod - def from_checkpoint(cls, - ckpt_dir: str, - rank: Optional[int] = None, - config: Optional["PretrainedConfig"] = None): - return super().from_checkpoint( - ckpt_dir, - rank, - config, - preprocess_weights_hook= - update_weights_following_modelopt_optimization, - ) - - def forward( - self, - input_ids: Tensor, - position_ids: Optional[Tensor] = None, - use_cache: bool = False, - last_token_ids: Optional[Tensor] = None, - attention_mask: Optional[Tensor] = None, - kv_cache_params: Optional[KeyValueCacheParams] = None, - attention_params: Optional[AttentionParams] = None, - hidden_states: Optional[Tensor] = None, - prompt_embedding_table: Optional[Tensor] = None, - prompt_tasks: Optional[Tensor] = None, - prompt_vocab_size: Optional[Tensor] = None, - lora_params: Optional[LoraParams] = None, - spec_decoding_params=None, - ): - # fill attention params. - attention_params = Attention.fill_attention_params( - self, attention_params) - - model_out = self.transformer.forward( - input_ids=input_ids, - position_ids=position_ids, - use_cache=use_cache, - attention_mask=attention_mask, - kv_cache_params=kv_cache_params, - attention_params=attention_params, - lora_params=lora_params, - hidden_states=hidden_states, - prompt_embedding_table=prompt_embedding_table, - prompt_tasks=prompt_tasks, - prompt_vocab_size=prompt_vocab_size, - spec_decoding_params=spec_decoding_params) - hidden_states = model_out.hidden_states - - if self.config.mapping.is_last_pp_rank(): - hidden_states = gather_last_token_logits( - hidden_states, - last_token_ids, - default_net().plugin_config.remove_input_padding, - ) - - lm_logits = self.lm_head(hidden_states) - lm_logits.mark_output("logits", self.config.logits_dtype) - else: - hidden_states.mark_output("hidden_states_output", self.config.dtype) - - if use_cache and not default_net().plugin_config.paged_kv_cache: - presents = model_out.present_kvs - for i, present in zip(self.transformer.layers.attention_layer_list, - presents): - present.mark_output(f"present_key_value_{i}", - self.config.kv_dtype) - if self.config.mapping.is_last_pp_rank(): - return (lm_logits, presents, hidden_states) - return (hidden_states, presents) - else: - if self.config.mapping.is_last_pp_rank(): - return lm_logits, hidden_states - return hidden_states - - def prepare_attention_inputs( - self, - *, - max_batch_size: int, - max_beam_width: int, - max_input_len: int, - max_seq_len: int, - num_kv_heads: int, - head_size: int, - num_layers: int, - kv_dtype: str, - kv_cache_type: KVCacheType, - num_profiles: int = 1, - enable_ctx_gen_opt_profiles: bool = False, - remove_input_padding: bool = False, - use_gpt_attention_plugin: bool = False, - paged_kv_cache: bool = False, - tokens_per_block: int = 32, - mapping: Mapping = Mapping(), - use_cache: bool = True, - streamingllm: bool = False, - attn_layer_idx: Optional[List[int]] = None, - opt_batch_size: Optional[int] = None, - num_kv_heads_per_layer: Optional[List[int]] = None): - - if attn_layer_idx is None: - attn_layer_idx, num_kv_heads_per_layer = [], [] - for layer_idx in range(self.config.num_hidden_layers): - layer_config = self.config.get_layer_config(layer_idx) - if layer_config.is_attention_layer: - attn_layer_idx.append(layer_idx) - num_kv_heads_per_layer.append( - layer_config.attention.num_key_value_heads) - - attention_inputs = super().prepare_attention_inputs( - max_batch_size=max_batch_size, - max_beam_width=max_beam_width, - max_input_len=max_input_len, - max_seq_len=max_seq_len, - num_kv_heads=num_kv_heads, - head_size=head_size, - num_layers=num_layers, - kv_dtype=kv_dtype, - num_profiles=num_profiles, - kv_cache_type=kv_cache_type, - enable_ctx_gen_opt_profiles=enable_ctx_gen_opt_profiles, - remove_input_padding=remove_input_padding, - use_gpt_attention_plugin=use_gpt_attention_plugin, - tokens_per_block=tokens_per_block, - mapping=mapping, - streamingllm=streamingllm, - attn_layer_idx=attn_layer_idx, - opt_batch_size=opt_batch_size, - num_kv_heads_per_layer=num_kv_heads_per_layer) - - return attention_inputs diff --git a/tensorrt_llm/models/opt/__init__.py b/tensorrt_llm/models/opt/__init__.py deleted file mode 100644 index 71bf6d298d42..000000000000 --- a/tensorrt_llm/models/opt/__init__.py +++ /dev/null @@ -1,14 +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. diff --git a/tensorrt_llm/models/opt/model.py b/tensorrt_llm/models/opt/model.py deleted file mode 100644 index 2a81a8c29b3e..000000000000 --- a/tensorrt_llm/models/opt/model.py +++ /dev/null @@ -1,181 +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. - -from ..._utils import pad_vocab_size -from ...functional import Tensor -from ...layers import (MLP, Attention, AttentionMaskType, ColumnLinear, - Embedding, LayerNorm) -from ...module import Module -from ..modeling_utils import (DecoderLayerList, DecoderModelForCausalLM, - PretrainedConfig) - - -class OPTDecoderLayer(Module): - - def __init__(self, config: PretrainedConfig, layer_idx: int): - super().__init__() - self.layer_idx = layer_idx - self.config = config - self.do_layer_norm_before = self.config.do_layer_norm_before - - hidden_size = config.hidden_size - dtype = config.dtype - tp_group = config.mapping.tp_group - tp_size = config.mapping.tp_size - - self.input_layernorm = LayerNorm(normalized_shape=hidden_size, - dtype=dtype) - - layers_range = config.mapping.pp_layers(config.num_hidden_layers) - local_layer_idx = layer_idx - layers_range[0] - self.attention = Attention( - local_layer_idx=local_layer_idx, - hidden_size=hidden_size, - num_attention_heads=config.num_attention_heads, - max_position_embeddings=config.max_position_embeddings, - attention_mask_type=AttentionMaskType.causal, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - quant_mode=config.quant_mode) - - mlp_hidden_size = hidden_size * 4 if config.intermediate_size is None else config.intermediate_size - - self.mlp = MLP(hidden_size=hidden_size, - ffn_hidden_size=mlp_hidden_size, - hidden_act=config.hidden_act, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - quant_mode=config.quant_mode) - self.post_layernorm = LayerNorm(normalized_shape=hidden_size, - dtype=dtype) - - def forward(self, - hidden_states: Tensor, - attention_mask=None, - use_cache=False, - kv_cache_params=None, - attention_params=None): - residual = hidden_states - - attention_input = hidden_states - if self.do_layer_norm_before: - attention_input = self.input_layernorm(hidden_states) - - # At this point the hidden_states object must be a Tensor. - assert isinstance(attention_input, Tensor) - - attention_output = self.attention(attention_input, - attention_mask=attention_mask, - use_cache=use_cache, - kv_cache_params=kv_cache_params, - attention_params=attention_params) - if use_cache: - attention_output, presents = attention_output - - hidden_states = residual + attention_output - if not self.do_layer_norm_before: - hidden_states = self.input_layernorm(hidden_states) - - residual = hidden_states - if self.do_layer_norm_before: - hidden_states = self.post_layernorm(hidden_states) - - hidden_states = self.mlp(hidden_states) - - hidden_states = residual + hidden_states - - if not self.do_layer_norm_before: - hidden_states = self.post_layernorm(hidden_states) - - if use_cache: - return (hidden_states, presents) - return hidden_states - - -class OPTModel(Module): - - def __init__(self, config: PretrainedConfig): - super().__init__() - self.do_layer_norm_before = config.do_layer_norm_before - - self.vocab_embedding = Embedding(config.vocab_size, - config.hidden_size, - dtype=config.dtype) - self.position_embedding = Embedding(config.max_position_embeddings, - config.hidden_size, - dtype=config.dtype) - - self.layers = DecoderLayerList(OPTDecoderLayer, config) - - if self.do_layer_norm_before: - self.ln_f = LayerNorm(normalized_shape=config.hidden_size, - dtype=config.dtype) - - def forward(self, - input_ids: Tensor, - position_ids=None, - use_cache=False, - attention_mask=None, - kv_cache_params=None, - attention_params=None, - prompt_embedding_table=None, - prompt_tasks=None, - prompt_vocab_size=None, - **kwargs): - - args = [prompt_embedding_table, prompt_tasks, prompt_vocab_size - ] if prompt_embedding_table is not None else [] - hidden_states = self.vocab_embedding(input_ids, *args) - hidden_states = hidden_states + self.position_embedding(position_ids) - - hidden_states = self.layers(hidden_states, - use_cache=use_cache, - attention_mask=attention_mask, - kv_cache_params=kv_cache_params, - attention_params=attention_params) - - if use_cache: - hidden_states, presents = hidden_states - - if self.do_layer_norm_before: - hidden_states = self.ln_f(hidden_states) - - if use_cache: - return (hidden_states, tuple(presents)) - return hidden_states - - -class OPTForCausalLM(DecoderModelForCausalLM): - - def __init__(self, config: PretrainedConfig): - self.check_config(config) - transformer = OPTModel(config) - vocab_size_padded = pad_vocab_size(config.vocab_size, - config.mapping.tp_size) - - lm_head = ColumnLinear(config.hidden_size, - vocab_size_padded, - bias=False, - dtype=config.dtype, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size, - gather_output=True) - - super().__init__(config, transformer, lm_head) - - def check_config(self, config): - config.set_if_not_exist('do_layer_norm_before', False) diff --git a/tensorrt_llm/models/phi/__init__.py b/tensorrt_llm/models/phi/__init__.py deleted file mode 100644 index 71bf6d298d42..000000000000 --- a/tensorrt_llm/models/phi/__init__.py +++ /dev/null @@ -1,14 +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. diff --git a/tensorrt_llm/models/phi/config.py b/tensorrt_llm/models/phi/config.py deleted file mode 100644 index 583de15fadf7..000000000000 --- a/tensorrt_llm/models/phi/config.py +++ /dev/null @@ -1,87 +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. - -from typing import Optional, Union - -from ..._utils import get_hf_rope_theta -from ...mapping import Mapping -from ..convert_utils import infer_dtype -from ..modeling_utils import PretrainedConfig, QuantConfig - - -class PhiConfig(PretrainedConfig): - - def __init__(self, - *, - rotary_base: float = 10000.0, - rotary_scaling: Optional[dict] = None, - **kwargs): - - self.rotary_base = rotary_base - self.rotary_scaling = rotary_scaling - - super().__init__(**kwargs) - - def to_dict(self): - output = super().to_dict() - # Serialize the fields added in PhiConfig - - output['rotary_base'] = self.rotary_base - output['rotary_scaling'] = self.rotary_scaling - - return output - - @classmethod - def from_hugging_face( - cls, - hf_config_or_dir: Union[str, 'transformers.PretrainedConfig'], - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - import transformers - trust_remote_code = kwargs.pop('trust_remote_code', True) - - if isinstance(hf_config_or_dir, transformers.PretrainedConfig): - hf_config = hf_config_or_dir - else: - hf_config_dir = str(hf_config_or_dir) - - hf_config = transformers.AutoConfig.from_pretrained( - hf_config_dir, trust_remote_code=trust_remote_code) - - num_key_value_heads = getattr(hf_config, "num_key_value_heads", - hf_config.num_attention_heads) - rotary_scaling = getattr(hf_config, "rope_scaling", None) - rotary_base = get_hf_rope_theta(hf_config, 10000.0) - dtype = infer_dtype(dtype, getattr(hf_config, 'torch_dtype', None)) - - return cls(architecture=hf_config.architectures[0], - dtype=dtype, - num_hidden_layers=hf_config.num_hidden_layers, - num_attention_heads=hf_config.num_attention_heads, - hidden_size=hf_config.hidden_size, - intermediate_size=hf_config.intermediate_size, - num_key_value_heads=num_key_value_heads, - vocab_size=hf_config.vocab_size, - position_embedding_type='rope_gpt_neox', - max_position_embeddings=hf_config.max_position_embeddings, - hidden_act=hf_config.hidden_act, - rotary_base=rotary_base, - rotary_scaling=rotary_scaling, - rotary_pct=hf_config.partial_rotary_factor, - mapping=mapping, - quantization=quant_config, - **kwargs) diff --git a/tensorrt_llm/models/phi/convert.py b/tensorrt_llm/models/phi/convert.py deleted file mode 100644 index 4bf3406c7262..000000000000 --- a/tensorrt_llm/models/phi/convert.py +++ /dev/null @@ -1,145 +0,0 @@ -import torch - -from ..._utils import get_hf_rope_theta, pad_vocab_size, str_dtype_to_torch - - -def split(v, tp_size, idx, dim=0): - if tp_size == 1: - return v - if len(v.shape) == 1: - return torch.chunk(v, tp_size)[idx].contiguous() - else: - return torch.chunk(v, tp_size, dim=dim)[idx].contiguous() - - -def load_weights_from_hf_model(hf_model, config): - torch_dtype = str_dtype_to_torch(config.dtype) - hf_state_dict = hf_model.state_dict() - weights = {} - is_weight_only = config.quant_mode.is_weight_only() - if config.quant_mode.is_int8_weight_only(): - plugin_weight_only_quant_type = torch.int8 - elif config.quant_mode.is_int4_weight_only(): - plugin_weight_only_quant_type = torch.quint4x2 - # replace key name - for key, value in hf_state_dict.items(): - # Decoder Layers - if "model.layers." in key: - key = key.replace("model.layers.", "transformer.layers.") - key = key.replace("self_attn.", "attention.") - key = key.replace("mlp.fc1.", "mlp.fc.") - key = key.replace("mlp.fc2.", "mlp.proj.") - # Embedding - key = key.replace("model.embed_tokens.weight", - "transformer.vocab_embedding.weight") - # Final Layer norm - key = key.replace("model.final_layernorm.", "transformer.ln_f.") - - weights[key] = value.to(torch_dtype).cpu() - - # merge qkv weights - qkv_keys = ["q_proj", "k_proj", "v_proj"] - - scales = {} - for key in hf_state_dict.keys(): - if 'self_attn.q_proj.weight' in key: - prefix = key.split('self_attn')[0].replace("model.layers.", - "transformer.layers.") - - # [(num_heads x q)|(num_heads x k)|(num_heads x v), hidden_size] - qkv_weights = [] - qkv_bias = [] - - for k in qkv_keys: - split_w = split(weights.pop(f"{prefix}attention.{k}.weight"), - config.mapping.tp_size, config.mapping.tp_rank) - - qkv_weights.append(split_w) - split_b = split(weights.pop(f"{prefix}attention.{k}.bias"), - config.mapping.tp_size, config.mapping.tp_rank) - qkv_bias.append(split_b) - v = torch.cat(qkv_weights, dim=0) - if is_weight_only: - processed_torch_weights, torch_weight_scales = \ - torch.ops.trtllm.symmetric_quantize_last_axis_of_batched_matrix( - v.t().contiguous().cpu(), plugin_weight_only_quant_type) - weights[ - f"{prefix}attention.qkv.weight"] = processed_torch_weights - scales[ - f"{prefix}attention.qkv.per_channel_scale"] = torch_weight_scales - else: - weights[f"{prefix}attention.qkv.weight"] = v - weights[f"{prefix}attention.qkv.bias"] = torch.cat(qkv_bias, dim=0) - - tp_rank = config.mapping.tp_rank - for weight_name in weights: - loaded_weight = weights[weight_name] - if "attention.dense.weight" in weight_name or "mlp.proj.weight" in weight_name: # RowLinear - v = split(loaded_weight, - config.mapping.tp_size, - config.mapping.tp_rank, - dim=1) - if is_weight_only: - processed_torch_weights, torch_weight_scales = \ - torch.ops.trtllm.symmetric_quantize_last_axis_of_batched_matrix( - v.t().contiguous().cpu(), plugin_weight_only_quant_type) - weights[weight_name] = processed_torch_weights - scales[weight_name.replace( - '.weight', '.per_channel_scale')] = torch_weight_scales - else: - weights[weight_name] = v - elif "mlp.fc." in weight_name: - - v = split(loaded_weight, config.mapping.tp_size, - config.mapping.tp_rank) - if is_weight_only and "mlp.fc.weight" in weight_name: - processed_torch_weights, torch_weight_scales = \ - torch.ops.trtllm.symmetric_quantize_last_axis_of_batched_matrix( - v.t().contiguous().cpu(), plugin_weight_only_quant_type) - weights[weight_name] = processed_torch_weights - scales[weight_name.replace( - '.weight', '.per_channel_scale')] = torch_weight_scales - else: - weights[weight_name] = v - elif "lm_head." in weight_name: - output_dim = 0 - shard_size = loaded_weight.shape[output_dim] - tp_rank * shard_size - vocab_size = loaded_weight.shape[output_dim] - - if shard_size % config.mapping.tp_size != 0: - pad_width = pad_vocab_size(vocab_size, - config.mapping.tp_size) - vocab_size - loaded_weight = torch.nn.functional.pad(loaded_weight, - (0, 0, 0, pad_width), - 'constant', - value=0) - weights[weight_name] = split(loaded_weight, config.mapping.tp_size, - config.mapping.tp_rank) - - weights.update(scales) - - return weights - - -def convert_hf_config(hf_config, dtype, args): - config = { - 'architecture': hf_config.architectures[0], - 'dtype': dtype, - 'num_hidden_layers': hf_config.num_hidden_layers, - 'num_attention_heads': hf_config.num_key_value_heads, - 'rotary_pct': hf_config.partial_rotary_factor, - 'rope_theta': get_hf_rope_theta(hf_config, 10000.0), - 'hidden_size': hf_config.hidden_size, - 'intermediate_size': hf_config.intermediate_size, - 'vocab_size': hf_config.vocab_size, - 'position_embedding_type': 'rope_gpt_neox', - 'max_position_embeddings': hf_config.max_position_embeddings, - 'hidden_act': hf_config.hidden_act, - 'mapping': { - 'world_size': args.tp_size * args.pp_size, - 'tp_size': args.tp_size, - 'pp_size': args.pp_size, - } - } - return config diff --git a/tensorrt_llm/models/phi/model.py b/tensorrt_llm/models/phi/model.py deleted file mode 100644 index fc5c6ccbc1f9..000000000000 --- a/tensorrt_llm/models/phi/model.py +++ /dev/null @@ -1,217 +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. -from typing import Optional, Union - -from transformers import AutoModelForCausalLM - -from ..._utils import pad_vocab_size -from ...functional import Tensor -from ...layers import (MLP, Attention, AttentionMaskType, ColumnLinear, - Embedding, LayerNorm) -from ...lora_helper import LoraConfig, use_lora -from ...mapping import Mapping -from ...module import Module -from ..modeling_utils import (DecoderLayerList, DecoderModelForCausalLM, - PretrainedConfig, QuantConfig) -from .config import PhiConfig -from .convert import load_weights_from_hf_model - - -class PhiDecoderLayer(Module): - - def __init__(self, config: PretrainedConfig, layer_idx: int): - super().__init__() - self.config = config - self.layer_idx = layer_idx - tp_group = config.mapping.tp_group - tp_size = config.mapping.tp_size - - self.input_layernorm = LayerNorm(normalized_shape=config.hidden_size, - dtype=config.dtype) - - layers_range = config.mapping.pp_layers(config.num_hidden_layers) - local_layer_idx = layer_idx - layers_range[0] - self.attention = Attention( - local_layer_idx=local_layer_idx, - hidden_size=config.hidden_size, - num_attention_heads=config.num_attention_heads, - rotary_embedding_percentage=config.rotary_pct, - position_embedding_type=config.position_embedding_type, - rotary_embedding_base=config.rotary_base, - max_position_embeddings=config.max_position_embeddings, - dtype=config.dtype, - attention_mask_type=AttentionMaskType.causal, - bias=True, - tp_group=tp_group, - tp_size=tp_size, - quant_mode=config.quant_mode) - - self.mlp = MLP(hidden_size=config.hidden_size, - ffn_hidden_size=config.intermediate_size, - hidden_act=config.hidden_act, - dtype=config.dtype, - tp_group=tp_group, - tp_size=tp_size, - quant_mode=config.quant_mode) - - def forward( - self, - hidden_states: Tensor, - attention_mask=None, - use_cache=False, - kv_cache_params=None, - attention_params=None, - lora_layer_params=None, - ): - residual = hidden_states - - input_layernorm_output = self.input_layernorm(hidden_states) - - attention_output = self.attention( - input_layernorm_output, - attention_mask=attention_mask, - use_cache=use_cache, - kv_cache_params=kv_cache_params, - attention_params=attention_params, - norm_before_bmm1=True, - lora_layer_params=lora_layer_params, - ) - - if use_cache: - attention_output, presents = attention_output - - feed_forward_hidden_states = self.mlp(input_layernorm_output, ) - hidden_states = attention_output + feed_forward_hidden_states + residual - if use_cache: - return (hidden_states, presents) - return hidden_states - - -class PhiModel(Module): - - def __init__(self, config: PretrainedConfig): - super().__init__() - self.vocab_embedding = Embedding(num_embeddings=config.vocab_size, - embedding_dim=config.hidden_size, - dtype=config.dtype) - - self.layers = DecoderLayerList(PhiDecoderLayer, config) - self.ln_f = LayerNorm(normalized_shape=config.hidden_size, - dtype=config.dtype) - - def forward( - self, - input_ids: Tensor, - position_ids=None, - use_cache=False, - attention_mask=None, - kv_cache_params=None, - attention_params=None, - prompt_embedding_table=None, - prompt_tasks=None, - prompt_vocab_size=None, - lora_params=None, - ): - args = [prompt_embedding_table, prompt_tasks, prompt_vocab_size - ] if prompt_embedding_table is not None else [] - hidden_states = self.vocab_embedding(input_ids, *args) - - hidden_states = self.layers(hidden_states, - use_cache=use_cache, - attention_mask=attention_mask, - kv_cache_params=kv_cache_params, - attention_params=attention_params, - lora_params=lora_params) - if use_cache: - hidden_states, presents = hidden_states - - hidden_states = self.ln_f(hidden_states) - - if use_cache: - return (hidden_states, tuple(presents)) - return hidden_states - - -class PhiForCausalLM(DecoderModelForCausalLM): - config_class = PhiConfig - - config_class = PhiConfig - - def __init__(self, config: PretrainedConfig): - self.check_config(config) - transformer = PhiModel(config) - vocab_size_padded = pad_vocab_size(config.vocab_size, - config.mapping.tp_size) - - lm_head = ColumnLinear(config.hidden_size, - vocab_size_padded, - bias=True, - dtype=config.dtype, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size, - gather_output=True) - - self.trtllm_modules_to_hf_modules = { - "attn_q": "q_proj", - "attn_k": "k_proj", - "attn_v": "v_proj" - } - - super().__init__(config, transformer, lm_head) - - def check_config(self, config): - config.set_if_not_exist('partial_rotary_factor', 0.4) - config.set_if_not_exist('rotary_base', 10000.0) - - @classmethod - def from_hugging_face( - cls, - hf_model_or_dir: Union[str, 'transformers.PreTrainedModel'], - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - import transformers - - assert hf_model_or_dir is not None - use_preloading = isinstance(hf_model_or_dir, - transformers.PreTrainedModel) - if use_preloading: - hf_model = hf_model_or_dir - hf_config_or_dir = hf_model.config - else: - hf_model_dir = hf_model_or_dir - hf_config_or_dir = hf_model_or_dir - config = PhiConfig.from_hugging_face(hf_config_or_dir, - dtype=dtype, - mapping=mapping, - quant_config=quant_config, - **kwargs) - if not use_preloading: - trust_remote_code = kwargs.pop('trust_remote_code', True) - - hf_model = AutoModelForCausalLM.from_pretrained( - hf_model_dir, dtype="auto", trust_remote_code=trust_remote_code) - - assert isinstance(hf_model, transformers.PreTrainedModel) - - weights = load_weights_from_hf_model(hf_model, config) - - model = cls(config) - model.load(weights) - return model - - def use_lora(self, lora_config: LoraConfig): - use_lora(self, lora_config, self.trtllm_modules_to_hf_modules) diff --git a/tensorrt_llm/models/phi3/__init__.py b/tensorrt_llm/models/phi3/__init__.py deleted file mode 100644 index 71bf6d298d42..000000000000 --- a/tensorrt_llm/models/phi3/__init__.py +++ /dev/null @@ -1,14 +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. diff --git a/tensorrt_llm/models/phi3/config.py b/tensorrt_llm/models/phi3/config.py deleted file mode 100644 index 42d3954092ec..000000000000 --- a/tensorrt_llm/models/phi3/config.py +++ /dev/null @@ -1,143 +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. - -from typing import Optional, Union - -from ..._utils import get_hf_rope_theta -from ...layers import MoeConfig -from ...mapping import Mapping -from ..convert_utils import infer_dtype -from ..modeling_utils import PretrainedConfig, QuantConfig - - -class Phi3Config(PretrainedConfig): - - def __init__(self, - *, - rotary_base: float = 10000.0, - rotary_scaling: Optional[dict] = None, - **kwargs): - - self.rotary_base = rotary_base - self.rotary_scaling = rotary_scaling - - super().__init__(**kwargs) - - def to_dict(self): - output = super().to_dict() - # Serialize the fields added in PhiConfig - - output['rotary_base'] = self.rotary_base - output['rotary_scaling'] = self.rotary_scaling - - return output - - @classmethod - def from_hugging_face( - cls, - hf_config_or_dir: Union[str, 'transformers.PretrainedConfig'], - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - import transformers - trust_remote_code = kwargs.pop('trust_remote_code', True) - - if isinstance(hf_config_or_dir, transformers.PretrainedConfig): - hf_config = hf_config_or_dir - else: - hf_config_dir = str(hf_config_or_dir) - - hf_config = transformers.AutoConfig.from_pretrained( - hf_config_dir, trust_remote_code=trust_remote_code) - if hasattr(hf_config, "llm_config"): - hf_config = hf_config.llm_config - num_key_value_heads = getattr(hf_config, "num_key_value_heads", - hf_config.num_attention_heads) - dtype = infer_dtype(dtype, getattr(hf_config, 'torch_dtype', None)) - - small_variant = hf_config.architectures[0] == "Phi3SmallForCausalLM" - - kwargs['rotary_pct'] = getattr(hf_config, "partial_rotary_factor", 1.0) - kwargs['tie_word_embeddings'] = getattr(hf_config, - "tie_word_embeddings", False) - if small_variant: - kwargs['gegelu_limit'] = getattr(hf_config, "gegelu_limit", None) - kwargs['rotary_base'] = hf_config.rope_embedding_base - kwargs['mup_attn_multiplier'] = getattr(hf_config, - "mup_attn_multiplier", None) - kwargs['mup_embedding_multiplier'] = getattr( - hf_config, "mup_embedding_multiplier", None) - kwargs['mup_use_scaling'] = getattr(hf_config, "mup_use_scaling", - None) - kwargs['mup_width_multiplier'] = getattr(hf_config, - "mup_width_multiplier", - None) - kwargs['blocksparse_block_size'] = getattr( - hf_config, "blocksparse_block_size", None) - kwargs['blocksparse_homo_head_pattern'] = getattr( - hf_config, "blocksparse_homo_head_pattern", None) - kwargs['blocksparse_num_local_blocks'] = getattr( - hf_config, "blocksparse_num_local_blocks", None) - kwargs['blocksparse_vertical_stride'] = getattr( - hf_config, "blocksparse_vert_stride", None) - kwargs['dense_attention_every_n_layers'] = getattr( - hf_config, "dense_attention_every_n_layers", None) - kwargs['norm_epsilon'] = hf_config.layer_norm_epsilon - else: - kwargs['rotary_base'] = get_hf_rope_theta(hf_config, 10000.0) - kwargs['norm_epsilon'] = hf_config.rms_norm_eps - moe_variant = hf_config.architectures[0] == "PhiMoEForCausalLM" - if moe_variant: - kwargs.update({ - 'moe': { - 'num_experts': hf_config.num_local_experts, - 'top_k': hf_config.num_experts_per_tok, - 'normalization_mode': - MoeConfig.ExpertScaleNormalizationMode.SPARSE_MIXER, - 'sparse_mixer_epsilon': hf_config.router_jitter_noise, - }, - 'attention_bias': hf_config.attention_bias - }) - - kwargs['position_embedding_type'] = 'rope_gpt_neox' - if hf_config.max_position_embeddings >= 128000: - kwargs[ - 'original_max_position_embeddings'] = hf_config.original_max_position_embeddings - kwargs['position_embedding_type'] = "long_rope" - kwargs['longrope_scaling_short_factors'] = hf_config.rope_scaling[ - "short_factor"] - kwargs['longrope_scaling_long_factors'] = hf_config.rope_scaling[ - "long_factor"] - if small_variant or moe_variant: - kwargs['longrope_long_mscale'] = hf_config.rope_scaling[ - "long_mscale"] - kwargs['longrope_short_mscale'] = hf_config.rope_scaling[ - "short_mscale"] - - return cls(architecture=hf_config.architectures[0], - dtype=dtype, - num_hidden_layers=hf_config.num_hidden_layers, - num_attention_heads=hf_config.num_attention_heads, - hidden_size=hf_config.hidden_size, - intermediate_size=hf_config.intermediate_size, - num_key_value_heads=num_key_value_heads, - vocab_size=hf_config.vocab_size, - max_position_embeddings=hf_config.max_position_embeddings, - hidden_act="swiglu" - if hf_config.hidden_act == 'silu' else hf_config.hidden_act, - mapping=mapping, - quantization=quant_config, - **kwargs) diff --git a/tensorrt_llm/models/phi3/convert.py b/tensorrt_llm/models/phi3/convert.py deleted file mode 100644 index 1f11408b6940..000000000000 --- a/tensorrt_llm/models/phi3/convert.py +++ /dev/null @@ -1,134 +0,0 @@ -import torch - -from ..._utils import str_dtype_to_torch -from .split_weights import shuffle_qkv_weights, split_weights_tp - - -def load_weights_from_hf_model(hf_model, config): - torch_dtype = str_dtype_to_torch(config.dtype) - hf_state_dict = hf_model.state_dict() - weights = {} - - config.quant_mode.is_weight_only() - if config.quant_mode.is_int8_weight_only(): - torch.int8 - elif config.quant_mode.is_int4_weight_only(): - torch.quint4x2 - # replace key name - for key, value in hf_state_dict.items(): - # Decoder Layers - orig_key = key - if "model.layers." in key: - key = key.replace("model.layers.", "transformer.layers.") - #Attention - key = key.replace("self_attn.", "attention.") - key = key.replace("query_key_value.", "qkv.") # small - key = key.replace("Wqkv.weight", "qkv.weight") - key = key.replace("qkv_proj.", "qkv.") #128k - #MLP - key = key.replace("mlp.fc1.", "mlp.fc.") - key = key.replace("mlp.fc2.", "mlp.proj.") - key = key.replace("mlp.gate_up_proj.", "mlp.fc.") - key = key.replace("mlp.up_proj.", "mlp.fc." if config.architecture - == 'Phi3SmallForCausalLM' else "mlp.gate.") #128k - key = key.replace("mlp.down_proj.", "mlp.proj.") #128k - key = key.replace("mlp.gate_proj.", "mlp.fc.") #128k - key = key.replace("o_proj.", "dense.") #128k - # LoRA - key = key.replace("base_layer.", "") - - #MoE - key = key.replace("block_sparse_moe.gate", "mlp.router") - key = key.replace("block_sparse_moe.experts.0.w3", "mlp.fc") - key = key.replace("block_sparse_moe.experts.0.w2", "mlp.proj") - - #Layer norm - key = key.replace("post_attention_layernorm.", - "post_layernorm.") #128k - - # Embedding - key = key.replace("model.embed_tokens.weight", - "transformer.vocab_embedding.weight") - # Final Layer norm - key = key.replace("model.final_layernorm.", "transformer.ln_f.") - key = key.replace("model.norm.", "transformer.ln_f.") #128k - - if "mlp.gate_up_proj." in orig_key: #4k - original_weights = value.contiguous().clone() - half_split = original_weights.shape[0] // 2 - first_half, second_half = original_weights[: - half_split, :], original_weights[ - half_split:, :] - # Swap the halves - value = torch.cat((second_half, first_half), dim=0) - - if config.architecture == "PhiMoEForCausalLM": - num_experts = config.moe["num_experts"] - mlp_hidden_size = config.intermediate_size - num_hidden = config.hidden_size - rank_experts = list(range(num_experts)) - if config.mapping.has_moe_ep(): - rank_experts = config.mapping.ep_experts(num_experts) - - def get_moe_weight(key, suffix): - param = [] - for expert in rank_experts: - name = key.replace(f"0.{suffix}", f"{expert}.{suffix}") - fc_value = hf_state_dict[name] - param.append(fc_value) - w = torch.stack(param) - return w.reshape(-1, mlp_hidden_size, num_hidden) - - if ".0.w3" in orig_key: - w3 = get_moe_weight(orig_key, 'w3') - w1 = get_moe_weight(orig_key.replace("w3", "w1"), 'w1') - value = torch.concat([w3, w1], dim=-2) - elif ".0.w2" in orig_key: - w2 = get_moe_weight(orig_key, 'w2') - value = w2.reshape(-1, num_hidden, mlp_hidden_size) - elif any([k in orig_key for k in ["w1", "w2", "w3"]]): - continue - - if "q_proj" in key: #128k - q_param = value - k_param = hf_state_dict[orig_key.replace("q_proj", "k_proj")] - v_param = hf_state_dict[orig_key.replace("q_proj", "v_proj")] - value = torch.cat([q_param, k_param, v_param], dim=0) - key = key.replace("q_proj", "qkv") - elif "k_proj" in key or "v_proj" in key: - continue - - dtype = torch.float if "router" in key else torch_dtype - weights[key] = value.to(dtype).cpu() - - #This is for InternVL-4B - if config.architecture == 'Phi3ForCausalLM': - keys_to_rename = [ - key for key in weights.keys() if 'language_model.' in key - ] - keys_to_delete = [ - key for key in weights.keys() if 'vision_model.' in key - ] - for key in keys_to_rename: - keys_rename = key.replace('language_model.', '') - weights[keys_rename] = weights[key] - del weights[key] - for key in keys_to_delete: - del weights[key] - - if config.tie_word_embeddings or config.architecture == 'Phi3SmallForCausalLM': - weights['lm_head.weight'] = weights[ - 'transformer.vocab_embedding.weight'].clone() - - if config.architecture == 'Phi3SmallForCausalLM': - # Transform QKV weights from custom Phi3Small format to TRT-LLM format - for key, value in weights.items(): - if "qkv." in key: - weights[key] = shuffle_qkv_weights(weights[key], config) - - if config.architecture in [ - 'Phi3SmallForCausalLM', "PhiMoEForCausalLM", "Phi3ForCausalLM" - ] and config.mapping.has_tp(): - weights = split_weights_tp(config, weights, torch_dtype) - - return weights diff --git a/tensorrt_llm/models/phi3/model.py b/tensorrt_llm/models/phi3/model.py deleted file mode 100644 index b8b1a3283848..000000000000 --- a/tensorrt_llm/models/phi3/model.py +++ /dev/null @@ -1,316 +0,0 @@ -from typing import Optional, Union - -import numpy as np -from transformers import AutoModelForCausalLM - -from ..._utils import pad_vocab_size -from ...functional import PositionEmbeddingType, Tensor -from ...layers import (MLP, MOE, Attention, AttentionMaskType, - BlockSparseAttnParams, ColumnLinear, Embedding, - LayerNorm, MoeConfig, RmsNorm) -from ...lora_helper import LoraConfig, use_lora -from ...mapping import Mapping -from ...module import Module -from ..modeling_utils import (DecoderLayerList, DecoderModelForCausalLM, - PretrainedConfig, QuantConfig) -from .config import Phi3Config -from .convert import load_weights_from_hf_model - - -class Phi3DecoderLayer(Module): - - def __init__(self, config: PretrainedConfig, layer_idx: int): - super().__init__() - self.config = config - self.layer_idx = layer_idx - tp_group = config.mapping.tp_group - tp_size = config.mapping.tp_size - - attention_mask_type = AttentionMaskType.causal - block_sparse_attn_params = BlockSparseAttnParams() - q_scaling = 1.0 - self.gegelu_limit = None - - self.small_variant = config.architecture == "Phi3SmallForCausalLM" - self.moe_variant = config.architecture == "PhiMoEForCausalLM" - if self.small_variant: - self.gegelu_limit = config.gegelu_limit - - # MuP uses norm_factor=attention_head_size (rather than sqrt(attention_head_size)) - # We achieve this using q_scaling = sqrt(attention_head_size) - hidden_size = config.hidden_size - num_attention_heads = config.num_attention_heads - attention_head_size = hidden_size / num_attention_heads - q_scaling = attention_head_size**.5 - - block_sparse = ((layer_idx + 1) % - config.dense_attention_every_n_layers) != 0 - attention_mask_type = AttentionMaskType.blocksparse if block_sparse else AttentionMaskType.causal - - block_sparse_attn_params = BlockSparseAttnParams( - config.blocksparse_block_size, - config.blocksparse_homo_head_pattern, - config.blocksparse_num_local_blocks, - config.blocksparse_vertical_stride) - - if self.small_variant or self.moe_variant: - self.input_layernorm = LayerNorm( - normalized_shape=config.hidden_size, - dtype=config.dtype, - eps=config.norm_epsilon) - self.post_layernorm = LayerNorm(normalized_shape=config.hidden_size, - dtype=config.dtype, - eps=config.norm_epsilon) - else: - self.input_layernorm = RmsNorm(normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - dtype=config.dtype) - self.post_layernorm = RmsNorm(normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - dtype=config.dtype) - - layers_range = config.mapping.pp_layers(config.num_hidden_layers) - local_layer_idx = layer_idx - layers_range[0] - position_embedding_type = PositionEmbeddingType.rope_gpt_neox - - rope_scaling_short_factors, rope_scaling_long_factors = None, None - rope_scaling_short_mscale, rope_scaling_long_mscale = None, None - original_max_position_embeddings = config.max_position_embeddings - - if hasattr(config, "longrope_scaling_short_factors"): - rope_scaling_short_factors = np.asarray( - config.longrope_scaling_short_factors).astype(np.float32) - rope_scaling_long_factors = np.asarray( - config.longrope_scaling_long_factors).astype(np.float32) - - original_max_position_embeddings = config.original_max_position_embeddings - position_embedding_type = PositionEmbeddingType.long_rope - - if self.small_variant or self.moe_variant: - rope_scaling_short_mscale = config.longrope_short_mscale - rope_scaling_long_mscale = config.longrope_long_mscale - - self.attention = Attention( - local_layer_idx=local_layer_idx, - hidden_size=config.hidden_size, - num_attention_heads=config.num_attention_heads, - num_kv_heads=config.num_key_value_heads, - position_embedding_type=position_embedding_type, - rotary_embedding_base=config.rotary_base, - max_position_embeddings=config.max_position_embeddings, - dtype=config.dtype, - attention_mask_type=attention_mask_type, - bias=self.small_variant or self.moe_variant, - q_scaling=q_scaling, - tp_group=tp_group, - tp_size=tp_size, - quant_mode=config.quant_mode, - rope_scaling_short_factors=rope_scaling_short_factors, - rope_scaling_long_factors=rope_scaling_long_factors, - rope_scaling_short_mscale=rope_scaling_short_mscale, - rope_scaling_long_mscale=rope_scaling_long_mscale, - original_max_position_embeddings=original_max_position_embeddings, - rotary_embedding_percentage=config.rotary_pct, - block_sparse_params=block_sparse_attn_params) - - ClsMLP = MLP - mlp_kwargs = {} - if hasattr(config, "moe"): - ClsMLP = MOE - moe_config = MoeConfig() - for key, value in config.moe.items(): - setattr(moe_config, key, value) - mlp_kwargs = { - "moe_config": moe_config, - "mapping": config.mapping, - } - - self.mlp = ClsMLP(hidden_size=config.hidden_size, - ffn_hidden_size=config.intermediate_size, - hidden_act=config.hidden_act, - dtype=config.dtype, - tp_group=tp_group, - tp_size=tp_size, - quant_mode=config.quant_mode, - bias=self.small_variant, - **mlp_kwargs) - - def forward( - self, - hidden_states: Tensor, - attention_mask=None, - use_cache=False, - kv_cache_params=None, - attention_params=None, - lora_layer_params=None, - ): - - input_layernorm_output = self.input_layernorm(hidden_states) - attention_output = self.attention( - input_layernorm_output, - attention_mask=attention_mask, - use_cache=use_cache, - kv_cache_params=kv_cache_params, - attention_params=attention_params, - norm_before_bmm1=not self.small_variant, - lora_layer_params=lora_layer_params, - ) - - if use_cache: - attention_output, presents = attention_output - - post_attention_input = hidden_states + attention_output - post_attention_output = self.post_layernorm(post_attention_input) - if self.small_variant: - feed_forward_hidden_states = self.mlp( - post_attention_output, - gegelu_limit=self.gegelu_limit, - lora_layer_params=lora_layer_params) - else: - feed_forward_hidden_states = self.mlp( - post_attention_output, lora_layer_params=lora_layer_params) - hidden_states = post_attention_input + feed_forward_hidden_states - if use_cache: - return (hidden_states, presents) - return hidden_states - - -class Phi3Model(Module): - - def __init__(self, config: PretrainedConfig): - super().__init__() - self.vocab_embedding = Embedding(num_embeddings=config.vocab_size, - embedding_dim=config.hidden_size, - dtype=config.dtype) - - self.layers = DecoderLayerList(Phi3DecoderLayer, config) - self.small_variant = config.architecture == "Phi3SmallForCausalLM" - self.moe_variant = config.architecture == "PhiMoEForCausalLM" - if self.small_variant or self.moe_variant: - self.ln_f = LayerNorm(normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - dtype=config.dtype) - if self.small_variant: - self.mup_embedding_multiplier = config.mup_embedding_multiplier - else: - self.ln_f = RmsNorm(normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - dtype=config.dtype) - - def forward( - self, - input_ids: Tensor, - position_ids=None, - use_cache=False, - attention_mask=None, - kv_cache_params=None, - attention_params=None, - prompt_embedding_table=None, - prompt_tasks=None, - prompt_vocab_size=None, - lora_params=None, - ): - args = [prompt_embedding_table, prompt_tasks, prompt_vocab_size - ] if prompt_embedding_table is not None else [] - hidden_states = self.vocab_embedding(input_ids, *args) - - if self.small_variant and self.mup_embedding_multiplier > 0.0: - hidden_states = hidden_states * self.mup_embedding_multiplier - - hidden_states = self.layers( - hidden_states, - use_cache=use_cache, - attention_mask=attention_mask, - kv_cache_params=kv_cache_params, - attention_params=attention_params, - lora_params=lora_params, - ) - if use_cache: - hidden_states, presents = hidden_states - - hidden_states = self.ln_f(hidden_states) - - if use_cache: - return (hidden_states, tuple(presents)) - return hidden_states - - -class Phi3ForCausalLM(DecoderModelForCausalLM): - config_class = Phi3Config - - def __init__(self, config: PretrainedConfig): - transformer = Phi3Model(config) - vocab_size_padded = pad_vocab_size(config.vocab_size, - config.mapping.tp_size) - - self.moe_variant = config.architecture == "PhiMoEForCausalLM" - lm_head = ColumnLinear(config.hidden_size, - vocab_size_padded, - bias=self.moe_variant, - dtype=config.dtype, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size, - gather_output=True) - - if self.moe_variant: - self.trtllm_modules_to_hf_modules = { - "attn_q": "q_proj", - "attn_k": "k_proj", - "attn_v": "v_proj", - "attn_dense": "o_proj", - "moe_h_to_4h": "w1", - "moe_4h_to_h": "w2", - "moe_gate": "w3", - "moe_router": "gate", - } - else: - self.trtllm_modules_to_hf_modules = { - "attn_qkv": ["qkv_proj", "query_key_value"], - "attn_dense": ["o_proj", "dense"], - "mlp_h_to_4h": ["gate_up_proj", "up_proj"], - "mlp_4h_to_h": "down_proj", - } - - super().__init__(config, transformer, lm_head) - - @classmethod - def from_hugging_face( - cls, - hf_model_or_dir: Union[str, 'transformers.PreTrainedModel'], - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - import transformers - - assert hf_model_or_dir is not None - use_preloading = isinstance(hf_model_or_dir, - transformers.PreTrainedModel) - if use_preloading: - hf_model = hf_model_or_dir - hf_config_or_dir = hf_model.config - else: - hf_model_dir = hf_model_or_dir - hf_config_or_dir = hf_model_or_dir - config = Phi3Config.from_hugging_face(hf_config_or_dir, - dtype=dtype, - mapping=mapping, - quant_config=quant_config, - **kwargs) - - if not use_preloading: - trust_remote_code = kwargs.pop('trust_remote_code', True) - - hf_model = AutoModelForCausalLM.from_pretrained( - hf_model_dir, dtype="auto", trust_remote_code=trust_remote_code) - - assert isinstance(hf_model, transformers.PreTrainedModel) - - weights = load_weights_from_hf_model(hf_model, config) - - model = cls(config) - model.load(weights) - return model - - def use_lora(self, lora_config: LoraConfig): - use_lora(self, lora_config, self.trtllm_modules_to_hf_modules) diff --git a/tensorrt_llm/models/phi3/split_weights.py b/tensorrt_llm/models/phi3/split_weights.py deleted file mode 100644 index bca33ba55116..000000000000 --- a/tensorrt_llm/models/phi3/split_weights.py +++ /dev/null @@ -1,218 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import torch - -from tensorrt_llm.models.convert_utils import (get_weight, get_weight_and_bias, - split, split_matrix_tp, - split_qkv_bias_tp, split_qkv_tp) - -from ..._utils import pad_vocab_size - - -def shuffle_qkv_weights(weights, config): - # Input weights are organized as - # (q00, q01, ... q0m, k0, v0), (q10, q11, ... q1m, k1, v1), ... (qn0, qn1, ... qnm, kn, vn) - # where n = num_kv_heads, m = num_attention_heads // num_kv_heads (i.e. #q_heads per kv_head) - # - # Output weights will be organized as - # (q00, q01, ..., qnm), (k0, k1, .., kn), (v0, v1, .., vn) - - num_heads = config.num_attention_heads - num_kv_heads = config.num_key_value_heads - num_q_per_kv = num_heads // num_kv_heads - - hidden_size = config.hidden_size - head_dim = hidden_size // num_heads - - input_shape = weights.shape - if weights.dim() < 2: - weights = weights.unsqueeze(1) - - weights = weights.reshape(num_kv_heads, (num_q_per_kv + 2), head_dim, - weights.shape[-1]) - q = weights[:, :-2, :, :] - k = weights[:, -2, :, :] - v = weights[:, -1, :, :] - - # num_heads x head_dim x hidden_size - q = q.reshape(-1, q.shape[2], q.shape[3]) - - # num_heads + (2 * num_kv_heads) x head_dim x hidden_size - weights = torch.cat([q, k, v], dim=0) - weights = weights.reshape(-1, weights.shape[2]) - - weights = weights.squeeze() - assert input_shape == weights.shape - - return weights - - -def split_embedding( - param: torch.Tensor, - tp_size: int, - tp_rank: int, - use_parallel_embedding: bool = False, - sharding_dim: int = 0, -) -> torch.Tensor: - if param is None: - return None - if not use_parallel_embedding: - return param - - vocab_size, hidden_size = param.size() - if sharding_dim == 0: - if vocab_size % tp_size != 0: - vocab_size_padded = pad_vocab_size(vocab_size, tp_size) - pad_width = vocab_size_padded - vocab_size - param = torch.nn.functional.pad(param, (0, 0, 0, pad_width), - value=0) - else: - assert hidden_size % tp_size == 0 - return split(param, tp_size, tp_rank, dim=sharding_dim) - - -def get_tllm_linear_weight(weight, - prefix, - bias=None, - use_weight_only=False, - plugin_weight_only_quant_type=torch.int8): - results = {} - if use_weight_only: - v = weight.t().contiguous() - processed_torch_weights, torch_weight_scales = \ - torch.ops.trtllm.symmetric_quantize_last_axis_of_batched_matrix( - v, plugin_weight_only_quant_type) - results[prefix + '.weight'] = processed_torch_weights - results[prefix + '.per_channel_scale'] = torch_weight_scales - else: - results[prefix + '.weight'] = weight.contiguous() - - if bias is not None: - results[prefix + '.bias'] = bias - - return results - - -def split_weights_tp(config, weights, dtype): - num_heads = config.num_attention_heads - num_kv_heads = config.num_key_value_heads - hidden_size = config.hidden_size - moe_variant = config.architecture == "PhiMoEForCausalLM" - - mha_mode = num_heads == num_kv_heads - tp_size = config.mapping.tp_size - rank = config.mapping.tp_rank - moe_tp_size = config.mapping.moe_tp_size - moe_tp_rank = config.mapping.moe_tp_rank - use_weight_only = config.quant_mode.is_weight_only() - plugin_weight_only_quant_type = None - if use_weight_only and config.quant_mode.is_int8_weight_only() == 'int8': - plugin_weight_only_quant_type = torch.int8 - elif use_weight_only and config.quant_mode.is_int4_weight_only() == 'int4': - plugin_weight_only_quant_type = torch.quint4x2 - - def get_quant_weight(weight, prefix, bias): - return get_tllm_linear_weight(weight, prefix, bias, use_weight_only, - plugin_weight_only_quant_type) - - for layer_id in range(config.num_hidden_layers): - layer_prefix = f"transformer.layers.{layer_id}." - - prefix = layer_prefix + 'attention.qkv' - qkv_weight, qkv_bias = get_weight_and_bias(weights, prefix, dtype) - - if not mha_mode: - num_q_per_kv = num_heads // num_kv_heads - - qkv_weight = qkv_weight.reshape(num_q_per_kv + 2, -1, hidden_size) - q = qkv_weight[:num_q_per_kv, :, :].reshape(-1, hidden_size) - k = qkv_weight[num_q_per_kv:num_q_per_kv + 1, :, :].reshape( - -1, hidden_size) - v = qkv_weight[num_q_per_kv + 1:num_q_per_kv + 2, :, :].reshape( - -1, hidden_size) - split_weight = torch.cat( - [split(x, tp_size, rank) for x in [q, k, v]], dim=0) - if qkv_bias is not None: - qkv_bias = qkv_bias.reshape(num_q_per_kv + 2, -1) - q = qkv_bias[:num_q_per_kv, :].reshape(-1) - k = qkv_bias[num_q_per_kv:num_q_per_kv + 1, :].reshape(-1) - v = qkv_bias[num_q_per_kv + 1:num_q_per_kv + 2, :].reshape(-1) - split_bias = torch.cat( - [split(x, tp_size, rank) for x in [q, k, v]], dim=0) - else: - split_bias = None - else: - split_weight = split_qkv_tp(qkv_weight, num_heads, hidden_size, - tp_size, rank) - if qkv_bias is not None: - split_bias = split_qkv_bias_tp(qkv_bias, num_heads, hidden_size, - tp_size, rank) - else: - split_bias = None - weights.update(get_quant_weight(split_weight, prefix, split_bias)) - - prefix = layer_prefix + 'attention.dense' - attn_dense_weight, attn_dense_bias = get_weight_and_bias( - weights, prefix, dtype) - split_v = split_matrix_tp(attn_dense_weight, tp_size, rank, dim=1) - weights.update(get_quant_weight(split_v, prefix, attn_dense_bias)) - - prefix = layer_prefix + 'mlp.fc' - if not moe_variant: - mlp_fc_weight, mlp_fc_bias = get_weight_and_bias( - weights, prefix, dtype) - split_v = split_matrix_tp(mlp_fc_weight, tp_size, rank, dim=0) - if mlp_fc_bias is not None: - bias = split_matrix_tp(mlp_fc_bias, tp_size, rank, dim=0) - else: - bias = None - weights.update(get_quant_weight(split_v, prefix, bias)) - else: - mlp_fc_weight = get_weight(weights, prefix, dtype) - w3 = split_matrix_tp(mlp_fc_weight, 2, 0, dim=1) - split_w3 = split_matrix_tp(w3, moe_tp_size, moe_tp_rank, dim=1) - w1 = split_matrix_tp(mlp_fc_weight, 2, 1, dim=1) - split_w1 = split_matrix_tp(w1, moe_tp_size, moe_tp_rank, dim=1) - split_v = torch.concat([split_w3, split_w1], dim=-2) - weights.update(get_quant_weight(split_v, prefix, None)) - - prefix = layer_prefix + 'mlp.proj' - if not moe_variant: - mlp_proj_weight, mlp_proj_bias = get_weight_and_bias( - weights, prefix, dtype) - split_v = split_matrix_tp(mlp_proj_weight, tp_size, rank, dim=1) - weights.update(get_quant_weight(split_v, prefix, mlp_proj_bias)) - else: - mlp_proj_weight = get_weight(weights, prefix, dtype) - split_v = split_matrix_tp(mlp_proj_weight, - moe_tp_size, - moe_tp_rank, - dim=2) - weights.update(get_quant_weight(split_v, prefix, None)) - - weights['transformer.vocab_embedding.weight'] = split_embedding( - weights['transformer.vocab_embedding.weight'], tp_size, rank) - weights['lm_head.weight'] = split_matrix_tp(weights['lm_head.weight'], - tp_size, - rank, - dim=0) - if moe_variant: - weights['lm_head.bias'] = split_matrix_tp(weights['lm_head.bias'], - tp_size, - rank, - dim=0) - - return weights diff --git a/tensorrt_llm/models/qwen/__init__.py b/tensorrt_llm/models/qwen/__init__.py deleted file mode 100644 index 71bf6d298d42..000000000000 --- a/tensorrt_llm/models/qwen/__init__.py +++ /dev/null @@ -1,14 +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. diff --git a/tensorrt_llm/models/qwen/config.py b/tensorrt_llm/models/qwen/config.py deleted file mode 100644 index 0f1bd34606be..000000000000 --- a/tensorrt_llm/models/qwen/config.py +++ /dev/null @@ -1,211 +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. -from typing import Optional, Union - -from ..._utils import get_hf_rope_theta -from ...layers import MoeConfig -from ...mapping import Mapping -from ..convert_utils import infer_dtype -from ..modeling_utils import PretrainedConfig, QuantConfig - - -class QWenConfig(PretrainedConfig): - - def __init__(self, - *, - mlp_bias: bool = False, - attn_bias: bool = True, - rotary_base: float = 10000.0, - rotary_scaling: Optional[dict] = None, - disable_weight_only_quant_plugin: bool = False, - use_logn_attn: bool = False, - moe: Optional[Union[MoeConfig, dict]] = None, - num_labels: int = 1, - mlp_only_layers: Optional[list] = None, - decoder_sparse_step: int = 1, - **kwargs): - self.mlp_bias = mlp_bias - self.attn_bias = attn_bias - self.rotary_base = rotary_base - self.rotary_scaling = rotary_scaling - self.disable_weight_only_quant_plugin = disable_weight_only_quant_plugin - self.num_labels = num_labels - self.use_logn_attn = use_logn_attn - self.mlp_only_layers = mlp_only_layers or [] - self.decoder_sparse_step = decoder_sparse_step - if moe is None: - # Legacy MOE config fields - moe = MoeConfig(num_experts=kwargs.pop('moe_num_experts', 0), - top_k=kwargs.pop('moe_top_k', 0), - normalization_mode=kwargs.pop( - 'moe_normalization_mode', - MoeConfig.ExpertScaleNormalizationMode.NONE)) - elif isinstance(moe, dict): - moe = MoeConfig.from_dict(moe) - assert isinstance(moe, MoeConfig) - self.moe = moe.validate() - - super().__init__(**kwargs) - - def to_dict(self): - output = super().to_dict() - # Serialize the fields added in QWenConfig - output['mlp_bias'] = self.mlp_bias - output['attn_bias'] = self.attn_bias - output['rotary_base'] = self.rotary_base - output['rotary_scaling'] = self.rotary_scaling - output[ - 'disable_weight_only_quant_plugin'] = self.disable_weight_only_quant_plugin - output['use_logn_attn'] = self.use_logn_attn - output['mlp_only_layers'] = self.mlp_only_layers - output['decoder_sparse_step'] = self.decoder_sparse_step - output['moe'] = self.moe.to_dict() - return output - - @classmethod - def from_hugging_face(cls, - hf_config_or_dir: Union[ - str, 'transformers.PretrainedConfig'], - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs) -> "QWenConfig": - import transformers - trust_remote_code = kwargs.pop('trust_remote_code', True) - - if isinstance(hf_config_or_dir, transformers.PretrainedConfig): - hf_config = hf_config_or_dir - else: - hf_config_dir = str(hf_config_or_dir) - - hf_config = transformers.AutoConfig.from_pretrained( - hf_config_dir, trust_remote_code=trust_remote_code) - if hasattr(hf_config, 'llm_config'): - hf_config = hf_config.llm_config - - qwen_type = hf_config.model_type - # lmms llava onevision qwen - if qwen_type == 'llava': - qwen_type = 'qwen2' - if hf_config.architectures and hf_config.architectures[ - 0] == 'LlavaQwenForCausalLM': - hf_config.architectures[0] = 'Qwen2ForCausalLM' - # hf llava onevision qwen - if qwen_type == 'llava_onevision': - hf_config = hf_config.text_config - qwen_type = f'{hf_config.model_type}_llava_onevision' - # Qwen2-Audio - if qwen_type == 'qwen2_audio': - hf_config = hf_config.text_config - hf_config.architectures = ['Qwen2ForCausalLM'] - - valid_types = ('qwen', 'qwen2', 'qwen2_moe', 'qwen2_llava_onevision', - 'qwen2_vl', 'qwen2_audio', 'qwen3', 'qwen3_moe') - assert qwen_type in valid_types, f"Unsupported Qwen type: {qwen_type}, only {valid_types} are acceptable." - num_key_value_heads = getattr(hf_config, "num_key_value_heads", - hf_config.num_attention_heads) - head_dim = getattr( - hf_config, "head_dim", - hf_config.hidden_size // hf_config.num_attention_heads) - head_size = getattr(hf_config, "kv_channels", head_dim) - hidden_act = getattr(hf_config, "hidden_act", "silu") - if qwen_type in ("qwen2_moe", "qwen3_moe"): - hidden_act = "swiglu" - - # Qwen3 models have no attention bias, while legacy models have bias - if qwen_type in ('qwen3', 'qwen3_moe'): - attn_bias = False # Qwen3 models have no attn bias - else: - attn_bias = True # Legacy Qwen models have attn bias - rotary_scaling = getattr(hf_config, "rope_scaling", None) - seq_length = getattr(hf_config, "seq_length", 8192) - use_logn_attn = getattr(hf_config, "use_logn_attn", False) - disable_weight_only_quant_plugin = kwargs.pop( - 'disable_weight_only_quant_plugin', False) - if qwen_type == "qwen": - rms_norm_eps = hf_config.layer_norm_epsilon - rotary_base = getattr(hf_config, "rotary_emb_base", 10000.0) - else: - rms_norm_eps = hf_config.rms_norm_eps - rotary_base = get_hf_rope_theta(hf_config, 100000.0) - - num_labels = 1 - if hf_config.architectures[0] == "Qwen2ForSequenceClassification": - num_labels = hf_config.num_labels - - moe_num_experts = getattr(hf_config, "num_experts", 0) - moe_top_k = getattr(hf_config, "num_experts_per_tok", 0) - moe_intermediate_size = getattr(hf_config, "moe_intermediate_size", 0) - moe_shared_expert_intermediate_size = getattr( - hf_config, "shared_expert_intermediate_size", 0) - moe_normalization_mode = MoeConfig.ExpertScaleNormalizationMode.NONE - - # Add support for mlp_only_layers and decoder_sparse_step (Qwen3 MoE) - mlp_only_layers = getattr(hf_config, "mlp_only_layers", []) - decoder_sparse_step = getattr(hf_config, "decoder_sparse_step", 1) - - moe_config = MoeConfig(num_experts=moe_num_experts, - top_k=moe_top_k, - normalization_mode=moe_normalization_mode) - moe_config.validate() - - dtype = infer_dtype(dtype, getattr(hf_config, 'torch_dtype', None)) - tie_word_embeddings = getattr(hf_config, 'tie_word_embeddings', False) - - if qwen_type == 'qwen2_vl': - pe_type = 'mrope' - rotary_embedding_percentage = getattr(hf_config, 'rotary_pct', 1.0) - rotary_embedding_dim = getattr( - hf_config, 'rotary_dim', - int(hf_config.hidden_size / hf_config.num_attention_heads * - rotary_embedding_percentage)) - rotary_scaling['type'] = 'mrope' - else: - pe_type = 'rope_gpt_neox' - rotary_embedding_dim = None - - return cls( - architecture=hf_config.architectures[0], - dtype=dtype, - num_hidden_layers=hf_config.num_hidden_layers, - num_attention_heads=hf_config.num_attention_heads, - hidden_size=hf_config.hidden_size, - intermediate_size=hf_config.intermediate_size, - num_key_value_heads=num_key_value_heads, - head_size=head_size, - vocab_size=hf_config.vocab_size, - position_embedding_type=pe_type, - max_position_embeddings=hf_config.max_position_embeddings, - rotary_embedding_dim=rotary_embedding_dim, - hidden_act=hidden_act, - norm_epsilon=rms_norm_eps, - attn_bias=attn_bias, - rotary_base=rotary_base, - rotary_scaling=rotary_scaling, - disable_weight_only_quant_plugin=disable_weight_only_quant_plugin, - seq_length=seq_length, - use_logn_attn=use_logn_attn, - qwen_type=qwen_type, - moe_intermediate_size=moe_intermediate_size, - moe_shared_expert_intermediate_size= - moe_shared_expert_intermediate_size, - mlp_only_layers=mlp_only_layers, - decoder_sparse_step=decoder_sparse_step, - moe=moe_config, - mapping=mapping, - quantization=quant_config, - num_labels=num_labels, - tie_word_embeddings=tie_word_embeddings, - **kwargs) diff --git a/tensorrt_llm/models/qwen/convert.py b/tensorrt_llm/models/qwen/convert.py deleted file mode 100644 index 23220538aed7..000000000000 --- a/tensorrt_llm/models/qwen/convert.py +++ /dev/null @@ -1,1284 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License.import functools - -import copy -import functools -import json -import os -import time -from collections import defaultdict -from typing import List, Optional - -import numpy as np -import safetensors -import torch -import torch.nn as nn -from tqdm import tqdm -from transformers import AutoConfig, AutoTokenizer -from transformers.pytorch_utils import Conv1D - -from ..._utils import pad_vocab_size, str_dtype_to_torch -from ...logger import logger -from ...mapping import Mapping -from ...quantization import QuantAlgo -from ..convert_utils import (dup_kv_bias, dup_kv_weight, generate_int8, - get_weight, get_weight_and_bias, - load_calib_dataset, smooth_gemm, - smooth_gemm_fc1_gate, split, split_matrix_tp, - split_qkv_bias_tp, split_qkv_tp) -from .config import QWenConfig -from .utils import get_qwen_key_list, make_context - - -@torch.no_grad() -def smooth_qwen_model(model, scales, alpha, qwen_qkv_para, qwen_smoother): - # Smooth the activation and weights with smoother = $\diag{s}$ - for name, module in model.named_modules(): - if not module._get_name() == "QWenBlock": - continue - # qkv_proj - layer_name = name + ".attn.c_attn" - smoother = smooth_gemm(module.attn.c_attn.weight, - scales[layer_name]["x"], module.ln_1.weight, - None, alpha) - - scales[layer_name]["x"] = scales[layer_name]["x"] / smoother - scales[layer_name]["w"] = module.attn.c_attn.weight.abs().max(dim=1)[0] - - # see transpose_weights function - qwen_qkv_para[layer_name] = module.attn.c_attn.weight.transpose( - 0, 1).contiguous() - - # ================================================================= - layer_name = name + ".attn.c_proj" - smoother = smooth_gemm( - module.attn.c_proj.weight, - scales[layer_name]["x"], - None, - None, - alpha=alpha, - ) - qwen_smoother[layer_name] = smoother.float() - - scales[layer_name]["x"] = scales[layer_name]["x"] / smoother - scales[layer_name]["w"] = module.attn.c_proj.weight.abs().max(dim=1)[0] - # ================================================================== - fc1_layer_name = name + ".mlp.w1" - gate_layer_name = name + ".mlp.w2" - - smoother = smooth_gemm_fc1_gate(module.mlp.w1.weight, - module.mlp.w2.weight, - scales[fc1_layer_name]["x"], - module.ln_2.weight, None, alpha) - - scales[fc1_layer_name]["x"] = scales[fc1_layer_name]["x"] / smoother - scales[fc1_layer_name]["w"] = module.mlp.w1.weight.abs().max(dim=1)[0] - - scales[gate_layer_name]["x"] = scales[gate_layer_name]["x"] / smoother - scales[gate_layer_name]["w"] = module.mlp.w2.weight.abs().max(dim=1)[0] - - # ================================================================== - layer_name = name + ".mlp.c_proj" - smoother = smooth_gemm(module.mlp.c_proj.weight, - scales[layer_name]["x"], None, None, alpha) - qwen_smoother[layer_name] = smoother.float() - scales[layer_name]["x"] = scales[layer_name]["x"] / smoother - scales[layer_name]["w"] = module.mlp.c_proj.weight.abs().max(dim=1)[0] - - -@torch.no_grad() -def smooth_qwen2_model(model, scales, alpha, qwen_qkv_para, qwen_smoother): - # Smooth the activation and weights with smoother = $\diag{s}$ - for name, module in model.named_modules(): - from transformers.models.qwen2.modeling_qwen2 import Qwen2DecoderLayer - from transformers.models.qwen2_vl.modeling_qwen2_vl import \ - Qwen2VLDecoderLayer - if not isinstance(module, Qwen2DecoderLayer) and not isinstance( - module, Qwen2VLDecoderLayer): - continue - # qkv_proj - layer_name_q = name + ".self_attn.q_proj" - layer_name_k = name + ".self_attn.k_proj" - layer_name_v = name + ".self_attn.v_proj" - layer_name_qkv = name + ".self_attn.qkv_proj" - - weight = torch.cat([ - module.self_attn.q_proj.weight, module.self_attn.k_proj.weight, - module.self_attn.v_proj.weight - ], - dim=0) - - smoother = smooth_gemm(weight, scales[layer_name_q]["x"], - module.input_layernorm.weight, None, alpha) - - scales[layer_name_qkv]["x"] = scales[layer_name_q]["x"] / smoother - scales[layer_name_qkv]["w"] = weight.abs().max(dim=1)[0] - scales[layer_name_qkv]["y"] = torch.cat([ - scales[layer_name_q]["y"], scales[layer_name_k]["y"], - scales[layer_name_v]["y"] - ], - dim=0) - - # see transpose_weights function - qwen_qkv_para[layer_name_qkv] = weight.transpose(0, 1).contiguous() - - # ================================================================= - layer_name = name + ".self_attn.o_proj" - smoother = smooth_gemm(module.self_attn.o_proj.weight, - scales[layer_name]["x"], None, None, alpha) - qwen_smoother[layer_name] = smoother.float() - - scales[layer_name]["x"] = scales[layer_name]["x"] / smoother - scales[layer_name]["w"] = module.self_attn.o_proj.weight.abs().max( - dim=1)[0] - - # ================================================================== - fc1_layer_name = name + ".mlp.gate_proj" - gate_layer_name = name + ".mlp.up_proj" - - smoother = smooth_gemm_fc1_gate(module.mlp.gate_proj.weight, - module.mlp.up_proj.weight, - scales[fc1_layer_name]["x"], - module.post_attention_layernorm.weight, - None, alpha) - - scales[fc1_layer_name]["x"] = scales[fc1_layer_name]["x"] / smoother - scales[fc1_layer_name]["w"] = module.mlp.gate_proj.weight.abs().max( - dim=1)[0] - - scales[gate_layer_name]["x"] = scales[gate_layer_name]["x"] / smoother - scales[gate_layer_name]["w"] = module.mlp.up_proj.weight.abs().max( - dim=1)[0] - - # ================================================================== - layer_name = name + ".mlp.down_proj" - smoother = smooth_gemm(module.mlp.down_proj.weight, - scales[layer_name]["x"], None, None, alpha) - qwen_smoother[layer_name] = smoother.float() - scales[layer_name]["x"] = scales[layer_name]["x"] / smoother - scales[layer_name]["w"] = module.mlp.down_proj.weight.abs().max( - dim=1)[0] - - scales_keys_to_rename = [ - key for key in scales.keys() if 'language_model.' in key - ] - - qwen_qkv_para_keys_to_rename = [ - key for key in qwen_qkv_para.keys() if 'language_model.' in key - ] - - qwen_smoother_keys_to_rename = [ - key for key in qwen_smoother.keys() if 'language_model.' in key - ] - - for key in scales_keys_to_rename: - scales[key.replace('language_model.', '')] = scales[key] - del scales[key] - - for key in qwen_qkv_para_keys_to_rename: - qwen_qkv_para[key.replace('language_model.', '')] = qwen_qkv_para[key] - del qwen_qkv_para[key] - - for key in qwen_smoother_keys_to_rename: - qwen_smoother[key.replace('language_model.', '')] = qwen_smoother[key] - del qwen_smoother[key] - - -@torch.no_grad() -def capture_activation_range(model, - qwen_type, - tokenizer, - dataset, - system_prompt, - chat_format, - num_samples=512, - seq_len=512): - model.eval() - device = next(model.parameters()).device - act_scales = defaultdict(lambda: {"x": None, "y": None, "w": None}) - - if qwen_type == 'qwen': - tokenizer.pad_token_id = tokenizer.im_end_id - else: - tokenizer.pad_token_id = tokenizer.eos_token_id - - def stat_tensor(name, tensor, act_scales, key): - hidden_dim = tensor.shape[-1] - tensor = tensor.view(-1, hidden_dim).abs().detach() - comming_max = torch.max(tensor, dim=0)[0].float() - - if act_scales[name][key] is None: - act_scales[name][key] = comming_max - else: - act_scales[name][key] = torch.max(act_scales[name][key], - comming_max) - - def stat_input_hook(m, x, y, name): - if isinstance(x, tuple): - x = x[0] - stat_tensor(name, x, act_scales, "x") - stat_tensor(name, y, act_scales, "y") - - if act_scales[name]["w"] is None: - act_scales[name]["w"] = m.weight.abs().clip(1e-8, - None).max(dim=1)[0] - - hooks = [] - for name, m in model.named_modules(): - if isinstance(m, nn.Linear) or isinstance(m, Conv1D): - hooks.append( - m.register_forward_hook( - functools.partial(stat_input_hook, name=name))) - - for i in tqdm(range(num_samples), desc="calibrating model"): - line = dataset[i] - line = line + ' TL;DR: ' - line = line.strip() - line = line.replace(" n't", "n't") - if qwen_type == 'qwen': - _, input_id_list = make_context(tokenizer=tokenizer, - query=line, - history=[], - system=system_prompt, - chat_format=chat_format, - max_input_length=seq_len) - line_encoded = torch.from_numpy( - np.array(input_id_list, - dtype=np.int32)).type(torch.int32).unsqueeze(0) - line_encoded = line_encoded.to(device) - else: - line_encoded = tokenizer(line, - return_tensors="pt", - max_length=seq_len, - padding=True, - truncation=True).input_ids.to(device) - model(line_encoded) - for h in hooks: - h.remove() - return act_scales - - -def get_tllm_linear_weight(weight, - prefix, - bias=None, - use_weight_only=False, - plugin_weight_only_quant_type=torch.int8, - dtype='float32', - use_gemm_woq_plugin=True, - postfix='weight', - quant_scale_name=None): - results = {} - if use_weight_only: - if weight.dim() > 2: - v = weight.transpose(1, 2).contiguous().clone() - else: - v = weight.t().contiguous().clone() - processed_torch_weights, torch_weight_scales = \ - torch.ops.trtllm.symmetric_quantize_last_axis_of_batched_matrix( - v.cpu(), plugin_weight_only_quant_type) - if not use_gemm_woq_plugin: - results[prefix + postfix] = v.to(dtype) - else: - results[prefix + postfix] = processed_torch_weights - if quant_scale_name is not None: - results[quant_scale_name] = torch_weight_scales - else: - results[prefix + 'per_channel_scale'] = torch_weight_scales - else: - results[prefix + postfix] = weight.clone() - - if bias is not None: - results[prefix + 'bias'] = bias - - return results - - -def get_tllm_linear_sq_weight(vals, - prefix, - shape, - tensor_parallel, - is_qkv=False, - per_token=False, - per_channel=False, - last_prefix=None, - bias=None, - smoother_value=None, - smoother_shape=None, - rank=0, - cat_dim=0, - multi_query_mode=False): - results = {} - - def multi_query_split(data, local_dim, head_size, tp_size, cur_rank): - - q, k, v = torch.split(data, [local_dim, head_size, head_size], dim=-1) - q_split = torch.split(q, q.shape[-1] // tp_size, dim=-1) - k_split = torch.split(k, k.shape[-1] // tp_size, dim=-1) - v_split = torch.split(v, v.shape[-1] // tp_size, dim=-1) - return [ - torch.concat((q_split[ii], k_split[ii], v_split[ii]), dim=-1) - for ii in range(tp_size) - ][cur_rank] - - col_shape = shape if (is_qkv or per_channel) else [1, 1] - - if per_token: - if per_channel: - original_weights = vals["weight.int8.col"] - else: - original_weights = vals["weight.int8"] - local_dim = original_weights.shape[0] - head_size = (original_weights.shape[1] - local_dim) // 2 - - if multi_query_mode: - cur_weights = multi_query_split(original_weights, local_dim, - head_size, tensor_parallel, rank) - else: - cur_weights = torch.chunk(original_weights, - tensor_parallel, - dim=cat_dim)[rank] - if is_qkv: - hidden_dim = cur_weights.shape[0] - cur_weights = cur_weights.reshape(hidden_dim, -1) - results[prefix + 'weight'] = cur_weights.t().contiguous() - if smoother_value is None: - results[last_prefix] = torch.from_numpy( - np.array([1.0], dtype=np.float32)) - - if per_channel: - cur_per_channel_value = vals["scale_w_quant_orig.col"] - if smoother_value is None: - if multi_query_mode: - - cur_per_channel_value = multi_query_split( - vals["scale_w_quant_orig.col"], local_dim, head_size, - tensor_parallel, rank) - else: - cur_per_channel_value = np.split( - vals["scale_w_quant_orig.col"], - tensor_parallel, - axis=cat_dim)[rank] - else: - cur_per_channel_value = vals["scale_w_quant_orig"] - if is_qkv: - if multi_query_mode: - cur_per_channel_value = multi_query_split( - vals["scale_w_quant_orig"], local_dim, head_size, - tensor_parallel, rank) - else: - cur_per_channel_value = torch.split( - vals["scale_w_quant_orig"], - tensor_parallel, - axis=cat_dim)[rank] - - results[prefix + - 'per_channel_scale'] = cur_per_channel_value.reshape(col_shape) - else: - if per_channel: - original_weights = vals["weight.int8.col"] - else: - original_weights = vals["weight.int8"] - local_dim = original_weights.shape[0] - head_size = (original_weights.shape[1] - local_dim) // 2 - - if multi_query_mode: - cur_weights = multi_query_split(original_weights, local_dim, - head_size, tensor_parallel, rank) - else: - cur_weights = torch.chunk(original_weights, - tensor_parallel, - dim=cat_dim)[rank] - if is_qkv: - hidden_dim = cur_weights.shape[0] - cur_weights = cur_weights.reshape(hidden_dim, -1) - results[prefix + 'weight'] = cur_weights.t().contiguous() - - if per_channel: - cur_per_channel_value = vals["scale_y_accum_quant.col"] - if smoother_value is None: - if multi_query_mode: - cur_per_channel_value = multi_query_split( - vals["scale_y_accum_quant.col"], local_dim, head_size, - tensor_parallel, rank) - else: - cur_per_channel_value = np.split( - vals["scale_y_accum_quant.col"], - tensor_parallel, - axis=cat_dim)[rank] - else: - cur_per_channel_value = vals["scale_y_accum_quant"] - # QKV is always per_channel - if is_qkv: - if multi_query_mode: - cur_per_channel_value = multi_query_split( - vals["scale_y_accum_quant"], local_dim, head_size, - tensor_parallel, rank) - else: - cur_per_channel_value = np.split( - vals["scale_y_accum_quant"], - tensor_parallel, - axis=cat_dim)[rank] - - results[prefix + 'per_channel_scale'] = cur_per_channel_value.reshape( - col_shape).contiguous() - - results[last_prefix] = vals['scale_x_orig_quant'].contiguous() - - results[prefix + 'act_scale'] = vals["scale_y_quant_orig"].contiguous() - - if smoother_value is not None: - cur_smoother_value = torch.split(smoother_value, - smoother_value.shape[-1] // - tensor_parallel, - dim=cat_dim)[rank] - - results[prefix + 'smoother'] = cur_smoother_value.reshape( - smoother_shape).contiguous().to(torch.float32) - - if bias is not None: - results[prefix + 'bias'] = bias - - return results - - -def load_hf_qwen(model_dir: str, load_model_on_cpu: bool = False): - config_path = os.path.join(model_dir, 'config.json') - with open(config_path, 'r') as f: - config = json.load(f) - if config['architectures'] == ['Qwen2ForSequenceClassification']: - from transformers import Qwen2ForSequenceClassification as model_cls - elif config['architectures'] == ['Qwen2VLForConditionalGeneration']: - from transformers import Qwen2VLForConditionalGeneration as model_cls - else: - from transformers import AutoModelForCausalLM as model_cls - - model = model_cls.from_pretrained( - model_dir, - device_map='auto' if not load_model_on_cpu else 'cpu', - dtype='auto', - trust_remote_code=True) - return model - - -def convert_hf_qwen(hf_model, - qwen_type, - mapping: Mapping, - vocab_size=32000, - dtype='float32', - use_parallel_embedding=False, - sharding_dim=0, - use_weight_only=False, - use_gemm_woq_plugin=False, - plugin_weight_only_quant_type=torch.int8, - use_smooth_quant=False, - per_channel=False, - per_token=False, - int8_kv_cache=False, - act_range=[], - qkv_para=[], - smoother=[], - moe_config=None): - weights = {} - tik = time.time() - tensor_parallel = mapping.tp_size - model_params = dict(hf_model.named_parameters()) - - dtype = getattr(torch, dtype) - hf_config = hf_model.config - if hasattr(hf_config, 'llm_config'): - hf_config = hf_config.llm_config - - #This is for InternVL2 - 1B - keys_to_rename = [ - key for key in model_params.keys() if 'language_model.' in key - ] - keys_to_delete = [ - key for key in model_params.keys() if 'vision_model.' in key - ] - for key in keys_to_rename: - keys_rename = key.replace('language_model.', '') - model_params[keys_rename] = model_params[key] - del model_params[key] - for key in keys_to_delete: - del model_params[key] - - num_attention_heads = hf_config.num_attention_heads - hidden_size = hf_config.hidden_size - head_size = hidden_size // num_attention_heads - if qwen_type == 'qwen': - intermediate_size = hf_config.intermediate_size // 2 # Qwen version 1 has actual intermediate_size one half of what's in hf_config - else: - intermediate_size = hf_config.intermediate_size - num_key_value_heads = hf_config.num_key_value_heads if hasattr( - hf_config, "num_key_value_heads") else num_attention_heads - mha_mode = (num_key_value_heads == num_attention_heads) - layers_range = mapping.pp_layers(hf_config.num_hidden_layers) - - layer_prefix = "transformer.h." if qwen_type == 'qwen' else "model.layers." - key_list = get_qwen_key_list(qwen_type) - - for l in layers_range: - prefix = layer_prefix + f'{l}.' - tllm_prex = f'transformer.layers.{l - layers_range[0]}.' - if qwen_type == 'qwen': - qkv_weight, qkv_bias = get_weight_and_bias(model_params, - prefix + key_list[0], - dtype) - qkv_w = split_qkv_tp(qkv_weight, num_attention_heads, hidden_size, - tensor_parallel, mapping.tp_rank) - qkv_b = split_qkv_bias_tp(qkv_bias, num_attention_heads, - hidden_size, tensor_parallel, - mapping.tp_rank) - else: - q_weight, q_bias = get_weight_and_bias( - model_params, prefix + key_list[0] + 'q_proj', dtype) - k_weight, k_bias = get_weight_and_bias( - model_params, prefix + key_list[0] + 'k_proj', dtype) - v_weight, v_bias = get_weight_and_bias( - model_params, prefix + key_list[0] + 'v_proj', dtype) - if not mha_mode: - if num_key_value_heads < tensor_parallel: - # duplicate the KV heads up to tensor_parallel - k_weight = dup_kv_weight(k_weight, num_key_value_heads, - tensor_parallel) - v_weight = dup_kv_weight(v_weight, num_key_value_heads, - tensor_parallel) - k_bias = dup_kv_bias(k_bias, num_key_value_heads, - tensor_parallel) - v_bias = dup_kv_bias(v_bias, num_key_value_heads, - tensor_parallel) - assert (k_weight.shape[0] % (mapping.tp_size * head_size)) == 0 - assert (v_weight.shape[0] % (mapping.tp_size * head_size)) == 0 - - if k_bias is not None and v_bias is not None: - assert (k_bias.shape[0] % - (mapping.tp_size * head_size)) == 0 - assert (v_bias.shape[0] % - (mapping.tp_size * head_size)) == 0 - - wq = split(q_weight, mapping.tp_size, mapping.tp_rank) - wk = split(k_weight, mapping.tp_size, mapping.tp_rank) - wv = split(v_weight, mapping.tp_size, mapping.tp_rank) - - qkv_w = torch.concat((wq, wk, wv)) - - if q_bias is not None and k_bias is not None and v_bias is not None: - bq = split(q_bias, mapping.tp_size, mapping.tp_rank) - bk = split(k_bias, mapping.tp_size, mapping.tp_rank) - bv = split(v_bias, mapping.tp_size, mapping.tp_rank) - qkv_b = torch.concat((bq, bk, bv)) - else: - qkv_b = None - else: - qkv_weight = torch.cat([q_weight, k_weight, v_weight], dim=0) - qkv_bias = torch.cat([q_bias, k_bias, v_bias], dim=0) - - qkv_w = split_qkv_tp(qkv_weight, num_attention_heads, - hidden_size, tensor_parallel, - mapping.tp_rank) - qkv_b = split_qkv_bias_tp(qkv_bias, num_attention_heads, - hidden_size, tensor_parallel, - mapping.tp_rank) - - if use_smooth_quant: - qkv_proj_key = key_list[ - 0] if qwen_type == 'qwen' else 'self_attn.qkv_proj' - qkv_weight = qkv_para[prefix + qkv_proj_key] - qkv_out_dim = qkv_weight.shape[1] - - if not mha_mode: - local_dim = qkv_weight.shape[0] - kv_hidden_size = (qkv_weight.shape[-1] - local_dim) // 2 - qkv_weight = qkv_weight.reshape(local_dim, - local_dim + 2 * kv_hidden_size) - else: - qkv_weight = qkv_weight.reshape(hidden_size, 3, hidden_size) - - int8_weights = generate_int8(qkv_weight, - act_range.get(prefix + qkv_proj_key), - is_qkv=True, - multi_query_mode=bool(not mha_mode)) - - weights.update( - get_tllm_linear_sq_weight(int8_weights, - tllm_prex + 'attention.qkv.', - [1, qkv_out_dim // tensor_parallel], - tensor_parallel, - is_qkv=True, - per_token=per_token, - per_channel=per_channel, - last_prefix=tllm_prex + - 'input_layernorm.scale_to_int', - bias=qkv_b, - smoother_value=None, - smoother_shape=None, - rank=mapping.tp_rank, - cat_dim=-1, - multi_query_mode=bool(not mha_mode))) - else: - weights.update( - get_tllm_linear_weight(qkv_w, tllm_prex + 'attention.qkv.', - qkv_b, use_weight_only, - plugin_weight_only_quant_type, dtype, - use_gemm_woq_plugin)) - if int8_kv_cache: - if qwen_type == 'qwen': - qkv_y = act_range.get(prefix + key_list[0])["y"] - else: - qkv_y = torch.cat([ - act_range.get(prefix + key_list[0] + 'q_proj')["y"], - act_range.get(prefix + key_list[0] + 'k_proj')["y"], - act_range.get(prefix + key_list[0] + 'v_proj')["y"] - ], - dim=0) - - int8_kv_scales = qkv_y.max() / 127. - - kv_cache_weights = {} - - kv_cache_weights[ - tllm_prex + - 'attention.kv_cache_scaling_factor'] = int8_kv_scales.reshape( - [1]) - - weights.update(kv_cache_weights) - - attn_dense_weight = get_weight(model_params, prefix + key_list[1], - dtype) - split_v = split_matrix_tp(attn_dense_weight, - tensor_parallel, - mapping.tp_rank, - dim=1) - if use_smooth_quant: - attn_dense_weight = attn_dense_weight.t() - int8_weights = generate_int8(attn_dense_weight, - act_range.get(prefix + key_list[1])) - weights.update( - get_tllm_linear_sq_weight( - int8_weights, - tllm_prex + 'attention.dense.', [1, hidden_size], - tensor_parallel, - is_qkv=False, - per_token=per_token, - per_channel=per_channel, - last_prefix=tllm_prex + - 'attention.quantization_scaling_factor', - smoother_value=smoother[(prefix + key_list[1])], - smoother_shape=[1, hidden_size // tensor_parallel], - rank=mapping.tp_rank, - cat_dim=0)) - else: - weights.update( - get_tllm_linear_weight(split_v, tllm_prex + 'attention.dense.', - None, use_weight_only, - plugin_weight_only_quant_type, dtype, - use_gemm_woq_plugin)) - - # Qwen3: Add q_norm and k_norm weight conversion - if qwen_type in ('qwen3', 'qwen3_moe'): - # Process q_norm.weight - q_norm_weight = get_weight(model_params, - prefix + key_list[0] + 'q_norm', dtype) - weights.update( - get_tllm_linear_weight( - q_norm_weight, - tllm_prex + 'attention.q_layernorm.', - None, - False, # LayerNorm should not be quantized - plugin_weight_only_quant_type, - dtype, - use_gemm_woq_plugin)) - - # Process k_norm.weight - k_norm_weight = get_weight(model_params, - prefix + key_list[0] + 'k_norm', dtype) - weights.update( - get_tllm_linear_weight( - k_norm_weight, - tllm_prex + 'attention.k_layernorm.', - None, - False, # LayerNorm should not be quantized - plugin_weight_only_quant_type, - dtype, - use_gemm_woq_plugin)) - - if moe_config and moe_config.has_moe(): - if qwen_type == "qwen2_moe": - # shared_expert for qwen2_moe - shared_expert_up_proj = model_params[ - f'model.layers.{l}.mlp.shared_expert.up_proj.weight'] - shared_expert_down_proj = model_params[ - f'model.layers.{l}.mlp.shared_expert.down_proj.weight'] - shared_expert_gate = model_params[ - f'model.layers.{l}.mlp.shared_expert.gate_proj.weight'] - shared_expert_up_proj = split(shared_expert_up_proj, - mapping.tp_size, - mapping.tp_rank, - dim=0) - shared_expert_down_proj = split(shared_expert_down_proj, - mapping.tp_size, - mapping.tp_rank, - dim=1) - shared_expert_gate = split(shared_expert_gate, - mapping.tp_size, - mapping.tp_rank, - dim=0) - shared_expert_gate_up_proj = torch.concat( - [shared_expert_up_proj, shared_expert_gate], - dim=-2).to(dtype) - - ## mlp.shared_expert.gate_up_proj.weight - weights.update( - get_tllm_linear_weight(shared_expert_gate_up_proj, - tllm_prex + 'mlp.shared_expert.fc.', - None, use_weight_only, - plugin_weight_only_quant_type, dtype, - use_gemm_woq_plugin)) - - ## mlp.shared_expert.down_proj.weight - weights.update( - get_tllm_linear_weight( - shared_expert_down_proj.to(dtype), - tllm_prex + 'mlp.shared_expert.proj.', None, - use_weight_only, plugin_weight_only_quant_type, dtype, - use_gemm_woq_plugin)) - - moe_shared_expert_gate_weights = get_weight( - model_params, prefix + 'mlp.shared_expert_gate', dtype) - weights.update( - get_tllm_linear_weight( - moe_shared_expert_gate_weights, - tllm_prex + 'mlp.shared_expert_gate.', - None, - False, # Router should never be quantized - plugin_weight_only_quant_type, - dtype, - use_gemm_woq_plugin)) - - ## fine-grained experts - rank_experts = list(range(moe_config.num_experts)) - if mapping.has_moe_ep(): - rank_experts = mapping.ep_experts(moe_config.num_experts) - for suffix in ["gate_proj", "down_proj", "up_proj"]: - model_params[f'model.layers.{l}.mlp.experts.{suffix}.weight'] = \ - torch.stack([model_params[f'model.layers.{l}.mlp.experts.{expert}.{suffix}.weight'].detach() - for expert in rank_experts]) - w3 = model_params[f'model.layers.{l}.mlp.experts.up_proj.weight'] - w2 = model_params[f'model.layers.{l}.mlp.experts.down_proj.weight'] - w1 = model_params[f'model.layers.{l}.mlp.experts.gate_proj.weight'] - if mapping.has_moe_tp(): - w3 = split(w3, mapping.moe_tp_size, mapping.moe_tp_rank, dim=1) - w2 = split(w2, mapping.moe_tp_size, mapping.moe_tp_rank, dim=2) - w1 = split(w1, mapping.moe_tp_size, mapping.moe_tp_rank, dim=1) - - moe_experts_w3w1_weights = torch.concat([w3, w1], dim=-2).to(dtype) - - ## mlp.experts.w2.weight - weights.update( - get_tllm_linear_weight(w2.to(dtype), tllm_prex + 'mlp.proj.', - None, use_weight_only, - plugin_weight_only_quant_type, dtype, - use_gemm_woq_plugin)) - - ## mlp.experts.w3w1.weight - weights.update( - get_tllm_linear_weight(moe_experts_w3w1_weights, - tllm_prex + 'mlp.fc.', None, - use_weight_only, - plugin_weight_only_quant_type, dtype, - use_gemm_woq_plugin)) - - moe_experts_gate_weights = get_weight(model_params, - prefix + 'mlp.gate', - torch.float32) - weights.update( - get_tllm_linear_weight( - moe_experts_gate_weights, - tllm_prex + 'mlp.router.', - None, - False, # Router should never be quantized - plugin_weight_only_quant_type, - dtype, - use_gemm_woq_plugin)) - - else: - mlp_gate_weight = get_weight(model_params, prefix + key_list[2], - dtype) - split_v = split_matrix_tp(mlp_gate_weight, - tensor_parallel, - mapping.tp_rank, - dim=0) - if use_smooth_quant: - mlp_gate_weight = mlp_gate_weight.t() - int8_weights = generate_int8( - mlp_gate_weight, act_range.get(prefix + key_list[2])) - - weights.update( - get_tllm_linear_sq_weight( - int8_weights, - tllm_prex + 'mlp.gate.', - [1, intermediate_size // tensor_parallel], - tensor_parallel, - is_qkv=False, - per_token=per_token, - per_channel=per_channel, - last_prefix=tllm_prex + 'post_layernorm.scale_to_int', - smoother_value=None, - smoother_shape=None, - rank=mapping.tp_rank, - cat_dim=-1)) - else: - weights.update( - get_tllm_linear_weight(split_v, tllm_prex + 'mlp.gate.', - None, use_weight_only, - plugin_weight_only_quant_type, dtype, - use_gemm_woq_plugin)) - - mlp_fc_weight = get_weight(model_params, prefix + key_list[3], - dtype) - split_v = split_matrix_tp(mlp_fc_weight, - tensor_parallel, - mapping.tp_rank, - dim=0) - - if use_smooth_quant: - mlp_fc_weight = mlp_fc_weight.t() #verified - int8_weights = generate_int8( - mlp_fc_weight, act_range.get(prefix + key_list[3])) - - weights.update( - get_tllm_linear_sq_weight( - int8_weights, - tllm_prex + 'mlp.fc.', - [1, intermediate_size // tensor_parallel], - tensor_parallel, - is_qkv=False, - per_token=per_token, - per_channel=per_channel, - last_prefix=tllm_prex + 'post_layernorm.scale_to_int', - smoother_value=None, - smoother_shape=None, - rank=mapping.tp_rank, - cat_dim=-1)) - else: - weights.update( - get_tllm_linear_weight(split_v, tllm_prex + 'mlp.fc.', None, - use_weight_only, - plugin_weight_only_quant_type, dtype, - use_gemm_woq_plugin)) - - mlp_proj_weight = get_weight(model_params, prefix + key_list[4], - dtype) - split_v = split_matrix_tp(mlp_proj_weight, - tensor_parallel, - mapping.tp_rank, - dim=1) - - if use_smooth_quant: - mlp_proj_weight = mlp_proj_weight.t() - int8_weights = generate_int8( - mlp_proj_weight, act_range.get(prefix + key_list[4])) - - weights.update( - get_tllm_linear_sq_weight( - int8_weights, - tllm_prex + 'mlp.proj.', [1, hidden_size], - tensor_parallel, - is_qkv=False, - per_token=per_token, - per_channel=per_channel, - last_prefix=tllm_prex + - 'mlp.quantization_scaling_factor', - smoother_value=smoother[prefix + key_list[4]], - smoother_shape=[ - 1, intermediate_size // tensor_parallel - ], - rank=mapping.tp_rank, - cat_dim=0)) - else: - weights.update( - get_tllm_linear_weight(split_v, tllm_prex + 'mlp.proj.', - None, use_weight_only, - plugin_weight_only_quant_type, dtype, - use_gemm_woq_plugin)) - - # Layer norms do not use tensor parallelism - input_ln_weight = get_weight(model_params, prefix + key_list[5], dtype) - weights[tllm_prex + 'input_layernorm.weight'] = input_ln_weight - - post_ln_weight = get_weight(model_params, prefix + key_list[6], dtype) - weights[tllm_prex + 'post_layernorm.weight'] = post_ln_weight - - v = get_weight(model_params, key_list[7], dtype) - - if mapping.is_last_pp_rank(): - if hf_config.tie_word_embeddings: - # lm_head.weight has the same weights as embedding - lm_head_weights = v.clone() - else: - lm_head_weights = get_weight(model_params, 'lm_head', dtype) - - if vocab_size % mapping.tp_size != 0: - # padding - vocab_size_padded = pad_vocab_size(vocab_size, mapping.tp_size) - pad_width = vocab_size_padded - vocab_size - - lm_head_weights = torch.from_numpy( - np.pad(lm_head_weights.detach().cpu().numpy(), - ((0, pad_width), (0, 0)), - 'constant', - constant_values=0)) - weights['lm_head.weight'] = split_matrix_tp(lm_head_weights, - tensor_parallel, - mapping.tp_rank, - dim=0) - - if use_parallel_embedding: - v = split_matrix_tp(v, - mapping.tp_size, - mapping.tp_rank, - dim=sharding_dim) - - if mapping.is_first_pp_rank(): - weights['transformer.vocab_embedding.weight'] = v - - if mapping.is_last_pp_rank(): - ln_f_w = get_weight(model_params, key_list[8], dtype) - weights['transformer.ln_f.weight'] = ln_f_w - - if hasattr(hf_model, 'score'): - score = get_weight(model_params, 'score', dtype) - weights['lm_head.weight'] = score - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Weights loaded. Total time: {t}') - return weights - - -def quantize(hf_model_dir: str, - output_dir: str, - config: QWenConfig, - calib_dataset='cnn_dailymail'): - ''' - Quantize the save the model as TRT-LLM checkpoint to output_dir - ''' - os.makedirs(output_dir, exist_ok=True) - config.to_json_file(os.path.join(output_dir, 'config.json')) - - mapping = config.mapping - assert mapping.rank == 0, "quantize should be called at rank 0 only" - - quant_config = config.quantization - use_smooth_quant = quant_config._use_plugin_sq - int8_kv_cache = quant_config.kv_cache_quant_algo == "INT8" - - assert use_smooth_quant or int8_kv_cache, "Call from_hugging_face when there is no quantization" - assert hf_model_dir is not None - ## only load and call smooth quant routine once for all ranks - hf_config = AutoConfig.from_pretrained(hf_model_dir, trust_remote_code=True) - if hf_config.architectures == ['Qwen2VLForConditionalGeneration']: - from transformers import Qwen2VLForConditionalGeneration as model_cls - else: - from transformers import AutoModelForCausalLM as model_cls - hf_model = model_cls.from_pretrained( - hf_model_dir, - device_map='auto', - dtype='auto' if not use_smooth_quant else torch.float16, - trust_remote_code=True).half() - - os.environ["TOKENIZERS_PARALLELISM"] = os.environ.get( - "TOKENIZERS_PARALLELISM", "false") - tokenizer = AutoTokenizer.from_pretrained(hf_model_dir, - trust_remote_code=True, - use_fast=False, - padding_side='left') - dataset = load_calib_dataset(calib_dataset) - - system_prompt = "You are a useful assistant, please directly output the corresponding summary according to the article entered by the user." - gen_config_path = os.path.join(hf_model_dir, 'generation_config.json') - with open(gen_config_path, 'r') as f: - gen_config = json.load(f) - chat_format = getattr(gen_config, 'chat_format', 'chatml') - act_range = capture_activation_range(hf_model, config.qwen_type, tokenizer, - dataset, system_prompt, chat_format) - qkv_para = {} - # smoother for inputs of self_attn.o_proj and mlp.down_proj - smoother = {} - if use_smooth_quant: - if config.qwen_type == 'qwen': - smooth_qwen_model(hf_model, act_range, quant_config.smoothquant_val, - qkv_para, smoother) - else: - smooth_qwen2_model(hf_model, act_range, - quant_config.smoothquant_val, qkv_para, smoother) - - for rank in range(mapping.world_size): - # To avoid changing the mapping arg in-place, also the given mapping from caller is rank agnostic, since quantize is called from only one rank - config = copy.deepcopy(config) - config.set_rank(rank) - weights = load_weights_from_hf_model(hf_model, - config=config, - act_range=act_range, - qkv_para=qkv_para, - smoother=smoother) - safetensors.torch.save_file( - weights, os.path.join(output_dir, f'rank{rank}.safetensors')) - del weights - - -def load_weights_from_hf_model(hf_model, - config: QWenConfig, - act_range: Optional[dict] = None, - qkv_para: Optional[dict] = None, - smoother: Optional[dict] = None): - #TODO: simplify the parameters here - - assert hf_model is not None - plugin_weight_only_quant_type = None # the value does not matter when use_weight_only is False - quant_algo = config.quantization.quant_algo - if quant_algo == QuantAlgo.W8A16: - plugin_weight_only_quant_type = torch.int8 - elif quant_algo == QuantAlgo.W4A16: - plugin_weight_only_quant_type = torch.quint4x2 - else: - plugin_weight_only_quant_type = None - use_gemm_woq_plugin = (not config.disable_weight_only_quant_plugin) - - mapping = config.mapping - moe_config = config.moe - - use_weight_only = quant_algo in [QuantAlgo.W8A16, QuantAlgo.W4A16] - use_smooth_quant = config.quantization._use_plugin_sq - per_channel = use_smooth_quant and 'PER_CHANNEL' in quant_algo - per_token = use_smooth_quant and 'PER_TOKEN' in quant_algo - int8_kv_cache = config.quantization.kv_cache_quant_algo == QuantAlgo.INT8 - qwen_type = config.qwen_type - weights = convert_hf_qwen( - hf_model, - qwen_type, - mapping, - vocab_size=config.vocab_size, - dtype=config.dtype, - use_weight_only=use_weight_only, - use_gemm_woq_plugin=use_gemm_woq_plugin, - plugin_weight_only_quant_type=plugin_weight_only_quant_type, - use_parallel_embedding=config.use_parallel_embedding, - sharding_dim=config.embedding_sharding_dim, - use_smooth_quant=use_smooth_quant, - per_channel=per_channel, - per_token=per_token, - int8_kv_cache=int8_kv_cache, - act_range=act_range, - qkv_para=qkv_para, - smoother=smoother, - moe_config=moe_config) - return weights - - -def load_weights_from_hf_gptq_model(hf_model, config: QWenConfig): - logger.info("loading weights from groupwise GPTQ QWen safetensors...") - weights = {} - tik = time.time() - - qwen_type = config.qwen_type - num_hidden_layers = config.num_hidden_layers - mapping = config.mapping - dtype = config.dtype - - model_params = {k: v for k, v in hf_model.state_dict().items()} - torch.cuda.empty_cache() - valid_types = ('qwen', 'qwen2', 'qwen2_vl', 'qwen3', 'qwen3_moe') - assert qwen_type in valid_types, f"Unsupported Qwen type: {qwen_type}, only {valid_types} are supported for GPTQ." - layer_prefix = "transformer.h." if qwen_type == 'qwen' else "model.layers." - key_list = get_qwen_key_list(qwen_type) - - def torch_split(v, dim): - if v.shape[dim] % mapping.tp_size != 0: - logger.error( - "Current weight shape is invalid for mapping.tp_size=" + - str(mapping.tp_size)) - assert False, "Invalid TP size" - return v.split(v.shape[dim] // mapping.tp_size, - dim=dim)[mapping.tp_rank] - - def unpack_int32_into_int8(w_packed): - # unpack inputs packed in int32/float32 into uint4 and store them in int8 format - w_packed_int4x2 = w_packed.contiguous().view(torch.uint8) - w_unpacked = torch.zeros(w_packed_int4x2.shape[0], - w_packed_int4x2.shape[1] * 2, - dtype=torch.int8) - w_unpacked[:, ::2] = w_packed_int4x2 % 16 - w_unpacked[:, 1::2] = w_packed_int4x2 // 16 - return w_unpacked.contiguous() - - def process_and_assign_weight(v: List[torch.Tensor], - tllm_prex: str, - tp_dim: int = -1): - if tp_dim == -1: - qweight_int32, qzeros_int32, scales_fp16 = [ - item.cpu() for item in v - ] - else: - qweight_int32, qzeros_int32, scales_fp16 = [ - torch_split(item, tp_dim).cpu() for item in v - ] - - USE_UINT4_INPUT = 1 # Set to true if checkpoint store UINT4 weights - USE_GPTQ_FOR_QWEN = 1 # GPTQ-for-QWEN added 1 to zeros - - qweight_unpacked_int8 = unpack_int32_into_int8( - qweight_int32.T).T.contiguous() - 8 - qweight_interleaved = preprocessor(packer(qweight_unpacked_int8), - torch.quint4x2, - torch.float16).view(torch.float16) - # zeros = zeros * scales - qzeros_unpacked_int32 = unpack_int32_into_int8(qzeros_int32) - if not USE_UINT4_INPUT: - # Correcting UINT4 values back to INT4 order - mask_negative = qzeros_unpacked_int32[qzeros_unpacked_int32 < 0] - mask_positive = qzeros_unpacked_int32[qzeros_unpacked_int32 >= 0] - qzeros_unpacked_int32 = qzeros_unpacked_int32 + 16 * mask_negative - 16 * mask_positive - zeros_x_scales_fp16 = (-qzeros_unpacked_int32 + 8 * USE_UINT4_INPUT - - USE_GPTQ_FOR_QWEN) * scales_fp16 - zeros_x_scales_fp16 = zeros_x_scales_fp16.half() - - results = { - f'{tllm_prex}.weight': qweight_interleaved, - f'{tllm_prex}.weights_scaling_factor': scales_fp16, - f'{tllm_prex}.zero': zeros_x_scales_fp16, - } - return results - - packer = torch.ops.trtllm.pack_int8_tensor_to_packed_int4 - preprocessor = torch.ops.trtllm.preprocess_weights_for_mixed_gemm - torch_dtype = str_dtype_to_torch(dtype) - - # Load weights from GPTQ checkpoint into TRT-LLM module - # 1. vocab_embedding - v = model_params[key_list[7] + '.weight'] - if mapping.is_first_pp_rank(): - weights['transformer.vocab_embedding.weight'] = v.to(torch_dtype) - - # 2. ln_f - v = model_params[key_list[8] + '.weight'] - if mapping.is_last_pp_rank(): - weights['transformer.ln_f.weight'] = v.to(torch_dtype) - - # 3. lm_head - v = model_params['lm_head.weight'] - if mapping.is_last_pp_rank(): - weights['lm_head.weight'] = torch_split(v, 0).to(torch_dtype) - - # 4. Weights inside each layer - layers_per_pipeline_stage = num_hidden_layers // mapping.pp_size - layers_range = list( - range(mapping.pp_rank * layers_per_pipeline_stage, - (mapping.pp_rank + 1) * layers_per_pipeline_stage, 1)) - suffixs = [".qweight", ".qzeros", ".scales"] - - for l in tqdm(layers_range, desc="loading weight in each layer..."): - layer_idx = l - mapping.pp_rank * layers_per_pipeline_stage - prefix = layer_prefix + str(layer_idx) + "." - tllm_prex = f'transformer.layers.{l-layers_range[0]}' - # 4.1 attention.qkv - qkv_weight_list = [] - if qwen_type == 'qwen': - for suf in suffixs: - qkv_part = model_params[prefix + key_list[0] + suf] - q_emb = qkv_part.shape[1] // 3 - model_emb = qkv_part.shape[0] - qkv_part = qkv_part.reshape(model_emb, 3, q_emb) - qkv_part = torch_split(qkv_part, 2) - qkv_part = qkv_part.reshape(model_emb, - 3 * (q_emb // mapping.tp_size)) - qkv_weight_list.append(qkv_part) - else: - for suf in suffixs: - qkv_list = [] - for comp in ["q_proj", "k_proj", "v_proj"]: - comp_part = model_params[prefix + key_list[0] + comp + suf] - comp_part = torch_split(comp_part, 1) - qkv_list.append(comp_part) - qkv_weight_list.append(torch.cat(qkv_list, dim=1)) - weights.update( - process_and_assign_weight(qkv_weight_list, - f'{tllm_prex}.attention.qkv')) - # 4.2 attention.bias - suf = ".bias" - if qwen_type == 'qwen': - qkv_bias = model_params[prefix + key_list[0] + - suf].to(torch_dtype).cpu().contiguous() - q_emb = qkv_bias.shape[0] // 3 - qkv_bias = qkv_bias.reshape(3, q_emb) - split_v = split(qkv_bias, mapping.tp_size, mapping.rank, dim=1) - qkv_bias = split_v.reshape(3 * (q_emb // mapping.tp_size)) - else: - qkv_bias_list = [] - for comp in ["q_proj", "k_proj", "v_proj"]: - comp_part = model_params[prefix + key_list[0] + comp + suf].to( - torch_dtype).cpu().contiguous() - comp_part = torch_split(comp_part, dim=0) - qkv_bias_list.append(comp_part) - qkv_bias = torch.cat(qkv_bias_list, dim=0) - weights[tllm_prex + ".attention.qkv.bias"] = qkv_bias - # 4.3 attention.dense - qkv_dense_list = [] - for suf in suffixs: - qkv_dense_part = model_params[prefix + key_list[1] + suf] - qkv_dense_list.append(qkv_dense_part) - weights.update( - process_and_assign_weight(qkv_dense_list, - f'{tllm_prex}.attention.dense', - tp_dim=0)) - # 4.4 mlp.gate - mlp_gate_list = [] - for suf in suffixs: - mlp_gate_part = model_params[prefix + key_list[2] + suf] - mlp_gate_list.append(mlp_gate_part) - weights.update( - process_and_assign_weight(mlp_gate_list, - f'{tllm_prex}.mlp.gate', - tp_dim=1)) - # 4.5 mlp.fc - mlp_fc_list = [] - for suf in suffixs: - mlp_fc_part = model_params[prefix + key_list[3] + suf] - mlp_fc_list.append(mlp_fc_part) - weights.update( - process_and_assign_weight(mlp_fc_list, - f'{tllm_prex}.mlp.fc', - tp_dim=1)) - # 4.6 mlp.proj - mlp_proj_list = [] - for suf in suffixs: - mlp_proj_part = model_params[prefix + key_list[4] + suf] - mlp_proj_list.append(mlp_proj_part) - weights.update( - process_and_assign_weight(mlp_proj_list, - f'{tllm_prex}.mlp.proj', - tp_dim=0)) - # 4.7 input_layernorm - v = model_params[prefix + key_list[5] + '.weight'] - weights[f'{tllm_prex}.input_layernorm.weight'] = v.to(torch_dtype) - # 4.8 post_layernorm - v = model_params[prefix + key_list[6] + '.weight'] - weights[f'{tllm_prex}.post_layernorm.weight'] = v.to(torch_dtype) - - tok = time.time() - t = time.strftime("%H:%M:%S", time.gmtime(tok - tik)) - logger.info(f"weights loaded. total time: {t}") - - return weights diff --git a/tensorrt_llm/models/qwen/model.py b/tensorrt_llm/models/qwen/model.py deleted file mode 100644 index 028d9e4e3e6f..000000000000 --- a/tensorrt_llm/models/qwen/model.py +++ /dev/null @@ -1,545 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import copy -import os -from typing import Optional, Union - -import torch -from tqdm import tqdm - -from ..._utils import pad_vocab_size -from ...functional import LayerNormType, Tensor, recv, send -from ...layers import (MOE, Attention, AttentionMaskType, ColumnLinear, - Embedding, GatedMLP, RmsNorm, SharedMoE) -from ...layers.moe import MOEWeightWrapper -from ...logger import logger -from ...lora_helper import (LoraConfig, - get_default_trtllm_modules_to_hf_modules, use_lora) -from ...mapping import Mapping -from ...module import Module -from ...quantization import QuantAlgo -from ..model_weights_loader import ModelWeightsLoader -from ..modeling_utils import (DecoderLayerList, DecoderModelForCausalLM, - QuantConfig) -from .config import QWenConfig -from .convert import (load_hf_qwen, load_weights_from_hf_gptq_model, - load_weights_from_hf_model) - - -class QWenDecoderLayer(Module): - - def __init__(self, config: QWenConfig, layer_idx: int): - super().__init__() - self.layer_idx = layer_idx - self.config = config - - dtype = config.dtype - self.tp_group = config.mapping.tp_group - self.tp_size = config.mapping.tp_size - - self.input_layernorm = RmsNorm(normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - dtype=dtype) - - layers_range = config.mapping.pp_layers(config.num_hidden_layers) - local_layer_idx = layer_idx - layers_range[0] - # Qwen3: Enable qk_layernorm for Q/K normalization (similar to Gemma3) - qk_layernorm = config.qwen_type in ('qwen3', 'qwen3_moe') - - self.attention = Attention( - local_layer_idx=local_layer_idx, - hidden_size=config.hidden_size, - attention_head_size=config.head_size, - num_attention_heads=config.num_attention_heads, - num_kv_heads=config.num_key_value_heads, - max_seqlen_for_logn_scaling=config.seq_length, - max_position_embeddings=config.max_position_embeddings, - dtype=dtype, - attention_mask_type=AttentionMaskType.causal, - bias=config.attn_bias, - position_embedding_type=config.position_embedding_type, - rotary_embedding_base=config.rotary_base, - rotary_embedding_scaling=config.rotary_scaling, - tp_rank=config.mapping.tp_rank, - tp_group=self.tp_group, - tp_size=self.tp_size, - cp_rank=config.mapping.cp_rank, - cp_size=config.mapping.cp_size, - cp_group=config.mapping.cp_group, - quant_mode=config.quant_mode, - use_logn_scaling=config.use_logn_attn, - dense_bias=False, - # Qwen3: Add Q/K layer normalization - qk_layernorm=qk_layernorm, - layernorm_type=LayerNormType.RmsNorm - if qk_layernorm else LayerNormType.LayerNorm) - - if config.moe.has_moe(): - mlp_kwargs = {'moe_config': config.moe, 'mapping': config.mapping} - if config.qwen_type == 'qwen2_moe': - # Qwen2 MoE uses SharedMoE with shared expert - ClsMLP = SharedMoE - mlp_kwargs['use_shared_gate'] = True - mlp_kwargs['use_side_stream'] = True - mlp_kwargs['moe_config'].shared_expert_intermediate_size = \ - config.moe_shared_expert_intermediate_size - elif config.qwen_type == 'qwen3_moe': - # Qwen3 MoE uses standard MOE without shared expert - ClsMLP = MOE - else: - ClsMLP = MOE - else: - ClsMLP = GatedMLP - mlp_kwargs = {} - - # Qwen's real inter_size depends on qwen_type - if self.config.qwen_type == 'qwen': - intermediate_size = config.intermediate_size // 2 - elif self.config.qwen_type in ('qwen2_moe', 'qwen3_moe'): - intermediate_size = config.moe_intermediate_size - else: - intermediate_size = config.intermediate_size - - self.mlp = ClsMLP(hidden_size=config.hidden_size, - ffn_hidden_size=intermediate_size, - hidden_act=config.hidden_act, - dtype=dtype, - bias=config.mlp_bias, - tp_group=self.tp_group, - tp_size=self.tp_size, - quant_mode=config.quant_mode, - **mlp_kwargs) - self.post_layernorm = RmsNorm(normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - dtype=dtype) - - def forward( - self, - hidden_states: Tensor, - attention_mask=None, - use_cache=False, - spec_decoding_params=None, - kv_cache_params=None, - attention_params=None, - lora_layer_params=None, - mrope_params=None, - ): - residual = hidden_states - hidden_states = self.input_layernorm(hidden_states) - attention_output = self.attention( - hidden_states, - attention_mask=attention_mask, - use_cache=use_cache, - spec_decoding_params=spec_decoding_params, - kv_cache_params=kv_cache_params, - attention_params=attention_params, - lora_layer_params=lora_layer_params, - mrope_params=mrope_params, - ) - if use_cache: - attention_output, presents = attention_output - - hidden_states = residual + attention_output - - residual = hidden_states - - hidden_states = self.post_layernorm(hidden_states) - - hidden_states = self.mlp(hidden_states, - lora_layer_params=lora_layer_params) - - hidden_states = residual + hidden_states - if use_cache: - return (hidden_states, presents) - return hidden_states - - -class QWenModel(Module): - - def __init__(self, config: QWenConfig) -> None: - super().__init__() - self.mapping = config.mapping - if self.mapping.is_first_pp_rank(): - self.vocab_embedding = Embedding(config.vocab_size, - config.hidden_size, - dtype=config.dtype) - - self.layers = DecoderLayerList(QWenDecoderLayer, config) - - if self.mapping.is_last_pp_rank(): - self.ln_f = RmsNorm(normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - dtype=config.dtype) - - def forward(self, - input_ids: Tensor, - position_ids=None, - use_cache=False, - spec_decoding_params=None, - attention_mask=None, - kv_cache_params=None, - attention_params=None, - mrope_params=None, - hidden_states=None, - prompt_embedding_table: Optional[Tensor] = None, - prompt_tasks: Optional[Tensor] = None, - prompt_vocab_size: Optional[Tensor] = None, - lora_params=None): - - ptuning_args = [ - prompt_embedding_table, prompt_tasks, prompt_vocab_size - ] if prompt_embedding_table is not None else [] - - if self.mapping.is_first_pp_rank(): - hidden_states = self.vocab_embedding(input_ids, *ptuning_args) - else: - hidden_states = recv(hidden_states, self.mapping.prev_pp_rank()) - - hidden_states = self.layers.forward( - hidden_states, - use_cache=use_cache, - spec_decoding_params=spec_decoding_params, - attention_mask=attention_mask, - kv_cache_params=kv_cache_params, - attention_params=attention_params, - lora_params=lora_params, - mrope_params=mrope_params) - - if use_cache: - hidden_states, presents = hidden_states - - if self.mapping.is_last_pp_rank(): - hidden_states = self.ln_f(hidden_states) - else: - hidden_states = send(hidden_states, self.mapping.next_pp_rank()) - - if use_cache: - return (hidden_states, tuple(presents)) - return hidden_states - - -class QWenForCausalLM(DecoderModelForCausalLM): - config_class = QWenConfig - - def __init__(self, config: QWenConfig): - transformer = QWenModel(config) - vocab_size_padded = pad_vocab_size(config.vocab_size, - config.mapping.tp_size) - - if config.mapping.is_last_pp_rank(): - if config.architecture == 'Qwen2ForSequenceClassification': - lm_head = ColumnLinear(config.hidden_size, - config.num_labels, - bias=False, - dtype=config.dtype, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size, - gather_output=True) - else: - lm_head = ColumnLinear(config.hidden_size, - vocab_size_padded, - bias=False, - dtype=config.dtype, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size, - gather_output=True) - else: - lm_head = None - self.quant_mode = config.quant_mode - self.mapping = config.mapping - if config.qwen_type == 'qwen': - self.trtllm_modules_to_hf_modules = { - "attn_qkv": "c_attn", - "attn_dense": "attn.c_proj", - "mlp_h_to_4h": "w2", - "mlp_4h_to_h": "mlp.c_proj", - "mlp_gate": "w1", - } - elif config.qwen_type in ('qwen2_moe', 'qwen3_moe'): - self.trtllm_modules_to_hf_modules = copy.copy( - get_default_trtllm_modules_to_hf_modules()) - # Common MoE expert mappings for both Qwen2 and Qwen3 MoE - self.trtllm_modules_to_hf_modules.update({ - "moe_h_to_4h": - "mlp.experts.gate_proj", - "moe_4h_to_h": - "mlp.experts.down_proj", - "moe_gate": - "mlp.experts.up_proj", - }) - # Qwen2 MoE additionally has shared expert - if config.qwen_type == 'qwen2_moe': - self.trtllm_modules_to_hf_modules.update({ - "mlp_h_to_4h": - "mlp.shared_expert.gate_proj", - "mlp_4h_to_h": - "mlp.shared_expert.down_proj", - "mlp_gate": - "mlp.shared_expert.up_proj", - "mlp_router": - "mlp.shared_expert_gate", - }) - else: - self.trtllm_modules_to_hf_modules = None - super().__init__(config, transformer, lm_head) - - @classmethod - def from_hugging_face( - cls, - hf_model_or_dir: Union[str, 'transformers.PreTrainedModel'], - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - **kwargs): - ''' Create a QWenForCausalLM object from give parameters - ''' - import transformers - - load_model_on_cpu = kwargs.pop('load_model_on_cpu', False) - use_autoawq = kwargs.pop('use_autoawq', False) - - assert hf_model_or_dir is not None - use_preloading = isinstance(hf_model_or_dir, - transformers.PreTrainedModel) - if use_preloading: - hf_model = hf_model_or_dir - hf_config_or_dir = hf_model.config - else: - hf_model_dir = hf_model_or_dir - hf_config_or_dir = hf_model_or_dir - - config = QWenConfig.from_hugging_face(hf_config_or_dir, - dtype=dtype, - mapping=mapping, - quant_config=quant_config, - **kwargs) - - if os.environ.get("TRTLLM_DISABLE_UNIFIED_CONVERTER") is None: - arg_dict = {"use_autoawq": True} if use_autoawq else {} - custom_dict = {} - - if config.qwen_type == "qwen": - custom_dict = { - "transformer": "transformer", - "vocab_embedding": "wte", - "ln_f": "ln_f", - "layers": "h", - "attention": "attn", - "qkv": "c_attn", - "dense": "c_proj", - "gate": "w1", - "proj": "c_proj", - "fc": "w2", - "input_layernorm": "ln_1", - "post_layernorm": "ln_2", - } - elif config.qwen_type == "qwen2_moe": - custom_dict = { - "mlp.shared_expert": "mlp.shared_expert", - "mlp.shared_expert_gate": "mlp.shared_expert_gate", - "fc": ["up_proj", "gate_proj"], - } - elif config.qwen_type == "qwen3_moe": - custom_dict = { - "fc": ["up_proj", "gate_proj"], - "q_layernorm": "q_norm", - "k_layernorm": "k_norm", - } - elif config.qwen_type in {"qwen2", "qwen2_vl" - } and config.tie_word_embeddings: - custom_dict = {"lm_head": "model.embed_tokens"} - elif config.architecture == "Qwen2ForSequenceClassification": - custom_dict = { - "lm_head": "score", - } - elif config.qwen_type == "qwen2_llava_onevision": - custom_dict = { - "transformer": "language_model.model", - "lm_head": "language_model.lm_head", - } - elif config.qwen_type == "qwen2_audio": - custom_dict = { - "transformer": "language_model.model", - "lm_head": "language_model.lm_head", - } - elif config.qwen_type == "qwen3": - custom_dict = { - "q_layernorm": "q_norm", - "k_layernorm": "k_norm", - } - loader = ModelWeightsLoader(hf_model_dir, custom_dict) - model = cls(config) - if config.qwen_type == "qwen" and model.config.mapping.has_tp(): - - def reshape_qkv(weights): - if weights is None: - return weights - mapping = model.config.mapping - unsqueeze = False - if isinstance(weights, torch.Tensor): - unsqueeze = True - weights = [weights] - - for idx, w in enumerate(weights): - if quant_config.quant_algo == QuantAlgo.W4A16_GPTQ: - w = w.reshape(-1, 3, w.shape[-1] // 3) - w = w.chunk(mapping.tp_size, 2)[mapping.tp_rank] - if w.shape[0] == 1: - weights[idx] = w.reshape(-1) - else: - weights[idx] = w.reshape(w.shape[0], -1) - else: - w = w.reshape(3, w.shape[0] // 3, -1) - w = w.chunk(mapping.tp_size, 1)[mapping.tp_rank] - if w.shape[-1] == 1: - weights[idx] = w.reshape(-1) - else: - weights[idx] = w.reshape(-1, w.shape[-1]) - if unsqueeze: - return weights[0] - else: - return weights - - loader.update_key_mapping(model) - tllm_weights = {} - for tllm_key, _ in tqdm(model.named_parameters()): - if "qkv" in tllm_key: - tllm_weights.update( - loader.load(tllm_key, - reshape_qkv, - skip_tp=True, - custom_postprocess_kwargs=arg_dict)) - else: - tllm_weights.update( - loader.load(tllm_key, - custom_postprocess_kwargs=arg_dict)) - loader.fill(tllm_weights) - elif config.qwen_type in ("qwen2_moe", "qwen3_moe"): - for tllm_key, _ in model.named_parameters(): - sub_module = model - for attr in tllm_key.split(".")[:-1]: - sub_module = getattr(sub_module, attr) - if "router" in tllm_key or isinstance( - sub_module, MOEWeightWrapper): - sub_module_dic = sub_module.tllm_to_externel_key_dict - sub_module_dic["mlp"] = "mlp" - if "fc" in sub_module_dic.keys(): - sub_module_dic["fc"] = [ - hf_keyword.replace("w1", "gate_proj") - for hf_keyword in sub_module_dic["fc"] - ] - sub_module_dic["fc"] = [ - hf_keyword.replace("w3", "up_proj") - for hf_keyword in sub_module_dic["fc"] - ] - if "proj" in sub_module_dic.keys(): - sub_module_dic["proj"] = [ - hf_keyword.replace("w2", "down_proj") - for hf_keyword in sub_module_dic["proj"] - ] - sub_module.tllm_to_externel_key_dict = sub_module_dic - - def concat_gate_up_proj(weights): - return torch.cat(weights, dim=-2) - - loader.update_key_mapping(model) - tllm_weights = {} - for tllm_key, _ in tqdm(model.named_parameters()): - if tllm_key.endswith("shared_expert.fc.weight"): - tllm_weights.update( - loader.load(tllm_key, - concat_gate_up_proj, - custom_postprocess_kwargs=arg_dict)) - else: - tllm_weights.update( - loader.load(tllm_key, - custom_postprocess_kwargs=arg_dict)) - loader.fill(tllm_weights) - else: - # For Qwen1 w/o TP, Qwen1.5 and Qwen2 w/o MoE - loader.generate_tllm_weights(model, arg_dict) - else: - if not use_preloading: - hf_model = load_hf_qwen(hf_model_dir, load_model_on_cpu) - - logger.debug(f"HuggingFace model: {hf_model}") - - model = QWenForCausalLM(config) - logger.debug(f"TensorRT LLM model: {model}") - - if quant_config.quant_algo == QuantAlgo.W4A16_GPTQ: - weights = load_weights_from_hf_gptq_model(hf_model, config) - else: - weights = load_weights_from_hf_model(hf_model, config) - model.load(weights) - return model - - def default_plugin_config(self, **kwargs): - plugin_config = super().default_plugin_config(**kwargs) - if self.quant_mode.is_int4_weight_only_per_group(): - plugin_config.weight_only_groupwise_quant_matmul_plugin = 'auto' - return plugin_config - - @classmethod - def quantize( - cls, - hf_model_dir: str, - output_dir: str, - dtype: str = 'auto', - mapping: Optional[Mapping] = None, - quant_config: Optional[QuantConfig] = None, - *, - calib_dataset='cnn_dailymail', - calib_batches=512, - calib_batch_size=1, - calib_max_seq_length=512, - random_seed=1234, - tokenizer_max_seq_length=2048, - **kwargs, - ): - if quant_config._requires_modelopt_quantization: - # modelopt quantization flow - super().quantize(hf_model_dir, - output_dir, - dtype=dtype, - mapping=mapping, - quant_config=quant_config, - calib_dataset=calib_dataset, - calib_batches=calib_batches, - calib_batch_size=calib_batch_size, - calib_max_seq_length=calib_max_seq_length, - random_seed=random_seed, - tokenizer_max_seq_length=tokenizer_max_seq_length) - elif quant_config._requires_calibration: - # non-modelopt quantization flow - from . import convert - - config = QWenConfig.from_hugging_face(hf_model_dir, - dtype=dtype, - mapping=mapping, - quant_config=quant_config, - **kwargs) - convert.quantize(hf_model_dir, - output_dir, - config=config, - calib_dataset=calib_dataset) - else: - raise ValueError( - f"The quant_config ({quant_config}) does not require calibration, try {cls.__name__}.from_hugging_face instead." - ) - - def use_lora(self, lora_config: LoraConfig): - use_lora(self, lora_config, self.trtllm_modules_to_hf_modules) diff --git a/tensorrt_llm/models/qwen/utils.py b/tensorrt_llm/models/qwen/utils.py deleted file mode 100644 index 1b484898f8d0..000000000000 --- a/tensorrt_llm/models/qwen/utils.py +++ /dev/null @@ -1,116 +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. - - -def make_context( - tokenizer, - query, - history, - system, - max_input_length, - max_window_size: int = 6144, - chat_format: str = "chatml", -): - if history is None: - history = [] - - if chat_format == "chatml": - im_start, im_end = "<|im_start|>", "<|im_end|>" - im_start_tokens = [tokenizer.im_start_id] - im_end_tokens = [tokenizer.im_end_id] - nl_tokens = tokenizer.encode("\n") - - def _tokenize_str(role, content): - return (f"{role}\n{content}", - tokenizer.encode( - role, - allowed_special=set(), - ) + nl_tokens + tokenizer.encode( - content, - allowed_special=set(), - )) - - system_text, system_tokens_part = _tokenize_str("system", system) - system_tokens = im_start_tokens + system_tokens_part + im_end_tokens - raw_text = "" - context_tokens = [] - - for turn_query, turn_response in reversed(history): - query_text, query_tokens_part = _tokenize_str("user", turn_query) - query_tokens = im_start_tokens + query_tokens_part + im_end_tokens - - response_text, response_tokens_part = _tokenize_str( - "assistant", turn_response) - response_tokens = im_start_tokens + response_tokens_part + im_end_tokens - next_context_tokens = nl_tokens + query_tokens + nl_tokens + response_tokens - prev_chat = ( - f"\n{im_start}{query_text}{im_end}\n{im_start}{response_text}{im_end}" - ) - - current_context_size = (len(system_tokens) + - len(next_context_tokens) + - len(context_tokens)) - if current_context_size < max_window_size: - context_tokens = next_context_tokens + context_tokens - raw_text = prev_chat + raw_text - else: - break - - context_tokens = system_tokens + context_tokens - raw_text = f"{im_start}{system_text}{im_end}" + raw_text - context_tokens += (nl_tokens + im_start_tokens + - _tokenize_str("user", query)[1] + im_end_tokens + - nl_tokens + im_start_tokens + - tokenizer.encode("assistant") + nl_tokens) - raw_text += f"\n{im_start}user\n{query}{im_end}\n{im_start}assistant\n" - - elif chat_format == "raw": - raw_text = query - context_tokens = tokenizer.encode(raw_text) - else: - raise NotImplementedError(f"Unknown chat format {chat_format!r}") - # truncate to max_input_length, truncate from the front - return raw_text, context_tokens[-max_input_length:] - - -def get_qwen_key_list(qwen_type): - qwen_key_list = [ - "attn.c_attn", # attention.qkv - "attn.c_proj", # attention.dense - "mlp.w1", # mlp.gate - "mlp.w2", # mlp.fc - "mlp.c_proj", # mlp.proj - "ln_1", # input_layernorm - "ln_2", # post_layernorm - "transformer.wte", # vocabulary embedding - "transformer.ln_f", # final layer norm - ] - qwen2_key_list = [ - "self_attn.", # attention.qkv - "self_attn.o_proj", # attention.dense - "mlp.up_proj", # mlp.gate - "mlp.gate_proj", # mlp.fc - "mlp.down_proj", # mlp.proj - "input_layernorm", # input_layernorm - "post_attention_layernorm", # post_layernorm - "model.embed_tokens", # vocabulary embedding - "model.norm", # final layer norm - ] - key_list = [] - if qwen_type == 'qwen': - key_list.extend(qwen_key_list) - else: - key_list.extend(qwen2_key_list) - return key_list diff --git a/tensorrt_llm/models/recurrentgemma/__init__.py b/tensorrt_llm/models/recurrentgemma/__init__.py deleted file mode 100644 index 2a36ca922710..000000000000 --- a/tensorrt_llm/models/recurrentgemma/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-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. diff --git a/tensorrt_llm/models/recurrentgemma/model.py b/tensorrt_llm/models/recurrentgemma/model.py deleted file mode 100644 index b5fc842f1359..000000000000 --- a/tensorrt_llm/models/recurrentgemma/model.py +++ /dev/null @@ -1,624 +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. -from collections import OrderedDict -from typing import List - -import tensorrt as trt - -from ..._common import default_net -from ..._utils import str_dtype_to_trt -from ...functional import (Tensor, arange, concat, expand, - gather_last_token_logits, shape, tanh, unsqueeze) -from ...layers import (Attention, AttentionMaskType, AttentionParams, - ColumnLinear, Embedding, GatedMLP, KeyValueCacheParams, - PositionEmbeddingType, Recurrent, RmsNorm) -from ...module import Module, ModuleList -from ...plugin import current_all_reduce_helper -from ..generation_mixin import GenerationMixin -from ..modeling_utils import (PretrainedConfig, PretrainedModel, - get_kv_cache_type_from_legacy) - - -class ResidualLayer(Module): - - def __init__(self, config: PretrainedConfig, layer_idx: int): - super().__init__() - layer_type_len = len(config.layer_types) - self.temporal_block_type = config.layer_types[layer_idx % - layer_type_len] - - self.input_layernorm = RmsNorm(normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - dtype=config.dtype) - - if self.temporal_block_type == 'recurrent': - self.recurrent = Recurrent(width=config.hidden_size, - lru_width=config.rnn_hidden_size, - d_conv=config.conv_kernel, - num_heads=config.num_attention_heads, - dtype=config.dtype, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size) - elif self.temporal_block_type == 'attention': - layer_types = config.layer_types * ( - (layer_idx + 1) // layer_type_len) - layer_types = layer_types + config.layer_types[0:( - (layer_idx + 1) % layer_type_len)] - attention_layer_idx = layer_types.count('attention') - 1 - - self.attention = Attention( - local_layer_idx=attention_layer_idx, - hidden_size=config.hidden_size, - num_attention_heads=config.num_attention_heads, - num_kv_heads=config.num_key_value_heads, - dtype=config.dtype, - attention_mask_type=AttentionMaskType.causal, - position_embedding_type=PositionEmbeddingType.rope_gpt_neox, - rotary_embedding_percentage=config.rotary_pct, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size, - tp_rank=config.mapping.tp_rank, - quant_mode=config.quant_mode, - bias=False, - dense_bias=True) - else: - raise ValueError( - 'RecurrentGemma only support "recurrent" and "attention" blocks.' - ) - - self.post_layernorm = RmsNorm(normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - dtype=config.dtype) - - self.mlp = GatedMLP(hidden_size=config.hidden_size, - ffn_hidden_size=config.intermediate_size, - hidden_act=config.hidden_act, - dtype=config.dtype, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size, - quant_mode=config.quant_mode) - - def forward(self, - hidden_states, - use_cache=False, - attention_mask=None, - kv_cache_params=None, - attention_params=None, - conv_state=None, - lru_state=None, - host_request_types=None, - last_token_ids=None, - host_context_lengths=None, - slot_mapping=None, - conv_indices=None): - - residual = hidden_states - - hidden_states = self.input_layernorm(hidden_states) - - if self.temporal_block_type == 'recurrent': - temporal_output, present_conv, present_lru = self.recurrent( - hidden_states, - conv_state=conv_state, - lru_state=lru_state, - host_request_types=host_request_types, - last_token_ids=last_token_ids, - host_context_lengths=host_context_lengths, - slot_mapping=slot_mapping, - conv_indices=conv_indices, - ) - else: - present_conv, present_lru = None, None - - if self.temporal_block_type == 'attention': - temporal_output, present_kv = self.attention( - hidden_states, - attention_mask=attention_mask, - use_cache=use_cache, - kv_cache_params=kv_cache_params, - attention_params=attention_params) - else: - present_kv = None - - hidden_states = residual + temporal_output - - residual = hidden_states - hidden_states = self.post_layernorm(hidden_states) - hidden_states = self.mlp(hidden_states) - hidden_states = residual + hidden_states - - return hidden_states, present_kv, present_conv, present_lru - - -class RecurrentGemmaModel(Module): - - def __init__(self, config: PretrainedConfig) -> None: - super().__init__() - self.d_conv = config.conv_kernel - self.lru_width = config.rnn_hidden_size - n_layer = config.num_hidden_layers - - self.vocab_embedding = Embedding(config.vocab_size, - config.hidden_size, - dtype=config.dtype) - self.layers = ModuleList( - [ResidualLayer(config, layer_idx=i) for i in range(n_layer)]) - - self.ln_f = RmsNorm(normalized_shape=config.hidden_size, - eps=config.norm_epsilon, - dtype=config.dtype) - - def forward(self, - input_ids, - use_cache=False, - attention_mask=None, - kv_cache_params=None, - attention_params=None, - conv_states=None, - lru_states=None, - host_request_types=None, - last_token_ids=None, - host_context_lengths=None, - slot_mapping=None): - - hidden_states = self.vocab_embedding(input_ids) - - # Get conv state indices - indices = None - if not default_net().plugin_config.mamba_conv1d_plugin: - batch_size = shape(input_ids, 0) - indices = expand( - unsqueeze(arange(0, self.d_conv - 1, dtype='int32'), 0), - concat([batch_size, self.d_conv - 1])) - offsets = expand(unsqueeze(last_token_ids, 1), - concat([batch_size, self.d_conv - 1])) - indices = unsqueeze(indices + offsets, 1) - indices = expand( - indices, concat([batch_size, self.lru_width, self.d_conv - 1])) - - present_kvs, present_convs, present_lrus = [], [], [] - for layer, past_kv, past_conv, past_lru in zip( - self.layers, kv_cache_params.past_key_value, conv_states, - lru_states): - hidden_states, present_kv, present_conv, present_lru = layer( - hidden_states, - use_cache, - attention_mask, - kv_cache_params=KeyValueCacheParams( - past_key_value=[past_kv], - host_past_key_value_lengths=kv_cache_params. - host_past_key_value_lengths, - host_max_attention_window_sizes=kv_cache_params. - host_max_attention_window_sizes, - host_sink_token_length=kv_cache_params. - host_sink_token_length, - kv_cache_block_offsets=kv_cache_params. - kv_cache_block_offsets, - host_kv_cache_block_offsets=kv_cache_params. - host_kv_cache_block_offsets, - host_kv_cache_pool_pointers=kv_cache_params. - host_kv_cache_pool_pointers, - host_kv_cache_pool_mapping=kv_cache_params. - host_kv_cache_pool_mapping, - cache_indirection=kv_cache_params.cache_indirection), - attention_params=attention_params, - conv_state=past_conv, - lru_state=past_lru, - host_request_types=host_request_types, - last_token_ids=last_token_ids, - host_context_lengths=host_context_lengths, - slot_mapping=slot_mapping, - conv_indices=indices) - present_kvs.append(present_kv) - present_convs.append(present_conv) - present_lrus.append(present_lru) - - hidden_states = self.ln_f(hidden_states) - return hidden_states, tuple(present_kvs), tuple(present_convs), tuple( - present_lrus) - - -class RecurrentGemmaForCausalLM(PretrainedModel): - - def __init__(self, config: PretrainedConfig): - super().__init__(config) - dtype = config.dtype - logits_dtype = config.logits_dtype - if isinstance(dtype, str): - self.dtype = str_dtype_to_trt(dtype) - else: - assert isinstance(dtype, trt.DataType) - self.dtype = dtype - - assert len(config.layer_types) > 0 - layer_types = config.layer_types - layer_types = layer_types * (config.num_hidden_layers // - len(layer_types)) - layer_types = layer_types + layer_types[0:(config.num_hidden_layers % - len(layer_types))] - self.layer_types = layer_types - - self.config = config - self.gather_context_logits = False - self.logits_soft_cap = config.logits_soft_cap - - # Create constant attention parameters to be reused by all layers. - Attention.create_attention_const_params(self, config) - self.position_embedding_type = config.position_embedding_type - - if isinstance(logits_dtype, str): - self._logits_dtype = str_dtype_to_trt(logits_dtype) - else: - assert isinstance(logits_dtype, trt.DataType) - self._logits_dtype = logits_dtype - - self.transformer = RecurrentGemmaModel(config) - self.lm_head = ColumnLinear(config.hidden_size, - config.vocab_size, - bias=False, - dtype=dtype, - tp_group=config.mapping.tp_group, - tp_size=config.mapping.tp_size, - gather_output=True) - - def forward(self, - input_ids, - position_ids=None, - use_cache=False, - attention_mask=None, - kv_cache_params=None, - attention_params=None, - conv_states=None, - rnn_states=None, - host_request_types=None, - last_token_ids=None, - last_token_ids_for_logits=None, - host_context_lengths=None, - slot_mapping=None): - - # fill attention params. - attention_params = Attention.fill_attention_params( - self, attention_params) - - hidden_states, present_kvs, present_convs, present_rnns = self.transformer( - input_ids, use_cache, attention_mask, kv_cache_params, - attention_params, conv_states, rnn_states, host_request_types, - last_token_ids, host_context_lengths, slot_mapping) - - if not self.gather_context_logits: - hidden_states = gather_last_token_logits( - hidden_states, last_token_ids_for_logits, - default_net().plugin_config.remove_input_padding) - - lm_logits = self.lm_head(hidden_states) - lm_logits = tanh( - lm_logits / self.logits_soft_cap) * self.logits_soft_cap - lm_logits.mark_output('logits', self._logits_dtype) - if not default_net().plugin_config.paged_kv_cache: - for i, present_kv in enumerate(present_kvs): - if present_kv is not None: - present_kv.mark_output(f'present_key_value_{i}', self.dtype) - - if not default_net().plugin_config.paged_state: - for i, present_conv in enumerate(present_convs): - if present_conv is not None: - present_conv.mark_output(f'present_conv_state_{i}', - self.dtype) - for i, present_rnn in enumerate(present_rnns): - if present_rnn is not None: - present_rnn.mark_output(f'present_rnn_state_{i}', - str_dtype_to_trt('float32')) - - return (lm_logits, present_kvs, present_convs, present_rnns) - - def prepare_recurrent_inputs(self, max_batch_size, num_profiles, mapping): - use_mamba_conv1d_plugin = default_net( - ).plugin_config.mamba_conv1d_plugin - - default_range = GenerationMixin.default_range - batch_range = [default_range(max_batch_size)] * num_profiles - - conv_states = [] - rnn_states = [] - dim = self.config.rnn_hidden_size // mapping.tp_size - if use_mamba_conv1d_plugin: - conv_state_dim_range = OrderedDict([ - ('batch_size', batch_range), - ('kernel_size', [self.config.conv_kernel - 1] * num_profiles), - ('dim_size', [dim] * num_profiles), - ]) - else: - conv_state_dim_range = OrderedDict([ - ('batch_size', batch_range), - ('dim_size', [dim] * num_profiles), - ('kernel_size', [self.config.conv_kernel - 1] * num_profiles), - ]) - - rnn_state_dim_range = OrderedDict([ - ('batch_size', batch_range), - ('state_size', [1] * num_profiles), - ('dim_size', [dim] * num_profiles), - ]) - one_dim_range = OrderedDict([ - ('buffer_count', [1] * num_profiles), - ]) - - for i in range(self.config.num_hidden_layers): - if self.layer_types[i] == 'recurrent': - if default_net().plugin_config.paged_state: - conv_state = Tensor(name=f'conv_state_ptr_{i}', - dtype=str_dtype_to_trt('int64'), - shape=[1], - dim_range=one_dim_range) - - rnn_state = Tensor(name=f'rnn_state_ptr_{i}', - dtype=str_dtype_to_trt('int64'), - shape=[1], - dim_range=one_dim_range) - else: - if use_mamba_conv1d_plugin: - conv_state = Tensor( - name=f'past_conv_state_{i}', - dtype=self.dtype, - shape=[-1, self.config.conv_kernel - 1, dim], - dim_range=conv_state_dim_range) - else: - conv_state = Tensor( - name=f'past_conv_state_{i}', - dtype=self.dtype, - shape=[-1, dim, self.config.conv_kernel - 1], - dim_range=conv_state_dim_range) - - rnn_state = Tensor(name=f'past_rnn_state_{i}', - dtype=str_dtype_to_trt('float32'), - shape=[-1, 1, dim], - dim_range=rnn_state_dim_range) - else: - conv_state, rnn_state = None, None - conv_states.append(conv_state) - rnn_states.append(rnn_state) - - slot_mapping = None - if default_net().plugin_config.paged_state: - slot_mapping = Tensor( - name='slot_mapping', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([('batch_size', batch_range)]), - ) - - return_dict = { - 'conv_states': conv_states, - 'rnn_states': rnn_states, - 'slot_mapping': slot_mapping, - } - return return_dict - - def prepare_inputs( - self, - max_batch_size, - max_input_len, - max_seq_len, - max_num_tokens, - use_cache, - max_beam_width: int = 1, - opt_num_tokens: int = None, - opt_batch_size: int = 0, - prompt_embedding_table_size: int = 0, - max_draft_len: int = 0, - gather_context_logits: bool = False, - lora_target_modules: List[str] = None, - speculative_decoding_draft_tokens_external: bool = False): - '''@brief: Prepare inputs Tensors for the model, the given sizes are used to determine the - ranges of the dimensions of when using TRT dynamic shapes. - - @return: a list contains values which can be fed into the self.forward() - ''' - assert speculative_decoding_draft_tokens_external == False, \ - "We don't support speculative decoding for the RecurrentGemma model." - assert max_beam_width == 1, "We don't support beam search for the RecurrentGemma model." - - remove_input_padding = default_net().plugin_config.remove_input_padding - use_gpt_attention_plugin = default_net( - ).plugin_config.gpt_attention_plugin - use_gemm_plugin = default_net().plugin_config.gemm_plugin - paged_kv_cache = default_net().plugin_config.paged_kv_cache - tokens_per_block = default_net().plugin_config.tokens_per_block - multiple_profiles = default_net().plugin_config.multiple_profiles - streamingllm = default_net().plugin_config.streamingllm - use_mamba_conv1d_plugin = default_net( - ).plugin_config.mamba_conv1d_plugin - - self.gather_context_logits = gather_context_logits - mapping = self.config.mapping - kv_cache_type = get_kv_cache_type_from_legacy(use_cache, paged_kv_cache) - - # basic inputs - enable_ctx_gen_opt_profiles = GenerationMixin.has_ctx_gen_opt_profiles( - use_gpt_attention_plugin=use_gpt_attention_plugin, - use_gemm_plugin=use_gemm_plugin, - remove_input_padding=remove_input_padding, - kv_cache_type=kv_cache_type) - num_profiles, ranges = GenerationMixin.get_profiles_ranges( - max_batch_size=max_batch_size, - max_beam_width=max_beam_width, - max_input_len=max_input_len, - max_num_tokens=max_num_tokens, - max_draft_len=max_draft_len, - opt_batch_size=opt_batch_size, - opt_num_tokens=opt_num_tokens, - enable_ctx_gen_opt_profiles=enable_ctx_gen_opt_profiles, - multiple_profiles=multiple_profiles, - kv_cache_type=kv_cache_type) - - if remove_input_padding: - assert use_mamba_conv1d_plugin, "mamba_conv1d_plugin is needed to support remove_input_padding" - input_ids = Tensor(name='input_ids', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([ - ('num_tokens', ranges['num_tokens_range']), - ])) - position_ids = Tensor(name='position_ids', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([ - ('position_ids_num_tokens_range', - ranges['num_tokens_range']), - ])) - else: - input_ids = Tensor(name='input_ids', - dtype=trt.int32, - shape=[-1, -1], - dim_range=OrderedDict([ - ('batch_size_beam_width', - ranges['bb_range']), - ('input_len', ranges['inlen_range']), - ])) - position_ids = Tensor(name='position_ids', - dtype=trt.int32, - shape=[-1, -1], - dim_range=OrderedDict([ - ('batch_size_beam_width', - ranges['bb_range']), - ('position_ids_inlen_range', - ranges['position_ids_inlen_range']), - ])) - if mapping.tp_size > 1: - current_all_reduce_helper().set_workspace_tensor( - mapping, num_profiles) - - # attention inputs - attn_layer_idx = [] - for i in range(self.config.num_hidden_layers): - if self.layer_types[i] == 'attention': - attn_layer_idx.append(i) - attention_inputs = self.prepare_attention_inputs( - max_batch_size=max_batch_size, - max_beam_width=max_beam_width, - max_input_len=max_input_len, - max_seq_len=max_seq_len, - num_kv_heads=self.config.num_key_value_heads, - head_size=self.config.head_size, - num_layers=self.config.num_hidden_layers, - kv_dtype=str_dtype_to_trt(self.config.kv_dtype), - num_profiles=num_profiles, - enable_ctx_gen_opt_profiles=enable_ctx_gen_opt_profiles, - remove_input_padding=remove_input_padding, - use_gpt_attention_plugin=use_gpt_attention_plugin, - kv_cache_type=kv_cache_type, - tokens_per_block=tokens_per_block, - mapping=mapping, - streamingllm=streamingllm, - attn_layer_idx=attn_layer_idx) - - # recurrent inputs - recurrent_inputs = self.prepare_recurrent_inputs( - max_batch_size=max_batch_size, - num_profiles=num_profiles, - mapping=mapping, - ) - - if use_gpt_attention_plugin: - host_request_types = attention_inputs['host_request_types'] - else: - host_request_types = Tensor( - name='host_request_types', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([('batch_size_beam_width', - ranges['bb_range'])]), - ) - - last_token_ids = Tensor( - name='last_token_ids', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([ - ('batch_size_last_token_ids', ranges['bbd_range']), - ]), - ) - last_token_ids_for_logits = None - if not gather_context_logits: - last_token_ids_for_logits = last_token_ids - - if use_gpt_attention_plugin and remove_input_padding: - host_context_lengths = attention_inputs['host_context_lengths'] - elif remove_input_padding: - host_context_lengths = Tensor( - name='host_context_lengths', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([('batch_size_beam_width', - ranges['bb_range'])]), - ) - else: - host_context_lengths = None - - return_dict = { - 'input_ids': - input_ids, - 'position_ids': - position_ids, - 'use_cache': - True, - 'attention_mask': - attention_inputs['attention_mask'], - 'kv_cache_params': - KeyValueCacheParams( - past_key_value=attention_inputs['past_key_value'], - host_past_key_value_lengths=attention_inputs[ - 'host_past_key_value_lengths'], - host_max_attention_window_sizes=attention_inputs[ - 'host_max_attention_window_sizes'], - host_sink_token_length=attention_inputs[ - 'host_sink_token_length'], - kv_cache_block_offsets=attention_inputs[ - 'kv_cache_block_offsets'], - host_kv_cache_block_offsets=attention_inputs[ - 'host_kv_cache_block_offsets'], - host_kv_cache_pool_pointers=attention_inputs[ - 'host_kv_cache_pool_pointers'], - host_kv_cache_pool_mapping=attention_inputs[ - 'host_kv_cache_pool_mapping'], - cache_indirection=attention_inputs['cache_indirection'], - ), - 'attention_params': - AttentionParams( - sequence_length=attention_inputs['sequence_length'], - context_lengths=attention_inputs['context_lengths'], - host_context_lengths=attention_inputs['host_context_lengths'], - max_context_length=max_input_len, - host_request_types=attention_inputs['host_request_types'], - host_runtime_perf_knobs=attention_inputs[ - 'host_runtime_perf_knobs'], - host_context_progress=attention_inputs['host_context_progress'], - ), - 'conv_states': - recurrent_inputs['conv_states'], - 'rnn_states': - recurrent_inputs['rnn_states'], - 'host_request_types': - host_request_types, - 'last_token_ids': - last_token_ids, - 'last_token_ids_for_logits': - last_token_ids_for_logits, - 'host_context_lengths': - host_context_lengths, - 'slot_mapping': - recurrent_inputs['slot_mapping'], - } - return return_dict diff --git a/tensorrt_llm/models/redrafter/__init__.py b/tensorrt_llm/models/redrafter/__init__.py deleted file mode 100644 index 2a36ca922710..000000000000 --- a/tensorrt_llm/models/redrafter/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-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. diff --git a/tensorrt_llm/models/redrafter/drafter.py b/tensorrt_llm/models/redrafter/drafter.py deleted file mode 100644 index c85560b3d473..000000000000 --- a/tensorrt_llm/models/redrafter/drafter.py +++ /dev/null @@ -1,117 +0,0 @@ -from typing import Optional - -from tensorrt_llm.functional import Tensor, silu -from tensorrt_llm.layers import ColumnLinear -from tensorrt_llm.mapping import Mapping -from tensorrt_llm.module import Module, ModuleList - -from ..._utils import str_dtype_to_trt - - -class ResBlock(Module): - - def __init__(self, - exit_dim: int, - dtype: Optional[str], - mapping: Mapping = Mapping()): - super().__init__() - self.linear = ColumnLinear( - exit_dim, - exit_dim, - bias=True, - dtype=dtype, - tp_group=mapping.tp_group, - tp_size=mapping.tp_size, - gather_output=True, - ) - - def forward(self, x: Tensor) -> Tensor: - return x + silu(self.linear(x)) - - -class Drafter(Module): - - def __init__( - self, - num_layers: int, - hidden_size: int, - exit_dim: int, - vocab_size: int, - dtype: Optional[str] = None, - is_rnn: bool = False, - mapping: Mapping = Mapping(), - ): - super().__init__() - self.num_layers = num_layers - self.is_rnn = is_rnn - self.dtype = str_dtype_to_trt(dtype) - - input_dim = 2 * hidden_size - self.input_proj = (None if input_dim == exit_dim else ColumnLinear( - input_dim, - exit_dim, - bias=True, - dtype=dtype, - tp_group=mapping.tp_group, - tp_size=mapping.tp_size, - gather_output=True, - )) - - self.layers = ModuleList([ - ResBlock(exit_dim, dtype, mapping) for _ in range(self.num_layers) - ]) - self.lm_head = ColumnLinear( - exit_dim, - vocab_size, - bias=False, - dtype=dtype, - tp_group=mapping.tp_group, - tp_size=mapping.tp_size, - gather_output=True, - ) - - if is_rnn: - self.rnn_u = ColumnLinear( - hidden_size, - hidden_size, - bias=True, - dtype=dtype, - tp_group=mapping.tp_group, - tp_size=mapping.tp_size, - gather_output=True, - ) - self.rnn_w = ColumnLinear( - hidden_size, - hidden_size, - bias=False, - dtype=dtype, - tp_group=mapping.tp_group, - tp_size=mapping.tp_size, - gather_output=True, - ) - return - - @classmethod - def from_config(cls, config, vocab_size_padded): - kwargs = { - "num_layers": config.redrafter_num_layers, - "hidden_size": config.redrafter_hidden_size, - "exit_dim": config.redrafter_exit_dim, - "vocab_size": vocab_size_padded, - "dtype": config.dtype, - "is_rnn": config.redrafter_is_rnn, - "mapping": config.mapping, - } - return cls(**kwargs) - - def forward(self, x: Tensor) -> Tensor: - hidden_states = self.input_proj(x) if self.input_proj is not None else x - for layer in self.layers: - hidden_states = layer(hidden_states) - - return self.lm_head(hidden_states) - - def rnn_embed(self, x: Tensor, prev: Tensor = None) -> Tensor: - assert self.is_rnn, "This function should not be called when redrafter_is_rnn is false." - w_embd = self.rnn_w(x) - return w_embd if prev is None else w_embd + self.rnn_u(prev) diff --git a/tensorrt_llm/models/redrafter/model.py b/tensorrt_llm/models/redrafter/model.py deleted file mode 100644 index 84c78cc79810..000000000000 --- a/tensorrt_llm/models/redrafter/model.py +++ /dev/null @@ -1,317 +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. - -from collections import OrderedDict - -import tensorrt as trt - -from tensorrt_llm._common import default_net -from tensorrt_llm.functional import Tensor, cast, categorical_sample -from tensorrt_llm.llmapi.kv_cache_type import KVCacheType -from tensorrt_llm.models import LLaMAForCausalLM, QWenForCausalLM -from tensorrt_llm.models.generation_mixin import GenerationMixin - -from ..._utils import pad_vocab_size, str_dtype_to_trt -from .drafter import Drafter -from .redrafter_helper import (_beam_search_candidates, _beams2tree, - _process_logits_and_hidden_states) - - -class ReDrafterMixin: - - def __init__(self, config): - - super().__init__(config) - self.dtype = str_dtype_to_trt(config.dtype) - self.vocab_size = config.vocab_size - vocab_size_padded = pad_vocab_size(self.vocab_size, - config.mapping.tp_size) - self.drafter = Drafter.from_config(config, vocab_size_padded) - self.num_beams = config.redrafter_num_beams - self.beam_candidate_length = config.redrafter_draft_len_per_beam - self.beam_length = self.beam_candidate_length + 1 # including true token - self.greedy_search = config.redrafter_greedy_search - self.is_rnn = config.redrafter_is_rnn - assert self.dtype == self.drafter.dtype, f"{self.dtype} != {self.drafter.dtype}" - - def _fwd_helper(self, hidden_states, lm_logits, embedding, drafter, - kwargs: dict): - ''' - Must enable remove_input_padding: - hidden_states [total_tokens, H] - lm_logits [total_tokens, V] - 1. process_logits: context vs gen - a. Context: just return the last hidden states, and logits/probs - b. Gen: - i. verify: use lm_logits, draft_probs, draft_indices, draft_tokens - ii. select hidden state and update probs - 3. Sample token based on probs - 4. Generate candidates using hidden_states, sampled token - 5. Using beams, generate validation buffers, mark them as output - 6. Mark all the outputs - ''' - - num_beams = self.num_beams - beam_length = self.beam_length - - # Get the inputs needed - rand_data_sample = kwargs['rand_data_sample'] - position_ids_base = kwargs['position_ids_base'] - - # Step 1: Process logits and hidden states - # process the base model output (verify for gen-phase) - probs, draft_input, num_accepted_tokens, \ - accepted_beam_index = _process_logits_and_hidden_states( - self, lm_logits, hidden_states, kwargs) - # NOTE: num_accepted_tokens doesn't include true token so add 1 here - num_accepted_tokens = num_accepted_tokens + 1 - - # At this point: - # probs : [bs, V] - # hidden_states : [bs, H] - - # Step 2: Sample token - next_token = categorical_sample(probs, rand_data_sample) - - # Step 3: beam search - new_draft_tokens, new_draft_logits = _beam_search_candidates( - draft_input, next_token, embedding, drafter, self.num_beams, - self.beam_length, self.is_rnn) - - # Step 4: tree processing - active_tokens_flattened, new_draft_token_indices, new_mask, \ - new_position_offsets, packed_position_ids, next_num_gen_tokens, max_gen_token, \ - total_gen_token = _beams2tree(new_draft_tokens, num_beams, beam_length, - position_ids_base + num_accepted_tokens) - - # Step 5: mark all the tensors we need - num_accepted_tokens.mark_output('num_accepted_tokens') - accepted_beam_index.mark_output('accepted_beam_index') - max_gen_token.mark_output('max_gen_token') - total_gen_token.mark_output('total_gen_token') - next_num_gen_tokens.mark_output('next_spec_decoding_generation_lengths') - active_tokens_flattened.mark_output('next_flat_tokens') - new_draft_tokens.mark_output('next_draft_tokens') - new_draft_logits.mark_output('next_draft_probs') - new_draft_token_indices.mark_output('next_draft_indices') - new_mask.mark_output('spec_decoding_mask') - new_position_offsets.mark_output('next_spec_decoding_position_offsets') - packed_position_ids.mark_output('packed_position_ids') - - return next_token, probs, draft_input - - def forward(self, *args, **kwargs): - """ - 0. run base model, get logits, hidden_states - """ - - extra_args = [ - 'draft_tokens', - 'draft_indices', - 'draft_probs', - 'device_request_types', - 'redrafter_inverted_temperature', - 'rand_data_validation', - 'rand_data_sample', - 'position_ids_base', - ] - use_cache = True - base_kwargs = {k: v for k, v in kwargs.items() if k not in extra_args} - if use_cache and default_net().plugin_config.paged_kv_cache is False: - lm_logits, presents, hidden_states = super().forward( - *args, **base_kwargs) - else: - lm_logits, hidden_states, _ = super().forward(*args, **base_kwargs) - - # lm_logits could be in fp32 - lm_logits_cast = cast(lm_logits, self.dtype) # no-op if same type - self.register_network_output("hidden_states", - hidden_states) # debugging - - new_draft_tokens, new_draft_logits, probs = self._fwd_helper( - hidden_states, - lm_logits_cast, - self.transformer.vocab_embedding, - self.drafter, - kwargs=kwargs) - - return new_draft_tokens, new_draft_logits, probs - - def prepare_inputs(self, *args, **kwargs): - """ - Inputs needed: - Assuming, max_gen_tokens = 1 + nb*(bl - 1), counting true token - device_request_types: [bs] - draft_tokens: [bs, nb, bl] - draft_indices: [bs, nb, bl] - draft_probs: [bs, nb, bl-1, V] - spec_decoding_generation_lengths: [bs] - spec_decoding_position_offsets: [bs, max_gen_tokens] - spec_decoding_packed_mask: [bs, max_gen_tokens, packed_length] ** - redrafter_inverted_temperature: [bs] - rand_data_sample: [bs] - rand_data_validation: [bs, nb, bl-1] - - ** The mask is tricky since the boolean mask will need to be - packed in runtime. So, the last dim will be: - packed_length = ceil(max_gen_tokens/32) - """ - default_range = GenerationMixin.default_range - remove_input_padding = default_net().plugin_config.remove_input_padding - use_gpt_attention_plugin = default_net( - ).plugin_config.gpt_attention_plugin - use_gemm_plugin = default_net().plugin_config.gemm_plugin - paged_kv_cache = default_net().plugin_config.paged_kv_cache - max_batch_size = kwargs['max_batch_size'] - assert max_batch_size is not None - bb_range = default_range(max_batch_size) - bb0_range = default_range(max_batch_size, min_range=0, opt_offset=1) - num_beam_tokens = self.num_beams * self.beam_length - max_draft_len = num_beam_tokens - self.num_beams # ignore the true token - max_gen_token_len = 1 + max_draft_len # for the true token - max_gen_token_len_range = default_range(max_gen_token_len) - bb_max_gen_token_len_range = default_range(max_gen_token_len * - max_batch_size, - min_range=0) - - kwargs['speculative_decoding_draft_tokens_external'] = False - kwargs['max_draft_len'] = max_draft_len - kwargs['spec_decoding_is_generation_length_variable'] = True - inputs = super().prepare_inputs(*args, **kwargs) - assert inputs['spec_decoding_params'] is not None - - enable_two_optimization_profiles = GenerationMixin.has_ctx_gen_opt_profiles( - use_gpt_attention_plugin=use_gpt_attention_plugin, - use_gemm_plugin=use_gemm_plugin, - remove_input_padding=remove_input_padding, - kv_cache_type=KVCacheType.PAGED - if paged_kv_cache else KVCacheType.CONTINUOUS) - if enable_two_optimization_profiles: - bb_range = [bb_range, bb_range] - bb0_range = [bb0_range, bb0_range] - max_gen_token_len_range = [ - max_gen_token_len_range, max_gen_token_len_range - ] - bb_max_gen_token_len_range = [ - bb_max_gen_token_len_range, bb_max_gen_token_len_range - ] - num_beams_range = [self.num_beams, self.num_beams] - beam_length_range = [self.beam_length, self.beam_length] - candidate_length_range = [ - self.beam_candidate_length, self.beam_candidate_length - ] - vocab_size_range = [self.vocab_size, self.vocab_size] - else: - bb_range = [bb_range] - bb0_range = [bb0_range] - max_gen_token_len_range = [max_gen_token_len_range] - bb_max_gen_token_len_range = [bb_max_gen_token_len_range] - num_beams_range = [self.num_beams] - beam_length_range = [self.beam_length] - candidate_length_range = [self.beam_candidate_length] - vocab_size_range = [self.vocab_size] - - device_request_types = Tensor(name='device_request_types', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([ - ('batch_size', bb_range), - ])) - draft_tokens = Tensor(name='draft_tokens', - dtype=trt.int32, - shape=[-1, self.num_beams, self.beam_length], - dim_range=OrderedDict([ - ('batch_size_wt0', bb0_range), - ('num_beams', num_beams_range), - ('beam_length', beam_length_range), - ])) - draft_indices = Tensor(name='draft_indices', - dtype=trt.int32, - shape=[-1, self.num_beams, self.beam_length], - dim_range=OrderedDict([ - ('batch_size_wt0', bb0_range), - ('num_beams', num_beams_range), - ('beam_length', beam_length_range), - ])) - draft_probs = Tensor( - name='draft_probs', - dtype=self.dtype, - shape=[-1, self.num_beams, self.beam_length - 1, self.vocab_size], - dim_range=OrderedDict([ - ('batch_size_wt0', bb0_range), - ('num_beams', num_beams_range), - ('candidate_length', candidate_length_range), - ('vocab_size', vocab_size_range), - ])) - redrafter_inverted_temperature = Tensor( - name='redrafter_inverted_temperature', - dtype=self.dtype, - shape=[-1], - dim_range=OrderedDict([ - ("batch_size", bb_range), - ])) - rand_data_validation = Tensor( - name='rand_data_validation', - dtype=self.dtype, - shape=[-1, self.num_beams, self.beam_length - 1], - dim_range=OrderedDict([ - ('batch_size_wt0', bb0_range), - ('num_beams', num_beams_range), - ('candidate_length', candidate_length_range), - ])) - rand_data_sample = Tensor(name='rand_data_sample', - dtype=self.dtype, - shape=[-1], - dim_range=OrderedDict([ - ('batch_size', bb_range), - ])) - position_ids_base = Tensor( - name="position_ids_base", - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([ - ("batch_size", bb_range), - ]), - ) - - inputs[ - 'device_request_types'] = device_request_types # needed by process_logits - inputs['draft_tokens'] = draft_tokens - inputs['draft_indices'] = draft_indices - inputs['draft_probs'] = draft_probs - inputs[ - 'redrafter_inverted_temperature'] = redrafter_inverted_temperature - inputs['rand_data_validation'] = rand_data_validation - inputs['rand_data_sample'] = rand_data_sample - inputs['position_ids_base'] = position_ids_base - return inputs - - -class ReDrafterForQWenLM(ReDrafterMixin, QWenForCausalLM): - """ReDrafter implementation for QWen models. - - Combines: - - Base QWen model functionality from QWenForCausalLM - - Drafting/speculative decoding logic from ReDrafterMixin - """ - - -class ReDrafterForLLaMALM(ReDrafterMixin, LLaMAForCausalLM): - """ReDrafter implementation for LLaMA models. - - Combines: - - Base LLaMA model functionality from LLaMAForCausalLM - - Drafting/speculative decoding logic from ReDrafterMixin - """ diff --git a/tensorrt_llm/models/redrafter/redrafter_helper.py b/tensorrt_llm/models/redrafter/redrafter_helper.py deleted file mode 100644 index 31e930f51a97..000000000000 --- a/tensorrt_llm/models/redrafter/redrafter_helper.py +++ /dev/null @@ -1,759 +0,0 @@ -from typing import Tuple - -import numpy as np - -from tensorrt_llm._common import default_net -from tensorrt_llm._utils import numpy_array - -# isort: off -from tensorrt_llm.functional import ( - Tensor, arange, argmax, cast, concat, constant, constant_to_tensor_, cumsum, - div, eq, exp, expand, expand_dims, floordiv, gather, gather_nd, - index_select, int32_array, log_softmax, lt, max, maximum, masked_select, - minimum, nonzero, not_op, op_and, rand, relu, scatter, select, shape, slice, - silu, softmax, squeeze, stack, sum, topk, transpose, unsqueeze, view, where) -# isort: on -from tensorrt_llm.layers import Embedding -from tensorrt_llm.module import Module - -INT_DTYPE_STR = "int32" -''' -NOTE: - Name differences from Apple's PyTorch Implementation: - `num_candidates` is mapped to `num_beams` and - `candidate_length` is mapped to `beam_length - 1`. - So for each sequence, the paths/beams to verify will be [num_beams, beam_length] tokens where - each beam is a path that includes the true token (1) and the candidate tokens (beam_length - 1). -''' - - -def _unpack_beams(x: Tensor, indices: Tensor, num_beams: int, - beam_length: int) -> Tensor: - """ - x: [bs, S, V] - indices: [bs, nb, bl] - output: - """ - assert x.rank() == 3 - d0 = shape(x, 0, INT_DTYPE_STR) - dl = shape(x, -1, INT_DTYPE_STR) - indices = view(indices, [-1, num_beams * beam_length, 1], False) - res_shape = concat([d0, num_beams, beam_length, dl]) - res = view(gather_nd(x, indices), res_shape, False) # [d0, nb, bl, dl] - return res - - -def _validate_draft_tokens(draft_log_probs: Tensor, - draft_tokens: Tensor, - draft_indices: Tensor, - flattened_logits: Tensor, - num_beams: int, - beam_length: int, - greedy_search: bool, - rand_data: Tensor = None): - ''' - draft_log_probs: [bs, nb, bl-1, V] - draft_tokens: [bs, nb, bl] - draft_indices: [bs, nb, bl] - flattened_logits: [bs, S, V], we need to unflatten it using draft_indices. - The unflattend_logits should be of shape [bs, nb, bl, V] by doing a gather on S. - ''' - batch_size = shape(flattened_logits, 0, INT_DTYPE_STR) - rand_shape = concat([batch_size, num_beams, beam_length - 1]) - if rand_data is None: - rand_data = rand(rand_shape, low=0, high=1, dtype=draft_log_probs.dtype) - - flat_log_probs = log_softmax(flattened_logits, dim=-1) - all_base_log_probs = _unpack_beams(flat_log_probs, draft_indices, num_beams, - beam_length) # [bs, nb, bl, V] - if greedy_search: - all_base_log_probs = _top_1_logits(all_base_log_probs) - - base_log_probs = index_select(all_base_log_probs, - dim=2, - index=constant( - np.arange(beam_length - 1, - dtype=np.int32))) - last_base_log_probs = select(all_base_log_probs, - dim=2, - index=beam_length - 1) - proposed_tokens = unsqueeze(slice(draft_tokens, [0, 0, 1], rand_shape), -1) - - token_base_log_probs = squeeze( - gather(base_log_probs, dim=-1, indices=proposed_tokens), -1) - token_draft_log_probs = squeeze( - gather(draft_log_probs, dim=-1, indices=proposed_tokens), -1) - diff_probs = exp(token_base_log_probs - token_draft_log_probs) - cmp = cast(lt(rand_data, diff_probs), dtype='int32') - ideal_sum = constant(np.arange(1, beam_length, dtype=np.int32)) - cum_sum = cumsum(cmp, dim=-1) - equality = cast((cum_sum == ideal_sum), dtype='int32') - num_accepted = sum(equality, dim=-1) - max_num_accepted_tokens, accepted_beam_index = topk( - num_accepted, k=1, - dim=-1) # need to use topk layer to get both value and index - return squeeze(max_num_accepted_tokens, -1), squeeze(accepted_beam_index, -1),\ - base_log_probs, last_base_log_probs, rand_data - - -def _get_prefix_match_indices(beams, beam_length): - ''' - beams: [bs, nb, bl] - ''' - prefix_target = constant( - np.expand_dims(np.arange(1, beam_length + 1, dtype=np.int32), - [0, 1, 2])) - matches = cast(expand_dims(beams, 1) == expand_dims(beams, 2), beams.dtype) - seq_matches = cast(cumsum(matches, dim=3) == prefix_target, - dtype=beams.dtype) - prefix_match_indices = argmax(seq_matches, dim=2) - return prefix_match_indices - - -def _get_draft_token_indices(prefix_match_indices, num_beams, beam_length): - ''' - prefix_match_indices: [bs, nb, bl] - ''' - pmi_dtype = prefix_match_indices.dtype - segments = cast( - constant(np.expand_dims(np.arange(0, num_beams, dtype=np.int32), - [0, 2])) == prefix_match_indices, pmi_dtype) - segment_lengths = sum(segments, dim=-1) - accum_lengths = cumsum(segment_lengths, dim=-1) - segment_lengths - segment_index = gather(accum_lengths, - dim=1, - indices=view(prefix_match_indices, - shape=[-1, num_beams * beam_length])) - segment_index = view(segment_index, [-1, num_beams, beam_length]) - match = cast( - expand_dims(segment_index, 3) == expand_dims(segment_index, 2), - pmi_dtype) - seq_index = constant(np.arange(beam_length, dtype=np.int32)) - lower_triangle = cast( - expand_dims(seq_index, 1) > expand_dims(seq_index, 0), pmi_dtype) - offset = sum(match * expand_dims(lower_triangle, [0, 1]), dim=-1) - draft_token_indices = segment_index + offset - return draft_token_indices - - -def _get_packed_position_ids( - active_indices: Tensor, - indices: Tensor, - total_lengths: Tensor, - position_ids_base: Tensor, -) -> Tensor: - expand_shape = concat([shape(total_lengths, 0), shape(indices, 0)]) - expanded_indices = expand(unsqueeze(indices, 0), expand_shape) - position_mask = expanded_indices < unsqueeze(total_lengths, 1) - position_ids = active_indices + unsqueeze(position_ids_base, 1) - packed_position_ids = masked_select(position_ids, position_mask) - return packed_position_ids - - -def _get_draft_token_array( - beams: Tensor, - prefix_match_indices: Tensor, - num_beams: int, - beam_length: int, - position_ids_base: Tensor = None, -) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor]: - ''' - beams: [bs, nb, bl] - prefix_match_indices: [bs, nb, bl] - ''' - prefix_ideal_indices = constant(np.arange(num_beams, dtype=np.int32)) - prefix_ideal_indices = expand_dims(prefix_ideal_indices, [0, 2]) - segments = cast(eq(prefix_match_indices, prefix_ideal_indices), - dtype=beams.dtype) - raw_draft_token_array = view(segments * beams + (segments - 1), - [-1, num_beams * beam_length], False) - raw_active_token_indices = transpose( - nonzero(not_op(raw_draft_token_array == -1)), 0, 1) - active_token_flattened = gather_nd(raw_draft_token_array, - raw_active_token_indices, 0) - - total_lengths = sum(view(segments, [-1, num_beams * beam_length], False), - dim=1) - slice_size = concat([shape(raw_active_token_indices, 0, INT_DTYPE_STR), 1]) - active_token_index_flattened = view( - slice(raw_active_token_indices, starts=[0, 1], sizes=slice_size), [-1], - False) - - max_len = max(total_lengths, dim=0) - total_gen_len = sum(total_lengths, dim=0) - # constant_0 = constant(int32_array(0)) - # offset = arange(constant_0, max_len, dtype='int32') - offset = slice(constant(np.arange(num_beams * beam_length, dtype=np.int32)), - constant_to_tensor_(0), unsqueeze(max_len, 0)) - idx_starts = cumsum(total_lengths, 0) - total_lengths - select_indices = unsqueeze(idx_starts, -1) + unsqueeze(offset, 0) - max_index_allowed = shape(active_token_flattened, 0, INT_DTYPE_STR) - 1 - select_indices = minimum(view(select_indices, [-1], False), - max_index_allowed) - compressed_shape = concat([shape(total_lengths, 0, INT_DTYPE_STR), max_len]) - # draft_token_array = view( - # gather(active_token_flattened, dim=0, indices=select_indices), - # compressed_shape, False) - active_token_indices = view( - gather(active_token_index_flattened, dim=0, indices=select_indices), - compressed_shape, False) - # adding position offsets here - position_offsets = active_token_indices % beam_length - packed_position_ids = constant_to_tensor_(0) # dummy initialization - if position_ids_base is not None: - packed_position_ids = _get_packed_position_ids(position_offsets, offset, - total_lengths, - position_ids_base) - return active_token_flattened, active_token_indices, total_lengths, max_len, total_gen_len, position_offsets, packed_position_ids - - -# FROM APPLE (minor changes by NV) -def _get_mask(draft_token_indices: Tensor, active_token_indices: Tensor, - num_beams: int, beam_length: int) -> Tensor: - """ - Return mask for candidates according to the flattened and compact index. - Args: - draft_token_indices: (batch_size, num_beams, beam_length) - A Mapping of draft candidates index from a stacked representation to a - flattened and compact representation. - active_token_indices: (batch_size, max_len) - A Mapping of draft candidates index from a flattened and compact representation - to a stacked representation. - Returns: - compact_candidate_mask: (batch_size, max_len, max_len) - Output a mask tensor for candidates with a flattened and compact indexing. - """ - - batch_size = shape(draft_token_indices, 0, INT_DTYPE_STR) - max_len = shape(active_token_indices, 1, INT_DTYPE_STR) - all_candidate_len = beam_length * num_beams - - arange_all_candidates = constant( - np.arange(all_candidate_len, dtype=np.int32)) - active_token_beam = div(active_token_indices, beam_length) - beam_blocks = div(arange_all_candidates, beam_length) - - lower_triangle_mask = (unsqueeze(arange_all_candidates, axis=-1) - - unsqueeze(arange_all_candidates, axis=0) >= 0) - block_diagonal_mask = unsqueeze(beam_blocks, axis=-1) - unsqueeze( - beam_blocks, axis=0) == 0 - # `candidates_mask` is the flattened candidates mask - candidates_mask = expand( - expand_dims(op_and(lower_triangle_mask, block_diagonal_mask), [0]), - concat([batch_size, all_candidate_len, all_candidate_len]), - ) - - expanded_active_token_indices = expand( - expand_dims(active_token_indices, [2]), - concat([batch_size, max_len, all_candidate_len])) - raw_token_mask = gather(candidates_mask, - dim=1, - indices=expanded_active_token_indices) - - src_idx = unsqueeze(active_token_beam, axis=-1) * beam_length + expand_dims( - constant(np.arange(beam_length, dtype=np.int32)), [0, 1]) - src_mask = gather(raw_token_mask, dim=2, indices=src_idx) - tgt_idx = gather( - draft_token_indices, - dim=1, - indices=expand(expand_dims(active_token_beam, [2]), - concat([batch_size, max_len, beam_length])), - ) - # `compact_candidate_mask` is the compact and flattened candidates mask - compact_candidate_mask = expand( - expand_dims(cast(constant_to_tensor_(0), dtype="bool"), [0, 1]), - concat([batch_size, max_len, max_len]), - ) - - updated_compact_candidate_mask = scatter( - compact_candidate_mask, - dim=2, - indices=tgt_idx, - updates=src_mask, - ) - - return updated_compact_candidate_mask - - -def _beams2tree( - beams: Tensor, - num_beams: int, - beam_length: int, - position_ids_base: Tensor = None, -) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor]: - ''' - beams: [bs, nb, bl] - ''' - prefix_match_indices = _get_prefix_match_indices(beams, beam_length) - draft_token_indices = _get_draft_token_indices(prefix_match_indices, - num_beams, beam_length) - active_tokens_flattened, active_token_indices, total_lengths, max_gen_len, \ - total_gen_len, position_offsets, packed_position_ids = _get_draft_token_array( - beams, prefix_match_indices, num_beams, beam_length, position_ids_base) - mask = _get_mask(draft_token_indices, active_token_indices, num_beams, - beam_length) - return active_tokens_flattened, draft_token_indices, mask, position_offsets, packed_position_ids, total_lengths, max_gen_len, total_gen_len - - -def _get_indices_for_gather_beams(batch_size: Tensor, beam_indices: Tensor, - num_beams: int) -> Tensor: - ''' - beam_indices: [bs, nb] - Returns: [bs*nb, 2] - ''' - constant_0 = constant(int32_array(0)) - batch_indices = arange(constant_0, batch_size * num_beams, dtype='int32') - batch_indices = floordiv(batch_indices, num_beams) - - indices = concat([ - view(batch_indices, [-1, 1], False), - view(beam_indices, [-1, 1], False) - ], - dim=1) - return indices - - -def _gather_beams(x: Tensor, indices: Tensor, batch_size: Tensor, - num_beams: int) -> Tensor: - ''' - x: [bs, nb, X] - beam_indices: [bs, nb] - Returns: [bs, nb, X] - ''' - target_shp = [batch_size, constant(int32_array(num_beams))] - for i in range(2, x.ndim()): - target_shp.append(shape(x, i, INT_DTYPE_STR)) - target_shp = concat(target_shp) - return view(gather_nd(x, indices, batch_dims=0), target_shp, False) - - -def _add_decoding_dim(x: Tensor, num_beams: int) -> Tensor: - assert x.ndim() == 1 or x.ndim() == 2 - x = unsqueeze(x, 1) - new_shp = [shape(x, 0, INT_DTYPE_STR), num_beams] if x.ndim() == 2 else [ - shape(x, 0, INT_DTYPE_STR), num_beams, - shape(x, 2, INT_DTYPE_STR) - ] - res = expand(x, concat(new_shp)) - return res - - -def _flatten_decoding_dim(x: Tensor) -> Tensor: - if x.ndim() > 1: - new_shp = [-1 - ] + [shape(x, i, INT_DTYPE_STR) for i in range(2, x.ndim())] - return view(x, concat(new_shp)) - return x - - -def _unflatten_decoding_dim(x: Tensor, num_beams: int) -> Tensor: - ''' - Unflattens the first, flat batch*decoding dimension of a non-scalar array. - x: [bs*num_beams, ...] - ''' - if x.ndim() > 0: - new_shp = [-1, num_beams - ] + [shape(x, i, INT_DTYPE_STR) for i in range(1, x.ndim())] - return view(x, concat(new_shp)) - return x - - -def _beam_search_candidates(prompt_state: Tensor, init_token: Tensor, - embedding: Embedding, drafter: Module, - num_beams: int, beam_length: int, - is_rnn: bool) -> Tuple[Tensor, Tensor]: - """ - This version of beam search matches with ReDrafter GitHub version as of 10/02/2024. - Link: https://github.com/apple/ml-recurrent-drafter/releases/tag/v1.1 - """ - - LOG_0 = -50000.0 - LOG_1 = 0.0 - - def maintain_logits(logits: Tensor) -> Tensor: - max_logits = max(logits, -1, keepdim=True) - max_logits = expand(max_logits, - shape(logits, cast_to_dtype=INT_DTYPE_STR)) - return logits - max_logits - - def warp_logits(logits: Tensor, - top_k: int = 50, - mask_value: float = LOG_0) -> Tensor: - top_k = minimum(top_k, shape(logits, - dim=-1, - cast_to_dtype=INT_DTYPE_STR)) - top_values, _ = topk(logits, k=top_k, dim=-1) # [bs, nb, top_k] - starts = concat([0, 0, top_k - 1]) - sizes = concat([shape(logits, 0), shape(logits, 1), 1]) - lt_mask = logits < slice(top_values, starts=starts, sizes=sizes) - logits = where(lt_mask, - constant_to_tensor_(mask_value, dtype=logits.dtype), - logits) - return logits - - def compute_logits(x: Tensor) -> Tensor: - """ - x: [bs, nb, 2*H] - """ - logits = drafter(x) # [bs, nb, 2*H] => [bs, nb, V] - logits = maintain_logits(logits) # [bs, nb, V] - logits = warp_logits(logits) # [bs, nb, V] - return logits - - assert prompt_state.ndim() == 2 - assert init_token.ndim() == 1 - assert beam_length > 1 - batch_size = shape(prompt_state, 0, INT_DTYPE_STR) - vocab_size = embedding.num_embeddings - dtype = prompt_state.dtype - - log_p_beam = expand( - unsqueeze( - constant( - numpy_array([LOG_1] + [LOG_0] * (num_beams - 1), - trt_dtype=dtype)), 0), # [1, nb] - concat([batch_size, num_beams])) # [bs, nb] - context = _add_decoding_dim(prompt_state, num_beams) # [bs, nb, H] - if init_token.ndim() == 1: - init_token = unsqueeze(init_token, -1) # [bs] => [bs, 1] - beams = _add_decoding_dim(init_token, num_beams) # [bs, nb, 1] - - last_tokens = squeeze(beams, -1) # [bs, nb] - state_shape = shape(context, cast_to_dtype=INT_DTYPE_STR) # [bs, nb, H] - state = expand(expand_dims(constant_to_tensor_(0.0, dtype=dtype), [0, 1]), - state_shape) # [bs, nb, H] - log_p_token_in_beam = None - candidate_length = beam_length - 1 - for _ in range(candidate_length): - state = ( - silu(drafter.rnn_w(embedding(last_tokens)) + - drafter.rnn_u(state)) if is_rnn else embedding(last_tokens) + - state) # [bs, nb, H] - - logits_new_token = compute_logits(concat([context, state], - -1)) # [bs, nb, V] - log_p_new_token = log_softmax(logits_new_token, -1) # [bs, nb, V] - - log_p_beam_new_token = log_p_new_token + unsqueeze(log_p_beam, - 2) # [bs, nb, V] - - tokens_times_beams = view(log_p_beam_new_token, - concat([batch_size, num_beams * vocab_size - ])) # [bs, nb*V] - log_p_beam, topk_indices = topk(tokens_times_beams, k=num_beams, - dim=-1) # [bs, nb] - top_beam_indices = topk_indices // vocab_size # [bs, nb] - # Avoid repeated division for: top_token_ids = topk_indices % vocab_size - top_token_ids = topk_indices - (top_beam_indices * vocab_size - ) # [bs, nb] - - # get the common indices to gather beams - gather_indices = _get_indices_for_gather_beams(batch_size, - top_beam_indices, - num_beams) - - # update running beams, state, logits, and last_tokens - prev_top_beams = _gather_beams(beams, gather_indices, batch_size, - num_beams) # [bs, nb] OR [bs, nb, 1+i] - if prev_top_beams.ndim() == 2: - prev_top_beams = unsqueeze(prev_top_beams, -1) # [bs, nb, 1] - new_tokens = unsqueeze(top_token_ids, -1) # [bs, nb, 1] - beams = concat([prev_top_beams, new_tokens], dim=-1) # [bs, nb, 1+i+1] - - state = _gather_beams(state, gather_indices, batch_size, - num_beams) # [bs, nb, H] - - cur_log_p_token_in_beam = unsqueeze( - _gather_beams(log_p_new_token, gather_indices, batch_size, - num_beams), 2) # [bs, nb, 1, V] - if log_p_token_in_beam is None: # first iteration - log_p_token_in_beam = cur_log_p_token_in_beam - else: - log_p_token_in_beam = concat( - [ - _gather_beams(log_p_token_in_beam, gather_indices, - batch_size, - num_beams), # prev_top_logits [bs, nb, i, V] - cur_log_p_token_in_beam - ], - dim=2) # [bs, nb, i+1, V] - last_tokens = top_token_ids # [bs, nb] - return beams, log_p_token_in_beam - - -def _top_1_logits(logits: Tensor, NINF=-50000.0) -> Tensor: - ''' - logits: [bs, S, V] - ''' - NEG_INF = constant_to_tensor_(NINF, logits.dtype) - # TODO: WAR for bug in max reduction: https://nvbugs/4714485 - # max_values = max(logits, dim=-1, keepdim=True) # [bs, S, 1] - max_values, _ = topk(logits, k=1, dim=-1) # [bs, S, 1] - cmp = not_op(logits == max_values) - res = cast(cmp, dtype=logits.dtype) * NEG_INF - return res - - -def _ctx_logits2probs(logits: Tensor, greedy_search: bool) -> Tensor: - """ - Inputs: - logits: [bs_ctx, V] - Returns: - probs: [bs_ctx, V] - """ - if greedy_search: - logits = _top_1_logits(logits) - probs = softmax(logits, dim=-1) - return probs - - -# Jointly developed with Apple -def _batch_index_select(x: Tensor, batch_index: Tensor) -> Tensor: - """select the tensor by index inside each batch - - Args: - x (Tensor): [batch, ..] - batch_index (Tensor): (batch_size) - - Returns: - Tensor: [batch, ..] Tensors selected by the indices - """ - expanded_shape = concat( - [shape(x, 0, INT_DTYPE_STR), 1] + - [shape(x, i, INT_DTYPE_STR) for i in range(2, x.rank())]) - batch_index = expand( - expand_dims(batch_index, range(1, - x.rank() - batch_index.rank() + 1)), - expanded_shape) - gathered_x = gather(x, dim=1, indices=batch_index) - return squeeze(gathered_x, dim=1) - - -# Jointly developed with Apple -def _prepare_drafter_input( - draft_log_probs: Tensor, - base_log_probs: Tensor, - last_base_log_probs: Tensor, - accepted_beam_index: Tensor, - num_accepted_tokens: Tensor, -) -> Tensor: - """ - Args: - num_accepted_tokens: (batch_size) - Highest count of accepted tokens. - accepted_beam_index: (batch_size) - Beam index with highest count of accepted tokens. - draft_log_probs: (batch_size, num_candidates, candidate_length, vocab_size) - Draft head log probs for draft_tokens. - base_log_probs: (batch_size, num_candidates, candidate_length, vocab_size) - LM log probs for draft_tokens. - last_base_log_probs: (batch_size, num_candidates, vocab_size) - Last token log probs for all candidates to predict the next token beyond each candidate. - Returns: - probs: (batch_size, vocab_size): - Predict next token probability. - - """ - # Select according to the chosen beam index. - candidate_length = shape(draft_log_probs, 2, INT_DTYPE_STR) - selected_draft_log_probs = _batch_index_select(draft_log_probs, - accepted_beam_index) - selected_base_log_probs = _batch_index_select(base_log_probs, - accepted_beam_index) - selected_last_base_log_probs = _batch_index_select(last_base_log_probs, - accepted_beam_index) - - # Check if the entire beam is accepted or not. - entire_beam_accept = unsqueeze(num_accepted_tokens == candidate_length, - axis=-1) - - # If the entire beam is accepted, we use maybe_last_probs to sample next token. - maybe_last_probs = exp(selected_last_base_log_probs) - - # Note the shape of selected_draft_log_probs and selected_base_log_probs is the same - # as [batch_size, candidate_length, vocab_size]. - # Thus, we clamp resample_index to be up to candidate_length - 1. - # Since when num_accepted_tokens == candidate_length, we use maybe_last_probs above. - resample_index = num_accepted_tokens - cast( - eq(num_accepted_tokens, candidate_length), dtype='int32') - sample_draft_log_probs = _batch_index_select(selected_draft_log_probs, - resample_index) - sample_base_log_probs = _batch_index_select(selected_base_log_probs, - resample_index) - # Rejection sampling probs. - probs = relu(exp(sample_base_log_probs) - exp(sample_draft_log_probs)) - probs = where(entire_beam_accept, maybe_last_probs, probs) - - return probs - - -def _process_gen_logits(logits: Tensor, - hidden: Tensor, - draft_probs: Tensor, - draft_tokens: Tensor, - draft_indices: Tensor, - num_beams: int, - beam_length: int, - greedy_search: bool, - rand_data: Tensor = None) -> Tensor: - num_accepted_tokens, accepted_beam_index,\ - base_log_probs, last_base_log_probs, _ = _validate_draft_tokens( - draft_probs, draft_tokens, draft_indices, logits, num_beams, beam_length, - greedy_search, rand_data) - - # need to retrieve flattened index from accepted_beam_index and num_accepted_tokens - indices = stack([accepted_beam_index, num_accepted_tokens], 1) - flat_indices = unsqueeze(gather_nd(draft_indices, indices, batch_dims=1), - -1) - filtered_probs = _prepare_drafter_input(draft_probs, base_log_probs, - last_base_log_probs, - accepted_beam_index, - num_accepted_tokens) - filtered_hidden = gather_nd(hidden, flat_indices, batch_dims=1) - return filtered_probs, filtered_hidden, num_accepted_tokens, accepted_beam_index - - -def _get_gen_token_indices_for_unpack( - num_gen_tokens: Tensor, num_beams: int, beam_length: int, - max_index_allowed: Tensor) -> Tuple[Tensor, Tensor]: - upper_bound = num_beams * beam_length - num_beams + 1 - max_gen_tokens = max(num_gen_tokens, dim=0) - max_gen_tokens = minimum(max_gen_tokens, upper_bound) - max_gen_tokens = maximum(max_gen_tokens, 0) - cum_gen_tokens = cumsum(num_gen_tokens, 0) - gen_token_starts = cum_gen_tokens - num_gen_tokens - gen_unpack_indxs = arange(constant_to_tensor_(0, to_array=False), - max_gen_tokens, - dtype='int32') - gen_unpack_indxs = unsqueeze(gen_unpack_indxs, 0) + unsqueeze( - gen_token_starts, 1) - gen_unpack_indxs = minimum(gen_unpack_indxs, max_index_allowed) - return gen_unpack_indxs, max_gen_tokens - - -def _unpack_gen_data(x: Tensor, num_gen_tokens: Tensor, - gen_unpack_indxs: Tensor, - max_gen_tokens: Tensor) -> Tensor: - """ - x: [sum(num_gen_tokens), V/H] - num_gen_tokens: [gen_bs] - gen_unpack_indxs: [bs, max(num_gen_tokens)] - Returns: - [gen_bs, max_gen_tokens, V/H] where max_gen_tokens = max(num_gen_tokens) - """ - unpacked_x = index_select(x, dim=0, index=view(gen_unpack_indxs, [-1])) - out_shape = concat([ - shape(num_gen_tokens, 0, INT_DTYPE_STR), max_gen_tokens, - shape(x, -1, INT_DTYPE_STR) - ]) - return unpacked_x.view(out_shape, zero_is_placeholder=False) - - -def _process_logits_and_hidden_states( - model: Module, logits: Tensor, hidden_states: Tensor, - kwargs: dict) -> Tuple[Tensor, Tensor, Tensor, Tensor]: - """ - Process the logits and hidden_states correctly. - For logits: - Can be all context, all gen or mixed. - For all context-phase: - the shape is [bs, V], just process to probs - For all gen-phase: - the shape is [sum(num_gen_tokens), V] - gather using num_gen_tokens => [gen_bs, max_gen_tokens, V] - then typical processing as above - For mixed case: - split the logits, do both ctx and gen phase processing - For hidden_states: - context phase: similar processing - gen-phase: filter based on accepted beams and their lengths. - """ - if model is not None: - num_beams = model.num_beams - beam_length = model.beam_length - greedy_search = model.greedy_search - else: - num_beams = kwargs['num_beams'] - beam_length = kwargs['beam_length'] - greedy_search = kwargs.get('greedy_search', False) - device_request_types = kwargs['device_request_types'] - inverted_temperature = kwargs['redrafter_inverted_temperature'] # [bs] - num_gen_tokens = kwargs[ - 'spec_decoding_params'].spec_decoding_generation_lengths - assert default_net( - ).plugin_config.remove_input_padding, "ReDrafter is only supported without input padding." - """ - Split the flattened data: context and generation - Process them separately. - NOTE: Involves processing 0-shaped tensors (if all context or all generation) - """ - # process context - const_0 = constant_to_tensor_(0, to_array=False) - bs = shape(device_request_types, 0, INT_DTYPE_STR) - num_gen = sum(device_request_types, -1) - num_gen = maximum(constant_to_tensor_(0, to_array=False), num_gen) - num_gen = minimum(bs, num_gen) - bs_ctx = bs - num_gen - ctx_idxs = arange(const_0, bs_ctx, dtype='int32') - assert bs_ctx.rank() == 0 - ctx_logits = index_select(logits, dim=0, index=ctx_idxs) - if not greedy_search: - ctx_temperature = index_select(inverted_temperature, - dim=0, - index=ctx_idxs) - ctx_temperature = unsqueeze(ctx_temperature, 1) - ctx_logits = ctx_logits * ctx_temperature - ctx_probs = _ctx_logits2probs(ctx_logits, greedy_search) - ctx_hidden_states = index_select(hidden_states, dim=0, index=ctx_idxs) - # we accept zero draft tokens for ctx-phase - ctx_num_accepted = expand(constant_to_tensor_(0), unsqueeze(bs_ctx, 0)) - ctx_accepted_beam_index = expand(constant_to_tensor_(0), - unsqueeze(bs_ctx, 0)) - - # process generation - # get the logits[bs_ctx:, :] and hidden_states[bs_ctx:, :] - gen_token_idxs = arange(bs_ctx, - shape(logits, 0, INT_DTYPE_STR), - dtype='int32') - gen_logits = index_select(logits, dim=0, index=gen_token_idxs) - gen_hidden = index_select(hidden_states, dim=0, index=gen_token_idxs) - max_index_allowed = shape(gen_logits, 0, INT_DTYPE_STR) - 1 - gen_unpack_idxs, max_gen_tokens = _get_gen_token_indices_for_unpack( - num_gen_tokens, num_beams, beam_length, max_index_allowed) - gen_logits = _unpack_gen_data(gen_logits, num_gen_tokens, gen_unpack_idxs, - max_gen_tokens) - if not greedy_search: - gen_temperature = index_select(inverted_temperature, - dim=0, - index=gen_token_idxs) - gen_temperature = expand_dims(gen_temperature, dim=[1, 2]) - expanded_gen_temperature = expand(gen_temperature, shape(gen_logits)) - gen_logits = gen_logits * expanded_gen_temperature - gen_hidden = _unpack_gen_data(gen_hidden, num_gen_tokens, gen_unpack_idxs, - max_gen_tokens) - - # verify the input draft tokens (from last step) using the gen_logits - gen_probs, gen_hidden_states, gen_num_accepted, gen_accepted_beam_index\ - = _process_gen_logits( - gen_logits, gen_hidden, kwargs['draft_probs'], - kwargs['draft_tokens'], kwargs['draft_indices'], - num_beams, beam_length, greedy_search, - kwargs.get('rand_data_validation', None) - ) - - # combine ctx and gen phase outputs - probs = concat([ctx_probs, gen_probs], dim=0) - drafter_input = concat([ctx_hidden_states, gen_hidden_states], dim=0) - num_accepted_tokens = concat([ctx_num_accepted, gen_num_accepted], dim=0) - accepted_beam_index = concat( - [ctx_accepted_beam_index, gen_accepted_beam_index], dim=0) - - # NOTE: This is needed with shape inference of data-dependent tensors - bs = shape(device_request_types, 0, INT_DTYPE_STR) - const_0 = constant_to_tensor_(0, to_array=False) - bidxs = arange(const_0, bs, dtype='int32') - probs = index_select(probs, dim=0, index=bidxs) - drafter_input = index_select(drafter_input, dim=0, index=bidxs) - num_accepted_tokens = index_select(num_accepted_tokens, dim=0, index=bidxs) - accepted_beam_index = index_select(accepted_beam_index, dim=0, index=bidxs) - return probs, drafter_input, num_accepted_tokens, accepted_beam_index diff --git a/tensorrt_llm/models/stdit/__init__.py b/tensorrt_llm/models/stdit/__init__.py deleted file mode 100644 index 0d106f45ce46..000000000000 --- a/tensorrt_llm/models/stdit/__init__.py +++ /dev/null @@ -1,14 +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. diff --git a/tensorrt_llm/models/stdit/config.py b/tensorrt_llm/models/stdit/config.py deleted file mode 100644 index ca16dbdfb490..000000000000 --- a/tensorrt_llm/models/stdit/config.py +++ /dev/null @@ -1,115 +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. -from typing import Any, Dict, Optional, Sequence - -from ...mapping import Mapping -from ..modeling_utils import PretrainedConfig, QuantConfig - - -class STDiTModelConfig(PretrainedConfig): - - def __init__(self, - architecture: str = 'STDiT3', - checkpoint_path: str = 'pretrained_ckpt/model.safetensors', - vae_type: str = "hpcai-tech/OpenSora-VAE-v1.2", - text_encoder_type: str = "DeepFloyd/t5-v1_1-xxl", - caption_channels: int = 4096, - num_hidden_layers: int = 28, - hidden_size: int = 1152, - width: int = 640, - height: int = 360, - num_frames: int = 102, - latent_size: Sequence[int] = [30, 45, 80], - stdit_patch_size: Sequence[int] = [1, 2, 2], - spatial_patch_size: Sequence[int] = [1, 8, 8], - temporal_patch_size: Sequence[int] = [4, 1, 1], - in_channels: int = 4, - input_sq_size: int = 512, - num_attention_heads: int = 16, - mlp_ratio: float = 4.0, - class_dropout_prob: float = 0.1, - model_max_length: int = 300, - learn_sigma: bool = True, - qk_norm: bool = True, - skip_y_embedder: bool = False, - dtype: Optional[str] = None, - mapping: Mapping = Mapping(), - quant_config: Optional[QuantConfig] = None, - **kwargs): - kwargs.update({ - 'architecture': architecture, - 'num_hidden_layers': num_hidden_layers, - 'num_attention_heads': num_attention_heads, - 'hidden_size': hidden_size, - 'dtype': dtype - }) - - super().__init__(**kwargs) - self.checkpoint_path = checkpoint_path - self.vae_type = vae_type - self.text_encoder_type = text_encoder_type - self.caption_channels = caption_channels - self.width = width - self.height = height - self.num_frames = num_frames - self.latent_size = latent_size - self.stdit_patch_size = stdit_patch_size - self.spatial_patch_size = spatial_patch_size - self.temporal_patch_size = temporal_patch_size - self.in_channels = in_channels - self.input_sq_size = input_sq_size - self.mlp_ratio = mlp_ratio - self.class_dropout_prob = class_dropout_prob - self.model_max_length = model_max_length - self.learn_sigma = learn_sigma - self.qk_norm = qk_norm - self.skip_y_embedder = skip_y_embedder - self.mapping = mapping - self.quant_config = quant_config - - @classmethod - def from_input_config(cls, - input_config: Dict[str, Any], - dtype: str = 'auto', - mapping: Mapping = Mapping(), - quant_config: Optional[QuantConfig] = None, - **kwargs): - return cls(architecture=input_config['architecture'], - checkpoint_path=input_config['checkpoint_path'], - vae_type=input_config['vae_type'], - text_encoder_type=input_config['text_encoder_type'], - caption_channels=input_config['caption_channels'], - num_hidden_layers=input_config['num_hidden_layers'], - width=input_config['width'], - height=input_config['height'], - num_frames=input_config['num_frames'], - latent_size=input_config['latent_size'], - hidden_size=input_config['hidden_size'], - stdit_patch_size=input_config['stdit_patch_size'], - spatial_patch_size=input_config['spatial_patch_size'], - temporal_patch_size=input_config['temporal_patch_size'], - in_channels=input_config['in_channels'], - input_sq_size=input_config['input_sq_size'], - num_attention_heads=input_config['num_attention_heads'], - mlp_ratio=input_config['mlp_ratio'], - class_dropout_prob=input_config['class_dropout_prob'], - model_max_length=input_config['model_max_length'], - learn_sigma=input_config['learn_sigma'], - qk_norm=input_config['qk_norm'], - skip_y_embedder=input_config['skip_y_embedder'], - dtype=dtype, - mapping=mapping, - quant_config=quant_config, - **kwargs) diff --git a/tensorrt_llm/models/stdit/model.py b/tensorrt_llm/models/stdit/model.py deleted file mode 100644 index 938f80186842..000000000000 --- a/tensorrt_llm/models/stdit/model.py +++ /dev/null @@ -1,1624 +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. - -import builtins -import collections -import functools -import json -import math -import os -import re -from collections import OrderedDict -from typing import Optional - -import numpy as np -import tensorrt as trt -import torch -from tqdm import tqdm - -import tensorrt_llm -from tensorrt_llm._common import default_net -from tensorrt_llm._utils import str_dtype_to_trt, trt_dtype_to_str -from tensorrt_llm.functional import (ACT2FN, AttentionMaskType, LayerNormType, - PositionEmbeddingType, Tensor, - constant_to_tensor_) -from tensorrt_llm.layers import (ColumnLinear, Conv3d, LayerNorm, Linear, - RowLinear) -from tensorrt_llm.layers.attention import (Attention, AttentionParams, - BertAttention, KeyValueCacheParams, - bert_attention, layernorm_map) -from tensorrt_llm.layers.normalization import RmsNorm -from tensorrt_llm.llmapi.kv_cache_type import KVCacheType -from tensorrt_llm.mapping import Mapping -from tensorrt_llm.models.generation_mixin import GenerationMixin -from tensorrt_llm.models.model_weights_loader import (ModelWeightsFormat, - ModelWeightsLoader) -from tensorrt_llm.models.modeling_utils import PretrainedConfig, PretrainedModel -from tensorrt_llm.module import Module, ModuleList -from tensorrt_llm.parameter import Parameter -from tensorrt_llm.plugin import current_all_reduce_helper -from tensorrt_llm.quantization import QuantMode - -from ...functional import (allgather, arange, cast, chunk, concat, constant, - cos, div, einsum, exp, expand, expand_dims, - expand_mask, masked_select, matmul, meshgrid2d, pad, - permute, pow, rearrange, repeat, repeat_interleave, - rms_norm, shape, sin, slice, softmax, split, squeeze, - stack, sum, unsqueeze, where) -from .config import STDiTModelConfig - -# [TODO] For now, we only support static shape, which might contains `-1` when inputs are with dynamic shape. -USE_STATIC_SHAPE = True - - -# From PyTorch internals -def _ntuple(n): - - def parse(x): - if isinstance(x, collections.abc.Iterable) and not isinstance(x, str): - return tuple(x) - return tuple([x] * n) - - return parse - - -to_1tuple = _ntuple(1) -to_2tuple = _ntuple(2) -to_3tuple = _ntuple(3) -to_4tuple = _ntuple(4) -to_ntuple = _ntuple - - -# [TODO] make constant `1` compatible with `scale` -def t2i_modulate(x, shift, scale): - return x * (1.0 + scale) + shift - - -class ModuleSequential(ModuleList): - - def __init__(self, modules) -> None: - super(ModuleSequential, self).__init__(modules=modules) - - def forward(self, *args, **kwargs): - module = self.__getitem__(0) - outputs = module(*args, **kwargs) - for idx in range(1, len(self._modules)): - module = self.__getitem__(idx) - outputs = module(outputs) - return outputs - - -class Activation(Module): - - def __init__(self, act_fn='silu'): - super().__init__() - self.act_fn = act_fn - - def forward(self, input: Tensor): - return ACT2FN[self.act_fn](input) - - -class RotaryEmbedder(Module): - - def __init__(self, - dim, - theta=10000, - interpolate_factor=1., - theta_rescale_factor=1., - seq_before_head_dim=False, - cache_if_possible=True, - use_xpos=False, - dtype=None, - mapping=Mapping(), - quant_mode=QuantMode(0)): - super().__init__() - theta *= theta_rescale_factor**(dim / (dim - 2)) - freqs = 1. / (theta - **(torch.arange(0, dim, 2)[:(dim // 2)].float() / dim)) - self.freqs = Parameter(freqs, dtype=dtype) - self.cached_freqs = None - self.seq_before_head_dim = seq_before_head_dim - self.default_seq_dim = -3 if seq_before_head_dim else -2 - self.scale = (torch.arange(0, dim, 2) + 0.4 * dim) / (1.4 * dim) - assert interpolate_factor >= 1. - self.interpolate_factor = interpolate_factor - - self.cache_if_possible = cache_if_possible - self.use_xpos = use_xpos - self.dtype = dtype - self.mapping = mapping - self.quant_mode = quant_mode - - def get_freqs(self, - t: Tensor, - seq_len: Optional[int] = None, - offset: int = 0): - should_cache = self.cache_if_possible and seq_len is not None - if should_cache and isinstance(self.cached_freqs, Tensor): - if (offset + seq_len) <= self.cached_freqs.shape[0]: - return slice(self.cached_freqs, - starts=[offset] + - [0] * len(self.cached_freqs.shape[1:]), - sizes=[seq_len, *self.cached_freqs.shape[1:]]) - freqs = self.freqs.value - freqs = unsqueeze(t, axis=-1) * unsqueeze(freqs, axis=0) - freqs = repeat_interleave(freqs, repeats=2, dim=(freqs.ndim() - 1)) - if should_cache: - self.cached_freqs = freqs - return freqs - - def get_seq_pos(self, seq_len: int, dtype: trt.DataType, offset: int = 0): - return (arange(start=0, end=seq_len, dtype=trt_dtype_to_str(dtype)) + - offset) / self.interpolate_factor - - def rotate_half(self, x: Tensor): - x = x.view([*x.shape[:-1], x.shape[-1] // 2, 2]) - x1, x2 = x.unbind(x.ndim() - 1) - x = stack([-1 * x2, x1], dim=-1) - x = x.view([*x.shape[:-2], x.shape[-2] * x.shape[-1]]) - return x - - def apply_rotary_emb(self, - freqs: Tensor, - t: Tensor, - start_index: int = 0, - scale: int = 1., - seq_dim: int = -2): - if t.ndim() == 3: - seq_len = t.shape[seq_dim] - # freqs = freqs[-seq_len:] - freqs = slice(starts=[freqs.shape[0] - seq_len], sizes=[seq_len]) - rot_dim = freqs.shape[-1] - end_index = start_index + rot_dim - assert rot_dim <= t.shape[-1], f'feature dimension {t.shape[-1]} is not of sufficient size ' + \ - 'to rotate in all the positions {rot_dim}' - t_left = slice(t, - starts=[0] * t.ndim(), - sizes=[*t.shape[:-1], start_index]) - t_right = slice(t, - starts=[0] * (t.ndim() - 1) + [end_index], - sizes=[*t.shape[:-1], t.shape[-1] - end_index]) - - t = (t * cos(freqs) * scale) + (self.rotate_half(t) * sin(freqs) * - scale) - return concat([t_left, t, t_right], dim=-1) - - def rotate_queries_or_keys(self, - t: Tensor, - seq_dim: Optional[int] = None, - offset: int = 0, - freq_seq_len: Optional[int] = None): - seq_dim = self.default_seq_dim if seq_dim is None else seq_dim - assert not self.use_xpos, 'you must use `.rotate_queries_and_keys` method ' + \ - 'instead and pass in both queries and keys, for ' + \ - 'length extrapolatable rotary embeddings' - - seq_len = t.shape[seq_dim] - if freq_seq_len is not None: - assert freq_seq_len >= seq_len - seq_len = freq_seq_len - - freqs = self.get_freqs(self.get_seq_pos(seq_len, - dtype=t.dtype, - offset=offset), - seq_len=seq_len, - offset=offset) - - if seq_dim == -3: - freqs = rearrange(freqs, 'n d -> n 1 d') - rope_output = self.apply_rotary_emb(freqs, t, seq_dim=seq_dim) - return rope_output - - -class STDiTRmsNorm(RmsNorm): - - def __init__(self, - normalized_shape, - num_groups=1, - eps=1e-06, - elementwise_affine=True, - dtype=None): - super().__init__(normalized_shape, num_groups, eps, elementwise_affine, - dtype) - - def forward(self, hidden_states): - weight = None if self.weight is None else self.weight.value - return rms_norm(input=hidden_states, - normalized_shape=self.normalized_shape, - num_groups=self.num_groups, - weight=weight, - eps=self.eps) - - -class STDiTAttention(BertAttention): - - def __init__(self, - hidden_size, - num_attention_heads, - qk_layernorm=True, - layernorm_type=LayerNormType.RmsNorm, - layernorm_eps=1e-06, - bias=True, - rotary_embedding_func=None, - dtype=None, - tp_group=None, - tp_size=1, - tp_rank=0, - cp_group=None, - cp_size=1, - quant_mode: QuantMode = QuantMode(0)): - assert hidden_size % num_attention_heads == 0, "hidden_size should be divisible by num_attention_heads" - super().__init__(hidden_size=hidden_size, - num_attention_heads=num_attention_heads, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - tp_rank=tp_rank, - cp_group=cp_group, - cp_size=cp_size, - quant_mode=quant_mode) - - self.qk_layernorm = qk_layernorm - if self.qk_layernorm: - ln_type = layernorm_map[layernorm_type] - self.q_layernorm = ln_type(self.attention_head_size, - eps=layernorm_eps, - dtype=dtype) - self.k_layernorm = ln_type(self.attention_head_size, - eps=layernorm_eps, - dtype=dtype) - self.rotary_embedding_func = rotary_embedding_func - - def forward(self, - hidden_states: Tensor, - attention_mask=None, - input_lengths=None, - max_input_length: int = None): - - assert isinstance(hidden_states, Tensor) - - B = shape(hidden_states, 0) - N = shape(hidden_states, 1) - C = shape(hidden_states, 2) - input_lengths = expand(unsqueeze(N, 0).cast('int32'), unsqueeze(B, 0)) - - assert (self.qkv is not None) - qkv = self.qkv(hidden_states) - - kv_size = self.attention_head_size * self.num_attention_kv_heads - query, key, value = split( - qkv, [self.attention_hidden_size, kv_size, kv_size], dim=2) - query = query.view( - concat([B, N, self.num_attention_heads, - self.attention_head_size])).permute(dims=[0, 2, 1, 3]) - key = key.view( - concat( - [B, N, self.num_attention_kv_heads, - self.attention_head_size])).permute(dims=[0, 2, 1, 3]) - - if self.qk_layernorm: - query = self.q_layernorm(query) - key = self.k_layernorm(key) - - if self.rotary_embedding_func is not None: - query = self.rotary_embedding_func(query) - key = self.rotary_embedding_func(key) - - # TODO deal with qkv - query = query.permute(dims=[0, 2, 1, 3]).view( - concat([B, N, self.attention_hidden_size])) - key = key.permute(dims=[0, 2, 1, 3]).view(concat([B, N, kv_size])) - qkv = concat([query, key, value], dim=2) - - if default_net().plugin_config.bert_attention_plugin: - # TRT plugin mode - assert input_lengths is not None - assert self.cp_size == 1 - if default_net().plugin_config.remove_input_padding: - qkv = qkv.view( - concat([-1, self.attention_hidden_size + 2 * kv_size])) - max_input_length = constant( - np.zeros([ - max_input_length, - ], dtype=np.int32)) - context = bert_attention(qkv, - input_lengths, - self.num_attention_heads, - self.attention_head_size, - q_scaling=self.q_scaling, - max_distance=self.max_distance, - max_input_length=max_input_length) - else: - # plain TRT mode - def transpose_for_scores(x): - new_x_shape = concat([ - shape(x, 0), - shape(x, 1), self.num_attention_heads, - self.attention_head_size - ]) - return x.view(new_x_shape).permute([0, 2, 1, 3]) - - kv_size = self.attention_head_size * self.num_attention_kv_heads - query, key, value = split( - qkv, [self.attention_hidden_size, kv_size, kv_size], dim=2) - if self.cp_size > 1 and self.cp_group is not None: - key = allgather(key, self.cp_group, gather_dim=1) - value = allgather(value, self.cp_group, gather_dim=1) - query = transpose_for_scores(query) - key = transpose_for_scores(key) - value = transpose_for_scores(value) - - key = key.permute([0, 1, 3, 2]) - attention_scores = matmul(query, key, use_fp32_acc=False) - attention_scores = attention_scores / (self.q_scaling * - self.norm_factor) - - if attention_mask is not None: - attention_mask = expand_mask(attention_mask, shape(query, 2)) - attention_mask = cast(attention_mask, attention_scores.dtype) - attention_scores = attention_scores + attention_mask - - attention_probs = softmax(attention_scores, dim=-1) - - context = matmul(attention_probs, value, - use_fp32_acc=False).permute([0, 2, 1, 3]) - context = context.view( - concat([ - shape(context, 0), - shape(context, 1), self.attention_hidden_size - ])) - - context = self.dense(context) - context = context.view(concat([B, N, C])) - return context - - -class STDiTCrossAttention(Attention): - - def __init__(self, - *, - local_layer_idx, - hidden_size, - num_attention_heads, - attention_mask_type=AttentionMaskType.causal, - qkv_bias=True, - dense_bias=True, - position_embedding_type=PositionEmbeddingType.learned_absolute, - dtype=None, - tp_group=None, - tp_size=1, - tp_rank=0, - cp_group=[0], - cp_size=1, - cp_rank=0, - quant_mode: QuantMode = QuantMode(0)): - assert hidden_size % num_attention_heads == 0, "hidden_size should be divisible by num_attention_heads" - super().__init__(local_layer_idx=local_layer_idx, - hidden_size=hidden_size, - num_attention_heads=num_attention_heads, - attention_mask_type=attention_mask_type, - bias=qkv_bias, - dense_bias=dense_bias, - cross_attention=True, - position_embedding_type=position_embedding_type, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - tp_rank=tp_rank, - cp_group=cp_group, - cp_size=cp_size, - cp_rank=cp_rank, - quant_mode=quant_mode) - - def forward(self, - hidden_states: Tensor, - encoder_output: Tensor, - use_cache=False, - attention_params: Optional[AttentionParams] = None, - kv_cache_params: Optional[KeyValueCacheParams] = None): - bs = shape(encoder_output, 0) - encoder_input_length = shape(encoder_output, 1) - encoder_hidden_size = shape(encoder_output, 2) - encoder_output = encoder_output.view( - concat([bs * 2, encoder_input_length // 2, encoder_hidden_size])) - - if default_net().plugin_config.remove_input_padding: - B = shape(hidden_states, 0) - N = shape(hidden_states, 1) - C = shape(hidden_states, 2) - hidden_states = hidden_states.view(concat([B * N, C])) - encoder_output = encoder_output.view( - concat([-1, encoder_hidden_size])) - - context = super().forward(hidden_states=hidden_states, - encoder_output=encoder_output, - use_cache=use_cache, - attention_params=attention_params, - kv_cache_params=kv_cache_params) - context = context.view(concat([B, -1, C])) - return context - - -class T2IFinalLayer(Module): - - def __init__(self, - hidden_size, - num_patch, - out_channels, - d_t=None, - d_s=None, - dtype=None, - mapping=Mapping(), - quant_mode=QuantMode(0)): - super().__init__() - self.norm_final = LayerNorm(hidden_size, - elementwise_affine=False, - eps=1e-6, - dtype=dtype) - self.linear = Linear(hidden_size, - num_patch * out_channels, - bias=True, - dtype=dtype) - self.scale_shift_table = Parameter(torch.randn(2, hidden_size) / - hidden_size**0.5, - dtype=dtype) - self.out_channels = out_channels - self.d_t = d_t - self.d_s = d_s - self.dtype = dtype - self.mapping = mapping - self.quant_mode = quant_mode - - def t_mask_select(self, x_mask, x, masked_x, T: int, S: int): - # x: [B, (T, S), C] - # mased_x: [B, (T, S), C] - # x_mask: [B, T] - x = rearrange(x, "B (T S) C -> B T S C", T=T, S=S) - masked_x = rearrange(masked_x, "B (T S) C -> B T S C", T=T, S=S) - x = where(expand_dims(x_mask, - [x_mask.ndim(), x_mask.ndim() + 1]), x, masked_x) - x = rearrange(x, "B T S C -> B (T S) C") - return x - - def forward(self, - x, - t, - x_mask=None, - t0=None, - T: Optional[int] = None, - S: Optional[int] = None): - if T is None: - T = self.d_t - if S is None: - S = self.d_s - shift, scale = chunk(expand_dims(self.scale_shift_table.value, 0) + - expand_dims(t, 1), - chunks=2, - dim=1) - x = t2i_modulate(self.norm_final(x), shift, scale) - if x_mask is not None: - shift_zero, scale_zero = chunk( - expand_dims(self.scale_shift_table.value, 0) + - expand_dims(t0, 1), - chunks=2, - dim=1) - x_zero = t2i_modulate(self.norm_final(x), shift_zero, scale_zero) - x = self.t_mask_select(x_mask, x, x_zero, T, S) - x = self.linear(x) - self.register_network_output('output', x) - return x - - -class PositionEmbedding2D(Module): - - def __init__(self, - dim: int, - dtype=None, - mapping=Mapping(), - quant_mode=QuantMode(0)): - super().__init__() - self.dim = dim - assert dim % 4 == 0, "dim must be divisible by 4" - half_dim = dim // 2 - self.inv_freq = Parameter( - 1.0 / (10000**(torch.arange(0, half_dim, 2).float() / half_dim)), - is_buffer=True, - dtype=dtype) - self.dtype = dtype - self.mapping = mapping - self.quant_mode = quant_mode - - def _get_sin_cos_emb(self, t): - out = einsum("i,d->id", [t, self.inv_freq.value]) - emb_cos = cos(out) - emb_sin = sin(out) - return concat([emb_sin, emb_cos], dim=-1) - - @functools.lru_cache(maxsize=512) - def _get_cached_emb( - self, - dtype, - h: int, - w: int, - scale: Tensor, - base_size: Optional[int] = None, - ): - grid_h = div(arange(0, h, 'float32'), scale.cast('float32')) - grid_w = div(arange(0, w, 'float32'), scale.cast('float32')) - if base_size is not None: - grid_h *= float(base_size) / h - grid_w *= float(base_size) / w - grid_h, grid_w = meshgrid2d(grid_w, grid_h) # here w goes first - grid_h = permute(grid_h, [1, 0]).flatten() - grid_w = permute(grid_w, [1, 0]).flatten() - emb_h = self._get_sin_cos_emb(grid_h) - emb_w = self._get_sin_cos_emb(grid_w) - return unsqueeze(concat([emb_h, emb_w], dim=-1), 0).cast(dtype) - - def forward(self, x, h: int, w: int, scale: Tensor, base_size=None): - pos_embedding = self._get_cached_emb(x.dtype, h, w, scale, base_size) - self.register_network_output('output', pos_embedding) - return pos_embedding - - -class TimestepEmbedder(Module): - - def __init__(self, - hidden_size, - frequency_embedding_size=256, - dtype=None, - mapping=Mapping(), - quant_mode=QuantMode(0)): - super().__init__() - self.mlp = ModuleSequential([ - Linear(frequency_embedding_size, - hidden_size, - bias=True, - dtype=dtype), - Activation('silu'), - Linear(hidden_size, hidden_size, bias=True, dtype=dtype) - ]) - self.frequency_embedding_size = frequency_embedding_size - self.dtype = dtype - self.mapping = mapping - self.quant_mode = quant_mode - - @staticmethod - def timestep_embedding(t, dim, max_period=10000, dtype=None): - half = dim // 2 - freqs = exp( - -math.log(max_period) * - arange(start=0, end=half, dtype=trt_dtype_to_str(trt.float32)) / - constant(np.array([half], dtype=np.float32))) - args = unsqueeze(t, -1).cast(trt.float32) * unsqueeze(freqs, 0) - embedding = concat([cos(args), sin(args)], dim=-1) - if dtype is not None: - embedding = embedding.cast(dtype) - if dim % 2: - embedding = pad(embedding, (0, 0, 0, 1)) - return embedding - - def forward(self, t, dtype): - t_freq = self.timestep_embedding( - t, self.frequency_embedding_size).cast(dtype) - t_emb = self.mlp(t_freq) - return t_emb - - -class SizeEmbedder(TimestepEmbedder): - - def __init__(self, - hidden_size, - frequency_embedding_size=256, - dtype=str_dtype_to_trt("float16"), - mapping=Mapping(), - quant_mode=QuantMode(0)): - super().__init__(hidden_size=hidden_size, - frequency_embedding_size=frequency_embedding_size, - dtype=dtype, - mapping=mapping, - quant_mode=quant_mode) - self.mlp = ModuleSequential([ - Linear(frequency_embedding_size, - hidden_size, - bias=True, - dtype=dtype), - Activation('silu'), - Linear(hidden_size, hidden_size, bias=True, dtype=dtype) - ]) - self.outdim = hidden_size - - def forward(self, s, bs: int): - if not USE_STATIC_SHAPE: - raise NotImplementedError('Only static shape is supported') - if s.ndim() == 1: - s = unsqueeze(s, 1) - assert s.ndim() == 2 - if s.shape[0] != bs: - s = repeat(s, [bs // s.shape[0], 1]) - assert s.shape[0] == bs - b, dims = s.shape[0], s.shape[1] - s = rearrange(s, "b d -> (b d)") - s_freq = self.timestep_embedding(s, self.frequency_embedding_size).cast( - self.dtype) - s_emb = self.mlp(s_freq) - s_emb = rearrange(s_emb, - "(b d) d2 -> b (d d2)", - b=b, - d=dims, - d2=self.outdim) - self.register_network_output('output', s_emb) - return s_emb - - -class CaptionMLP(Module): - - def __init__( - self, - in_features, - hidden_features=None, - out_features=None, - act_layer="gelu", - bias=True, - dtype=None, - mapping=Mapping(), - quant_mode=QuantMode(0), - inner_layernorm=False, - eps=1e-05, - ): - super().__init__() - hidden_act = act_layer - out_features = out_features or in_features - hidden_features = hidden_features or in_features - bias = to_2tuple(bias) - if hidden_act not in ACT2FN: - raise ValueError( - 'unsupported activation function: {}'.format(hidden_act)) - fc_output_size = 2 * hidden_features if hidden_act in [ - 'swiglu', 'gegelu' - ] else hidden_features - self.inner_layernorm = LayerNorm(hidden_features, dtype=dtype, - eps=eps) if inner_layernorm else None - - self.fc1 = ColumnLinear(in_features, - fc_output_size, - bias=bias[0], - dtype=dtype, - tp_group=mapping.tp_group, - tp_size=mapping.tp_size, - gather_output=False) - self.fc2 = RowLinear(hidden_features, - out_features, - bias=bias[1], - dtype=dtype, - tp_group=mapping.tp_group, - tp_size=mapping.tp_size) - - self.in_features = in_features - self.hidden_features = hidden_features - self.out_features = out_features - self.hidden_act = hidden_act - self.dtype = dtype - self.bias = bias - self.mapping = mapping - self.quant_mode = quant_mode - self.eps = eps - - def forward(self, hidden_states, gegelu_limit=None): - inter = self.fc1(hidden_states) - if self.hidden_act == 'gegelu': - inter = ACT2FN[self.hidden_act](inter, gegelu_limit) - else: - inter = ACT2FN[self.hidden_act](inter) - if self.inner_layernorm is not None: - inter = self.inner_layernorm(inter) - output = self.fc2(inter) - return output - - -class CaptionEmbedder(Module): - - def __init__(self, - in_channels, - hidden_size, - uncond_prob, - act_layer='gelu', - token_num=120, - dtype=None, - mapping=Mapping(), - quant_mode=QuantMode(0)): - super().__init__() - self.y_proj = CaptionMLP( - in_features=in_channels, - hidden_features=hidden_size, - out_features=hidden_size, - act_layer=act_layer, - mapping=mapping, - dtype=dtype, - ) - self.y_embedding = Parameter(torch.randn(token_num, in_channels) / - in_channels**0.5, - dtype=dtype) - self.uncond_prob = uncond_prob - self.dtype = dtype - self.mapping = mapping - self.quant_mode = quant_mode - - def token_drop(self, caption, force_drop_ids=None): - if not USE_STATIC_SHAPE: - raise NotImplementedError('Only static shape is supported') - assert (isinstance(force_drop_ids, torch.Tensor) - or isinstance(force_drop_ids, np.array)) - if force_drop_ids is None: - drop_ids = torch.rand(caption.shape[0]).cuda() < self.uncond_prob - else: - drop_ids = torch.Tensor(force_drop_ids) == 1 - drop_ids = constant(drop_ids.cpu().numpy()) - caption = where(expand_dims(drop_ids, [1, 2, 3]), - self.y_embedding.value, caption) - return caption - - def forward(self, caption, force_drop_ids=None): - if force_drop_ids is not None: - caption = self.token_drop(caption, force_drop_ids) - caption = self.y_proj(caption) - self.register_network_output('output', caption) - return caption - - -class PatchEmbed3D(Module): - - def __init__(self, - patch_size=(2, 4, 4), - in_chans=3, - embed_dim=96, - norm_layer=None, - flatten=True, - dtype=None, - mapping=Mapping(), - quant_mode=QuantMode(0)): - super().__init__() - self.patch_size = patch_size - self.flatten = flatten - self.in_chans = in_chans - self.embed_dim = embed_dim - - self.proj = Conv3d(in_chans, - embed_dim, - kernel_size=patch_size, - stride=patch_size, - dtype=dtype) - if norm_layer is not None: - self.norm = norm_layer(embed_dim) - else: - self.norm = None - self.dtype = dtype - self.mapping = mapping - self.quant_mode = quant_mode - - def forward(self, x): - if not USE_STATIC_SHAPE: - raise NotImplementedError('Only static shape is supported') - _, _, D, H, W = x.shape - if W % self.patch_size[2] != 0: - x = pad(x, (0, self.patch_size[2] - W % self.patch_size[2])) - if H % self.patch_size[1] != 0: - x = pad(x, (0, 0, 0, self.patch_size[1] - H % self.patch_size[1])) - if D % self.patch_size[0] != 0: - x = pad( - x, (0, 0, 0, 0, 0, self.patch_size[0] - D % self.patch_size[0])) - x = self.proj(x) # (B C T H W) - if self.norm is not None: - D = shape(x, 2) - Wh = shape(x, 3) - Ww = shape(x, 4) - x = x.flatten(2).transpose(1, 2) - x = self.norm(x) - x = x.transpose(1, 2).view([-1, self.embed_dim, D, Wh, Ww]) - if self.flatten: - x = x.flatten(2).transpose(1, 2) # BCTHW -> BNC - self.register_network_output('output', x) - return x - - -class STDiT3Block(Module): - - def __init__(self, - hidden_size, - num_heads, - mlp_ratio=4.0, - rope=None, - qk_norm=False, - temporal=False, - dtype=None, - local_layer_idx=0, - mapping=Mapping(), - quant_mode=QuantMode(0)): - super().__init__() - self.temporal = temporal - self.hidden_size = hidden_size - - attn_cls = STDiTAttention - mha_cls = STDiTCrossAttention - - self.norm1 = LayerNorm(hidden_size, - eps=1e-6, - elementwise_affine=False, - dtype=dtype) - self.attn = attn_cls(hidden_size=hidden_size, - num_attention_heads=num_heads, - qk_layernorm=qk_norm, - bias=True, - rotary_embedding_func=rope, - dtype=dtype, - tp_group=mapping.tp_group, - tp_size=mapping.tp_size, - tp_rank=mapping.tp_rank, - cp_group=mapping.cp_group, - cp_size=mapping.cp_size, - quant_mode=quant_mode) - self.cross_attn = mha_cls( - local_layer_idx=local_layer_idx, - hidden_size=hidden_size, - num_attention_heads=num_heads, - attention_mask_type=tensorrt_llm.layers.AttentionMaskType.causal, - qkv_bias=True, - dense_bias=True, - dtype=dtype, - tp_group=mapping.tp_group, - tp_size=mapping.tp_size, - tp_rank=mapping.tp_rank, - cp_group=mapping.cp_group, - cp_size=mapping.cp_size, - quant_mode=quant_mode) - self.norm2 = LayerNorm(hidden_size, - eps=1e-6, - elementwise_affine=False, - dtype=dtype) - self.mlp = CaptionMLP( - in_features=hidden_size, - hidden_features=int(hidden_size * mlp_ratio), - act_layer='gelu', - mapping=mapping, - dtype=dtype, - ) - self.scale_shift_table = Parameter(torch.randn(6, hidden_size) / - hidden_size**0.5, - dtype=dtype) - self.dtype = dtype - self.mapping = mapping - self.quant_mode = quant_mode - - def t_mask_select(self, x_mask, x, masked_x, T: int, S: int): - # x: [B, (T, S), C] - # mased_x: [B, (T, S), C] - # x_mask: [B, T] - x = rearrange(x, "B (T S) C -> B T S C", T=T, S=S) - masked_x = rearrange(masked_x, "B (T S) C -> B T S C", T=T, S=S) - x = where(expand_dims(x_mask, [2, 3]), x, masked_x) - x = rearrange(x, "B T S C -> B (T S) C") - return x - - def forward( - self, - x, - y, - t, - x_mask=None, # temporal mask - t0=None, # t with timestamp=0 - T: Optional[int] = None, # number of frames - S: Optional[int] = None, # number of pixel patches - attention_params: Optional[AttentionParams] = None, - kv_cache_params: Optional[KeyValueCacheParams] = None, - ): - if not USE_STATIC_SHAPE: - raise NotImplementedError('Only static shape is supported') - # prepare modulate parameters - B, N, C = x.shape - shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = chunk( - expand_dims(self.scale_shift_table.value, 0) + t.view([B, 6, -1]), - chunks=6, - dim=1) - if x_mask is not None: - shift_msa_zero, scale_msa_zero, gate_msa_zero, shift_mlp_zero, scale_mlp_zero, gate_mlp_zero = chunk( - expand_dims(self.scale_shift_table.value, 0) + - t0.view([B, 6, -1]), - chunks=6, - dim=1) - - # modulate (attention) - x_m = t2i_modulate(self.norm1(x), shift_msa, scale_msa) - if x_mask is not None: - x_m_zero = t2i_modulate(self.norm1(x), shift_msa_zero, - scale_msa_zero) - x_m = self.t_mask_select(x_mask, x_m, x_m_zero, T, S) - - # attention - if self.temporal: - x_m = rearrange(x_m, "B (T S) C -> (B S) T C", T=T, S=S) - x_m = self.attn( - x_m, max_input_length=attention_params.max_context_length) - x_m = rearrange(x_m, "(B S) T C -> B (T S) C", T=T, S=S) - else: - x_m = rearrange(x_m, "B (T S) C -> (B T) S C", T=T, S=S) - x_m = self.attn( - x_m, max_input_length=attention_params.max_context_length) - x_m = rearrange(x_m, "(B T) S C -> B (T S) C", T=T, S=S) - - # modulate (attention) - x_m_s = gate_msa * x_m - if x_mask is not None: - x_m_s_zero = gate_msa_zero * x_m - x_m_s = self.t_mask_select(x_mask, x_m_s, x_m_s_zero, T, S) - - # residual - x = x + x_m_s - - # cross attention - cattn = self.cross_attn( - hidden_states=x, - encoder_output=y, - attention_params=AttentionParams( - sequence_length=attention_params.sequence_length, - context_lengths=attention_params.context_lengths, - host_context_lengths=attention_params.host_context_lengths, - max_context_length=attention_params.max_context_length, - host_request_types=attention_params.host_request_types, - encoder_input_lengths=attention_params.encoder_input_lengths, - encoder_max_input_length=attention_params. - encoder_max_input_length, - host_runtime_perf_knobs=attention_params. - host_runtime_perf_knobs, - host_context_progress=attention_params.host_context_progress, - ), - kv_cache_params=KeyValueCacheParams( - past_key_value=kv_cache_params.past_key_value, - host_past_key_value_lengths=kv_cache_params. - host_past_key_value_lengths, - host_max_attention_window_sizes=kv_cache_params. - host_max_attention_window_sizes, - host_sink_token_length=kv_cache_params.host_sink_token_length, - cache_indirection=kv_cache_params.cache_indirection, - kv_cache_block_offsets=None, - host_kv_cache_block_offsets=None, - host_kv_cache_pool_pointers=None, - host_kv_cache_pool_mapping=None, - cross_kv_cache_block_offsets=None, - host_cross_kv_cache_block_offsets=None, - host_cross_kv_cache_pool_pointers=None, - host_cross_kv_cache_pool_mapping=None, - )) - x = x + cattn - - # modulate (MLP) - x_m = t2i_modulate(self.norm2(x), shift_mlp, scale_mlp) - if x_mask is not None: - x_m_zero = t2i_modulate(self.norm2(x), shift_mlp_zero, - scale_mlp_zero) - x_m = self.t_mask_select(x_mask, x_m, x_m_zero, T, S) - - # MLP - x_m = self.mlp(x_m) - - # modulate (MLP) - x_m_s = gate_mlp * x_m - if x_mask is not None: - x_m_s_zero = gate_mlp_zero * x_m - x_m_s = self.t_mask_select(x_mask, x_m_s, x_m_s_zero, T, S) - - # residual - x = x + x_m_s - - return x - - -class STDiT3Model(PretrainedModel): - - def __init__(self, config: STDiTModelConfig): - self.check_config(config) - super().__init__(config) - self.learn_sigma = config.learn_sigma - self.in_channels = config.in_channels - self.out_channels = config.in_channels * 2 if config.learn_sigma else config.in_channels - self.caption_channels = config.caption_channels - self.depth = config.num_hidden_layers - self.mlp_ratio = config.mlp_ratio - self.hidden_size = config.hidden_size - self.num_heads = config.num_attention_heads - self.model_max_length = config.model_max_length - self.latent_size = config.latent_size - self.input_sq_size = config.input_sq_size - self.patch_size = config.stdit_patch_size - self.class_dropout_prob = config.class_dropout_prob - self.qk_norm = config.qk_norm - self.dtype = config.dtype - self.mapping = config.mapping - - self.pos_embed = PositionEmbedding2D(self.hidden_size, dtype=self.dtype) - self.rope = RotaryEmbedder(dim=self.hidden_size // self.num_heads, - dtype=self.dtype) - self.x_embedder = PatchEmbed3D(self.patch_size, - self.in_channels, - self.hidden_size, - dtype=self.dtype) - self.t_embedder = TimestepEmbedder(self.hidden_size, dtype=self.dtype) - self.fps_embedder = SizeEmbedder(self.hidden_size, dtype=self.dtype) - self.t_block = ModuleSequential([ - Activation('silu'), - Linear(self.hidden_size, - 6 * self.hidden_size, - bias=True, - dtype=self.dtype) - ]) - self.y_embedder = CaptionEmbedder(in_channels=self.caption_channels, - hidden_size=self.hidden_size, - uncond_prob=self.class_dropout_prob, - act_layer='gelu', - token_num=self.model_max_length, - dtype=self.dtype) - self.spatial_blocks = ModuleList([ - STDiT3Block(hidden_size=self.hidden_size, - num_heads=self.num_heads, - mlp_ratio=self.mlp_ratio, - qk_norm=self.qk_norm, - dtype=self.dtype, - local_layer_idx=idx, - mapping=self.mapping) for idx in range(self.depth) - ]) - self.temporal_blocks = ModuleList([ - STDiT3Block(hidden_size=self.hidden_size, - num_heads=self.num_heads, - mlp_ratio=self.mlp_ratio, - qk_norm=self.qk_norm, - temporal=True, - rope=self.rope.rotate_queries_or_keys, - dtype=self.dtype, - local_layer_idx=idx, - mapping=self.mapping) for idx in range(self.depth) - ]) - self.final_layer = T2IFinalLayer(self.hidden_size, - np.prod(self.patch_size), - self.out_channels, - dtype=self.dtype, - mapping=self.mapping) - - def check_config(self, config: PretrainedConfig): - config.set_if_not_exist('caption_channels', 4096) - config.set_if_not_exist('num_hidden_layers', 28) - config.set_if_not_exist('latent_size', [30, 45, 80]) - config.set_if_not_exist('hidden_size', 1152) - config.set_if_not_exist('stdit_patch_size', [1, 2, 2]) - config.set_if_not_exist('in_channels', 4) - config.set_if_not_exist('input_sq_size', 512) - config.set_if_not_exist('num_attention_heads', 16) - config.set_if_not_exist('mlp_ratio', 4.0) - config.set_if_not_exist('class_dropout_prob', 0.1) - config.set_if_not_exist('model_max_length', 300) - config.set_if_not_exist('learn_sigma', True) - config.set_if_not_exist('dtype', None) - config.set_if_not_exist('qk_norm', True) - config.set_if_not_exist('skip_y_embedder', False) - - def __post_init__(self): - return - - def get_dynamic_size(self, x: Tensor): - if not USE_STATIC_SHAPE: - raise NotImplementedError('Only static shape is supported') - _, _, T, H, W = x.shape - if T % self.patch_size[0] != 0: - T += self.patch_size[0] - T % self.patch_size[0] - if H % self.patch_size[1] != 0: - H += self.patch_size[1] - H % self.patch_size[1] - if W % self.patch_size[2] != 0: - W += self.patch_size[2] - W % self.patch_size[2] - T = T // self.patch_size[0] - H = H // self.patch_size[1] - W = W // self.patch_size[2] - return (T, H, W) - - def encode_text(self, y: Tensor, mask: Optional[Tensor] = None): - y = self.y_embedder(y) # [B, 1, N_token, C] - if mask is not None: - if mask.shape[0] != y.shape[0]: - mask = repeat(mask, sizes=(y.shape[0] // mask.shape[0], 1)) - mask = squeeze(squeeze(mask, 1), 1) - y = masked_select( - squeeze(y, 1), - where( - unsqueeze(mask, -1).__eq__( - constant_to_tensor_(0, dtype=mask.dtype)), - constant_to_tensor_(False), - constant_to_tensor_(True))).view((1, -1, self.hidden_size)) - # [TODO] how to convert y_lens to list? - # y_lens = mask.sum(dim=1).tolist() - y_lens = sum(mask, dim=1) - else: - y_lens = constant( - np.array([y.shape[2]] * y.shape[0], dtype=np.int64)) - y = squeeze(y, 1).view((1, -1, self.hidden_size)) - self.register_network_output('encode_text.output.y', y) - self.register_network_output('encode_text.output.y_lens', y_lens) - return y, y_lens - - def unpatchify(self, x: Tensor, N_t: int, N_h: int, N_w: int, R_t: int, - R_h: int, R_w: int): - T_p, H_p, W_p = self.patch_size - x = rearrange( - x, - "B (N_t N_h N_w) (T_p H_p W_p C_out) -> B C_out (N_t T_p) (N_h H_p) (N_w W_p)", - N_t=N_t, - N_h=N_h, - N_w=N_w, - T_p=T_p, - H_p=H_p, - W_p=W_p, - C_out=self.out_channels, - ) - # unpad - x = slice(x, - starts=[0] * x.ndim(), - sizes=concat([ - shape(x, 0), - shape(x, 1), - constant(np.array([R_t, R_h, R_w]).astype(np.int64)) - ])) - self.register_network_output('unpatchify.output', x) - return x - - def forward(self, - x: Tensor, - timestep: Tensor, - y: Tensor, - fps: Tensor, - height: Tensor, - width: Tensor, - mask: Optional[Tensor] = None, - x_mask: Optional[Tensor] = None, - attention_params: Optional[AttentionParams] = None, - kv_cache_params: Optional[KeyValueCacheParams] = None, - **kwargs): - if not USE_STATIC_SHAPE: - raise NotImplementedError('Only static shape is supported') - assert tuple(x.shape[2:]) == tuple( - self.latent_size), "For now only static shape is supported." - B = x.shape[0] - x = x.cast(self.dtype) - timestep = timestep.cast(self.dtype) - y = y.cast(self.dtype) - fps = fps.cast(self.dtype) - - # === get pos embed === - _, _, Tx, Hx, Wx = x.shape - T, H, W = self.get_dynamic_size(x) - S = H * W - base_size = round(S**0.5) - resolution_sq = pow( - height.cast('float32') * width.cast('float32'), - constant_to_tensor_(0.5, dtype='float32')) - scale = (resolution_sq / self.input_sq_size).cast(self.dtype) - pos_emb = self.pos_embed(x, h=H, w=W, scale=scale, base_size=base_size) - - # === get timestep embed === - t = self.t_embedder(timestep, dtype=x.dtype) # [B, C] - fps = self.fps_embedder(unsqueeze(fps, 1), bs=B) - t = t + fps - t_mlp = self.t_block(t) - t0 = t0_mlp = None - if x_mask is not None: - t0_timestep = constant( - np.zeros(shape=timestep.shape).astype(np.float32)).cast(x.dtype) - t0 = self.t_embedder(t0_timestep, dtype=x.dtype) - t0 = t0 + fps - t0_mlp = self.t_block(t0) - - # === get y embed === - if self.config.skip_y_embedder: - y_lens = mask - if isinstance(y_lens, Tensor): - y_lens = y_lens.cast('int64') - else: - y, y_lens = self.encode_text(y, mask) - y_lens = None #[11, 11] - - # === get x embed === - x = self.x_embedder(x) # [B, N, C] - x = rearrange(x, "B (T S) C -> B T S C", T=T, S=S) - x = x + pos_emb - x = rearrange(x, "B T S C -> B (T S) C", T=T, S=S) - - # === blocks === - cnt = 0 - for spatial_block, temporal_block in zip(self.spatial_blocks, - self.temporal_blocks): - x = spatial_block( - x, - y, - t_mlp, - x_mask=x_mask, - t0=t0_mlp, - T=T, - S=S, - attention_params=attention_params, - kv_cache_params=kv_cache_params, - ) - x = temporal_block( - x, - y, - t_mlp, - x_mask=x_mask, - t0=t0_mlp, - T=T, - S=S, - attention_params=attention_params, - kv_cache_params=kv_cache_params, - ) - cnt += 1 - - # === final layer === - x = self.final_layer(x, t, x_mask, t0, T, S) - x = self.unpatchify(x, T, H, W, Tx, Hx, Wx) - - # cast to float32 for better accuracy - output = x.cast('float32') - output.mark_output('output', 'float32') - return output - - def prepare_inputs(self, max_batch_size, **kwargs): - '''@brief: Prepare inputs Tensors for the model, the given sizes are used to determine the - ranges of the dimensions of when using TRT dynamic shapes. - - @return: a list contains values which can be fed into the self.forward() - ''' - - mapping = self.config.mapping - if mapping.tp_size > 1: - current_all_reduce_helper().set_workspace_tensor(mapping, 1) - - def stdit_default_batch_range(max_batch_size): - return [max_batch_size, max_batch_size, max_batch_size] - - default_range = stdit_default_batch_range - # [NOTE] For now only static batch size is supported, so we run the model with max_batch_size. - batch_size = max_batch_size - - x = Tensor(name='x', - dtype=self.dtype, - shape=[batch_size, self.in_channels, *self.latent_size], - dim_range=OrderedDict([ - ('batch_size', [default_range(max_batch_size)]), - ('in_channels', [[self.in_channels] * 3]), - ('latent_frames', [[self.latent_size[0]] * 3]), - ('latent_height', [[self.latent_size[1]] * 3]), - ('latent_width', [[self.latent_size[2]] * 3]), - ])) - timestep = Tensor(name='timestep', - dtype=self.dtype, - shape=[batch_size], - dim_range=OrderedDict([ - ('batch_size', [default_range(max_batch_size)]), - ])) - y = Tensor( - name='y', - dtype=self.dtype, - shape=[batch_size, 1, self.model_max_length, self.caption_channels], - dim_range=OrderedDict([ - ('batch_size', [default_range(max_batch_size)]), - ('mask_batch_size', [[1, 1, 1]]), - ('num_tokens', [[self.model_max_length] * 3]), - ('caption_channels', [[self.caption_channels] * 3]), - ])) - mask = Tensor(name='mask', - dtype=trt.int32, - shape=[1, self.model_max_length], - dim_range=OrderedDict([ - ('mask_batch_size', [[1, 1, 1]]), - ('num_tokens', [[self.model_max_length] * 3]), - ])) - x_mask = Tensor(name='x_mask', - dtype=tensorrt_llm.str_dtype_to_trt('bool'), - shape=[batch_size, self.latent_size[0]], - dim_range=OrderedDict([ - ('batch_size', [default_range(max_batch_size)]), - ('latent_frames', [[self.latent_size[0]] * 3]), - ])) - fps = Tensor(name='fps', dtype=self.dtype, shape=[ - 1, - ]) - height = Tensor(name='height', dtype=self.dtype, shape=[ - 1, - ]) - width = Tensor(name='width', dtype=self.dtype, shape=[ - 1, - ]) - - use_gpt_attention_plugin = default_net( - ).plugin_config.gpt_attention_plugin - remove_input_padding = default_net().plugin_config.remove_input_padding - - cross_attn_batch_size = batch_size - max_cattn_seq_len = int( - np.prod([ - np.ceil(d / p) - for d, p in zip(self.latent_size, self.patch_size) - ])) - max_cattn_enc_len = self.model_max_length - attn_inputs = GenerationMixin().prepare_attention_inputs( - max_batch_size=cross_attn_batch_size, - opt_batch_size=cross_attn_batch_size, - max_beam_width=1, - max_input_len=max_cattn_seq_len, - max_seq_len=max_cattn_seq_len, - num_kv_heads=self.num_heads, - head_size=self.hidden_size // self.num_heads, - num_layers=self.depth, - kv_dtype=self.dtype, - kv_cache_type=KVCacheType.DISABLED, - remove_input_padding=remove_input_padding, - use_gpt_attention_plugin=use_gpt_attention_plugin, - enable_ctx_gen_opt_profiles=False, - mapping=self.mapping, - ) - - sequence_length = attn_inputs['sequence_length'] - host_context_lengths = attn_inputs['host_context_lengths'] - host_max_attention_window_sizes = attn_inputs[ - 'host_max_attention_window_sizes'] - host_sink_token_length = attn_inputs['host_sink_token_length'] - context_lengths = attn_inputs['context_lengths'] - host_request_types = attn_inputs['host_request_types'] - - host_past_key_value_lengths = attn_inputs['host_past_key_value_lengths'] - past_key_value = attn_inputs['past_key_value'] - if past_key_value: - past_key_value = past_key_value[0] - cache_indirection = attn_inputs['cache_indirection'] - host_runtime_perf_knobs_tensor = attn_inputs['host_runtime_perf_knobs'] - host_context_progress = attn_inputs['host_context_progress'] - cross_encoder_input_lengths = Tensor(name='encoder_input_lengths', - shape=(cross_attn_batch_size, ), - dtype=str_dtype_to_trt('int32')) - cross_max_encoder_seq_len = Tensor( - name='encoder_max_input_length', - shape=[-1], - dim_range=OrderedDict([ - ("encoder_max_input_length", - [[1, (max_cattn_enc_len + 1) // 2, max_cattn_enc_len]]) - ]), - dtype=str_dtype_to_trt('int32')) - - attention_params = AttentionParams( - sequence_length=sequence_length, - context_lengths=context_lengths, - host_context_lengths=host_context_lengths, - max_context_length=max_cattn_seq_len, - host_request_types=host_request_types, - encoder_input_lengths=cross_encoder_input_lengths, - encoder_max_input_length=cross_max_encoder_seq_len, - host_runtime_perf_knobs=host_runtime_perf_knobs_tensor, - host_context_progress=host_context_progress) - - kv_cache_params = KeyValueCacheParams( - past_key_value=past_key_value, - host_past_key_value_lengths=host_past_key_value_lengths, - host_max_attention_window_sizes=host_max_attention_window_sizes, - host_sink_token_length=host_sink_token_length, - cache_indirection=cache_indirection, - kv_cache_block_offsets=None, - host_kv_cache_block_offsets=None, - host_kv_cache_pool_pointers=None, - host_kv_cache_pool_mapping=None, - cross_kv_cache_block_offsets=None, - host_cross_kv_cache_block_offsets=None, - host_cross_kv_cache_pool_pointers=None, - host_cross_kv_cache_pool_mapping=None, - ) - - return { - 'x': x, - 'timestep': timestep, - 'y': y, - 'mask': mask, - 'x_mask': x_mask, - 'fps': fps, - 'height': height, - 'width': width, - 'attention_params': attention_params, - 'kv_cache_params': kv_cache_params, - } - - @classmethod - def from_pretrained(cls, - pretrained_model_dir: str, - dtype='float16', - mapping=Mapping(), - **kwargs): - - quant_ckpt_path = kwargs.pop('quant_ckpt_path', None) - assert os.path.exists(f"{pretrained_model_dir}/config.json") - with open(f"{pretrained_model_dir}/config.json", 'r') as f: - hf_config = json.load(f) - - hf_tllm_config_remapping = { - 'model_type': 'architecture', - 'depth': 'num_hidden_layers', - 'num_heads': 'num_attention_heads', - 'patch_size': 'stdit_patch_size', - 'pred_sigma': 'learn_sigma', - } - for hf_key, tllm_key in hf_tllm_config_remapping.items(): - hf_config[tllm_key] = hf_config.pop(hf_key) - hf_config.update(kwargs) - - model_config = STDiTModelConfig.from_input_config(hf_config, - dtype=dtype, - mapping=mapping) - - model_dir = pretrained_model_dir - custom_dict = {} - if quant_ckpt_path is not None: - model_dir = quant_ckpt_path - - loader = STDiT3ModelWeightsLoader(model_dir, custom_dict) - model = cls(model_config) - loader.generate_tllm_weights(model) - return model - - -class STDiT3ModelWeightsLoader(ModelWeightsLoader): - - def translate_to_external_key(self, tllm_key: str, - tllm_to_externel_key_dict: dict): - """Convert and load external checkpoint into a TensorRT LLM model. - """ - trtllm_to_hf_name = { - r"spatial_blocks.(\d+).attn.q_layernorm.weight": - "spatial_blocks.*.attn.q_norm.weight", - r"spatial_blocks.(\d+).attn.k_layernorm.weight": - "spatial_blocks.*.attn.k_norm.weight", - r"spatial_blocks.(\d+).attn.dense.weight": - "spatial_blocks.*.attn.proj.weight", - r"spatial_blocks.(\d+).attn.dense.bias": - "spatial_blocks.*.attn.proj.bias", - r"temporal_blocks.(\d+).attn.q_layernorm.weight": - "temporal_blocks.*.attn.q_norm.weight", - r"temporal_blocks.(\d+).attn.k_layernorm.weight": - "temporal_blocks.*.attn.k_norm.weight", - r"temporal_blocks.(\d+).attn.dense.weight": - "temporal_blocks.*.attn.proj.weight", - r"temporal_blocks.(\d+).attn.dense.bias": - "temporal_blocks.*.attn.proj.bias", - r"spatial_blocks.(\d+).cross_attn.dense.weight": - "spatial_blocks.*.cross_attn.proj.weight", - r"spatial_blocks.(\d+).cross_attn.dense.bias": - "spatial_blocks.*.cross_attn.proj.bias", - r"temporal_blocks.(\d+).cross_attn.dense.weight": - "temporal_blocks.*.cross_attn.proj.weight", - r"temporal_blocks.(\d+).cross_attn.dense.bias": - "temporal_blocks.*.cross_attn.proj.bias", - } - - for k, v in trtllm_to_hf_name.items(): - m = re.match(k, tllm_key) - if m is not None: - matched_pos = m.groups() - placeholders = v.count('*') - assert len(matched_pos) == placeholders - for i in range(len(matched_pos)): - v = v.replace('*', matched_pos[i], 1) - return v - return tllm_key - - def load_tensor(self, key, tp_size=1, tp_dim=-1, tp_rank=0): - hidden_size = self.model.config.hidden_size - - if "attn.qkv" in key: - is_cross_attn = "cross_attn.qkv" in key - if is_cross_attn: - # process for cross attention - process_qkv_names = [ - 'q_linear'.join(key.split('qkv')), - 'kv_linear'.join(key.split('qkv')) - ] - else: - process_qkv_names = [key] - qkv_tensors = [] - for qkv_key in process_qkv_names: - # Retrieve shard index - assert qkv_key in self.shard_map - ptr_idx = self.shard_map[qkv_key] - if self.format == ModelWeightsFormat.SAFETENSORS: - # Force to load Pytorch tensor - tensor = self.shards[ptr_idx].get_tensor(qkv_key) - else: - tensor = self.shards[ptr_idx][qkv_key] - qkv_tensors.append(tensor) - if is_cross_attn: - tensor = torch.concat(qkv_tensors, dim=0) - else: - tensor = qkv_tensors[0] - # Post-process weight and bias if tp_size > 1 - if tp_size > 1: - if "weight" in key: - tensor = tensor.reshape(3, hidden_size, hidden_size) - elif "bias" in key: - tensor = tensor.reshape(3, hidden_size) - tp_dim = 1 - tensor_shape = tensor.shape - else: - # Retrieve shard index - if key in self.shard_map: - ptr_idx = self.shard_map[key] - else: - return None - - if self.format == ModelWeightsFormat.SAFETENSORS: - tensor = self.shards[ptr_idx].get_slice(key) - tensor_shape = tensor.get_shape() - if tensor_shape == []: - tensor = self.shards[ptr_idx].get_tensor(key).unsqueeze(0) - tensor_shape = tensor.shape - else: - tensor = self.shards[ptr_idx][key] - tensor_shape = tensor.shape - - if tp_size <= 1 or tp_dim < 0: - return tensor[:] - else: - if len(tensor_shape) == 1 and (tp_dim > 0 or tensor_shape[0] == 1): - return tensor[:] - else: - width = tensor_shape[tp_dim] - if width == 1: - return tensor[:] - slice_width = math.ceil(width / tp_size) - slice_start = tp_rank * slice_width - slice_end = builtins.min((tp_rank + 1) * slice_width, width) - slice_obj = [builtins.slice(None)] * len(tensor_shape) - slice_obj[tp_dim] = builtins.slice(slice_start, slice_end) - res = tensor[tuple(slice_obj)] - if "qkv.weight" in key: - res = res.reshape(3 * (hidden_size // tp_size), hidden_size) - elif "qkv.bias" in key: - res = res.reshape(3 * (hidden_size // tp_size)) - return res - - def generate_tllm_weights(self, - model, - custom_postprocess_kwargs: dict = {}): - self.update_key_mapping(model) - tp_module_patterns = [ - r'.*_blocks.*.attn.qkv.weight$', - r'.*_blocks.*.attn.qkv.bias$', - r'.*_blocks.*.attn.dense.weight$', - r'.*_blocks.*.cross_attn.qkv.weight$', - r'.*_blocks.*.cross_attn.qkv.bias$', - r'.*_blocks.*.cross_attn.dense.weight$', - r'.*_blocks.*.mlp.fc1.weight$', - r'.*_blocks.*.mlp.fc1.bias$', - r'.*_blocks.*.mlp.fc2.weight$', - ] - tllm_weights = {} - for tllm_key, _ in tqdm(model.named_parameters()): - skip_tp = not any([ - re.match(pattern, tllm_key) is not None - for pattern in tp_module_patterns - ]) - tllm_weights.update( - self.load(tllm_key, - custom_postprocess_kwargs=custom_postprocess_kwargs, - skip_tp=skip_tp)) - self.fill(tllm_weights) diff --git a/tensorrt_llm/models/unet/__init__.py b/tensorrt_llm/models/unet/__init__.py deleted file mode 100644 index 71bf6d298d42..000000000000 --- a/tensorrt_llm/models/unet/__init__.py +++ /dev/null @@ -1,14 +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. diff --git a/tensorrt_llm/models/unet/attention.py b/tensorrt_llm/models/unet/attention.py deleted file mode 100644 index c56605106362..000000000000 --- a/tensorrt_llm/models/unet/attention.py +++ /dev/null @@ -1,332 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -import math -from typing import Optional - -from ..._common import precision -from ...functional import geglu, matmul, softmax, split -from ...layers import Conv2d, GroupNorm, LayerNorm, Linear -from ...module import Module, ModuleList - - -class AttentionBlock(Module): - - def __init__(self, - channels: int, - num_head_channels: Optional[int] = None, - num_groups: int = 32, - rescale_output_factor: float = 1.0, - eps: float = 1e-5): - super().__init__() - self.channels = channels - - self.num_heads = channels // num_head_channels if num_head_channels is not None else 1 - self.num_head_size = num_head_channels - self.group_norm = GroupNorm(num_channels=channels, - num_groups=num_groups, - eps=eps, - affine=True) - - self.qkv = Linear(channels, channels * 3) - self.rescale_output_factor = rescale_output_factor - self.proj_attn = Linear(channels, channels, 1) - - def transpose_for_scores(self, projection): - new_projection_shape = projection.size()[:-1] + (self.num_heads, - self.num_head_size) - # move heads to 2nd position (B, T, H * D) -> (B, T, H, D) -> (B, H, T, D) - new_projection = projection.view(new_projection_shape).permute( - [0, 2, 1, 3]) - return new_projection - - def forward(self, hidden_states): - assert not hidden_states.is_dynamic() - - residual = hidden_states - batch, channel, height, width = hidden_states.size() - - # norm - hidden_states = self.group_norm(hidden_states) - hidden_states = hidden_states.view([batch, channel, - height * width]).transpose(1, 2) - - # proj to q, k, v - qkv_proj = self.qkv(hidden_states) - - query_proj, key_proj, value_proj = split(qkv_proj, channel, dim=2) - - # transpose - query_states = self.transpose_for_scores(query_proj) - key_states = self.transpose_for_scores(key_proj) - value_states = self.transpose_for_scores(value_proj) - - # get scores - with precision('float32'): - attention_scores = matmul(query_states, - (key_states).transpose(-1, -2)) - attention_scores = attention_scores / math.sqrt( - self.channels / self.num_heads) - attention_probs = softmax(attention_scores, dim=-1) - - # compute attention output - hidden_states = matmul(attention_probs, value_states) - hidden_states = hidden_states.permute([0, 2, 1, 3]) - - new_hidden_states_shape = hidden_states.size()[:-2] + (self.channels, ) - hidden_states = hidden_states.view(new_hidden_states_shape) - - # compute next hidden_states - hidden_states = self.proj_attn(hidden_states) - hidden_states = hidden_states.transpose(-1, -2).view( - [batch, channel, height, width]) - - # res connect and rescale - hidden_states = (hidden_states + residual) / self.rescale_output_factor - return hidden_states - - -def _transpose_for_scores(tensor, heads): - batch_size, seq_len, dim = tensor.size() - tensor = tensor.view([batch_size, seq_len, heads, dim // heads]) - tensor = tensor.permute([0, 2, 1, 3]) - return tensor - - -def _attention(query, key, value, scale): - # Multiply scale first to avoid overflow - # Do not use use_fp32_acc or it will be very slow - attention_scores = matmul(query * math.sqrt(scale), - key.transpose(-1, -2) * math.sqrt(scale), - use_fp32_acc=False) - attention_probs = softmax(attention_scores, dim=-1) - hidden_states = matmul(attention_probs, value, use_fp32_acc=False) - hidden_states = hidden_states.permute([0, 2, 1, 3]) - return hidden_states - - -class SelfAttention(Module): - - def __init__(self, - query_dim: int, - heads: int = 8, - dim_head: int = 64, - dtype=None): - super().__init__() - self.inner_dim = dim_head * heads - self.scale = dim_head**-0.5 - self.heads = heads - self._slice_size = None - - self.to_qkv = Linear(query_dim, - 3 * self.inner_dim, - bias=False, - dtype=dtype) - self.to_out = Linear(self.inner_dim, query_dim, dtype=dtype) - - def forward(self, hidden_states, mask=None): - assert not hidden_states.is_dynamic() - - qkv = self.to_qkv(hidden_states) - - query, key, value = split(qkv, self.inner_dim, dim=2) - query = _transpose_for_scores(query, self.heads) - key = _transpose_for_scores(key, self.heads) - value = _transpose_for_scores(value, self.heads) - hidden_states = _attention(query, key, value, self.scale) - - batch_size, seq_len, head_size, head_dim = hidden_states.size() - hidden_states = hidden_states.view( - [batch_size, seq_len, head_size * head_dim]) - return self.to_out(hidden_states) - - -class CrossAttention(Module): - - def __init__(self, - query_dim: int, - context_dim: Optional[int] = None, - heads: int = 8, - dim_head: int = 64, - dtype=None): - super().__init__() - self.inner_dim = dim_head * heads - context_dim = context_dim if context_dim is not None else query_dim - self.scale = dim_head**-0.5 - self.heads = heads - self._slice_size = None - - self.to_q = Linear(query_dim, self.inner_dim, bias=False, dtype=dtype) - self.to_kv = Linear(context_dim, - 2 * self.inner_dim, - bias=False, - dtype=dtype) - self.to_out = Linear(self.inner_dim, query_dim, dtype=dtype) - - def forward(self, hidden_states, context=None, mask=None): - assert not hidden_states.is_dynamic() - query = self.to_q(hidden_states) - is_cross_attn = context is not None - context = context if is_cross_attn else hidden_states - assert not context.is_dynamic() - kv = self.to_kv(context) - - query = _transpose_for_scores(query, self.heads) - key, value = split(kv, self.inner_dim, dim=2) - key = _transpose_for_scores(key, self.heads) - value = _transpose_for_scores(value, self.heads) - hidden_states = _attention(query, key, value, self.scale) - - batch_size, seq_len, head_size, head_dim = hidden_states.size() - hidden_states = hidden_states.view( - [batch_size, seq_len, head_size * head_dim]) - return self.to_out(hidden_states) - - -class FeedForward(Module): - - def __init__(self, - dim: int, - dim_out: Optional[int] = None, - mult: int = 4, - dtype=None): - super().__init__() - inner_dim = int(dim * mult) - dim_out = dim_out if dim_out is not None else dim - self.proj_in = Linear(dim, inner_dim * 2, dtype=dtype) - self.proj_out = Linear(inner_dim, dim_out, dtype=dtype) - - def forward(self, hidden_states): - x = self.proj_in(hidden_states) - x = geglu(x) - return self.proj_out(x) - - -class BasicTransformerBlock(Module): - - def __init__( - self, - dim: int, - n_heads: int, - d_head: int, - context_dim: Optional[int] = None, - dtype=None, - ): - super().__init__() - self.attn1 = SelfAttention(query_dim=dim, - heads=n_heads, - dim_head=d_head, - dtype=dtype) # is a self-attention - self.ff = FeedForward(dim, dtype=dtype) - self.attn2 = CrossAttention( - query_dim=dim, - context_dim=context_dim, - heads=n_heads, - dim_head=d_head, - dtype=dtype) # is self-attn if context is none - self.norm1 = LayerNorm(dim, dtype=dtype) - self.norm2 = LayerNorm(dim, dtype=dtype) - self.norm3 = LayerNorm(dim, dtype=dtype) - - def forward(self, hidden_states, context=None): - hidden_states = self.attn1(self.norm1(hidden_states)) + hidden_states - hidden_states = self.attn2(self.norm2(hidden_states), - context=context) + hidden_states - hidden_states = self.ff(self.norm3(hidden_states)) + hidden_states - return hidden_states - - -class Transformer2DModel(Module): - - def __init__( - self, - num_attention_heads: int = 16, - attention_head_dim: int = 88, - in_channels: Optional[int] = None, - num_layers: int = 1, - norm_num_groups: int = 32, - cross_attention_dim: Optional[int] = None, - use_linear_projection: bool = False, - dtype=None, - ): - super().__init__() - self.use_linear_projection = use_linear_projection - self.num_attention_heads = num_attention_heads - self.attention_head_dim = attention_head_dim - inner_dim = num_attention_heads * attention_head_dim - - self.norm = GroupNorm(num_groups=norm_num_groups, - num_channels=in_channels, - eps=1e-6, - affine=True, - dtype=dtype) - - if use_linear_projection: - self.proj_in = Linear(in_channels, inner_dim, dtype=dtype) - else: - self.proj_in = Conv2d(in_channels, - inner_dim, - kernel_size=(1, 1), - stride=(1, 1), - padding=(0, 0), - dtype=dtype) - - self.transformer_blocks = ModuleList([ - BasicTransformerBlock(inner_dim, - num_attention_heads, - attention_head_dim, - context_dim=cross_attention_dim, - dtype=dtype) for d in range(num_layers) - ]) - - if use_linear_projection: - self.proj_out = Linear(inner_dim, in_channels, dtype=dtype) - else: - self.proj_out = Conv2d(inner_dim, - in_channels, - kernel_size=(1, 1), - stride=(1, 1), - padding=(0, 0), - dtype=dtype) - - def forward(self, hidden_states, context=None): - assert not hidden_states.is_dynamic() - batch, _, height, weight = hidden_states.size() - residual = hidden_states - hidden_states = self.norm(hidden_states) - - if not self.use_linear_projection: - hidden_states = self.proj_in(hidden_states) - inner_dim = hidden_states.size()[1] - hidden_states = hidden_states.permute([0, 2, 3, 1]).view( - [batch, height * weight, inner_dim]) - else: - inner_dim = hidden_states.size()[1] - hidden_states = hidden_states.permute([0, 2, 3, 1]).view( - [batch, height * weight, inner_dim]) - hidden_states = self.proj_in(hidden_states) - - for block in self.transformer_blocks: - hidden_states = block(hidden_states, context=context) - - if not self.use_linear_projection: - hidden_states = hidden_states.view( - [batch, height, weight, inner_dim]).permute([0, 3, 1, 2]) - hidden_states = self.proj_out(hidden_states) - else: - hidden_states = self.proj_out(hidden_states) - hidden_states = hidden_states.view( - [batch, height, weight, inner_dim]).permute([0, 3, 1, 2]) - - return hidden_states + residual diff --git a/tensorrt_llm/models/unet/embeddings.py b/tensorrt_llm/models/unet/embeddings.py deleted file mode 100644 index 8bfe16a408d6..000000000000 --- a/tensorrt_llm/models/unet/embeddings.py +++ /dev/null @@ -1,117 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -import math - -import tensorrt as trt - -from ..._utils import fp16_array, fp32_array -from ...functional import concat, constant, cos, exp, silu, sin -from ...layers import Linear -from ...module import Module - - -def get_timestep_embedding(timesteps, - embedding_dim, - flip_sin_to_cos=False, - downscale_freq_shift=1.0, - scale=1.0, - max_period=10000, - dtype=None): - """ - This matches the implementation in Denoising Diffusion Probabilistic Models: Create sinusoidal timestep embeddings. - :param timesteps: a 1-D Tensor of N indices, one per batch element. - These may be fractional. - :param embedding_dim: the dimension of the output. :param max_period: controls the minimum frequency of the - embeddings. :return: an [N x dim] Tensor of positional embeddings. - """ - assert timesteps.rank() == 1, "Timesteps should be a 1d-array" - - half_dim = embedding_dim // 2 - - exponent = [ - i * -math.log(max_period) / (half_dim - downscale_freq_shift) - for i in range(half_dim) - ] - - if dtype == trt.DataType.HALF: - emb = exp(constant(fp16_array(exponent))) - else: - emb = exp(constant(fp32_array(exponent))) - - ts_shape = list(timesteps.size()) - ts_shape.append(1) - emb_shape = list(emb.size()) - emb_shape.insert(0, 1) - - emb = timesteps.view(ts_shape) * emb.view(emb_shape) - - emb = scale * emb - # concat sine and cosine embeddings - - # flip sine and cosine embeddings - if flip_sin_to_cos: - emb = concat([cos(emb), sin(emb)], dim=1) - else: - emb = concat([sin(emb), cos(emb)], dim=1) - - #TODO Enable below logic when TensorRT LLM supports pad feature. - # zero pad - # if embedding_dim % 2 == 1: - # emb = torch.nn.functional.pad(emb, (0, 1, 0, 0)) - return emb - - -class TimestepEmbedding(Module): - - def __init__(self, channel, time_embed_dim, act_fn="silu", dtype=None): - super().__init__() - - self.linear_1 = Linear(channel, time_embed_dim, dtype=dtype) - self.act = None - if act_fn == "silu": - self.act = silu - self.linear_2 = Linear(time_embed_dim, time_embed_dim, dtype=dtype) - - def forward(self, sample): - sample = self.linear_1(sample) - - if self.act is not None: - sample = self.act(sample) - - sample = self.linear_2(sample) - return sample - - -class Timesteps(Module): - - def __init__(self, - num_channels, - flip_sin_to_cos, - downscale_freq_shift, - dtype=None): - super().__init__() - self.num_channels = num_channels - self.flip_sin_to_cos = flip_sin_to_cos - self.downscale_freq_shift = downscale_freq_shift - self.dtype = dtype - - def forward(self, timesteps): - t_emb = get_timestep_embedding( - timesteps, - self.num_channels, - flip_sin_to_cos=self.flip_sin_to_cos, - downscale_freq_shift=self.downscale_freq_shift, - dtype=self.dtype) - return t_emb diff --git a/tensorrt_llm/models/unet/pp/attention.py b/tensorrt_llm/models/unet/pp/attention.py deleted file mode 100755 index f64d666b0891..000000000000 --- a/tensorrt_llm/models/unet/pp/attention.py +++ /dev/null @@ -1,107 +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. -from ....functional import allgather, split -from ....mapping import Mapping -from ....module import Module -from ..attention import CrossAttention, SelfAttention, _attention - - -class DistriSelfAttentionPP(Module): - - def __init__(self, module: SelfAttention, mapping: Mapping = Mapping()): - super().__init__() - self.mapping = mapping - self.module = module - - def forward(self, hidden_states): - mapping = self.mapping - attn = self.module - - batch_size, sequence_length, _ = hidden_states.shape - - qkv = attn.to_qkv(hidden_states) - - query, kv = split(qkv, [attn.inner_dim, attn.inner_dim * 2], dim=2) - - if mapping.tp_size == 1: - full_kv = kv - else: - full_kv = allgather(kv, group=mapping.tp_group, gather_dim=1) - - key, value = split(full_kv, full_kv.shape[-1] // 2, dim=-1) - - inner_dim = key.shape[-1] - head_dim = inner_dim // attn.heads - - query = query.view([batch_size, -1, attn.heads, - head_dim]).transpose(1, 2) - key = key.view([batch_size, -1, attn.heads, head_dim]).transpose(1, 2) - value = value.view([batch_size, -1, attn.heads, - head_dim]).transpose(1, 2) - - hidden_states = _attention(query, key, value, attn.scale) - - hidden_states = hidden_states.view( - [batch_size, -1, attn.heads * head_dim]) - - # linear proj - hidden_states = attn.to_out(hidden_states) - - return hidden_states - - -class DistriCrossAttentionPP(Module): - - def __init__(self, module: CrossAttention, mapping: Mapping = Mapping()): - super().__init__() - self.mapping = mapping - self.module = module - self.kv_cache = None - - def forward(self, hidden_states, context): - attn = self.module - recompute_kv = self.kv_cache is None - - if context is None: - context = hidden_states - - batch_size, sequence_length, _ = context.shape - - query = attn.to_q(hidden_states) - - if recompute_kv or self.kv_cache is None: - kv = attn.to_kv(context) - self.kv_cache = kv - else: - kv = self.kv_cache - key, value = split(kv, kv.shape[-1] // 2, dim=-1) - - inner_dim = key.shape[-1] - head_dim = inner_dim // attn.heads - - query = query.view([batch_size, -1, attn.heads, - head_dim]).transpose(1, 2) - key = key.view([batch_size, -1, attn.heads, head_dim]).transpose(1, 2) - value = value.view([batch_size, -1, attn.heads, - head_dim]).transpose(1, 2) - - hidden_states = _attention(query, key, value, scale=attn.scale) - - hidden_states = hidden_states.view( - [batch_size, -1, attn.heads * head_dim]) - - # linear proj - hidden_states = attn.to_out(hidden_states) - return hidden_states diff --git a/tensorrt_llm/models/unet/pp/conv2d.py b/tensorrt_llm/models/unet/pp/conv2d.py deleted file mode 100755 index a1aa8f7b72c5..000000000000 --- a/tensorrt_llm/models/unet/pp/conv2d.py +++ /dev/null @@ -1,109 +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. - -from ....functional import allgather, concat, conv2d, slice, stack, unsqueeze -from ....layers import Conv2d -from ....mapping import Mapping -from ....module import Module - - -class DistriConv2dPP(Module): - - def __init__(self, - conv: Conv2d, - mapping: Mapping = Mapping(), - is_first_layer: bool = False): - super().__init__() - self.mapping = mapping - self.conv = conv - self.is_first_layer = is_first_layer - - def sliced_forward(self, x): - mapping = self.mapping - b, c, h, w = x.shape - assert h % mapping.tp_size == 0 - - stride = self.conv.stride[0] - padding = self.conv.padding[0] - - output_h = x.shape[2] // stride // mapping.tp_size - idx = mapping.tp_rank - h_begin = output_h * idx * stride - padding - h_end = output_h * (idx + 1) * stride + padding - pre_padding = [0, padding] - post_padding = [0, padding] - if h_begin < 0: - h_begin = 0 - pre_padding[0] = padding - if h_end > h: - h_end = h - post_padding[0] = padding - sliced_input = slice(x, [0, 0, h_begin, 0], [b, c, h_end - h_begin, w]) - return conv2d(sliced_input, - self.conv.weight.value, - None if self.conv.bias is None else self.conv.bias.value, - stride=self.conv.stride, - padding=(0, 0), - pre_padding=tuple(pre_padding), - post_padding=tuple(post_padding)) - - def forward(self, x, *args, **kwargs): - mapping = self.mapping - if self.is_first_layer: - full_x = x - output = self.sliced_forward(full_x) - else: - boundary_size = self.conv.padding[0] - - def create_padded_x(x, boundaries): - preH = 0 - postH = 0 - if mapping.tp_rank == 0: - b = boundaries.select(0, mapping.tp_rank + 1).select(0, 0) - padded_x = concat([x, b], dim=2) - preH = boundary_size - elif mapping.tp_rank == mapping.tp_size - 1: - b = boundaries.select(0, mapping.tp_rank - 1).select(0, 1) - padded_x = concat([b, x], dim=2) - postH = boundary_size - else: - b0 = boundaries.select(0, mapping.tp_rank - 1).select(0, 1) - b1 = boundaries.select(0, mapping.tp_rank + 1).select(0, 0) - padded_x = concat( - [ - b0, - x, - b1, - ], - dim=2, - ) - return padded_x, preH, postH - - n, c, h, w = x.shape - b0 = slice(x, [0, 0, 0, 0], [n, c, boundary_size, w]) - b1 = slice(x, [0, 0, h - boundary_size, 0], - [n, c, boundary_size, w]) - boundary = stack([b0, b1], dim=0) - - boundaries = allgather(unsqueeze(boundary, 0), - group=mapping.tp_group) - padded_x, preH, postH = create_padded_x(x, boundaries) - output = conv2d(padded_x, - self.conv.weight.value, - self.conv.bias.value, - stride=self.conv.stride, - pre_padding=(preH, self.conv.padding[1]), - post_padding=(postH, self.conv.padding[1])) - return output diff --git a/tensorrt_llm/models/unet/pp/groupnorm.py b/tensorrt_llm/models/unet/pp/groupnorm.py deleted file mode 100755 index 6cbf4d2d0228..000000000000 --- a/tensorrt_llm/models/unet/pp/groupnorm.py +++ /dev/null @@ -1,57 +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. -from ....functional import allreduce, pow, select, stack -from ....layers import GroupNorm -from ....mapping import Mapping -from ....module import Module - - -class DistriGroupNorm(Module): - - def __init__(self, - module: GroupNorm, - mapping: Mapping = Mapping(), - is_first_layer: bool = False): - super().__init__() - self.mapping = mapping - self.module = module - - def forward(self, x, *args, **kwargs): - mapping = self.mapping - module = self.module - n, c, h, w = x.shape - num_groups = module.num_groups - group_size = c // num_groups - - x = x.view([n, num_groups, group_size, h, w]) - x_mean = x.mean(dim=4, keepdim=True).mean(dim=(3, 2), keepdim=True) - x2_mean = pow(x, 2.0).mean(dim=4, keepdim=True).mean(dim=(3, 2), - keepdim=True) - mean = stack([x_mean, x2_mean], dim=0) - mean = allreduce(mean, mapping.tp_group) - mean = mean / (mapping.tp_size * 1.0) - x_mean = select(mean, 0, 0) - x2_mean = select(mean, 0, 1) - var = x2_mean - pow(x_mean, 2.0) - num_elements = group_size * h * w - var = var * (num_elements / (num_elements - 1)) - std = (var + module.eps).sqrt() - output = (x - x_mean) / std - output = output.view([n, c, h, w]) - if module.affine: - output = output * module.weight.value.view([1, -1, 1, 1]) - output = output + module.bias.value.view([1, -1, 1, 1]) - - return output diff --git a/tensorrt_llm/models/unet/pp/unet_pp.py b/tensorrt_llm/models/unet/pp/unet_pp.py deleted file mode 100755 index 9b1a7d64a60f..000000000000 --- a/tensorrt_llm/models/unet/pp/unet_pp.py +++ /dev/null @@ -1,131 +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. -from ....functional import allgather, concat, slice, stack -from ....layers import Conv2d, GroupNorm -from ....mapping import Mapping -from ....module import Module -from ..attention import CrossAttention, SelfAttention -from ..unet_2d_condition import UNet2DConditionModel -from .attention import DistriCrossAttentionPP, DistriSelfAttentionPP -from .conv2d import DistriConv2dPP -from .groupnorm import DistriGroupNorm - - -class DistriUNetPP(Module): - - def __init__(self, - model: UNet2DConditionModel, - mapping: Mapping = Mapping()): - super().__init__() - self.mapping = mapping - self.model = model - if mapping.tp_size > 1: - for name, module in model.named_modules(): - if isinstance(module, DistriConv2dPP) or isinstance(module, DistriSelfAttentionPP) \ - or isinstance(module, DistriCrossAttentionPP) or isinstance(module, DistriGroupNorm): - continue - for subname, submodule in module.named_children(): - if isinstance(submodule, Conv2d): - kernel_size = submodule.kernel_size - if kernel_size == (1, 1) or kernel_size == 1: - continue - wrapped_submodule = DistriConv2dPP( - submodule, - mapping, - is_first_layer=subname == "conv_in") - setattr(module, subname, wrapped_submodule) - elif isinstance(submodule, SelfAttention): - wrapped_submodule = DistriSelfAttentionPP( - submodule, mapping) - setattr(module, subname, wrapped_submodule) - elif isinstance(submodule, CrossAttention): - wrapped_submodule = DistriCrossAttentionPP( - submodule, mapping) - setattr(module, subname, wrapped_submodule) - elif isinstance(submodule, GroupNorm): - wrapped_submodule = DistriGroupNorm(submodule, mapping) - setattr(module, subname, wrapped_submodule) - - def forward(self, - sample, - timesteps, - encoder_hidden_states, - text_embeds=None, - time_ids=None): - mapping = self.mapping - b, c, h, w = sample.shape - - if mapping.world_size == 1: - output = self.model( - sample, - timesteps, - encoder_hidden_states, - text_embeds=text_embeds, - time_ids=time_ids, - ) - elif mapping.pp_size > 1: - assert b == 2 and mapping.pp_size == 2 - batch_idx = mapping.pp_rank - # sample[batch_idx : batch_idx + 1] - sample = slice(sample, [batch_idx, 0, 0, 0], [1, c, h, w]) - e_shape = encoder_hidden_states.shape - encoder_hidden_states = slice( - encoder_hidden_states, [batch_idx, 0, 0], - [1, e_shape[1], e_shape[2] - ]) # encoder_hidden_states[batch_idx : batch_idx + 1] - if text_embeds: - t_shape = text_embeds.shape - # text_embeds[batch_idx : batch_idx + 1] - text_embeds = slice(text_embeds, [batch_idx, 0], - [1, t_shape[1]]) - if time_ids: - t_shape = time_ids.shape - # time_ids[batch_idx : batch_idx + 1] - time_ids = slice(time_ids, [batch_idx, 0], [1, t_shape[1]]) - output = self.model( - sample, - timesteps, - encoder_hidden_states, - text_embeds=text_embeds, - time_ids=time_ids, - ) - output = allgather( - output, - [i for i in range(mapping.world_size)], - ) - patch_list = [] - for i in range(mapping.tp_size): - patch_list.append(output.select(dim=0, index=i)) - b1 = concat(patch_list, dim=1) - patch_list = [] - for i in range(mapping.tp_size, mapping.world_size): - patch_list.append(output.select(dim=0, index=i)) - b2 = concat(patch_list, dim=1) - output = stack([b1, b2], dim=0) - else: - output = self.model( - sample, - timesteps, - encoder_hidden_states, - text_embeds=text_embeds, - time_ids=time_ids, - ) - output = allgather(output, mapping.tp_group, 2) - - return output - - @property - def add_embedding(self): - return self.model.add_embedding diff --git a/tensorrt_llm/models/unet/resnet.py b/tensorrt_llm/models/unet/resnet.py deleted file mode 100644 index 453e55324f04..000000000000 --- a/tensorrt_llm/models/unet/resnet.py +++ /dev/null @@ -1,254 +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. -from functools import partial - -from ...functional import avg_pool2d, interpolate, silu -from ...layers import (AvgPool2d, Conv2d, ConvTranspose2d, GroupNorm, Linear, - Mish) -from ...module import Module - - -class Upsample2D(Module): - - def __init__(self, - channels: int, - use_conv=False, - use_conv_transpose=False, - out_channels=None, - dtype=None) -> None: - super().__init__() - - self.channels = channels - self.out_channels = out_channels - - self.use_conv_transpose = use_conv_transpose - self.use_conv = use_conv - if self.use_conv_transpose: - self.conv = ConvTranspose2d(channels, - self.out_channels, - 4, - 2, - 1, - dtype=dtype) - elif use_conv: - self.conv = Conv2d(self.channels, - self.out_channels, (3, 3), - padding=(1, 1), - dtype=dtype) - else: - self.conv = None - - def forward(self, hidden_states, output_size=None): - assert not hidden_states.is_dynamic() - batch, channels, _, _ = hidden_states.size() - assert channels == self.channels - - if self.use_conv_transpose: - return self.conv(hidden_states) - - if output_size is None: - hidden_states = interpolate(hidden_states, - scale_factor=2.0, - mode="nearest") - else: - hidden_states = interpolate(hidden_states, - size=output_size, - mode="nearest") - - if self.use_conv: - hidden_states = self.conv(hidden_states) - - return hidden_states - - -class Downsample2D(Module): - - def __init__(self, - channels, - use_conv=False, - out_channels=None, - padding=1, - dtype=None) -> None: - super().__init__() - self.channels = channels - self.out_channels = out_channels or channels - self.use_conv = use_conv - self.padding = padding - stride = (2, 2) - - if use_conv: - self.conv = Conv2d(self.channels, - self.out_channels, (3, 3), - stride=stride, - padding=(padding, padding), - dtype=dtype) - else: - assert self.channels == self.out_channels - self.conv = AvgPool2d(kernel_size=stride, stride=stride) - - def forward(self, hidden_states): - assert not hidden_states.is_dynamic() - batch, channels, _, _ = hidden_states.size() - assert channels == self.channels - - #TODO add the missing pad function - hidden_states = self.conv(hidden_states) - - return hidden_states - - -class ResnetBlock2D(Module): - - def __init__( - self, - *, - in_channels, - out_channels=None, - conv_shortcut=False, - dropout=0.0, - temb_channels=512, - groups=32, - groups_out=None, - pre_norm=True, - eps=1e-6, - non_linearity="swish", - time_embedding_norm="default", - kernel=None, - output_scale_factor=1.0, - use_in_shortcut=None, - up=False, - down=False, - dtype=None, - ): - super().__init__() - self.pre_norm = pre_norm - self.pre_norm = True - self.in_channels = in_channels - out_channels = in_channels if out_channels is None else out_channels - self.out_channels = out_channels - self.use_conv_shortcut = conv_shortcut - self.time_embedding_norm = time_embedding_norm - self.up = up - self.down = down - self.output_scale_factor = output_scale_factor - - if groups_out is None: - groups_out = groups - - self.norm1 = GroupNorm(num_groups=groups, - num_channels=in_channels, - eps=eps, - affine=True, - dtype=dtype) - self.conv1 = Conv2d(in_channels, - out_channels, - kernel_size=(3, 3), - stride=(1, 1), - padding=(1, 1), - dtype=dtype) - - if temb_channels is not None: - self.time_emb_proj = Linear(temb_channels, - out_channels, - dtype=dtype) - else: - self.time_emb_proj = None - - self.norm2 = GroupNorm(num_groups=groups_out, - num_channels=out_channels, - eps=eps, - affine=True, - dtype=dtype) - self.conv2 = Conv2d(out_channels, - out_channels, - kernel_size=(3, 3), - stride=(1, 1), - padding=(1, 1), - dtype=dtype) - - if non_linearity == "swish": - self.nonlinearity = lambda x: silu(x) - elif non_linearity == "mish": - self.nonlinearity = Mish() - elif non_linearity == "silu": - self.nonlinearity = silu - - self.upsample = self.downsample = None - #TODO add the fir kernel supporting. - if self.up: - if kernel == "sde_vp": - self.upsample = partial(interpolate, - scale_factor=2.0, - mode="nearest") - else: - self.upsample = Upsample2D(in_channels, - use_conv=False, - dtype=dtype) - elif self.down: - - if kernel == "sde_vp": - self.downsample = partial(avg_pool2d, kernel_size=2, stride=2) - else: - self.downsample = Downsample2D(in_channels, - use_conv=False, - padding=1, - name="op", - dtype=dtype) - - self.use_in_shortcut = self.in_channels != self.out_channels if use_in_shortcut is None else use_in_shortcut - - if self.use_in_shortcut: - self.conv_shortcut = Conv2d(in_channels, - out_channels, - kernel_size=(1, 1), - stride=(1, 1), - padding=(0, 0), - dtype=dtype) - else: - self.conv_shortcut = None - - def forward(self, input_tensor, temb): - hidden_states = input_tensor - hidden_states = self.norm1(hidden_states) - hidden_states = self.nonlinearity(hidden_states) - - if self.upsample is not None: - input_tensor = self.upsample(input_tensor) - hidden_states = self.upsample(hidden_states) - elif self.downsample is not None: - input_tensor = self.downsample(input_tensor) - hidden_states = self.downsample(hidden_states) - - hidden_states = self.conv1(hidden_states) - if self.time_emb_proj is not None: - temb = self.time_emb_proj(self.nonlinearity(temb)) - new_shape = list(temb.size()) - new_shape.extend([1, 1]) #[:,:,None,None] ->view - temb = temb.view(new_shape) - - assert self.time_embedding_norm == "default" - if temb is not None: - hidden_states = hidden_states + temb - - hidden_states = self.norm2(hidden_states) - hidden_states = self.nonlinearity(hidden_states) - hidden_states = self.conv2(hidden_states) - - if self.conv_shortcut is not None: - input_tensor = self.conv_shortcut(input_tensor) - - output_tensor = (input_tensor + - hidden_states) / self.output_scale_factor - return output_tensor diff --git a/tensorrt_llm/models/unet/unet_2d_blocks.py b/tensorrt_llm/models/unet/unet_2d_blocks.py deleted file mode 100644 index 02479221292e..000000000000 --- a/tensorrt_llm/models/unet/unet_2d_blocks.py +++ /dev/null @@ -1,636 +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. -from typing import Tuple, Union - -from ...functional import concat -from ...module import Module, ModuleList -from .attention import AttentionBlock, Transformer2DModel -from .resnet import Downsample2D, ResnetBlock2D, Upsample2D - - -def get_down_block( - down_block_type, - num_layers, - in_channels, - out_channels, - temb_channels, - add_downsample, - resnet_eps, - resnet_act_fn, - attn_num_head_channels, - transformer_layers_per_block=1, - cross_attention_dim=None, - downsample_padding=None, - use_linear_projection=False, - dtype=None, -): - down_block_type = down_block_type[7:] if down_block_type.startswith( - "UNetRes") else down_block_type - if down_block_type == "DownBlock2D": - return DownBlock2D( - num_layers=num_layers, - in_channels=in_channels, - out_channels=out_channels, - temb_channels=temb_channels, - add_downsample=add_downsample, - resnet_eps=resnet_eps, - resnet_act_fn=resnet_act_fn, - downsample_padding=downsample_padding, - dtype=dtype, - ) - elif down_block_type == "CrossAttnDownBlock2D": - if cross_attention_dim is None: - raise ValueError( - "cross_attention_dim must be specified for CrossAttnUpBlock2D") - return CrossAttnDownBlock2D( - num_layers=num_layers, - transformer_layers_per_block=transformer_layers_per_block, - in_channels=in_channels, - out_channels=out_channels, - temb_channels=temb_channels, - add_downsample=add_downsample, - resnet_eps=resnet_eps, - resnet_act_fn=resnet_act_fn, - downsample_padding=downsample_padding, - cross_attention_dim=cross_attention_dim, - attn_num_head_channels=attn_num_head_channels, - use_linear_projection=use_linear_projection, - dtype=dtype, - ) - - raise ValueError(f"{down_block_type} does not exist.") - - -def get_up_block( - up_block_type, - num_layers, - in_channels, - out_channels, - prev_output_channel, - temb_channels, - add_upsample, - resnet_eps, - resnet_act_fn, - attn_num_head_channels, - transformer_layers_per_block=1, - cross_attention_dim=None, - use_linear_projection=False, - dtype=None, -): - up_block_type = up_block_type[7:] if up_block_type.startswith( - "UNetRes") else up_block_type - if up_block_type == "UpBlock2D": - return UpBlock2D( - num_layers=num_layers, - in_channels=in_channels, - out_channels=out_channels, - prev_output_channel=prev_output_channel, - temb_channels=temb_channels, - add_upsample=add_upsample, - resnet_eps=resnet_eps, - resnet_act_fn=resnet_act_fn, - dtype=dtype, - ) - elif up_block_type == "CrossAttnUpBlock2D": - if cross_attention_dim is None: - raise ValueError( - "cross_attention_dim must be specified for CrossAttnUpBlock2D") - return CrossAttnUpBlock2D( - num_layers=num_layers, - transformer_layers_per_block=transformer_layers_per_block, - in_channels=in_channels, - out_channels=out_channels, - prev_output_channel=prev_output_channel, - temb_channels=temb_channels, - add_upsample=add_upsample, - resnet_eps=resnet_eps, - resnet_act_fn=resnet_act_fn, - cross_attention_dim=cross_attention_dim, - attn_num_head_channels=attn_num_head_channels, - use_linear_projection=use_linear_projection, - dtype=dtype, - ) - - raise ValueError(f"{up_block_type} does not exist.") - - -class UpBlock2D(Module): - - def __init__( - self, - in_channels: int, - prev_output_channel: int, - out_channels: int, - temb_channels: int, - dropout: float = 0.0, - num_layers: int = 1, - resnet_eps: float = 1e-6, - resnet_time_scale_shift: str = "default", - resnet_act_fn: str = "swish", - resnet_groups: int = 32, - resnet_pre_norm: bool = True, - output_scale_factor=1.0, - add_upsample=True, - dtype=None, - ): - super().__init__() - resnets = [] - - for i in range(num_layers): - res_skip_channels = in_channels if (i == num_layers - - 1) else out_channels - resnet_in_channels = prev_output_channel if i == 0 else out_channels - - resnets.append( - ResnetBlock2D( - in_channels=resnet_in_channels + res_skip_channels, - out_channels=out_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - dtype=dtype, - )) - - self.resnets = ModuleList(resnets) - - if add_upsample: - self.upsamplers = ModuleList([ - Upsample2D(out_channels, - use_conv=True, - out_channels=out_channels, - dtype=dtype) - ]) - else: - self.upsamplers = None - - def forward(self, - hidden_states, - res_hidden_states_tuple, - temb=None, - upsample_size=None): - for resnet in self.resnets: - # pop res hidden states - res_hidden_states = res_hidden_states_tuple[-1] - res_hidden_states_tuple = res_hidden_states_tuple[:-1] - hidden_states = concat([hidden_states, res_hidden_states], dim=1) - - hidden_states = resnet(hidden_states, temb) - - if self.upsamplers is not None: - for upsampler in self.upsamplers: - hidden_states = upsampler(hidden_states, upsample_size) - - return hidden_states - - -class DownBlock2D(Module): - - def __init__( - self, - in_channels: int, - out_channels: int, - temb_channels: int, - dropout: float = 0.0, - num_layers: int = 1, - resnet_eps: float = 1e-6, - resnet_time_scale_shift: str = "default", - resnet_act_fn: str = "swish", - resnet_groups: int = 32, - resnet_pre_norm: bool = True, - output_scale_factor=1.0, - add_downsample=True, - downsample_padding=1, - dtype=None, - ): - super().__init__() - resnets = [] - - for i in range(num_layers): - in_channels = in_channels if i == 0 else out_channels - resnets.append( - ResnetBlock2D( - in_channels=in_channels, - out_channels=out_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - dtype=dtype, - )) - - self.resnets = ModuleList(resnets) - - if add_downsample: - self.downsamplers = ModuleList([ - Downsample2D(out_channels, - use_conv=True, - out_channels=out_channels, - padding=downsample_padding, - dtype=dtype) - ]) - else: - self.downsamplers = None - - def forward(self, hidden_states, temb=None): - output_states = () - for resnet in self.resnets: - hidden_states = resnet(hidden_states, temb) - - output_states += (hidden_states, ) - if self.downsamplers is not None: - for downsampler in self.downsamplers: - hidden_states = downsampler(hidden_states) - - output_states += (hidden_states, ) - return hidden_states, output_states - - -class UNetMidBlock2D(Module): - - def __init__( - self, - in_channels: int, - temb_channels: int, - dropout: float = 0.0, - num_layers: int = 1, - resnet_eps: float = 1e-6, - resnet_time_scale_shift: str = "default", - resnet_act_fn: str = "swish", - resnet_groups: int = 32, - resnet_pre_norm: bool = True, - attn_num_head_channels=1, - attention_type="default", - output_scale_factor=1.0, - dtype=None, - **kwargs, - ): - super().__init__() - - self.attention_type = attention_type - resnet_groups = resnet_groups if resnet_groups is not None else min( - in_channels // 4, 32) - - # there is always at least one resnet - resnets = [ - ResnetBlock2D( - in_channels=in_channels, - out_channels=in_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - dtype=dtype, - ) - ] - attentions = [] - - for _ in range(num_layers): - attentions.append( - AttentionBlock( - in_channels, - num_head_channels=attn_num_head_channels, - rescale_output_factor=output_scale_factor, - eps=resnet_eps, - num_groups=resnet_groups, - dtype=dtype, - )) - resnets.append( - ResnetBlock2D( - in_channels=in_channels, - out_channels=in_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - dtype=dtype, - )) - - self.attentions = ModuleList(attentions) - self.resnets = ModuleList(resnets) - - def forward(self, hidden_states, temb=None, encoder_states=None): - hidden_states = self.resnets[0](hidden_states, temb) - for attn, resnet in zip(self.attentions, self.resnets[1:]): - if self.attention_type == "default": - hidden_states = attn(hidden_states) - else: - hidden_states = attn(hidden_states, encoder_states) - hidden_states = resnet(hidden_states, temb) - - return hidden_states - - -class CrossAttnUpBlock2D(Module): - - def __init__( - self, - in_channels: int, - out_channels: int, - prev_output_channel: int, - temb_channels: int, - dropout: float = 0.0, - num_layers: int = 1, - transformer_layers_per_block: Union[int, Tuple[int]] = 1, - resnet_eps: float = 1e-6, - resnet_time_scale_shift: str = "default", - resnet_act_fn: str = "swish", - resnet_groups: int = 32, - resnet_pre_norm: bool = True, - attn_num_head_channels=1, - cross_attention_dim=1280, - attention_type="default", - output_scale_factor=1.0, - add_upsample=True, - use_linear_projection: bool = False, - dtype=None, - ): - super().__init__() - resnets = [] - attentions = [] - - self.attention_type = attention_type - self.attn_num_head_channels = attn_num_head_channels - - # support for variable transformer layers per block - if isinstance(transformer_layers_per_block, int): - transformer_layers_per_block = [transformer_layers_per_block - ] * num_layers - - for i in range(num_layers): - res_skip_channels = in_channels if (i == num_layers - - 1) else out_channels - resnet_in_channels = prev_output_channel if i == 0 else out_channels - - resnets.append( - ResnetBlock2D( - in_channels=resnet_in_channels + res_skip_channels, - out_channels=out_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - dtype=dtype, - )) - - attentions.append( - Transformer2DModel(in_channels=out_channels, - num_layers=transformer_layers_per_block[i], - num_attention_heads=attn_num_head_channels, - attention_head_dim=out_channels // - attn_num_head_channels, - norm_num_groups=resnet_groups, - use_linear_projection=use_linear_projection, - cross_attention_dim=cross_attention_dim, - dtype=dtype)) - self.attentions = ModuleList(attentions) - self.resnets = ModuleList(resnets) - - if add_upsample: - self.upsamplers = ModuleList([ - Upsample2D(out_channels, - use_conv=True, - out_channels=out_channels, - dtype=dtype) - ]) - else: - self.upsamplers = None - - def forward( - self, - hidden_states, - res_hidden_states_tuple, - temb=None, - encoder_hidden_states=None, - upsample_size=None, - ): - - for resnet, attn in zip(self.resnets, self.attentions): - # pop res hidden states - res_hidden_states = res_hidden_states_tuple[-1] - res_hidden_states_tuple = res_hidden_states_tuple[:-1] - hidden_states = concat([hidden_states, res_hidden_states], dim=1) - - hidden_states = resnet(hidden_states, temb) - hidden_states = attn(hidden_states, context=encoder_hidden_states) - - if self.upsamplers is not None: - for upsampler in self.upsamplers: - hidden_states = upsampler(hidden_states, upsample_size) - - return hidden_states - - -class CrossAttnDownBlock2D(Module): - - def __init__( - self, - in_channels: int, - out_channels: int, - temb_channels: int, - dropout: float = 0.0, - num_layers: int = 1, - transformer_layers_per_block: Union[int, Tuple[int]] = 1, - resnet_eps: float = 1e-6, - resnet_time_scale_shift: str = "default", - resnet_act_fn: str = "swish", - resnet_groups: int = 32, - resnet_pre_norm: bool = True, - attn_num_head_channels=1, - cross_attention_dim=1280, - attention_type="default", - output_scale_factor=1.0, - downsample_padding=1, - add_downsample=True, - use_linear_projection: bool = False, - dtype=None, - ): - super().__init__() - resnets = [] - attentions = [] - - self.attention_type = attention_type - self.attn_num_head_channels = attn_num_head_channels - - # support for variable transformer layers per block - if isinstance(transformer_layers_per_block, int): - transformer_layers_per_block = [transformer_layers_per_block - ] * num_layers - - for i in range(num_layers): - in_channels = in_channels if i == 0 else out_channels - resnets.append( - ResnetBlock2D( - in_channels=in_channels, - out_channels=out_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - dtype=dtype, - )) - attentions.append( - Transformer2DModel(in_channels=out_channels, - num_layers=transformer_layers_per_block[i], - num_attention_heads=attn_num_head_channels, - attention_head_dim=out_channels // - attn_num_head_channels, - norm_num_groups=resnet_groups, - use_linear_projection=use_linear_projection, - cross_attention_dim=cross_attention_dim, - dtype=dtype)) - self.attentions = ModuleList(attentions) - self.resnets = ModuleList(resnets) - - if add_downsample: - self.downsamplers = ModuleList([ - Downsample2D(out_channels, - use_conv=True, - out_channels=out_channels, - padding=downsample_padding, - dtype=dtype) - ]) - else: - self.downsamplers = None - - def forward(self, hidden_states, temb=None, encoder_hidden_states=None): - output_states = () - - for resnet, attn in zip(self.resnets, self.attentions): - - hidden_states = resnet(hidden_states, temb) - - hidden_states = attn(hidden_states, context=encoder_hidden_states) - - output_states += (hidden_states, ) - - if self.downsamplers is not None: - for downsampler in self.downsamplers: - hidden_states = downsampler(hidden_states) - - output_states += (hidden_states, ) - - return hidden_states, output_states - - -class UNetMidBlock2DCrossAttn(Module): - - def __init__( - self, - in_channels: int, - temb_channels: int, - dropout: float = 0.0, - num_layers: int = 1, - transformer_layers_per_block: Union[int, Tuple[int]] = 1, - resnet_eps: float = 1e-6, - resnet_time_scale_shift: str = "default", - resnet_act_fn: str = "swish", - resnet_groups: int = 32, - resnet_pre_norm: bool = True, - attn_num_head_channels=1, - attention_type="default", - output_scale_factor=1.0, - cross_attention_dim=1280, - use_linear_projection: bool = False, - dtype=None, - **kwargs, - ): - super().__init__() - - self.attention_type = attention_type - self.attn_num_head_channels = attn_num_head_channels - resnet_groups = resnet_groups if resnet_groups is not None else min( - in_channels // 4, 32) - - # support for variable transformer layers per block - if isinstance(transformer_layers_per_block, int): - transformer_layers_per_block = [transformer_layers_per_block - ] * num_layers - - # there is always at least one resnet - resnets = [ - ResnetBlock2D(in_channels=in_channels, - out_channels=in_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - dtype=dtype) - ] - attentions = [] - - for i in range(num_layers): - attentions.append( - Transformer2DModel(in_channels=in_channels, - num_layers=transformer_layers_per_block[i], - num_attention_heads=attn_num_head_channels, - attention_head_dim=in_channels // - attn_num_head_channels, - norm_num_groups=resnet_groups, - use_linear_projection=use_linear_projection, - cross_attention_dim=cross_attention_dim, - dtype=dtype)) - resnets.append( - ResnetBlock2D(in_channels=in_channels, - out_channels=in_channels, - temb_channels=temb_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - time_embedding_norm=resnet_time_scale_shift, - non_linearity=resnet_act_fn, - output_scale_factor=output_scale_factor, - pre_norm=resnet_pre_norm, - dtype=dtype)) - - self.attentions = ModuleList(attentions) - self.resnets = ModuleList(resnets) - - def forward(self, hidden_states, temb=None, encoder_hidden_states=None): - hidden_states = self.resnets[0](hidden_states, temb) - - for attn, resnet in zip(self.attentions, self.resnets[1:]): - hidden_states = attn(hidden_states, encoder_hidden_states) - hidden_states = resnet(hidden_states, temb) - - return hidden_states diff --git a/tensorrt_llm/models/unet/unet_2d_condition.py b/tensorrt_llm/models/unet/unet_2d_condition.py deleted file mode 100644 index 96531df8db7e..000000000000 --- a/tensorrt_llm/models/unet/unet_2d_condition.py +++ /dev/null @@ -1,249 +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. -from typing import Optional - -from ...functional import cast, concat, silu -from ...layers import Conv2d, GroupNorm -from ...module import Module, ModuleList -from .embeddings import TimestepEmbedding, Timesteps -from .unet_2d_blocks import (UNetMidBlock2DCrossAttn, get_down_block, - get_up_block) - - -class UNet2DConditionModel(Module): - - def __init__( - self, - sample_size=None, - in_channels=4, - out_channels=4, - center_input_sample=False, - flip_sin_to_cos=True, - freq_shift=0, - down_block_types=("CrossAttnDownBlock2D", "CrossAttnDownBlock2D", - "CrossAttnDownBlock2D", "DownBlock2D"), - up_block_types=("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", - "CrossAttnUpBlock2D"), - block_out_channels=(320, 640, 1280, 1280), - layers_per_block=2, - downsample_padding=1, - mid_block_scale_factor=1.0, - act_fn="silu", - norm_num_groups=32, - norm_eps=1e-5, - cross_attention_dim=1280, - transformer_layers_per_block=1, - attention_head_dim=8, - use_linear_projection=False, - addition_embed_type: Optional[str] = None, - addition_time_embed_dim: Optional[int] = None, - projection_class_embeddings_input_dim: Optional[int] = None, - dtype=None, - ): - super().__init__() - - self.sample_size = sample_size - self.addition_embed_type = addition_embed_type - time_embed_dim = block_out_channels[0] * 4 - - # input - self.conv_in = Conv2d(in_channels, - block_out_channels[0], - kernel_size=(3, 3), - padding=(1, 1), - dtype=dtype) - # time - self.time_proj = Timesteps(block_out_channels[0], - flip_sin_to_cos, - freq_shift, - dtype=dtype) - timestep_input_dim = block_out_channels[0] - - self.time_embedding = TimestepEmbedding(timestep_input_dim, - time_embed_dim, - dtype=dtype) - - if addition_embed_type == "text_time": - self.add_time_proj = Timesteps(addition_time_embed_dim, - flip_sin_to_cos, - freq_shift, - dtype=dtype) - self.add_embedding = TimestepEmbedding( - projection_class_embeddings_input_dim, - time_embed_dim, - dtype=dtype) - - down_blocks = [] - up_blocks = [] - - if isinstance(attention_head_dim, int): - attention_head_dim = (attention_head_dim, ) * len(down_block_types) - - if isinstance(transformer_layers_per_block, int): - transformer_layers_per_block = [transformer_layers_per_block - ] * len(down_block_types) - - # down - output_channel = block_out_channels[0] - for i, down_block_type in enumerate(down_block_types): - input_channel = output_channel - output_channel = block_out_channels[i] - is_final_block = i == len(block_out_channels) - 1 - - down_block = get_down_block( - down_block_type, - num_layers=layers_per_block, - transformer_layers_per_block=transformer_layers_per_block[i], - in_channels=input_channel, - out_channels=output_channel, - temb_channels=time_embed_dim, - add_downsample=not is_final_block, - resnet_eps=norm_eps, - resnet_act_fn=act_fn, - cross_attention_dim=cross_attention_dim, - attn_num_head_channels=attention_head_dim[i], - downsample_padding=downsample_padding, - use_linear_projection=use_linear_projection, - dtype=dtype) - down_blocks.append(down_block) - self.down_blocks = ModuleList(down_blocks) - # mid - self.mid_block = UNetMidBlock2DCrossAttn( - in_channels=block_out_channels[-1], - temb_channels=time_embed_dim, - resnet_eps=norm_eps, - resnet_act_fn=act_fn, - output_scale_factor=mid_block_scale_factor, - transformer_layers_per_block=transformer_layers_per_block[-1], - resnet_time_scale_shift="default", - cross_attention_dim=cross_attention_dim, - attn_num_head_channels=attention_head_dim[-1], - resnet_groups=norm_num_groups, - use_linear_projection=use_linear_projection, - dtype=dtype, - ) - # up - reversed_block_out_channels = list(reversed(block_out_channels)) - reversed_attention_head_dim = list(reversed(attention_head_dim)) - reversed_transformer_layers_per_block = list( - reversed(transformer_layers_per_block)) - output_channel = reversed_block_out_channels[0] - for i, up_block_type in enumerate(up_block_types): - prev_output_channel = output_channel - output_channel = reversed_block_out_channels[i] - input_channel = reversed_block_out_channels[min( - i + 1, - len(block_out_channels) - 1)] - - is_final_block = i == len(block_out_channels) - 1 - - up_block = get_up_block( - up_block_type, - num_layers=layers_per_block + 1, - transformer_layers_per_block= - reversed_transformer_layers_per_block[i], - in_channels=input_channel, - out_channels=output_channel, - prev_output_channel=prev_output_channel, - temb_channels=time_embed_dim, - add_upsample=not is_final_block, - resnet_eps=norm_eps, - resnet_act_fn=act_fn, - cross_attention_dim=cross_attention_dim, - attn_num_head_channels=reversed_attention_head_dim[i], - use_linear_projection=use_linear_projection, - dtype=dtype, - ) - up_blocks.append(up_block) - prev_output_channel = output_channel - self.up_blocks = ModuleList(up_blocks) - # out - self.conv_norm_out = GroupNorm(num_channels=block_out_channels[0], - num_groups=norm_num_groups, - eps=norm_eps, - dtype=dtype) - self.conv_act = silu - self.conv_out = Conv2d(block_out_channels[0], - out_channels, (3, 3), - padding=(1, 1), - dtype=dtype) - - def forward(self, - sample, - timesteps, - encoder_hidden_states, - text_embeds=None, - time_ids=None): - # time - t_emb = self.time_proj(timesteps) - emb = self.time_embedding(t_emb) - - aug_emb = None - if self.addition_embed_type == "text_time": - assert text_embeds is not None and time_ids is not None - time_embeds = self.add_time_proj(time_ids.view([-1])) - time_embeds = time_embeds.view([text_embeds.shape[0], -1]) - add_embeds = concat([text_embeds, time_embeds], dim=1) - add_embeds = cast(add_embeds, emb.dtype) - aug_emb = self.add_embedding(add_embeds) - - emb = emb + aug_emb if aug_emb is not None else emb - - sample = self.conv_in(sample) - - down_block_res_samples = (sample, ) - for downsample_block in self.down_blocks: - - if hasattr( - downsample_block, - "attentions") and downsample_block.attentions is not None: - - sample, res_samples = downsample_block( - hidden_states=sample, - temb=emb, - encoder_hidden_states=encoder_hidden_states) - else: - sample, res_samples = downsample_block(hidden_states=sample, - temb=emb) - down_block_res_samples += res_samples - - sample = self.mid_block(sample, - emb, - encoder_hidden_states=encoder_hidden_states) - - for upsample_block in self.up_blocks: - - res_samples = down_block_res_samples[-len(upsample_block.resnets):] - down_block_res_samples = down_block_res_samples[:-len(upsample_block - .resnets)] - - if hasattr(upsample_block, - "attentions") and upsample_block.attentions is not None: - sample = upsample_block( - hidden_states=sample, - temb=emb, - res_hidden_states_tuple=res_samples, - encoder_hidden_states=encoder_hidden_states, - ) - else: - sample = upsample_block(hidden_states=sample, - temb=emb, - res_hidden_states_tuple=res_samples) - - sample = self.conv_norm_out(sample) - sample = self.conv_act(sample) - sample = self.conv_out(sample) - - return sample diff --git a/tensorrt_llm/models/unet/weights.py b/tensorrt_llm/models/unet/weights.py deleted file mode 100644 index 9a8e4c7c63d1..000000000000 --- a/tensorrt_llm/models/unet/weights.py +++ /dev/null @@ -1,197 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -import time - -import numpy as np - -from ...logger import logger - - -def update_timestep_weight(src, dst): - dst.linear_1.update_parameters(src.linear_1) - dst.linear_2.update_parameters(src.linear_2) - - -def update_crossattn_downblock_2d_weight(src, dst): - for index, value in enumerate(src.resnets): - update_resnet_block_weight(value, dst.resnets[index]) - - for index, value in enumerate(src.attentions): - update_transformer_2d_model_weight(dst.attentions[index], value) - - if src.downsamplers: - for index, value in enumerate(src.downsamplers): - dst.downsamplers[index].conv.update_parameters(value.conv) - - -def update_transformer_2d_model_weight(gm, m): - gm.norm.update_parameters(m.norm) - gm.proj_in.update_parameters(m.proj_in) - for i in range(len(gm.transformer_blocks)): - gm.transformer_blocks[i].attn1.to_qkv.weight.value = np.concatenate( - (m.transformer_blocks[i].attn1.to_q.weight.detach().cpu().numpy(), - m.transformer_blocks[i].attn1.to_k.weight.detach().cpu().numpy(), - m.transformer_blocks[i].attn1.to_v.weight.detach().cpu().numpy())) - gm.transformer_blocks[i].attn1.to_out.update_parameters( - m.transformer_blocks[i].attn1.to_out[0]) - - gm.transformer_blocks[i].attn2.to_q.update_parameters( - m.transformer_blocks[i].attn2.to_q) - gm.transformer_blocks[i].attn2.to_kv.weight.value = np.concatenate( - (m.transformer_blocks[i].attn2.to_k.weight.detach().cpu().numpy(), - m.transformer_blocks[i].attn2.to_v.weight.detach().cpu().numpy())) - gm.transformer_blocks[i].attn2.to_out.update_parameters( - m.transformer_blocks[i].attn2.to_out[0]) - - gm.transformer_blocks[i].ff.proj_in.update_parameters( - m.transformer_blocks[i].ff.net[0].proj) - gm.transformer_blocks[i].ff.proj_out.update_parameters( - m.transformer_blocks[i].ff.net[2]) - - gm.transformer_blocks[i].norm1.update_parameters( - m.transformer_blocks[i].norm1) - gm.transformer_blocks[i].norm2.update_parameters( - m.transformer_blocks[i].norm2) - gm.transformer_blocks[i].norm3.update_parameters( - m.transformer_blocks[i].norm3) - - gm.proj_out.update_parameters(m.proj_out) - - -def update_upblock_2d_weight(src, dst): - if src.upsamplers: - for index, value in enumerate(src.upsamplers): - dst.upsamplers[index].conv.update_parameters(value.conv) - - for index, value in enumerate(src.resnets): - dst.resnets[index].norm1.update_parameters(value.norm1) - dst.resnets[index].conv1.update_parameters(value.conv1) - - dst.resnets[index].norm2.update_parameters(value.norm2) - dst.resnets[index].conv2.update_parameters(value.conv2) - - if value.conv_shortcut: - dst.resnets[index].conv_shortcut.update_parameters( - value.conv_shortcut) - - dst.resnets[index].time_emb_proj.update_parameters(value.time_emb_proj) - - -def update_downblock_2d_weight(src, dst): - if src.downsamplers: - for index, value in enumerate(src.downsamplers): - dst.downsamplers[index].conv.update_parameters(value.conv) - - for index, value in enumerate(src.resnets): - dst.resnets[index].norm1.update_parameters(value.norm1) - dst.resnets[index].conv1.update_parameters(value.conv1) - - dst.resnets[index].norm2.update_parameters(value.norm2) - dst.resnets[index].conv2.update_parameters(value.conv2) - - if value.conv_shortcut: - dst.resnets[index].conv_shortcut.update_parameters( - value.conv_shortcut) - - dst.resnets[index].time_emb_proj.update_parameters(value.time_emb_proj) - - -def update_unet_mid_block_2d_weight(src, dst): - for index, value in enumerate(src.resnets): - update_resnet_block_weight(value, dst.resnets[index]) - - for index, value in enumerate(src.attentions): - update_transformer_2d_model_weight(dst.attentions[index], value) - - -def update_crossattn_upblock_2d_weight(src, dst): - for index, value in enumerate(src.resnets): - update_resnet_block_weight(value, dst.resnets[index]) - - for index, value in enumerate(src.attentions): - update_transformer_2d_model_weight(dst.attentions[index], value) - if src.upsamplers: - for index, value in enumerate(src.upsamplers): - dst.upsamplers[index].conv.update_parameters(value.conv) - - -def update_resnet_block_weight(src, dst): - dst.norm1.update_parameters(src.norm1) - dst.conv1.update_parameters(src.conv1) - - dst.norm2.update_parameters(src.norm2) - dst.conv2.update_parameters(src.conv2) - - dst.time_emb_proj.update_parameters(src.time_emb_proj) - if src.conv_shortcut: - dst.conv_shortcut.update_parameters(src.conv_shortcut) - - -def update_unetmidblock_2d_weight(src, dst): - for index, value in enumerate(src.attentions): - dst.attentions[index].group_norm.update_parameters(value.group_norm) - dst.attentions[index].proj_attn.update_parameters(value.proj_attn) - - dst.attentions[index].qkv.weight.value = np.concatenate( - (value.query.weight.detach().cpu().numpy(), - value.key.weight.detach().cpu().numpy(), - value.value.weight.detach().cpu().numpy())) - dst.attentions[index].qkv.bias.value = np.concatenate( - (value.query.bias.detach().cpu().numpy(), - value.key.bias.detach().cpu().numpy(), - value.value.bias.detach().cpu().numpy())) - - for index, value in enumerate(src.resnets): - update_resnet_block_weight(value, dst.resnets[index]) - - -def update_unet_2d_condition_model_weights(src, dst): - dst.conv_in.update_parameters(src.conv_in) - - dst.time_embedding.update_parameters(src.time_embedding) - if src.config.addition_embed_type: - dst.add_embedding.update_parameters(src.add_embedding) - - for index, type in enumerate(src.config.down_block_types): - if type == 'CrossAttnDownBlock2D': - update_crossattn_downblock_2d_weight(src.down_blocks[index], - dst.down_blocks[index]) - elif type == 'DownBlock2D': - update_downblock_2d_weight(src.down_blocks[index], - dst.down_blocks[index]) - - update_unet_mid_block_2d_weight(src.mid_block, dst.mid_block) - - for index, type in enumerate(src.config.up_block_types): - if type == 'CrossAttnUpBlock2D': - update_crossattn_upblock_2d_weight(src.up_blocks[index], - dst.up_blocks[index]) - elif type == 'UpBlock2D': - update_upblock_2d_weight(src.up_blocks[index], dst.up_blocks[index]) - - dst.conv_norm_out.update_parameters(src.conv_norm_out) - - dst.conv_out.update_parameters(src.conv_out) - - -def load_from_hf_unet(src, dst): - logger.info('Loading weights from HF Unet...') - tik = time.time() - - update_unet_2d_condition_model_weights(src, dst) - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - logger.info(f'Weights loaded. Total time: {t}') diff --git a/tensorrt_llm/module.py b/tensorrt_llm/module.py deleted file mode 100644 index c67674ec15b2..000000000000 --- a/tensorrt_llm/module.py +++ /dev/null @@ -1,288 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -import operator - -from ._common import default_net -from .logger import logger - - -def _addindent(s_, numSpaces): - s = s_.split("\n") - # don't do anything for single-line stuff - if len(s) == 1: - return s_ - first = s.pop(0) - s = [(numSpaces * " ") + line for line in s] - s = "\n".join(s) - s = first + "\n" + s - return s - - -class Module(object): - def __init__(self) -> None: - self._modules = {} - self._parameters = {} - self._network_outputs = {} - - def forward(self, *args, **kwargs): - raise NotImplementedError - - def __call__(self, *args, **kwargs): - current_net = default_net() - if not current_net._module_call_stack.module_names_set(): - logger.debug("Initializing top level module") - current_net._module_call_stack.set_module_names(self) - unique_name = current_net._module_call_stack.get_mod_name(self) - with current_net._module_call_stack.call_stack_mgr() as stack: - stack.append(unique_name) - start_layer_idx = current_net.trt_network.num_layers - output = self.forward(*args, **kwargs) - end_layer_idx = current_net.trt_network.num_layers - current_net._module_call_stack.set_layer_range( - self, range(start_layer_idx, end_layer_idx) - ) - return output - - def __getattr__(self, name): - parameters = self.__dict__.get("_parameters") - if parameters is not None and name in parameters: - return parameters[name] - - modules = self.__dict__.get("_modules") - if modules is not None and name in modules: - return modules[name] - - raise AttributeError("'{}' object has no attribute '{}'".format(type(self).__name__, name)) - - def __setattr__(self, name, value) -> None: - from .parameter import Parameter - - # Improved module setattr to handle one edge case: - # attribute could be first set to None and later reset to Parameter / Module class - - try: - super().__getattribute__(name) - - except AttributeError: - # if base class doesn't have the attribute, no matter we init or reset: - # - keep Parameter and Module attrs in this Module class - # - leave all other attrs in base class - if isinstance(value, Parameter): - self.__dict__.get("_parameters")[name] = value - elif isinstance(value, Module): - self.__dict__.get("_modules")[name] = value - else: - super().__setattr__(name, value) - - else: - # if base class has the attribute, reset as follows: - # - when reset as Parameter or Module attr, remove from base & add to this Module class - # - other types reset and remain in base class - if isinstance(value, Parameter): - super().__delattr__(name) - self.__dict__.get("_parameters")[name] = value - elif isinstance(value, Module): - super().__delattr__(name) - self.__dict__.get("_modules")[name] = value - else: - super().__setattr__(name, value) - - def named_modules(self, memo=None, prefix="", remove_duplicate=True): - if memo is None: - memo = set() - if self not in memo: - if remove_duplicate: - memo.add(self) - yield prefix, self - for name, module in self._modules.items(): - if module is None: - continue - submodule_prefix = prefix + ("." if prefix else "") + name - for m in module.named_modules(memo, submodule_prefix, remove_duplicate): - yield m - - def named_modules_with_parent(self, memo=None, prefix="", parent=None, remove_duplicate=True): - if memo is None: - memo = set() - if self not in memo: - if remove_duplicate: - memo.add(self) - yield prefix, self, parent - - if parent: - # Use the up-to-date module from the parent, to allow replacing - # layers while iterating this generator. - module_name = prefix.rsplit(".", 1)[-1] - module = getattr(parent, module_name) - if module is None: - return - else: - module = self - - for child_name, child_module in module._modules.items(): - if child_module is None: - continue - submodule_prefix = prefix + ("." if prefix else "") + child_name - for m in child_module.named_modules_with_parent( - memo, submodule_prefix, module, remove_duplicate - ): - yield m - - def named_children(self): - memo = set() - for name, module in self._modules.items(): - if module is not None and module not in memo: - memo.add(module) - yield name, module - - def _named_members(self, get_members_fn, prefix="", recurse=True): - memo = set() - modules = self.named_modules(prefix=prefix) if recurse else [(prefix, self)] - for module_prefix, module in modules: - members = get_members_fn(module) - for k, v in members: - if v is None or v in memo: - continue - memo.add(v) - name = module_prefix + ("." if module_prefix else "") + k - yield name, v - - def parameters(self, recurse=True): - for name, param in self.named_parameters(): - yield param - - def named_parameters(self, prefix="", recurse=True): - gen = self._named_members( - lambda module: module._parameters.items(), prefix=prefix, recurse=recurse - ) - for elem in gen: - yield elem - - def children(self): - for _, module in self.named_children(): - yield module - - def apply(self, fn): - for module in self.children(): - module.apply(fn) - fn(self) - return self - - def _get_name(self): - return self.__class__.__name__ - - def register_parameter(self, name, param): - if param is None: - self._parameters[name] = None - else: - self._parameters[name] = param - - def register_network_output(self, name, value): - self._network_outputs[name] = value - - def named_network_outputs(self): - for name, module in self.named_modules(): - for n, output in module._network_outputs.items(): - yield name + ("." if name else "") + n, output - - def update_parameters(self, torch_module): - m = {k: v for k, v in self.named_parameters()} - tm = {k: v for k, v in torch_module.named_parameters()} - - assert sorted(m.keys()) == sorted(tm.keys()), ( - "The parameter names of the TensorRT LLM module must be the same with the torch module" - ) - - for k, v in self.named_parameters(): - v.value = tm[k].detach().cpu().numpy() - - def _get_name(self): - return self.__class__.__name__ - - def __repr__(self): - # We treat the extra repr like the sub-module, one item per line - child_lines = [] - for key, module in self._modules.items(): - mod_str = repr(module) - mod_str = _addindent(mod_str, 2) - child_lines.append("(" + key + "): " + mod_str) - main_str = self._get_name() + "(" - if child_lines: - # simple one-liner info, which most builtin Modules will use - main_str += "\n " + "\n ".join(child_lines) + "\n" - main_str += ")" - return main_str - - -class ModuleList(Module): - def __init__(self, modules) -> None: - super(ModuleList, self).__init__() - offset = len(self) - for i, module in enumerate(modules): - self._modules[str(offset + i)] = module - - def _get_abs_string_index(self, idx): - """Get the absolute index for the list of modules.""" - idx = operator.index(idx) - if not (-len(self) <= idx < len(self)): - raise IndexError("index {} is out of range".format(idx)) - if idx < 0: - idx += len(self) - return str(idx) - - def __getitem__(self, idx): - if isinstance(idx, slice): - return self.__class__(list(self._modules.values())[idx]) - else: - return self._modules[self._get_abs_string_index(idx)] - - def __setitem__(self, idx, module) -> None: - idx = self._get_abs_string_index(idx) - return setattr(self, str(idx), module) - - def __len__(self): - return len(self._modules) - - def __repr__(self): - """Return a custom repr for ModuleList that compresses repeated module representations.""" - list_of_reprs = [repr(item) for item in self] - if len(list_of_reprs) == 0: - return self._get_name() + "()" - - start_end_indices = [[0, 0]] - repeated_blocks = [list_of_reprs[0]] - for i, r in enumerate(list_of_reprs[1:], 1): - if r == repeated_blocks[-1]: - start_end_indices[-1][1] += 1 - continue - - start_end_indices.append([i, i]) - repeated_blocks.append(r) - - lines = [] - main_str = self._get_name() + "(" - for (start_id, end_id), b in zip(start_end_indices, repeated_blocks): - local_repr = f"({start_id}): {b}" # default repr - - if start_id != end_id: - n = end_id - start_id + 1 - local_repr = f"({start_id}-{end_id}): {n} x {b}" - - local_repr = _addindent(local_repr, 2) - lines.append(local_repr) - - main_str += "\n " + "\n ".join(lines) + "\n" - main_str += ")" - return main_str diff --git a/tensorrt_llm/network.py b/tensorrt_llm/network.py deleted file mode 100644 index e2d3dc453f6b..000000000000 --- a/tensorrt_llm/network.py +++ /dev/null @@ -1,959 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -import collections -import contextlib -import hashlib -import inspect -import weakref -from collections import OrderedDict, defaultdict -from dataclasses import dataclass, field -from pathlib import Path -from typing import (Any, Dict, Iterable, List, Optional, OrderedDict, Set, - Tuple, Union) - -import numpy as np -import onnx -import onnx_graphsurgeon as gs -import tensorrt as trt - -from tensorrt_llm.module import Module - -from ._common import set_network -from ._utils import get_extra_attr, has_extra_attr, set_extra_attr -from .logger import logger -from .plugin import PluginConfig - - -class _UniqueNameGenerator(object): - - def __init__(self, prefix=''): - self.ids = collections.defaultdict(int) - self.prefix = prefix - - def __call__(self, key, module_name=''): - if module_name != '': - module_name = module_name.replace(".", "/") - key = module_name + '/' + key - tmp = self.ids[key] - self.ids[key] += 1 - return f"{self.prefix}{key}_{tmp}" - - -class PluginInfo: - plugin_creator: trt.IPluginCreator - plugin_name: str - pfc: trt.PluginFieldCollection - - def __init__(self, plugin_creator: trt.IPluginCreator, plugin_name: str, - pfc: trt.PluginFieldCollection): - self.plugin_creator = plugin_creator - self.plugin_name = plugin_name - self.pfc = pfc - self._parse_pfc(pfc) - - def _parse_pfc(self, pfc: trt.PluginFieldCollection): - self.pfc_as_ndarray = {} - self.pfc_as_list = {} - for i in range(len(pfc)): - name, data = pfc[i].name, pfc[i].data - array_data = data - self.pfc_as_ndarray[name] = array_data.copy() - list_data = array_data.tolist() - self.pfc_as_list[name] = list_data - - -def get_plugin_info(trt_network: trt.INetworkDefinition, - layer_name: str) -> PluginInfo: - if not has_extra_attr(trt_network, "plugin_infos"): - return None - plugin_infos = get_extra_attr(trt_network, "plugin_infos") - if layer_name not in plugin_infos: - return None - return plugin_infos[layer_name] - - -def set_plugin_info(trt_network: trt.INetworkDefinition, layer_name: str, - plugin_info: PluginInfo): - if not has_extra_attr(trt_network, "plugin_infos"): - set_extra_attr(trt_network, "plugin_infos", {}) - plugin_infos = get_extra_attr(trt_network, "plugin_infos") - plugin_infos[layer_name] = plugin_info - - -def delete_plugin_info(trt_network: trt.INetworkDefinition, layer_name: str): - if not has_extra_attr(trt_network, "plugin_infos"): - return - plugin_infos = get_extra_attr(trt_network, "plugin_infos") - if layer_name not in plugin_infos: - return - del plugin_infos[layer_name] - - -# TODO: remove this WAR after https://nvbugs/4359151 fixed. -def get_np_weight(trt_network: trt.INetworkDefinition, - layer_name: str) -> np.array: - if not has_extra_attr(trt_network, "np_weights"): - return None - np_weights = get_extra_attr(trt_network, "np_weights") - if layer_name not in np_weights: - return None - return np_weights[layer_name] - - -# TODO: remove this WAR after https://nvbugs/4359151 fixed. -def set_np_weight(trt_network: trt.INetworkDefinition, layer_name: str, - np_weight: np.array): - if not has_extra_attr(trt_network, "np_weights"): - set_extra_attr(trt_network, "np_weights", {}) - np_weights = get_extra_attr(trt_network, "np_weights") - np_weights[layer_name] = np_weight - - -class Network(object): - - def __init__(self, **kwargs): - # intentionally use **kwargs, user should never call this ctor directly - # use Builder.create_network() instead - - # Holds the removed layers and disable them in graph rewriting and other phases. - # This is a hacky way since INetwork python API doesn't provide a way to remove a layer. - # TODO: remove this when TensorRT provides a better way to remove a layer - self._removed_layers: Set[str] = set() - - self.is_graph_altered = False - - from .graph_rewriting import FLayerInfoMemo - self.flayer_memo = FLayerInfoMemo() # holds the functional metadata - self._parameter_tensors = {} # holds the parameter tensors - - def _init(self, trt_network): - self._trt_network = trt_network - self._inputs = {} - self._named_parameters = None - # layer precision of a given scope, this is used together with precision(dtype) context manager - self._dtype = None - self._name_generator = _UniqueNameGenerator() - self._plugin_config = PluginConfig() - self._module_call_stack = _TrtLlmModuleCallStack() - self._registered_ndarrays = [] - self._strongly_typed = trt.INetworkDefinition.get_flag( - self._trt_network, trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED) - self._unfilled_weights: Dict[str, Tuple[np.array, np.array]] = {} - - return self - - def _register_unfilled_weights(self, layer_name: str, weights: np.array, - values: np.array): - self._unfilled_weights[layer_name] = (weights, values) - - def _fill_weights(self): - from tensorrt_llm.parameter import Parameter - - for layer_name in list(self._unfilled_weights.keys()): - weights, values = self._unfilled_weights.pop(layer_name) - self.register_ndarray(weights) - if values is not None: - np.copyto(weights, values, casting='no') - else: - Parameter.xavier_init(weights) - - @property - def parameter_tensors(self): - return self._parameter_tensors - - def get_parameter_tensor(self, param): - return self.parameter_tensors.get(param, None) - - def set_parameter_tensor(self, param, tensor): - assert param not in self.parameter_tensors - self.parameter_tensors[param] = tensor - - @property - def dtype(self) -> trt.DataType: - return self._dtype - - @dtype.setter - def dtype(self, dtype: trt.DataType): - assert isinstance(dtype, trt.DataType) or dtype is None - self._dtype = dtype - - @property - def trt_network(self) -> trt.INetworkDefinition: - return self._trt_network - - @property - def plugin_config(self) -> PluginConfig: - return self._plugin_config - - @plugin_config.setter - def plugin_config(self, cfg: PluginConfig): - assert isinstance( - cfg, - PluginConfig), f"Expecting a PluginConfig object, got {type(cfg)}" - self._plugin_config = cfg - - @property - def strongly_typed(self) -> bool: - return self._strongly_typed - - def _add_input(self, - tensor, - name, - dtype, - shape, - dim_range: OrderedDict = None): - assert isinstance(dtype, trt.DataType) - tensor.trt_tensor = self.trt_network.add_input( - name=name, - shape=shape, - dtype=dtype, - ) - assert tensor.trt_tensor is not None, f"Couldn't create TRT tensor for {name} {dtype} {shape}" - if dim_range is not None: - logger.debug( - f'Add input: {name}, shape: {shape}, dtype: {dtype}, dimension names:{list(dim_range.keys())}' - ) - for i, dim_name in enumerate(dim_range.keys()): - tensor.trt_tensor.set_dimension_name(i, str(dim_name)) - else: - logger.debug(f'Add input: {name}, shape: {shape}, dtype: {dtype}') - self._inputs[name] = tensor - - def _mark_output(self, tensor, name, dtype): - from .functional import cast - - # In strongly_typed, if tensor output is not the same, add a cast - if dtype is not None and self.strongly_typed: - tensor = cast(tensor, dtype) - self.trt_network.mark_output(tensor.trt_tensor) - tensor.trt_tensor.name = name - if not self.strongly_typed: - tensor.trt_tensor.dtype = dtype or tensor.trt_tensor.dtype - logger.debug(f'Mark output: {name}, dtype: {dtype}') - - def set_named_parameters(self, named_parameters): - self._named_parameters = named_parameters - - @property - def named_parameters(self): - return self._named_parameters - - def _set_layer_name(self, layer): - original_layer_name = layer.name - layer_name = str(layer.type).split('.')[-1] - current_module = self._module_call_stack.get_current_module() - - func_stack = [] - frame = inspect.currentframe().f_back.f_back - while frame: - func_name = frame.f_code.co_name - line_num = frame.f_lineno - if func_name == "forward": - break - func_stack.insert(0, f"{func_name}_L{line_num}") - if len(func_stack) >= 10: - # NOTE: TRT error messages has a character limit. - # Limiting to only 10 levels helps retain - # the true error message from TRT. - break - frame = frame.f_back - current_module = f"{current_module}.{'.'.join(func_stack)}" - - if layer.type == trt.LayerType.PLUGIN_V2: - layer_name = '_'.join( - [layer_name, - str(layer.plugin.plugin_type).split('.')[-1]]) - elif layer.type in [ - trt.LayerType.UNARY, trt.LayerType.REDUCE, - trt.LayerType.ELEMENTWISE - ]: - layer_name = '_'.join([layer_name, str(layer.op).split('.')[-1]]) - - layer.name = self._name_generator(layer_name, current_module) - for idx in range(layer.num_outputs): - # TRT initializes tensor names from the initial layer's name when the layer is created, - # and does not update tensor names when layer name changed by application, needs to - # change the tensor name to align with the new layer name for better debugging - layer.get_output(idx).name = f"{layer.name}_output_{idx}" - if original_layer_name != layer.name: - if layer.type == trt.LayerType.PLUGIN_V2: - plugin_info = get_plugin_info(self.trt_network, - original_layer_name) - if plugin_info is not None: - set_plugin_info(self.trt_network, layer.name, plugin_info) - delete_plugin_info(self.trt_network, original_layer_name) - - # Set layer metadata to the same as the layer name so that it can show up in NVTX. - layer.metadata = layer.name - - def register_ndarray(self, ndarray: np.ndarray) -> None: - ''' When the functional APIs need to create local numpy array and use as weights for constant or other layers, - they need to register the ndarray objects to the TRT-LLM Network to prolong the lifetime of the ndarray, such that weights are - still valid when functional API returned. - All the weights referenced by the trt Network are weak referenced, it's TRT-LLM's responsibility to keep the weights alive - during the TRT network construction and TRT engine building process. - ''' - self._registered_ndarrays.append(ndarray) - - def _generate_optimization_profiles(self) -> List[trt.IOptimizationProfile]: - input_tensors = self._inputs - if len(input_tensors) == 0: - return [] - num_profiles = len(list(input_tensors.values())[0].profiles) - profiles = [] - for i in range(num_profiles): - logger.debug(f'Adding optimization profile {i+1}/{num_profiles}') - profile = self._trt_network.builder.create_optimization_profile() - for input_name, input_tensor in input_tensors.items(): - shape_profile = input_tensor.profiles[i] - min_shape = list(shape_profile.min) - opt_shape = list(shape_profile.opt) - max_shape = list(shape_profile.max) - if input_tensor.trt_tensor.is_shape_tensor: - profile.set_shape_input(input_name, min_shape, opt_shape, - max_shape) - else: - profile.set_shape(input_name, min_shape, opt_shape, - max_shape) - logger.debug( - f'{input_name}, min: {min_shape}, opt: {opt_shape}, max: {max_shape}' - ) - profiles.append(profile) - return profiles - - def get_inputs(self): - ''' - Get the inputs of the network. - - Returns: - Iterable[Tensor] - ''' - return self._inputs.values() - - def get_outputs(self): - ''' - Get the outputs of the network. - - Returns: - Iterable[Tensor] - ''' - from .functional import Tensor - for i in range(self._trt_network.num_outputs): - tensor = self._trt_network.get_output(i) - yield Tensor(trt_tensor=tensor, - network=self, - is_network_input=False) - - def is_input(self, tensor) -> bool: - ''' - Tell if a tensor is a input of the network. - - Parameters: - tensor: Union[Tensor, str, trt.ITensor] - ''' - from .functional import Tensor - - if isinstance(tensor, str): - tensor_name = tensor - elif isinstance(tensor, (trt.ITensor, Tensor)): - tensor_name = tensor.name - else: - raise ValueError( - f"tensor should be Tensor, str or ITensor, got {tensor}") - - return self._inputs.get(tensor_name, False) - - def is_output(self, tensor) -> bool: - ''' - Tell if a tensor is a output of the network. - - Parameters: - tensor: Tensor - ''' - for i in range(self._trt_network.num_outputs): - if tensor.trt_tensor is self._trt_network.get_output(i): - return True - return False - - def get_layers(self) -> Iterable["Layer"]: - ''' - Get all the layers of network. - - Returns: - Iterable[Layer] - ''' - from .graph_rewriting import Layer - for i in range(self._trt_network.num_layers): - layer = Layer(network=self, - trt_layer=self._trt_network.get_layer(i)) - yield layer - - def get_layer_by_name(self, name: str) -> Optional["Layer"]: - state = self._get_graph() - return state.name_to_layer.get(name, None) - - def get_tensor_users(self, tensor) -> Iterable["Layer"]: - ''' - Get the layers those consumes this tensor. - ''' - state = self._get_graph() - for layer in state.tensor_to_consumers[tensor]: - yield layer - - def get_tensor_parent(self, tensor) -> Optional["Layer"]: - ''' - Get the layer that produces this tensor. - ''' - state = self._get_graph() - return state.tensor_to_producer.get(tensor, None) - - def mark_removed_layer(self, layer: "Layer"): - from .graph_rewriting import FLayerInfoMemo - self._removed_layers.add(layer.name) - - # Try to delete the layer if it is a Plugin - FLayerInfoMemo.instance().remove(layer.name) - - def is_removed_layer(self, layer: "Layer") -> bool: - return layer.name in self._removed_layers - - @property - def removed_layers(self) -> Iterable["Layer"]: - for layer_name in self._removed_layers: - layer = self.get_layer_by_name(layer_name) - assert layer, "Invalid layer name" - yield layer - - def to_dot(self, path: Union[str, Path] = None) -> Optional[str]: - ''' - Get a graphviz representation of the network. - - NOTE, the graph might be redundancy since TRT's INetwork won't clean the unused inputs and layers - automatically. - TODO: add an flag to hide all the removed layers and their output tensors - TODO: replace this when TensorRT provides a better way to get the graph of INetworkDefinition - TODO: a little feature, add blocks in the figure to highlight the subgraphes of Modules - - Parameters: - path: the path to save the graphviz file, if not provided, will return the graphviz source code - ''' - format = 'text' if not path else path.split('.')[-1] - - try: - import graphviz - except ImportError: - logger.error( - "Failed to import graphviz, please install graphviz to enable Network.to_dot()" - ) - return - - dot = graphviz.Digraph( - comment= - f'TensorRT Graph of {self._get_network_hash(lightweight=False)}', - format=format if format != 'text' else None) - - inputs_names = set([x.name for x in self.get_inputs()]) - output_names = set([x.name for x in self.get_outputs()]) - - node_style = dict( - shape='box', - style='rounded,filled,bold', - fontname='Arial', - fillcolor='#ffffff', - color='#303A3A', - width='1.3', - height='0.84', - ) - - hl_node_style = dict( - shape='box', - style='rounded,filled,bold', - fontname='Arial', - fillcolor='lightblue', - color='#303A3A', - width='1.3', - height='0.84', - ) - - state = self._get_graph() - nodes = set() - tensor_to_alias = {} - tensor_id = [0] - - def get_alias(tensor, tensor_id): - if tensor not in tensor_to_alias: - if (tensor not in inputs_names) and (tensor - not in output_names): - tensor_to_alias[tensor] = f"t{tensor_id[0]}" - tensor_id[0] += 1 - else: - tensor_to_alias[tensor] = tensor - - return tensor_to_alias[tensor] - - def create_tensor_node(tensor: str, dtype=None, shape=None): - tensor_alias = get_alias(tensor, tensor_id) - if tensor_alias not in nodes: - dot.node(tensor_alias, - str(dtype) + "\n" + tensor_alias + "\n" + str(shape), - **node_style) - nodes.add(tensor_alias) - return tensor_alias - - def create_layer_node(layer: str): - if layer not in nodes: - dot.node(layer, layer, **hl_node_style) - nodes.add(layer) - - for tensor, layer in state.tensor_to_producer.items(): - tensor_alias = create_tensor_node(tensor.name, tensor.dtype, - tensor.shape) - create_layer_node(layer.name) - dot.edge(layer.name, tensor_alias) - for tensor, layers in state.tensor_to_consumers.items(): - tensor_alias = create_tensor_node(tensor.name, tensor.dtype, - tensor.shape) - for layer in layers: - create_layer_node(layer.name) - dot.edge(tensor_alias, layer.name) - - if format == "text": - return dot.source - dot.save(path) - - def to_onnx(self, path: str = None) -> None: - ''' - Export the network into a "ONNX-like" file for visualization. - - Parameters: - path: the path to the output file - ''' - if path is None: - return - - network = self.trt_network - b_onnx_type = True # Use ONNX style operator names - - def layer_type_to_class(layer: trt.ILayer = None) -> trt.ILayer: - ''' - Convert trt.ILayer into a exact TensorRT layer - ''' - layer_type_name = str(layer.type)[10:] - # Some special cases - if layer_type_name == "ELEMENTWISE": - return trt.IElementWiseLayer - if layer_type_name == "LRN": - return trt.ILRNLayer - if layer_type_name == "NMS": - return trt.INMSLayer - if layer_type_name == "PARAMETRIC_RELU": - return trt.IParametricReLULayer - if layer_type_name == "PLUGIN": - return None # IPluginLayer is not supported any more - if layer_type_name == "RAGGED_SOFTMAX": - return trt.IRaggedSoftMaxLayer - if layer_type_name == "SOFTMAX": - return trt.ISoftMaxLayer - if layer_type_name == "TOPK": - return trt.ITopKLayer - - # general cases, e.g. MATRIX_MULTIPLY -> MatrixMultiply - name = "".join(name[0] + name[1:].lower() - for name in layer_type_name.split("_")) - return trt.__builtins__["getattr"](trt, f"I{name}Layer") - - def convert_type_to_onnx(node_type: str = "", - attribution: OrderedDict = OrderedDict()): - ''' - Convert layer names of TensorRT style into ONNX style - ''' - if node_type == "ACTIVATION": - convert_list = { - "RELU": "Relu", - "SIGMOID": "Sigmoid", - "TANH": "Tanh", - "LEAKY_RELU": "LeakyRelu", - "ELU": "Elu", - "SELU": "Selu", - "SOFTSIGN": "Softsign", - "SOFTPLUS": "Softplus", - "CLIP": "Clip", - "HARD_SIGMOID": "HardSigmoid", - "SCALED_TANH": "ScaledTanh", - "THRESHOLDED_RELU": "ThresholdedRelu", - } - # No corresponding operator for GELU_ERF, GELU_TANH - if "algo-type" in attribution.keys() and attribution[ - "algo-type"].split(".")[-1] in convert_list.keys(): - return convert_list[attribution["algo-type"].split(".")[-1]] - return node_type - if node_type == "CAST": - return "Cast" - if node_type == "CONCATENATION": - return "Concat" - if node_type == "CONSTANT": - return "Constant" - if node_type == "CONVOLUTION": - return "Conv" - if node_type == "DECONVOLUTION": - return "Deconv" - if node_type == "ELEMENTWISE": - convert_list = { - "SUM": "Add", - "PROD": "Mul", - "MAX": "Max", - "MIN": "Min", - "SUB": "Sub", - "DIV": "Div", - "POW": "Pow", - "AND": "And", - "OR": "Or", - "XOR": "Xor", - "EQUAL": "Equal", - "GREATER": "Greater", - "LESS": "Less", - } - if "op" in attribution.keys() and attribution["op"].split( - ".")[-1] in convert_list.keys(): - return convert_list[attribution["op"].split(".")[-1]] - return node_type - if node_type == "GATHER": - return "Gather" - if node_type == "LOOP": - return "Loop" - if node_type == "MATRIX_MULTIPLY": - return "Gemm" - if node_type == "POOLING": - convert_list = {"MAX": "MaxPool", "AVERAGE": "AveragePool"} - if "algo-type" in attribution.keys() and attribution[ - "algo-type"].split(".")[-1] in convert_list.keys(): - return convert_list[attribution["algo-type"].split(".")[-1]] - return node_type - if node_type == "REDUCE": - convert_list = { - "SUM": "ReduceSum", - "PROD": "ReduceProd", - "MAX": "ReduceMax", - "MIN": "ReduceMin", - "AVG": "ReduceMean", - } - if "op" in attribution.keys() and attribution["op"].split( - ".")[-1] in convert_list.keys(): - return convert_list[attribution["op"].split(".")[-1]] - return node_type - if node_type == "SELECT": - return "Where" - if node_type == "SHUFFLE": - return "Reshape" - if node_type == "SHAPE": - return "Shape" - if node_type == "SLICE": - return "Slice" - if node_type == "SOFTMAX": - return "Softmax" - if node_type == "TOPK": - return "TopK" - if node_type == "UNARY": - convert_list = {"SQRT": "Sqrt", "NOT": "Not"} - if "op" in attribution.keys() and attribution["op"].split( - ".")[-1] in convert_list.keys(): - return convert_list[attribution["op"].split(".")[-1]] - return node_type - return node_type - - def add_node_for_trt_network( - graph: gs.Graph = None, - node_name: str = "", - node_type: str = "", - input_list: List[gs.Variable] = [], - attribution: OrderedDict = OrderedDict(), - name_list: Union[str, List[str]] = "", - datatype_list: Union[np.dtype, List[np.dtype]] = [], - shape_list: Union[list, List[list]] = [], - number: int = 0, - b_onnx_type: bool = False, - ) -> Tuple[gs.Variable, int]: - ''' - Simplify version of function `add_node`, and we do some beautify to it. - ''' - if isinstance(name_list, list) or isinstance( - datatype_list, list) or isinstance( - shape_list, list): # Case of multi-output - assert len(name_list) == len(datatype_list) - assert len(name_list) == len(shape_list) - else: # Case of single-output - name_list = [name_list] - datatype_list = [datatype_list] - shape_list = [shape_list] - - n_output = len(name_list) - output_list = [] - for i in range(n_output): - tensor = gs.Variable(name_list[i], datatype_list[i], - shape_list[i]) - output_list.append(tensor) - - if b_onnx_type: - node_type = convert_type_to_onnx(node_type, attribution) - - node = gs.Node(node_type, - node_name, - inputs=input_list, - outputs=output_list, - attrs=attribution) - graph.nodes.append(node) # Update graph inside `add_node` - - if len(output_list) == 1: # Case of single-output - output_list = output_list[0] - return output_list, number + 1 - - graph = gs.Graph(nodes=[], inputs=[], outputs=[]) - graph.name = "" if network.name == "Unnamed Network 0" else network.name - n = 0 - - # mapping from TRT tensor (trt.ITensor) to GS tensor (gs.Variable) - global_tensor_map = {} - - for i in range(network.num_inputs): - trt_tensor = network.get_input(i) - gs_tensor = gs.Variable(trt_tensor.name, - trt.nptype(trt_tensor.dtype), - trt_tensor.shape) - global_tensor_map[trt_tensor] = gs_tensor - if gs_tensor not in graph.inputs: - graph.inputs.append(gs_tensor) - - for i in range(network.num_layers): - layer = network.get_layer(i) - - input_tensor_list = [] - for j in range(layer.num_inputs): - trt_tensor = layer.get_input(j) - if trt_tensor is None: # Useful for constant layer - gs_tensor = None - elif trt_tensor in global_tensor_map.keys( - ): # already in the map - gs_tensor = global_tensor_map[trt_tensor] - else: - logger.debug( - f"[ExportONNX]Layer input tensor not in global_tensor_map: {trt_tensor.name}" - ) # ■ - gs_tensor = gs.Variable(trt_tensor.name, - trt.nptype(trt_tensor.dtype), - trt_tensor.shape) - global_tensor_map[trt_tensor] = gs_tensor - input_tensor_list.append(gs_tensor) - - output_name_list = [] - output_datatype_list = [] - output_shape_list = [] - for i in range(layer.num_outputs): - trt_tensor = layer.get_output(i) - # Don't do this check because we need this trt_tensor to overwrite the placeholder tensor in ■ - # if trt_tensor in global_tensor_map.keys(): - # gs_tensor = global_tensor_map[trt_tensor] - output_name_list.append(trt_tensor.name) - output_datatype_list.append(trt.nptype(trt_tensor.dtype)) - output_shape_list.append(trt_tensor.shape) - - attr = OrderedDict() - # Set attribution of ILayer - for key in dir(layer): - if not (key.startswith("_") - or callable(layer.__getattribute__(key))): - attr[key] = str(layer.__getattribute__(key)) - # Set attribution of exact layer type - layer.__class__ = layer_type_to_class(layer) - for key in dir(layer): - if key in dir(trt.ILayer) and key != "type": - continue - if key == "type" and not isinstance(layer.type, trt.LayerType): - attr["algo-type"] = str(layer.type) - continue - value = layer.__getattribute__(key) - if isinstance(value, np.ndarray): - # Convert all attributions into string besides weights - value = value.astype(np.float32) # In case of overflow - ss = f"shape={value.shape}, SumAbs={np.sum(abs(value)):.5e}, Var={np.var(value):.5f}, " - ss += f"Max={np.max(value):.5f}, Min={np.min(value):.5f}, SAD={np.sum(np.abs(np.diff(value.reshape(-1)))):.5f}, " - ss += f"[:5]={value.reshape(-1)[:5]}, [-5:]={value.reshape(-1)[-5:]}" - attr[key] = ss - else: - attr[key] = str(value) - - output_tensor_list, n = add_node_for_trt_network(graph, layer.name, attr["type"][10:], input_tensor_list, attr, \ - output_name_list, output_datatype_list, output_shape_list, n, b_onnx_type) - - if layer.num_outputs == 1: - global_tensor_map[layer.get_output(0)] = output_tensor_list - else: - for i in range(layer.num_outputs): - global_tensor_map[layer.get_output( - i)] = output_tensor_list[i] - - for i in range(network.num_outputs): - gs_tensor = global_tensor_map[network.get_output(i)] - if gs_tensor not in graph.outputs: - graph.outputs.append(gs_tensor) - - onnx.save(gs.export_onnx(graph), - path + "/network.onnx", - save_as_external_data=False) - return - - def _get_graph(self) -> "Network._GraphState": - ''' - Get the graph of the network. - - Returns: - Network._GraphState - ''' - return self._get_graph_impl(self._get_network_hash()) - - #TODO: using one LRU cache here can cause the Network object to be leaked, need a way to speed this function w/o using global lru cache. - def _get_graph_impl(self, network_hash: bytes) -> "Network._GraphState": - graph = Network._GraphState() - graph.build(self) - return graph - - @dataclass - class _GraphState: - # Tensor to Layers - tensor_to_consumers: Dict[Any, List["Layer"]] = field( - default_factory=lambda: defaultdict(list)) - # Tensor to Layer - tensor_to_producer: Dict[Any, "Layer"] = field(default_factory=dict) - inputs: Dict[str, Any] = field(default_factory=OrderedDict) - outputs: Dict[str, Any] = field(default_factory=OrderedDict) - name_to_layer: Dict[str, "Layer"] = field(default_factory=dict) - - def build(self, network: "Network") -> None: - from .graph_rewriting import Layer - self.inputs = network.get_inputs() - self.outputs = network.get_outputs() - - for layer in network.get_layers(): - self.name_to_layer[layer.name] = Layer( - network=network, trt_layer=layer.trt_layer) - for i in range(layer.num_inputs): - input_tensor = layer.get_inputs(i)[0] - if input_tensor.is_trt_wrapper(): - self.tensor_to_consumers[input_tensor].append(layer) - for i in range(layer.num_outputs): - output_tensor = layer.get_outputs(i)[0] - if output_tensor.is_trt_wrapper(): - self.tensor_to_producer[output_tensor] = layer - - def _get_network_hash(self, lightweight=True) -> bytes: - # TODO: Ask TensorRT team to add a hash function for INetworkDefinition instead of using this hacky way - num_layers = self.trt_network.num_layers - - # Some special layers, such as slice, may be associated with tensors that do not have the `trt_tensor` member. - get_tensor_tag = lambda tensor: tensor.trt_tensor.name if tensor.is_trt_wrapper( - ) else 'None' - - if lightweight and not self.is_graph_altered: - return num_layers - self.is_graph_altered = False - - data = hashlib.sha256() - # network layer count - data.update(str(num_layers).encode()) - # network inputs - data.update(','.join( - [get_tensor_tag(tensor) for tensor in self.get_inputs()]).encode()) - # network outputs - data.update(','.join( - [get_tensor_tag(tensor) for tensor in self.get_outputs()]).encode()) - # layer names - data.update(','.join( - [layer.trt_layer.name for layer in self.get_layers()]).encode()) - - # layer -> output - data.update(','.join([ - f'{layer.trt_layer.name}->{get_tensor_tag(tensor)}' - for layer in self.get_layers() for tensor in layer.get_outputs() - ]).encode()) - - # input -> layer - data.update(','.join([ - f'{get_tensor_tag(tensor)}->{layer.trt_layer.name}' - for layer in self.get_layers() for tensor in layer.get_inputs() - ]).encode()) - - return data.hexdigest() - - -@contextlib.contextmanager -def net_guard(network): - from ._common import net - assert isinstance( - network, Network - ), f"Invalid network, can only guard Network instance, got: {network}" - - old_net = net - set_network(network) - yield - set_network(old_net) - - -class _TrtLlmModuleCallStack(object): - - def __init__(self): - super().__init__() - self.call_stack = [] - self.module_name_map = weakref.WeakKeyDictionary() - self.module_to_layer_range_map: Dict[str, range] = {} - self.mod_names_set = False - - def module_names_set(self): - return self.mod_names_set - - def set_module_names(self, top_level_module): - assert top_level_module, "Expected a top level module" - for name, mod in top_level_module.named_modules( - prefix=top_level_module._get_name()): - if mod not in self.module_name_map: - self.module_name_map[mod] = name - self.mod_names_set = True - return - - def get_current_module(self): - mod_name = '' - if len(self.call_stack): - mod_name = self.call_stack[-1] - return mod_name - - def get_mod_name(self, mod_obj): - name = '' - if mod_obj in self.module_name_map: - name = self.module_name_map[mod_obj] - return name - - def set_layer_range(self, mod_obj: Module, layer_range: range): - if mod_obj in self.module_name_map: - name = self.module_name_map[mod_obj] - self.module_to_layer_range_map[name] = layer_range - - def get_stack(self): - return self.call_stack - - @contextlib.contextmanager - def call_stack_mgr(self): - call_stack = self.get_stack() - try: - yield call_stack - finally: - call_stack.pop() diff --git a/tensorrt_llm/parameter.py b/tensorrt_llm/parameter.py deleted file mode 100644 index d740cbb0bdb9..000000000000 --- a/tensorrt_llm/parameter.py +++ /dev/null @@ -1,281 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -import math -import weakref -from typing import Optional, Sequence, Union - -import numpy as np - -# isort: off -import torch -import tensorrt as trt -# isort: on - -from ._common import default_net -from ._utils import (copy_torch_to_numpy, np_dtype_to_trt, str_dtype_to_trt, - torch_to_numpy, trt_dtype_to_np, trt_dtype_to_torch) -from .functional import Tensor, constant -from .logger import logger -from .network import Network - - -class Parameter: - _DEFAULT_DTYPE = trt.DataType.FLOAT - - def __init__(self, - value: Optional[Union[np.ndarray, torch.Tensor]] = None, - shape: Sequence[int] = None, - dtype: Union[str, trt.DataType] = None, - is_buffer: bool = False, - prefer_managed=False): - if dtype is None: - logger.warning( - f'Parameter dtype is None, using default dtype: {self._DEFAULT_DTYPE}, it is recommended to always specify dtype explicitly' - ) - dtype = self._DEFAULT_DTYPE if dtype is None else dtype - if isinstance(dtype, str): - dtype = str_dtype_to_trt(dtype) - self._dtype: trt.DataType = dtype - if value is None: - assert isinstance(shape, ( - list, - tuple)), f"shape must be list or tuple, receive {(type(shape))}" - self._shape = tuple(shape) - self._value = None - else: - self._shape = value.shape - self._value = self._regularize_value(value) - self.is_buffer = is_buffer - self._prefer_managed = prefer_managed - self._tensor: Tensor = None - self._network: weakref.ref = None - self._name = None - self.need_transpose = False - - @property - def shape(self): - return self._shape - - @property - def dtype(self): - return self._dtype - - @property - def name(self): - return self._name - - def _create_managed_tensor(self, network) -> Tensor: - num = len(network._inputs) - self._name = f"managed_constant_{num}" - - if self._value is None or (isinstance(self._value, np.ndarray) - and not self._value.flags['C_CONTIGUOUS']): - value_old = self._value - self._value = np.empty(self._shape, trt_dtype_to_np(self._dtype)) - network._register_unfilled_weights( - # use updated self._shape here - self._name, - self._value, - value_old) - return Tensor(name=self._name, dtype=self._dtype, shape=self._shape) - - def get_managed_tensor(self, network: Network) -> Tensor: - if self._network is None or self._network() != network: - self._network = weakref.ref(network) - self._tensor = network.get_parameter_tensor(self) - if self._tensor is None: - self._tensor = self._create_managed_tensor(network) - network.set_parameter_tensor(self, self._tensor) - return self._tensor - - def _create_constant_tensor(self) -> Tensor: - if (self._value is not None and isinstance(self._value, np.ndarray) - and self._value.flags['C_CONTIGUOUS']): - lower_type = None - lower_shape = None - # workaround for reinterpreted data type - dtype = self._value.dtype - if (self.dtype == trt.fp4 or self.dtype - == trt.fp8) and (dtype == np.uint8 or dtype == np.int8 - or dtype == np.int32 or dtype == np.int64): - lower_type = self.dtype - lower_shape = self.shape - - self._value = constant(self._value, lower_type, lower_shape) - return self._value - elif self._value is None or isinstance(self._value, np.ndarray): - if self._dtype == trt.fp4: - shape = list(self._shape) - assert shape[ - -1] % 16 == 0, "For FP4, the last dimension of the shape should be multiple of 16" - shape[-1] = shape[-1] // 16 - dtype = np.int64 - else: - shape = self._shape - dtype = trt_dtype_to_np(self._dtype) - ndarray = np.empty(shape, dtype) - tensor = constant(ndarray, self._dtype, self._shape) - default_net()._register_unfilled_weights(tensor.producer.name, - ndarray, self._value) - return tensor - - def get_constant_tensor(self, network: Network) -> Tensor: - if self._network is None or self._network() != network: - self._network = weakref.ref(network) - self._tensor = network.get_parameter_tensor(self) - if self._tensor is None: - self._tensor = self._create_constant_tensor() - self._name = self._tensor.producer.name - network.set_parameter_tensor(self, self._tensor) - return self._tensor - - def get_tensor(self, network) -> Tensor: - if self.is_managed(network): - return self.get_managed_tensor(network) - else: - return self.get_constant_tensor(network) - - def is_managed(self, network): - if network is None: - network = default_net() - return self._prefer_managed and network.plugin_config.manage_weights - - @property - def value(self) -> Tensor: - return self.get_tensor(default_net()) - - @classmethod - def xavier_init(cls, weights: np.ndarray): - shape = weights.shape - dtype = np_dtype_to_trt(weights.dtype) - if len(shape) == 2: - # Xavier initialization see https://paperswithcode.com/method/xavier-initialization - v_range = math.sqrt(6) / math.sqrt(shape[0] + shape[1]) - else: - v_range = 0.1 - - if dtype == trt.DataType.INT8 or dtype == trt.DataType.INT32 or dtype == trt.DataType.INT64: - range_map = { - trt.DataType.INT8: 128, - trt.DataType.INT32: 2**31, - trt.DataType.INT64: 2**63 - } - upper = math.ceil(range_map[dtype] * v_range) - value = torch.randint(-upper, - upper, (shape), - dtype=trt_dtype_to_torch(dtype), - device='cuda') - # value ~ U[int(-128 * v_range), int(128 * v_range)] - elif dtype == trt.DataType.FP8: - value = torch.rand((shape), device='cuda') * 2 - 1 - # value ~ U[-v_range, v_range] - value = value * v_range - value = value.to(trt_dtype_to_torch(dtype)) - else: - value = torch.rand( - (shape), dtype=trt_dtype_to_torch(dtype), device='cuda') * 2 - 1 - # value ~ U[-v_range, v_range] - value = value * v_range - - copy_torch_to_numpy(value, weights) - - def is_inited(self) -> bool: - return self._value is not None - - @property - def raw_value(self) -> np.ndarray: - if self._value is None: - dtype = trt_dtype_to_np(self.dtype) - self._value = np.empty(self.shape, dtype) - Parameter.xavier_init(self._value) - assert isinstance( - self._value, np.ndarray - ), "Must be np.ndarray. Proper usage: get parameter.raw_value before getting parameter.value" - return self._value - - @value.setter - def value(self, v: Union[np.ndarray, torch.Tensor]): - v = self._regularize_value(v) - - if v.shape != self.shape and v.ndim == 0 and max(self.shape) == 1: - # convert the scalar into a tensor which each dim is 1. - v = v.reshape(self.shape) - - if self.dtype == trt.fp4: - assert v.shape[:-1] == self.shape[:-1] and v.shape[-1] == self.shape[-1] // 2 // v.dtype.itemsize, \ - f'For FP4, the shape of the value should be the same as the original shape, ' \ - f'except the last dimension should be half of the original shape. ' \ - f'Updated: {v.shape}, original: {self.shape}' - else: - assert v.shape == self.shape, \ - f'The value updated is not the same shape as the original. ' \ - f'Updated: {v.shape}, original: {self.shape}' - if (self.dtype == trt.fp4 or self.dtype - == trt.fp8) and (v.dtype == np.int8 or v.dtype == np.uint8 - or v.dtype == np.int32 or v.dtype == np.int64): - pass - else: - dtype = np_dtype_to_trt(v.dtype) - if self.dtype != dtype: - logger.warning( - f"Parameter was initialized as {self.dtype} but set to {dtype}" - ) - self._dtype = dtype - self._value = v - - def set_value_or_dummy(self, v: Union[np.ndarray, torch.Tensor]): - v = self._regularize_value(v) - if v.shape != self._shape: - self.value = np.empty(self._shape, trt_dtype_to_np(self._dtype)) - return - - self.value = v - - def set_name(self, name: str, network: Network): - self._name = name - if self.is_managed(network): - self._get_weights(network).name = name - return True - else: - weights = self._get_weights(network) - # TensorRT bindings may return numpy array instead of trt.Weights - if isinstance(weights, np.ndarray): - trt_dtype = np_dtype_to_trt( - weights.dtype - ) if weights.dtype != np.object_ else self._dtype - trt_count = int(np.prod(weights.shape)) - weights = trt.Weights(trt_dtype, weights.ctypes.data, trt_count) - return network.trt_network.set_weights_name(weights, name) - - def _get_weights(self, network: Network) -> trt.Weights | Tensor | None: - tensor = network.get_parameter_tensor(self) - if self.is_managed(network): - return tensor - elif tensor is not None: - tensor.producer.__class__ = trt.IConstantLayer - return tensor.producer.weights - else: - return None - - def _regularize_value(self, value): - if isinstance(value, np.ndarray): - return value - - elif isinstance(value, torch.distributed.tensor.DTensor): - return value.to_local().cpu().numpy() - elif isinstance(value, torch.Tensor): - return torch_to_numpy(value) - raise TypeError( - f'Expected numpy.ndarray or torch.Tensor, got {type(value)}') diff --git a/tensorrt_llm/plugin/__init__.py b/tensorrt_llm/plugin/__init__.py deleted file mode 100644 index 55bffdee6b85..000000000000 --- a/tensorrt_llm/plugin/__init__.py +++ /dev/null @@ -1,23 +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. -from .plugin import (TRT_LLM_PLUGIN_NAMESPACE, PluginConfig, _load_plugin_lib, - add_plugin_argument, current_all_reduce_helper, - init_all_reduce_helper, plugin_lib_path) - -__all__ = [ - 'TRT_LLM_PLUGIN_NAMESPACE', '_load_plugin_lib', 'PluginConfig', - 'add_plugin_argument', 'plugin_lib_path', "current_all_reduce_helper", - "init_all_reduce_helper" -] diff --git a/tensorrt_llm/plugin/plugin.py b/tensorrt_llm/plugin/plugin.py deleted file mode 100644 index f8434cf11bec..000000000000 --- a/tensorrt_llm/plugin/plugin.py +++ /dev/null @@ -1,780 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -import argparse -import ctypes -import os -import platform -from collections import OrderedDict -from enum import IntEnum -from pathlib import Path -from textwrap import dedent -from typing import (Any, List, Literal, Optional, Tuple, Union, get_args, - get_origin) - -import tensorrt as trt -from pydantic import (ConfigDict, Field, PrivateAttr, ValidationInfo, - field_validator, model_validator) - -from .._ipc_utils import IpcMemory, can_access_peer -from .._utils import get_sm_version -from ..bindings.internal.runtime import (lamport_initialize, - lamport_initialize_all, - max_workspace_size_lowprecision) -from ..llmapi.utils import StrictBaseModel -from ..logger import logger -from ..mapping import Mapping - -TRT_LLM_PLUGIN_NAMESPACE = 'tensorrt_llm' - - -def plugin_lib_path() -> str: - project_dir = Path(__file__).parent.parent.absolute() - dyn_lib = "libnvinfer_plugin_tensorrt_llm.so" if platform.system( - ) != "Windows" else "nvinfer_plugin_tensorrt_llm.dll" - return str(project_dir.joinpath("libs", dyn_lib)) - - -def _load_plugin_lib(): - on_windows = platform.system() == "Windows" - winmode = 0 if on_windows else None - handle = ctypes.CDLL(plugin_lib_path(), - mode=ctypes.RTLD_GLOBAL, - winmode=winmode) - try: - handle.initTrtLlmPlugins.argtypes = [ctypes.c_void_p, ctypes.c_char_p] - handle.initTrtLlmPlugins.restype = ctypes.c_bool - except AttributeError as err: - raise ImportError('TensorRT LLM Plugin is unavailable') from err - - try: - assert handle.initTrtLlmPlugins( - None, TRT_LLM_PLUGIN_NAMESPACE.encode('utf-8')) - except OSError as e: - windows_err = """ - The error above may be caused by an outdated Microsoft Visual C++ Redistributable Version. - Please install the latest MSVC from the link below and re-launch. - - https://learn.microsoft.com/en-us/cpp/windows/latest-supported-vc-redist?view=msvc-170#latest-microsoft-visual-c-redistributable-version - """ - err_msg = dedent(windows_err if on_windows else "Unknown error") - raise RuntimeError(err_msg) from e - except Exception as e: - raise e - - -class ContextFMHAType(IntEnum): - disabled = 0 - # FP16 I/O, FP16 Accumulation - enabled = 1 - # FP16 I/O, FP32 Accumulation - enabled_with_fp32_acc = 2 - - -DefaultPluginDtype = Literal["auto", "float16", "float32", "bfloat16", "int32", - None] - - -class PluginConfig(StrictBaseModel): - """The config that manages plugin-related options. - - There are two option categories: - * Plugin options (typically with xxx_plugin naming). These options can be assigned with: - * "float16"/"bfloat16"/"float32"/"int32", which means the plugin is enabled with the specified precision; (Some plugins only support limited dtype, i.e., gemm_swiglu_plugin and low_latency_gemm_swiglu_plugin only supports fp8 now) - * "auto", which means the plugin is enabled with the precision of `dtype` field (the `dtype` field must be same to model dtype, i.e., the one in PretrainedConfig); - * None, which means the plugin is disabled. - * Other features. These options can be assigned with boolean: - * True, which means the plugin is enabled; - * False, which means the plugin is disabled. - """ - model_config = ConfigDict(validate_assignment=True, extra="forbid") - - dtype: str = Field(default="float16", - description="Base dtype for the model and plugins") - - # Plugins - bert_attention_plugin: Optional[DefaultPluginDtype] = Field( - default="auto", - description= - "The plugin that uses efficient kernels and enables an in-place update of the KV cache for attention layer of BERT-like encoder models." - ) - gpt_attention_plugin: Optional[DefaultPluginDtype] = Field( - default="auto", - description= - "The plugin that uses efficient kernels and enables an in-place update of the KV cache for attention layer of GPT-like decoder models." - ) - gemm_plugin: Optional[Literal[ - "auto", "float16", "float32", "bfloat16", "int32", "fp8", "nvfp4", - None]] = Field( - default=None, - description= - "The GEMM plugin that utilizes NVIDIA cuBLASLt to perform GEMM operations. " - "Note: it's only affective for non-quantized gemm operations (except FP8)." - "Note: For FP8, it also requires same calibration in checkpoint.") - _explicitly_disable_gemm_plugin: bool = PrivateAttr(default=False) - gemm_swiglu_plugin: Optional[Literal["fp8", None]] = Field( - default=None, - description= - "The GEMM + SwiGLU fusion in Gated-MLP combines two Matmul operations and " - "one SwiGLU operation into a single kernel. Currently this is only supported for FP8 precision on Hopper." - ) - fp8_rowwise_gemm_plugin: Optional[DefaultPluginDtype] = Field( - default=None, - description= - "The quantized GEMM for fp8, which uses per token dynamic scales for " - "activation and per channel static scales for weights." - "Note: It also requires same calibration in checkpoint.") - qserve_gemm_plugin: Optional[DefaultPluginDtype] = Field( - default=None, - description= - "The quantized GEMM from [QServe](https://arxiv.org/abs/2405.04532), " - "which employs 4-bit quantization for weights and 8-bit quantization for activations." - ) - identity_plugin: Optional[DefaultPluginDtype] = Field( - default=None, - description= - "The identity plugin simply copies inputs to outputs, it's used mostly for debugging purpose." - ) - nccl_plugin: Optional[DefaultPluginDtype] = Field( - default="auto", - description= - "The NCCL plugin wraps NCCL operators to support multi-GPU and even multi-nodes." - ) - lora_plugin: Optional[DefaultPluginDtype] = Field( - default=None, description="Enable LoRA.") - dora_plugin: bool = Field(default=False, description="Enable DoRA.") - weight_only_groupwise_quant_matmul_plugin: Optional[ - DefaultPluginDtype] = Field( - default=None, - description= - "Enable weight-only groupwise quantization matmul operators.") - weight_only_quant_matmul_plugin: Optional[DefaultPluginDtype] = Field( - default=None, - description="Enable weight-only quantization matmul operators.") - smooth_quant_plugins: bool = Field( - default=True, - description="Enable a group of plugins to support smooth quantization.") - smooth_quant_gemm_plugin: Optional[DefaultPluginDtype] = Field( - default=None, - description= - "Enable plugin that supports smooth quantization gemm kernels.") - layernorm_quantization_plugin: Optional[DefaultPluginDtype] = Field( - default=None, - description="Enable plugin that supports layernorm quantization kernels." - ) - rmsnorm_quantization_plugin: Optional[DefaultPluginDtype] = Field( - default=None, - description="Enable plugin that supports rmsnorm quantization kernels.") - quantize_per_token_plugin: bool = Field( - default=False, - description="Enable plugin that supports per-token quantization.") - quantize_tensor_plugin: bool = Field( - default=False, - description="Enable plugin that supports per-tensor quantization.") - moe_plugin: Optional[DefaultPluginDtype] = Field( - default="auto", - description= - "Enable some customized kernels to speed up the MoE layer of MoE models." - ) - mamba_conv1d_plugin: Optional[DefaultPluginDtype] = Field( - default="auto", - description= - "Enable customized kernels to speed up conv1d operator for Mamba.") - low_latency_gemm_plugin: Optional[Literal["fp8", None]] = Field( - default=None, - description= - "The GEMM plugin that optimized specially for low latency scenarios.") - low_latency_gemm_swiglu_plugin: Optional[Literal["fp8", None]] = Field( - default=None, - description= - "The GEMM + SwiGLU fusion plugin that optimized specially for low latency scenarios." - ) - gemm_allreduce_plugin: Optional[Literal[ - "float16", "bfloat16", - None]] = Field(default=None, - description="The GEMM + AllReduce kernel fusion plugin.") - - # Features - context_fmha: bool = Field( - default=True, - description= - "Enable the fused multi-head attention during the context phase, " - "will trigger a kernel that performs the MHA/MQA/GQA block using a single kernel." - ) - bert_context_fmha_fp32_acc: bool = Field( - default=False, - description= - "Enable the FP32 accumulator for context FMHA in the bert_attention_plugin. " - "If disabled, FP16 is used, better performance but potentially worse accuracy is expected." - ) - paged_kv_cache: Optional[bool] = Field( - default=None, - description= - "Enable paged KV cache, which helps manage memory for the KV cache more efficiently, " - "and usually leads to an increase in the batch size and an improved efficiency." - ) - remove_input_padding: bool = Field( - default=True, - description= - "Pack different tokens together, which reduces both the amount of computations and memory consumption." - ) - norm_quant_fusion: bool = Field( - default=False, - description= - "Fuse the LayerNorm and quantization kernels into a single kernel, " - "resulting in improved end-to-end performance.") - reduce_fusion: bool = Field( - default=False, - description= - "Fuse the ResidualAdd and LayerNorm kernels after AllReduce into a single kernel, " - "resulting in improved end-to-end performance.") - user_buffer: bool = Field( - default=False, - description= - "Eliminate extra copies from the local buffer to the shared buffer " - "in the communication kernel, leading to improved end-to-end performance. " - "This feature must be enabled with `--reduce_fusion enable` and " - "is currently only supported for the FP8 LLAMA model.") - tokens_per_block: int = Field( - default=32, - description= - "Define how many tokens are contained in each paged kv cache block.") - use_paged_context_fmha: bool = Field( - default=True, - description= - "Allow advanced features like KV cache reuse and chunked context.") - use_fp8_context_fmha: bool = Field( - default=True, - description= - "When FP8 quantization is activated, the attention can be further accelerated by enabling FP8 Context FMHA" - ) - fuse_fp4_quant: bool = Field( - default=False, - description="Whether to fuse FP4 quantization into attention kernel.") - multiple_profiles: bool = Field( - default=False, - description= - "Enables multiple TensorRT optimization profiles in the built engines, " - "will benefits the performance especially when GEMM plugin is disabled, " - "because more optimization profiles help TensorRT have more chances to select better kernels. " - "Note: This feature increases engine build time but no other adverse effects are expected." - ) - paged_state: bool = Field( - default=True, - description= - "Enable paged state, which helps manage memory for the RNN state more efficiently." - ) - streamingllm: bool = Field( - default=False, - description= - "Enable [StreamingLLM](https://arxiv.org/abs/2309.17453), which uses a window attention to perform efficient and stable LLM on long texts." - ) - manage_weights: bool = Field( - default=False, - description= - "Enable TensorRT LLM managed weights to speed up engine building process." - ) - use_fused_mlp: bool = Field( - default=True, - description= - "Enable horizontal fusion in Gated-MLP that combines two Matmul " - "operations into a single one followed by a separate SwiGLU kernel.") - pp_reduce_scatter: bool = Field( - default=False, - description="Enable a pipeline parallelism optimization with " - "ReduceScatter + AllGather targeting large MoE models.") - - def __getattribute__(self, name: str) -> Any: - """Override to resolve 'auto' values to dtype field. - - When a plugin field has value 'auto', return the value of dtype instead. - """ - # Use object.__getattribute__ to avoid infinite recursion - value = object.__getattribute__(self, name) - - if name != "dtype" and value == "auto": - return self.dtype - - return value - - @field_validator("dtype") - @classmethod - def validate_dtype_not_auto(cls, v: str) -> str: - if v == "auto": - raise ValueError("Plugin dtype cannot be 'auto'") - return v - - @field_validator("*", mode="before") - @classmethod - def convert_enable_disable(cls, value, info: ValidationInfo): - """Allow passing enable/disable strings which map to boolean/None values.""" - if value == "enable": - return True - elif value == "disable": - annotation = cls.model_fields[info.field_name].annotation - if annotation is bool or (get_origin(annotation) is Union - and bool in get_args(annotation)): - return False - return None - return value - - @field_validator("*", mode="after") - @classmethod - def log_field_changes(cls, v: Any, info: ValidationInfo) -> Any: - """Log all field changes for debugging.""" - logger.info(f"Set {cls.__name__}.{info.field_name} to {v}.") - return v - - @classmethod - def from_arguments(cls, args: argparse.Namespace): - """Create a PluginConfig from argparse arguments.""" - args = vars(args) - # Filter to only include fields that are part of PluginConfig - valid_fields = set(cls.model_fields.keys()) - filtered_args = {k: v for k, v in args.items() if k in valid_fields} - obj = cls(**filtered_args) - - # We want to know if the user explicitly disabled the gemm_plugin - # because nvfp4 gemm uses plugin by default currently - if 'gemm_plugin' in args and args['gemm_plugin'] == 'disable': - obj._explicitly_disable_gemm_plugin = True - - return obj - - def to_legacy_setting(self): - """Legacy setting means that all of the plugins and features are - disabled, this is needed for the legacy `build.py` script, which will be - migrated to the centralized building script `tensorrt_llm/commands/build.py`. - - After the migration is done, this function may or may not be deleted. - """ - for field_name, field_value in self: - if field_name == "dtype": - continue - elif isinstance(field_value, str): - setattr(self, field_name, None) - elif isinstance(field_value, - bool) or field_name == "paged_kv_cache": - setattr(self, field_name, False) - - @model_validator(mode="after") - def _validate_sm_compatibility(self) -> "PluginConfig": - unsupported_plugins = { - # bert_attention_plugin is handled within BertAttention - 100: [ - 'gemm_swiglu_plugin', 'fp8_rowwise_gemm_plugin', - 'low_latency_gemm_plugin', 'low_latency_gemm_swiglu_plugin', - 'bert_context_fmha_fp32_acc' - ] - } - sm = get_sm_version() - if sm in unsupported_plugins: - for plugin in unsupported_plugins[sm]: - val = getattr(self, plugin, None) - if val is not None and val != False: - raise ValueError( - f"{plugin}={val} is not supported on SM {sm}.") - return self - - @property - def context_fmha_type(self): - if self.bert_context_fmha_fp32_acc: - return ContextFMHAType.enabled_with_fp32_acc - elif self.context_fmha: - return ContextFMHAType.enabled - else: - return ContextFMHAType.disabled - - def is_context_fmha_enabled(self): - return self.context_fmha_type != ContextFMHAType.disabled - - @context_fmha_type.setter - def context_fmha_type(self, value): - if value == ContextFMHAType.disabled: - self.context_fmha = False - self.bert_context_fmha_fp32_acc = False - else: - self.context_fmha = True - if value == ContextFMHAType.enabled: - self.bert_context_fmha_fp32_acc = False - elif value == ContextFMHAType.enabled_with_fp32_acc: - self.bert_context_fmha_fp32_acc = True - - def set_smooth_quant_plugins(self, dtype: str = "auto"): - self.smooth_quant_gemm_plugin = dtype - self.rmsnorm_quantization_plugin = dtype - self.layernorm_quantization_plugin = dtype - self.quantize_per_token_plugin = True - self.quantize_tensor_plugin = True - return self - - def set_qserve_plugins(self, dtype: str = "auto"): - self.qserve_gemm_plugin = dtype - self.rmsnorm_quantization_plugin = dtype - self.quantize_per_token_plugin = True - return self - - def set_fp8_rowwise_quant_plugins(self, dtype: str = "auto"): - self.fp8_rowwise_gemm_plugin = dtype - self.rmsnorm_quantization_plugin = dtype - self.layernorm_quantization_plugin = dtype - self.quantize_per_token_plugin = True - self.quantize_tensor_plugin = True - return self - - def set_context_fmha(self, context_fmha_type=ContextFMHAType.enabled): - assert type(context_fmha_type) == ContextFMHAType - self.context_fmha_type = context_fmha_type - return self - - def enable_paged_kv_cache(self, tokens_per_block: int = 32): - self.paged_kv_cache = True - self.tokens_per_block = tokens_per_block - return self - - def set_nccl_plugin(self, dtype: str = "auto"): - self.nccl_plugin = dtype - init_all_reduce_helper() - return self - - def set_lora_plugin(self, dtype: str = None): - self.lora_plugin = dtype - return self - - def set_dora_plugin(self, enable: bool = False): - self.dora_plugin = enable - return self - - -# Only plugin configs in this list will be exposed as `trtllm-build` arguments, -# others are automatically enabled when needed, no need for users to control. -cli_plugin_args = [ - # Plugins - "bert_attention_plugin", - "gpt_attention_plugin", - "gemm_plugin", - "gemm_swiglu_plugin", - "fp8_rowwise_gemm_plugin", - "lora_plugin", - "dora_plugin", - "moe_plugin", - "mamba_conv1d_plugin", - "nccl_plugin", - "low_latency_gemm_plugin", - "low_latency_gemm_swiglu_plugin", - "gemm_allreduce_plugin", - - # Features - "context_fmha", - "bert_context_fmha_fp32_acc", - "remove_input_padding", - "tokens_per_block", - "use_paged_context_fmha", - "use_fp8_context_fmha", - "fuse_fp4_quant", - "multiple_profiles", - "paged_state", - "streamingllm", - "norm_quant_fusion", - "reduce_fusion", - "user_buffer", - "use_fused_mlp", - "pp_reduce_scatter", -] - - -def add_plugin_argument(parser: argparse.ArgumentParser): - for field_name, field_info in PluginConfig.model_fields.items(): - if field_name not in cli_plugin_args: - continue - help_message = field_info.description - if not help_message: - raise AttributeError(f"Please add help message for {field_name}.") - annotation = field_info.annotation - - # Extract choices from the Optional[Literal[...]] type - plugin_dtype_options = None - if get_origin(annotation) is Union: - args = get_args(annotation) - for arg in args: - if get_origin(arg) is Literal: - plugin_dtype_options = list(get_args(arg)) - if type(None) in args: - plugin_dtype_options.append(None) - break - - if plugin_dtype_options is not None: - if field_name == "gemm_plugin": - default = field_info.default - else: - default = field_info.default if field_info.default else "disable" - parser.add_argument( - "--" + field_name, - type=str, - default=default, - choices=[x if x else "disable" for x in plugin_dtype_options], - help=help_message) - elif annotation is bool: - parser.add_argument( - "--" + field_name, - type=str, - default="enable" if field_info.default else "disable", - choices=["enable", "disable"], - help=help_message) - else: - parser.add_argument("--" + field_name, - type=annotation, - default=field_info.default, - help=help_message) - return parser - - -def force_all_reduce_deterministic(): - return os.getenv("FORCE_DETERMINISTIC", "0") == "1" or os.getenv( - "FORCE_ALL_REDUCE_DETERMINISTIC", "0") == "1" - - -class CustomAllReduceHelper: - """ - Globally visible class to help usage of custom_all_reduce plugin. - Provides the following utilities: - - workspace: Tensor - When using CUSTOM or AUTO mode, a tensor containing pointers to memory - visible to all GPUs. It should be 3 pointers per TP rank - - ptr to data buffer, ptr to barriers in, ptr to barriers out. - It must be initialized using IpcMemory class. - - Usage: - - Set custom_all_reduce_helper.workspace with the required tensor. - Then, each instance of allreduce will reference that tensor automatically. - """ - POINTERS_PER_RANK = 7 - POINTERS_OF_COUNTER = 3 - - def __init__(self) -> None: - self.workspace: Optional[Tensor] = None - - def set_workspace_tensor(self, - mapping: Mapping, - num_profiles: Optional[int] = None): - from ..functional import Tensor - workspace_size = self.POINTERS_PER_RANK * mapping.tp_size + self.POINTERS_OF_COUNTER - - dim_range = None - if num_profiles is not None: - dim_range = OrderedDict([('all_reduce_size', - [workspace_size] * num_profiles)]) - - self.workspace = Tensor( - name='all_reduce_workspace', - dtype=trt.int64, - shape=[workspace_size], - dim_range=dim_range, - ) - - @staticmethod - def max_workspace_size_auto(tp_size: int, - support_deterministic=True) -> int: - """Calculate workspace size for allreduce fusion kernel. - - The workspace is used for lamport buffers in the fusion kernel. - Required size calculation: - - Each GPU needs 3 sub-buffers (for triple buffering) - - Each sub-buffer stores: max_num_tokens * hidden_size * dtype_size (bf16=2) - - The lamport allocation multiplies by tp_size, so: - lamport_size = 3 * size * tp_size (per GPU) - - Example: Llama 8B (hidden=4096), max_tokens=8192, bf16, TP=4 - - Data per sub-buffer: 8192 * 4096 * 2 = 64 MiB - - Total lamport: 3 * 64MB * 4 = 768 MiB per GPU - - Required 'size' parameter: 64 MiB (gets multiplied by tp_size in allocation) - - Default (67,108,864 = 64 MiB) supports: - - Models up to hidden_size=4096 with max_num_tokens=8192 - - Or hidden_size=8192 with max_num_tokens=4096 - - Override with TRTLLM_ALLREDUCE_FUSION_WORKSPACE_SIZE env var if needed for larger models. - """ - if force_all_reduce_deterministic() and support_deterministic: - workspace_size = os.getenv("FORCE_ALLREDUCE_KERNEL_WORKSPACE_SIZE", - "1000000000") - return int(workspace_size) - - # Allow override via environment variable for edge cases - workspace_size_env = os.getenv("TRTLLM_ALLREDUCE_FUSION_WORKSPACE_SIZE") - if workspace_size_env: - size = int(workspace_size_env) - logger.info( - f"Using custom allreduce fusion workspace size: {size} bytes ({size / (1024**2):.1f} MiB)" - ) - return size - - # Default: 64 MiB - supports most common model configurations - # Increase via env var if you see CUDA illegal memory access errors with large models - default_size = 67_108_864 # Exactly 64 MiB - return default_size - - @staticmethod - def max_workspace_size_lowprecision(tp_size: int) -> int: - return max_workspace_size_lowprecision(tp_size) - - @staticmethod - def initialize_lowprecision_buffers(workspace: "torch.tensor", - tp_size: int) -> None: - import torch - return torch.ops.trtllm.initialize_static_lowprecision_buffers( - workspace, tp_size) - - @staticmethod - def allocate_workspace(mapping: Mapping, - size: int) -> Tuple[List[IpcMemory], "torch.tensor"]: - import torch - - # Force pull mode and disable lamport when force deterministic is enabled, for reducing device memory usage. - force_deterministic = force_all_reduce_deterministic() - is_p2p_supported = can_access_peer(mapping) - ipc_buffers_size = size if force_deterministic else size * mapping.tp_size - ipc_buffers_ping = IpcMemory(mapping, ipc_buffers_size, - is_p2p_supported) - ipc_buffers_pong = IpcMemory(mapping, ipc_buffers_size, - is_p2p_supported) - ipc_barriers_in = IpcMemory( - mapping, IpcMemory.IPC_BARRIERS_SIZE_PER_GPU * mapping.tp_size * 2 * - mapping.tp_size, is_p2p_supported) - ipc_barriers_out = IpcMemory( - mapping, IpcMemory.IPC_BARRIERS_SIZE_PER_GPU * mapping.tp_size * 2 * - mapping.tp_size, is_p2p_supported) - lamport_buffers_size = 1 if force_deterministic else size * mapping.tp_size - lamport_buffers_0 = IpcMemory(mapping, lamport_buffers_size, - is_p2p_supported) - lamport_buffers_1 = IpcMemory(mapping, lamport_buffers_size, - is_p2p_supported) - lamport_buffers_2 = IpcMemory(mapping, lamport_buffers_size, - is_p2p_supported) - # TODO: it seems we may need to initialize lamport buffers for all tp groups - # just like its cpp counterpart (AllReduceBuffers::AllReduceBuffers()) does. - if is_p2p_supported: - lamport_initialize_all( - lamport_buffers_0.local_ptr, - lamport_buffers_1.local_ptr, - lamport_buffers_2.local_ptr, - lamport_buffers_size, - ) - buffers = [ - ipc_buffers_ping, - ipc_buffers_pong, - ipc_barriers_in, - ipc_barriers_out, - lamport_buffers_0, - lamport_buffers_1, - lamport_buffers_2, - # 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. - torch.tensor([1, 1, 0], dtype=torch.int64, device="cuda") - ] - - return buffers, torch.tensor( - ipc_buffers_ping.serialize() + ipc_buffers_pong.serialize() + - ipc_barriers_in.serialize() + ipc_barriers_out.serialize() + - lamport_buffers_0.serialize() + lamport_buffers_1.serialize() + - lamport_buffers_2.serialize() + [buffers[-1].data_ptr()] + - [buffers[-1][1:].data_ptr()] + [buffers[-1][2:].data_ptr()], - dtype=torch.int64, - device="cpu") - - @staticmethod - def allocate_lowprecision_workspace( - mapping: Mapping, - size: int) -> Tuple[List[IpcMemory], "torch.tensor"]: - import torch - - # Force pull mode and disable lamport when force deterministic is enabled, for reducing device memory usage. - is_p2p_supported = can_access_peer(mapping) - ipc_buffers_size = size - ipc_buffers_ping = IpcMemory(mapping, ipc_buffers_size, - is_p2p_supported) - ipc_buffers_pong = IpcMemory(mapping, ipc_buffers_size, - is_p2p_supported) - ipc_barriers_in = IpcMemory( - mapping, IpcMemory.IPC_BARRIERS_SIZE_PER_GPU * mapping.tp_size * 2, - is_p2p_supported) - ipc_barriers_out = IpcMemory( - mapping, IpcMemory.IPC_BARRIERS_SIZE_PER_GPU * mapping.tp_size * 2, - is_p2p_supported) - buffers = [ - ipc_buffers_ping, ipc_buffers_pong, ipc_barriers_in, - ipc_barriers_out - ] - - return buffers, torch.tensor( - ipc_buffers_ping.serialize() + ipc_buffers_pong.serialize() + - ipc_barriers_in.serialize() + ipc_barriers_out.serialize() + [0] + - [0], - dtype=torch.int64, - device="cpu") - - @staticmethod - def allocate_allreduce_fusion_workspace( - mapping: Mapping, - size: int) -> Tuple[List[IpcMemory], "torch.tensor"]: - import torch - is_p2p_supported = can_access_peer(mapping) - ipc_buffers_size = size * mapping.tp_size - ipc_buffers = IpcMemory(mapping, ipc_buffers_size, is_p2p_supported) - ipc_barriers = IpcMemory(mapping, 256 * mapping.tp_size, - is_p2p_supported) - lamport_buffers_size = size * mapping.tp_size - lamport_buffers = IpcMemory(mapping, 3 * lamport_buffers_size, - is_p2p_supported) - if is_p2p_supported: - lamport_initialize( - lamport_buffers.local_ptr, - 3 * lamport_buffers_size, - ) - # flag_buffer[0], atomic flag read counter - # flag_buffer[1], non-lamport flag - # flag_buffer[2], lamport flag - flag_buffer = torch.tensor([0, 0, 0], dtype=torch.int, device="cuda") - # layout_buffer[0], clear size for next lamport kernel - # layout_buffer[1], triple buffer offset for lamport kernel - layout_buffer = torch.tensor([0, lamport_buffers_size], - dtype=torch.int64, - device="cuda") - - buffers = [ - ipc_buffers, ipc_barriers, lamport_buffers, flag_buffer, - layout_buffer - ] - - return buffers, torch.tensor( - ipc_buffers.serialize() + ipc_barriers.serialize() + - lamport_buffers.serialize() + [flag_buffer.data_ptr()] + - [layout_buffer.data_ptr()], - dtype=torch.int64, - device="cuda") - - -custom_all_reduce_helper = None - - -def init_all_reduce_helper(): - global custom_all_reduce_helper - custom_all_reduce_helper = CustomAllReduceHelper() - - -def current_all_reduce_helper(): - global custom_all_reduce_helper - assert custom_all_reduce_helper is not None, "You must call `init_all_reduce_helper` first" - return custom_all_reduce_helper diff --git a/tensorrt_llm/profiler.py b/tensorrt_llm/profiler.py index 090657274191..15ad21936185 100644 --- a/tensorrt_llm/profiler.py +++ b/tensorrt_llm/profiler.py @@ -17,10 +17,7 @@ from functools import partial from typing import Literal, Optional, Tuple, Union -# isort: off import torch -import tensorrt as trt -# isort: on try: import psutil @@ -30,12 +27,9 @@ import pynvml except ImportError: pynvml = None -import traceback from tensorrt_llm.logger import logger -from ._common import _is_building - if psutil is None: logger.warning( "A required package 'psutil' is not installed. Will not " @@ -210,76 +204,3 @@ def print_memory_usage( f"Device {_format(alloc_device_mem, unit)}" ) _print_mem_message(msg, tag) - - -@_is_building -def check_gpt_mem_usage( - engine, - kv_dtype, - use_gpt_attention_plugin, - paged_kv_cache, - max_batch_size, - max_beam_width, - max_seq_len, - local_num_kv_heads, - head_size, - num_layers, -) -> int: - # Get the amount of memory - runtime = trt.Runtime(logger.trt_logger) - # 1. TensorRT engine activation memory - activation_size = 0 - try: - cuda_engine = runtime.deserialize_cuda_engine(engine) - assert cuda_engine is not None - activation_size = cuda_engine.device_memory_size_v2 / 1024 / 1024 - del cuda_engine - except Exception: - logger.warning(f"Exception when deserializing engine: {traceback.format_exc()}") - logger.warning("Activation memory size will be regarded as 0.") - logger.info(f"Activation memory size: {activation_size:.2f} MiB") - - # 2. Weights - weights_size = bytes_to_target_unit(engine.nbytes, "MiB") - logger.info(f"Weights memory size: {weights_size:.2f} MiB") - - # 3. Estimated max KV Cache size - kv_cache_size = ( - max_batch_size - * max_beam_width - * 2 - * local_num_kv_heads - * max_seq_len - * head_size - * num_layers - * kv_dtype.itemsize - ) - # without plugin, we need two set of kv cache buffers, - # one for inputs, and the other for outputs. - if not use_gpt_attention_plugin: - kv_cache_size *= 2 - kv_cache_size = bytes_to_target_unit(kv_cache_size, "MiB") - logger.info(f"Max KV Cache memory size: {kv_cache_size:.2f} MiB") - - # Estimated total amount of memory - est_memory_size = activation_size + weights_size + kv_cache_size - logger.info(f"Estimated max memory usage on runtime: {est_memory_size:.2f} MiB") - _, _, total_mem = device_memory_info(torch.cuda.current_device()) - total_mem = bytes_to_target_unit(total_mem, "MiB") - if est_memory_size > total_mem: - logger.warning( - f"Engine is successfully built, but GPU Memory ({total_mem:.2f} MB)" - " may not be enough when running inference on max shape." - ) - if paged_kv_cache: - logger.warning( - "Since paged_kv_cache is enabled, the max KV Cache " - "memory size is a estimate for very extreme cases, " - "it's possible that most cases won't meet OOM." - ) - else: - logger.warning( - "Enabling `--paged_kv_cache` could help reduce the GPU memory usage on runtime." - ) - - return est_memory_size diff --git a/tensorrt_llm/python_plugin.py b/tensorrt_llm/python_plugin.py deleted file mode 100644 index 3c6e4bc62d03..000000000000 --- a/tensorrt_llm/python_plugin.py +++ /dev/null @@ -1,578 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import inspect -import pickle # nosec B403 -import typing -from copy import deepcopy -from dataclasses import dataclass -from typing import Sequence, Type, Union - -import numpy as np -import tensorrt as trt -import torch - -from ._common import default_trtnet -from ._utils import ( - TensorWrapper, - np_dtype_to_trt, - str_dtype_to_trt, - torch_dtype_to_trt, - trt_dtype_to_torch, -) -from .functional import Tensor, _create_tensor -from .plugin.plugin import TRT_LLM_PLUGIN_NAMESPACE - -_plugin_registered = dict() - - -@dataclass(slots=True, frozen=True) -class PluginInfo: - trt_plugin_version: int - plugin_namespace: str - plugin_name: str - plugin_version: str - plugin_num_outputs: int - - def __hash__(self): - return hash((self.plugin_name, self.plugin_namespace, self.plugin_version)) - - def __eq__(self, obj): - if not isinstance(obj, PluginInfo): - return False - return ( - self.plugin_name == obj.plugin_name - and self.plugin_namespace == obj.plugin_namespace - and self.plugin_version == obj.plugin_version - ) - - -def make_expr( - exprBuilder: Union[trt.IExprBuilder, Type[None]], - dim: Union["DimensionExpr", trt.IDimensionExpr, int, Type[None]], -) -> Union[trt.IDimensionExpr, Type[None]]: - """Make a dimension expression. - - Parameters: - exprBuilder: The trt.exprBuilder object. Using it to check whether dim has the same exprBuilder - or to create trt.IDimensionExpr if necessary. - dim: The input dim. - - Returns: - A trt.IDimensionExpr object. - """ - if isinstance(dim, DimensionExpr): - assert exprBuilder == dim.exprBuilder - return dim.expr - elif isinstance(dim, int): - return exprBuilder.constant(dim) - elif dim is None: - return None - elif isinstance(dim, trt.IDimensionExpr): - return dim - else: - raise Exception - - -def expr_operation( - a: Union[trt.IDimensionExpr, Type[None]], - b: Union[trt.IDimensionExpr, Type[None]], - operation: trt.DimensionOperation, - exprBuilder: trt.IExprBuilder, -): - """The function to do expr operation with None support.""" - if exprBuilder is None or a is None or b is None: - expr = None - else: - expr = exprBuilder.operation(operation, a, b) - return DimensionExpr(expr, exprBuilder) - - -class DimensionExpr: - """The class to wrap `trt.IDimensionExpr` to support more pythonic methods.""" - - def __init__( - self, - expr: Union[trt.IDimensionExpr, int, Type[None]], - exprBuilder: Union[trt.IExprBuilder, Type[None]], - ): - self.exprBuilder = exprBuilder - self.expr = expr - - @property - def expr(self): - return self._expr - - @expr.setter - def expr(self, expr: Union["DimensionExpr", trt.IDimensionExpr, int, Type[None]]): - self._expr = make_expr(self.exprBuilder, expr) - - def __add__(self, expr: Union["DimensionExpr", trt.IDimensionExpr, int, Type[None]]): - expr = make_expr(self.exprBuilder, expr) - return expr_operation(self.expr, expr, trt.DimensionOperation.SUM, self.exprBuilder) - - def __radd__(self, expr: Union["DimensionExpr", trt.IDimensionExpr, int, Type[None]]): - return self.__add__(expr) - - def __mul__(self, expr: Union["DimensionExpr", trt.IDimensionExpr, int, Type[None]]): - expr = make_expr(self.exprBuilder, expr) - return expr_operation(self.expr, expr, trt.DimensionOperation.PROD, self.exprBuilder) - - def __rmul__(self, expr: Union["DimensionExpr", trt.IDimensionExpr, int, Type[None]]): - return self.__mul__(expr) - - def __sub__(self, expr: Union["DimensionExpr", trt.IDimensionExpr, int, Type[None]]): - expr = make_expr(self.exprBuilder, expr) - return expr_operation(self.expr, expr, trt.DimensionOperation.SUB, self.exprBuilder) - - def __rsub__(self, expr: Union["DimensionExpr", trt.IDimensionExpr, int, Type[None]]): - expr = make_expr(self.exprBuilder, expr) - return expr_operation(expr, self.expr, trt.DimensionOperation.SUB, self.exprBuilder) - - def __eq__(self, expr: Union["DimensionExpr", trt.IDimensionExpr, int, Type[None]]): - expr = make_expr(self.exprBuilder, expr) - return expr_operation(self.expr, expr, trt.DimensionOperation.EQUAL, self.exprBuilder) - - def __lt__(self, expr: Union["DimensionExpr", trt.IDimensionExpr, int, Type[None]]): - expr = make_expr(self.exprBuilder, expr) - return expr_operation(self.expr, expr, trt.DimensionOperation.LESS, self.exprBuilder) - - def __floordiv__(self, expr: Union["DimensionExpr", trt.IDimensionExpr, int, Type[None]]): - expr = make_expr(self.exprBuilder, expr) - return expr_operation(self.expr, expr, trt.DimensionOperation.FLOOR_DIV, self.exprBuilder) - - def __rfloordiv__(self, expr: Union["DimensionExpr", trt.IDimensionExpr, int, Type[None]]): - expr = make_expr(self.exprBuilder, expr) - return expr_operation(expr, self.expr, trt.DimensionOperation.FLOOR_DIV, self.exprBuilder) - - def __truediv__(self, expr: Union["DimensionExpr", trt.IDimensionExpr, int, Type[None]]): - expr = make_expr(self.exprBuilder, expr) - return expr_operation(self.expr, expr, trt.DimensionOperation.CEIL_DIV, self.exprBuilder) - - def __rtruediv__(self, expr: Union["DimensionExpr", trt.IDimensionExpr, int, Type[None]]): - expr = make_expr(self.exprBuilder, expr) - return expr_operation(expr, self.expr, trt.DimensionOperation.CEIL_DIV, self.exprBuilder) - - def max(self, expr: Union["DimensionExpr", trt.IDimensionExpr, int, Type[None]]): - expr = make_expr(self.exprBuilder, expr) - return expr_operation(self.expr, expr, trt.DimensionOperation.MAX, self.exprBuilder) - - def min(self, expr: Union["DimensionExpr", trt.IDimensionExpr, int, Type[None]]): - expr = make_expr(self.exprBuilder, expr) - return expr_operation(self.expr, expr, trt.DimensionOperation.MIN, self.exprBuilder) - - -class ShapeExpr: - """The class to Wrap `trt.DimsExprs` to support more pythonic methods.""" - - def __init__( - self, - dims: Union[Sequence[trt.IDimensionExpr], Sequence[int], Sequence[type[None]]], - exprBuilder: Union[trt.IExprBuilder, type[None]], - ): - self.exprBuilder = exprBuilder - self.dims = dims - - @property - def dims(self): - return self._dims - - @dims.setter - def dims( - self, - dims: Sequence[Union["DimensionExpr", trt.IDimensionExpr, int, Type[None]]], - ): - if dims is not None: - self._dims = [ - DimensionExpr(make_expr(self.exprBuilder, i), self.exprBuilder) for i in dims - ] - else: - self._dims = None - - def __getitem__(self, index: int): - if self._dims is not None: - return self._dims[index] - else: - return DimensionExpr(None, self.exprBuilder) - - def __setitem__( - self, - index: int, - value: Union["DimensionExpr", trt.IDimensionExpr, int, Type[None]], - ): - if self._dims is None: - return - assert index < len(self._dims) - value = DimensionExpr(make_expr(self.exprBuilder, value), self.exprBuilder) - self._dims[index] = value - - def __len__(self): - if self._dims is None: - return 0 - else: - return len(self._dims) - - def to_trt(self) -> trt.DimsExprs: - return trt.DimsExprs([i.expr for i in self.dims]) - - -class SymTensor: - """The class to represent symbolic tensors. - - Only contains dtype and shape information for users to write their own shape/dtype inference function. - """ - - def __init__( - self, - dtype: Union[torch.dtype, np.dtype, str, trt.DataType, Type[None]], - shape: Union[ShapeExpr, Sequence[int]], - ): - self.dtype = dtype - self.shape = shape - - @property - def shape(self) -> Union[ShapeExpr, Sequence[int]]: - return self._shape - - @shape.setter - def shape(self, shape: Union[ShapeExpr, Sequence[int]]): - assert isinstance(shape, (ShapeExpr, list, tuple)) - if isinstance(shape, (list, tuple)): - for i in shape: - assert isinstance(i, int) - self._shape = shape - - @property - def dtype(self) -> Union[trt.DataType, Type[None]]: - return self._dtype - - @dtype.setter - def dtype(self, dtype: Union[torch.dtype, str, np.dtype, trt.DataType, Type[None]]): - if isinstance(dtype, torch.dtype): - self._dtype = torch_dtype_to_trt(dtype) - elif isinstance(dtype, str): - self._dtype = str_dtype_to_trt(dtype) - elif isinstance(dtype, np.dtype): - self._dtype = np_dtype_to_trt(dtype) - elif isinstance(dtype, trt.DataType): - self._dtype = dtype - elif dtype is None: - self._dtype = None - else: - raise TypeError(f"Unsupported dtype: {dtype}") - - -def _convert_return_value_to_list(ret): - if not isinstance(ret, (list, tuple)): - return [ret] - assert isinstance(ret, (list, tuple)) - return ret - - -class PluginBase( - trt.IPluginV3, trt.IPluginV3OneCore, trt.IPluginV3OneBuild, trt.IPluginV3OneRuntime -): - """The base class of TRT-LLM plugin. - - All TRT-LLM plugin should inherit this class and at least rewrite `forward` and `shape_dtype_inference` - function. `forward` defines the plugin's compute flow while `shape_dtype_inference` defines how would - the output tensor's shape and dtype be inferenced from the input tensor. - """ - - _plugin_creator = None - _no_serialize_attr = {"_current_stream", "_workspace"} - - def __init__(self): - cls = type(self) - # Runtime check for plugin decorator - assert cls._plugin_creator is not None, ( - "Please make sure the plugin is registered through `@trtllm_plugin`" - ) - assert cls != PluginBase - - trt.IPluginV3.__init__(self) - trt.IPluginV3OneCore.__init__(self) - trt.IPluginV3OneBuild.__init__(self) - trt.IPluginV3OneRuntime.__init__(self) - - self.plugin_phase = trt.TensorRTPhase.BUILD - self.num_outputs = self._num_outputs - self.plugin_namespace = self._plugin_namespace - self.plugin_name = self._plugin_name - self.plugin_version = self._plugin_version - self.current_stream = -1 - self.workspace = 0 # nullptr - - @property - def current_stream(self): - if self._current_stream == -1: - return torch.cuda.current_stream().cuda_stream - else: - return self._current_stream - - @current_stream.setter - def current_stream(self, stream: int): - assert isinstance(stream, int) - self._current_stream = stream - - @property - def workspace(self) -> int: - buffer = self._workspace - return buffer if isinstance(buffer, int) else buffer.data_ptr() - - @workspace.setter - def workspace(self, workspace: Union[int, torch.Tensor]): - assert isinstance(workspace, (int, torch.Tensor)) - self._workspace = workspace - - def clone(self): - cls = type(self) - cloned_plugin = cls.__new__(cls) - super(cls, cloned_plugin).__init__() - cloned_plugin.__dict__.update(self._get_dict_to_serialize()) - return cloned_plugin - - def get_capability_interface(self, type): - return self - - def configure_plugin(self, input_desc, output_desc): - pass - - def get_output_data_types(self, input_types): - ret = self.shape_dtype_inference([SymTensor(i, ShapeExpr(None, None)) for i in input_types]) - - ret = _convert_return_value_to_list(ret) - assert len(ret) == self.num_outputs - for i in ret: - assert isinstance(i, SymTensor) - - return [i.dtype for i in ret] - - def get_output_shapes(self, inputs, shape_inputs, exprBuilder): - assert len(shape_inputs) == 0, "Currently we do not support shape inputs" - - ret = self.shape_dtype_inference( - [SymTensor(None, ShapeExpr(i, exprBuilder)) for i in inputs] - ) - - ret = _convert_return_value_to_list(ret) - assert len(ret) == self.num_outputs - for i in ret: - assert isinstance(i, SymTensor) - - return [i.shape.to_trt() for i in ret] - - def supports_format_combination(self, pos, in_out, num_inputs): - """By default, TRT-LLM plugin supports all dtype and linear format. - - It is the users responsibility to check the dtype the plugin supported in `forward` function. - """ - assert pos < len(in_out) - - desc = in_out[pos].desc - if desc.format != trt.TensorFormat.LINEAR: - return False - - return True - - def attach_to_context(self, context): - return self.clone() - - def get_fields_to_serialize(self): - buffer = pickle.dumps(self._get_dict_to_serialize()) - return trt.PluginFieldCollection( - [trt.PluginField("__plugin_pickle_obj__", buffer, trt.PluginFieldType.UNKNOWN)] - ) - - def enqueue(self, input_desc, output_desc, inputs, outputs, workspace, stream): - torch_stream = torch.cuda.ExternalStream(stream_ptr=stream) - self.workspace = workspace - self.current_stream = stream - - with torch.cuda.stream(torch_stream): - self.forward( - tuple( - TensorWrapper.from_trt_desc(input_desc[i], inputs[i]) - for i in range(len(input_desc)) - ), - tuple( - TensorWrapper.from_trt_desc(output_desc[i], outputs[i]) - for i in range(len(output_desc)) - ), - ) - - self.current_stream = -1 - - def __call__(self, *args: Union[Sequence[TensorWrapper], Sequence[torch.Tensor]]): - is_trtllm = True - for i in args: - is_trtllm &= isinstance(i, Tensor) - - if not is_trtllm: - for i in args: - assert isinstance(i, torch.Tensor), ( - "Plugin inputs must be `tensorrt_llm.Tensor`s or `torch.Tensor`s" - ) - sym_tensors = self.shape_dtype_inference( - [SymTensor(i.dtype, [j for j in i.shape]) for i in args] - ) - sym_tensors = _convert_return_value_to_list(sym_tensors) - ret = [ - torch.empty(sym_tensor.shape, dtype=trt_dtype_to_torch(sym_tensor.dtype)) - for sym_tensor in sym_tensors - ] - self.current_stream = torch.cuda.current_stream().cuda_stream - self.workspace = torch.empty(self.workspace) - self.forward(args, ret) - else: - args = [i.trt_tensor for i in args] - layer_plugin = default_trtnet().add_plugin_v3(args, [], self) - ret = [ - _create_tensor(layer_plugin.get_output(i), layer_plugin) - for i in range(self.num_outputs) - ] - - if len(ret) == 1: - return ret[0] - - return ret - - def on_shape_change(self, input_desc, output_desc): - pass - - def get_valid_tactics(self): - return [] - - def set_tactic(self, index): - if index != 0: - raise RuntimeError( - "By default TRT should not set tactics since PluginBase do not provide custom tactic." - ) - - def forward(self, inputs: Sequence[TensorWrapper], outputs: Sequence[TensorWrapper]): - """Expect users to rewrite this function to define the compute flow. - - There are a few special attributes for users to get access to some resources. - - `self.workspace`: The workspace address of TRT managed workspace. - `self.current_stream`: The CUDA stream this plugin is expected to execute on. By default - `PluginBase` set the torch.cuda.current_stream() to this stream. This attribute is for the - toolkit that doesn't work with torch's stream. - """ - raise NotImplementedError - - def shape_dtype_inference(self, inputs: Sequence[SymTensor]): - """Expect users to rewrite this function to define the shape dtype inference for output tensors.""" - raise NotImplementedError - - def _get_dict_to_serialize(self): - ret = {} - for k, v in self.__dict__.items(): - if k not in self._no_serialize_attr: - ret[k] = deepcopy(v) if self.deepcopy_clone else v - return ret - - -class PluginCreatorBase(trt.IPluginCreatorV3One): - def __init__(self): - super().__init__() - - def create_plugin(self, name, fc, phase): - if len(fc) == 1 and fc[0].name == "__plugin_pickle_obj__": - data = fc[0].data - plugin_dict = pickle.loads(data) # nosec B301 - plugin = self.plugin_cls.__new__(self.plugin_cls) - super(self.plugin_cls, plugin).__init__() - plugin.__dict__.update(plugin_dict) - else: - raise RuntimeError("Expect to be called by TRT") - plugin.plugin_phase = phase - return plugin - - -def trtllm_plugin( - plugin_name: str, - *, - plugin_version: str = "1", - plugin_namespace: str = TRT_LLM_PLUGIN_NAMESPACE, - plugin_num_outputs: Union[int, Type[None]] = None, - deepcopy_clone: bool = True, - no_serialize_attr: Sequence[str] = set(), -): - def plugin_registration(plugin_cls): - assert issubclass(plugin_cls, PluginBase) - assert hasattr(plugin_cls, "__dict__"), ( - "Plugin wrapper uses `__dict__` to track plugin states" - ) - nonlocal plugin_num_outputs - - annotation = inspect.signature(plugin_cls.shape_dtype_inference).return_annotation - origin_annotation = typing.get_origin(annotation) - if origin_annotation is tuple or annotation is SymTensor: - if origin_annotation is tuple: - element_types = typing.get_args(annotation) - for ty in element_types: - assert ty == SymTensor, ( - f"Plugin {plugin_name}'s `shape_dtype_inference` return annotation must be SymTensor " - "or a tuple of SymTensor" - ) - infered_num_outputs = len(element_types) - else: - infered_num_outputs = 1 - if plugin_num_outputs is not None: - assert plugin_num_outputs == infered_num_outputs, ( - f"Plugin {plugin_name}'s `_num_outputs` and return annotation mismatch, " - f"{plugin_cls._num_outputs} != {infered_num_outputs}" - ) - plugin_num_outputs = infered_num_outputs - else: - assert plugin_num_outputs is not None, ( - "Must specify `num_outputs` or valid `shape_dtype_inference` return annotation for " - f"{plugin_name}. The valid types are SymTensor or a tuple of SymTensor, got {annotation}." - ) - - plugin_info = PluginInfo( - 3, plugin_namespace, plugin_name, plugin_version, plugin_num_outputs - ) - assert plugin_info not in _plugin_registered, ( - f"Redefine plugin with info: {plugin_info} which is previously defined as " - f"{_plugin_registered[plugin_info]}" - ) - - _plugin_registered[plugin_info] = plugin_info - plugin_cls._plugin_name = plugin_name - plugin_cls._plugin_version = plugin_version - plugin_cls._plugin_namespace = plugin_namespace - plugin_cls._num_outputs = plugin_num_outputs - plugin_cls.deepcopy_clone = deepcopy_clone - plugin_cls._no_serialize_attr.update(no_serialize_attr) - - plugin_registry = trt.get_plugin_registry() - - plugin_creator = PluginCreatorBase() - plugin_creator.name = plugin_cls._plugin_name - plugin_creator.plugin_namespace = plugin_cls._plugin_namespace - plugin_creator.plugin_version = plugin_cls._plugin_version - plugin_creator.field_names = trt.PluginFieldCollection([]) - plugin_creator.plugin_cls = plugin_cls - - plugin_cls._plugin_creator = plugin_creator - ret = plugin_registry.register_creator(plugin_creator, plugin_cls._plugin_namespace) - - assert ret, f"Plugin: {plugin_cls} register failed, please check the error log." - - return plugin_cls - - return plugin_registration diff --git a/tensorrt_llm/quantization/functional.py b/tensorrt_llm/quantization/functional.py index ddaa7394eb5b..2eb1eab32624 100644 --- a/tensorrt_llm/quantization/functional.py +++ b/tensorrt_llm/quantization/functional.py @@ -12,950 +12,19 @@ # WITHOUT 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 typing import Optional, Tuple, Union -import numpy as np -import tensorrt as trt import torch -import torch.nn.functional as F -from .._common import default_net, default_trtnet -from .._utils import (get_sm_version, str_dtype_to_np, str_dtype_to_trt, - trt_dtype_to_np) -from ..functional import (Tensor, _add_plugin_info, _create_tensor, cast, clip, - constant, flatten, layer_norm, matmul, - repeat_interleave, rms_norm, round, sum, view) -from ..layers.linear import ColumnLinear -from ..parameter import Parameter -from ..plugin import TRT_LLM_PLUGIN_NAMESPACE -from .mode import QuantMode - - -def smooth_quant_gemm(input: Tensor, weights: Tensor, scales_a: Tensor, - scales_b: Tensor, per_token_scaling: bool, - per_channel_scaling: bool, dtype: str) -> Tensor: - if not default_net().plugin_config.smooth_quant_gemm_plugin: - if per_token_scaling and input.size(0) == -1: - # WAR for DQ per-token scaling doesn't support dynamic shapes - - scale_one = constant(np.array(1.0, dtype=np.float32)) - input = dequantize(input, scale_one, 0, 'float32') - weights = dequantize(weights, scale_one, 0, 'float32') - result = matmul(input, weights, False, True, False) - scales = matmul(scales_a, scales_b, False, False, False) - result = result * scales - result = cast(result, dtype) - return result - else: - if not per_token_scaling: - scales_a = view(scales_a, []) - else: - scales_a = flatten(scales_a) - if not per_channel_scaling: - scales_b = view(scales_b, []) - else: - scales_b = flatten(scales_b) - input = dequantize(input, scales_a, 0, dtype) - weights = dequantize(weights, scales_b, 0, dtype) - result = matmul(input, weights, False, True, False) - return result - else: - plg_creator = trt.get_plugin_registry().get_plugin_creator( - 'SmoothQuantGemm', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert plg_creator is not None - - per_channel_scaling = 1 if per_channel_scaling else 0 - per_channel_scaling = trt.PluginField( - "has_per_channel_scaling", - np.array(per_channel_scaling, dtype=np.int32), - trt.PluginFieldType.INT32) - - per_token_scaling = 1 if per_token_scaling else 0 - per_token_scaling = trt.PluginField( - "has_per_token_scaling", np.array(per_token_scaling, - dtype=np.int32), - trt.PluginFieldType.INT32) - - p_dtype = default_net().plugin_config.smooth_quant_gemm_plugin - pf_type = trt.PluginField( - "type_id", np.array([int(str_dtype_to_trt(p_dtype))], np.int32), - trt.PluginFieldType.INT32) - - pfc = trt.PluginFieldCollection( - [per_channel_scaling, per_token_scaling, pf_type]) - gemm_plug = plg_creator.create_plugin("sq_gemm", pfc) - plug_inputs = [ - input.trt_tensor, weights.trt_tensor, scales_a.trt_tensor, - scales_b.trt_tensor - ] - layer = default_trtnet().add_plugin_v2(plug_inputs, gemm_plug) - _add_plugin_info(layer, plg_creator, "sq_gemm", pfc) - if not default_net().strongly_typed: - layer.get_input(0).set_dynamic_range(-127, 127) - layer.get_input(1).set_dynamic_range(-127, 127) - return _create_tensor(layer.get_output(0), layer) - - -def qserve_gemm_per_group(input: Tensor, - act_scales: Tensor, - weights: Tensor, - s1_scales: Tensor, - s2_scales: Tensor, - s2_zeros: Tensor, - group_size: int = 128) -> Tensor: - if not default_net().plugin_config.qserve_gemm_plugin: - raise TypeError("QServe Quant GEMM is only supported with plugin") - else: - plg_creator = trt.get_plugin_registry().get_plugin_creator( - 'QServeGemm', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert plg_creator is not None - - p_dtype = default_net().plugin_config.qserve_gemm_plugin - pf_type = trt.PluginField( - "type_id", np.array([int(str_dtype_to_trt(p_dtype))], np.int32), - trt.PluginFieldType.INT32) - pf_group_size = trt.PluginField("group_size", - np.array([group_size], np.int32), - trt.PluginFieldType.INT32) - - pfc = trt.PluginFieldCollection([pf_type, pf_group_size]) - gemm_plug = plg_creator.create_plugin("qserve_gemm", pfc) - plug_inputs = [ - input.trt_tensor, weights.trt_tensor, s2_zeros.trt_tensor, - s2_scales.trt_tensor, s1_scales.trt_tensor, act_scales.trt_tensor - ] - layer = default_trtnet().add_plugin_v2(plug_inputs, gemm_plug) - _add_plugin_info(layer, plg_creator, "qserve_gemm", pfc) - if not default_net().strongly_typed: - # Useless. But must be kept otherwise leads to the following TRT API Usage error: - # input/output with DataType Int8 in network without Q/DQ layers must have dynamic range set when no calibrator is used - layer.get_input(0).set_dynamic_range(-128, 127) - layer.get_input(1).set_dynamic_range(-128, 127) - layer.get_input(2).set_dynamic_range(-128, 127) - layer.get_input(3).set_dynamic_range(-128, 127) - return _create_tensor(layer.get_output(0), layer) - - -def qserve_gemm_per_channel(input: Tensor, act_scales: Tensor, act_sums: Tensor, - weights: Tensor, s1_scales: Tensor, - s1_szeros: Tensor) -> Tensor: - if not default_net().plugin_config.qserve_gemm_plugin: - raise TypeError("QServe Quant GEMM is only supported with plugin") - else: - plg_creator = trt.get_plugin_registry().get_plugin_creator( - 'QServeGemm', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert plg_creator is not None - - p_dtype = default_net().plugin_config.qserve_gemm_plugin - pf_type = trt.PluginField( - "type_id", np.array([int(str_dtype_to_trt(p_dtype))], np.int32), - trt.PluginFieldType.INT32) - pf_group_size = trt.PluginField("group_size", np.array([-1], np.int32), - trt.PluginFieldType.INT32) - - pfc = trt.PluginFieldCollection([pf_type, pf_group_size]) - gemm_plug = plg_creator.create_plugin("qserve_gemm", pfc) - - plug_inputs = [ - input.trt_tensor, weights.trt_tensor, s1_scales.trt_tensor, - s1_szeros.trt_tensor, act_sums.trt_tensor, act_scales.trt_tensor - ] - layer = default_trtnet().add_plugin_v2(plug_inputs, gemm_plug) - _add_plugin_info(layer, plg_creator, "qserve_gemm", pfc) - - if not default_net().strongly_typed: - # Useless. But must be kept otherwise leads to the following TRT API Usage error: - # input/output with DataType Int8 in network without Q/DQ layers must have dynamic range set when no calibrator is used - layer.get_input(0).set_dynamic_range(-128, 127) - layer.get_input(1).set_dynamic_range(-128, 127) - - return _create_tensor(layer.get_output(0), layer) - - -def fp8_rowwise_gemm(input: Tensor, weights: Tensor, scales_a: Tensor, - scales_b: Tensor, per_token_scaling: bool, - per_channel_scaling: bool) -> Tensor: - if not default_net().plugin_config.fp8_rowwise_gemm_plugin: - raise TypeError("Fp8 Rowwise GEMM is only supported with plugin") - else: - plg_creator = trt.get_plugin_registry().get_plugin_creator( - 'Fp8RowwiseGemm', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert plg_creator is not None - - per_channel_scaling = 1 if per_channel_scaling else 0 - per_channel_scaling = trt.PluginField( - "has_per_channel_scaling", - np.array(per_channel_scaling, dtype=np.int32), - trt.PluginFieldType.INT32) - - per_token_scaling = 1 if per_token_scaling else 0 - per_token_scaling = trt.PluginField( - "has_per_token_scaling", np.array(per_token_scaling, - dtype=np.int32), - trt.PluginFieldType.INT32) - - p_dtype = default_net().plugin_config.fp8_rowwise_gemm_plugin - pf_type = trt.PluginField( - "type_id", np.array([int(str_dtype_to_trt(p_dtype))], np.int32), - trt.PluginFieldType.INT32) - - pfc = trt.PluginFieldCollection( - [per_channel_scaling, per_token_scaling, pf_type]) - gemm_plug = plg_creator.create_plugin("fp8_rowwise_gemm", pfc) - plug_inputs = [ - input.trt_tensor, weights.trt_tensor, scales_a.trt_tensor, - scales_b.trt_tensor - ] - layer = default_trtnet().add_plugin_v2(plug_inputs, gemm_plug) - _add_plugin_info(layer, plg_creator, "fp8_rowwise_gemm", pfc) - if not default_net().strongly_typed: - layer.get_input(0).set_dynamic_range(-448, 448) - layer.get_input(1).set_dynamic_range(-448, 448) - return _create_tensor(layer.get_output(0), layer) - - -def weight_only_quant_matmul(input: Tensor, - weights: Tensor, - scales: Tensor, - weightTypeId: int, - dtype: str = 'float16', - transa: bool = False, - transb: bool = False) -> Tensor: - if not default_net( - ).plugin_config.weight_only_quant_matmul_plugin or transa or transb: - scale_axis = 0 if transb else 1 - if weights.dtype != trt.int8: - # Q->DQ - weights = quantize(weights, scales, dtype='int8', axis=1) - weights = dequantize(weights, scales, scale_axis, input.dtype) - else: - weights = dequantize(weights, scales, scale_axis, input.dtype) - - res = matmul(input, weights, transa=transa, transb=transb) - return cast(res, dtype) - else: - plg_creator = trt.get_plugin_registry().get_plugin_creator( - 'WeightOnlyQuantMatmul', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert plg_creator is not None - - weight_type_id = trt.PluginField("weight_type_id", - np.array(weightTypeId, dtype=np.int32), - trt.PluginFieldType.INT32) - - p_dtype = default_net().plugin_config.weight_only_quant_matmul_plugin - pf_type = trt.PluginField( - "type_id", np.array([int(str_dtype_to_trt(p_dtype))], np.int32), - trt.PluginFieldType.INT32) - - pfc = trt.PluginFieldCollection([pf_type, weight_type_id]) - matmul_plug = plg_creator.create_plugin("woq_matmul", pfc) - plug_inputs = [input.trt_tensor, weights.trt_tensor, scales.trt_tensor] - layer = default_trtnet().add_plugin_v2(plug_inputs, matmul_plug) - _add_plugin_info(layer, plg_creator, "woq_matmul", pfc) - if not default_net().strongly_typed: - layer.get_input(1).set_dynamic_range(-127, 127) - return _create_tensor(layer.get_output(0), layer) - - -def weight_only_groupwise_quant_matmul(input: Tensor, - pre_quant_scale: Tensor, - weights: Tensor, - scales: Tensor, - zeros: Tensor, - biases: Tensor, - alpha: Parameter, - quant_algo: int, - group_size: int, - dtype: str = 'float16') -> Tensor: - if not default_net( - ).plugin_config.weight_only_groupwise_quant_matmul_plugin: - scales = repeat_interleave(scales, group_size, 0) - weights = quantize(weights, scales, dtype='int8', axis=1) - weights = dequantize(weights, scales, 1, input.dtype) - - if quant_algo & 8: - # fp8_alpha - input = input * alpha.value - if quant_algo & 4: - # pre quant - input = input * pre_quant_scale - elif quant_algo & 2: - # zero - zeros = repeat_interleave(zeros, group_size, 0) - weights += zeros - res = matmul(input, weights) - if quant_algo & 1: - # bias - res += biases - - return cast(res, dtype) - else: - plg_creator = trt.get_plugin_registry().get_plugin_creator( - 'WeightOnlyGroupwiseQuantMatmul', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert plg_creator is not None - - quant_algo_ = trt.PluginField("quant_algo", - np.array(quant_algo, dtype=np.int32), - trt.PluginFieldType.INT32) - group_size_ = trt.PluginField("group_size", - np.array(group_size, dtype=np.int32), - trt.PluginFieldType.INT32) - - if alpha: - alpha.is_buffer = True - alpha_value = alpha.raw_value[0] - else: - alpha_value = 1.0 - - alpha_ = trt.PluginField("alpha", np.array(alpha_value, - dtype=np.float32), - trt.PluginFieldType.FLOAT32) - - p_dtype = default_net( - ).plugin_config.weight_only_groupwise_quant_matmul_plugin - pf_type_ = trt.PluginField( - "type_id", np.array([int(str_dtype_to_trt(p_dtype))], np.int32), - trt.PluginFieldType.INT32) - - pfc = trt.PluginFieldCollection( - [pf_type_, quant_algo_, group_size_, alpha_]) - - matmul_plug = plg_creator.create_plugin("woq_groupwise_matmul", pfc) - - # quant_algo = use_int8_weight * 16 + fp8_alpha * 8 + pre_quant_scale * 4 + zero * 2 + bias - plug_inputs = [input.trt_tensor] - - # Flags for indicating whether the corresponding inputs are applied in quant_algo - # quant_algo = use_int8_weight * INT8_WEIGHT + fp8_alpha * FP8_ALPHA + pre_quant_scale * PRE_QUANT_SCALE + zero * ZERO + bias * BIAS - # Here use_int8_weight, pre_quant_scale, zero and bias are boolean type - BIAS = 1 - ZERO = 2 - PRE_QUANT_SCALE = 4 - - if quant_algo & PRE_QUANT_SCALE: - plug_inputs += [pre_quant_scale.trt_tensor] - - plug_inputs += [weights.trt_tensor, scales.trt_tensor] - - if quant_algo & ZERO: - plug_inputs += [zeros.trt_tensor] - if quant_algo & BIAS: - plug_inputs += [biases.trt_tensor] - - layer = default_trtnet().add_plugin_v2(plug_inputs, matmul_plug) - _add_plugin_info(layer, plg_creator, "woq_groupwise_matmul", pfc) - - return _create_tensor(layer.get_output(0), layer) - - -# TODO: Should be renamed to layer_norm_quantize. -def smooth_quant_layer_norm(input: Tensor, - normalized_shape: Union[int, Tuple[int]], - weight: Optional[Tensor] = None, - bias: Optional[Tensor] = None, - scale: Optional[Tensor] = None, - eps: float = 1e-05, - use_diff_of_squares: bool = True, - dynamic_act_scaling: bool = False) -> Tensor: - if not default_net().plugin_config.layernorm_quantization_plugin: - dtype = trt_dtype_to_np(input.dtype) - if weight is None: - weight = constant(np.ones(normalized_shape, dtype=dtype)) - if bias is None: - bias = constant(np.zeros(normalized_shape, dtype=dtype)) - result = layer_norm(input, normalized_shape, weight, bias, eps, - use_diff_of_squares) - if not dynamic_act_scaling: - return quantize_tensor(result, scale) - else: - return quantize_per_token(result) - else: - plg_creator = trt.get_plugin_registry().get_plugin_creator( - 'LayernormQuantization', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert plg_creator is not None - - output_type = trt.PluginField("out_type_id", - np.array([int(trt.int8)], np.int32), - trt.PluginFieldType.INT32) - quant_mode = trt.PluginField( - "quant_mode", - np.array([int(QuantMode.use_smooth_quant(per_token=True))], - np.int32), trt.PluginFieldType.INT32) - eps = trt.PluginField("eps", np.array(eps, dtype=np.float32), - trt.PluginFieldType.FLOAT32) - use_diff_of_squares = trt.PluginField( - "use_diff_of_squares", - np.array([int(use_diff_of_squares)], dtype=np.int32), - trt.PluginFieldType.INT32) - - dyn_act_scaling = trt.PluginField( - "dyn_act_scaling", np.array([int(dynamic_act_scaling)], np.int32), - trt.PluginFieldType.INT32) - - p_dtype = default_net().plugin_config.layernorm_quantization_plugin - pf_type = trt.PluginField( - "type_id", np.array([int(str_dtype_to_trt(p_dtype))], np.int32), - trt.PluginFieldType.INT32) - pfc = trt.PluginFieldCollection([ - eps, use_diff_of_squares, dyn_act_scaling, pf_type, output_type, - quant_mode - ]) - layernorm_plug = plg_creator.create_plugin("layernorm_quantized", pfc) - normalized_shape = [normalized_shape] if isinstance( - normalized_shape, int) else normalized_shape - if weight is None: - weight = constant( - np.ones(normalized_shape, dtype=str_dtype_to_np(p_dtype))) - if bias is None: - bias = constant( - np.zeros(normalized_shape, dtype=str_dtype_to_np(p_dtype))) - - # LayerNorm plugin only supports float32 scale - scale = cast(scale, "float32") - plug_inputs = [ - input.trt_tensor, weight.trt_tensor, bias.trt_tensor, - scale.trt_tensor - ] - layer = default_trtnet().add_plugin_v2(plug_inputs, layernorm_plug) - if not default_net().strongly_typed: - layer.get_output(0).set_dynamic_range(-127, 127) - _add_plugin_info(layer, plg_creator, "layernorm_quantized", pfc) - if not dynamic_act_scaling: - return _create_tensor(layer.get_output(0), layer) - - return _create_tensor(layer.get_output(0), - layer), _create_tensor(layer.get_output(1), layer) - - -# TODO: Should be renamed to rms_norm_quantize. This is also used by QServe. -def smooth_quant_rms_norm( - input: Tensor, - normalized_shape: Union[int, Tuple[int]], - weight: Optional[Tensor] = None, - bias: Optional[Tensor] = None, - scale: Optional[Tensor] = None, - clamp_val: Optional[Tensor] = None, - eps: float = 1e-05, - dynamic_act_scaling: bool = False, - scale_dtype='float32', - sum_per_token: bool = False, - sum_dtype='float32' -) -> Tensor | tuple[Tensor, Tensor] | tuple[Tensor, Tensor, Tensor]: - if sum_per_token and not dynamic_act_scaling: - raise ValueError( - "sum_per_token is only allowed if dynamic_act_scaling is enabled!") - - if not default_net().plugin_config.rmsnorm_quantization_plugin: - result = rms_norm(input, normalized_shape, 1, weight, eps) - if bias is not None: - result += bias - if not dynamic_act_scaling: - return quantize_tensor(result, scale) - else: - return quantize_per_token(result, clamp_val, scale_dtype, - sum_per_token, sum_dtype) - else: - plg_creator = trt.get_plugin_registry().get_plugin_creator( - 'RmsnormQuantization', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert plg_creator is not None - - output_type = trt.PluginField("out_type_id", - np.array([int(trt.int8)], np.int32), - trt.PluginFieldType.INT32) - quant_mode = trt.PluginField( - "quant_mode", - np.array([int(QuantMode.use_smooth_quant(per_token=True))], - np.int32), trt.PluginFieldType.INT32) - clamp_enabled = trt.PluginField( - "clamp_enabled", np.array([clamp_val is not None], np.int32), - trt.PluginFieldType.INT32) - - eps = trt.PluginField("eps", np.array(eps, dtype=np.float32), - trt.PluginFieldType.FLOAT32) - - dyn_act_scaling = trt.PluginField( - "dyn_act_scaling", np.array([int(dynamic_act_scaling)], np.int32), - trt.PluginFieldType.INT32) - - sum_per_token_pf = trt.PluginField( - "sum_per_token", np.array([int(sum_per_token)], np.int32), - trt.PluginFieldType.INT32) - - p_dtype = default_net().plugin_config.rmsnorm_quantization_plugin - pf_type = trt.PluginField( - "type_id", np.array([int(str_dtype_to_trt(p_dtype))], np.int32), - trt.PluginFieldType.INT32) - pfc = trt.PluginFieldCollection([ - eps, dyn_act_scaling, sum_per_token_pf, clamp_enabled, quant_mode, - pf_type, output_type - ]) - rmsnorm_plug = plg_creator.create_plugin("rmsnorm_quantized", pfc) - normalized_shape = [normalized_shape] if isinstance( - normalized_shape, int) else normalized_shape - if weight is None: - weight = constant( - np.ones(normalized_shape, dtype=str_dtype_to_np(p_dtype))) - if bias is None: - bias = constant( - np.zeros(normalized_shape, dtype=str_dtype_to_np(p_dtype))) - - # TODO: Why not fuse scale (which seems to be a per-tensor scaling factor of the original values) into weight? - if scale is None: - scale = constant(np.ones(1, dtype=str_dtype_to_np(p_dtype))) - - # RMS Norm Plugin only supports float32 scale - scale = cast(scale, "float32") - - plug_inputs = [ - input.trt_tensor, weight.trt_tensor, bias.trt_tensor, - scale.trt_tensor - ] - if clamp_val: - plug_inputs += [clamp_val.trt_tensor] - layer = default_trtnet().add_plugin_v2(plug_inputs, rmsnorm_plug) - if not default_net().strongly_typed: - layer.get_output(0).set_dynamic_range(-127, 127) - _add_plugin_info(layer, plg_creator, "rmsnorm_quantized", pfc) - if not dynamic_act_scaling: - return _create_tensor(layer.get_output(0), layer) - - output_quantized = _create_tensor(layer.get_output(0), layer) - output_scales = _create_tensor(layer.get_output(1), layer) - - # TODO: The plugin should be able to directly output float16 scales - if str_dtype_to_trt(scale_dtype) != output_scales.dtype: - output_scales = cast(output_scales, scale_dtype) - - if not sum_per_token: - return output_quantized, output_scales - - output_sums = _create_tensor(layer.get_output(2), layer) - # TODO: The plugin should be able to directly output float16 sums - if str_dtype_to_trt(sum_dtype) != output_sums.dtype: - output_sums = cast(output_sums, sum_dtype) - - return output_quantized, output_scales, output_sums - - -def fp8_rowwise_rms_norm(input: Tensor, - normalized_shape: Union[int, Tuple[int]], - weight: Optional[Tensor] = None, - bias: Optional[Tensor] = None, - scale: Optional[Tensor] = None, - clamp_val: Optional[Tensor] = None, - eps: float = 1e-05, - dynamic_act_scaling: bool = True) -> Tensor: - if not default_net().plugin_config.rmsnorm_quantization_plugin: - raise TypeError("Fp8 Rowwise Rms Norm is only supported with plugin") - else: - plg_creator = trt.get_plugin_registry().get_plugin_creator( - 'RmsnormQuantization', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert plg_creator is not None - - output_type = trt.PluginField("out_type_id", - np.array([int(trt.fp8)], np.int32), - trt.PluginFieldType.INT32) - quant_mode = trt.PluginField( - "quant_mode", - np.array([int(QuantMode.from_description(use_fp8_rowwise=True))], - np.int32), trt.PluginFieldType.INT32) - clamp_enabled = trt.PluginField( - "clamp_enabled", np.array([clamp_val is not None], np.int32), - trt.PluginFieldType.INT32) - - eps = trt.PluginField("eps", np.array(eps, dtype=np.float32), - trt.PluginFieldType.FLOAT32) - - dyn_act_scaling = trt.PluginField( - "dyn_act_scaling", np.array([int(dynamic_act_scaling)], np.int32), - trt.PluginFieldType.INT32) - sum_per_token_pf = trt.PluginField("sum_per_token", - np.array([int(False)], np.int32), - trt.PluginFieldType.INT32) - - p_dtype = default_net().plugin_config.rmsnorm_quantization_plugin - pf_type = trt.PluginField( - "type_id", np.array([int(str_dtype_to_trt(p_dtype))], np.int32), - trt.PluginFieldType.INT32) - - pfc = trt.PluginFieldCollection([ - eps, dyn_act_scaling, sum_per_token_pf, clamp_enabled, quant_mode, - pf_type, output_type - ]) - - rmsnorm_plug = plg_creator.create_plugin("rmsnorm_quantized", pfc) - normalized_shape = [normalized_shape] if isinstance( - normalized_shape, int) else normalized_shape - if weight is None: - weight = constant( - np.ones(normalized_shape, dtype=str_dtype_to_np(p_dtype))) - if bias is None: - bias = constant( - np.zeros(normalized_shape, dtype=str_dtype_to_np(p_dtype))) - if scale is None: - scale = constant(np.ones((1, ), dtype=str_dtype_to_np(p_dtype))) - - # RMS Norm Plugin only supports float32 scale - scale = cast(scale, "float32") - plug_inputs = [ - input.trt_tensor, weight.trt_tensor, bias.trt_tensor, - scale.trt_tensor - ] - if clamp_val: - plug_inputs += [clamp_val.trt_tensor] - layer = default_trtnet().add_plugin_v2(plug_inputs, rmsnorm_plug) - if not default_net().strongly_typed: - layer.get_output(0).set_dynamic_range(-448, 448) - _add_plugin_info(layer, plg_creator, "rmsnorm_quantized", pfc) - if not dynamic_act_scaling: - return _create_tensor(layer.get_output(0), layer) - - return _create_tensor(layer.get_output(0), - layer), _create_tensor(layer.get_output(1), layer) - - -def fp8_rowwise_layer_norm(input: Tensor, - normalized_shape: Union[int, Tuple[int]], - weight: Optional[Tensor] = None, - bias: Optional[Tensor] = None, - scale: Optional[Tensor] = None, - clamp_val: Optional[Tensor] = None, - eps: float = 1e-05, - dynamic_act_scaling: bool = True) -> Tensor: - if not default_net().plugin_config.layernorm_quantization_plugin: - raise TypeError("Fp8 Rowwise Layer Norm is only supported with plugin") - else: - plg_creator = trt.get_plugin_registry().get_plugin_creator( - 'LayernormQuantization', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert plg_creator is not None - - output_type = trt.PluginField("out_type_id", - np.array([int(trt.fp8)], np.int32), - trt.PluginFieldType.INT32) - quant_mode = trt.PluginField( - "quant_mode", - np.array([int(QuantMode.from_description(use_fp8_rowwise=True))], - np.int32), trt.PluginFieldType.INT32) - clamp_enabled = trt.PluginField( - "clamp_enabled", np.array([clamp_val is not None], np.int32), - trt.PluginFieldType.INT32) - - eps = trt.PluginField("eps", np.array(eps, dtype=np.float32), - trt.PluginFieldType.FLOAT32) - - dyn_act_scaling = trt.PluginField( - "dyn_act_scaling", np.array([int(dynamic_act_scaling)], np.int32), - trt.PluginFieldType.INT32) - sum_per_token_pf = trt.PluginField("sum_per_token", - np.array([int(False)], np.int32), - trt.PluginFieldType.INT32) - - p_dtype = default_net().plugin_config.layernorm_quantization_plugin - pf_type = trt.PluginField( - "type_id", np.array([int(str_dtype_to_trt(p_dtype))], np.int32), - trt.PluginFieldType.INT32) - - pfc = trt.PluginFieldCollection([ - eps, dyn_act_scaling, sum_per_token_pf, clamp_enabled, quant_mode, - pf_type, output_type - ]) - - layernorm_plug = plg_creator.create_plugin("layernorm_quantized", pfc) - normalized_shape = [normalized_shape] if isinstance( - normalized_shape, int) else normalized_shape - if weight is None: - weight = constant( - np.ones(normalized_shape, dtype=str_dtype_to_np(p_dtype))) - if bias is None: - bias = constant( - np.zeros(normalized_shape, dtype=str_dtype_to_np(p_dtype))) - if scale is None: - scale = constant(np.ones((1, ), dtype=str_dtype_to_np(p_dtype))) - - # Layer Norm Plugin only supports float32 scale - scale = cast(scale, "float32") - plug_inputs = [ - input.trt_tensor, weight.trt_tensor, bias.trt_tensor, - scale.trt_tensor - ] - if clamp_val: - plug_inputs += [clamp_val.trt_tensor] - layer = default_trtnet().add_plugin_v2(plug_inputs, layernorm_plug) - if not default_net().strongly_typed: - layer.get_output(0).set_dynamic_range(-448, 448) - _add_plugin_info(layer, plg_creator, "layernorm_quantized", pfc) - if not dynamic_act_scaling: - return _create_tensor(layer.get_output(0), layer) - - return _create_tensor(layer.get_output(0), - layer), _create_tensor(layer.get_output(1), layer) - - -def fused_layernorm( - input: Tensor, - normalized_shape: Union[int, Tuple[int]], - residual: Optional[Tensor] = None, - weight: Optional[Tensor] = None, - # beta: Optional[Tensor] = None, - # bias: Optional[Tensor] = None, - scale: Optional[Tensor] = None, - eps: float = 1e-05, - p_dtype: str = 'float16', - need_fp32_output: bool = False) -> Tensor: - plg_creator = trt.get_plugin_registry().get_plugin_creator( - 'FusedLayernorm', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert plg_creator is not None - eps = trt.PluginField("eps", np.array(eps, dtype=np.float32), - trt.PluginFieldType.FLOAT32) - pf_type = trt.PluginField( - "type_id", np.array([int(str_dtype_to_trt(p_dtype))], np.int32), - trt.PluginFieldType.INT32) - need_fp32_output_value = need_fp32_output - need_fp32_output = trt.PluginField( - "need_fp32_output", np.array([int(need_fp32_output_value)], np.int32), - trt.PluginFieldType.INT32) - need_quantize_value = scale is not None - need_quantize = trt.PluginField( - "need_quantize", np.array([int(need_quantize_value)], np.int32), - trt.PluginFieldType.INT32) - pfc = trt.PluginFieldCollection( - [eps, need_fp32_output, need_quantize, pf_type]) - fused_layernorm_plug = plg_creator.create_plugin("fused_layernorm", pfc) - normalized_shape = [normalized_shape] if isinstance( - normalized_shape, int) else normalized_shape - if weight is None: - weight = constant( - np.ones(normalized_shape, dtype=str_dtype_to_np(p_dtype))) - # if beta is None: - # beta = constant( - # np.zeros(normalized_shape, dtype=str_dtype_to_np(p_dtype))) - # if bias is None: - # bias = constant( - # np.zeros(normalized_shape, dtype=str_dtype_to_np(p_dtype))) - if need_quantize_value: - plug_inputs = [ - input.trt_tensor, residual.trt_tensor, weight.trt_tensor, - scale.trt_tensor - ] - else: - plug_inputs = [ - input.trt_tensor, - residual.trt_tensor, - weight.trt_tensor, - ] - layer = default_trtnet().add_plugin_v2(plug_inputs, fused_layernorm_plug) - _add_plugin_info(layer, plg_creator, "fused_layernorm", pfc) - if not need_quantize_value: - return _create_tensor(layer.get_output(0), - layer), _create_tensor(layer.get_output(1), layer) - return _create_tensor(layer.get_output(0), layer), _create_tensor( - layer.get_output(1), layer), _create_tensor(layer.get_output(2), layer) - - -def quantize(input: Tensor, - scale_factor: Tensor, - dtype: str, - axis: int = -1) -> Tensor: - layer = default_trtnet().add_quantize(input.trt_tensor, - scale_factor.trt_tensor, - str_dtype_to_trt(dtype)) - layer.axis = axis - - output = _create_tensor(layer.get_output(0), layer) - - return output - - -def dequantize(input: Tensor, - scale_factor: Tensor, - axis: int = -1, - output_type: Union[str, trt.DataType] = 'float16') -> Tensor: - - if isinstance(output_type, str): - output_type = str_dtype_to_trt(output_type) - - layer = default_trtnet().add_dequantize(input.trt_tensor, - scale_factor.trt_tensor, - output_type) - layer.axis = axis - - if not default_net().strongly_typed: - layer.precision = input.dtype - - output = _create_tensor(layer.get_output(0), layer) - - return output - - -def quantize_per_token( - x: Tensor, - clamp_val: Optional[Tensor] = None, - scale_dtype='float32', - sum_per_token: bool = False, - sum_dtype='float32', -) -> tuple[Tensor, Tensor] | tuple[Tensor, Tensor, Tensor]: - if not default_net().plugin_config.quantize_per_token_plugin: - x = cast(x, 'float32') - xmax = x.abs().max(-1, keepdim=True) - scales = xmax / 127.0 - out = x * 127.0 / xmax - out = round(out) - out = clip(out, -128, 127) - quantized = cast(out, 'int8') - if not sum_per_token: - return quantized, scales - sums = sum(x, -1, keepdim=True) - if sum_dtype is not None and str_dtype_to_trt(sum_dtype) != sums.dtype: - sums = cast(sums, sum_dtype) - return quantized, scales, sums - - plg_creator = trt.get_plugin_registry().get_plugin_creator( - 'QuantizePerToken', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert plg_creator is not None - - output_type = trt.PluginField("type_id", np.array([int(trt.int8)], - np.int32), - trt.PluginFieldType.INT32) - quant_mode = trt.PluginField( - "quant_mode", - np.array([int(QuantMode.use_smooth_quant(per_token=True))], np.int32), - trt.PluginFieldType.INT32) - clamp_enabled = trt.PluginField("clamp_enabled", - np.array([clamp_val is not None], np.int8), - trt.PluginFieldType.INT8) - - sum_per_token_pf = trt.PluginField("sum_per_token", - np.array([int(sum_per_token)], np.int32), - trt.PluginFieldType.INT32) - - pfc = trt.PluginFieldCollection( - [output_type, quant_mode, clamp_enabled, sum_per_token_pf]) - quantize_plug = plg_creator.create_plugin("quantize_per_token_plugin", pfc) - - plug_inputs = [x.trt_tensor] - if clamp_val: - plug_inputs += [clamp_val.trt_tensor] - layer = default_trtnet().add_plugin_v2(plug_inputs, quantize_plug) - if not default_net().strongly_typed: - layer.get_output(0).set_dynamic_range(-127, 127) - _add_plugin_info(layer, plg_creator, "quantize_per_token_plugin", pfc) - - quantized = _create_tensor(layer.get_output(0), layer) - scales = _create_tensor(layer.get_output(1), layer) - - # TODO: The plugin should be able to directly output float16 scales to avoid a cast - if scale_dtype is not None and str_dtype_to_trt( - scale_dtype) != scales.dtype: - scales = cast(scales, scale_dtype) - if not sum_per_token: - return quantized, scales - - sums = _create_tensor(layer.get_output(2), layer) - # TODO: The plugin should be able to directly output float16 sums to avoid a cast - if sum_dtype is not None and str_dtype_to_trt(sum_dtype) != sums.dtype: - sums = cast(sums, sum_dtype) - - return quantized, scales, sums - - -def quantize_fp8_per_token(x: Tensor, - clamp_val: Optional[Tensor] = None) -> Tuple[Tensor]: - if not default_net().plugin_config.quantize_per_token_plugin: - x = cast(x, 'float32') - xmax = x.abs().max(-1, keepdim=True) - scale = xmax / 448.0 - out = x * 448.0 / xmax - out = round(out) - out = clip(out, -448, 448) - quantized_out = cast(out, 'fp8') - return quantized_out, scale - else: - plg_creator = trt.get_plugin_registry().get_plugin_creator( - 'QuantizePerToken', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert plg_creator is not None - - output_type = trt.PluginField("type_id", - np.array([int(trt.fp8)], np.int32), - trt.PluginFieldType.INT32) - quant_mode = trt.PluginField( - "quant_mode", - np.array([int(QuantMode.from_description(use_fp8_rowwise=True))], - np.int32), trt.PluginFieldType.INT32) - clamp_enabled = trt.PluginField( - "clamp_enabled", np.array([clamp_val is not None], np.int8), - trt.PluginFieldType.INT8) - sum_per_token_pf = trt.PluginField("sum_per_token", - np.array([int(False)], np.int32), - trt.PluginFieldType.INT32) - pfc = trt.PluginFieldCollection( - [output_type, quant_mode, clamp_enabled, sum_per_token_pf]) - quantize_plug = plg_creator.create_plugin("quantize_per_token_plugin", - pfc) - - plug_inputs = [x.trt_tensor] - if clamp_val: - plug_inputs += [clamp_val.trt_tensor] - layer = default_trtnet().add_plugin_v2(plug_inputs, quantize_plug) - if not default_net().strongly_typed: - layer.get_output(0).set_dynamic_range(-448, 448) - _add_plugin_info(layer, plg_creator, "quantize_per_token_plugin", pfc) - - quantized = _create_tensor(layer.get_output(0), layer) - scales = _create_tensor(layer.get_output(1), layer) - - return quantized, scales - - -def quantize_tensor(x, scale): - if not default_net().plugin_config.quantize_tensor_plugin: - if scale.dtype == str_dtype_to_trt('float32'): - x = cast(x, 'float32') - scaled = x * scale - rounded = round(scaled) - clipped = clip(rounded, -128, 127) - quantized = cast(clipped, 'int8') - else: - scale = cast(scale, 'float32') - - plg_creator = trt.get_plugin_registry().get_plugin_creator( - 'QuantizeTensor', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert plg_creator is not None - - pfc = trt.PluginFieldCollection([]) - quantize_plug = plg_creator.create_plugin("quantize_tensor_plugin", pfc) - - plug_inputs = [x.trt_tensor, scale.trt_tensor] - layer = default_trtnet().add_plugin_v2(plug_inputs, quantize_plug) - if not default_net().strongly_typed: - layer.get_output(0).set_dynamic_range(-127, 127) - _add_plugin_info(layer, plg_creator, "quantize_tensor_plugin", pfc) - - quantized = _create_tensor(layer.get_output(0), layer) - return quantized - - -def symmetric_quantize_last_axis_of_batched_matrix(weight, quant_mode): - amax = weight.abs().max(dim=0)[0].to(weight.dtype) - if quant_mode == torch.int8: - scale = amax / 128. - qweight = torch.clamp((weight / scale).round(), -128, 127).char() - qweight = qweight.T.reshape(weight.shape) - else: - scale = amax / 8. - qweight = torch.clamp((weight / scale).round(), -8, 7).char() - qweight[qweight < 0] += 16 - qweight = qweight.T.view(torch.uint8) - qweight = (qweight[:, 1::2] * 16 + qweight[:, ::2]).view(torch.int8) - qweight = qweight.reshape(weight.shape[0], weight.shape[1] // 2) - return qweight, scale +from .._utils import get_sm_version def preprocess_weights_for_mixed_gemm( - tensor: torch.Tensor, - quant_mode: torch.dtype, - act_dtype: torch.dtype, - sm_: int = -1, - do_weight_interleave: bool = True) -> torch.Tensor: + tensor: torch.Tensor, + quant_mode: torch.dtype, + act_dtype: torch.dtype, + sm_: int = -1, + do_weight_interleave: bool = True, +) -> torch.Tensor: sm_ = sm_ if sm_ > 0 else get_sm_version() # 3-D inputs (MoE) on Hopper+ and any input on SM120/SM121 reuse the SM80 # interleaved layout. Check the original rank before unsqueeze. @@ -969,13 +38,73 @@ def preprocess_weights_for_mixed_gemm( permutation_map = { "16_8": [0, 1, 8, 9, 2, 3, 10, 11, 4, 5, 12, 13, 6, 7, 14, 15], "16_4": [ - 0, 1, 8, 9, 16, 17, 24, 25, 2, 3, 10, 11, 18, 19, 26, 27, 4, 5, 12, - 13, 20, 21, 28, 29, 6, 7, 14, 15, 22, 23, 30, 31 + 0, + 1, + 8, + 9, + 16, + 17, + 24, + 25, + 2, + 3, + 10, + 11, + 18, + 19, + 26, + 27, + 4, + 5, + 12, + 13, + 20, + 21, + 28, + 29, + 6, + 7, + 14, + 15, + 22, + 23, + 30, + 31, ], "8_4": [ - 0, 1, 2, 3, 16, 17, 18, 19, 4, 5, 6, 7, 20, 21, 22, 23, 8, 9, 10, - 11, 24, 25, 26, 27, 12, 13, 14, 15, 28, 29, 30, 31 - ] + 0, + 1, + 2, + 3, + 16, + 17, + 18, + 19, + 4, + 5, + 6, + 7, + 20, + 21, + 22, + 23, + 8, + 9, + 10, + 11, + 24, + 25, + 26, + 27, + 12, + 13, + 14, + 15, + 28, + 29, + 30, + 31, + ], } # permute_B_rows_for_mixed_gemm @@ -988,9 +117,9 @@ def preprocess_weights_for_mixed_gemm( num_rows = tensor.shape[1] num_cols = tensor.shape[2] - assert (sm_ >= 75) - assert (num_rows % B_ROWS_PER_MMA == 0) - assert (num_cols % MMA_SHAPE_N == 0) + assert sm_ >= 75 + assert num_rows % B_ROWS_PER_MMA == 0 + assert num_cols % MMA_SHAPE_N == 0 if do_weight_interleave and sm_ < 100: row_idx_list = [(row_idx // B_ROWS_PER_MMA) * B_ROWS_PER_MMA + @@ -1020,12 +149,16 @@ def preprocess_weights_for_mixed_gemm( rows_per_tile = 128 * 8 // BITS_PER_ELT_A elts_in_int32 = 32 // BITS_PER_ELT_B - assert (num_rows % elts_in_int32 == 0) - assert (num_rows % rows_per_tile == 0) + assert num_rows % elts_in_int32 == 0 + assert num_rows % rows_per_tile == 0 - tensor = tensor.reshape(num_experts, -1, interleave, - num_rows // rows_per_tile, - rows_per_tile * 4 // elts_in_int32) + tensor = tensor.reshape( + num_experts, + -1, + interleave, + num_rows // rows_per_tile, + rows_per_tile * 4 // elts_in_int32, + ) tensor = tensor.permute(0, 1, 3, 2, 4).reshape(original_shape) # add_bias_and_interleave_quantized_tensor_inplace @@ -1049,412 +182,3 @@ def preprocess_weights_for_mixed_gemm( raise NotImplementedError return tensor.squeeze(0).contiguous() - - -def get_weight_scale_interleave_factor(interleaved_dim: int, - group_size: int = 128) -> int: - # Calculate the weight_scale interleave factor for W4A8 groupwise MoE quant - # only Hopper w4a8 does interleave for weight scale, other arch or Hopper w4a16 default to 1 - factor = 1 - if get_sm_version() == 90: - if interleaved_dim % (4 * group_size) == 0: - factor = 4 - elif interleaved_dim % (2 * group_size) == 0: - factor = 2 - elif interleaved_dim % group_size == 0: - factor = 1 - else: - raise NotImplementedError( - f"Interleaved dimension must be a multiple of group_size ({group_size}), received {interleaved_dim}." - ) - return factor - - -def validate_group_size(layer): - # TODO: Remove this function and its usage after W4A8-AWQ with group_size = 64 is implemented. - W4A8_AWQ = 8 - if layer.quant_algo & W4A8_AWQ and layer.group_size == 64: - raise NotImplementedError( - "W4A8_AWQ with group_size = 64 is not implemented yet!") - - -def unpack_int32_into_int8(w_packed, autoawq_reorder=False): - # Unpack inputs packed in int32/float32 into uint4 and store them in int8 format - w_packed_int4x2 = w_packed.contiguous().view(torch.uint8) - w_unpacked = torch.zeros(w_packed_int4x2.shape[0], - w_packed_int4x2.shape[1] * 2, - dtype=torch.int8) - w_unpacked[:, ::2] = w_packed_int4x2 % 16 - w_unpacked[:, 1::2] = w_packed_int4x2 // 16 - if autoawq_reorder: - w_unpacked = w_unpacked.view(-1, 8)[:, [0, 4, 1, 5, 2, 6, 3, 7]].view( - w_unpacked.shape) - return w_unpacked.contiguous() - - -def change_qkv_leading_dim(w, num_heads): - if w.dim() == 1: - w = w.reshape(num_heads, 3, -1) - w = w.transpose(0, 1).reshape(-1) - else: - shape = w.shape - head_dim = shape[1] // (3 * num_heads) - w = w.reshape(-1, num_heads, 3, head_dim) - w = w.transpose(1, 2).reshape(shape[0], -1) - return w - - -def pad_like(w, target_shape, value=0): - if w.shape != target_shape: - pad_dim = [] - for dim in range(len(target_shape)): - current_dim = -1 - dim - pad_dim.append(0) - pad_dim.append( - max(0, target_shape[current_dim] - w.shape[current_dim])) - res = F.pad(w, pad_dim, value=value) - return res - else: - return w - - -def postprocess_weight_only(tllm_key, weights, quant_mode, layer): - if weights.dim() > 2: - v = weights.transpose(-1, -2) - else: - v = weights.t() - - tp_dim = 1 if isinstance(layer, ColumnLinear) else 0 - if "weight" in tllm_key: - if layer.is_padded: - split_size = layer.out_features if tp_dim == 1 else layer.in_features - v = torch.split(v, split_size, tp_dim)[layer.tp_rank] - v = pad_like(v, (layer.in_features, layer.out_features)) - processed_torch_weights, torch_weight_scales = \ - torch.ops.trtllm.symmetric_quantize_last_axis_of_batched_matrix( - v.contiguous(), quant_mode) - return { - tllm_key: processed_torch_weights, - tllm_key.replace("weight", "per_channel_scale"): - torch_weight_scales, - } - else: - if layer.is_padded and tp_dim == 1: - weights = torch.split(weights, layer.out_features, - tp_dim)[layer.tp_rank] - weights = pad_like(weights, (layer.out_features, )) - return {tllm_key: weights} # Bias - - -def postprocess_weight_only_groupwise(tllm_key, weights, torch_dtype, layer, - **kwargs): - using_head_as_leading_dim = kwargs.get("using_head_as_leading_dim", False) - config = kwargs.get("config", None) - use_autoawq = kwargs.get("use_autoawq", None) - num_heads = config.num_attention_heads - USE_GPTQ = layer.prequant_scaling_factor is None and use_autoawq is None - USE_HF_AWQ = layer.prequant_scaling_factor is None and use_autoawq is not None - USE_MODELOPT_AWQ = layer.prequant_scaling_factor is not None - USE_INT8_WEIGHT = layer.quant_algo & 16 - - tp_dim = 1 if isinstance(layer, ColumnLinear) else 0 - is_qkv = layer.is_qkv if hasattr(layer, "is_qkv") else False - - if using_head_as_leading_dim: - assert config.num_attention_heads == config.num_key_value_heads, "using_head_as_leading_dim require head_size to be multiple of 3." - if tllm_key.endswith("weights_scaling_factor"): - # TODO: Remove reshaping after modelopt optimizes scale shape - if is_qkv: - for idx, w in enumerate(weights): - scales = w.to(torch_dtype) - scales = scales.reshape(-1, - layer.weights_scaling_factor.shape[0]).T - scales = scales.chunk(layer.tp_size, 1)[layer.tp_rank] - weights[idx] = scales - weights = torch.cat(weights, dim=1) - else: - scales = weights.to(torch_dtype) - scales_shape = [ - layer.weights_scaling_factor.shape[1], - layer.weights_scaling_factor.shape[0] - ] - scales_shape[1 - tp_dim] *= layer.tp_size - scales = scales.reshape(scales_shape).T - weights = scales.chunk(layer.tp_size, tp_dim)[layer.tp_rank] - if is_qkv and isinstance(weights, list) and len(weights) >= 3: - if USE_MODELOPT_AWQ: - if tllm_key.endswith("prequant_scaling_factor"): - weights = weights[0] - else: - weights = torch.cat(weights, dim=0) - elif len(weights) > 3: - weights = [ - torch.cat(weights[i::len(weights) // 3], dim=1) - for i in range(len(weights) // 3) - ] - - if tllm_key.endswith("bias"): - if is_qkv and isinstance(weights, list): - weights = torch.cat(weights) - if layer.is_padded: - weights = pad_like(weights, layer.bias.shape) - if using_head_as_leading_dim: - weights = change_qkv_leading_dim(weights, num_heads) - results = {tllm_key: weights.to(torch_dtype)} - elif tllm_key.endswith("weight"): - if not USE_INT8_WEIGHT: - # 4 bit quantization - if USE_GPTQ: - qweight = unpack_int32_into_int8(weights[0].T).T - 8 - elif USE_HF_AWQ: - qweight = unpack_int32_into_int8(weights[0], True) - 8 - else: - qweight = unpack_int32_into_int8(weights.T) - qweight -= (qweight >> 4) << 4 - qweight = qweight.view(torch.uint8) - elif USE_INT8_WEIGHT and USE_GPTQ: - # 8 bit quantization (only consider INT8 GPTQ here) - qweight = ( - weights[0].T.contiguous().view(torch.uint8).T.contiguous() - - 128).to(torch.int8) - else: - raise NotImplementedError( - "Unsupported quantization mode for weight.") - - if using_head_as_leading_dim: - qweight = change_qkv_leading_dim(qweight, num_heads) - if layer.is_padded: - qweight = torch.split(qweight, layer.out_features, - tp_dim)[layer.tp_rank] - qweight = pad_like(qweight, (layer.in_features, layer.out_features)) - # pack int8 tensor to packed int4 - if not USE_INT8_WEIGHT: - qweight = (qweight[:, 1::2] * 16 + qweight[:, ::2]).view(torch.int8) - weight_type = torch.int8 if USE_INT8_WEIGHT else torch.quint4x2 - qweight = preprocess_weights_for_mixed_gemm( - qweight, weight_type, torch.float16).view(torch_dtype) - results = {tllm_key: qweight} - - # scales and zeros for GPTQ and HF-AWQ - if USE_GPTQ or USE_HF_AWQ: - scales = weights[1].to(torch_dtype) - if USE_INT8_WEIGHT: - qzeros = weights[2].view(torch.uint8) - else: - qzeros = unpack_int32_into_int8(weights[2], USE_HF_AWQ) - if using_head_as_leading_dim: - scales = change_qkv_leading_dim(scales, num_heads) - qzeros = change_qkv_leading_dim(qzeros, num_heads) - if layer.is_padded: - scales = torch.split(scales, - layer.weights_scaling_factor.shape[tp_dim], - tp_dim)[layer.tp_rank] - scales = pad_like(scales, layer.weights_scaling_factor.shape, 1) - qzeros = torch.split(qzeros, - layer.weights_scaling_factor.shape[tp_dim], - tp_dim)[layer.tp_rank] - qzeros = pad_like(qzeros, layer.zero.shape, 7) - if USE_INT8_WEIGHT: - zeros_x_scales = (-qzeros + 128 - 1 * USE_GPTQ) * scales - else: - zeros_x_scales = (-qzeros + 8 - 1 * USE_GPTQ) * scales - zeros_x_scales = zeros_x_scales.to(torch_dtype) - results.update({ - tllm_key.replace("weight", "weights_scaling_factor"): - scales, - tllm_key.replace("weight", "zero"): - zeros_x_scales, - }) - elif tllm_key.endswith("weights_scaling_factor"): - # TODO: Remove reshaping after modelopt optimizes scale shape - if layer.is_padded: - raise NotImplementedError( - "Auto-padding is not Implemented for ModelOpt HF-AWQ.") - results = {tllm_key: weights} - elif tllm_key.endswith("prequant_scaling_factor"): - prequant_scale = weights.to(torch_dtype).reshape(1, -1) - if layer.is_padded and tp_dim == 1: - prequant_scale = torch.split(prequant_scale, - layer.prequant_scaling_factor.shape[1], - 1)[layer.tp_rank] - prequant_scale = pad_like(prequant_scale, - layer.prequant_scaling_factor.shape, 0) - results = {tllm_key: prequant_scale} - - return results - - -def postprocess_fp8_rowwise(tllm_key, weights, **kwargs): - if tllm_key.endswith("per_channel_scale"): - return {} - - config = kwargs.get("config", None) - weights, scales = weights[0::2], weights[1::2] - - if scales[0] is not None: - assert all(w.dtype == torch.float8_e4m3fn for w in weights) - weights = torch.cat(weights, dim=0) - scales = torch.cat([s.to(torch.float32).flatten() for s in scales]) - return { - tllm_key: weights, - tllm_key.replace("weight", "per_channel_scale"): scales - } - else: - x = torch.cat(weights, dim=0).to(torch.float32) - clamp_val = config.quantization.clamp_val - if clamp_val is not None: - # activation range bound. - x = x.clamp(clamp_val[0], clamp_val[1]) - xmax = x.abs().max(-1, keepdim=True).values - # minimum scaling factor. - torch_weight_scales = (xmax / 448.0).clamp(min=1.0 / (448.0 * 512.0)) - out = x / torch_weight_scales - torch_weight_scales = torch_weight_scales.reshape(-1) - out = torch.clamp(out, -448, 448) - processed_torch_weights = out.to(torch.float8_e4m3fn) - processed_torch_weights = processed_torch_weights.to( - torch.float8_e4m3fn) - return { - tllm_key: processed_torch_weights, - tllm_key.replace("weight", "per_channel_scale"): torch_weight_scales - } - - -def fp4_gemm(input: Tensor, - input_sf: Tensor, - weight: Tensor, - weight_sf: Tensor, - global_sf: Tensor, - output_dtype: str | trt.DataType, - scaling_vector_size: int = 16): - ''' - Parameters: - input : Tensor (On GPU) - The input tensor. Its shape is [batch_size, seq_len, input_dim] or [num_tokens, input_dim] for remove_input_padding, should be fp4 - input_sf : Tensor (On GPU) - The input scaling factor tensor. Its shape is [batch_size, seq_len, input_dim / scaling_vector_size] or [num_tokens, input_dim / scaling_vector_size] for remove_input_padding, should be int32 (4 packed) - weight : Tensor (On GPU) - The weight tensor. Its shape is [output_dim, input_dim], should be fp4 - weight_sf : Tensor (On GPU) - The weight scaling factor tensor. Its shape is [output_dim, input_dim / scaling_vector_size], should be fp8 - global_sf : Tensor (On GPU) - The global scaling factor tensor. Its shape is [1,], should be float32, used as alpha of Gemm. - output_dtype: str - output data type - scaling_vector_size: int - scaling vector block size - ''' - if isinstance(output_dtype, str): - output_dtype = str_dtype_to_trt(output_dtype) - - fp4_gemm_plg_creator = trt.get_plugin_registry().get_plugin_creator( - 'Fp4Gemm', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert fp4_gemm_plg_creator is not None - sv_vec_size = trt.PluginField("sv_vec_size", - np.array(scaling_vector_size, dtype=np.int32), - trt.PluginFieldType.INT32) - output_dtype = trt.PluginField("output_type_id", - np.array([int(output_dtype)], np.int32), - trt.PluginFieldType.INT32) - pfc = trt.PluginFieldCollection([sv_vec_size, output_dtype]) - fp4_gemm_plug = fp4_gemm_plg_creator.create_plugin("fp4_gemm", pfc) - plug_inputs = [input, input_sf, weight, weight_sf, global_sf] - plug_inputs = [i.trt_tensor for i in plug_inputs] - layer = default_trtnet().add_plugin_v2(plug_inputs, fp4_gemm_plug) - _add_plugin_info(layer, fp4_gemm_plg_creator, "fp4_gemm", pfc) - output = _create_tensor(layer.get_output(0), layer) - return output - - -def quantize_to_fp4_tensor(input: Tensor, sf_scale: Tensor): - ''' - Parameters: - input : Tensor (On GPU) - The input tensor. Its shape is [batch_size, seq_len, input_dim] or [num_tokens, input_dim] for remove_input_padding, should be fp16 - sf_scale : Tensor (On GPU) - The global per-tensor scaling factor. Its shape is [1,], should be float32. - used to scale SF from input range to fp8 range (448.f / (MaxVal of input / 6.f)). - output : Tensor (On GPU) - The output tensor. Its shape is [batch_size, seq_len, input_dim] or [num_tokens, input_dim] for remove_input_padding, should be FP4 - output_sf : Tensor (On GPU) - The input scaling factor tensor. Its shape is [batch_size, seq_len, input_dim / scaling_vector_size] or [num_tokens, input_dim / scaling_vector_size] for remove_input_padding, should be FP8 - ''' - plg_creator = trt.get_plugin_registry().get_plugin_creator( - 'QuantizeToFP4', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert plg_creator is not None - - pfc = trt.PluginFieldCollection([]) - quantize_plug = plg_creator.create_plugin("quantize_to_fp4_plugin", pfc) - - plug_inputs = [input.trt_tensor, sf_scale.trt_tensor] - layer = default_trtnet().add_plugin_v2(plug_inputs, quantize_plug) - _add_plugin_info(layer, plg_creator, "quantize_to_fp4_plugin", pfc) - - quantized = _create_tensor(layer.get_output(0), layer) - scales = _create_tensor(layer.get_output(1), layer) - return quantized, scales - - -def dynamic_quantize( - x: Tensor, - double_scale: Tensor, - axis: int = -1, - block_size: int = 16, - data_qtype: trt.DataType = trt.fp4, - scale_qtype: trt.DataType = trt.fp8) -> Tuple[Tensor, Tensor]: - ''' - Parameters: - x : Tensor (On GPU) - The input tensor. - double_scale : Tensor (On GPU) - The global per-tensor scaling factor. It should contain only 1 element. - axis : int - The axis to quantize. Default is -1 (the last axis). - block_size : int - The block size for quantization. Default is 16. - data_qtype : trt.DataType - The data type for quantized data. Default is FP4. - scale_qtype : trt.DataType - The data type for block scale. Default is FP8. - Returns: - A tuple of two tensors: quantized tensor and block scale tensor. - ''' - if axis < 0: - axis = len(x.shape) + axis - dynq = default_trtnet().add_dynamic_quantize(x.trt_tensor, axis, block_size, - data_qtype, scale_qtype) - dynq.set_input(1, double_scale.trt_tensor) - quantized = _create_tensor(dynq.get_output(0), dynq) - scale = _create_tensor(dynq.get_output(1), dynq) - return quantized, scale - - -def block_double_dequantize(x: Tensor, - scale: Tensor, - double_scale: Tensor, - dtype: trt.DataType | str = 'float16') -> Tensor: - ''' - Parameters: - x : Tensor (On GPU) - The input tensor. - scale : Tensor (On GPU) - The block scale tensor. - double_scale : Tensor (On GPU) - The global per-tensor scaling factor. It should contain only 1 element. - dtype : trt.DataType | str - The data type for dequantized data. Default is float32. - Returns: - The dequantized tensor. - ''' - if isinstance(dtype, str): - dtype = str_dtype_to_trt(dtype) - dequantize_scale_layer = default_trtnet().add_dequantize( - scale.trt_tensor, double_scale.trt_tensor, dtype) - scale = _create_tensor(dequantize_scale_layer.get_output(0), - dequantize_scale_layer) - - dequantize_data_layer = default_trtnet().add_dequantize( - x.trt_tensor, scale.trt_tensor, dtype) - dequantize_data = _create_tensor(dequantize_data_layer.get_output(0), - dequantize_data_layer) - return dequantize_data diff --git a/tensorrt_llm/quantization/layers.py b/tensorrt_llm/quantization/layers.py deleted file mode 100644 index 28f4cbf2e8cf..000000000000 --- a/tensorrt_llm/quantization/layers.py +++ /dev/null @@ -1,3383 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -import math -from typing import Optional - -import numpy as np -import tensorrt as trt -import torch - -from .._common import default_net, precision -from .._utils import is_same_dtype, str_dtype_to_torch, trt_dtype_to_torch -from ..functional import (ACT2FN, AllReduceFusionOp, AllReduceParams, - AttentionMaskType, PositionEmbeddingType, - RotaryScalingType, Tensor, allgather, allreduce, cast, - concat, constant, div, embedding, gemm_allreduce, - generate_alibi_slopes, gpt_attention, matmul, mul, - shape, slice, softmax, split, where) -from ..layers import MropeParams, SpecDecodingParams -from ..layers.embedding import Embedding -from ..layers.linear import Linear, RowLinear -from ..module import Module -from ..parameter import Parameter -from .utils import fp4_utils - -# isort: off -from .functional import ( - block_double_dequantize, dequantize, dynamic_quantize, fp4_gemm, - quantize_to_fp4_tensor, fp8_rowwise_gemm, fp8_rowwise_rms_norm, - postprocess_fp8_rowwise, postprocess_weight_only, - postprocess_weight_only_groupwise, quantize, quantize_fp8_per_token, - quantize_per_token, quantize_tensor, validate_group_size, smooth_quant_gemm, - smooth_quant_layer_norm, smooth_quant_rms_norm, - weight_only_groupwise_quant_matmul, weight_only_quant_matmul, - qserve_gemm_per_group, qserve_gemm_per_channel, fp8_rowwise_layer_norm) -# isort: on -from .mode import GroupwiseQuantAlgo, QuantMode - - -class Quantize(Module): - """ - Quantize Layer - For per-tensor mode, the scaling factor is a scalar. - For per-channel mode, the scaling factor is a vector. - """ - - def __init__( - self, - output_dtype: str = 'int8', - scaling_factor_dtype: str = 'float32', - in_channels: int = -1, - axis=-1, - ) -> None: - super().__init__() - self.scaling_factor = Parameter(shape=(in_channels, ) if axis != -1 else - (), - dtype=scaling_factor_dtype) - self.output_dtype = output_dtype - self.axis = axis - - def forward(self, x): - return quantize(x, self.scaling_factor.value, self.output_dtype, - self.axis) - - -class QuantizePerToken(Module): - """ - Quantize Per Token and compute dynamic scales for SmoothQuant - """ - - def forward(self, x): - return quantize_per_token(x) - - -class Dequantize(Module): - """ - Dequantize Layer. - """ - - def __init__(self, axis: int = -1) -> None: - super().__init__() - self.scaling_factor = Parameter(shape=(), dtype='float32') - self.axis = axis - - def forward(self, input): - return dequantize(input, self.scaling_factor.value, self.axis) - - -class SmoothQuantLinear(Linear): - - def __init__(self, - in_features, - out_features, - bias=True, - dtype=None, - tp_group=None, - tp_size=1, - gather_output=True, - quant_mode=QuantMode(0), - prefer_managed_weight=True): - super().__init__(in_features, - out_features, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - gather_output=gather_output, - prefer_managed_weight=prefer_managed_weight) - - if not quant_mode.has_act_and_weight_quant(): - raise ValueError( - "SmoothQuant Linear has to have act+weight quantization mode set" - ) - - weights_dtype = dtype - if quant_mode.has_act_and_weight_quant(): - weights_dtype = "int8" - - self.weight = Parameter(shape=(self.out_features, self.in_features), - dtype=weights_dtype, - prefer_managed=self.prefer_managed_weight) - - if quant_mode.has_act_and_weight_quant(): - scale_shape = (1, self.out_features - ) if quant_mode.has_per_channel_scaling() else (1, 1) - self.per_channel_scale = Parameter(shape=scale_shape, - dtype="float32") - - if quant_mode.has_act_static_scaling(): - self.act_scale = Parameter(shape=(1, 1), dtype="float32") - - self.quant_mode = quant_mode - - def forward(self, x, lora_runtime_params=None): - assert lora_runtime_params is None, "lora is not supported on SmoothQuantLinear now" - if self.quant_mode.has_act_static_scaling(): - per_token_scale = self.act_scale.value - else: - # If we are in SmoothQuant with dynamic activation scaling, - # input x has to be a tuple of int8 tensor and fp32 scaling factors - x, per_token_scale = x - x = smooth_quant_gemm(x, self.weight.value, per_token_scale, - self.per_channel_scale.value, - self.quant_mode.has_per_token_dynamic_scaling(), - self.quant_mode.has_per_channel_scaling(), - self.dtype) - - if self.bias is not None: - x = x + self.bias.value - - if self.gather_output and self.tp_size > 1 and self.tp_group is not None: - # [dim0, local_dim] -> [dim0 * tp_size, local_dim] --> [dim0, local_dim * tp_size] - x = allgather(x, self.tp_group, gather_dim=1) - - return x - - -SmoothQuantColumnLinear = SmoothQuantLinear - - -class SmoothQuantRowLinear(RowLinear): - - def __init__( - self, - in_features, - out_features, - bias=True, - dtype=None, - tp_group=None, - tp_size=1, - quant_mode=QuantMode(0), - prefer_managed_weight=True, - ): - super().__init__(in_features, - out_features, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - prefer_managed_weight=prefer_managed_weight) - if not quant_mode.has_act_and_weight_quant(): - raise ValueError( - "SmoothQuant Linear has to have act+weight quantization mode set" - ) - weights_dtype = dtype - if quant_mode.has_act_and_weight_quant(): - weights_dtype = "int8" - - self.weight = Parameter(shape=(self.out_features, self.in_features), - dtype=weights_dtype, - prefer_managed=self.prefer_managed_weight) - self.smoother = Parameter(shape=(1, self.in_features), dtype="float32") - if quant_mode.has_act_and_weight_quant(): - scale_shape = (1, self.out_features - ) if quant_mode.has_per_channel_scaling() else (1, 1) - self.per_channel_scale = Parameter(shape=scale_shape, - dtype="float32") - - if quant_mode.has_act_static_scaling(): - self.act_scale = Parameter(shape=(1, 1), dtype="float32") - - self.quant_mode = quant_mode - - def forward(self, x, lora_runtime_params=None, all_reduce_params=None): - assert lora_runtime_params is None, "lora is not supported on SmoothQuantRowLinear now" - if self.quant_mode.has_act_static_scaling(): - per_token_scale = self.act_scale.value - else: - x, per_token_scale = x - x = smooth_quant_gemm(x, self.weight.value, per_token_scale, - self.per_channel_scale.value, - self.quant_mode.has_per_token_dynamic_scaling(), - self.quant_mode.has_per_channel_scaling(), - self.dtype) - - if self.tp_size > 1 and self.tp_group is not None: - need_bias = self.bias is not None - fuse_bias_into_all_reduce = need_bias and ( - all_reduce_params - is not None) and (all_reduce_params.fusion_op - == AllReduceFusionOp.RESIDUAL_RMS_NORM) - if fuse_bias_into_all_reduce: - all_reduce_params.bias = self.bias.value - x = allreduce(x, self.tp_group, all_reduce_params=all_reduce_params) - if need_bias and not fuse_bias_into_all_reduce: - x = x + self.bias.value - return x - - if self.bias is not None: - x = x + self.bias.value - - return x - - -class SmoothQuantLayerNorm(Module): - - def __init__( - self, - normalized_shape, - eps=1e-05, - elementwise_affine=True, - bias=True, - dtype=None, - quant_mode=QuantMode(0), - ): - super().__init__() - if isinstance(normalized_shape, int): - normalized_shape = (normalized_shape, ) - if not quant_mode.has_act_and_weight_quant(): - raise ValueError( - "SmoothQuant layer norm has to have some quantization mode set") - self.normalized_shape = tuple(normalized_shape) - self.elementwise_affine = elementwise_affine - if self.elementwise_affine: - self.weight = Parameter(shape=self.normalized_shape, dtype=dtype) - if bias: - self.bias = Parameter(shape=self.normalized_shape, dtype=dtype) - else: - self.register_parameter('bias', None) - else: - self.register_parameter('weight', None) - self.register_parameter('bias', None) - - self.eps = eps - self.dtype = dtype - self.quant_mode = quant_mode - - if self.quant_mode.has_act_and_weight_quant(): - self.scale_to_int = Parameter(shape=(1, ), dtype=dtype) - else: - self.register_parameter('scale_to_int', None) - - def forward(self, x): - weight = None if self.weight is None else self.weight.value - bias = None if self.bias is None else self.bias.value - scale = None if self.scale_to_int is None else self.scale_to_int.value - return smooth_quant_layer_norm( - x, - self.normalized_shape, - weight, - bias, - scale, - self.eps, - dynamic_act_scaling=self.quant_mode.has_per_token_dynamic_scaling()) - - -class SmoothQuantRmsNorm(Module): - - def __init__( - self, - normalized_shape, - eps=1e-06, - elementwise_affine=True, - dtype=None, - quant_mode=QuantMode(0), - bias=False, - clamp_val=None, - ): - super().__init__() - if isinstance(normalized_shape, int): - normalized_shape = (normalized_shape, ) - if not quant_mode.has_act_and_weight_quant(): - raise ValueError( - "SmoothQuant Rms norm has to have some quantization mode set") - self.normalized_shape = tuple(normalized_shape) - self.elementwise_affine = elementwise_affine - if self.elementwise_affine: - self.weight = Parameter(shape=self.normalized_shape, dtype=dtype) - else: - self.register_parameter('weight', None) - - if bias: - self.bias = Parameter(shape=self.normalized_shape, dtype=dtype) - else: - self.register_parameter('bias', None) - if clamp_val: - if not (isinstance(clamp_val, list) and len(clamp_val) == 2): - raise ValueError(f'unsupported clamp_val {clamp_val}') - self.clamp_val = Parameter(np.array(clamp_val, dtype=np.float32), - dtype='float32', - is_buffer=True) - else: - self.register_parameter('clamp_val', None) - - self.eps = eps - self.dtype = dtype - self.quant_mode = quant_mode - - if self.quant_mode.has_act_and_weight_quant(): - self.scale_to_int = Parameter(shape=(1, ), dtype=dtype) - else: - self.register_parameter('scale_to_int', None) - - def forward(self, x): - weight = None if self.weight is None else self.weight.value - bias = None if self.bias is None else self.bias.value - scale = None if self.scale_to_int is None else self.scale_to_int.value - clamp_val = None if self.clamp_val is None else self.clamp_val.value - return smooth_quant_rms_norm( - x, - self.normalized_shape, - weight, - bias, - scale, - clamp_val, - self.eps, - dynamic_act_scaling=self.quant_mode.has_per_token_dynamic_scaling()) - - -class QServeW4A8Linear(Linear): - - def __init__(self, - in_features, - out_features, - bias=True, - dtype=None, - tp_group=None, - tp_size=1, - gather_output=True, - quant_mode=QuantMode(0)): - assert dtype == "float16" # Currently the kernel only supports float16 output - - super().__init__(in_features, - out_features, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - gather_output=gather_output) - - self.quant_mode = quant_mode - assert self.quant_mode.is_qserve_w4a8() - - # Only support g128 now. - if self.quant_mode.has_per_group_scaling(): - self.group_size = 128 - else: - self.group_size = -1 - - self.weight = Parameter(shape=(self.out_features, - self.in_features // 2), - dtype="int8") - - self.s1_scales = Parameter(shape=(self.out_features, ), dtype="float16") - - if self.group_size == -1: - self.s1_szeros = Parameter(shape=(self.out_features, ), - dtype="float16") - else: - self.s2_scales = Parameter( - shape=(self.in_features // self.group_size, self.out_features), - dtype="int8") - self.s2_szeros = Parameter( - shape=(self.in_features // self.group_size, self.out_features), - dtype="int8") - - def forward(self, x): - if self.group_size == -1: - x, per_token_scale, per_token_sum = x - x = qserve_gemm_per_channel(x, per_token_scale, per_token_sum, - self.weight.value, self.s1_scales.value, - self.s1_szeros.value) - - else: - x, per_token_scale = x - x = qserve_gemm_per_group(x, per_token_scale, self.weight.value, - self.s1_scales.value, - self.s2_scales.value, - self.s2_szeros.value, self.group_size) - - if self.bias is not None: - x = x + self.bias.value - - if self.gather_output and self.tp_size > 1 and self.tp_group is not None: - # [dim0, local_dim] -> [dim0 * tp_size, local_dim] --> [dim0, local_dim * tp_size] - x = allgather(x, self.tp_group, gather_dim=1) - return x - - -QServeW4A8ColumnLinear = QServeW4A8Linear - - -class QServeW4A8RowLinear(RowLinear): - - def __init__( - self, - in_features, - out_features, - bias=True, - dtype=None, - tp_group=None, - tp_size=1, - quant_mode=QuantMode(0), - ): - assert dtype == "float16" # Currently the kernel only supports float16 output - - super().__init__(in_features, - out_features, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size) - - self.quant_mode = quant_mode - assert self.quant_mode.is_qserve_w4a8() - # Only supports 128g now. - if self.quant_mode.has_per_group_scaling(): - self.group_size = 128 - else: - self.group_size = -1 - - self.weight = Parameter(shape=(self.out_features, - self.in_features // 2), - dtype="int8") - - self.s1_scales = Parameter(shape=(self.out_features, ), dtype="float16") - - if self.group_size == -1: - self.s1_szeros = Parameter(shape=(self.out_features, ), - dtype="float16") - else: - self.s2_scales = Parameter( - shape=(self.in_features // self.group_size, self.out_features), - dtype="int8") - self.s2_szeros = Parameter( - shape=(self.in_features // self.group_size, self.out_features), - dtype="int8") - - def forward(self, x, all_reduce_params=None): - if self.group_size == -1: - x, per_token_scale, per_token_sum = x - x = qserve_gemm_per_channel(x, per_token_scale, per_token_sum, - self.weight.value, self.s1_scales.value, - self.s1_szeros.value) - else: - x, per_token_scale = x - x = qserve_gemm_per_group(x, per_token_scale, self.weight.value, - self.s1_scales.value, - self.s2_scales.value, - self.s2_szeros.value, self.group_size) - - if self.tp_size > 1 and self.tp_group is not None: - need_bias = self.bias is not None - fuse_bias_into_all_reduce = need_bias and ( - all_reduce_params - is not None) and (all_reduce_params.fusion_op - == AllReduceFusionOp.RESIDUAL_RMS_NORM) - if fuse_bias_into_all_reduce: - all_reduce_params.bias = self.bias.value - x = allreduce(x, self.tp_group, all_reduce_params=all_reduce_params) - if need_bias and not fuse_bias_into_all_reduce: - x = x + self.bias.value - return x - - if self.bias is not None: - x = x + self.bias.value - - return x - - -class Fp8RowwiseRmsNorm(Module): - - def __init__( - self, - normalized_shape, - eps=1e-06, - elementwise_affine=True, - dtype=None, - quant_mode=QuantMode(0), - bias=False, - clamp_val=None, - ): - super().__init__() - if isinstance(normalized_shape, int): - normalized_shape = (normalized_shape, ) - if not quant_mode.has_fp8_rowwise(): - raise ValueError( - "Fp8 Rowwise Rms norm has to have some quantization mode set") - self.normalized_shape = tuple(normalized_shape) - self.elementwise_affine = elementwise_affine - if self.elementwise_affine: - self.weight = Parameter(shape=self.normalized_shape, dtype=dtype) - else: - self.register_parameter('weight', None) - - if bias: - self.bias = Parameter(shape=self.normalized_shape, dtype=dtype) - else: - self.register_parameter('bias', None) - - if clamp_val: - if not (isinstance(clamp_val, list) and len(clamp_val) == 2): - raise ValueError(f'unsupported clamp_val {clamp_val}') - self.clamp_val = Parameter(np.array(clamp_val, dtype=np.float32), - dtype='float32', - is_buffer=True) - else: - self.register_parameter('clamp_val', None) - - self.eps = eps - self.dtype = dtype - self.quant_mode = quant_mode - - def forward(self, x): - weight = None if self.weight is None else self.weight.value - bias = None if self.bias is None else self.bias.value - scale = None - clamp_val = None if self.clamp_val is None else self.clamp_val.value - return fp8_rowwise_rms_norm( - x, - self.normalized_shape, - weight, - bias, - scale, - clamp_val, - self.eps, - dynamic_act_scaling=self.quant_mode.has_fp8_rowwise()) - - -class Fp8RowwiseLinear(Linear): - - def __init__(self, - in_features, - out_features, - bias=True, - dtype=None, - tp_group=None, - tp_size=1, - gather_output=True, - quant_mode=QuantMode(0)): - super().__init__(in_features, - out_features, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - gather_output=gather_output) - - if not quant_mode.has_fp8_rowwise(): - raise ValueError( - "Fp8 Rowwise Linear has to have act+weight quantization mode set" - ) - - weights_dtype = dtype - if quant_mode.has_fp8_rowwise(): - weights_dtype = "fp8" - - self.weight = Parameter(shape=(self.out_features, self.in_features), - dtype=weights_dtype) - - if quant_mode.has_fp8_rowwise(): - self.per_channel_scale = Parameter(shape=(self.out_features, ), - dtype="float32") - - self.quant_mode = quant_mode - self.tllm_to_externel_key_dict = {"weight": ["weight", "weight_scale"]} - - def forward(self, x, lora_runtime_params=None): - assert lora_runtime_params is None, "lora is not supported on SmoothQuantLinear now" - x, per_token_scale = x - x = fp8_rowwise_gemm(x, self.weight.value, per_token_scale, - self.per_channel_scale.value, - self.quant_mode.has_per_token_dynamic_scaling(), - self.quant_mode.has_per_channel_scaling()) - - if self.bias is not None: - x = x + self.bias.value - - if self.gather_output and self.tp_size > 1 and self.tp_group is not None: - # [dim0, local_dim] -> [dim0 * tp_size, local_dim] --> [dim0, local_dim * tp_size] - x = allgather(x, self.tp_group, gather_dim=1) - - return x - - def postprocess(self, tllm_key, weights, **kwargs): - if "per_channel_scale" in tllm_key: - return {} - return postprocess_fp8_rowwise(tllm_key, weights, **kwargs) - - -Fp8RowwiseColumnLinear = Fp8RowwiseLinear - - -class Fp8RowwiseRowLinear(RowLinear): - - def __init__( - self, - in_features, - out_features, - bias=True, - dtype=None, - tp_group=None, - tp_size=1, - quant_mode=QuantMode(0), - ): - super().__init__(in_features, - out_features, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size) - if not quant_mode.has_fp8_rowwise(): - raise ValueError( - "Fp8 Rowwise Linear has to have act+weight quantization mode set" - ) - weights_dtype = dtype - if quant_mode.has_fp8_rowwise(): - weights_dtype = "fp8" - - self.weight = Parameter(shape=(self.out_features, self.in_features), - dtype=weights_dtype) - if quant_mode.has_fp8_rowwise(): - self.per_channel_scale = Parameter(shape=(self.out_features, ), - dtype="float32") - - self.quant_mode = quant_mode - self.tllm_to_externel_key_dict = {"weight": ["weight", "weight_scale"]} - - def forward(self, x, lora_runtime_params=None, all_reduce_params=None): - assert lora_runtime_params is None, "lora is not supported on SmoothQuantRowLinear now" - x, per_token_scale = x - x = fp8_rowwise_gemm(x, self.weight.value, per_token_scale, - self.per_channel_scale.value, - self.quant_mode.has_fp8_rowwise(), - self.quant_mode.has_fp8_rowwise()) - - if self.tp_size > 1 and self.tp_group is not None: - need_bias = self.bias is not None - fuse_bias_into_all_reduce = need_bias and ( - all_reduce_params - is not None) and (all_reduce_params.fusion_op - == AllReduceFusionOp.RESIDUAL_RMS_NORM) - if fuse_bias_into_all_reduce: - all_reduce_params.bias = self.bias.value - x = allreduce(x, self.tp_group, all_reduce_params=all_reduce_params) - if need_bias and not fuse_bias_into_all_reduce: - x = x + self.bias.value - return x - - if self.bias is not None: - x = x + self.bias.value - - return x - - def postprocess(self, tllm_key, weights, **kwargs): - if "per_channel_scale" in tllm_key: - return {} - return postprocess_fp8_rowwise(tllm_key, weights, **kwargs) - - -class WeightOnlyQuantLinear(Linear): - - def __init__( - self, - in_features, - out_features, - bias=True, - dtype=None, - tp_group=None, - tp_size=1, - tp_rank=0, - gather_output=True, - quant_mode=QuantMode.use_weight_only(), - transa=False, - transb=False, - is_qkv=False, - prefer_managed_weight=True, - ): - multiple = 64 * tp_size - self.is_padded = False - if in_features % multiple > 0: - in_features = math.ceil(in_features / multiple) * multiple - self.is_padded = True - if out_features % multiple > 0: - out_features = math.ceil(out_features / multiple) * multiple - self.is_padded = True - - super().__init__(in_features, - out_features, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - gather_output=gather_output, - is_qkv=is_qkv, - prefer_managed_weight=prefer_managed_weight) - if quant_mode.is_int8_weight_only(): - self.weight_only_quant_mode = 1 - quant_type_size_in_bits = 8 - elif quant_mode.is_int4_weight_only(): - self.weight_only_quant_mode = 2 - quant_type_size_in_bits = 4 - # we use a fake tensor with data_type = int8 - self.weight = Parameter(shape=(self.in_features, - int(self.out_features * - quant_type_size_in_bits / 8)), - dtype="int8", - prefer_managed=self.prefer_managed_weight) - - scale_shape = (self.out_features, ) - self.per_channel_scale = Parameter(shape=scale_shape, dtype=dtype) - - self.transa = transa - self.transb = transb - self.tp_rank = tp_rank - if self.is_padded: - self.tp_dim = -1 - self.quant_mode = quant_mode - - def forward(self, x, lora_runtime_params=None): - # ootb has not supported int4 yet. - if self.weight_only_quant_mode == 2 and not default_net( - ).plugin_config.weight_only_quant_matmul_plugin: - raise TypeError( - "Int4 Weight-only Quant MatMul is only supported with plugin") - hidden_state = x - x = weight_only_quant_matmul(x, self.weight.value, - self.per_channel_scale.value, - self.weight_only_quant_mode, self.dtype, - self.transa, self.transb) - - if default_net( - ).plugin_config.lora_plugin and lora_runtime_params is not None: - x = x + self.lora(hidden_state, - lora_runtime_params=lora_runtime_params) - - if self.bias is not None: - x = x + self.bias.value - - if self.gather_output and self.tp_size > 1 and self.tp_group is not None: - # [dim0, local_dim] -> [dim0 * tp_size, local_dim] --> [dim0, local_dim * tp_size] - x = allgather(x, self.tp_group, gather_dim=1) - - return x - - def postprocess(self, tllm_key, weights, **kwargs): - if "per_channel_scale" in tllm_key: - return {} - weights = super().postprocess(tllm_key, weights, **kwargs)[tllm_key] - weights = weights.to(str_dtype_to_torch(self.dtype)) - return postprocess_weight_only( - tllm_key, weights, - torch.int8 if self.weight_only_quant_mode == 1 else torch.quint4x2, - self) - - -WeightOnlyQuantColumnLinear = WeightOnlyQuantLinear - - -class WeightOnlyQuantRowLinear(RowLinear): - - def __init__( - self, - in_features, - out_features, - bias=True, - dtype=None, - tp_group=None, - tp_size=1, - tp_rank=0, - quant_mode=QuantMode.use_weight_only(), - prefer_managed_weight=True, - is_expert=False, - ): - multiple = 64 * tp_size - self.is_padded = False - if in_features % multiple > 0: - in_features = math.ceil(in_features / multiple) * multiple - self.is_padded = True - if out_features % multiple > 0: - out_features = math.ceil(out_features / multiple) * multiple - self.is_padded = True - - super().__init__(in_features, - out_features, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - prefer_managed_weight=prefer_managed_weight, - is_expert=is_expert) - if quant_mode.is_int8_weight_only(): - self.weight_only_quant_mode = 1 - elif quant_mode.is_int4_weight_only(): - self.weight_only_quant_mode = 2 - #we use a fake tensor with data_type = int8 - self.weight = Parameter(shape=(self.in_features, - int(self.out_features / - self.weight_only_quant_mode)), - dtype="int8", - prefer_managed=prefer_managed_weight) - self.per_channel_scale = Parameter(shape=(self.out_features, ), - dtype=dtype) - self.tp_rank = tp_rank - if self.is_padded: - self.tp_dim = -1 - self.quant_mode = quant_mode - - def forward(self, x, lora_runtime_params=None, all_reduce_params=None): - hidden_state = x - x = weight_only_quant_matmul(x, self.weight.value, - self.per_channel_scale.value, - self.weight_only_quant_mode, self.dtype) - - if default_net( - ).plugin_config.lora_plugin and lora_runtime_params is not None: - x = x + self.lora(hidden_state, - lora_runtime_params=lora_runtime_params) - - if self.tp_size > 1 and self.tp_group is not None: - need_bias = self.bias is not None - fuse_bias_into_all_reduce = need_bias and ( - all_reduce_params - is not None) and (all_reduce_params.fusion_op - == AllReduceFusionOp.RESIDUAL_RMS_NORM) - if fuse_bias_into_all_reduce: - all_reduce_params.bias = self.bias.value - if not self.is_expert: - x = allreduce(x, - self.tp_group, - all_reduce_params=all_reduce_params) - if need_bias and not fuse_bias_into_all_reduce: - bias = cast(self.bias.value, x.dtype) - x = x + bias - else: - if need_bias and not fuse_bias_into_all_reduce: - bias = cast(self.bias.value, x.dtype) - x = x + bias / self.tp_size - return x - - if self.bias is not None: - x = x + self.bias.value - - return x - - def postprocess(self, tllm_key, weights, **kwargs): - if "per_channel_scale" in tllm_key: - return {} - weights = weights.to(str_dtype_to_torch(self.dtype)) - return postprocess_weight_only( - tllm_key, weights, - torch.int8 if self.weight_only_quant_mode == 1 else torch.quint4x2, - self) - - -class WeightOnlyQuantEmbedding(Embedding): - - def __init__( - self, - num_embeddings: int, - embedding_dim: int, - dtype: Optional[str] = None, - tp_size: int = 1, - tp_group: Optional[list] = None, - sharding_dim: int = 0, - tp_rank: Optional[int] = None, - quant_mode=QuantMode.use_weight_only(), - ): - super().__init__( - num_embeddings, - embedding_dim, - dtype, # dtype, - tp_size, - tp_group, - sharding_dim, - tp_rank) - # only support int8 wo now - # TODO support int4 wo - self.quant_mode = quant_mode - self.per_token_scale = Parameter(shape=(self.num_embeddings, ), - dtype=dtype) - - if sharding_dim == 1: - self.weight = Parameter(shape=(self.num_embeddings, - self.embedding_dim // self.tp_size), - dtype="int8") - elif sharding_dim == 0: - self.weight = Parameter(shape=(math.ceil( - self.num_embeddings / self.tp_size), self.embedding_dim), - dtype="int8") - - def forward(self, x): - result = embedding(x, - self.weight.value, - tp_size=self.tp_size, - tp_group=self.tp_group, - sharding_dim=self.sharding_dim, - tp_rank=self.tp_rank, - per_token_scale=self.per_token_scale.value) - - return result - - -def unpack_int32_into_int8(w_packed): - # Unpack inputs packed in int32/float32 into uint4 and store them in int8 format - w_packed_int4x2 = w_packed.contiguous().view(torch.uint8) - w_unpacked = torch.zeros(w_packed_int4x2.shape[0], - w_packed_int4x2.shape[1] * 2, - dtype=torch.int8) - w_unpacked[:, ::2] = w_packed_int4x2 % 16 - w_unpacked[:, 1::2] = w_packed_int4x2 // 16 - w_unpacked = w_unpacked.view(-1, 8)[:, [0, 4, 1, 5, 2, 6, 3, 7]].view( - w_unpacked.shape) - return w_unpacked.contiguous() - - -class WeightOnlyGroupwiseQuantLinear(Linear): - - def __init__( - self, - in_features, - out_features, - group_size=128, - pre_quant_scale=False, - zero=False, - bias=False, - dtype=None, - tp_group=None, - tp_size=1, - tp_rank=0, - gather_output=True, - use_w4a8_awq=False, - use_int8_weight=False, - is_qkv=False, - prefer_managed_weight=True, - ): - multiple = max((128 if use_w4a8_awq else 64), group_size) * tp_size - self.is_padded = False - if in_features % multiple > 0: - in_features = math.ceil(in_features / multiple) * multiple - self.is_padded = True - if out_features % multiple > 0: - out_features = math.ceil(out_features / multiple) * multiple - self.is_padded = True - - super().__init__(in_features, - out_features, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - gather_output=gather_output, - is_qkv=is_qkv, - prefer_managed_weight=prefer_managed_weight) - - self.quant_algo = ( - use_int8_weight * GroupwiseQuantAlgo.INT8_WEIGHT + - use_w4a8_awq * GroupwiseQuantAlgo.W4A8_ALPHA + - pre_quant_scale * GroupwiseQuantAlgo.PRE_QUANT_SCALE + - zero * GroupwiseQuantAlgo.ZERO + bias * GroupwiseQuantAlgo.BIAS) - self.group_size = group_size - # packed in FP16 format (INT4*4 -> FP16, INT8*2 -> FP16) - pack_ratio = 2 if use_int8_weight else 4 - self.weight = Parameter(shape=(self.in_features, - self.out_features // pack_ratio), - dtype=dtype, - prefer_managed=self.prefer_managed_weight) - - scale_shape = (self.in_features // group_size, self.out_features) - self.weights_scaling_factor = Parameter(shape=scale_shape, dtype=dtype) - self.tp_rank = tp_rank - if self.is_padded: - self.tp_dim = -1 - self.pre_quant_scale = pre_quant_scale - self.use_w4a8_awq = use_w4a8_awq - - if pre_quant_scale: - self.prequant_scaling_factor = Parameter(shape=(1, - self.in_features), - dtype=dtype) - else: - self.register_parameter('prequant_scaling_factor', None) - - if zero: - self.zero = Parameter(shape=scale_shape, dtype=dtype) - else: - self.register_parameter('zero', None) - - if use_w4a8_awq: - self.alpha = Parameter(shape=(1, ), dtype="float32") - else: - self.register_parameter('alpha', None) - - validate_group_size(self) - if pre_quant_scale: - self.tllm_to_externel_key_dict = { - "weights_scaling_factor": "weight_scale", - "prequant_scaling_factor": "input_quantizer._pre_quant_scale", - } # AWQ - else: - self.tllm_to_externel_key_dict = { - "weight": ["qweight", "scales", "qzeros"] - } # GPTQ - - def forward(self, x, lora_runtime_params=None): - pre_quant_scale = self.prequant_scaling_factor.value if self.prequant_scaling_factor else None - zero = self.zero.value if self.zero else None - bias = self.bias.value if self.bias else None - alpha = self.alpha if self.alpha else None - - hidden_state = x - x = weight_only_groupwise_quant_matmul( - x, pre_quant_scale, self.weight.value, - self.weights_scaling_factor.value, zero, bias, alpha, - self.quant_algo, self.group_size, self.dtype) - - if default_net( - ).plugin_config.lora_plugin and lora_runtime_params is not None: - x = x + self.lora(hidden_state, - lora_runtime_params=lora_runtime_params) - - if self.gather_output and self.tp_size > 1 and self.tp_group is not None: - # [dim0, local_dim] -> [dim0 * tp_size, local_dim] --> [dim0, local_dim * tp_size] - x = allgather(x, self.tp_group, gather_dim=1) - - return x - - def postprocess(self, tllm_key, weights, **kwargs): - if tllm_key.endswith("zero") or ( - self.prequant_scaling_factor is None - and tllm_key.endswith("weights_scaling_factor")): - return {} - torch_dtype = str_dtype_to_torch(self.dtype) - return postprocess_weight_only_groupwise(tllm_key, weights, torch_dtype, - self, **kwargs) - - -WeightOnlyGroupwiseQuantColumnLinear = WeightOnlyGroupwiseQuantLinear - - -class WeightOnlyGroupwiseQuantRowLinear(RowLinear): - - def __init__(self, - in_features, - out_features, - group_size=128, - pre_quant_scale=False, - zero=False, - bias=False, - dtype=None, - tp_group=None, - tp_size=1, - tp_rank=0, - use_w4a8_awq=False, - use_int8_weight=False, - prefer_managed_weight=True): - multiple = max((128 if use_w4a8_awq else 64), group_size) * tp_size - self.is_padded = False - if in_features % multiple > 0: - in_features = math.ceil(in_features / multiple) * multiple - self.is_padded = True - if out_features % multiple > 0: - out_features = math.ceil(out_features / multiple) * multiple - self.is_padded = True - super().__init__(in_features, - out_features, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - prefer_managed_weight=prefer_managed_weight) - - self.quant_algo = ( - use_int8_weight * GroupwiseQuantAlgo.INT8_WEIGHT + - use_w4a8_awq * GroupwiseQuantAlgo.W4A8_ALPHA + - pre_quant_scale * GroupwiseQuantAlgo.PRE_QUANT_SCALE + - zero * GroupwiseQuantAlgo.ZERO + bias * GroupwiseQuantAlgo.BIAS) - self.group_size = group_size - # packed in FP16 format (INT4*4 -> FP16, INT8*2 -> FP16) - pack_ratio = 2 if use_int8_weight else 4 - self.weight = Parameter(shape=(self.in_features, - self.out_features // pack_ratio), - dtype=dtype, - prefer_managed=self.prefer_managed_weight) - scale_shape = (self.in_features // group_size, self.out_features) - self.weights_scaling_factor = Parameter(shape=scale_shape, dtype=dtype) - self.tp_rank = tp_rank - if self.is_padded: - self.tp_dim = -1 - self.pre_quant_scale = pre_quant_scale - self.use_w4a8_awq = use_w4a8_awq - - if pre_quant_scale: - self.prequant_scaling_factor = Parameter(shape=(1, - self.in_features), - dtype=dtype) - else: - self.register_parameter('prequant_scaling_factor', None) - - if zero: - self.zero = Parameter(shape=scale_shape, dtype=dtype) - else: - self.register_parameter('zero', None) - - if use_w4a8_awq: - self.alpha = Parameter(shape=(1, ), dtype="float32") - self.activation_scaling_factor = Parameter(shape=(1, ), - dtype=trt.float32) - else: - self.register_parameter('alpha', None) - self.register_parameter('activation_scaling_factor', None) - - validate_group_size(self) - if pre_quant_scale: - self.tllm_to_externel_key_dict = { - "weights_scaling_factor": "weight_scale", - "prequant_scaling_factor": "input_quantizer._pre_quant_scale", - } # AWQ - else: - self.tllm_to_externel_key_dict = { - "weight": ["qweight", "scales", "qzeros"] - } # GPTQ - - def forward(self, x, lora_runtime_params=None, all_reduce_params=None): - pre_quant_scale = self.prequant_scaling_factor.value if self.prequant_scaling_factor else None - zero = self.zero.value if self.zero else None - bias = self.bias.value if self.bias else None - alpha = self.alpha if self.alpha else None - activation_scaling_factor = self.activation_scaling_factor.value if self.activation_scaling_factor else None - - if self.alpha and x.dtype == trt.fp8: - x = dequantize(x, activation_scaling_factor, -1, self.dtype) - - hidden_state = x - x = weight_only_groupwise_quant_matmul( - x, pre_quant_scale, self.weight.value, - self.weights_scaling_factor.value, zero, bias, alpha, - self.quant_algo, self.group_size, self.dtype) - - if default_net( - ).plugin_config.lora_plugin and lora_runtime_params is not None: - x = x + self.lora(hidden_state, - lora_runtime_params=lora_runtime_params) - - if self.tp_size > 1 and self.tp_group is not None: - x = allreduce(x, self.tp_group, all_reduce_params=all_reduce_params) - - return x - - def postprocess(self, tllm_key, weights, **kwargs): - if tllm_key.endswith("zero") or ( - self.prequant_scaling_factor is None - and tllm_key.endswith("weights_scaling_factor")): - return {} - torch_dtype = str_dtype_to_torch(self.dtype) - return postprocess_weight_only_groupwise(tllm_key, weights, torch_dtype, - self, **kwargs) - - -class SmoothQuantMLP(Module): - - def __init__( - self, - hidden_size, - ffn_hidden_size, - hidden_act, - bias=True, - dtype=None, - tp_group=None, - tp_size=1, - quant_mode=QuantMode(0), - ): - super().__init__() - if hidden_act not in ACT2FN: - raise ValueError( - 'unsupported activation function: {}'.format(hidden_act)) - fc_output_size = 2 * ffn_hidden_size if hidden_act == 'swiglu' else ffn_hidden_size - self.fc = SmoothQuantColumnLinear(hidden_size, - fc_output_size, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - gather_output=False, - quant_mode=quant_mode) - - self.proj = SmoothQuantRowLinear(ffn_hidden_size, - hidden_size, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - quant_mode=quant_mode) - - self.hidden_act = hidden_act - self.quant_mode = quant_mode - self.dtype = dtype - - if self.quant_mode.has_act_static_scaling(): - self.quantization_scaling_factor = Parameter(shape=(1, ), - dtype='float32') - else: - self.register_parameter('quantization_scaling_factor', None) - - def forward(self, hidden_states, lora_layer_params=None): - - inter = self.fc(hidden_states) - inter = ACT2FN[self.hidden_act](inter) - value = cast(self.proj.smoother.value, inter.dtype) - inter = inter / value - if self.quant_mode.has_act_and_weight_quant(): - if self.quant_mode.has_act_static_scaling(): - # Avoid quantization layers as it breaks int8 plugins - inter = quantize_tensor(inter, - self.quantization_scaling_factor.value) - else: - # Quantize per token outputs tuple: - # quantized tensor and scaling factors per token - inter = quantize_per_token(inter) - output = self.proj(inter) - return output - - -class Int8SmoothQuantRowLinear(RowLinear): - - def __init__(self, - in_features, - out_features, - bias=True, - dtype=None, - tp_group=None, - tp_size=1, - prefer_managed_weight=True): - super().__init__(in_features, - out_features, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - prefer_managed_weight=prefer_managed_weight) - self.activation_scaling_factor = Parameter(shape=(1, ), - dtype=trt.float32) - self.weights_scaling_factor = Parameter(shape=(self.out_features, ), - dtype=trt.float32) - self.prequant_scaling_factor = Parameter(shape=(self.in_features, ), - dtype=dtype) - self.weight = Parameter(shape=(self.out_features, self.in_features), - dtype=trt.int8, - prefer_managed=self.prefer_managed_weight) - - def forward(self, x, lora_runtime_params=None, all_reduce_params=None): - lora_hidden_state = x if lora_runtime_params is not None else None - if default_net().strongly_typed: - assert is_same_dtype( - x.dtype, - self.dtype), f"Got input type {x.dtype}, expecting {self.dtype}" - x = mul(x, self.prequant_scaling_factor.value) - - x = cast(x, self.activation_scaling_factor.value.dtype) - - quantized_out = quantize(x, self.activation_scaling_factor.value, - 'int8') - - dequantized_out = dequantize(quantized_out, - self.activation_scaling_factor.value, -1, - self.activation_scaling_factor.value.dtype) - - dequantized_out = cast(dequantized_out, self.dtype) - - w_deq_out = dequantize(self.weight.value, - self.weights_scaling_factor.value, 0, - self.weights_scaling_factor.value.dtype) - - w_deq_out = cast(w_deq_out, self.dtype) - return self.multiply_collect(dequantized_out, - w_deq_out, - gemm_plugin=None, - all_reduce_params=all_reduce_params, - lora_runtime_params=lora_runtime_params, - lora_hidden_state=lora_hidden_state) - - -class Int8SmoothQuantLinear(Linear): - - def __init__( - self, - in_features, - out_features, - bias=True, - dtype=None, - tp_group=None, - tp_size=1, - gather_output=True, - prefer_managed_weight=True, - ): - super().__init__(in_features, - out_features, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - gather_output=gather_output, - prefer_managed_weight=prefer_managed_weight) - self.activation_scaling_factor = Parameter(shape=(1, ), - dtype=trt.float32) - - self.weights_scaling_factor = Parameter(shape=(self.out_features, ), - dtype=trt.float32) - self.prequant_scaling_factor = Parameter(shape=(self.in_features, ), - dtype=dtype) - self.weight = Parameter(shape=(self.out_features, self.in_features), - dtype=trt.int8, - prefer_managed=self.prefer_managed_weight) - - def forward(self, x, lora_runtime_params=None): - lora_hidden_state = x if lora_runtime_params is not None else None - if default_net().strongly_typed: - assert is_same_dtype( - x.dtype, - self.dtype), f"Got input type {x.dtype}, expecting {self.dtype}" - x = mul(x, self.prequant_scaling_factor.value) - x = cast(x, self.activation_scaling_factor.value.dtype) - - quantized_out = quantize(x, self.activation_scaling_factor.value, - 'int8') - - dequantized_out = dequantize(quantized_out, - self.activation_scaling_factor.value, -1, - self.activation_scaling_factor.value.dtype) - - dequantized_out = cast(dequantized_out, self.dtype) - - w_deq_out = dequantize(self.weight.value, - self.weights_scaling_factor.value, 0, - self.weights_scaling_factor.value.dtype) - w_deq_out = cast(w_deq_out, self.dtype) - - return self.multiply_collect(dequantized_out, - w_deq_out, - gemm_plugin=None, - lora_runtime_params=lora_runtime_params, - lora_hidden_state=lora_hidden_state) - - -class FP8Linear(Linear): - - def __init__(self, - in_features, - out_features, - bias=True, - dtype=None, - tp_group=None, - tp_size=1, - gather_output=True, - prefer_managed_weight=True, - is_qkv=False): - super().__init__(in_features, - out_features, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - gather_output=gather_output, - prefer_managed_weight=prefer_managed_weight, - is_qkv=is_qkv) - self.weight = Parameter(shape=(self.out_features, self.in_features), - dtype='fp8', - prefer_managed=self.prefer_managed_weight) - self.activation_scaling_factor = Parameter(shape=(1, ), - dtype=trt.float32) - self.weights_scaling_factor = Parameter(shape=(1, ), dtype=trt.float32) - if self.is_qkv: - self.tllm_to_externel_key_dict = { - "activation_scaling_factor": "input_scale", - "weight": ["weight", "weight_scale"], - "weights_scaling_factor": "", - } - else: - self.tllm_to_externel_key_dict = { - "activation_scaling_factor": "input_scale", - "weights_scaling_factor": "weight_scale", - } - - def forward(self, x, lora_runtime_params=None): - assert lora_runtime_params is None or default_net( - ).plugin_config.lora_plugin == self.dtype - - if default_net().strongly_typed: - assert default_net().plugin_config.user_buffer or is_same_dtype( - x.dtype, - self.dtype), f"Got input type {x.dtype}, expecting {self.dtype}" - - alpha = self.weights_scaling_factor.raw_value * self.activation_scaling_factor.raw_value - activation_scaling_factor = constant( - self.activation_scaling_factor.raw_value) - activation_scaling_factor = cast(activation_scaling_factor, self.dtype) - if x.dtype != trt.fp8: - quantized_out = quantize(x, activation_scaling_factor, 'fp8') - lora_hidden_state = x if lora_runtime_params is not None else None - else: - quantized_out = x - # TODO: add fp8 LoRA support - lora_hidden_state = dequantize( - x, activation_scaling_factor, -1, - self.dtype) if lora_runtime_params is not None else None - - weights_scaling_factor = cast(self.weights_scaling_factor.value, - self.dtype) - - if self.weight.value.dtype != trt.fp8: - w_quant_out = quantize(self.weight.value, weights_scaling_factor, - 'fp8') - else: - w_quant_out = self.weight.value - - gemm_plugin = default_net().plugin_config.gemm_plugin - - low_latency_gemm_plugin = default_net( - ).plugin_config.low_latency_gemm_plugin - if (low_latency_gemm_plugin == "fp8"): - return self.multiply_collect( - quantized_out, - w_quant_out, - gemm_plugin=None, - low_latency_gemm_plugin=low_latency_gemm_plugin, - use_fp8=True, - alpha=alpha, - lora_runtime_params=lora_runtime_params, - lora_hidden_state=lora_hidden_state) - elif gemm_plugin == 'fp8': - return self.multiply_collect( - quantized_out, - w_quant_out, - gemm_plugin=gemm_plugin, - use_fp8=True, - alpha=alpha, - lora_runtime_params=lora_runtime_params, - lora_hidden_state=lora_hidden_state) - else: - dequantized_out = dequantize(quantized_out, - activation_scaling_factor, -1, - self.dtype) - w_deq_out = dequantize(w_quant_out, weights_scaling_factor, -1, - self.dtype) - return self.multiply_collect( - dequantized_out, - w_deq_out, - gemm_plugin=None, - use_fp8=True, - lora_runtime_params=lora_runtime_params, - lora_hidden_state=lora_hidden_state) - - def postprocess(self, tllm_key, weights, **kwargs): - if self.is_qkv: - if tllm_key.endswith("activation_scaling_factor"): - return max(weights).reshape(1, ).to(torch.float32) - elif tllm_key.endswith("weights_scaling_factor"): - return {} - elif tllm_key.endswith("weight"): - assert len(weights) == 6 - weight_scaling_factors = weights[1::2] - new_amax = max(weight_scaling_factors).reshape(1, ).to( - torch.float32) - for qkv_idx in range(3): - idx = qkv_idx * 2 - weights[idx] = weights[idx].view(torch.float8_e4m3fn).to( - torch.float32) - weights[idx] *= weights[idx + 1] - weights[idx] /= new_amax - weights[idx] = weights[idx].to(torch.float8_e4m3fn) - weights = torch.cat(weights[::2]) - scales = new_amax - return { - tllm_key: weights, - tllm_key.replace("weight", "weights_scaling_factor"): scales - } - else: - return super().postprocess(tllm_key, weights, **kwargs) - if tllm_key.endswith("scaling_factor"): - return weights.reshape(1, ).to(torch.float32) - elif tllm_key.endswith("weight"): - return weights.view(torch.float8_e4m3fn) - elif tllm_key.endswith("bias"): - return weights.to(trt_dtype_to_torch(self.bias.dtype)) - - -class FP8RowLinear(RowLinear): - - def __init__(self, - in_features, - out_features, - bias=True, - dtype=None, - tp_group=None, - tp_size=1, - prefer_managed_weight=True, - is_expert=False): - super().__init__(in_features, - out_features, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - prefer_managed_weight=prefer_managed_weight, - is_expert=is_expert) - self.weight = Parameter( - shape=(self.out_features, self.in_features), - dtype="fp8", - prefer_managed=self.prefer_managed_weight, - ) - self.activation_scaling_factor = Parameter(shape=(1, ), - dtype=trt.float32) - self.weights_scaling_factor = Parameter(shape=(1, ), dtype=trt.float32) - self.tllm_to_externel_key_dict = { - "activation_scaling_factor": "input_scale", - "weights_scaling_factor": "weight_scale", - } - - def forward(self, x, lora_runtime_params=None, all_reduce_params=None): - assert lora_runtime_params is None or default_net( - ).plugin_config.lora_plugin == self.dtype - - alpha = self.weights_scaling_factor.raw_value * self.activation_scaling_factor.raw_value - activation_scaling_factor = cast(self.activation_scaling_factor.value, - self.dtype) - if x.dtype != trt.fp8: - quantized_out = quantize(x, activation_scaling_factor, 'fp8') - lora_hidden_state = x if lora_runtime_params is not None else None - else: - quantized_out = x - # TODO: add fp8 LoRA support - lora_hidden_state = dequantize( - x, activation_scaling_factor, -1, - self.dtype) if lora_runtime_params is not None else None - - weights_scaling_factor = cast(self.weights_scaling_factor.value, - self.dtype) - if self.weight.value.dtype != trt.fp8: - w_quant_out = quantize(self.weight.value, weights_scaling_factor, - 'fp8') - else: - w_quant_out = self.weight.value - - gemm_plugin = default_net().plugin_config.gemm_plugin - low_latency_gemm_plugin = default_net( - ).plugin_config.low_latency_gemm_plugin - gemm_allreduce_plugin = default_net( - ).plugin_config.gemm_allreduce_plugin - if gemm_allreduce_plugin: - ret = self.multiply_collect(quantized_out, - w_quant_out, - gemm_plugin=None, - use_fp8=True, - alpha=alpha, - lora_runtime_params=lora_runtime_params, - lora_hidden_state=lora_hidden_state, - all_reduce_params=all_reduce_params) - elif (low_latency_gemm_plugin == "fp8"): - ret = self.multiply_collect( - quantized_out, - w_quant_out, - gemm_plugin=None, - low_latency_gemm_plugin=low_latency_gemm_plugin, - use_fp8=True, - alpha=alpha, - lora_runtime_params=lora_runtime_params, - lora_hidden_state=lora_hidden_state, - all_reduce_params=all_reduce_params) - elif gemm_plugin == 'fp8': - ret = self.multiply_collect(quantized_out, - w_quant_out, - gemm_plugin=gemm_plugin, - use_fp8=True, - alpha=alpha, - lora_runtime_params=lora_runtime_params, - lora_hidden_state=lora_hidden_state, - all_reduce_params=all_reduce_params) - else: - dequantized_out = dequantize(quantized_out, - activation_scaling_factor, -1, - self.dtype) - w_deq_out = dequantize(w_quant_out, weights_scaling_factor, -1, - self.dtype) - ret = self.multiply_collect(dequantized_out, - w_deq_out, - gemm_plugin=None, - use_fp8=True, - lora_runtime_params=lora_runtime_params, - lora_hidden_state=lora_hidden_state, - all_reduce_params=all_reduce_params) - return ret - - def postprocess(self, tllm_key, weights, **kwargs): - if tllm_key.endswith("scaling_factor"): - return weights.reshape(1, ).to(torch.float32) - elif tllm_key.endswith("weight"): - return weights.view(torch.float8_e4m3fn) - elif tllm_key.endswith("bias"): - return weights.to(str_dtype_to_torch(self.bias.dtype)) - - -class Fp8RowwiseMLP(Module): - - def __init__( - self, - hidden_size, - ffn_hidden_size, - hidden_act, - bias=True, - dtype=None, - tp_group=None, - tp_size=1, - quant_mode=QuantMode(0), - clamp_val=None, - ): - super().__init__() - if hidden_act not in ACT2FN: - raise ValueError( - 'unsupported activation function: {}'.format(hidden_act)) - fc_output_size = 2 * ffn_hidden_size if hidden_act == 'swiglu' else ffn_hidden_size - self.fc = Fp8RowwiseColumnLinear(hidden_size, - fc_output_size, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - gather_output=False, - quant_mode=quant_mode) - - self.proj = Fp8RowwiseRowLinear(ffn_hidden_size, - hidden_size, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - quant_mode=quant_mode) - - self.hidden_size = hidden_size - self.ffn_hidden_size = ffn_hidden_size - self.hidden_act = hidden_act - self.bias = bias - self.dtype = dtype - self.tp_group = tp_group - self.tp_size = tp_size - self.quant_mode = quant_mode - - if clamp_val: - if not (isinstance(clamp_val, list) and len(clamp_val) == 2): - raise ValueError(f'unsupported clamp_val {clamp_val}') - self.clamp_val = Parameter(np.array(clamp_val, dtype=np.float32), - dtype='float32', - is_buffer=True) - else: - self.register_parameter('clamp_val', None) - - def forward(self, hidden_states, lora_layer_params=None): - assert lora_layer_params is None, f"lora is not supported on {self.__class__.__name__} now" - inter = self.fc(hidden_states) - inter = ACT2FN[self.hidden_act](inter) - if self.quant_mode.has_fp8_rowwise(): - # Quantize per token outputs tuple: - # quantized tensor and scaling factors per token - clamp_val = None if self.clamp_val is None else self.clamp_val.value - inter = quantize_fp8_per_token(inter, clamp_val) - output = self.proj(inter) - return output - - -class Fp8RowwiseGatedMLP(Fp8RowwiseMLP): - - def __init__( - self, - hidden_size, - ffn_hidden_size, - hidden_act, - bias=True, - dtype=None, - tp_group=None, - tp_size=1, - quant_mode=QuantMode(0), - clamp_val=None, - ): - super().__init__(hidden_size, - ffn_hidden_size, - hidden_act, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - quant_mode=quant_mode, - clamp_val=clamp_val) - if hidden_act not in ACT2FN: - raise ValueError( - 'unsupported activation function: {}'.format(hidden_act)) - self.gate = Fp8RowwiseColumnLinear(hidden_size, - ffn_hidden_size, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - gather_output=False, - quant_mode=quant_mode) - - def forward(self, hidden_states, lora_layer_params=None): - assert lora_layer_params is None, f"lora is not supported on {self.__class__.__name__} now" - inter = self.fc(hidden_states) - inter = ACT2FN[self.hidden_act](inter) - gate = self.gate(hidden_states) - inter_x_gate = inter * gate - if self.quant_mode.has_fp8_rowwise(): - # Quantize per token outputs tuple: - # quantized tensor and scaling factors per token - clamp_val = None if self.clamp_val is None else self.clamp_val.value - inter_x_gate = quantize_fp8_per_token(inter_x_gate, clamp_val) - output = self.proj(inter_x_gate) - return output - - -class Fp8RowwiseFusedGatedMLP(Module): - - def __init__( - self, - hidden_size, - ffn_hidden_size, - hidden_act, - bias=True, - dtype=None, - tp_group=None, - tp_size=1, - quant_mode=QuantMode(0), - clamp_val=None, - ): - super().__init__() - self.hidden_size = hidden_size - self.ffn_hidden_size = ffn_hidden_size - self.hidden_act = hidden_act - self.bias = bias - self.dtype = dtype - self.tp_group = tp_group - self.tp_size = tp_size - self.quant_mode = quant_mode - - self.fused_fc = Fp8RowwiseColumnLinear(hidden_size, - ffn_hidden_size * 2, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - gather_output=False, - quant_mode=quant_mode) - - self.proj = Fp8RowwiseRowLinear(ffn_hidden_size, - hidden_size, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - quant_mode=quant_mode) - - if clamp_val: - if not (isinstance(clamp_val, list) and len(clamp_val) == 2): - raise ValueError(f'unsupported clamp_val {clamp_val}') - self.clamp_val = Parameter(np.array(clamp_val, dtype=np.float32), - dtype='float32', - is_buffer=True) - else: - self.register_parameter('clamp_val', None) - - def forward(self, hidden_states, lora_layer_params=None): - assert lora_layer_params is None, f"lora is not supported on {self.__class__.__name__} now" - inter = self.fused_fc(hidden_states) - - if self.hidden_act == 'silu': - inter = ACT2FN['swiglu'](inter) - elif self.hidden_act == 'gelu': - inter = ACT2FN['geglu'](inter) - else: - raise NotImplementedError( - f"Activation {self.hidden_act} not yet implemented for {self.__class__.__name__}." - ) - - if self.quant_mode.has_fp8_rowwise(): - # Quantize per token outputs tuple: - # quantized tensor and scaling factors per token - clamp_val = None if self.clamp_val is None else self.clamp_val.value - inter = quantize_fp8_per_token(inter, clamp_val) - output = self.proj(inter) - return output - - -class Fp8RowwiseAttention(Module): - - def __init__(self, - *, - local_layer_idx, - hidden_size, - num_attention_heads, - num_kv_heads=None, - max_position_embeddings=1024, - num_layers=1, - apply_query_key_layer_scaling=False, - attention_head_size=None, - attention_mask_type=AttentionMaskType.padding, - bias=True, - dense_bias=None, - dtype=None, - position_embedding_type=PositionEmbeddingType.learned_absolute, - rotary_embedding_base=10000.0, - rotary_embedding_scaling=None, - rotary_embedding_percentage=1.0, - tp_group=None, - tp_size=1, - tp_rank=0, - scale_alibi_bias=False, - paged_kv_cache=False, - quant_mode=QuantMode(0), - clamp_val=None): - super().__init__() - self.local_layer_idx = local_layer_idx - self.attention_mask_type = attention_mask_type - self.attention_head_size = hidden_size // num_attention_heads if attention_head_size is None else attention_head_size - self.num_attention_heads = num_attention_heads // tp_size - self.num_kv_heads = num_kv_heads - self.num_attention_kv_heads = ( - num_kv_heads + tp_size - 1 - ) // tp_size if num_kv_heads is not None else self.num_attention_heads - self.hidden_size = hidden_size // tp_size - self.max_position_embeddings = 0 if max_position_embeddings is None else max_position_embeddings - self.tp_size = tp_size - self.tp_rank = tp_rank - self.dense_bias = dense_bias - if dense_bias is None: - self.dense_bias = bias - - self.num_layers = num_layers - self.apply_query_key_layer_scaling = apply_query_key_layer_scaling - self.norm_factor = math.sqrt(self.attention_head_size) - self.q_scaling = 1 - if self.apply_query_key_layer_scaling: - self.norm_factor *= self.num_layers - self.q_scaling *= self.num_layers - # Whether to scale ALiBi bias. Mathematically, it's equivalent to - # normalizing QK after adding bias. - # - False, inv_sqrt_Dh * Q*K^T + alibi_bias - # - True, inv_sqrt_Dh * Q*K^T + inv_sqrt_Dh * alibi_bias - self.scale_alibi_bias = scale_alibi_bias - - self.position_embedding_type = position_embedding_type - self.paged_kv_cache = paged_kv_cache - - self.rotary_embedding_base = rotary_embedding_base - self.rotary_embedding_scale_type = RotaryScalingType.none - self.rotary_embedding_scale = 1.0 - self.rotary_embedding_dim = 0 - - if rotary_embedding_scaling is not None: - rotary_scaling_type = rotary_embedding_scaling.get( - "type", rotary_embedding_scaling.get("rope_type")) - self.rotary_embedding_scale_type = RotaryScalingType.from_string( - rotary_scaling_type) - self.rotary_embedding_scale = rotary_embedding_scaling.get( - "factor", 1.0) - - if self.position_embedding_type.is_rope(): - self.rotary_embedding_dim = int(self.attention_head_size * - rotary_embedding_percentage) - elif self.position_embedding_type.is_alibi(): - alibi_scale = 1. / self.norm_factor if self.scale_alibi_bias else 1. - alibi_slopes = generate_alibi_slopes(self.num_attention_heads * - self.tp_size, - tp_size=self.tp_size, - tp_rank=self.tp_rank, - alibi_scale=alibi_scale) - self.register_parameter( - 'alibi_slopes', - Parameter(alibi_slopes, dtype='float32', is_buffer=True)) - - self.quant_mode = quant_mode - self.dtype = dtype - - self.register_parameter('kv_cache_scaling_factor', None) - - self.qkv = Fp8RowwiseColumnLinear( - hidden_size, - tp_size * self.num_attention_heads * self.attention_head_size + - (2 * tp_size * self.num_attention_kv_heads * - self.attention_head_size), - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - gather_output=False, - quant_mode=quant_mode) - - self.dense = Fp8RowwiseRowLinear(tp_size * self.num_attention_heads * - self.attention_head_size, - hidden_size, - bias=self.dense_bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - quant_mode=quant_mode) - - if clamp_val: - if not (isinstance(clamp_val, list) and len(clamp_val) == 2): - raise ValueError(f'unsupported clamp_val {clamp_val}') - self.clamp_val = Parameter(np.array(clamp_val, dtype=np.float32), - dtype='float32', - is_buffer=True) - else: - self.register_parameter('clamp_val', None) - - self.use_lora = False - - def forward( - self, - hidden_states: Tensor, - attention_mask=None, - use_cache=False, - kv_cache_params=None, - attention_params=None, - spec_decoding_params=None, - encoder_output=None, - position_embedding=None, - norm_before_bmm1=False, - lora_layer_params=None, - all_reduce_params: Optional[AllReduceParams] = None, - ): - assert lora_layer_params is None, ( - f"LoRA is not supported by {self.__class__.__name__} (e.g., --use_fp8_rowwise). " - "If you need LoRA support, please use a non-quantized (e.g., bf16) attention implementation. " - "See https://github.com/NVIDIA/TensorRT-LLM/issues/2603 for details." - ) - qkv = self.qkv(hidden_states) - - alibi_slopes = None - if self.position_embedding_type == PositionEmbeddingType.alibi: - alibi_slopes = self.alibi_slopes.value - dtype = trt.float32 - if default_net().plugin_config.gpt_attention_plugin or default_net( - ).plugin_config.inflight_batching_gpt_attention_plugin: - dtype = hidden_states.dtype if self.quant_mode.has_act_static_scaling( - ) else hidden_states[0].dtype - if dtype == trt.int8: - dtype = trt.float16 - alibi_slopes = cast(alibi_slopes, dtype) - - if spec_decoding_params is None: - spec_decoding_params = SpecDecodingParams() - - assert default_net().plugin_config.gpt_attention_plugin - - assert attention_params.is_valid( - default_net().plugin_config.gpt_attention_plugin, - default_net().plugin_config.remove_input_padding, use_cache) - if use_cache: - assert kv_cache_params.is_valid( - default_net().plugin_config.gpt_attention_plugin) - assert self.attention_mask_type == AttentionMaskType.causal, \ - 'Plugin only support masked MHA.' - if self.kv_cache_scaling_factor is not None: - kv_orig_quant_scale = self.kv_cache_rcp_scaling_factor.value - kv_quant_orig_scale = self.kv_cache_scaling_factor.value - else: - kv_orig_quant_scale = None - kv_quant_orig_scale = None - if self.position_embedding_type.is_rope(): - rotary_inv_freq = attention_params.rotary_inv_freq - rotary_cos_sin = attention_params.embed_positions_for_gpt_attention - else: - rotary_inv_freq = None - rotary_cos_sin = None - context, past_key_value = gpt_attention( - qkv=qkv, - past_key_value=kv_cache_params.get_first_past_key_value(), - sequence_length=attention_params.sequence_length, - host_past_key_value_lengths=kv_cache_params. - host_past_key_value_lengths, - host_max_attention_window_sizes=kv_cache_params. - host_max_attention_window_sizes, - host_sink_token_length=kv_cache_params.host_sink_token_length, - context_lengths=attention_params.context_lengths, - cache_indirection=kv_cache_params.cache_indirection, - host_request_types=attention_params.host_request_types, - layer_idx=self.local_layer_idx, - num_heads=self.num_attention_heads, - num_kv_heads=self.num_attention_kv_heads, - num_kv_heads_origin=self.num_kv_heads, - hidden_size_per_head=self.attention_head_size, - q_scaling=self.q_scaling, - rotary_embedding_dim=self.rotary_embedding_dim, - rotary_embedding_base=self.rotary_embedding_base, - rotary_embedding_scale_type=self.rotary_embedding_scale_type, - rotary_embedding_scale=self.rotary_embedding_scale, - rotary_embedding_max_positions=self.max_position_embeddings, - position_embedding_type=self.position_embedding_type, - rotary_inv_freq=rotary_inv_freq, - rotary_cos_sin=rotary_cos_sin, - kv_orig_quant_scale=kv_orig_quant_scale, - kv_quant_orig_scale=kv_quant_orig_scale, - kv_cache_quant_mode=self.quant_mode, - max_context_length=attention_params.max_context_length, - alibi_slopes=alibi_slopes, - tp_size=self.tp_size, - tp_rank=self.tp_rank, - kv_cache_block_offsets=kv_cache_params.kv_cache_block_offsets, - host_kv_cache_block_offsets=kv_cache_params. - host_kv_cache_block_offsets, - host_kv_cache_pool_pointers=kv_cache_params. - host_kv_cache_pool_pointers, - host_kv_cache_pool_mapping=kv_cache_params. - host_kv_cache_pool_mapping, - host_context_lengths=attention_params.host_context_lengths, - use_cache=use_cache, - spec_decoding_generation_lengths=spec_decoding_params. - spec_decoding_generation_lengths, - spec_decoding_position_offsets=spec_decoding_params. - spec_decoding_position_offsets, - spec_decoding_packed_mask=spec_decoding_params. - spec_decoding_packed_mask, - host_runtime_perf_knobs=attention_params.host_runtime_perf_knobs, - host_context_progress=attention_params.host_context_progress) - - if self.quant_mode.has_fp8_rowwise(): - # Quantize per token outputs tuple: - # quantized tensor and scaling factors per token - clamp_val = None if self.clamp_val is None else self.clamp_val.value - context = quantize_fp8_per_token(context, clamp_val) - - context = self.dense( - context, - all_reduce_params=all_reduce_params, - ) - - if use_cache: - return (context, past_key_value) - - return context - - def postprocess(self, tllm_key, weights, **kwargs): - if tllm_key.endswith("kv_cache_scaling_factor") and weights is None: - return {tllm_key: torch.ones(1, )} - else: - return {tllm_key: weights} - - -class FP4Linear(Linear): - - def __init__( - self, - in_features, - out_features, - bias=True, - dtype=None, - tp_group=None, - tp_size=1, - gather_output=True, - prefer_managed_weight=True, - is_qkv=False, - ): - super().__init__(in_features, - out_features, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - gather_output=gather_output, - prefer_managed_weight=prefer_managed_weight, - is_qkv=is_qkv) - self.scaling_vector_size = 16 - assert self.in_features % self.scaling_vector_size == 0, \ - "Input features must be a multiple of 16 for FP4 GEMM" - - self.weight = Parameter(shape=(self.out_features, self.in_features), - dtype=trt.fp4) - self.weights_block_scaling_factor = Parameter( - shape=(self.out_features, - self.in_features // self.scaling_vector_size), - dtype=trt.fp8) - nrows = fp4_utils.pad_up(self.out_features, 128) - ncols = fp4_utils.pad_up(self.in_features // self.scaling_vector_size, - 4) - self.weights_block_scaling_factor_interleaved = Parameter(shape=(nrows, - ncols), - dtype=trt.fp8) - self.weights_global_scaling_factor = Parameter(shape=(1, ), - dtype=trt.float32) - self.activation_global_scaling_factor = Parameter(shape=(1, ), - dtype=trt.float32) - # alpha = 1.0 / (weight_global_scale * act_global_scale) - self.alpha = Parameter(shape=(1, ), dtype=trt.float32) - if self.is_qkv: - self.tllm_to_externel_key_dict = { - "weight": - ["weight", "weight_scale", "weight_scale_2", "input_scale"], - "weights_block_scaling_factor": - "", - "weights_block_scaling_factor_interleaved": - "", - "weights_global_scaling_factor": - "", - "activation_global_scaling_factor": - "", - "alpha": - "", - } - else: - self.tllm_to_externel_key_dict = { - "weight": ["weight"], - "weights_block_scaling_factor": "weight_scale", - "weights_block_scaling_factor_interleaved": "weight_scale", - "weights_global_scaling_factor": "weight_scale_2", - "activation_global_scaling_factor": "input_scale", - "alpha": ["weight_scale_2", "input_scale"], - } - - def forward(self, x, lora_runtime_params=None): - assert lora_runtime_params is None, "lora is not supported on FP4Linear now" - if isinstance(x, (tuple, list)): - fp4_x, act_per_block_scale = x - else: - if default_net().plugin_config.gemm_plugin == 'nvfp4': - fp4_x, act_per_block_scale = quantize_to_fp4_tensor( - x, div(1, self.activation_global_scaling_factor.value)) - else: - fp4_x, act_per_block_scale = dynamic_quantize( - x, self.activation_global_scaling_factor.value) - if default_net().plugin_config.gemm_plugin == 'nvfp4': - x = fp4_gemm(fp4_x, act_per_block_scale, self.weight.value, - self.weights_block_scaling_factor_interleaved.value, - self.alpha.value, self.dtype) - else: - quant_w = self.weight.value - scale_w = self.weights_block_scaling_factor.value - dequant_w = block_double_dequantize( - quant_w, - scale_w, - self.weights_global_scaling_factor.value, - dtype=trt.float16) - dequant_x = block_double_dequantize( - fp4_x, - act_per_block_scale, - self.activation_global_scaling_factor.value, - dtype=trt.float16) - x = matmul(dequant_x, dequant_w, transb=True).cast(self.dtype) - - if self.bias is not None: - x = x + self.bias.value - - if self.gather_output and self.tp_size > 1 and self.tp_group is not None: - # [dim0, local_dim] -> [dim0 * tp_size, local_dim] --> [dim0, local_dim * tp_size] - x = allgather(x, self.tp_group, gather_dim=1) - - return x - - def postprocess(self, tllm_key, weights, **kwargs): - if not any([ - tllm_key.endswith(suffix) - for suffix in self.tllm_to_externel_key_dict - ]): - return super().postprocess(tllm_key, weights, **kwargs) - if self.is_qkv: - if tllm_key.endswith("weight"): - assert len(weights) == 12 - qkv_weight = weights[0::4] - qkv_block_scale = weights[1::4] - qkv_global_scale = weights[2::4] - qkv_input_scale = weights[3::4] - - # Ckpt uses qkv max is guaranteed. So no need to re-quantize. - qkv_input_scale = max(qkv_input_scale).reshape(1, ).to( - torch.float32) - qkv_global_scale = max(qkv_global_scale).reshape(1, ).to( - torch.float32) - qkv_weight = torch.cat(qkv_weight) - qkv_block_scale = torch.cat(qkv_block_scale) - return { - tllm_key: - qkv_weight, - tllm_key.replace("weight", "weights_block_scaling_factor"): - qkv_block_scale, - tllm_key.replace( - 'weight', "weights_block_scaling_factor_interleaved"): - torch.ops.trtllm.block_scale_interleave( - qkv_block_scale.view( - torch.uint8).cpu().contiguous()).reshape( - qkv_block_scale.shape).view( - torch.float8_e4m3fn), - tllm_key.replace("weight", "weights_global_scaling_factor"): - qkv_global_scale, - tllm_key.replace("weight", "activation_global_scaling_factor"): - qkv_input_scale, - tllm_key.replace("weight", "alpha"): - qkv_input_scale * qkv_global_scale, - } - else: - return {} - else: - if tllm_key.endswith("weight"): - return weights - elif tllm_key.endswith("weights_block_scaling_factor"): - return weights - elif tllm_key.endswith("weights_block_scaling_factor_interleaved"): - return torch.ops.trtllm.block_scale_interleave( - weights.view(torch.uint8).cpu().contiguous()).reshape( - weights.shape).view(torch.float8_e4m3fn) - elif tllm_key.endswith("weights_global_scaling_factor"): - return weights.float() - elif tllm_key.endswith('activation_global_scaling_factor'): - return weights.float() - elif tllm_key.endswith('alpha'): - weight_global_sf = weights[0].float() - act_global_sf = weights[1].float() - return act_global_sf * weight_global_sf - - -class FP4RowLinear(RowLinear): - - def __init__( - self, - in_features, - out_features, - bias=True, - dtype=None, - tp_group=None, - tp_size=1, - ): - super().__init__(in_features, - out_features, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size) - self.scaling_vector_size = 16 - assert self.in_features % self.scaling_vector_size == 0, \ - "Input features must be a multiple of 16 for FP4 GEMM" - self.weight = Parameter(shape=(self.out_features, self.in_features), - dtype=trt.fp4) - self.weights_block_scaling_factor = Parameter( - shape=(self.out_features, - self.in_features // self.scaling_vector_size), - dtype=trt.fp8) - nrows = fp4_utils.pad_up(self.out_features, 128) - ncols = fp4_utils.pad_up(self.in_features // self.scaling_vector_size, - 4) - self.weights_block_scaling_factor_interleaved = Parameter(shape=(nrows, - ncols), - dtype=trt.fp8) - self.weights_global_scaling_factor = Parameter(shape=(1, ), - dtype=trt.float32) - self.activation_global_scaling_factor = Parameter(shape=(1, ), - dtype=trt.float32) - # alpha = 1.0 / (weight_global_scale * act_global_scale) - self.alpha = Parameter(shape=(1, ), dtype=trt.float32) - - self.tllm_to_externel_key_dict = { - "weight": ["weight"], - "weights_block_scaling_factor": "weight_scale", - "weights_block_scaling_factor_interleaved": "weight_scale", - "weights_global_scaling_factor": "weight_scale_2", - "activation_global_scaling_factor": "input_scale", - "alpha": ["weight_scale_2", "input_scale"], - } - - def forward(self, x, lora_runtime_params=None, all_reduce_params=None): - assert lora_runtime_params is None, "lora is not supported on FP4Linear now" - - if isinstance(x, (tuple, list)): - fp4_x, act_per_block_scale = x - else: - if default_net().plugin_config.gemm_plugin == "nvfp4": - fp4_x, act_per_block_scale = quantize_to_fp4_tensor( - x, div(1.0, self.activation_global_scaling_factor.value)) - else: - # WAR for FP8 output attention - if x.dtype == trt.fp8: - # Since the scale is NVFP4 scale, we need to make it back to fp8 scale - new_scale_factor = self.activation_global_scaling_factor.raw_value - new_scale_factor = constant(new_scale_factor * 6) - x = dequantize(x, new_scale_factor, 0, - new_scale_factor.dtype) - fp4_x, act_per_block_scale = dynamic_quantize( - x, self.activation_global_scaling_factor.value) - - if default_net().plugin_config.gemm_allreduce_plugin: - x = gemm_allreduce( - a=fp4_x, - b=self.weight.value, - a_sf=act_per_block_scale, - b_sf=self.weights_block_scaling_factor_interleaved.value, - transa=False, # row-major - transb=True, # col-major - alpha=self.alpha.value, - group=self.tp_group, # ranks participating - fp8_inputs_override=False) - else: - if default_net().plugin_config.gemm_plugin == "nvfp4": - x = fp4_gemm( - fp4_x, act_per_block_scale, self.weight.value, - self.weights_block_scaling_factor_interleaved.value, - self.alpha.value, self.dtype) - else: - quant_w = self.weight.value - scale_w = self.weights_block_scaling_factor.value - dequant_x = block_double_dequantize( - fp4_x, act_per_block_scale, - self.activation_global_scaling_factor.value, trt.float16) - dequant_w = block_double_dequantize( - quant_w, scale_w, self.weights_global_scaling_factor.value, - trt.float16) - x = matmul(dequant_x, dequant_w, transb=True).cast(self.dtype) - - if self.tp_size > 1 and self.tp_group is not None: - need_bias = self.bias is not None - fuse_bias_into_all_reduce = need_bias and ( - all_reduce_params - is not None) and (all_reduce_params.fusion_op - == AllReduceFusionOp.RESIDUAL_RMS_NORM) - if fuse_bias_into_all_reduce: - all_reduce_params.bias = self.bias.value - x = allreduce(x, - self.tp_group, - all_reduce_params=all_reduce_params) - if need_bias and not fuse_bias_into_all_reduce: - x = x + self.bias.value - return x - - if self.bias is not None: - x = x + self.bias.value - - return x - - def postprocess(self, tllm_key, weights, **kwargs): - if not any([ - tllm_key.endswith(suffix) - for suffix in self.tllm_to_externel_key_dict - ]): - return super().postprocess(tllm_key, weights, **kwargs) - - if tllm_key.endswith("weight"): - return weights - elif tllm_key.endswith("weights_block_scaling_factor"): - return weights - elif tllm_key.endswith("weights_block_scaling_factor_interleaved"): - return torch.ops.trtllm.block_scale_interleave( - weights.view(torch.uint8).cpu().contiguous()).reshape( - weights.shape).view(torch.float8_e4m3fn) - elif tllm_key.endswith("weights_global_scaling_factor"): - return weights.float() - elif tllm_key.endswith('activation_global_scaling_factor'): - return weights.float() - elif tllm_key.endswith('alpha'): - weight_global_sf = weights[0].float() - act_global_sf = weights[1].float() - return act_global_sf * weight_global_sf - - -class SmoothQuantGatedMLP(SmoothQuantMLP): - - def __init__( - self, - hidden_size, - ffn_hidden_size, - hidden_act, - bias=True, - dtype=None, - tp_group=None, - tp_size=1, - quant_mode=QuantMode(0), - ): - super().__init__(hidden_size, - ffn_hidden_size, - hidden_act, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - quant_mode=quant_mode) - if hidden_act not in ACT2FN: - raise ValueError( - 'unsupported activation function: {}'.format(hidden_act)) - self.gate = SmoothQuantColumnLinear(hidden_size, - ffn_hidden_size, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - gather_output=False, - quant_mode=quant_mode) - - if self.quant_mode.has_act_static_scaling(): - self.quantization_scaling_factor = Parameter(shape=(1, ), - dtype='float32') - else: - self.register_parameter('quantization_scaling_factor', None) - - def forward(self, hidden_states, lora_layer_params=None): - assert lora_layer_params is None, f"lora is not supported on {self.__class__.__name__} now" - inter = self.fc(hidden_states) - inter = ACT2FN[self.hidden_act](inter) - gate = self.gate(hidden_states) - inter_x_gate = inter * gate - smoother = cast(self.proj.smoother.value, self.dtype) - inter_x_gate = inter_x_gate / smoother - if self.quant_mode.has_act_and_weight_quant(): - if self.quant_mode.has_act_static_scaling(): - # Avoid quantization layers as it breaks int8 plugins - inter_x_gate = quantize_tensor( - inter_x_gate, self.quantization_scaling_factor.value) - else: - # Quantize per token outputs tuple: - # quantized tensor and scaling factors per token - inter_x_gate = quantize_per_token(inter_x_gate) - - output = self.proj(inter_x_gate) - return output - - -class SmoothQuantAttention(Module): - - def __init__(self, - *, - local_layer_idx, - hidden_size, - num_attention_heads, - num_kv_heads=None, - max_position_embeddings=1024, - num_layers=1, - apply_query_key_layer_scaling=False, - attention_head_size=None, - attention_mask_type=AttentionMaskType.padding, - bias=True, - dense_bias=None, - dtype=None, - position_embedding_type=PositionEmbeddingType.learned_absolute, - rotary_embedding_base=10000.0, - rotary_embedding_scaling=None, - rotary_embedding_percentage=1.0, - tp_group=None, - tp_size=1, - tp_rank=0, - scale_alibi_bias=False, - paged_kv_cache=False, - quant_mode=QuantMode(0)): - super().__init__() - self.local_layer_idx = local_layer_idx - self.attention_mask_type = attention_mask_type - self.attention_head_size = hidden_size // num_attention_heads if attention_head_size is None else attention_head_size - self.num_attention_heads = num_attention_heads // tp_size - self.num_kv_heads = num_kv_heads - self.num_attention_kv_heads = ( - num_kv_heads + tp_size - 1 - ) // tp_size if num_kv_heads is not None else self.num_attention_heads - self.hidden_size = hidden_size // tp_size - self.max_position_embeddings = 0 if max_position_embeddings is None else max_position_embeddings - self.tp_size = tp_size - self.tp_rank = tp_rank - self.dense_bias = dense_bias - if dense_bias is None: - self.dense_bias = bias - - self.num_layers = num_layers - self.apply_query_key_layer_scaling = apply_query_key_layer_scaling - self.norm_factor = math.sqrt(self.attention_head_size) - self.q_scaling = 1 - if self.apply_query_key_layer_scaling: - self.norm_factor *= self.num_layers - self.q_scaling *= self.num_layers - # Whether to scale ALiBi bias. Mathematically, it's equivalent to - # normalizing QK after adding bias. - # - False, inv_sqrt_Dh * Q*K^T + alibi_bias - # - True, inv_sqrt_Dh * Q*K^T + inv_sqrt_Dh * alibi_bias - self.scale_alibi_bias = scale_alibi_bias - - self.position_embedding_type = position_embedding_type - self.paged_kv_cache = paged_kv_cache - - self.rotary_embedding_base = rotary_embedding_base - self.rotary_embedding_scale_type = RotaryScalingType.none - self.rotary_embedding_scale = 1.0 - self.rotary_embedding_dim = 0 - - if rotary_embedding_scaling is not None: - rotary_scaling_type = rotary_embedding_scaling.get( - "type", rotary_embedding_scaling.get("rope_type")) - self.rotary_embedding_scale_type = RotaryScalingType.from_string( - rotary_scaling_type) - self.rotary_embedding_scale = rotary_embedding_scaling.get( - "factor", 1.0) - - if self.position_embedding_type.is_rope(): - self.rotary_embedding_dim = int(self.attention_head_size * - rotary_embedding_percentage) - elif self.position_embedding_type.is_alibi(): - alibi_scale = 1. / self.norm_factor if self.scale_alibi_bias else 1. - alibi_slopes = generate_alibi_slopes(self.num_attention_heads * - self.tp_size, - tp_size=self.tp_size, - tp_rank=self.tp_rank, - alibi_scale=alibi_scale) - self.register_parameter( - 'alibi_slopes', - Parameter(alibi_slopes, dtype='float32', is_buffer=True)) - - self.quant_mode = quant_mode - self.dtype = dtype - - if self.quant_mode.has_act_static_scaling(): - self.quantization_scaling_factor = Parameter(shape=(1, ), - dtype='float32') - else: - self.register_parameter('quantization_scaling_factor', None) - - qkv_quant_mode = quant_mode - if self.quant_mode.has_act_and_weight_quant(): - # We need to hijack quant_mode for QKV because QKV always uses per channel scaling - qkv_quant_mode = QuantMode.from_description( - True, True, quant_mode.has_per_token_dynamic_scaling(), True) - - self.register_parameter('kv_cache_scaling_factor', None) - - self.qkv = SmoothQuantColumnLinear( - hidden_size, - tp_size * self.num_attention_heads * self.attention_head_size + - (2 * tp_size * self.num_attention_kv_heads * - self.attention_head_size), - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - gather_output=False, - quant_mode=qkv_quant_mode) - - self.dense = SmoothQuantRowLinear(tp_size * self.num_attention_heads * - self.attention_head_size, - hidden_size, - bias=self.dense_bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - quant_mode=quant_mode) - - self.use_lora = False - - def forward( - self, - hidden_states: Tensor, - attention_mask=None, - use_cache=False, - kv_cache_params=None, - attention_params=None, - spec_decoding_params=None, - mrope_params=None, - encoder_output=None, - position_embedding=None, - norm_before_bmm1=False, - lora_layer_params=None, - all_reduce_params: Optional[AllReduceParams] = None, - ): - assert lora_layer_params is None, f"lora is not supported on {self.__class__.__name__} now" - qkv = self.qkv(hidden_states) - - alibi_slopes = None - if self.position_embedding_type == PositionEmbeddingType.alibi: - alibi_slopes = self.alibi_slopes.value - dtype = trt.float32 - if default_net().plugin_config.gpt_attention_plugin or default_net( - ).plugin_config.inflight_batching_gpt_attention_plugin: - dtype = hidden_states.dtype if self.quant_mode.has_act_static_scaling( - ) else hidden_states[0].dtype - if dtype == trt.int8: - dtype = trt.float16 - alibi_slopes = cast(alibi_slopes, dtype) - - if spec_decoding_params is None: - spec_decoding_params = SpecDecodingParams() - - if mrope_params is None: - mrope_params = MropeParams() - - if default_net().plugin_config.gpt_attention_plugin: - - assert attention_params.is_valid( - default_net().plugin_config.gpt_attention_plugin, - default_net().plugin_config.remove_input_padding, use_cache) - if use_cache: - assert kv_cache_params.is_valid( - default_net().plugin_config.gpt_attention_plugin) - assert self.attention_mask_type == AttentionMaskType.causal, \ - 'Plugin only support masked MHA.' - if self.kv_cache_scaling_factor is not None: - kv_orig_quant_scale = self.kv_cache_rcp_scaling_factor.value - kv_quant_orig_scale = self.kv_cache_scaling_factor.value - else: - kv_orig_quant_scale = None - kv_quant_orig_scale = None - if self.position_embedding_type.is_rope(): - rotary_inv_freq = attention_params.rotary_inv_freq - rotary_cos_sin = attention_params.embed_positions_for_gpt_attention - else: - rotary_inv_freq = None - rotary_cos_sin = None - context, past_key_value = gpt_attention( - qkv=qkv, - past_key_value=kv_cache_params.get_first_past_key_value(), - sequence_length=attention_params.sequence_length, - host_past_key_value_lengths=kv_cache_params. - host_past_key_value_lengths, - host_max_attention_window_sizes=kv_cache_params. - host_max_attention_window_sizes, - host_sink_token_length=kv_cache_params.host_sink_token_length, - context_lengths=attention_params.context_lengths, - cache_indirection=kv_cache_params.cache_indirection, - host_request_types=attention_params.host_request_types, - layer_idx=self.local_layer_idx, - num_heads=self.num_attention_heads, - num_kv_heads=self.num_attention_kv_heads, - num_kv_heads_origin=self.num_kv_heads, - hidden_size_per_head=self.attention_head_size, - q_scaling=self.q_scaling, - rotary_embedding_dim=self.rotary_embedding_dim, - rotary_embedding_base=self.rotary_embedding_base, - rotary_embedding_scale_type=self.rotary_embedding_scale_type, - rotary_embedding_scale=self.rotary_embedding_scale, - rotary_embedding_max_positions=self.max_position_embeddings, - position_embedding_type=self.position_embedding_type, - rotary_inv_freq=rotary_inv_freq, - rotary_cos_sin=rotary_cos_sin, - kv_orig_quant_scale=kv_orig_quant_scale, - kv_quant_orig_scale=kv_quant_orig_scale, - kv_cache_quant_mode=self.quant_mode, - max_context_length=attention_params.max_context_length, - alibi_slopes=alibi_slopes, - tp_size=self.tp_size, - tp_rank=self.tp_rank, - kv_cache_block_offsets=kv_cache_params.kv_cache_block_offsets, - host_kv_cache_block_offsets=kv_cache_params. - host_kv_cache_block_offsets, - host_kv_cache_pool_pointers=kv_cache_params. - host_kv_cache_pool_pointers, - host_kv_cache_pool_mapping=kv_cache_params. - host_kv_cache_pool_mapping, - host_context_lengths=attention_params.host_context_lengths, - use_cache=use_cache, - spec_decoding_generation_lengths=spec_decoding_params. - spec_decoding_generation_lengths, - spec_decoding_position_offsets=spec_decoding_params. - spec_decoding_position_offsets, - spec_decoding_packed_mask=spec_decoding_params. - spec_decoding_packed_mask, - mrope_rotary_cos_sin=mrope_params.mrope_rotary_cos_sin, - mrope_position_deltas=mrope_params.mrope_position_deltas, - host_runtime_perf_knobs=attention_params. - host_runtime_perf_knobs, - host_context_progress=attention_params.host_context_progress, - ) - else: - assert self.paged_kv_cache == False - - def transpose_for_scores(x): - new_x_shape = concat([ - shape(x, 0), - shape(x, 1), self.num_attention_heads, - self.attention_head_size - ]) - return x.view(new_x_shape).permute([0, 2, 1, 3]) - - query, key, value = split(qkv, self.hidden_size, dim=2) - query = transpose_for_scores(query) - key = transpose_for_scores(key) - value = transpose_for_scores(value) - - past_key_value = kv_cache_params.get_first_past_key_value() - if past_key_value is not None: - - def dequantize_tensor(x, scale): - # Cast from int8 to dtype - casted_x = cast(x, self.dtype) - return casted_x * scale - - if self.quant_mode.has_int8_kv_cache(): - past_key_value = dequantize_tensor( - past_key_value, self.kv_dequantization_scale.value) - - # past_key_value [bs, 2, num_heads, max_seq_len, head_dim] - past_key, past_value = split(past_key_value, 1, dim=1) - - key_shape = concat([ - shape(past_key, 0), - shape(past_key, 2), - shape(past_key, 3), - shape(past_key, 4) - ]) - past_key = past_key.view(key_shape, zero_is_placeholder=False) - past_value = past_value.view(key_shape, - zero_is_placeholder=False) - key = concat([past_key, key], dim=2) - value = concat([past_value, value], dim=2) - - def merge_caches(): - key_inflated_shape = concat([ - shape(key, 0), 1, - shape(key, 1), - shape(key, 2), - shape(key, 3) - ]) - inflated_key = key.view(key_inflated_shape, - zero_is_placeholder=False) - inflated_value = value.view(key_inflated_shape, - zero_is_placeholder=False) - past_key_value = concat([inflated_key, inflated_value], dim=1) - return past_key_value - - if self.attention_mask_type == AttentionMaskType.causal: - query_length = shape(query, 2) - key_length = shape(key, 2) - starts = concat([0, 0, key_length - query_length, 0]) - sizes = concat([1, 1, query_length, key_length]) - buffer = constant( - np.expand_dims( - np.tril( - np.ones( - (self.max_position_embeddings, - self.max_position_embeddings))).astype(bool), - (0, 1))) - causal_mask = slice(buffer, starts, sizes) - - key = key.permute([0, 1, 3, 2]) - with precision("float32"): - attention_scores = matmul(query, key) - - if self.attention_mask_type == AttentionMaskType.causal: - attention_scores = where(causal_mask, attention_scores, - -10000.0) - - attention_scores = attention_scores / self.norm_factor - attention_probs = softmax(attention_scores, dim=-1) - - context = matmul(attention_probs, value, - use_fp32_acc=False).permute([0, 2, 1, 3]) - context = context.view( - concat([shape(context, 0), - shape(context, 1), self.hidden_size])) - - past_key_value = merge_caches() - - if use_cache and self.quant_mode.has_int8_kv_cache(): - past_key_value = quantize_tensor( - past_key_value, self.kv_quantization_scale.value) - value = cast(self.dense.smoother.value, context.dtype) - context = context / value - if self.quant_mode.has_act_and_weight_quant(): - if self.quant_mode.has_act_static_scaling(): - # Avoid quantization layers as it breaks int8 plugins - context = quantize_tensor( - context, self.quantization_scaling_factor.value) - else: - # Quantize per token outputs tuple: - # quantized tensor and scaling factors per token - context = quantize_per_token(context) - - context = self.dense( - context, - all_reduce_params=all_reduce_params, - ) - - if use_cache: - return (context, past_key_value) - - return context - - -# TODO: Duplicates SmoothQuantRmsNorm -class QServeRmsNorm(Module): - - def __init__(self, - normalized_shape, - eps=1e-06, - elementwise_affine=False, - dtype=None, - quant_mode=QuantMode(0), - bias=False, - clamp_val=None): - super().__init__() - if isinstance(normalized_shape, int): - normalized_shape = (normalized_shape, ) - assert quant_mode.is_qserve_w4a8() - self.normalized_shape = tuple(normalized_shape) - self.elementwise_affine = elementwise_affine - if self.elementwise_affine: - self.weight = Parameter(shape=self.normalized_shape, dtype=dtype) - else: - self.register_parameter('weight', None) - - if bias: - self.bias = Parameter(shape=self.normalized_shape, dtype=dtype) - else: - self.register_parameter('bias', None) - if clamp_val: - if not (isinstance(clamp_val, list) and len(clamp_val) == 2): - raise ValueError(f'unsupported clamp_val {clamp_val}') - self.clamp_val = Parameter(np.array(clamp_val, dtype=np.float32), - dtype='float32', - is_buffer=True) - else: - self.register_parameter('clamp_val', None) - - self.eps = eps - self.dtype = dtype - self.quant_mode = quant_mode - - if self.quant_mode.has_act_and_weight_quant(): - self.scale_to_int = Parameter(shape=(1, ), dtype=dtype) - else: - self.register_parameter('scale_to_int', None) - - def forward(self, x): - weight = None if self.weight is None else self.weight.value - bias = None if self.bias is None else self.bias.value - scale = None if self.scale_to_int is None else self.scale_to_int.value - clamp_val = None if self.clamp_val is None else self.clamp_val.value - return smooth_quant_rms_norm( - x, - self.normalized_shape, - weight, - bias, - scale, - clamp_val, - self.eps, - dynamic_act_scaling=True, - scale_dtype='float16', - sum_per_token=not self.quant_mode.has_per_group_scaling(), - sum_dtype='float16') - - -# TODO: Mostly duplicates SmoothQuantMLP. -# TODO: MLP could represent GatedMLP if hidden_act=='swiglu'. -class QServeMLP(Module): - - def __init__( - self, - hidden_size, - ffn_hidden_size, - hidden_act, - bias=True, - dtype=None, - tp_group=None, - tp_size=1, - quant_mode=QuantMode(0), - ): - super().__init__() - if hidden_act not in ACT2FN: - raise ValueError( - 'unsupported activation function: {}'.format(hidden_act)) - fc_output_size = 2 * ffn_hidden_size if hidden_act == 'swiglu' else ffn_hidden_size - self.fc = QServeW4A8ColumnLinear(hidden_size, - fc_output_size, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - gather_output=False, - quant_mode=quant_mode) - - self.proj = QServeW4A8RowLinear(ffn_hidden_size, - hidden_size, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - quant_mode=quant_mode) - - self.hidden_act = hidden_act - self.quant_mode = quant_mode - self.dtype = dtype - - def forward(self, hidden_states): - inter = self.fc(hidden_states) - inter = ACT2FN[self.hidden_act](inter) - inter = quantize_per_token( - inter, - scale_dtype='float16', - sum_per_token=not self.quant_mode.has_per_group_scaling(), - sum_dtype='float16') - output = self.proj(inter) - return output - - -# TODO: Mostly duplicates SmoothQuantGatedMLP. -class QServeGatedMLP(QServeMLP): - - def __init__( - self, - hidden_size, - ffn_hidden_size, - hidden_act, - bias=True, - dtype=None, - tp_group=None, - tp_size=1, - quant_mode=QuantMode(0), - ): - super().__init__(hidden_size, - ffn_hidden_size, - hidden_act, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - quant_mode=quant_mode) - if hidden_act not in ACT2FN: - raise ValueError( - 'unsupported activation function: {}'.format(hidden_act)) - self.gate = QServeW4A8Linear(hidden_size, - ffn_hidden_size, - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - gather_output=False, - quant_mode=quant_mode) - - def forward(self, hidden_states, lora_layer_params=None): - assert lora_layer_params is None, "lora_layer_params not supported" - inter = self.fc(hidden_states) - inter = ACT2FN[self.hidden_act](inter) - gate = self.gate(hidden_states) - - inter_x_gate = inter * gate - inter_x_gate = quantize_per_token( - inter_x_gate, - scale_dtype='float16', - sum_per_token=not self.quant_mode.has_per_group_scaling(), - sum_dtype='float16') - - output = self.proj(inter_x_gate) - return output - - -# TODO: Duplicates SmoothQuantAttention. -class QServeAttention(Module): - - def __init__(self, - *, - local_layer_idx, - hidden_size, - num_attention_heads, - num_kv_heads=None, - max_position_embeddings=1024, - num_layers=1, - apply_query_key_layer_scaling=False, - attention_head_size=None, - attention_mask_type=AttentionMaskType.padding, - bias=True, - dense_bias=None, - dtype=None, - position_embedding_type=PositionEmbeddingType.learned_absolute, - rotary_embedding_base=10000.0, - rotary_embedding_scaling=None, - rotary_embedding_percentage=1.0, - tp_group=None, - tp_size=1, - tp_rank=0, - scale_alibi_bias=False, - paged_kv_cache=False, - quant_mode=QuantMode(0)): - super().__init__() - self.local_layer_idx = local_layer_idx - self.attention_mask_type = attention_mask_type - self.attention_head_size = hidden_size // num_attention_heads if attention_head_size is None else attention_head_size - self.num_attention_heads = num_attention_heads // tp_size - self.num_kv_heads = num_kv_heads - self.num_attention_kv_heads = ( - num_kv_heads + tp_size - 1 - ) // tp_size if num_kv_heads is not None else self.num_attention_heads - self.hidden_size = hidden_size // tp_size - self.max_position_embeddings = 0 if max_position_embeddings is None else max_position_embeddings - self.tp_size = tp_size - self.tp_rank = tp_rank - self.dense_bias = dense_bias - if dense_bias is None: - self.dense_bias = bias - - self.num_layers = num_layers - self.apply_query_key_layer_scaling = apply_query_key_layer_scaling - self.norm_factor = math.sqrt(self.attention_head_size) - self.q_scaling = 1 - if self.apply_query_key_layer_scaling: - self.norm_factor *= self.num_layers - self.q_scaling *= self.num_layers - # Whether to scale ALiBi bias. Mathematically, it's equivalent to - # normalizing QK after adding bias. - # - False, inv_sqrt_Dh * Q*K^T + alibi_bias - # - True, inv_sqrt_Dh * Q*K^T + inv_sqrt_Dh * alibi_bias - self.scale_alibi_bias = scale_alibi_bias - - self.position_embedding_type = position_embedding_type - self.paged_kv_cache = paged_kv_cache - - self.rotary_embedding_base = rotary_embedding_base - self.rotary_embedding_scale_type = RotaryScalingType.none - self.rotary_embedding_scale = 1.0 - self.rotary_embedding_dim = 0 - - if rotary_embedding_scaling is not None: - rotary_scaling_type = rotary_embedding_scaling.get( - "type", rotary_embedding_scaling.get("rope_type")) - self.rotary_embedding_scale_type = RotaryScalingType.from_string( - rotary_scaling_type) - self.rotary_embedding_scale = rotary_embedding_scaling.get( - "factor", 1.0) - - if self.position_embedding_type.is_rope(): - self.rotary_embedding_dim = int(self.attention_head_size * - rotary_embedding_percentage) - elif self.position_embedding_type.is_alibi(): - alibi_scale = 1. / self.norm_factor if self.scale_alibi_bias else 1. - alibi_slopes = generate_alibi_slopes(self.num_attention_heads * - self.tp_size, - tp_size=self.tp_size, - tp_rank=self.tp_rank, - alibi_scale=alibi_scale) - self.register_parameter( - 'alibi_slopes', - Parameter(alibi_slopes, dtype='float32', is_buffer=True)) - - self.quant_mode = quant_mode - self.dtype = dtype - - # QServe does not use act static scaling - # if self.quant_mode.has_act_static_scaling(): - # self.quantization_scaling_factor = Parameter(shape=(1, ), - # dtype='float32') - # else: - # self.register_parameter('quantization_scaling_factor', None) - - qkv_quant_mode = quant_mode - if self.quant_mode.has_act_and_weight_quant(): - # We need to hijack quant_mode for QKV because QKV always uses per channel scaling - qkv_quant_mode = QuantMode.from_description( - True, True, quant_mode.has_per_token_dynamic_scaling(), True) - - self.register_parameter('kv_cache_scaling_factor', None) - - self.qkv = QServeW4A8ColumnLinear( - hidden_size, - tp_size * self.num_attention_heads * self.attention_head_size + - (2 * tp_size * self.num_attention_kv_heads * - self.attention_head_size), - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - gather_output=False, - quant_mode=qkv_quant_mode) - - self.dense = QServeW4A8RowLinear(tp_size * self.num_attention_heads * - self.attention_head_size, - hidden_size, - bias=self.dense_bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - quant_mode=quant_mode) - - self.use_lora = False - - def forward( - self, - hidden_states: Tensor, - attention_mask=None, - use_cache=False, - kv_cache_params=None, - attention_params=None, - spec_decoding_params=None, - encoder_output=None, - position_embedding=None, - norm_before_bmm1=False, - lora_layer_params=None, - all_reduce_params: Optional[AllReduceParams] = None, - ): - assert lora_layer_params is None, "lora is not supported on SmoothQuantAttention now" - if default_net().plugin_config.qserve_gemm_plugin: - qkv = self.qkv(hidden_states) - else: - raise ValueError("qserve_gemm_plugin is not set") - - alibi_slopes = None - if self.position_embedding_type == PositionEmbeddingType.alibi: - alibi_slopes = self.alibi_slopes.value - dtype = trt.float32 - if default_net().plugin_config.gpt_attention_plugin or default_net( - ).plugin_config.inflight_batching_gpt_attention_plugin: - dtype = hidden_states.dtype if self.quant_mode.has_act_static_scaling( - ) else hidden_states[0].dtype - if dtype == trt.int8: - dtype = trt.float16 - alibi_slopes = cast(alibi_slopes, dtype) - - if spec_decoding_params is None: - spec_decoding_params = SpecDecodingParams() - - if default_net().plugin_config.gpt_attention_plugin: - - assert attention_params.is_valid( - default_net().plugin_config.gpt_attention_plugin, - default_net().plugin_config.remove_input_padding, use_cache) - if use_cache: - assert kv_cache_params.is_valid( - default_net().plugin_config.gpt_attention_plugin) - assert self.attention_mask_type == AttentionMaskType.causal, \ - 'Plugin only support masked MHA.' - if self.kv_cache_scaling_factor is not None: - kv_orig_quant_scale = self.kv_cache_rcp_scaling_factor.value - kv_quant_orig_scale = self.kv_cache_scaling_factor.value - else: - kv_orig_quant_scale = None - kv_quant_orig_scale = None - if self.position_embedding_type.is_rope(): - rotary_inv_freq = attention_params.rotary_inv_freq - rotary_cos_sin = attention_params.embed_positions_for_gpt_attention - else: - rotary_inv_freq = None - rotary_cos_sin = None - context, past_key_value = gpt_attention( - qkv=qkv, - past_key_value=kv_cache_params.get_first_past_key_value(), - sequence_length=attention_params.sequence_length, - host_past_key_value_lengths=kv_cache_params. - host_past_key_value_lengths, - host_max_attention_window_sizes=kv_cache_params. - host_max_attention_window_sizes, - host_sink_token_length=kv_cache_params.host_sink_token_length, - context_lengths=attention_params.context_lengths, - cache_indirection=kv_cache_params.cache_indirection, - host_request_types=attention_params.host_request_types, - layer_idx=self.local_layer_idx, - num_heads=self.num_attention_heads, - num_kv_heads=self.num_attention_kv_heads, - num_kv_heads_origin=self.num_kv_heads, - hidden_size_per_head=self.attention_head_size, - q_scaling=self.q_scaling, - rotary_embedding_dim=self.rotary_embedding_dim, - rotary_embedding_base=self.rotary_embedding_base, - rotary_embedding_scale_type=self.rotary_embedding_scale_type, - rotary_embedding_scale=self.rotary_embedding_scale, - rotary_embedding_max_positions=self.max_position_embeddings, - position_embedding_type=self.position_embedding_type, - rotary_inv_freq=rotary_inv_freq, - rotary_cos_sin=rotary_cos_sin, - kv_orig_quant_scale=kv_orig_quant_scale, - kv_quant_orig_scale=kv_quant_orig_scale, - kv_cache_quant_mode=self.quant_mode, - max_context_length=attention_params.max_context_length, - alibi_slopes=alibi_slopes, - tp_size=self.tp_size, - tp_rank=self.tp_rank, - kv_cache_block_offsets=kv_cache_params.kv_cache_block_offsets, - host_kv_cache_block_offsets=kv_cache_params. - host_kv_cache_block_offsets, - host_kv_cache_pool_pointers=kv_cache_params. - host_kv_cache_pool_pointers, - host_kv_cache_pool_mapping=kv_cache_params. - host_kv_cache_pool_mapping, - host_context_lengths=attention_params.host_context_lengths, - use_cache=use_cache, - spec_decoding_generation_lengths=spec_decoding_params. - spec_decoding_generation_lengths, - spec_decoding_position_offsets=spec_decoding_params. - spec_decoding_position_offsets, - spec_decoding_packed_mask=spec_decoding_params. - spec_decoding_packed_mask, - host_runtime_perf_knobs=attention_params. - host_runtime_perf_knobs, - host_context_progress=attention_params.host_context_progress) - else: - assert self.paged_kv_cache == False - - def transpose_for_scores(x): - new_x_shape = concat([ - shape(x, 0), - shape(x, 1), self.num_attention_heads, - self.attention_head_size - ]) - return x.view(new_x_shape).permute([0, 2, 1, 3]) - - query, key, value = split(qkv, self.hidden_size, dim=2) - query = transpose_for_scores(query) - key = transpose_for_scores(key) - value = transpose_for_scores(value) - - past_key_value = kv_cache_params.get_first_past_key_value() - if past_key_value is not None: - - def dequantize_tensor(x, scale): - # Cast from int8 to dtype - casted_x = cast(x, self.dtype) - return casted_x * scale - - if self.quant_mode.has_int8_kv_cache(): - past_key_value = dequantize_tensor( - past_key_value, self.kv_dequantization_scale.value) - - # past_key_value [bs, 2, num_heads, max_seq_len, head_dim] - past_key, past_value = split(past_key_value, 1, dim=1) - - key_shape = concat([ - shape(past_key, 0), - shape(past_key, 2), - shape(past_key, 3), - shape(past_key, 4) - ]) - past_key = past_key.view(key_shape, zero_is_placeholder=False) - past_value = past_value.view(key_shape, - zero_is_placeholder=False) - key = concat([past_key, key], dim=2) - value = concat([past_value, value], dim=2) - - def merge_caches(): - key_inflated_shape = concat([ - shape(key, 0), 1, - shape(key, 1), - shape(key, 2), - shape(key, 3) - ]) - inflated_key = key.view(key_inflated_shape, - zero_is_placeholder=False) - inflated_value = value.view(key_inflated_shape, - zero_is_placeholder=False) - past_key_value = concat([inflated_key, inflated_value], dim=1) - return past_key_value - - if self.attention_mask_type == AttentionMaskType.causal: - query_length = shape(query, 2) - key_length = shape(key, 2) - starts = concat([0, 0, key_length - query_length, 0]) - sizes = concat([1, 1, query_length, key_length]) - buffer = constant( - np.expand_dims( - np.tril( - np.ones( - (self.max_position_embeddings, - self.max_position_embeddings))).astype(bool), - (0, 1))) - causal_mask = slice(buffer, starts, sizes) - - key = key.permute([0, 1, 3, 2]) - with precision("float32"): - attention_scores = matmul(query, key) - - if self.attention_mask_type == AttentionMaskType.causal: - attention_scores = where(causal_mask, attention_scores, - -10000.0) - - attention_scores = attention_scores / self.norm_factor - attention_probs = softmax(attention_scores, dim=-1) - - context = matmul(attention_probs, value, - use_fp32_acc=False).permute([0, 2, 1, 3]) - context = context.view( - concat([shape(context, 0), - shape(context, 1), self.hidden_size])) - - past_key_value = merge_caches() - - if use_cache and self.quant_mode.has_int8_kv_cache(): - past_key_value = quantize_tensor( - past_key_value, self.kv_quantization_scale.value) - - # Quantize per token outputs tuple: - # quantized tensor and scaling factors per token - context = quantize_per_token( - context, - scale_dtype='float16', - sum_per_token=not self.quant_mode.has_per_group_scaling(), - sum_dtype='float16') - - context = self.dense(context, all_reduce_params=all_reduce_params) - - if use_cache: - return (context, past_key_value) - - return context - - -class Fp8RowwiseLayerNorm(Module): - - def __init__( - self, - normalized_shape, - eps=1e-05, - elementwise_affine=True, - dtype=None, - quant_mode=QuantMode(0), - bias=False, - clamp_val=None, - ): - super().__init__() - if isinstance(normalized_shape, int): - normalized_shape = (normalized_shape, ) - if not quant_mode.has_fp8_rowwise(): - raise ValueError( - "Fp8 Rowwise Layer norm has to have some quantization mode set") - self.normalized_shape = tuple(normalized_shape) - self.elementwise_affine = elementwise_affine - if self.elementwise_affine: - self.weight = Parameter(shape=self.normalized_shape, dtype=dtype) - else: - self.register_parameter('weight', None) - - if bias: - self.bias = Parameter(shape=self.normalized_shape, dtype=dtype) - else: - self.register_parameter('bias', None) - - if clamp_val: - if not (isinstance(clamp_val, list) and len(clamp_val) == 2): - raise ValueError(f'unsupported clamp_val {clamp_val}') - self.clamp_val = Parameter(np.array(clamp_val, dtype=np.float32), - dtype='float32', - is_buffer=True) - else: - self.register_parameter('clamp_val', None) - - self.eps = eps - self.dtype = dtype - self.quant_mode = quant_mode - - def forward(self, x): - weight = None if self.weight is None else self.weight.value - bias = None if self.bias is None else self.bias.value - scale = None - clamp_val = None if self.clamp_val is None else self.clamp_val.value - return fp8_rowwise_layer_norm( - x, - self.normalized_shape, - weight, - bias, - scale, - clamp_val, - self.eps, - dynamic_act_scaling=self.quant_mode.has_fp8_rowwise()) diff --git a/tensorrt_llm/quantization/quantize.py b/tensorrt_llm/quantization/quantize.py deleted file mode 100644 index a04a223914f2..000000000000 --- a/tensorrt_llm/quantization/quantize.py +++ /dev/null @@ -1,611 +0,0 @@ -import fnmatch -from typing import Union - -import torch - -from .._utils import get_init_params -from ..layers import (MLP, Attention, ColumnLinear, Embedding, GatedMLP, - LayerNorm, RmsNorm, RowLinear) -from ..layers.moe import MixtureOfExperts -from ..models.modeling_utils import LayerQuantConfig, QuantConfig -from ..parameter import Parameter - -# isort: off -from .layers import ( - FP4Linear, FP4RowLinear, FP8Linear, FP8RowLinear, Fp8RowwiseAttention, - Fp8RowwiseGatedMLP, Fp8RowwiseLayerNorm, Fp8RowwiseMLP, Fp8RowwiseRmsNorm, - Int8SmoothQuantLinear, Int8SmoothQuantRowLinear, QServeAttention, - QServeGatedMLP, QServeMLP, QServeRmsNorm, SmoothQuantAttention, - SmoothQuantGatedMLP, SmoothQuantLayerNorm, SmoothQuantMLP, - SmoothQuantRmsNorm, WeightOnlyGroupwiseQuantColumnLinear, - WeightOnlyGroupwiseQuantRowLinear, WeightOnlyQuantColumnLinear, - WeightOnlyQuantEmbedding, WeightOnlyQuantRowLinear) -# isort: on -from .mode import W8A8_SQ_PLUGIN_LIST, QuantAlgo, QuantMode - - -def quantize_layers( - model, - quant_config: QuantConfig, - quant_map, - preprocess_init_params=None, -): - exclude_modules = quant_config.exclude_modules - if exclude_modules is None: - exclude_modules = [ - '*lm_head', - '*router', - '*vocab_embedding', - '*position_embedding', - '*block_embedding', - '*shared_expert_gate', - ] - - for name, module, parent in model.named_modules_with_parent(): - module_name = name.rsplit('.', 1)[-1] - is_excluded = False - quant_cls = None - - # handle exclusion - for exclude_module in exclude_modules: - if fnmatch.fnmatchcase(name, exclude_module): - is_excluded = True - break - - # MoE modules are quantized on their constructor, so they must always - # be re-created with the appropriate quant_mode. When excluded, - # re-create with quant_mode 0. - # We need to handle it specially, we may want to redesign MoE implementation - if isinstance(module, MixtureOfExperts): - quant_cls = type(module) - elif not is_excluded: - for cls in quant_map: - if isinstance(module, cls): - quant_cls = quant_map[cls] - break - - if quant_cls: - init_params = get_init_params(module, quant_cls) - if isinstance(module, MixtureOfExperts): - if is_excluded: - quant_mode = QuantMode(0) - else: - quant_mode = quant_config.quant_mode - init_params["quant_mode"] = quant_mode - - # Auto-detect pre_quant_scale based on quant_algo - # For AWQ-based quantization methods that use pre_quant_scale - if quant_config.quant_algo in [ - QuantAlgo.W4A16_AWQ, QuantAlgo.NVFP4_AWQ, - QuantAlgo.W4A8_AWQ - ]: - init_params["pre_quant_scale"] = True - if "bias" in init_params and not isinstance(module, - MixtureOfExperts): - init_params["bias"] = init_params["bias"] is not None - if isinstance(module, ColumnLinear): - init_params[ - "out_features"] = module.out_features * module.tp_size - elif isinstance(module, RowLinear): - init_params["in_features"] = module.in_features * module.tp_size - if preprocess_init_params is not None: - preprocess_init_params(init_params, name, module) - quant_layer = quant_cls(**init_params) - if parent is not None: - setattr(parent, module_name, quant_layer) - else: - model = quant_layer - - setattr(model, 'quant_mode', quant_config.quant_mode) - return model - - -def weight_only_quantize(model, quant_config: QuantConfig, model_config=None): - assert quant_config.quant_mode.is_weight_only() - - try: - model_cfg = model.config - except AttributeError: - model_cfg = model_config - - quant_map = { - ColumnLinear: WeightOnlyQuantColumnLinear, - RowLinear: WeightOnlyQuantRowLinear, - Embedding: WeightOnlyQuantEmbedding, - } - - def preprocess_init_params(init_params, name, module): - init_params["quant_mode"] = quant_config.quant_mode - if isinstance(module, ColumnLinear): - module_name = name.rsplit('.', 1)[-1] - init_params["transb"] = module_name == "lm_head" - if "tp_rank" in init_params: - init_params["tp_rank"] = model_cfg.mapping.tp_rank - - model = quantize_layers( - model, - quant_config, - quant_map, - preprocess_init_params, - ) - return model - - -def weight_only_groupwise_quantize(model, - quant_config: QuantConfig, - model_config=None): - assert quant_config.quant_mode.is_weight_only() - - try: - model_cfg = model.config - except AttributeError: - model_cfg = model_config - - quant_map = { - ColumnLinear: WeightOnlyGroupwiseQuantColumnLinear, - RowLinear: WeightOnlyGroupwiseQuantRowLinear, - MixtureOfExperts: MixtureOfExperts, - } - - def preprocess_init_params(init_params, name, module): - init_params["group_size"] = quant_config.group_size - init_params["pre_quant_scale"] = quant_config.pre_quant_scale - init_params["zero"] = quant_config.has_zero_point - init_params[ - "use_w4a8_awq"] = quant_config.quant_algo == QuantAlgo.W4A8_AWQ - init_params[ - "use_int8_weight"] = quant_config.quant_algo == QuantAlgo.W8A16_GPTQ - if "tp_rank" in init_params: - init_params["tp_rank"] = model_cfg.mapping.tp_rank - - model = quantize_layers( - model, - quant_config, - quant_map, - preprocess_init_params, - ) - return model - - -def smooth_quantize_ootb( - model, - quant_config: QuantConfig, -): - quant_map = { - ColumnLinear: Int8SmoothQuantLinear, - RowLinear: Int8SmoothQuantRowLinear, - } - - model = quantize_layers( - model, - quant_config, - quant_map, - ) - return model - - -def smooth_quantize_plugin(model, quant_mode): - quant_map = { - RmsNorm: SmoothQuantRmsNorm, - LayerNorm: SmoothQuantLayerNorm, - GatedMLP: SmoothQuantGatedMLP, - MLP: SmoothQuantMLP, - Attention: SmoothQuantAttention, - } - for name, layer, parent in model.named_modules_with_parent(): - layer_name = name.rsplit('.', 1)[-1] - if layer_name in ['ln_f', 'ln_embed']: - continue - - quant_cls = None - for cls in quant_map: - if isinstance(layer, cls): - quant_cls = quant_map[cls] - break - - if quant_cls is None: - continue - - init_params = get_init_params(layer, quant_cls) - init_params["quant_mode"] = quant_mode - if isinstance(layer, Attention): - init_params[ - "num_attention_heads"] = layer.num_attention_heads * layer.tp_size - quant_layer = quant_cls(**init_params) - if parent is not None: - setattr(parent, layer_name, quant_layer) - else: - model = quant_layer - - setattr(model, 'quant_mode', quant_mode) - return model - - -def smooth_quantize(model, quant_config: QuantConfig): - assert quant_config.quant_mode.has_act_and_weight_quant() - if quant_config.quant_algo in W8A8_SQ_PLUGIN_LIST: - return smooth_quantize_plugin(model, quant_config.quant_mode) - else: - return smooth_quantize_ootb(model, quant_config) - - -def fp8_quantize(model, quant_config: QuantConfig): - assert quant_config.quant_mode.has_fp8_qdq() - - quant_map = { - ColumnLinear: FP8Linear, - RowLinear: FP8RowLinear, - MixtureOfExperts: MixtureOfExperts, - } - - model = quantize_layers( - model, - quant_config, - quant_map, - ) - return model - - -def fp8_rowwise_quantize(model, quant_config: QuantConfig): - assert quant_config.quant_mode.has_fp8_rowwise() - - quant_cls_map = { - RmsNorm: Fp8RowwiseRmsNorm, - LayerNorm: Fp8RowwiseLayerNorm, - GatedMLP: Fp8RowwiseGatedMLP, - MLP: Fp8RowwiseMLP, - Attention: Fp8RowwiseAttention, - } - - exclude_modules = quant_config.exclude_modules - if exclude_modules is None: - exclude_modules = [] - # Always exclude these modules for FP8 rowwise - exclude_modules = list( - set(exclude_modules + ['*ln_f', '*ln_embed', '*lm_head'])) - - def extract_layer_idx(name): - ss = name.split('.') - for s in ss: - if s.isdigit(): - return int(s) - return None - - # Meta's LLaMA 3.1 recipe: - # (1) Skip quantization for the first and last Transformer layers - # (2) Skip quantization for the Attention layers - if quant_config.use_meta_recipe: - exclude_modules.extend(['*input_layernorm', '*attention']) - - for name, layer, parent in model.named_modules_with_parent(): - module_name = name.rsplit('.', 1)[-1] - - if quant_config.use_meta_recipe: - local_layer_idx = extract_layer_idx(name) - mapping = model.config.mapping - layers_range = mapping.pp_layers(model.config.num_hidden_layers) - if mapping.is_first_pp_rank() and local_layer_idx == 0: - continue - if mapping.is_last_pp_rank( - ) and local_layer_idx == len(layers_range) - 1: - continue - - quant_cls = None - for cls in quant_cls_map: - if isinstance(layer, cls): - quant_cls = quant_cls_map[cls] - break - if quant_cls is None: - continue - - is_excluded = False - for exclude_module in exclude_modules: - if fnmatch.fnmatchcase(name, exclude_module): - is_excluded = True - break - if is_excluded: - continue - - init_params = get_init_params(layer, quant_cls) - init_params["quant_mode"] = quant_config.quant_mode - if isinstance(layer, Attention): - init_params[ - "num_attention_heads"] = layer.num_attention_heads * layer.tp_size - quant_layer = quant_cls(**init_params, clamp_val=quant_config.clamp_val) - if parent is not None: - setattr(parent, module_name, quant_layer) - else: - model = quant_layer - - setattr(model, 'quant_mode', quant_config.quant_mode) - return model - - -# TODO: These functions should be moved to ModelOpt. -def qserve_quantize_weight_per_group(linear_weight: torch.HalfTensor, - s1_scales: torch.FloatTensor, - s2_scales: torch.FloatTensor, - s2_szeros: torch.FloatTensor, - group_size: int) -> torch.CharTensor: - out_features = linear_weight.shape[0] - in_features = linear_weight.shape[1] - - # Step 1: Quantize the weights to int8 - linear_weight = linear_weight.div( - s1_scales.reshape(out_features, 1).to(linear_weight.device)) - linear_weight = linear_weight.round() - # assert linear_weight.min() >= -119 and linear_weight.max() <= 119, "Stage 1: Quantized weight out of range" # 119 is the "magic" number - assert (linear_weight.min() >= -128 and linear_weight.max() - <= 127), "Stage 1: Quantized weight out of range" - - # Step 2: Quantize the weights to int4 - linear_weight = linear_weight.reshape(out_features, - in_features // group_size, group_size) - s2_szeros = s2_szeros.reshape(out_features, in_features // group_size, - 1).to(torch.float16).to(linear_weight.device) - s2_scales = s2_scales.reshape(out_features, in_features // group_size, - 1).to(torch.float16).to(linear_weight.device) - linear_weight = linear_weight.add(s2_szeros).div(s2_scales).round() - assert (linear_weight.min() >= 0 and linear_weight.max() - <= 15), "Stage 2: Quantized weight out of range" - - qweight = linear_weight.reshape(out_features, in_features).to(torch.int8) - return qweight - - -def qserve_quantize_weight_per_channel( - linear_weight: torch.HalfTensor, s1_scales: torch.FloatTensor, - s1_szeros: torch.FloatTensor) -> torch.CharTensor: - out_features = linear_weight.shape[0] - in_features = linear_weight.shape[1] - - # Step 1: Quantize the weights to int4 - s1_scales = s1_scales.reshape(out_features, 1).to(linear_weight.device) - s1_szeros = s1_szeros.reshape(out_features, 1).to(linear_weight.device) - - qweight = linear_weight.add(s1_szeros).div(s1_scales).round() - assert (qweight.min() >= 0 - and qweight.max() <= 15), "Quantized weight out of range" - - return qweight.reshape(out_features, in_features).to(torch.int8) - - -# Pack the quantized weights, scales and zeros and apply the reordering required by QServe kernels. -# Return: processed [qweight, s1_scales, s2_scales, s2_zeros] -def qserve_pack_reorder_per_group(qweight: torch.CharTensor, - s1_scales: torch.FloatTensor, - s2_scales: torch.FloatTensor, - s2_szeros: torch.FloatTensor, group_size): - out_features = qweight.shape[0] - in_features = qweight.shape[1] - - outputs = [] - - s1_scales = s1_scales.reshape(out_features).to(torch.float16) - s2_szeros = s2_szeros.reshape(out_features, - in_features // group_size).to(torch.int8) - s2_scales = s2_scales.reshape(out_features, - in_features // group_size).to(torch.int8) - - # Step 3: Pack the quantized weights to real quantized weights - # ---- Repack the weight ---- # - assert qweight.dtype == torch.int8 - # pack to M // 32, K // 32, (8, 4), ([2], 2, 2, 4) - W_unpack_reorder = (qweight.reshape( - out_features // 32, - 2, - 2, - 8, - in_features // 32, - 2, - 4, - 4, - ).permute(0, 4, 3, 6, 1, 5, 2, 7).contiguous()) - W_unpack_reorder = (W_unpack_reorder.permute(0, 1, 2, 3, 5, 6, 7, - 4).contiguous().to(torch.int8)) - # B_fp16_reorder = B_fp16_reorder[:, :, :, :, :, :, [3, 2, 1, 0]].contiguous() - # [16, 0, 17, 1, ...] - W_unpack_repacked = (W_unpack_reorder[..., 1] << 4) + W_unpack_reorder[..., - 0] - W_unpack_repacked = W_unpack_repacked.reshape(out_features // 32, - in_features // 32, 32, 16) - W_unpack_repacked = W_unpack_repacked.reshape(out_features, - in_features // 2) - - outputs.append(W_unpack_repacked) - - # for the last dimension, organize as 0, 8, 16, 24, 1, 9, 17, 25, ... following the requirement of tensor core gemm - # ---- Pack the scales ---- # - outputs.append(s1_scales.reshape(out_features)) - - s2_scales = (s2_scales.reshape(out_features, in_features // - group_size).transpose(0, 1).contiguous()) - s2_scales = s2_scales.reshape(in_features // group_size, out_features // 32, - 32) - s2_scales = (s2_scales.reshape(in_features // group_size, - out_features // 32, 4, - 8).transpose(-2, -1).contiguous()) - s2_scales = s2_scales.reshape(in_features // group_size, - out_features).contiguous() - outputs.append(s2_scales) - - # ---- Pack the zeros ---- # - s2_szeros = (s2_szeros.reshape(out_features, in_features // - group_size).transpose(0, 1).contiguous()) - s2_szeros = s2_szeros.reshape(in_features // group_size, out_features // 32, - 32) - s2_szeros = (s2_szeros.reshape(in_features // group_size, - out_features // 32, 4, - 8).transpose(-2, -1).contiguous()) - s2_szeros = (s2_szeros.reshape(in_features // group_size, - out_features).contiguous()) - - # (q - s2_zeros) * s2_scales = q * s2_scales - s2_zeros * s2_scales, - # We convert the s2_zeros -> -s2_zeros * s2_scales - s2_szeros = (-s2_szeros).int() # It has been pre-scaled in DeepCompressor - s2_szeros = s2_szeros.to(torch.int8) - - outputs.append(s2_szeros) - - return outputs - - -def qserve_pack_reorder_per_channel(qweight: torch.CharTensor, - s1_scales: torch.FloatTensor, - s1_szeros: torch.FloatTensor): - out_features = qweight.shape[0] - in_features = qweight.shape[1] - - outputs = [] - - # ---- Repack the weight ---- # - assert qweight.dtype == torch.int8 - # pack to M // 32, K // 32, (8, 4), ([2], 2, 2, 4) - W_unpack_reorder = (qweight.reshape( - out_features // 32, - 2, - 2, - 8, - in_features // 32, - 2, - 4, - 4, - ).permute(0, 4, 3, 6, 1, 5, 2, 7).contiguous()) - W_unpack_reorder = (W_unpack_reorder.permute(0, 1, 2, 3, 5, 6, 7, - 4).contiguous()) - # B_fp16_reorder = B_fp16_reorder[:, :, :, :, :, :, [3, 2, 1, 0]].contiguous() - # [16, 0, 17, 1, ...] - W_unpack_repacked = (W_unpack_reorder[..., 1] << 4) + W_unpack_reorder[..., - 0] - W_unpack_repacked = W_unpack_repacked.reshape(out_features // 32, - in_features // 32, 32, 16) - W_unpack_repacked = W_unpack_repacked.reshape(out_features, in_features // - 2).contiguous() - - outputs.append(W_unpack_repacked) - - # ---- Pack the scales and zeros ---- # - s1_scales = s1_scales.reshape(out_features).contiguous() - outputs.append(s1_scales.half()) - - s1_szeros = s1_szeros.reshape(out_features).contiguous().half() - outputs.append(s1_szeros) - - return outputs - - -# TODO: Duplicates smooth_quantize and quantize_layers -def qserve_quantize(model, quant_config: QuantConfig): - quant_mode = quant_config.quant_mode - assert quant_config.quant_mode.is_qserve_w4a8() - - quant_map = { - RmsNorm: QServeRmsNorm, - LayerNorm: QServeRmsNorm, - GatedMLP: QServeGatedMLP, - MLP: QServeMLP, - Attention: QServeAttention, - } - - for name, layer, parent in model.named_modules_with_parent(): - layer_name = name.rsplit('.', 1)[-1] - if layer_name in ['ln_f', 'ln_embed']: - continue - - quant_cls = None - for cls in quant_map: - if isinstance(layer, cls): - quant_cls = quant_map[cls] - break - - if quant_cls is None: - continue - - init_params = get_init_params(layer, quant_cls) - init_params["quant_mode"] = quant_mode - if isinstance(layer, Attention): - init_params[ - "num_attention_heads"] = layer.num_attention_heads * layer.tp_size - quant_layer = quant_cls(**init_params) - if parent is not None: - setattr(parent, layer_name, quant_layer) - else: - model = quant_layer - - setattr(model, 'quant_mode', quant_mode) - return model - - -def fp4_quantize(model, quant_config: QuantConfig): - assert quant_config.quant_mode.has_nvfp4() - quant_map = { - ColumnLinear: FP4Linear, - RowLinear: FP4RowLinear, - MixtureOfExperts: MixtureOfExperts, - } - - model = quantize_layers( - model, - quant_config, - quant_map, - ) - return model - - -# Now consider the kv cache is enabled for all layers -def kv_cache_quantize(model): - for name, module in model.named_modules(): - if isinstance(module, - (Attention, SmoothQuantAttention, Fp8RowwiseAttention)): - # for dequant - module.kv_cache_scaling_factor = Parameter(shape=(1, ), - dtype='float32') - # for quant - module.kv_cache_rcp_scaling_factor = Parameter(shape=(1, ), - dtype='float32') - return model - - -def quantize(model, quant_config: Union[QuantConfig, LayerQuantConfig]): - - for name, module, parent in model.named_modules_with_parent(): - - if quant_config.quant_algo == QuantAlgo.MIXED_PRECISION: - layer_quant_mode = quant_config.layer_quant_mode(name) - else: - layer_quant_mode = quant_config.layer_quant_mode - if layer_quant_mode == QuantMode(0): - continue - - layer_quant_cfg = quant_config._get_quant_cfg(name) - - if layer_quant_mode.has_fp8_qdq(): - module = fp8_quantize(module, layer_quant_cfg) - elif layer_quant_mode.has_fp8_rowwise(): - module = fp8_rowwise_quantize(module, layer_quant_cfg) - elif layer_quant_mode.is_qserve_w4a8(): - module = qserve_quantize(module, quant_config) - elif layer_quant_mode.has_nvfp4(): - module = fp4_quantize(module, layer_quant_cfg) - elif layer_quant_mode.has_act_and_weight_quant(): - module = smooth_quantize(module, layer_quant_cfg) - elif layer_quant_mode.is_weight_only(): - if layer_quant_mode.has_per_group_scaling(): - module = weight_only_groupwise_quantize(module, layer_quant_cfg, - model.config) - else: - module = weight_only_quantize(module, layer_quant_cfg, - model.config) - - if parent is not None: # for per layer - module_name = name.rsplit('.', 1)[-1] - setattr(parent, module_name, module) - else: # for all layer - model = module - break - - if quant_config.quant_mode.has_kv_cache_quant(): - model = kv_cache_quantize(model) - - setattr(model, 'quant_mode', quant_config.quant_mode) - return model diff --git a/tensorrt_llm/quantization/quantize_by_modelopt.py b/tensorrt_llm/quantization/quantize_by_modelopt.py index fbb50310f893..be9236aa6cef 100755 --- a/tensorrt_llm/quantization/quantize_by_modelopt.py +++ b/tensorrt_llm/quantization/quantize_by_modelopt.py @@ -29,7 +29,6 @@ from accelerate.hooks import remove_hook_from_module from datasets import load_dataset from modelopt.torch.utils import print_rank_0 -from safetensors.torch import load_file, save_file from torch import nn from torch.utils.data import DataLoader from transformers import (AutoConfig, AutoModelForCausalLM, AutoProcessor, @@ -37,7 +36,6 @@ from .._utils import get_hf_rope_theta, release_gc, str_dtype_to_torch from ..logger import logger -from ..mapping import Mapping from .image_processing import MllamaImageProcessor from .mode import QuantAlgo @@ -722,73 +720,6 @@ def calibrate_loop(): return model -def combine_medusa_weight(tp_size, pp_size, base_model_output_dir, - num_medusa_heads, num_medusa_layers, max_draft_len, - medusa_hidden_act, medusa_model_dir, - quant_medusa_head): - - with open(f"{medusa_model_dir}/config.json", "r") as fp: - medusa_config = json.load(fp) - - num_medusa_heads_from_config = medusa_config.get('medusa_num_heads', - num_medusa_heads) - num_medusa_layers = medusa_config.get('medusa_num_layers', - num_medusa_layers) - if num_medusa_heads is None: - num_medusa_heads = num_medusa_heads_from_config - - assert max_draft_len > 0, "should have max_draft_len > 0" - - world_size = tp_size * pp_size - # Process for each rank - for rank in range(world_size): - mapping = Mapping(world_size=world_size, - rank=rank, - tp_size=tp_size, - pp_size=pp_size) - # 1. Load medusa weight for each rank - from tensorrt_llm.models.medusa.weight import load_medusa_hf - medusa_weights = load_medusa_hf(medusa_path=medusa_model_dir, - num_medusa_heads=num_medusa_heads, - num_medusa_layers=num_medusa_layers, - mapping=mapping, - dtype="float16") - # 2. Load base model safetensors (after quant) - base_model_weights = load_file( - f"{base_model_output_dir}/rank{rank}.safetensors") - - # 3. Combine and save weight - base_model_weights.update(medusa_weights) - save_file(base_model_weights, - f"{base_model_output_dir}/rank{rank}.safetensors") - - # 4. Add medusa config into config.json - with open(f"{base_model_output_dir}/config.json", 'r') as f: - base_model_config = json.load(f) - f.close() - - with open(f"{base_model_output_dir}/config.json", 'w') as f: - base_model_config['architecture'] = "MedusaForCausalLM" - base_model_config['quantization']['exclude_modules'] = [ - 'lm_head', - '*router', - '*vocab_embedding', - '*position_embedding', - '*block_embedding', - ] - if not quant_medusa_head: - base_model_config['quantization']['exclude_modules'].append( - '*medusa_heads*') - - base_model_config['max_draft_len'] = max_draft_len - base_model_config['num_medusa_heads'] = num_medusa_heads - base_model_config['num_medusa_layers'] = num_medusa_layers - json.dump(base_model_config, f, indent=4) - - torch.cuda.empty_cache() - logger.info("Combine medusa heads' weight, done.") - - def quantize_and_export(*, model_dir, device, @@ -1082,28 +1013,6 @@ def quantize_and_export(*, with open(f"{export_path}/config.json", "w") as f: json.dump(tensorrt_llm_config, f, indent=4) - # Workaround for combining medusa head - # TODO: move these integration into modelopt to avoid redundant reading and writing - if medusa_model_dir is not None: - combine_medusa_weight(tp_size, pp_size, export_path, - num_medusa_heads, num_medusa_layers, - max_draft_len, medusa_hidden_act, - medusa_model_dir, quant_medusa_head) - - # Workaround for mllama - if model_type == 'mllama': - from tensorrt_llm.models.mllama.config import MLLaMAConfig - config = MLLaMAConfig.from_hugging_face( - model_dir, - dtype=dtype, - ) - for key, value in config.to_dict().items(): - if key not in tensorrt_llm_config: - tensorrt_llm_config[key] = value - - with open(f"{export_path}/config.json", "w") as f: - json.dump(tensorrt_llm_config, f, indent=4) - end_time = time.time() logger.info( "Quantized model exported to {} \nTotal time used {:.2f} s.".format( diff --git a/tensorrt_llm/runtime/__init__.py b/tensorrt_llm/runtime/__init__.py index d6f018ff49aa..eaad2f1d3ab9 100644 --- a/tensorrt_llm/runtime/__init__.py +++ b/tensorrt_llm/runtime/__init__.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# 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"); @@ -12,6 +12,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. + import os import sys from contextlib import contextmanager @@ -40,16 +41,7 @@ def temporary_sys_path(path: str) -> Iterator[None]: with temporary_sys_path(os.path.dirname(os.path.abspath(__file__))): import kv_cache_manager_v2 -from .enc_dec_model_runner import EncDecModelRunner -from .generation import SamplingConfig # autoflake: skip -from .generation import (ChatGLMGenerationSession, GenerationSession, - LogitsProcessor, LogitsProcessorList, ModelConfig, - QWenForCausalLMGenerationSession, StoppingCriteria, - StoppingCriteriaList, decode_words_list) -from .kv_cache_manager import GenerationSequence, KVCacheManager -from .model_runner import ModelRunner -from .multimodal_model_runner import MultimodalModelRunner -from .session import Session, TensorInfo +from .model_config import ModelConfig try: import tensorrt_llm.bindings # NOQA @@ -57,28 +49,8 @@ def temporary_sys_path(path: str) -> Iterator[None]: except ImportError: PYTHON_BINDINGS = False -if PYTHON_BINDINGS: - from .model_runner_cpp import ModelRunnerCpp - __all__ = [ 'ModelConfig', - 'GenerationSession', - 'GenerationSequence', - 'KVCacheManager', - 'SamplingConfig', - 'Session', - 'TensorInfo', - 'ChatGLMGenerationSession', - 'QWenForCausalLMGenerationSession', - 'decode_words_list', - 'LogitsProcessorList', - 'LogitsProcessor', - 'StoppingCriteriaList', - 'StoppingCriteria', - 'ModelRunner', - 'ModelRunnerCpp', - 'EncDecModelRunner', - 'MultimodalModelRunner', 'PYTHON_BINDINGS', 'kv_cache_manager_v2', ] diff --git a/tensorrt_llm/runtime/enc_dec_model_runner.py b/tensorrt_llm/runtime/enc_dec_model_runner.py deleted file mode 100644 index 44ce54f03a26..000000000000 --- a/tensorrt_llm/runtime/enc_dec_model_runner.py +++ /dev/null @@ -1,537 +0,0 @@ -import json -import time -from pathlib import Path - -# isort: off -import torch -import tensorrt as trt - -from .._deprecation import emit_engine_arch_deprecation -from .._utils import torch_to_numpy, trt_dtype_to_torch, mpi_world_size, mpi_rank -from ..logger import logger -from ..plugin.plugin import CustomAllReduceHelper -from .generation import ModelConfig, SamplingConfig, LoraManager, GenerationSession -from ..mapping import Mapping -from .session import Session -from ..models.modeling_utils import get_kv_cache_type_from_legacy - - -def get_engine_name(rank): - return 'rank{}.engine'.format(rank) - - -def read_config(config_path: Path): - with open(config_path, "r") as f: - config = json.load(f) - - builder_config = config['build_config'] - plugin_config = builder_config['plugin_config'] - pretrained_config = config['pretrained_config'] - lora_config = builder_config['lora_config'] - use_gpt_attention_plugin = plugin_config["gpt_attention_plugin"] - remove_input_padding = plugin_config["remove_input_padding"] - use_lora_plugin = plugin_config["lora_plugin"] - tp_size = pretrained_config['mapping']['tp_size'] - pp_size = pretrained_config['mapping']['pp_size'] - gpus_per_node = pretrained_config['mapping']['gpus_per_node'] - world_size = tp_size * pp_size - assert world_size == mpi_world_size(), \ - f'Engine world size ({world_size}) != Runtime world size ({mpi_world_size()})' - num_heads = pretrained_config["num_attention_heads"] - hidden_size = pretrained_config["hidden_size"] - head_size = pretrained_config["head_size"] - vocab_size = pretrained_config["vocab_size"] - max_batch_size = builder_config["max_batch_size"] - max_beam_width = builder_config["max_beam_width"] - num_layers = pretrained_config["num_hidden_layers"] - num_kv_heads = pretrained_config.get('num_kv_heads', num_heads) - - assert (num_heads % tp_size) == 0 - num_heads = num_heads // tp_size - hidden_size = hidden_size // tp_size - num_kv_heads = (num_kv_heads + tp_size - 1) // tp_size - - cross_attention = pretrained_config["architecture"] == "DecoderModel" - skip_cross_kv = pretrained_config.get('skip_cross_kv', False) - has_position_embedding = pretrained_config["has_position_embedding"] - has_token_type_embedding = hasattr(pretrained_config, "type_vocab_size") - dtype = pretrained_config["dtype"] - - paged_kv_cache = plugin_config['paged_kv_cache'] - tokens_per_block = plugin_config['tokens_per_block'] - - gather_context_logits = builder_config.get('gather_context_logits', False) - gather_generation_logits = builder_config.get('gather_generation_logits', - False) - max_prompt_embedding_table_size = builder_config.get( - 'max_prompt_embedding_table_size', 0) - - kv_cache_type = get_kv_cache_type_from_legacy(True, paged_kv_cache) - language_adapter_config = pretrained_config.get("language_adapter_config", - None) - - model_config = ModelConfig( - num_heads=num_heads, - num_kv_heads=num_kv_heads, - hidden_size=hidden_size, - head_size=head_size, - max_batch_size=max_batch_size, - max_beam_width=max_beam_width, - vocab_size=vocab_size, - num_layers=num_layers, - gpt_attention_plugin=use_gpt_attention_plugin, - remove_input_padding=remove_input_padding, - kv_cache_type=kv_cache_type, - tokens_per_block=tokens_per_block, - cross_attention=cross_attention, - has_position_embedding=has_position_embedding, - has_token_type_embedding=has_token_type_embedding, - dtype=dtype, - gather_context_logits=gather_context_logits, - gather_generation_logits=gather_generation_logits, - max_prompt_embedding_table_size=max_prompt_embedding_table_size, - lora_plugin=use_lora_plugin, - lora_target_modules=lora_config.get('lora_target_modules'), - trtllm_modules_to_hf_modules=lora_config.get( - 'trtllm_modules_to_hf_modules'), - skip_cross_kv=skip_cross_kv, - language_adapter_config=language_adapter_config) - - return model_config, tp_size, pp_size, gpus_per_node, dtype - - -class EncDecModelRunner: - - def __init__(self, - engine_name, - engine_dir, - lora_dir=None, - lora_task_uids=None, - debug_mode=False, - skip_encoder=False, - stream: torch.cuda.Stream = None, - enable_context_fmha_fp32_acc: bool = None): - emit_engine_arch_deprecation("EncDecModelRunner") - # in multi-node setup, it's important to set_device at the very beginning so .to('cuda') refers to current device - # accordingly, all input & output tensors should be moved to current device - # otherwise, it's default to 'cuda:0' - self.runtime_rank = mpi_rank() - device_id = self.runtime_rank % torch.cuda.device_count() - torch.cuda.set_device(device_id) - self.device = torch.cuda.current_device() - self.skip_encoder = skip_encoder - self.lora_task_uids = lora_task_uids - self.enable_context_fmha_fp32_acc = enable_context_fmha_fp32_acc - - # when enc-dec runs by itself, stream can be None and we create new stream here - # when enc-dec has to run as a component in a bigger workflow (e.g., multimodal), earlier components in the workflow may have results in its stream, which we should pass that stream in to avoid unnecessary stream sync - self.stream = stream - if self.stream is None: - self.stream = torch.cuda.Stream(self.device) - torch.cuda.set_stream(self.stream) - - engine_dir = Path(engine_dir) - - def engine_setup(component): - # model config - config_path = engine_dir / component / "config.json" - logger.info(f"Using config path {config_path}") - model_config, tp_size, pp_size, gpus_per_node, dtype = read_config( - config_path) - - # MGMN config - world_size = tp_size * pp_size - runtime_rank = mpi_rank() - assert runtime_rank < world_size, "Runtime GPU rank exceeds MPI world size. Did you launch more MPI processes than required?" - runtime_mapping = Mapping(world_size, - runtime_rank, - tp_size=tp_size, - pp_size=pp_size, - gpus_per_node=gpus_per_node) - - # load engine - engine_fname = get_engine_name(runtime_rank) - with open(engine_dir / component / engine_fname, "rb") as f: - engine_buffer = f.read() - - return model_config, runtime_mapping, engine_buffer - - # Note: encoder and decoder doesn't necessarily have the same TP & PP config - - if not skip_encoder: - self.encoder_model_config, self.encoder_runtime_mapping, encoder_engine_buffer = engine_setup( - component='encoder') - - self.nccl_comm = None - if self.encoder_runtime_mapping.has_pp(): - # for Pipeline Parallelism in encoder - self.nccl_comm = torch.classes.trtllm.NcclCommunicatorOp( - self.encoder_runtime_mapping.world_size, - self.encoder_runtime_mapping.rank) - - # session setup - self.encoder_session = Session.from_serialized_engine( - encoder_engine_buffer) - - # encoder lora manager setup - if self.encoder_model_config.lora_plugin: - self.encoder_lora_manager = LoraManager( - mapping=self.encoder_runtime_mapping, - model_config=self.encoder_model_config, - ) - # TODO: this is only for bart - self.encoder_lora_manager.load_from_hf( - model_dirs=lora_dir, - model_config=self.encoder_model_config, - component='encoder', - ) - else: - self.encoder_lora_manager = None - else: - self.encoder_model_config, self.encoder_runtime_mapping, encoder_engine_buffer = None, None, None - self.nccl_comm, self.encoder_session = None, None - - self.decoder_model_config, self.decoder_runtime_mapping, decoder_engine_buffer = engine_setup( - component='decoder') - self.decoder_session = GenerationSession(self.decoder_model_config, - decoder_engine_buffer, - self.decoder_runtime_mapping, - debug_mode=debug_mode) - - # decoder lora manager setup - if self.decoder_model_config.lora_plugin: - self.decoder_lora_manager = LoraManager( - mapping=self.decoder_runtime_mapping, - model_config=self.decoder_model_config, - ) - # TODO: this is only for bart - self.decoder_lora_manager.load_from_hf( - model_dirs=lora_dir, - model_config=self.decoder_model_config, - component='decoder', - ) - else: - self.decoder_lora_manager = None - - @classmethod - def from_engine(cls, - engine_name, - engine_dir, - lora_dir=None, - lora_task_uids=None, - debug_mode=False, - skip_encoder=False, - stream=None, - enable_context_fmha_fp32_acc=None): - return cls(engine_name, - engine_dir, - lora_dir, - lora_task_uids, - debug_mode=debug_mode, - skip_encoder=skip_encoder, - stream=stream, - enable_context_fmha_fp32_acc=enable_context_fmha_fp32_acc) - - def process_input(self, - input_ids, - remove_input_padding=False, - pad_token_id=0, - prompt_tasks=None, - language_adapter_routings=None): - if remove_input_padding: - # in remove padding mode --> flatten input, calculate actual length and max length - # Note: 1st token should never be removed, even if it is pad_token_id - first_ids = input_ids[:, 0] - input_ids = input_ids[:, 1:] - input_lengths = 1 + (input_ids != pad_token_id).sum(dim=1).type( - torch.IntTensor).to(self.device) # [batch_size] - new_ids = [] - for i in range(len(input_ids)): - row = input_ids[i, :] - row = row[row != pad_token_id] - new_ids.append( - torch.cat( - (torch.IntTensor([first_ids[i]]).to(self.device), row))) - input_ids = torch.cat(new_ids) # [num_tokens] - if prompt_tasks is not None: - prompt_tasks = prompt_tasks[:input_ids.shape[0]] - else: - # in padding mode --> keep input, just calculate actual length and max length - # Note: 1st token should always count, even if it is pad_token_id. e.g., decoder start id in enc-dec models could be a single pad_token_id, we should count - input_lengths = torch.tensor( - 1 + (input_ids[:, 1:] != pad_token_id).sum(dim=1).type( - torch.IntTensor).to(self.device), - dtype=torch.int32, - device=self.device) - max_input_length = torch.max(input_lengths).item() - if language_adapter_routings is not None: - language_adapter_routings = language_adapter_routings.to( - self.device) - return input_ids, input_lengths, max_input_length, prompt_tasks, language_adapter_routings - - def encoder_run(self, - input_ids, - input_lengths, - max_input_length, - position_ids=None, - token_type_ids=None, - debug_mode=False, - prompt_embedding_table=None, - prompt_tasks=None, - prompt_vocab_size=None, - attention_mask=None, - language_adapter_routings=None): - - # each engine has hidden_dim/TP, don't forget to multiply TP - hidden_size = self.encoder_model_config.hidden_size * self.encoder_runtime_mapping.tp_size - if input_ids.dim() == 1: - hidden_states_shape = (input_ids.shape[0], hidden_size - ) # [num_tokens,D] - else: - hidden_states_shape = (input_ids.shape[0], input_ids.shape[1], - hidden_size) # [BS,seqlen,D] - hidden_states_dtype = lambda name: trt_dtype_to_torch( - self.encoder_session.engine.get_tensor_dtype(name)) - - # input tensors. only first PP rank has id input, others are hidden_states input - inputs = {} - if self.encoder_runtime_mapping.is_first_pp_rank(): - inputs['input_ids'] = input_ids.contiguous() - if self.encoder_model_config.has_position_embedding: - if position_ids is None: - if self.encoder_model_config.remove_input_padding: - position_ids = [ - torch.arange(sample_length, - dtype=torch.int32, - device=input_ids.device) - for sample_length in torch_to_numpy(input_lengths) - ] - position_ids = torch.cat(position_ids) - else: - bsz, seq_len = input_ids.shape[:2] - position_ids = torch.arange( - seq_len, dtype=torch.int32, - device=input_ids.device).expand(bsz, -1) - inputs['position_ids'] = position_ids.contiguous() - if self.encoder_model_config.has_token_type_embedding: - inputs['token_type_ids'] = token_type_ids.contiguous() - - if self.encoder_model_config.max_prompt_embedding_table_size > 0: - inputs[ - 'prompt_embedding_table'] = prompt_embedding_table.contiguous( - ) - inputs['tasks'] = prompt_tasks.contiguous() - inputs['prompt_vocab_size'] = prompt_vocab_size.contiguous() - else: - # just need a placeholder, engine will call NCCL to recv and fill data from previous rank - inputs['hidden_states_input'] = torch.empty( - hidden_states_shape, - dtype=hidden_states_dtype('hidden_states_input'), - device=self.device).contiguous() - if attention_mask is not None and not self.encoder_model_config.gpt_attention_plugin: - inputs['attention_mask'] = attention_mask.contiguous() - - inputs['input_lengths'] = input_lengths - # use shape info to pass max length info in remove padding mode - inputs['max_input_length'] = torch.empty( - (max_input_length, ), - dtype=hidden_states_dtype('max_input_length'), - device=self.device).contiguous() - - if self.encoder_runtime_mapping.tp_size > 1: - ipc_buffers, all_reduce_workspace = CustomAllReduceHelper.allocate_workspace( - self.encoder_runtime_mapping, - CustomAllReduceHelper.max_workspace_size_auto( - self.encoder_runtime_mapping.tp_size)) - inputs['all_reduce_workspace'] = all_reduce_workspace - - if self.encoder_model_config.lora_plugin: - inputs.update( - self.encoder_lora_manager.input_buffers( - self.lora_task_uids, - self.encoder_runtime_mapping, - self.encoder_model_config.num_layers, - )) - batch_size = input_lengths.size(0) - inputs['host_request_types'] = torch.IntTensor([0] * - batch_size).to('cpu') - if self.encoder_model_config.remove_input_padding: - inputs['host_context_lengths'] = input_lengths.to('cpu') - if language_adapter_routings is not None: - inputs['language_adapter_routings'] = language_adapter_routings - # Note: runtime.Session's run() method will set input/output tensor address, here we only need to provide tensor shape - self.encoder_session.set_shapes(inputs) - - # output tensors. only last PP rank final encoder output, others are intermediate hidden_states output. Need broadcast later - outputs = {} - if self.encoder_runtime_mapping.is_last_pp_rank(): - outputs['encoder_output'] = torch.empty( - hidden_states_shape, - dtype=hidden_states_dtype('encoder_output'), - device=self.device).contiguous() - else: - outputs['hidden_states_output'] = torch.empty( - hidden_states_shape, - dtype=hidden_states_dtype('hidden_states_output'), - device=self.device).contiguous() - - # ------------------------------------------- - if debug_mode: - engine = self.encoder_session.engine - context = self.encoder_session.context - # setup debugging buffer for the encoder - for i in range(self.encoder_session.engine.num_io_tensors): - name = engine.get_tensor_name(i) - if engine.get_tensor_mode( - name - ) == trt.TensorIOMode.OUTPUT and name not in outputs.keys(): - dtype = engine.get_tensor_dtype(name) - shape = context.get_tensor_shape(name) - outputs[name] = torch.zeros(tuple(shape), - dtype=trt_dtype_to_torch(dtype), - device=self.device) - context.set_tensor_address(name, outputs[name].data_ptr()) - # ------------------------------------------- - - # TRT session run - # Note: need cuda stream ID, not a torch Stream - ok = self.encoder_session.run(inputs, outputs, self.stream.cuda_stream) - assert ok, "Runtime execution failed" - self.stream.synchronize() - - # Tensor Parallelism is handled by model/engine definition - # But we need to broadcast among PP group at the end of encoder's Pipeline Parallelism - # After this, all ranks should recv the encoder output, and world might be re-configured using decoder's TP-PP config - def pp_communicate_encoder_output(encoder_output): - if self.encoder_runtime_mapping.is_last_pp_rank(): - for pp_rank in self.encoder_runtime_mapping.pp_group: - if pp_rank != self.encoder_runtime_mapping.rank: - self.nccl_comm.send(encoder_output, pp_rank) - return encoder_output - else: - self.nccl_comm.recv(encoder_output, - self.encoder_runtime_mapping.pp_group[-1]) - return encoder_output - - if self.encoder_runtime_mapping.has_pp(): - # use hidden_states output buffer to receive output as the shapes are same - encoder_output_buf = outputs[ - 'encoder_output'] if self.encoder_runtime_mapping.is_last_pp_rank( - ) else outputs['hidden_states_output'] - encoder_output = pp_communicate_encoder_output(encoder_output_buf) - else: - encoder_output = outputs['encoder_output'] - - return encoder_output - - def generate(self, - encoder_input_ids, - decoder_input_ids, - max_new_tokens, - num_beams=1, - pad_token_id=None, - eos_token_id=None, - bos_token_id=None, - debug_mode=False, - return_dict=False, - prompt_embedding_table=None, - prompt_tasks=None, - prompt_vocab_size=None, - attention_mask=None, - time_encoder=False, - return_encoder_output=False, - encoder_language_adapter_routings=None, - decoder_language_adapter_routings=None): - ## ensure all externally provided tensors are on the correct device. - encoder_input_ids = encoder_input_ids.to(self.device) - decoder_input_ids = decoder_input_ids.to(self.device) - - if attention_mask is not None: - attention_mask = torch.tensor(attention_mask, - dtype=torch.int32, - device=self.device) - - ## encoder run - encoder_remove_input_padding = self.encoder_model_config.remove_input_padding if self.encoder_model_config else self.decoder_model_config.remove_input_padding - - encoder_input_ids, encoder_input_lengths, encoder_max_input_length, prompt_tasks, encoder_language_adapter_routings = self.process_input( - encoder_input_ids, encoder_remove_input_padding, pad_token_id, - prompt_tasks, encoder_language_adapter_routings) - - if not self.skip_encoder: - logger.info(f"Rank {self.runtime_rank} Running encoder engine ...") - if time_encoder: - tik = time.time() - encoder_output = self.encoder_run( - encoder_input_ids, - encoder_input_lengths, - encoder_max_input_length, - debug_mode=debug_mode, - prompt_embedding_table=prompt_embedding_table, - prompt_tasks=prompt_tasks, - prompt_vocab_size=prompt_vocab_size, - attention_mask=attention_mask, - language_adapter_routings=encoder_language_adapter_routings) - if time_encoder: - tok = time.time() - print(f"TRT-LLM Encoder time {(tok-tik)*1000}ms") - else: - encoder_output = prompt_embedding_table - if encoder_input_ids.dim() > 1: - encoder_output = encoder_output.unsqueeze(0) - - ## decoder run - logger.info(f"Rank {self.runtime_rank} Running decoder engine ...") - decoder_input_ids, decoder_input_lengths, decoder_max_input_length, _, decoder_language_adapter_routings = self.process_input( - decoder_input_ids, self.decoder_model_config.remove_input_padding, - pad_token_id, None, decoder_language_adapter_routings) - # `cross_attention_mask` in context phase [batch_size, query_len, encoder_input_len] - # where query_len happens to be 1 in current cases, but not necessarily always, and - # `cross_attention_mask` in generation phase [batch_size, 1, encoder_input_len] where - # the query_len is always 1 since we have kv cache. But we use - # cross_attention_mask[:, step, :] during generation - cross_attention_mask = None - if attention_mask is not None: - cross_attention_mask = torch.tensor(attention_mask, - dtype=torch.int32, - device=self.device).reshape( - attention_mask.shape[0], 1, - attention_mask.shape[1]) - cross_attention_mask = cross_attention_mask.repeat( - [1, decoder_max_input_length + max_new_tokens, 1]) - - # generation config - sampling_config = SamplingConfig(end_id=eos_token_id, - pad_id=pad_token_id, - num_beams=num_beams, - min_length=1, - return_dict=return_dict) - sampling_config.update(output_cum_log_probs=return_dict, - output_log_probs=return_dict) - - # decoder autoregressive generation - self.decoder_session.setup( - decoder_input_lengths.size(0), - decoder_max_input_length, - max_new_tokens, - num_beams, - max_attention_window_size=None, - encoder_max_input_length=encoder_max_input_length, - lora_manager=self.decoder_lora_manager, - lora_uids=self.lora_task_uids, - enable_context_fmha_fp32_acc=self.enable_context_fmha_fp32_acc) - - output = self.decoder_session.decode( - decoder_input_ids, - decoder_input_lengths, - sampling_config, - encoder_output=encoder_output, - encoder_input_lengths=encoder_input_lengths, - return_dict=return_dict, - cross_attention_mask=cross_attention_mask, - language_adapter_routings=decoder_language_adapter_routings) - - if return_dict and return_encoder_output: - output['encoder_output'] = encoder_output - - return output diff --git a/tensorrt_llm/runtime/generation.py b/tensorrt_llm/runtime/generation.py deleted file mode 100755 index 6a49587d3054..000000000000 --- a/tensorrt_llm/runtime/generation.py +++ /dev/null @@ -1,4785 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import copy -import math -import os -import platform -from collections import Counter -from dataclasses import dataclass, field -from functools import reduce, wraps -from pathlib import Path -from typing import Dict, Iterable, List, Optional, Sequence, Set, Union - -import numpy as np - -# isort: off -import torch -import tensorrt as trt -# isort: on -try: - from cuda.bindings import runtime as cudart -except ImportError: - from cuda import cudart - -from tensorrt_llm.runtime.memory_pools.memory_pools_allocator import \ - MemoryPoolsAllocator -from tensorrt_llm.runtime.memory_pools.pools_kv_cache_manager import \ - PoolsKVCacheManager -from tensorrt_llm.runtime.redrafter_utils import * - -from .._utils import (binding_layer_type_to_str, binding_to_str_dtype, - pad_vocab_size, str_dtype_to_torch, torch_to_numpy, - trt_dtype_to_torch) -from ..bindings import ipc_nvls_allocate, ipc_nvls_free -from ..layers import LanguageAdapterConfig -from ..llmapi.kv_cache_type import KVCacheType -from ..logger import logger -from ..lora_manager import LoraManager -from ..mapping import Mapping -from ..plugin.plugin import CustomAllReduceHelper -from ..quantization import QuantMode -from .kv_cache_manager import GenerationSequence, KVCacheUpdater -from .session import _scoped_stream - -# When variable is set, this will disable torch.cuda.set_device(...) calls -# Useful in situations where device is already assigned by another library, i.e., megatron. -DISABLE_TORCH_DEVICE_SET = os.environ.get("DISABLE_TORCH_DEVICE_SET", False) - - -def decode_words_list(word_dict: List[List[str]], - tokenizer=None, - add_special_tokens=False): - ''' - format of word_dict - len(word_dict) should be same to batch_size - word_dict[i] means the words for batch i - len(word_dict[i]) >= 1, which means it must contain at least 1 string - For example, word_dict[2] = [" I am happy", " I am sad"]. - ''' - assert tokenizer != None, "need to set tokenizer" - - decoded_words_batch = [] - for word_dict_item in word_dict: - decoded_words_request = [] - - for item in word_dict_item: - if isinstance(item, bytes): - item = [item.decode()] - - ids = tokenizer.encode(item, add_special_tokens=add_special_tokens) - - if len(ids) == 0: - continue - - decoded_words_request.append(ids) - decoded_words_batch.append(decoded_words_request) - - return decoded_words_batch - - -def to_word_list_format(word_dict: List[List[List[int]]]): - ''' - format of word_dict - len(word_dict) should be same to batch_size - word_dict[i] means the words for batch i - len(word_dict[i]) >= 1, which means it must contain at least 1 word - For example, word_dict[2] = [[1, 267], [534]] has two words. - ''' - - flat_ids = [] - offsets = [] - for word_dict_item in word_dict: - items_flat_ids = [] - items_offsets = [] - - for ids in word_dict_item: - items_flat_ids += ids - items_offsets.append(len(ids)) - - flat_ids.append(np.array(items_flat_ids)) - offsets.append(np.cumsum(np.array(items_offsets))) - - pad_to = max(1, max(len(ids) for ids in flat_ids)) - - for i, (ids, offs) in enumerate(zip(flat_ids, offsets)): - flat_ids[i] = np.pad(ids, (0, pad_to - len(ids)), constant_values=0) - offsets[i] = np.pad(offs, (0, pad_to - len(offs)), constant_values=-1) - - return np.array([flat_ids, offsets], dtype="int32").transpose((1, 0, 2)) - - -def _prepare_input_ids(tensors: Sequence[torch.Tensor]): - tensors = [torch.flatten(t) for t in tensors] - data = torch.concat(tensors) - row_lengths = [t.size(0) for t in tensors] - row_lengths = torch.tensor(row_lengths, - dtype=torch.int32, - device=data.device) - return (data, row_lengths) - - -def CUASSERT(cuda_ret): - err = cuda_ret[0] - if err != cudart.cudaError_t.cudaSuccess: - raise RuntimeError( - f"CUDA ERROR: {err}, error code reference: https://nvidia.github.io/cuda-python/module/cudart.html#cuda.cudart.cudaError_t" - ) - if len(cuda_ret) > 1: - return cuda_ret[1:] - return None - - -def _update_cuda_graph_instance(instance, graph): - err = cudart.cudaGraphExecUpdate(instance, graph) - if err != cudart.cudaError_t.cudaSuccess: - # When updating cuda graph failed, destroy and instantiate one. - CUASSERT(cudart.cudaGraphExecDestroy(instance)) - instance = CUASSERT(cudart.cudaGraphInstantiate(graph, 0))[0] - return instance - - -def _prepare_attention_mask(input_ids: torch.Tensor, - pad_id: Optional[int] = None): - is_pad_id_in_inputs = (pad_id is not None) and (pad_id in input_ids) - if input_ids is not None and is_pad_id_in_inputs: - mask = input_ids.ne(pad_id).int() - # for enc-dec models, pad_id could be the start token and should be always counted - # as valid token rather than padded token, so we force its mask to be 1. - # This doesn't impact the existing behavior - mask[:, 0] = 1 - return mask - else: - return torch.ones(input_ids.shape, - dtype=torch.int32, - device=input_ids.device) - - -def _tile_beam_width(tensor: torch.Tensor, num_beams: int): - new_shape = np.array(tensor.shape) - new_shape[0] = new_shape[0] * num_beams - - tile_size = np.ones(new_shape.shape, dtype=np.int32) - tile_size = np.insert(tile_size, 1, num_beams) - - new_tensor = torch.unsqueeze(tensor, 1) - new_tensor = new_tensor.tile(tile_size.tolist()) - new_tensor = new_tensor.reshape(new_shape.tolist()) - return new_tensor - - -class _Profiler(trt.IProfiler): - - def __init__(self): - super().__init__() - self.results = [] - - def report_layer_time(self, layer_name, ms): - self.results.append((layer_name, ms)) - - -def _contiguous_tile_beam_width(tensor: torch.Tensor, size: int, - num_beams: int): - new_shape = list(tensor.shape) - new_shape[0] *= num_beams - - numel = tensor.numel() - new_tensor = torch.empty(num_beams * numel, - device=tensor.device, - dtype=tensor.dtype) - - # Take the first 'size' values to tile and skip the others. - vals = tensor.view(-1)[:size] - for i in range(num_beams): - new_tensor[i * size:(i + 1) * size] = vals - - return new_tensor.view(new_shape) - - -class _Runtime(object): - runtime_rank: int - runtime: trt.Runtime - engine: trt.ICudaEngine - ctx_context: trt.IExecutionContext - context_0: trt.IExecutionContext - context_1: trt.IExecutionContext - profiler: _Profiler - engine_inspector: trt.EngineInspector - cuda_graph_instances: List[cudart.cudaGraphExec_t] - input_tensor_names: Set[str] - output_tensor_names: Set[str] - - def __init__(self, engine_buffer, mapping: Mapping): - self.address = None - self.device_memory_size = 0 - self.__prepare(mapping, engine_buffer) - - def _serialize_engine(self) -> trt.IHostMemory: - return self.engine.serialize() - - def __create_and_setup_context(self, address, size, profile_idx, - stream) -> trt.IExecutionContext: - context = self.engine.create_execution_context_without_device_memory() - assert context is not None, "Failed to create an execution context with the provided device memory!" - context.set_device_memory(address, size) - context.set_optimization_profile_async(profile_idx, stream) - # If nvtx verbosity is DETAILED, change it to LAYER_NAMES_ONLY for inference performance - if context.nvtx_verbosity == trt.ProfilingVerbosity.DETAILED: - context.nvtx_verbosity = trt.ProfilingVerbosity.LAYER_NAMES_ONLY - return context - - def _set_profiler(self): - if self.profiler is not None: - return - assert self.context_0 is not None - assert self.context_1 is not None - self.profiler = _Profiler() - self.context_0.profiler = self.profiler - self.context_0.enqueue_emits_profile = False - self.context_1.profiler = self.profiler - self.context_1.enqueue_emits_profile = False - if self.engine.num_optimization_profiles == 2: - assert self.ctx_context is not None - self.ctx_context.profiler = self.profiler - self.ctx_context.enqueue_emits_profile = False - - def __prepare(self, mapping: Mapping, engine_buffer): - self.runtime_rank = mapping.rank - local_rank = self.runtime_rank % mapping.gpus_per_node - if DISABLE_TORCH_DEVICE_SET: - CUASSERT(cudart.cudaSetDevice(torch.cuda.current_device())) - else: - torch.cuda.set_device(local_rank) - CUASSERT(cudart.cudaSetDevice(local_rank)) - - self.runtime = trt.Runtime(logger.trt_logger) - self.engine = self.runtime.deserialize_cuda_engine(engine_buffer) - assert self.engine is not None - - self.input_tensor_names = set() - self.output_tensor_names = set() - for i in range(self.engine.num_io_tensors): - name = self.engine.get_tensor_name(i) - if self.engine.get_tensor_mode(name) == trt.TensorIOMode.OUTPUT: - self.output_tensor_names.add(name) - else: - self.input_tensor_names.add(name) - - self.profiler = None - self.engine_inspector = self.engine.create_engine_inspector() - # cuda graph ping-pong instances - self.cuda_graph_instances = [None for _ in range(2)] - if not self.engine.streamable_weights_size: - # engine does not have weight streaming enabled - self.__prepare_execution_contexts() - else: - self.engine.weight_streaming_budget_v2 = 0 # avoid OOM when print engine info - - if logger.level == "verbose": - self.__print_engine_info() - - def __prepare_execution_contexts(self): - self.context_0 = None - self.context_1 = None - self.ctx_context = None - - # The device_memory_size_v2 stores the memory required by the largest profile. - # When weight streaming is enable, it must be queried after the weight streaming budget set. - if self.address: - if self.device_memory_size != self.engine.device_memory_size_v2: - self.device_memory_size = self.engine.device_memory_size_v2 - CUASSERT(cudart.cudaFree(self.address)) - address = CUASSERT(cudart.cudaMalloc( - self.device_memory_size))[0] - self.address = address - else: - self.device_memory_size = self.engine.device_memory_size_v2 - address = CUASSERT(cudart.cudaMalloc(self.device_memory_size))[0] - self.address = address - - with _scoped_stream() as stream: - if self.engine.num_optimization_profiles == 1: - # At step = 0, context_1 is active - # At step = 1, context_0 is active - # At step = 2, context_1 is active - self.context_0 = self.__create_and_setup_context( - self.address, self.device_memory_size, 0, stream) - self.context_1 = self.__create_and_setup_context( - self.address, self.device_memory_size, 0, stream) - self.ctx_context = self.context_1 - elif self.engine.num_optimization_profiles == 2: - # At step = 0, ctx_context is active - # At step = 1, context_0 is active - # At step = 2, context_1 is active - self.ctx_context = self.__create_and_setup_context( - self.address, self.device_memory_size, 0, stream) - self.context_0 = self.__create_and_setup_context( - self.address, self.device_memory_size, 1, stream) - self.context_1 = self.__create_and_setup_context( - self.address, self.device_memory_size, 1, stream) - else: - logger.error( - f"Number of optimization profiles: {self.engine.num_optimization_profiles}" - ) - raise NotImplementedError( - "Python runtime only support 1 or 2 optimization profiles, " - "set --multiple_profiles=disable when calling trtllm-build " - "to disable the feature.") - - def __print_engine_info(self) -> None: - engine = self.engine - context = engine.create_execution_context( - trt.ExecutionContextAllocationStrategy.USER_MANAGED) - n_op = engine.num_optimization_profiles - max_name_width = 0 # Maximum Width of tensor Name - max_shape_width = 0 # Maximum Width of tensor Shape - tensor_name_list = [ - engine.get_tensor_name(i) for i in range(engine.num_io_tensors) - ] - - # Get information of engine input / output - tid = {} # Tensor Information Dictionary - for name in tensor_name_list: - item = dict() - max_name_width = max(max_name_width, len(name)) - item["mode"] = 'I' if engine.get_tensor_mode( - name) == trt.TensorIOMode.INPUT else 'O' - item["location"] = 'GPU' if engine.get_tensor_location( - name) else 'CPU' - item["data_type"] = str(engine.get_tensor_dtype(name))[9:] - item["build_shape"] = str(engine.get_tensor_shape(name)) - item["profile_list"] = [[] for _ in range(n_op)] - if item["mode"] == "I": - for k in range(n_op): - if item["location"] == "GPU": - shape = engine.get_tensor_profile_shape(name, k) - else: - shape = engine.get_tensor_profile_value(k, name) - item["profile_list"][k].extend(shape) - max_shape_width = max(max_shape_width, - *[len(str(s)) for s in shape]) - tid[name] = item - # Set input shape to get output shape - for k in range(n_op): - for j in range(3): # Min, Opt, Max - for name in tid.keys(): - if tid[name]["mode"] == "I": - if tid[name]["location"] == "GPU": - context.set_input_shape( - name, tid[name]["profile_list"][k][j]) - else: - context.set_tensor_address( - name, - tid[name]["profile_list"][k][j].ctypes.data) - elif tid[name]["mode"] == "O": - assert context.all_binding_shapes_specified and context.all_shape_inputs_specified - shape = context.get_tensor_shape(name) - tid[name]["profile_list"][k].append(shape) - - # Print information of engine input / output - logger.debug("Information of engine input / output.") - logger.debug(f"{'='*(max_name_width + max_shape_width + 24)}") - logger.debug( - f"{'Name':^{max_name_width}}|I/O|Location|DataType|{'Shape':^{max_shape_width}}|" - ) - logger.debug(f"{'-'*(max_name_width + max_shape_width + 24)}") - for name in tensor_name_list: - item = tid[name] - info = f"{name:<{max_name_width}}|{item['mode']:^3s}|{item['location']:^8s}|{item['data_type']:^8s}|" - info += f"{item['build_shape']:^{max_shape_width}}|" - logger.debug(info) - logger.debug(f"{'='*(max_name_width + max_shape_width + 24)}") - # Print information of optimization profile - logger.debug("Information of optimization profile.") - for k in range(n_op): - logger.debug(f"Optimization Profile {k}:") - logger.debug(f"{'='*(max_name_width + max_shape_width * 3 + 4)}") - logger.debug( - f"{'Name':^{max_name_width}}|{'Min':^{max_shape_width}}|{'Opt':^{max_shape_width}}|{'Max':^{max_shape_width}}|" - ) - logger.debug(f"{'-'*(max_name_width + max_shape_width * 3 + 4)}") - for name in tensor_name_list: - item = tid[name] - info = f"{name:<{max_name_width}}|" - info += f"{str(item['profile_list'][k][0]):^{max_shape_width}}|" - info += f"{str(item['profile_list'][k][1]):^{max_shape_width}}|" - info += f"{str(item['profile_list'][k][2]):^{max_shape_width}}|" - logger.debug(info) - logger.debug(f"{'='*(max_name_width + max_shape_width * 3 + 4)}") - - def print_context_info(self, context, context_index) -> None: - n_io = self.engine.num_io_tensors - max_name_width = 0 # Maximum Width of tensor Name - max_shape_width = 0 # Maximum Width of tensor Shape - tensorInfo = {} - for i in range(n_io): - name = self.engine.get_tensor_name(i) - b_input = self.engine.get_tensor_mode( - name) == trt.TensorIOMode.INPUT - shape = str(self.engine.get_tensor_shape(name)) - tensorInfo[i] = [name, b_input, shape] - max_name_width = max(max_name_width, len(name)) - max_shape_width = max(max_shape_width, len(shape)) - # Shape input tensor is not used in TRT-LLM yet - - logger.debug(f"Information of context input / output.") - logger.debug(f"Using Optimization Profile: {context_index}") - logger.debug(f"{'='*(max_name_width + max_shape_width + 6)}") - logger.debug( - f"{'Name':^{max_name_width}}|I/O|{'Shape':^{max_shape_width}}|") - logger.debug(f"{'-'*(max_name_width + max_shape_width + 6)}") - for i in range(n_io): - name, b_input, shape = tensorInfo[i] - info = f"{name:<{max_name_width}}|{'I' if b_input else 'O':^3s}|{shape:^{max_shape_width}}|" - logger.debug(info) - logger.debug(f"{'='*(max_name_width + max_shape_width + 6)}") - - def _set_shape(self, context: trt.IExecutionContext, - shape_dict: Dict[str, List[int]]): - for i in range(self.engine.num_io_tensors): - name = self.engine.get_tensor_name(i) - if name not in shape_dict: - # shape and buffer can be set by calling _set_tensors API - continue - if self.engine.get_tensor_mode(name) == trt.TensorIOMode.INPUT: - ok = context.set_input_shape(name, shape_dict[name]) - dtype = self.engine.get_tensor_dtype(name) - logger.debug( - f"setting input tensor {name} with shape {shape_dict[name]} and type {dtype}" - ) - if not ok: - raise ValueError( - f"Couldn't assign {name} with shape {shape_dict[name]}, " - f"engine supports [min, opt, max] = {self.engine.get_tensor_profile_shape(name, context.active_optimization_profile)}" - ) - - def _set_buffer(self, context: trt.IExecutionContext, - buffer_dict: Dict[str, torch.Tensor]): - for i in range(self.engine.num_io_tensors): - name = self.engine.get_tensor_name(i) - if name not in buffer_dict.keys(): - dtype = self.engine.get_tensor_dtype(name) - shape = context.get_tensor_shape(name) - buffer_dict[name] = torch.zeros(tuple(shape), - dtype=trt_dtype_to_torch(dtype), - device='cuda') - assert buffer_dict[name].is_contiguous( - ), f"{name} is not contiguous()" - context.set_tensor_address(name, buffer_dict[name].data_ptr()) - - def _set_tensors(self, context: trt.IExecutionContext, - tensors: Dict[str, "RuntimeTensor"]): - for name in self.input_tensor_names: - # it's allowed to call set_tensors multi times with different tensors - # each time only set some of the engine tensors, so it is valid to skip the ones not in the current given tensors dict - if name not in tensors: - continue - - tensor = tensors[name] - if context.get_tensor_address(name) != tensor.data: - context.set_tensor_address(name, tensor.data) - - if list(context.get_tensor_shape(name)) != tensor.shape: - context.set_input_shape(name, tensor.shape) - - for name in self.output_tensor_names: - if name not in tensors: - dtype = self.engine.get_tensor_dtype(name) - shape = context.get_tensor_shape(name) - tensors[name] = RuntimeTensor.from_torch( - name, - torch.zeros(tuple(shape), - dtype=trt_dtype_to_torch(dtype), - device='cuda')) - t = tensors[name] - # output's shape is inference by TRT, no need to set the shape here - context.set_tensor_address(t.name, t.data) - - def _set_weight_streaming(self, gpu_weights_percent): - if not self.engine.streamable_weights_size: - assert gpu_weights_percent == 1, "Engine built without weight streaming. Cannot set gpu_weights_percent to a value other than 1." - return - - assert self.engine is not None - self.context_0 = None - self.context_1 = None - self.ctx_context = None - - min = 0 - max = self.engine.streamable_weights_size - budget = int(gpu_weights_percent * max) - self.engine.weight_streaming_budget_v2 = budget - assert self.engine.weight_streaming_budget_v2 == budget, "Failed to set weight streaming budget!" - logger.info( - f"Set gpu weights percent to {gpu_weights_percent}, which is {budget} bytes. Valid range: {min} bytes ~ {max} bytes." - ) - - try: - self.__prepare_execution_contexts() - except: - free_mem = torch.cuda.mem_get_info()[0] - if free_mem < budget: - print( - f"Failed to create context. Possibly out of memory: Memory budget is {budget} bytes but only {free_mem} bytes are available on the GPU." - ) - raise - - def _check_tensors(self, context: trt.IExecutionContext) -> None: - tensors = [] - for i in range(self.engine.num_io_tensors): - name = self.engine.get_tensor_name(i) - ptr = context.get_tensor_address(name) - if ptr == 0: - raise RuntimeError(f"Engine I/O tensor {name} is unbound") - shp = list(context.get_tensor_shape(name)) - if any([s < 0 for s in shp]): # skip if shape is not available - continue - dt = self.engine.get_tensor_dtype(name) - tdt = trt_dtype_to_torch(dt) - sz = torch.tensor([], dtype=tdt).element_size() * np.prod(shp) - tensors.append((ptr, ptr + sz, name, shp, sz)) - tensors.sort() # sort by start address - starts, ends, names, _, _ = zip(*tensors) - starts = torch.tensor(starts) - ends = torch.tensor(ends) - overalps = (torch.nonzero((starts[1:] < ends[:-1]).int()) + 1).squeeze() - if overalps.ndim == 0: - # unsqueeze if there is a single value so it became scalar - overalps = torch.unsqueeze(overalps, 0) - if overalps.numel() > 0: - assert overalps.ndim == 1 - for i in list(overalps): - left_name = names[i] - right_name = names[i - 1] - if "key_value" in left_name and "key_value" in right_name: # kv - left_names = left_name.split("_") - right_names = right_name.split("_") - if left_names[-1] == right_names[-1]: # same kv layer - assert (left_names[0] == "past" and right_names[0] == "present") or ( - left_names[0] == "present" and right_names[0] == "past"), \ - f"Overlap found between {tensors[i]} and {tensors[i-1]}" - continue - logger.warning( - f"TENSOR BUFFER OVERLAP DETECTED: {tensors[i]} and {tensors[i-1]} !!!" - ) - return - - def _insert_step_to_profiler(self, step: int): - if not self.profiler: - raise RuntimeError("Profiler is disable") - self.profiler.results.append(("step", step)) - - def _is_profiling(self): - return self.profiler is not None - - def _run(self, - context: trt.IExecutionContext, - stream: Union[int, torch.cuda.Stream] = None) -> bool: - if stream is None: - stream = torch.cuda.current_stream().cuda_stream - elif isinstance(stream, torch.cuda.Stream): - stream = stream.cuda_stream - ok = context.execute_async_v3(stream) - return ok - - def __del__(self): - try: - if self.address is not None: - cudart.cudaFree(self.address) - except TypeError: - pass - - @property - def context_mem_size(self) -> int: - return self.engine.device_memory_size_v2 - - -@dataclass -class ModelConfig: - max_batch_size: int - max_beam_width: int - vocab_size: int - num_layers: int - num_heads: int - num_kv_heads: int - hidden_size: int - gpt_attention_plugin: bool - gemm_allreduce_plugin: str = None - remove_input_padding: bool = False - model_name: str = "" - kv_cache_type: KVCacheType = KVCacheType.CONTINUOUS - cross_attention: bool = False - head_size: int = None - has_position_embedding: bool = True - has_token_type_embedding: bool = False - tokens_per_block: int = 32 - max_prompt_embedding_table_size: int = 0 - quant_mode: QuantMode = QuantMode(0) - gather_context_logits: bool = False - gather_generation_logits: bool = False - dtype: str = "" - lora_plugin: bool = False - lora_target_modules: List[str] = field(default_factory=list) - trtllm_modules_to_hf_modules: dict = None - skip_cross_kv: bool = False - num_medusa_heads: int = 0 - max_medusa_tokens: int = 0 - paged_state: bool = True - mamba_conv1d_plugin: bool = True - conv_kernel: int = 0 - layer_types: List[str] = field(default_factory=list) - rnn_hidden_size: int = 0 - rnn_head_size: int = 0 - rnn_conv_dim_size: int = 0 - state_size: int = 0 - state_dtype: str = "" - gpu_weights_percent: float = 1.0 - # ReDrafter - redrafter_num_beams: int = 0 - redrafter_draft_len_per_beam: int = 0 - num_kv_heads_per_layer: Optional[List[int]] = None - num_kv_heads_per_cross_attn_layer: Optional[List[int]] = None - skip_cross_attn_blocks: bool = False - # language adapter - language_adapter_config: Optional[LanguageAdapterConfig] = None - - @classmethod - def from_model_config_cpp(cls, model_config_cpp) -> 'ModelConfig': - """Create a partially initialized ModelConfig instance from a given ModelConfig CPP binding instance. - - Note that each of these classes have fields that don't exist in the other, so the created ModelConfigPython - won't have all of its fields initialized. - """ - return cls( - max_batch_size=model_config_cpp.max_batch_size, - max_beam_width=model_config_cpp.max_beam_width, - vocab_size=model_config_cpp.vocab_size, - num_layers=model_config_cpp.num_layers(), - num_heads=model_config_cpp.num_heads, - num_kv_heads=model_config_cpp.num_kv_heads(0), - hidden_size=model_config_cpp.hidden_size, - remove_input_padding=model_config_cpp.use_packed_input, - kv_cache_type=model_config_cpp.kv_cache_type, - cross_attention=model_config_cpp.use_cross_attention, - head_size=model_config_cpp.head_size, - max_prompt_embedding_table_size=model_config_cpp. - max_prompt_embedding_table_size, - quant_mode=QuantMode(model_config_cpp.quant_mode.value), - gather_context_logits=model_config_cpp.compute_context_logits, - gather_generation_logits=model_config_cpp.compute_generation_logits, - gpt_attention_plugin=model_config_cpp.use_gpt_attention_plugin, - dtype=binding_to_str_dtype(model_config_cpp.data_type), - num_kv_heads_per_layer=model_config_cpp.num_kv_heads_per_layer, - tokens_per_block=model_config_cpp.tokens_per_block, - lora_plugin=model_config_cpp.use_lora_plugin, - layer_types=[ - binding_layer_type_to_str(lt) - for lt in model_config_cpp.layer_types - ], - ) - - -@dataclass -class SamplingConfig: - end_id: int - pad_id: int - - max_new_tokens: int = field(default=20) - num_beams: int = field(default=1) - num_return_sequences: Optional[int] = field(default=None) - max_attention_window_size: Optional[int] = field(default=None) - sink_token_length: Optional[int] = field(default=None) - output_sequence_lengths: bool = field(default=False) - return_dict: bool = field(default=False) - stop_words_list: Optional[Union[list, np.ndarray, - torch.Tensor]] = field(default=None) - bad_words_list: Optional[Union[list, np.ndarray, - torch.Tensor]] = field(default=None) - - temperature: Union[float, torch.Tensor] = field(default=1.0) - top_k: Union[int, torch.Tensor] = field(default=1) - top_p: Union[float, torch.Tensor] = field(default=0.0) - top_p_decay: Optional[torch.Tensor] = field(default=None) # float - top_p_min: Optional[torch.Tensor] = field(default=None) # float - top_p_reset_ids: Optional[torch.Tensor] = field(default=None) # int - random_seed: Union[int, torch.Tensor] = field(default=None) - - length_penalty: Union[float, torch.Tensor] = field(default=1.0) - early_stopping: Union[int, torch.Tensor] = field(default=1) - repetition_penalty: Union[float, torch.Tensor] = field(default=1.0) - min_length: Union[int, torch.Tensor] = field(default=1) - presence_penalty: Union[float, torch.Tensor] = field(default=0.0) - frequency_penalty: Union[float, torch.Tensor] = field(default=0.0) - prompt_ignore_length: Union[int, torch.Tensor] = field(default=0) - use_beam_hyps: bool = field(default=True) - - # None here means user didn't set it, and dynamicDecodeOp.cpp take optional value - # The real default value is set in dynamicDecodeOp.cpp when it's None - beam_search_diversity_rate: Union[float, torch.Tensor] = field(init=False, - default=0.0) - output_cum_log_probs: bool = field(init=False, default=False) - output_log_probs: bool = field(init=False, default=False) - no_repeat_ngram_size: Union[int, torch.Tensor] = field(init=False, - default=None) - min_p: Union[float, torch.Tensor] = field(default=0.0) - - def update(self, **kwargs): - unused_kwargs = dict() - for key, value in kwargs.items(): - if hasattr(self, key): - setattr(self, key, value) - else: - unused_kwargs[key] = value - return unused_kwargs - - -class LogitsProcessor: - """ - Base class for all logit processors that can be applied during generation. - """ - - def __call__(self, step: int, input_ids: torch.Tensor, - scores: torch.Tensor) -> torch.Tensor: - raise NotImplementedError( - f"{self.__class__} is an abstract class. Only classes inheriting this class can be called." - ) - - -class LogitsProcessorList(list, LogitsProcessor): - - def __call__(self, step: int, input_ids: torch.Tensor, - scores: torch.Tensor) -> torch.Tensor: - for processor in self: - scores = processor(step, input_ids, scores) - return scores - - -class StoppingCriteria: - """ - Base class for all stopping criteria that can be applied during generation. - """ - - def __call__(self, step: int, input_ids: torch.Tensor, - scores: torch.Tensor) -> bool: - raise NotImplementedError("StoppingCriteria needs to be subclassed") - - -class StoppingCriteriaList(list, StoppingCriteria): - - def __call__(self, step: int, input_ids: torch.Tensor, - scores: torch.Tensor) -> bool: - return any(criteria(step, input_ids, scores) for criteria in self) - - -class RuntimeTensor: - - def __init__(self): - self._name = "" - # shape is the one sent to TRT, the actual torch tensor can be larger than the shape - # this is useful when allocating a big KV cache tensor at the beginning and incremental seq length dim of TRT engine's input tensor - self._shape = None - self._torch_tensor = None - # Used when pointer specified - self._data_ptr = None - self._dtype = None - - @staticmethod - def from_pointer(name: str, pointer, shape, - str_dtype: str) -> 'RuntimeTensor': - t = RuntimeTensor() - t._name = name - t._data_ptr = pointer - t._shape = shape - t._dtype = str_dtype_to_torch(str_dtype) - return t - - @staticmethod - def from_torch( - name: str, - data: torch.Tensor, - override_shape: Optional[Iterable] = None) -> 'RuntimeTensor': - assert (isinstance(data, torch.Tensor)), f"data {name} is {type(data)}" - t = RuntimeTensor() - t._name = name - # need to hold the torch tensor for memory life time - t._torch_tensor = data.contiguous() - t._dtype = t._torch_tensor.dtype - t._data_ptr = t._torch_tensor.data_ptr() - torch_shape = list(data.size()) - if override_shape is not None: - t._shape = override_shape - assert isinstance(override_shape, list) or isinstance( - override_shape, tuple) - assert all([lambda x: x >= 0 for x in override_shape - ]), f"Expect all dimensions >=0, got {override_shape}" - - def volume_func(dims): - return reduce(lambda x, y: x * y, dims, 1) - assert volume_func(override_shape) <= volume_func(torch_shape), \ - f"Override the shape to be larger than the underlying torch Tensor, got {override_shape}, torch tensor shape {torch_shape}" - else: - t._shape = torch_shape - return t - - def to_torch(self) -> torch.Tensor: - if self._torch_tensor is None: - raise RuntimeError( - 'RuntimeTensor cannot be converted to torch tensor as constructed from pointer' - ) - return self._torch_tensor - - @property - def shape(self) -> Iterable[int]: - return self._shape - - @property - def data(self): - return self._data_ptr - - @property - def name(self) -> str: - return self._name - - @property - def dtype(self) -> torch.dtype: - return self._dtype - - -class GenerationSession(object): - - _model_config: ModelConfig - mapping: Mapping - runtime: _Runtime - device: torch.device - batch_size: int - buffer_allocated: bool - debug_mode: bool - quant_mode: QuantMode - cuda_graph_mode: bool - dtype: trt.DataType - debug_tensors_to_save: None - num_draft_tokens: int = 0 - medusa_topks: List[int] = None - medusa_paths: List[List[int]] = None - medusa_tree_ids: List[int] = None - medusa_position_offsets: List[int] = None - medusa_temperature: float = 0.0 - - def __init__(self, - model_config: ModelConfig, - engine_buffer, - mapping: Mapping, - debug_mode=False, - debug_tensors_to_save=None, - cuda_graph_mode=False, - stream: torch.cuda.Stream = None): - assert isinstance(model_config, ModelConfig) - self._model_config = model_config - self.mapping = mapping - self.runtime = _Runtime(engine_buffer, mapping) - if DISABLE_TORCH_DEVICE_SET: - self.device = torch.device(f'cuda:{torch.cuda.current_device()}') - else: - self.device = torch.device( - f'cuda:{self.runtime.runtime_rank % mapping.gpus_per_node}') - torch.cuda.set_device(self.device) - # dynamic_decoder currently use torch's current stream, so must let TRT enqueue use same stream here - self.stream = stream - if self.stream is None: - self.stream = torch.cuda.Stream(self.device) - torch.cuda.set_stream(self.stream) - self.debug_mode = debug_mode - self.debug_tensors_to_save = debug_tensors_to_save - - self.cuda_graph_mode = cuda_graph_mode - # Optional inputs for dynamic decoder - self.top_p_decay = None - self.top_p_min = None - self.top_p_reset_ids = None - # TODO: in tensorrt_llm/cpp/tensorrt_llm/thop/dynamicDecodeOp.cpp it's T, can be float or half? - self.embedding_bias_opt = None - - self.buffer = None - self.buffer_allocated = False - - self.vocab_size_padded = pad_vocab_size(self.vocab_size, - self.mapping.tp_size) - if len(model_config.layer_types) == 0: - self.layer_types = ['attention'] * model_config.num_layers - else: - layer_types = model_config.layer_types - layer_types = layer_types * (model_config.num_layers // - len(layer_types)) - layer_types = layer_types + layer_types[0:(model_config.num_layers % - len(layer_types))] - self.layer_types = layer_types - self.num_attn_layers = \ - self.layer_types[self.first_layer:self.last_layer].count('attention') - self.has_attn_layers = self.num_attn_layers > 0 - self.has_rnn_layers = 'recurrent' in self.layer_types[ - self.first_layer:self.last_layer] - - self.attn_to_general_idx = {} - self.general_to_attn_idx = {} - attn_layer_idx = 0 - for i in range(self.first_layer, self.last_layer): - if self.layer_types[i] == 'attention': - self.attn_to_general_idx[attn_layer_idx] = i - self.general_to_attn_idx[i] = attn_layer_idx - attn_layer_idx += 1 - - # Cyclic KV cache buffer names. - if self.attn_to_general_idx: - self.kv_cache_buffer_names = [ - f'present_key_value_{layer_idx}' - for _, layer_idx in self.attn_to_general_idx.items() - ] + [f'1_present_key_value_{self.attn_to_general_idx[0]}'] - else: - self.kv_cache_buffer_names = [] - - if self.paged_kv_cache: - logger.warning( - "The paged KV cache in Python runtime is experimental. For performance and correctness, please, use C++ runtime." - ) - - if self.mapping.has_pp(): - self.nccl_comm = torch.classes.trtllm.NcclCommunicatorOp( - self.mapping.world_size, self.mapping.rank) - - if self.mapping.is_last_pp_rank(): - self.decoder_logits_dtype = self._tensor_dtype('logits') - if self.decoder_logits_dtype not in [torch.float16, torch.float32]: - logger.warning( - "Logits dtype not supported by decoder. Falling back to float32. You may want to change the logits dtype to float16 in your model definition." - ) - self.decoder_logits_dtype = torch.float32 - self.dynamic_decoder = torch.classes.trtllm.DynamicDecodeOp( - model_config.max_batch_size, model_config.max_beam_width, - self.vocab_size, self.vocab_size_padded, self.mapping.tp_size, - self.mapping.pp_size, self.decoder_logits_dtype) - - expected_tensor_names = [] - if self.mapping.tp_size > 1: - self.ipc_buffers, self.all_reduce_workspace = CustomAllReduceHelper.allocate_workspace( - self.mapping, - CustomAllReduceHelper.max_workspace_size_auto( - self.mapping.tp_size)) - - self.gather_tree = torch.ops.tensorrt_llm.gather_tree - - if self.mapping.is_first_pp_rank(): - expected_tensor_names += ['input_ids'] - else: - expected_tensor_names += ['hidden_states_input'] - - if self.mapping.is_last_pp_rank(): - expected_tensor_names += ['logits'] - if not model_config.gather_context_logits or self.has_rnn_layers: - expected_tensor_names += ['last_token_ids'] - else: - expected_tensor_names += ['hidden_states_output'] - - if self.has_attn_layers: - if model_config.has_position_embedding and self.mapping.is_first_pp_rank( - ): - expected_tensor_names += ['position_ids'] - if model_config.has_token_type_embedding and self.mapping.is_first_pp_rank( - ): - expected_tensor_names += ['token_type_ids'] - - if self.use_kv_cache: - expected_tensor_names += ['cache_indirection'] - - if self.paged_kv_cache and self.has_attn_layers: - expected_tensor_names += [f'kv_cache_block_offsets'] - expected_tensor_names += [f'host_kv_cache_block_offsets'] - expected_tensor_names += [f'host_kv_cache_pool_pointers'] - expected_tensor_names += [f'host_kv_cache_pool_mapping'] - if self.cross_attention: - expected_tensor_names += [f'cross_kv_cache_block_offsets'] - expected_tensor_names += [f'host_cross_kv_cache_block_offsets'] - expected_tensor_names += [f'host_cross_kv_cache_pool_pointers'] - expected_tensor_names += [f'host_cross_kv_cache_pool_mapping'] - expected_tensor_names += [f'cross_attention_mask'] - expected_tensor_names += [f'cross_attention_packed_mask'] - else: - # Refer to gpt_attention() inside functional.py - if self.use_kv_cache and not self.paged_kv_cache: - for i in range(self.first_layer, self.last_layer): - if self.layer_types[i] == 'attention': - expected_tensor_names += [ - f'past_key_value_{i}', f'present_key_value_{i}' - ] - if model_config.cross_attention: - if model_config.gpt_attention_plugin: - for i in range(self.first_layer, self.last_layer): - if self.layer_types[i] == 'attention': - expected_tensor_names += [ - f'cross_present_key_value_{i}', - f'cross_past_key_value_{i}' - ] - expected_tensor_names += [ - 'cross_attention_mask', - ] - expected_tensor_names += [f'cross_attention_packed_mask'] - else: - expected_tensor_names += [ - 'cross_attention_mask', - ] - - if self.paged_state and self.has_rnn_layers: - for i in range(self.first_layer, self.last_layer): - if self.layer_types[i] == 'recurrent': - expected_tensor_names += [ - f'conv_state_ptr_{i}', f'rnn_state_ptr_{i}' - ] - expected_tensor_names += ['slot_mapping'] - else: - for i in range(self.first_layer, self.last_layer): - if self.layer_types[i] == 'recurrent': - expected_tensor_names += [ - f'past_conv_state_{i}', f'present_conv_state_{i}', - f'past_rnn_state_{i}', f'present_rnn_state_{i}' - ] - - if model_config.gpt_attention_plugin and self.has_attn_layers: - if self.use_kv_cache: - expected_tensor_names += [ - 'sequence_length', 'host_past_key_value_lengths' - ] - - expected_tensor_names += [ - 'context_lengths', 'host_request_types', - 'host_sink_token_length', 'host_runtime_perf_knobs', - 'host_context_progress' - ] - expected_tensor_names += [f'host_max_attention_window_sizes'] - if model_config.remove_input_padding: - expected_tensor_names.append('host_context_lengths') - else: - if self.has_rnn_layers: - expected_tensor_names += ['host_request_types'] - if model_config.mamba_conv1d_plugin and model_config.remove_input_padding: - expected_tensor_names.append('host_context_lengths') - if self.has_attn_layers: - expected_tensor_names += ['attention_mask'] - - if model_config.max_prompt_embedding_table_size > 0: - expected_tensor_names += [ - 'prompt_embedding_table', 'tasks', 'prompt_vocab_size' - ] - - if model_config.cross_attention: - expected_tensor_names += [ - 'encoder_output', - 'encoder_input_lengths', - 'encoder_max_input_length', - 'cross_kv_cache_gen', - ] - if model_config.skip_cross_attn_blocks: - expected_tensor_names += ['skip_cross_attn_blocks'] - self.skip_cross_kv = model_config.skip_cross_kv - if self.skip_cross_kv: - expected_tensor_names += ['cross_kv_reuse'] - - if self.mapping.tp_size > 1: - expected_tensor_names += ['all_reduce_workspace'] - - self.lora_target_modules = model_config.lora_target_modules - self.missing_qkv_modules = LoraManager.get_missing_qkv_modules( - self.lora_target_modules) - if model_config.lora_plugin: - for lora_module in (self.lora_target_modules + - self.missing_qkv_modules): - for i in range(self.first_layer, self.last_layer): - expected_tensor_names += [ - f'{lora_module}_lora_ranks_{i}', - f'{lora_module}_lora_weights_pointers_{i}' - ] - if self.cross_attention and self.remove_input_padding: - expected_tensor_names += ['host_encoder_input_lengths'] - - if model_config.num_medusa_heads > 0: - expected_tensor_names += [ - 'spec_decoding_generation_lengths', - 'spec_decoding_position_offsets', 'spec_decoding_packed_mask', - 'spec_decoding_use', 'medusa_logits' - ] - - if self.is_redrafter_mode: - expected_tensor_names += get_redrafter_tensor_names() - - # language adapter - if model_config.language_adapter_config: - expected_tensor_names += ['language_adapter_routings'] - - found_tensor_names = [ - self.runtime.engine.get_tensor_name(i) - for i in range(self.runtime.engine.num_io_tensors) - ] - for name in found_tensor_names: - if name.startswith("allreduce_ub_") or name.startswith( - "gemm_allreduce"): - expected_tensor_names += [name] - if not self.debug_mode and set(expected_tensor_names) != set( - found_tensor_names): - logger.error( - f"The following expected tensors are not found: {set(expected_tensor_names).difference(set(found_tensor_names))}" - ) - logger.error( - f"Those tensors in engine are not expected: {set(found_tensor_names).difference(set(expected_tensor_names))}" - ) - logger.error(f"Expected tensor names: {expected_tensor_names}") - logger.error(f"Found tensor names: {found_tensor_names}") - raise RuntimeError( - "Tensor names in engine are not the same as expected, to use this GenerationSession, " - "you need to use PretrainedModel.prepare_inputs to create TRT Network inputs." - ) - if self.debug_mode: - self.debug_tensors = list( - set(found_tensor_names) - set(expected_tensor_names)) - if self.debug_tensors_to_save is None: - self.debug_tensors_to_save = self.debug_tensors - logger.info(f"Debug tensors found: {self.debug_tensors}") - logger.info(f"Debug tensors to save: {self.debug_tensors_to_save}") - - def __del__(self): - try: - if self.use_gemm_allreduce_plugin: - assert self.gemm_allreduce_output_handle is not None - ipc_nvls_free(self.gemm_allreduce_output_handle) - except TypeError: - pass - - @property - def context_mem_size(self) -> int: - return self.runtime.context_mem_size - - @property - def vocab_size(self): - return self._model_config.vocab_size - - @property - def num_layers(self): - assert self._model_config.num_layers % self.mapping.pp_size == 0, \ - f"num_layers {self._model_config.num_layers} must be a multiple of pipeline parallelism size {self.mapping.pp_size}" - return self._model_config.num_layers // self.mapping.pp_size - - @property - def first_layer(self): - return self.num_layers * self.mapping.pp_rank - - @property - def last_layer(self): - return self.first_layer + self.num_layers - - @property - def num_heads(self): - return self._model_config.num_heads - - @property - def hidden_size(self): - # For linear layer in attention block - return self._model_config.hidden_size - - @property - def use_gpt_attention_plugin(self): - return self._model_config.gpt_attention_plugin - - @property - def use_mamba_conv1d_plugin(self): - return self._model_config.mamba_conv1d_plugin - - @property - def paged_kv_cache(self): - return self._model_config.kv_cache_type == KVCacheType.PAGED - - @property - def kv_cache_type(self): - return self._model_config.kv_cache_type - - @property - def use_kv_cache(self): - return self._model_config.kv_cache_type != KVCacheType.DISABLED - - @property - def tokens_per_block(self): - return self._model_config.tokens_per_block - - @property - def remove_input_padding(self): - return self._model_config.remove_input_padding - - def get_num_heads_kv(self, layer_idx: Optional[int] = None) -> int: - if layer_idx is None or self._model_config.num_kv_heads_per_layer is None: - return self._model_config.num_kv_heads - - if self._model_config.layer_types: - assert self._model_config.layer_types[ - layer_idx] == "attention", f"Layer {layer_idx} is not an attention layer" - - if self._model_config.num_kv_heads_per_layer: - return self._model_config.num_kv_heads_per_layer[layer_idx] - - return self._model_config.num_kv_heads - - @property - def head_size(self): - return self.hidden_size // self.num_heads if self._model_config.head_size is None else self._model_config.head_size - - @property - def max_prompt_embedding_table_size(self): - return self._model_config.max_prompt_embedding_table_size - - @property - def quant_mode(self): - return self._model_config.quant_mode - - @property - def gather_context_logits(self): - return self._model_config.gather_context_logits - - @property - def gather_generation_logits(self): - return self._model_config.gather_generation_logits - - @property - def dtype(self): - return str_dtype_to_torch(self._model_config.dtype) - - @property - def profiler(self): - return self.runtime.profiler - - @property - def engine_inspector(self): - return self.runtime.engine_inspector - - def cuda_stream_guard(func): - """Sync external stream and set current stream to the one bound to the session. Reset on exit. - """ - - @wraps(func) - def wrapper(self, *args, **kwargs): - external_stream = torch.cuda.current_stream() - if external_stream != self.stream: - external_stream.synchronize() - torch.cuda.set_stream(self.stream) - ret = func(self, *args, **kwargs) - if external_stream != self.stream: - self.stream.synchronize() - torch.cuda.set_stream(external_stream) - return ret - - return wrapper - - @property - def cross_attention(self): - return self._model_config.cross_attention - - @property - def has_position_embedding(self): - return self._model_config.has_position_embedding - - @property - def has_token_type_embedding(self): - return self._model_config.has_token_type_embedding - - @property - def use_lora_plugin(self): - return self._model_config.lora_plugin - - @property - def use_gemm_allreduce_plugin(self): - return bool(self._model_config.gemm_allreduce_plugin) - - @property - def gemm_allreduce_plugin(self): - return self._model_config.gemm_allreduce_plugin - - @property - def is_medusa_mode(self): - return self.num_medusa_heads > 0 - - @property - def is_redrafter_mode(self): - return self._model_config.redrafter_num_beams > 0 and self._model_config.redrafter_draft_len_per_beam > 0 - - @property - def max_draft_tokens(self): - if self.is_redrafter_mode: - return self._model_config.redrafter_num_beams * self._model_config.redrafter_draft_len_per_beam - return self._model_config.max_medusa_tokens - - @property - def num_medusa_heads(self): - return self._model_config.num_medusa_heads - - @property - def paged_state(self): - return self._model_config.paged_state - - @property - def conv_kernel(self): - return self._model_config.conv_kernel - - @property - def rnn_hidden_size(self): - return self._model_config.rnn_hidden_size - - @property - def rnn_head_size(self): - return self._model_config.rnn_head_size - - @property - def rnn_conv_dim_size(self): - return self._model_config.rnn_conv_dim_size - - @property - def state_size(self): - return self._model_config.state_size - - @property - def state_dtype(self): - if self._model_config.state_dtype == "": - return str_dtype_to_torch(self._model_config.dtype) - return str_dtype_to_torch(self._model_config.state_dtype) - - def _capture_cuda_graph_and_instantiate(self, context, stream, step): - instance_idx = (step + 1) % 2 - if not self.has_attn_layers: - # Create two cuda graph once.If cuda graph has already existed, skip it. - if self.runtime.cuda_graph_instances[instance_idx] is not None: - return - # capture cuda graph - CUASSERT( - cudart.cudaStreamBeginCapture( - stream, - cudart.cudaStreamCaptureMode.cudaStreamCaptureModeGlobal)) - context.execute_async_v3(stream) - next_graph = CUASSERT(cudart.cudaStreamEndCapture(stream))[0] - - if self.runtime.cuda_graph_instances[instance_idx] is not None: - self.runtime.cuda_graph_instances[ - instance_idx] = _update_cuda_graph_instance( - self.runtime.cuda_graph_instances[instance_idx], next_graph) - else: - self.runtime.cuda_graph_instances[instance_idx] = CUASSERT( - cudart.cudaGraphInstantiate(next_graph, 0))[0] - - # Pre-upload cuda graph to stream - CUASSERT( - cudart.cudaGraphUpload( - self.runtime.cuda_graph_instances[instance_idx], stream)) - - def __setup_decoder(self, input_ids: torch.Tensor, - sampling_config: SamplingConfig, - host_context_lengths: torch.Tensor): - '''Allocate buffers and setup the post-processing decoder kernel - ''' - batch_size = host_context_lengths.shape[0] - scfg = sampling_config # just to make a shorter name, no other meaning - if isinstance(scfg.top_k, torch.Tensor): - assert scfg.top_k.dtype == torch.int32, f"scfg.top_k.dtype ({scfg.top_k.dtype}) must be torch.int32" - assert scfg.top_k.shape[ - 0] == batch_size, f"scfg.top_k.shape[0] ({scfg.top_k.shape[0]}) must equal to batch_size ({batch_size})" - self.top_k = scfg.top_k - else: - self.top_k = torch.full([batch_size], scfg.top_k, dtype=torch.int32) - - if isinstance(scfg.top_p, torch.Tensor): - assert scfg.top_p.dtype == torch.float32, f"scfg.top_p.dtype ({scfg.top_p.dtype}) must be torch.float32" - assert scfg.top_p.shape[ - 0] == batch_size, f"scfg.top_p.shape[0] ({scfg.top_p.shape[0]}) must equal to batch_size ({batch_size})" - self.top_p = scfg.top_p - else: - self.top_p = torch.full([batch_size], - scfg.top_p, - dtype=torch.float32) - - if isinstance(scfg.temperature, torch.Tensor): - assert scfg.temperature.dtype == torch.float32, f"scfg.temperature.dtype ({scfg.temperature.dtype}) must be torch.float32" - assert scfg.temperature.shape[ - 0] == batch_size, f"scfg.temperature.shape[0] ({scfg.temperature.shape[0]}) must equal to batch_size ({batch_size})" - self.temperature = scfg.temperature - else: - self.temperature = torch.full([batch_size], - scfg.temperature, - dtype=torch.float32) - - if isinstance(scfg.repetition_penalty, torch.Tensor): - assert scfg.repetition_penalty.dtype == torch.float32, f"scfg.repetition_penalty.dtype ({scfg.repetition_penalty.dtype}) must be torch.float32" - assert scfg.repetition_penalty.shape[ - 0] == batch_size, f"scfg.repetition_penalty.shape[0] ({scfg.repetition_penalty.shape[0]}) must equal to batch_size ({batch_size})" - self.repetition_penalty = scfg.repetition_penalty - elif scfg.repetition_penalty == 1.0: - self.repetition_penalty = None - else: - self.repetition_penalty = torch.full([batch_size], - scfg.repetition_penalty, - dtype=torch.float32) - - if isinstance(scfg.length_penalty, torch.Tensor): - assert scfg.length_penalty.dtype == torch.float32, f"scfg.length_penalty.dtype ({scfg.length_penalty.dtype}) must be torch.float32" - assert scfg.length_penalty.shape[ - 0] == batch_size, f"scfg.length_penalty.shape[0] ({scfg.length_penalty.shape[0]}) must equal to batch_size ({batch_size})" - self.host_length_penalty = scfg.length_penalty - else: - self.host_length_penalty = torch.full([batch_size], - scfg.length_penalty, - dtype=torch.float32) - self.length_penalty = self.host_length_penalty.to(self.device) - - if isinstance(scfg.early_stopping, torch.Tensor): - assert scfg.early_stopping.dtype == torch.int32, f"scfg.early_stopping.dtype ({scfg.early_stopping.dtype}) must be torch.int32" - assert scfg.early_stopping.shape[ - 0] == batch_size, f"scfg.early_stopping.shape[0] ({scfg.early_stopping.shape[0]}) must equal to batch_size ({batch_size})" - self.host_early_stopping = scfg.early_stopping - else: - self.host_early_stopping = torch.full([batch_size], - scfg.early_stopping, - dtype=torch.int32) - - if isinstance(scfg.presence_penalty, torch.Tensor): - assert scfg.presence_penalty.dtype == torch.float32, f"scfg.presence_penalty.dtype ({scfg.presence_penalty.dtype}) must be torch.float32" - assert scfg.presence_penalty.shape[ - 0] == batch_size, f"scfg.presence_penalty.shape[0] ({scfg.presence_penalty.shape[0]}) must equal to batch_size ({batch_size})" - self.presence_penalty = scfg.presence_penalty - elif scfg.presence_penalty == 0.0: - self.presence_penalty = None - else: - self.presence_penalty = torch.full([batch_size], - scfg.presence_penalty, - dtype=torch.float32) - - if isinstance(scfg.frequency_penalty, torch.Tensor): - assert scfg.frequency_penalty.dtype == torch.float32, f"scfg.frequency_penalty.dtype ({scfg.frequency_penalty.dtype}) must be torch.float32" - assert scfg.frequency_penalty.shape[ - 0] == batch_size, f"scfg.frequency_penalty.shape[0] ({scfg.frequency_penalty.shape[0]}) must equal to batch_size ({batch_size})" - self.frequency_penalty = scfg.frequency_penalty - elif scfg.frequency_penalty == 0.0: - self.frequency_penalty = None - else: - self.frequency_penalty = torch.full([batch_size], - scfg.frequency_penalty, - dtype=torch.float32) - - if isinstance(scfg.prompt_ignore_length, torch.Tensor): - assert scfg.prompt_ignore_length.dtype == torch.int32, f"scfg.prompt_ignore_length.dtype ({scfg.prompt_ignore_length.dtype}) must be torch.int32" - assert scfg.prompt_ignore_length.shape[ - 0] == batch_size, f"scfg.prompt_ignore_length.shape[0] ({scfg.prompt_ignore_length.shape[0]}) must equal to batch_size ({batch_size})" - self.prompt_ignore_length = scfg.prompt_ignore_length - else: - self.prompt_ignore_length = torch.full([batch_size], - scfg.prompt_ignore_length, - dtype=torch.int32) - - if isinstance(scfg.min_length, torch.Tensor): - assert scfg.min_length.dtype == torch.int32, f"scfg.min_length.dtype ({scfg.min_length.dtype}) must be torch.int32" - assert scfg.min_length.shape[ - 0] == batch_size, f"scfg.min_length.shape[0] ({scfg.min_length.shape[0]}) must equal to batch_size ({batch_size})" - self.min_length = scfg.min_length - else: - self.min_length = torch.full([batch_size], - scfg.min_length, - dtype=torch.int32) - - if isinstance(scfg.beam_search_diversity_rate, torch.Tensor): - assert scfg.beam_search_diversity_rate.dtype == torch.float32, f"scfg.beam_search_diversity_rate.dtype ({scfg.beam_search_diversity_rate.dtype}) must be torch.float32" - assert scfg.beam_search_diversity_rate.shape[ - 0] == batch_size, f"scfg.beam_search_diversity_rate.shape[0] ({scfg.beam_search_diversity_rate.shape[0]}) must equal to batch_size ({batch_size})" - self.beam_search_diversity_rate = scfg.beam_search_diversity_rate - elif scfg.beam_search_diversity_rate is not None: - self.beam_search_diversity_rate = torch.full( - [batch_size], - scfg.beam_search_diversity_rate, - dtype=torch.float32) - else: - self.beam_search_diversity_rate = None - - if isinstance(scfg.random_seed, torch.Tensor): - assert scfg.random_seed.dtype == torch.int64, f"scfg.random_seed.dtype ({scfg.random_seed.dtype}) must be torch.int64" - assert scfg.random_seed.shape[ - 0] == batch_size, f"scfg.random_seed.shape[0] ({scfg.random_seed.shape[0]}) must equal to batch_size ({batch_size})" - self.random_seed = scfg.random_seed - elif scfg.random_seed is not None: - self.random_seed = torch.full([batch_size], - scfg.random_seed, - dtype=torch.int64) - else: - self.random_seed = None - - if isinstance(scfg.no_repeat_ngram_size, torch.Tensor): - assert scfg.no_repeat_ngram_size.dtype == torch.int32, f"scfg.no_repeat_ngram_size.dtype ({scfg.no_repeat_ngram_size.dtype}) must be torch.int32" - assert scfg.no_repeat_ngram_size.shape[ - 0] == batch_size, f"scfg.no_repeat_ngram_size.shape[0] ({scfg.no_repeat_ngram_size.shape[0]}) must equal to batch_size ({batch_size})" - self.no_repeat_ngram_size = scfg.no_repeat_ngram_size - elif scfg.no_repeat_ngram_size is not None: - self.no_repeat_ngram_size = torch.full([batch_size], - scfg.no_repeat_ngram_size, - dtype=torch.int32) - else: - self.no_repeat_ngram_size = None - - if isinstance(scfg.min_p, torch.Tensor): - assert scfg.min_p.dtype == torch.float32, f"scfg.min_p.dtype ({scfg.min_p.dtype}) must be torch.float32" - assert scfg.min_p.shape[ - 0] == batch_size, f"scfg.min_p.shape[0] ({scfg.min_p.shape[0]}) must equal to batch_size ({batch_size})" - self.min_p = scfg.min_p - elif scfg.min_p == 1.0: - self.min_p = None - else: - self.min_p = torch.full([batch_size], - scfg.min_p, - dtype=torch.float32) - - if self.mapping.is_last_pp_rank(): - self.dynamic_decoder.setup( - batch_size, - scfg.num_beams, - self.top_k, - self.top_p, - self.temperature, - self.repetition_penalty, - self.presence_penalty, - self.frequency_penalty, - self.prompt_ignore_length, - self.min_length, - self.host_length_penalty, - self.host_early_stopping, - self.beam_search_diversity_rate, - self.random_seed, - self.top_p_decay, - self.top_p_min, - self.top_p_reset_ids, - self.no_repeat_ngram_size, - self.min_p, - scfg.output_log_probs, - scfg.num_beams > 1 or scfg.output_cum_log_probs, - ) - - assert scfg.end_id is not None, "end_id cannot be none" - assert scfg.pad_id is not None, 'pad_id cannot be none' - self.end_ids = torch.full((batch_size, ), - scfg.end_id, - dtype=torch.int32, - device=self.device) - max_context_length = host_context_lengths.max() - - # setup output ids buffer - if input_ids.dim() == 1: - # input_ids only have one dimension, which means remove_padding is enabled - split_ids_list = list( - torch.split(input_ids.unsqueeze(0), - host_context_lengths.numpy().tolist(), - dim=1)) - padded_input_ids = torch.nested.to_padded_tensor( - torch.nested.nested_tensor(split_ids_list, - dtype=torch.int32, - device='cuda'), - scfg.pad_id).reshape(batch_size, max_context_length) - else: - padded_input_ids = input_ids - if scfg.num_beams > 1: - tiled_input_ids = _tile_beam_width(padded_input_ids, scfg.num_beams) - tiled_input_ids = tiled_input_ids.reshape(batch_size, - scfg.num_beams, - max_context_length) - tiled_input_ids.permute(2, 0, 1) # TODO: delete? - self.output_ids = torch.cat( - (tiled_input_ids, - torch.full((batch_size, scfg.num_beams, - self.max_seq_length - max_context_length), - scfg.end_id, - dtype=padded_input_ids.dtype, - device=padded_input_ids.device)), - axis=-1) - else: - self.output_ids = torch.cat( - (padded_input_ids, - torch.full( - (batch_size, self.max_seq_length - max_context_length), - scfg.end_id, - dtype=padded_input_ids.dtype, - device=padded_input_ids.device)), - axis=-1) - - # Note: we still allocate max_seq_length size of parent ids (not max_attention_window_size). - self.parent_ids = torch.zeros( - (batch_size, scfg.num_beams, self.max_seq_length), - dtype=torch.int32, - device=self.device) - - if self.is_redrafter_mode: - self.new_tokens = torch.zeros([ - batch_size, self._model_config.redrafter_draft_len_per_beam + 1 - ], - dtype=torch.int32, - device=self.device) - self.accept_lengths = torch.ones([batch_size], - dtype=torch.int32, - device=self.device) - self.buffer["redrafter_inverted_temperature"] = torch.reciprocal( - self.temperature).to(device=self.device, dtype=self.dtype) - elif self.is_medusa_mode: - self.new_tokens = torch.zeros( - [batch_size, self.num_medusa_heads + 1], - dtype=torch.int32, - device=self.device) - self.medusa_output_tokens = torch.zeros( - [batch_size, self.num_draft_tokens], - dtype=torch.int32, - device=self.device) - self.generation_input_ids = torch.zeros( - [batch_size, self.num_draft_tokens + 1], - dtype=torch.int32, - device=self.device) - self.accept_lengths = torch.ones([batch_size], - dtype=torch.int32, - device=self.device) - if self.medusa_temperature != 0: - self.medusa_output_logits = torch.empty( - [batch_size, self.num_medusa_heads, self.vocab_size_padded], - dtype=self._tensor_dtype('logits'), - device=self.device) - elif scfg.num_beams > 1: - self.new_tokens = torch.zeros([batch_size, scfg.num_beams, 1], - dtype=torch.int32, - device=self.device) - else: - self.new_tokens = torch.zeros([batch_size, 1], - dtype=torch.int32, - device=self.device) - - if scfg.num_beams > 1 or scfg.output_cum_log_probs: - self.cum_log_probs = torch.full((batch_size, scfg.num_beams), - -1e20, - dtype=torch.float32, - device=self.device) - self.cum_log_probs[:, 0] = 0.0 - else: - self.cum_log_probs = None - - if scfg.output_log_probs: - self.log_probs = torch.zeros( - (batch_size, scfg.num_beams, self.max_seq_length), - dtype=torch.float32, - device=self.device) - self.log_probs_tiled = torch.zeros( - (self.max_seq_length, self._model_config.max_batch_size, - scfg.num_beams), - dtype=torch.float32, - device=self.device) - else: - self.log_probs = None - self.log_probs_tiled = None - - self.finished = torch.zeros((batch_size, scfg.num_beams), - dtype=torch.uint8, - device=self.device) - - if scfg.use_beam_hyps: - self.beam_hyps_output_ids_cba = torch.full( - size=[batch_size, scfg.num_beams * 2, self.max_seq_length], - fill_value=scfg.end_id, - dtype=torch.int32, - device=self.device) - self.beam_hyps_seq_len_cba = torch.zeros( - [batch_size, scfg.num_beams * 2], - dtype=torch.int32, - device=self.device) - self.beam_hyps_cum_log_probs_cba = torch.zeros( - [batch_size, scfg.num_beams * 2], - dtype=torch.float, - device=self.device) - self.beam_hyps_normed_scores_cba = torch.zeros( - [batch_size, scfg.num_beams * 2], - dtype=torch.float, - device=self.device) - self.beam_hyps_log_probs_cba = torch.zeros( - [batch_size, scfg.num_beams * 2, self.max_seq_length], - dtype=torch.float, - device=self.device) - self.beam_hyps_min_normed_scores = torch.zeros([batch_size], - dtype=torch.float, - device=self.device) - self.beam_hyps_num_beams = torch.zeros([batch_size], - dtype=torch.int32, - device=self.device) - self.beam_hyps_is_done = torch.zeros([batch_size], - dtype=torch.bool, - device=self.device) - else: - self.beam_hyps_output_ids_cba = None - self.beam_hyps_seq_len_cba = None - self.beam_hyps_cum_log_probs_cba = None - self.beam_hyps_normed_scores_cba = None - self.beam_hyps_log_probs_cba = None - self.beam_hyps_min_normed_scores = None - self.beam_hyps_num_beams = None - self.beam_hyps_is_done = None - - self.cross_kv_reuse = None - - def _tensor_dtype(self, name): - # return torch dtype given tensor name for convenience - dtype = trt_dtype_to_torch(self.runtime.engine.get_tensor_dtype(name)) - return dtype - - def _init_medusa(self, medusa_choices: List[List[int]]): - from tensorrt_llm.runtime.medusa_utils import (_medusa_setup, - expand_choices_if_needed) - medusa_choices = expand_choices_if_needed(medusa_choices) - self.num_draft_tokens = len(medusa_choices) - assert self.num_draft_tokens > 0 and self.num_draft_tokens <= self.max_draft_tokens - medusa_info = _medusa_setup(medusa_choices, self.num_medusa_heads) - self.medusa_topks = medusa_info.medusa_topks - self.medusa_mask = medusa_info.medusa_mask[1:, 1:].to( - torch.bool - ) # convert to bool, original mask includes true token as well - - # Expand medusa position offsets to number of batch size in order to be compatible with the new Medusa. - target_shape = list(medusa_info.medusa_packed_mask.unsqueeze(0).shape) - target_shape[0] = self.batch_size - # Note: spec_decoding_packed_mask has no paddings in the first dimension. - self.spec_decoding_packed_mask = medusa_info.medusa_packed_mask.unsqueeze( - 0).expand(target_shape).reshape(-1, target_shape[-1]).cuda() - self.spec_decoding_use = medusa_info.medusa_spec_decoding_use - - self.medusa_paths = medusa_info.medusa_paths - self.medusa_tree_ids = medusa_info.medusa_tree_ids - - # Expand medusa position offsets to number of batch size in order to be compatible with the new Medusa. - target_shape = list( - medusa_info.medusa_position_offsets.unsqueeze(0).shape) - target_shape[0] = self.batch_size - # Note: medusa_position_offsets still keeps the paddings in order to get max_gen_input_length from the shape info. - self.spec_decoding_position_offsets = medusa_info.medusa_position_offsets.unsqueeze( - 0).expand(target_shape).int().cuda() - # Fixed sequence lengths currently. - # Support variable sequence lengths later. - self.spec_decoding_generation_lengths = (torch.ones( - (self.batch_size)) * (self.num_draft_tokens + 1)).int().cuda() - if not self.use_gpt_attention_plugin: - medusa_fp_mask = torch.zeros_like(self.medusa_mask, - dtype=torch.float32) - medusa_fp_mask[torch.logical_not(self.medusa_mask)] = float('-inf') - self.medusa_mask = medusa_fp_mask - return - - def _get_num_paged_blocks(self, max_attention_window_size, - sink_token_length): - bubble_len = 0 - if sink_token_length % self.tokens_per_block > 0: - bubble_len += (self.tokens_per_block - - sink_token_length % self.tokens_per_block) - max_blocks_per_seq = math.ceil( - (max_attention_window_size + bubble_len) / self.tokens_per_block) - num_blocks = self.batch_size * self.beam_width * max_blocks_per_seq - - return num_blocks, max_blocks_per_seq - - def setup(self, - batch_size: int, - max_context_length: int, - max_new_tokens: int, - beam_width: int = 1, - max_attention_window_size: Optional[int] = None, - sink_token_length: Optional[int] = None, - encoder_max_input_length: Optional[int] = None, - lora_manager: LoraManager = None, - lora_uids: List[str] = None, - medusa_choices: List[List[int]] = None, - multi_block_mode: bool = True, - enable_context_fmha_fp32_acc: bool = None): - # Store these params related to buffer size to check against - # the input shape with the params given in decode() - self.batch_size = batch_size - self.max_context_length = max_context_length - self.max_new_tokens = max_new_tokens - self.max_seq_length = max_context_length + max_new_tokens - if medusa_choices is not None or self.is_redrafter_mode: - self.max_seq_length += self.max_draft_tokens - self.beam_width = beam_width - self.encoder_max_input_length = encoder_max_input_length - self.multi_block_mode = multi_block_mode - self.enable_context_fmha_fp32_acc = enable_context_fmha_fp32_acc - if max_attention_window_size is None: - self.max_attention_window_size = self.max_seq_length - logger.debug( - "The max_attention_window_size is not set, we will use max_seq_length by default." - ) - self.host_max_attention_window_sizes = torch.ones( - (self.num_attn_layers, ), - dtype=torch.int32) * self.max_attention_window_size - - elif isinstance(max_attention_window_size, int): - if max_attention_window_size > self.max_seq_length: - logger.warning( - "The value of max_attention_window_size should ideally not exceed max_seq_length. " - "Therefore, it has been adjusted to match the value of max_seq_length." - ) - self.max_attention_window_size = min(max_attention_window_size, - self.max_seq_length) - self.host_max_attention_window_sizes = torch.ones( - (self.num_attn_layers, ), - dtype=torch.int32) * self.max_attention_window_size - - elif isinstance(max_attention_window_size, (torch.Tensor, list)): - if isinstance(max_attention_window_size, list): - max_attention_window_size = torch.tensor( - max_attention_window_size, dtype=torch.int32) - self.max_attention_window_size = int( - torch.max(max_attention_window_size).item()) - attn_win_size_len = max_attention_window_size.shape[0] - num_total_attn_layers = self.layer_types.count('attention') - if attn_win_size_len < num_total_attn_layers: - repeat_num = num_total_attn_layers // attn_win_size_len - remain_num = num_total_attn_layers % attn_win_size_len - warning_info = "The size of max_attention_window_size tensor/list is less than num_attn_layers, " \ - + "and it will be repeated to num_attn_layers. So the actual max_attention_window_size " \ - + f"is {max_attention_window_size.tolist()} * {repeat_num}" - warning_info += f" + {max_attention_window_size.tolist()[0:remain_num]}. " if remain_num > 0 else ". " - warning_info += "Note that num_attn_layers is the number of total attention layers." - logger.warning(warning_info) - elif attn_win_size_len > num_total_attn_layers: - logger.error( - "The size of max_attention_window_size tensor/list is larger than num_attn_layers! " - "Note that num_attn_layers is the number of total attention layers." - ) - assert False - if self.max_attention_window_size > self.max_seq_length: - logger.warning( - "The value of max_attention_window_size should ideally not exceed max_seq_length. " - "Therefore, it has been adjusted to match the value of max_seq_length." - ) - self.max_attention_window_size = min(self.max_attention_window_size, - self.max_seq_length) - max_attention_window_size = torch.minimum( - max_attention_window_size.to(torch.int32), - torch.IntTensor([self.max_seq_length] * attn_win_size_len)) - self.host_max_attention_window_sizes = torch.ones( - (self.num_attn_layers, ), dtype=torch.int32) - for i in range(self.num_attn_layers): - self.host_max_attention_window_sizes[ - i] = max_attention_window_size[ - (self.layer_types[0:self.first_layer].count('attention') - + i) % attn_win_size_len] - else: - assert False, "invalid max_attention_window_size!" - - if sink_token_length is None: - self.sink_token_length = 0 - self.host_sink_token_length = torch.zeros((1, ), dtype=torch.int32) - elif isinstance(sink_token_length, int): - self.sink_token_length = sink_token_length - self.host_sink_token_length = torch.ones( - (1, ), dtype=torch.int32) * self.sink_token_length - else: - assert False, "invalid sink_token_length!" - - self.lora_manager = lora_manager - if medusa_choices is not None: - self._init_medusa(medusa_choices) - - self.buffer = {} - if self.mapping.is_last_pp_rank(): - if self.is_redrafter_mode: - init_allocate_redrafter_tensors(self, batch_size) - self.buffer['logits'] = torch.empty( - (batch_size, self.max_draft_tokens + 1, - self.vocab_size_padded) - if not self.gather_context_logits else - (batch_size, max_context_length, self.vocab_size_padded), - dtype=self._tensor_dtype('logits'), - device=self.device) - elif self.is_medusa_mode: - self.buffer['logits'] = torch.empty( - (batch_size, self.num_draft_tokens + 1, - self.vocab_size_padded) - if not self.gather_context_logits else - (batch_size, max_context_length, self.vocab_size_padded), - dtype=self._tensor_dtype('logits'), - device=self.device) - medusa_logits_shape = (self.num_medusa_heads, batch_size, - (self.num_draft_tokens + 1), - self.vocab_size_padded) - if self.remove_input_padding: - medusa_logits_shape = (self.num_medusa_heads, batch_size * - (self.num_draft_tokens + 1), - self.vocab_size_padded) - - self.buffer['medusa_logits'] = torch.empty( - medusa_logits_shape if not self.gather_context_logits else - (self.num_medusa_heads, batch_size, max_context_length, - self.vocab_size_padded), - dtype=self._tensor_dtype('medusa_logits'), - device=self.device) - else: - self.buffer['logits'] = torch.empty( - (batch_size, self.vocab_size_padded) - if not self.gather_context_logits else - (batch_size, max_context_length, self.vocab_size_padded), - dtype=self._tensor_dtype('logits'), - device=self.device) - - if self.cross_attention: - # use shape info to pass max length info in remove padding mode - self.buffer['encoder_max_input_length'] = torch.empty( - (encoder_max_input_length, ), - dtype=self._tensor_dtype('encoder_max_input_length'), - device=self.device) - - if self.quant_mode.has_kv_cache_quant(): - # Since torch does not support fp8 now, using int8 here. - kv_cache_type = torch.int8 - else: - if self.use_kv_cache and self.has_attn_layers: - first_atten_layer = self.layer_types[ - self.first_layer:self.last_layer].index( - 'attention') + self.first_layer - kv_cache_type = self.dtype if self.paged_kv_cache else self._tensor_dtype( - f'present_key_value_{first_atten_layer}') - else: - kv_cache_type = None - - if self.use_kv_cache: - if self.paged_kv_cache and self.has_attn_layers: - num_blocks, _ = self._get_num_paged_blocks( - self.max_attention_window_size, self.sink_token_length) - self._memory_pool_allocator = MemoryPoolsAllocator( - num_blocks=num_blocks, - tokens_per_block=self.tokens_per_block, - head_size=self.head_size) - if self._model_config.num_kv_heads_per_layer is None: - num_kv_heads_per_layer = MemoryPoolsAllocator.prepare_num_kv_heads_per_layer( - self.get_num_heads_kv(), self.num_attn_layers) - else: - num_kv_heads_per_layer = self._model_config.num_kv_heads_per_layer - - self._memory_pool_allocator.allocate(kv_cache_type, - num_kv_heads_per_layer) - - if self.cross_attention: # As for now we enable cross paged kv and self paged kv to share the same tokens_per_block - cross_num_blocks, _ = self._get_num_paged_blocks( - self.encoder_max_input_length, sink_token_length=0) - - num_kv_heads_per_layer = MemoryPoolsAllocator.prepare_num_kv_heads_per_layer( - self.get_num_heads_kv(), self.num_attn_layers) - - self._cross_memory_pool_allocator = MemoryPoolsAllocator( - num_blocks=cross_num_blocks, - tokens_per_block=self.tokens_per_block, - head_size=self.head_size) - if self._model_config.num_kv_heads_per_cross_attn_layer is None: - num_kv_heads_per_cross_attn_layer = MemoryPoolsAllocator.prepare_num_kv_heads_per_layer( - self.get_num_heads_kv(), self.num_attn_layers) - else: - num_kv_heads_per_cross_attn_layer = self._model_config.num_kv_heads_per_cross_attn_layer - - self._cross_memory_pool_allocator.allocate( - kv_cache_type, num_kv_heads_per_cross_attn_layer) - - elif self.has_attn_layers: - - for i in range(self.first_layer, self.last_layer): - if self.layer_types[i] == 'attention': - cache_shape = ( - batch_size, - 2, - self.get_num_heads_kv(i), - self.max_attention_window_size, - self.head_size, - ) - self.buffer[f'present_key_value_{i}'] = torch.empty( - cache_shape, - dtype=kv_cache_type, - device=self.device) - - if self.cross_attention: - cross_cache_shape = ( - batch_size, - 2, - self.get_num_heads_kv(), - self.encoder_max_input_length, - self.head_size, - ) - for i in range(self.first_layer, self.last_layer): - if self.layer_types[i] == 'attention': - self.buffer[ - f'cross_present_key_value_{i}'] = torch.empty( - cross_cache_shape, - dtype=kv_cache_type, - device=self.device) - - if self.use_gpt_attention_plugin: - self.sequence_length_buffer = torch.ones((batch_size, ), - dtype=torch.int32, - device=self.device) - else: - # Without plugin, we need extra kv cache buffers. - # Because we don't support inplace update, so we need separate buffer for inputs and outputs. - # We can do reuse between different layers' inputs and outputs, i.e. current layer's output can - # reuse previous layer's input memory. But this need one extra buffer as the guard. - if self.use_kv_cache and self.has_attn_layers: # Not applicable to cross KV buffers as it's constant - i = self.attn_to_general_idx[0] - trt_dtype = self.runtime.engine.get_tensor_dtype( - f'present_key_value_{i}') - - if trt_dtype == trt.fp8: - # PyTorch doesn't support fp8 datatype, use int8 instead of it because int8 datatype size is same with fp8. - # TODO: Remove this section when PyTorch support fp8 datatype - dtype = torch.int8 - else: - dtype = self._tensor_dtype(f'present_key_value_{i}') - self.buffer[f'1_present_key_value_{i}'] = torch.empty( - cache_shape, dtype=dtype, device=self.device) - - if self.use_mamba_conv1d_plugin: - conv_state_shape = ( - batch_size, - self.conv_kernel - 1, - self.rnn_conv_dim_size, - ) - else: - conv_state_shape = ( - batch_size, - self.rnn_conv_dim_size, - self.conv_kernel - 1, - ) - - if self.rnn_head_size > 1: - rnn_state_shape = ( - batch_size, - self.rnn_hidden_size // self.rnn_head_size, - self.state_size, - self.rnn_head_size, - ) - else: - rnn_state_shape = ( - batch_size, - self.state_size, - self.rnn_hidden_size, - ) - - for i in range(self.first_layer, self.last_layer): - if self.layer_types[i] == 'recurrent': - dtype = self.dtype - self.buffer[f'present_conv_state_{i}'] = torch.empty( - conv_state_shape, dtype=dtype, device=self.device) - self.buffer[f'1_present_conv_state_{i}'] = torch.empty( - conv_state_shape, dtype=dtype, device=self.device) - self.buffer[f'present_rnn_state_{i}'] = torch.empty( - rnn_state_shape, dtype=self.state_dtype, device=self.device) - if self.paged_state: - conv_state_ptr = torch.tensor( - [self.buffer[f'present_conv_state_{i}'].data_ptr()], - dtype=torch.int64, - device='cpu') - rnn_state_ptr = torch.tensor( - [self.buffer[f'present_rnn_state_{i}'].data_ptr()], - dtype=torch.int64, - device='cpu') - self.buffer[f'conv_state_ptr_{i}'] = conv_state_ptr - self.buffer[f'rnn_state_ptr_{i}'] = rnn_state_ptr - - if self.use_lora_plugin and self.lora_manager is not None: - lora_uids = lora_uids or ["-1"] - self.buffer.update( - self.lora_manager.input_buffers( - lora_uids, - self.mapping, - self._model_config.num_layers, - )) - - if self.use_gemm_allreduce_plugin: - max_num_tokens = max(batch_size * beam_width, - batch_size * self.max_seq_length) - M = max_num_tokens - N = self.hidden_size - self.gemm_allreduce_output_size = M * N - itemsize = str_dtype_to_torch(self.gemm_allreduce_plugin).itemsize - alloc_bytes = self.gemm_allreduce_output_size * itemsize - self.gemm_allreduce_output_handle = ipc_nvls_allocate( - alloc_bytes, set(self.mapping.tp_group)) - logger.debug(f'Allocated NVLS IPC memory: {alloc_bytes} bytes') - - if self.is_medusa_mode: - self.buffer[ - 'spec_decoding_packed_mask'] = self.spec_decoding_packed_mask - self.buffer[ - 'spec_decoding_position_offsets'] = self.spec_decoding_position_offsets - self.buffer[ - 'spec_decoding_generation_lengths'] = self.spec_decoding_generation_lengths - self.buffer['spec_decoding_use'] = self.spec_decoding_use - self.buffer_allocated = True - if self.is_medusa_mode: - return self.num_draft_tokens - - def _allocate_empty_kv_cache_pools(self, kv_cache_type, num_blocks): - # Layers are homogeneous, use old kv cache shape - unique_cache_pools = [] - if self._model_config.num_kv_heads_per_layer is None: - cache_shape = ( - num_blocks, - self.num_attn_layers, - 2, - self.get_num_heads_kv(), - self.tokens_per_block, - self.head_size, - ) - unique_cache_pools.append( - torch.empty(cache_shape, - dtype=kv_cache_type, - device=self.device)) - - # Layers are not homogeneous, use new kv cache shape - else: - kv_heads_unique_counter = Counter( - self._model_config.num_kv_heads_per_layer) - for kv_head, num_layers in kv_heads_unique_counter.items(): - cache_shape = ( - num_blocks, - num_layers, - 2, - kv_head, - self.tokens_per_block, - self.head_size, - ) - unique_cache_pools.append( - torch.empty(cache_shape, - dtype=kv_cache_type, - device=self.device)) - - return unique_cache_pools - - def _get_context_shape_buffer( - self, - input_ids: torch.Tensor, - context_lengths: torch.Tensor, - host_context_lengths: torch.Tensor, - position_ids: torch.Tensor, - last_token_ids: torch.Tensor, - attention_mask: torch.Tensor, - cross_attention_mask: torch.Tensor, - cache_indirection: torch.Tensor, - kv_cache_block_offsets: torch.Tensor, - host_kv_cache_block_offsets: torch.Tensor, - cross_kv_cache_block_offsets: torch.Tensor = None, - host_cross_kv_cache_block_offsets: torch.Tensor = None, - hidden_states_input: torch.Tensor = None, - prompt_embedding_table: torch.Tensor = None, - tasks: torch.Tensor = None, - prompt_vocab_size: torch.Tensor = None, - encoder_output: torch.Tensor = None, - encoder_input_lengths: torch.Tensor = None, - host_runtime_perf_knobs: torch.Tensor = None, - host_context_progress: torch.Tensor = None, - skip_cross_attn_blocks: torch.Tensor = None, - language_adapter_routings: torch.Tensor = None, - ) -> Dict[str, RuntimeTensor]: - tensors = {} - - def sym(x, name): - return RuntimeTensor.from_torch(name, x) - - def add_tensor_from_pointer(pointer, name, shape, str_dtype): - return tensors.update({ - name: - RuntimeTensor.from_pointer(name, pointer, shape, str_dtype) - }) - - def add_tensor(x, name): - return tensors.update({name: sym(x, name)}) - - def add_tensor_with_shape(x, name, shape): - return tensors.update( - {name: RuntimeTensor.from_torch(name, x, override_shape=shape)}) - - def add_tensor_with_bs(x, name, bs): - # this assumes dim0 to be bs and only overrides dim0 with given bs - shape = list(x.shape) - shape[0] = bs - return tensors.update( - {name: RuntimeTensor.from_torch(name, x, override_shape=shape)}) - - if self.has_attn_layers: - if self.use_gpt_attention_plugin: - add_tensor(context_lengths, 'context_lengths') - assert host_runtime_perf_knobs != None, "gpt_attention_plugin needs to set host_runtime_perf_knobs" - add_tensor(host_runtime_perf_knobs, 'host_runtime_perf_knobs') - add_tensor(host_context_progress, 'host_context_progress') - add_tensor(cache_indirection, 'cache_indirection') - - if self.has_position_embedding: - add_tensor(position_ids, 'position_ids') - - if self.cross_attention: - # in context phase, need to generate cross kv cache, set to True - add_tensor(torch.ones(1, dtype=torch.bool, device=self.device), - 'cross_kv_cache_gen') - if self._model_config.skip_cross_attn_blocks: - add_tensor(skip_cross_attn_blocks, 'skip_cross_attn_blocks') - if self.skip_cross_kv: - if self.cross_kv_reuse is None: - # see Attention's self.qkv output dim - cross_kv_out_dim = 2 * self.get_num_heads_kv( - ) * self.head_size - cross_kv_shape = encoder_output.shape[:-1] + ( - cross_kv_out_dim, ) - cross_kv_reuse = torch.empty(cross_kv_shape, - dtype=encoder_output.dtype, - device=encoder_output.device) - self.cross_kv_reuse = cross_kv_reuse - add_tensor(self.cross_kv_reuse, 'cross_kv_reuse') - add_tensor(encoder_output, 'encoder_output') - add_tensor(encoder_input_lengths, 'encoder_input_lengths') - if language_adapter_routings is not None: - add_tensor(language_adapter_routings, - 'language_adapter_routings') - add_tensor(self.buffer['encoder_max_input_length'], - 'encoder_max_input_length') - if not self.use_gpt_attention_plugin: - add_tensor(cross_attention_mask, 'cross_attention_mask') - else: - if cross_attention_mask != None: - # cross-attention packed mask (used by fmha). - cross_attention_packed_mask = torch.ops.tensorrt_llm.pack_fmha_mask_by_input( - cross_attention_mask, context_lengths, - encoder_input_lengths, 1.0) - add_tensor(cross_attention_mask, 'cross_attention_mask') - add_tensor(cross_attention_packed_mask, - 'cross_attention_packed_mask') - else: - # create a full 1 cross_attention_mask because it is necessary - batch_size = context_lengths.shape[0] - cross_attention_mask = torch.ones( - (np.asarray(input_ids.shape).prod(), - np.asarray(list(encoder_output.shape)[:-1]).prod()), - dtype=torch.bool, - device=self.device) - add_tensor(cross_attention_mask, "cross_attention_mask") - cross_attention_packed_mask = torch.ops.tensorrt_llm.pack_fmha_mask_by_input( - cross_attention_mask, context_lengths, - encoder_input_lengths, 1.0) - add_tensor(cross_attention_packed_mask, - "cross_attention_packed_mask") - - if self.mapping.has_pp(): - hidden_size = self.hidden_size * self.mapping.tp_size - if input_ids.dim() == 2: - hidden_states_input = hidden_states_input.resize_( - input_ids.shape[0], input_ids.shape[1], hidden_size) - else: - hidden_states_input = hidden_states_input.resize_( - input_ids.shape[0], hidden_size) - - if self.mapping.is_last_pp_rank(): - if self.is_redrafter_mode: - set_redrafter_ctx_tensors(self, add_tensor, add_tensor_with_bs) - add_tensor(self.buffer['logits'], 'logits') - if self.is_medusa_mode: - add_tensor(self.buffer['medusa_logits'], 'medusa_logits') - - if not self.gather_context_logits or self.has_rnn_layers: - add_tensor(last_token_ids, 'last_token_ids') - else: - add_tensor(hidden_states_input, 'hidden_states_output') - - if self.mapping.is_first_pp_rank(): - add_tensor(input_ids, 'input_ids') - else: - add_tensor(hidden_states_input, 'hidden_states_input') - - if prompt_embedding_table is not None: - add_tensor(prompt_embedding_table, 'prompt_embedding_table') - - if self.remove_input_padding: - tasks_generation = torch.concat([ - torch.full([context_lengths[b].item()], - tasks[b].item(), - dtype=torch.int32) - for b in range(context_lengths.size(0)) - ]).cuda() - else: - tasks_generation = tasks.unsqueeze(-1) - add_tensor(tasks_generation, 'tasks') - add_tensor(prompt_vocab_size, 'prompt_vocab_size') - - if self.paged_kv_cache and self.has_attn_layers: - buffer = kv_cache_block_offsets.contiguous() - shape = kv_cache_block_offsets.shape - shape = [shape[0], shape[1] * shape[2], *shape[3:]] - add_tensor_with_shape(buffer, f'kv_cache_block_offsets', shape) - add_tensor_with_shape(host_kv_cache_block_offsets, - f'host_kv_cache_block_offsets', shape) - pool_pointers = f'host_kv_cache_pool_pointers' - pool_mapping = f'host_kv_cache_pool_mapping' - add_tensor(self.buffer[pool_pointers], pool_pointers) - add_tensor(self.buffer[pool_mapping], pool_mapping) - if self.cross_attention: - cross_buffer = cross_kv_cache_block_offsets.contiguous() - cross_shape = cross_kv_cache_block_offsets.shape - cross_shape = [ - cross_shape[0], cross_shape[1] * cross_shape[2], - *cross_shape[3:] - ] - add_tensor_with_shape(cross_buffer, - f'cross_kv_cache_block_offsets', - cross_shape) - add_tensor_with_shape(host_cross_kv_cache_block_offsets, - f'host_cross_kv_cache_block_offsets', - cross_shape) - cross_pool_pointers = f'host_cross_kv_cache_pool_pointers' - cross_pool_mapping = f'host_cross_kv_cache_pool_mapping' - add_tensor(self.buffer[cross_pool_pointers], - cross_pool_pointers) - add_tensor(self.buffer[cross_pool_mapping], cross_pool_mapping) - - batch_size = context_lengths.shape[0] - if self.use_kv_cache and not self.paged_kv_cache: - for idx in range(self.first_layer, self.last_layer): - if not self.use_gpt_attention_plugin and self.layer_types[ - idx] == 'attention': - kv_cache_shape = (batch_size, 2, - self.get_num_heads_kv( - self.general_to_attn_idx[idx]), 0, - self.head_size) - # for empty tensor, TRT does not really use the tensor data, so any dtype is fine - kv_cache_buffer = torch.zeros((1, ), - dtype=torch.float32, - device=self.device) - add_tensor_with_shape(kv_cache_buffer, - f'past_key_value_{idx}', - kv_cache_shape) - present = f'present_key_value_{idx}' - add_tensor(self.buffer[present], present) - - if self.cross_attention: - cross_kv_cache_shape = (batch_size, 2, - self.get_num_heads_kv(), 0, - self.head_size) - # for empty tensor, TRT does not really use the tensor data, so any dtype is fine - cross_kv_cache_buffer = torch.zeros((1, ), - dtype=torch.float32, - device=self.device) - add_tensor_with_shape(cross_kv_cache_buffer, - f'cross_past_key_value_{idx}', - cross_kv_cache_shape) - cross_present = f'cross_present_key_value_{idx}' - add_tensor(self.buffer[cross_present], cross_present) - elif self.layer_types[idx] == 'attention': - key_value_cache = self.buffer[f'present_key_value_{idx}'] - # when plugin is used, past_ket_value tensor does not need to be empty tensor - # because plugin does not care, and does not use this shape. - add_tensor(key_value_cache, f'past_key_value_{idx}') - add_tensor(key_value_cache, f'present_key_value_{idx}') - - if self.cross_attention: - cross_cache_buffer = self.buffer[ - f'cross_present_key_value_{idx}'] - add_tensor(cross_cache_buffer, - f'cross_past_key_value_{idx}') - add_tensor(cross_cache_buffer, - f'cross_present_key_value_{idx}') - - for idx in range(self.first_layer, self.last_layer): - if self.layer_types[idx] != 'recurrent': - continue - if self.paged_state: - add_tensor(self.buffer[f'conv_state_ptr_{idx}'], - f'conv_state_ptr_{idx}') - add_tensor(self.buffer[f'rnn_state_ptr_{idx}'], - f'rnn_state_ptr_{idx}') - else: - # conv state - dtype = self._tensor_dtype(f'present_conv_state_{idx}') - if self.use_mamba_conv1d_plugin: - conv_state_shape = (batch_size, self.conv_kernel - 1, - self.rnn_conv_dim_size) - else: - conv_state_shape = (batch_size, self.rnn_conv_dim_size, - self.conv_kernel - 1) - - conv_state = torch.zeros(conv_state_shape, - dtype=dtype, - device=self.device) - add_tensor(conv_state, f'past_conv_state_{idx}') - present = f'present_conv_state_{idx}' - add_tensor(self.buffer[present], present) - # rnn state - rnn_state = self.buffer[f'present_rnn_state_{idx}'] - add_tensor(rnn_state, f'past_rnn_state_{idx}') - add_tensor(rnn_state, f'present_rnn_state_{idx}') - - if self.paged_state and self.has_rnn_layers: - slot_mapping = torch.arange(0, - batch_size, - device='cuda', - dtype=torch.int32) - add_tensor(slot_mapping, 'slot_mapping') - - if self.use_gpt_attention_plugin and self.has_attn_layers: - # context request - host_request_types = torch.zeros_like(context_lengths, - device='cpu').int() - self.sequence_length_buffer = context_lengths.detach().clone() - if self.is_redrafter_mode: - device_request_types = torch.zeros_like( - context_lengths, device=self.device).int() - add_tensor(device_request_types, 'device_request_types') - add_tensor_with_shape(self.sequence_length_buffer, - 'sequence_length', (batch_size, )) - - # field 0: past_key_value_length, field 1: is_context (deprecated). changed to [0], otherwise affects batch padded input mode - add_tensor_with_shape(host_context_lengths.clone(), - 'host_past_key_value_lengths', (batch_size, )) - add_tensor_with_shape(self.host_sink_token_length, - 'host_sink_token_length', (1, )) - add_tensor(host_request_types, 'host_request_types') - add_tensor_with_shape(self.host_max_attention_window_sizes, - f'host_max_attention_window_sizes', - (self.num_attn_layers, )) - if self.remove_input_padding: - add_tensor(host_context_lengths, 'host_context_lengths') - else: - if self.has_rnn_layers: - host_request_types = torch.zeros_like(context_lengths, - device='cpu').int() - add_tensor(host_request_types, 'host_request_types') - if self.remove_input_padding: - add_tensor(host_context_lengths, 'host_context_lengths') - if self.has_attn_layers: - add_tensor(attention_mask, 'attention_mask') - - if self.mapping.tp_size > 1: - add_tensor(self.all_reduce_workspace, 'all_reduce_workspace') - if self.use_gemm_allreduce_plugin: - found_tensor_names = [ - self.runtime.engine.get_tensor_name(i) - for i in range(self.runtime.engine.num_io_tensors) - ] - for name in found_tensor_names: - if name.startswith("gemm_allreduce_uc_out"): - add_tensor_from_pointer( - self.gemm_allreduce_output_handle.uc_ptr, - name, - shape=(self.gemm_allreduce_output_size), - str_dtype=self.gemm_allreduce_plugin) - if name.startswith("gemm_allreduce_mc_out"): - add_tensor_from_pointer( - self.gemm_allreduce_output_handle.mc_ptr, - name, - shape=(self.gemm_allreduce_output_size), - str_dtype=self.gemm_allreduce_plugin) - if name.startswith("gemm_allreduce_ipc_out"): - add_tensor_from_pointer( - self.gemm_allreduce_output_handle.get_ipc_ptrs(), - name, - shape=(self.gemm_allreduce_output_size), - str_dtype=self.gemm_allreduce_plugin) - - if self.use_lora_plugin: - for idx in range(self.num_layers): - for lora_module in (self.lora_target_modules + - self.missing_qkv_modules): - layer_idx = idx + self.first_layer - lora_ranks = f'{lora_module}_lora_ranks_{layer_idx}' - add_tensor(self.buffer[lora_ranks], lora_ranks) - lora_weights = f'{lora_module}_lora_weights_pointers_{layer_idx}' - add_tensor(self.buffer[lora_weights], lora_weights) - if self.cross_attention and self.remove_input_padding: - add_tensor(encoder_input_lengths.to('cpu'), - 'host_encoder_input_lengths') - if self.is_medusa_mode: - # Medusa mask and position offsets are fixed for the whole session. - add_tensor(self.buffer['spec_decoding_packed_mask'], - 'spec_decoding_packed_mask') - add_tensor(self.buffer['spec_decoding_position_offsets'], - 'spec_decoding_position_offsets') - add_tensor(self.buffer['spec_decoding_generation_lengths'], - 'spec_decoding_generation_lengths') - add_tensor(self.buffer['spec_decoding_use'], 'spec_decoding_use') - - return tensors - - def _get_next_step_shape_buffer( - self, - batch_size: int, - beam_width: int, - max_context_length: int, - step: int, - context_lengths: torch.Tensor, - host_context_lengths: torch.Tensor, - position_ids: torch.Tensor, - last_token_ids: torch.Tensor, - attention_mask: torch.Tensor, - cross_attention_mask: torch.Tensor, - cache_indirection: torch.Tensor, - kv_cache_block_offsets: torch.Tensor, - host_kv_cache_block_offsets: torch.Tensor, - cross_kv_cache_block_offsets: torch.Tensor = None, - host_cross_kv_cache_block_offsets: torch.Tensor = None, - hidden_states_input: torch.Tensor = None, - prompt_embedding_table: torch.Tensor = None, - tasks: torch.Tensor = None, - prompt_vocab_size: torch.Tensor = None, - encoder_output: torch.Tensor = None, - encoder_input_lengths: torch.Tensor = None, - host_runtime_perf_knobs: torch.Tensor = None, - host_context_progress: torch.Tensor = None, - skip_cross_attn_blocks: torch.Tensor = None, - language_adapter_routings: torch.Tensor = None, - ): - torch.cuda.nvtx.range_push("_get_next_step_shape_buffer") - tensors = {} # Dict[str, RuntimeTensor] - - def add_tensor_from_pointer(pointer, name, shape, str_dtype): - return tensors.update({ - name: - RuntimeTensor.from_pointer(name, pointer, shape, str_dtype) - }) - - def sym(x, name): - return RuntimeTensor.from_torch(name, x) - - def add_tensor(x, name): - return tensors.update({name: sym(x, name)}) - - def add_tensor_with_shape(x, name, shape): - return tensors.update( - {name: RuntimeTensor.from_torch(name, x, override_shape=shape)}) - - context_lengths_local = context_lengths.clone() - host_context_lengths_local = host_context_lengths.clone() - if self.has_attn_layers: - if self.use_gpt_attention_plugin: - add_tensor(context_lengths_local, 'context_lengths') - assert host_runtime_perf_knobs != None, "gpt_attention_plugin needs to set host_runtime_perf_knobs" - add_tensor(host_runtime_perf_knobs, 'host_runtime_perf_knobs') - add_tensor(host_context_progress, 'host_context_progress') - add_tensor(cache_indirection, 'cache_indirection') - if self.has_position_embedding: - add_tensor(position_ids, 'position_ids') - - if self.mapping.has_pp(): - hidden_size = self.hidden_size * self.mapping.tp_size - shape = (batch_size * beam_width, - hidden_size) if self.remove_input_padding else ( - batch_size * beam_width, 1, hidden_size) - hidden_states_input = hidden_states_input.resize_(*shape) - - if self.mapping.is_last_pp_rank(): - add_tensor(self.buffer['logits'], 'logits') - if self.is_medusa_mode: - add_tensor(self.buffer['medusa_logits'], 'medusa_logits') - - if not self.gather_context_logits or self.has_rnn_layers: - add_tensor(last_token_ids, 'last_token_ids') - else: - add_tensor(hidden_states_input, 'hidden_states_output') - - if self.mapping.is_first_pp_rank(): - if self.is_redrafter_mode: - input_ids_shape = (self.host_total_gen_token, ) - else: - input_ids_shape = ( - batch_size * beam_width * (self.num_draft_tokens + 1), - ) if self.remove_input_padding else (batch_size * beam_width, - self.num_draft_tokens + 1) - if self.is_redrafter_mode: - add_tensor_with_shape(self.buffer['flat_tokens'], 'input_ids', - input_ids_shape) - elif self.is_medusa_mode: - add_tensor_with_shape(self.generation_input_ids, 'input_ids', - input_ids_shape) - else: - add_tensor_with_shape(self.new_tokens, 'input_ids', - input_ids_shape) - else: - add_tensor(hidden_states_input, 'hidden_states_input') - - if self.cross_attention: - if self.use_gpt_attention_plugin: - # disable (or minimize) cross qkv computation at generation phase - if self.skip_cross_kv: - # disable - encoder_output_shape = encoder_output.shape - add_tensor(self.cross_kv_reuse, 'cross_kv_reuse') - else: - # minimize - # use TensorRT Empty Tensor to skip redundant computation - # 0 for generation phase, >0 for context phase - encoder_output_shape = list(encoder_output.shape) - if self.remove_input_padding: - encoder_output_shape[-2] = 0 - else: - encoder_output_shape = [1, 0, encoder_output.shape[-1]] - else: - # OOTB path doesn't have kv cache for now, so this encoder_output is - # a must-have input. We just use the encoder_output - encoder_output_shape = encoder_output.shape - - # in generation phase, cross kv cache is already filled during context phase, set to False - add_tensor(torch.zeros(1, dtype=torch.bool, device=self.device), - 'cross_kv_cache_gen') - if self._model_config.skip_cross_attn_blocks: - add_tensor(skip_cross_attn_blocks, 'skip_cross_attn_blocks') - add_tensor_with_shape(encoder_output, 'encoder_output', - encoder_output_shape) - add_tensor(encoder_input_lengths, 'encoder_input_lengths') - if language_adapter_routings is not None: - add_tensor(language_adapter_routings, - 'language_adapter_routings') - add_tensor(self.buffer['encoder_max_input_length'], - 'encoder_max_input_length') - if not self.use_gpt_attention_plugin: - add_tensor(cross_attention_mask, 'cross_attention_mask') - else: - if cross_attention_mask != None: - cross_attention_mask = _tile_beam_width( - cross_attention_mask, beam_width) - # Empty packed mask is passed in the generation phase as it is not used. - cross_attention_packed_mask = torch.empty( - (batch_size, - (cross_attention_mask.shape[1] + 31) // 32), - dtype=torch.int32, - device=self.device) - add_tensor(cross_attention_mask, 'cross_attention_mask') - add_tensor(cross_attention_packed_mask, - 'cross_attention_packed_mask') - else: - # create a full 1 cross_attention_mask because it is necessary in generation phase - add_tensor( - torch.ones((batch_size, - np.asarray(list( - encoder_output.shape)[:-1]).prod()), - dtype=torch.bool, - device=self.device), "cross_attention_mask") - # Empty packed mask is passed in the generation phase as it is not used. - add_tensor( - torch.empty((batch_size, 1), - dtype=torch.int32, - device=self.device), - "cross_attention_packed_mask") - - if self.paged_kv_cache and self.has_attn_layers: - shape = kv_cache_block_offsets.shape - shape = [shape[0], shape[1] * shape[2], *shape[3:]] - add_tensor_with_shape(kv_cache_block_offsets, - f'kv_cache_block_offsets', shape) - add_tensor_with_shape(host_kv_cache_block_offsets, - f'host_kv_cache_block_offsets', shape) - pool_pointers = f'host_kv_cache_pool_pointers' - pool_mapping = f'host_kv_cache_pool_mapping' - add_tensor(self.buffer[pool_pointers], pool_pointers) - add_tensor(self.buffer[pool_mapping], pool_mapping) - if self.cross_attention: - cross_shape = cross_kv_cache_block_offsets.shape - cross_shape = [ - cross_shape[0], cross_shape[1] * cross_shape[2], - *cross_shape[3:] - ] - add_tensor_with_shape(cross_kv_cache_block_offsets, - f'cross_kv_cache_block_offsets', - cross_shape) - add_tensor_with_shape(host_cross_kv_cache_block_offsets, - f'host_cross_kv_cache_block_offsets', - cross_shape) - cross_pool_pointers = f'host_cross_kv_cache_pool_pointers' - cross_pool_mapping = f'host_cross_kv_cache_pool_mapping' - add_tensor(self.buffer[cross_pool_pointers], - cross_pool_pointers) - add_tensor(self.buffer[cross_pool_mapping], cross_pool_mapping) - - if prompt_embedding_table is not None: - add_tensor(prompt_embedding_table, 'prompt_embedding_table') - - if self.remove_input_padding: - gen_tasks = tasks - else: - gen_tasks = tasks.unsqueeze(-1) - add_tensor(gen_tasks, 'tasks') - add_tensor(prompt_vocab_size, 'prompt_vocab_size') - - if not self.paged_kv_cache: - for attn_idx, layer_idx in self.attn_to_general_idx.items(): - if not self.use_gpt_attention_plugin: - next_shape = (batch_size * beam_width, 2, - self.get_num_heads_kv(), - max_context_length + step, self.head_size) - # We will make current layer's output KV-cache overwrite previous layers input KV-cache - # buffer id: ... 5, 6, 7, 8, 9, ... - # layer n: out in - # layer n+1: out in - # layer n+2 out in - # And when finish a step, we will make every layer's in/out buffer index subtract 1 in - # a circular buffer way to make sure current outputs become next step's inputs. - num_buffers = self.num_attn_layers + 1 - input_idx = (attn_idx - (step % num_buffers)) % num_buffers - output_idx = (input_idx - 1) % num_buffers - input_name = self.kv_cache_buffer_names[input_idx] - output_name = self.kv_cache_buffer_names[output_idx] - - add_tensor_with_shape(self.buffer[input_name], - f'past_key_value_{layer_idx}', - next_shape) - add_tensor(self.buffer[output_name], - f'present_key_value_{layer_idx}') - else: - key_value_cache = self.buffer[ - f'present_key_value_{layer_idx}'] - add_tensor(key_value_cache, f'past_key_value_{layer_idx}') - add_tensor(key_value_cache, - f'present_key_value_{layer_idx}') - - if self.cross_attention: - cross_cache_buffer = self.buffer[ - f'cross_present_key_value_{layer_idx}'] - add_tensor(cross_cache_buffer, - f'cross_past_key_value_{layer_idx}') - add_tensor(cross_cache_buffer, - f'cross_present_key_value_{layer_idx}') - - for idx in range(self.first_layer, self.last_layer): - if self.layer_types[idx] != 'recurrent': - continue - if self.paged_state: - add_tensor(self.buffer[f'conv_state_ptr_{idx}'], - f'conv_state_ptr_{idx}') - add_tensor(self.buffer[f'rnn_state_ptr_{idx}'], - f'rnn_state_ptr_{idx}') - else: - # conv state - if self.use_mamba_conv1d_plugin: - conv_state_shape = (batch_size, self.conv_kernel - 1, - self.rnn_conv_dim_size) - else: - conv_state_shape = (batch_size, self.rnn_conv_dim_size, - self.conv_kernel - 1) - if step % 2: - add_tensor_with_shape( - self.buffer[f'1_present_conv_state_{idx}'], - f'past_conv_state_{idx}', conv_state_shape) - add_tensor(self.buffer[f'present_conv_state_{idx}'], - f'present_conv_state_{idx}') - else: - add_tensor_with_shape( - self.buffer[f'present_conv_state_{idx}'], - f'past_conv_state_{idx}', conv_state_shape) - add_tensor(self.buffer[f'1_present_conv_state_{idx}'], - f'present_conv_state_{idx}') - # rnn state - rnn_state = self.buffer[f'present_rnn_state_{idx}'] - add_tensor(rnn_state, f'past_rnn_state_{idx}') - add_tensor(rnn_state, f'present_rnn_state_{idx}') - - if self.paged_state and self.has_rnn_layers: - slot_mapping = torch.arange(0, - batch_size, - device='cuda', - dtype=torch.int32) - add_tensor(slot_mapping, 'slot_mapping') - - if self.use_gpt_attention_plugin and self.has_attn_layers: - # generation requests - host_request_types = torch.ones_like(context_lengths, - device='cpu').int() - if self.is_redrafter_mode: - torch.cuda.nvtx.range_push("device_request_types") - device_request_types = torch.ones_like( - context_lengths, device=self.device).int() - add_tensor(device_request_types, 'device_request_types') - torch.cuda.nvtx.range_pop() - if self.is_medusa_mode or self.is_redrafter_mode: - host_past_key_value_lengths = self.sequence_length_buffer.cpu() - else: - # previous [past_kv_length, is_context] has been deprecated. only past_kv_length should be given here - # Note we should use max_context_length here to align to max -- but isn't this done in attn plugin's max_element() already? - host_past_key_value_lengths = torch.tensor( - [max_context_length + step] * (batch_size * beam_width), - dtype=torch.int32, - device='cpu') - add_tensor(host_past_key_value_lengths, - 'host_past_key_value_lengths') - add_tensor(host_request_types, 'host_request_types') - # Sequence lengths are not used in the context phase actually. - sequence_length = self.sequence_length_buffer - - add_tensor_with_shape(sequence_length, 'sequence_length', - (batch_size * beam_width, )) - add_tensor_with_shape(self.host_sink_token_length, - 'host_sink_token_length', (1, )) - add_tensor_with_shape(self.host_max_attention_window_sizes, - f'host_max_attention_window_sizes', - (self.num_attn_layers, )) - if self.remove_input_padding: - add_tensor(host_context_lengths_local, 'host_context_lengths') - else: - if self.has_rnn_layers: - host_request_types = torch.ones_like(context_lengths, - device='cpu').int() - add_tensor(host_request_types, 'host_request_types') - if self.remove_input_padding: - add_tensor(host_context_lengths_local, - 'host_context_lengths') - if self.has_attn_layers: - add_tensor(attention_mask, 'attention_mask') - - if self.mapping.tp_size > 1: - add_tensor(self.all_reduce_workspace, 'all_reduce_workspace') - if self.use_gemm_allreduce_plugin: - found_tensor_names = [ - self.runtime.engine.get_tensor_name(i) - for i in range(self.runtime.engine.num_io_tensors) - ] - for name in found_tensor_names: - if name.startswith("gemm_allreduce_uc_out"): - add_tensor_from_pointer( - self.gemm_allreduce_output_handle.uc_ptr, - name, - shape=(self.gemm_allreduce_output_size), - str_dtype=self.gemm_allreduce_plugin) - if name.startswith("gemm_allreduce_mc_out"): - add_tensor_from_pointer( - self.gemm_allreduce_output_handle.mc_ptr, - name, - shape=(self.gemm_allreduce_output_size), - str_dtype=self.gemm_allreduce_plugin) - if name.startswith("gemm_allreduce_ipc_out"): - add_tensor_from_pointer( - self.gemm_allreduce_output_handle.get_ipc_ptrs(), - name, - shape=(self.gemm_allreduce_output_size), - str_dtype=self.gemm_allreduce_plugin) - - # Since we are using a ping-pong context design and the lora weight remains constant within the same request, - # it is only necessary to set the lora weight for the first two steps. - if self.use_lora_plugin and step < 2: - for idx in range(self.num_layers): - layer_idx = idx + self.first_layer - for lora_module in (self.lora_target_modules + - self.missing_qkv_modules): - lora_ranks = f'{lora_module}_lora_ranks_{layer_idx}' - add_tensor(self.buffer[lora_ranks], lora_ranks) - lora_module = f'{lora_module}_lora_weights_pointers_{layer_idx}' - add_tensor(self.buffer[lora_module], lora_module) - if self.cross_attention and self.remove_input_padding: - add_tensor(encoder_input_lengths.to('cpu'), - 'host_encoder_input_lengths') - - if self.is_medusa_mode: - # Spec Decoding mask and position offsets are fixed for the whole session for Medusa. - add_tensor(self.buffer['spec_decoding_packed_mask'], - 'spec_decoding_packed_mask') - add_tensor(self.buffer['spec_decoding_position_offsets'], - 'spec_decoding_position_offsets') - add_tensor(self.buffer['spec_decoding_generation_lengths'], - 'spec_decoding_generation_lengths') - add_tensor(self.buffer['spec_decoding_use'], 'spec_decoding_use') - - if self.is_redrafter_mode: - set_redrafter_gen_tensors(self, batch_size, add_tensor, - add_tensor_with_shape) - torch.cuda.nvtx.range_pop() - - return tensors - - def _prepare_context_inputs(self, batch_size, context_lengths, - host_context_lengths, use_gpt_attention_plugin, - remove_input_padding, **kwargs): - - last_token_ids = context_lengths.detach().clone() - if (self.is_medusa_mode - or self.is_redrafter_mode) and not remove_input_padding: - # For Medusa, last_token_ids should contain the actual indices - last_token_ids = last_token_ids - 1 # sub 1 from context_lengths for indices - last_token_ids = last_token_ids.reshape([batch_size, -1]) - if (use_gpt_attention_plugin - or self.has_rnn_layers) and remove_input_padding: - last_token_ids = torch.cumsum(last_token_ids, dim=0).int() - ret = {'last_token_ids': last_token_ids} - - if use_gpt_attention_plugin: - max_context_length = kwargs.pop('max_context_length') - if remove_input_padding: - position_ids = torch.concat([ - torch.arange(0, - host_context_lengths[i], - dtype=torch.int32, - device='cuda') for i in range(batch_size) - ]) - else: - position_ids = torch.tensor(range(max_context_length), - dtype=torch.int32, - device='cuda').reshape( - [1, - -1]).expand([batch_size, -1]) - - perf_knob_tensor_size = 16 - context_runtime_perf_knobs = torch.tensor([-1] * - perf_knob_tensor_size, - dtype=torch.int64) - if self.multi_block_mode: - context_runtime_perf_knobs[0] = 1 # multi_block_mode - if self.enable_context_fmha_fp32_acc: - context_runtime_perf_knobs[ - 1] = 1 # enable_context_fmha_fp32_acc - ret['host_runtime_perf_knobs'] = context_runtime_perf_knobs - else: - if self.has_attn_layers: - input_ids = kwargs.pop('input_ids') - pad_id = kwargs.pop('pad_id', None) - attention_mask = _prepare_attention_mask(input_ids, pad_id) - position_ids = attention_mask.long().cumsum(-1) - 1 - position_ids.masked_fill_(attention_mask == 0, 1) - position_ids = position_ids.int() - ret['attention_mask'] = attention_mask - - if self.has_position_embedding and self.has_attn_layers: - ret['position_ids'] = position_ids - - if self.is_redrafter_mode: - self.buffer['position_ids_base'] = context_lengths.clone() - # NOTE: Generate random tensors using torch - redrafter_prepare_random_tensors(self, batch_size, initialize=True) - - return ret - - def _prepare_generation_inputs(self, batch_size, context_lengths, - use_gpt_attention_plugin, - remove_input_padding, **kwargs): - torch.cuda.nvtx.range_push("_prepare_generation_inputs") - - step = kwargs.pop('step') - last_token_ids = torch.ones_like(context_lengths) - if use_gpt_attention_plugin and (self.is_medusa_mode - or self.is_redrafter_mode): - if remove_input_padding: - if self.is_medusa_mode: - # For Medusa, last_token_ids should be [bs * seq] and should contain the actual indices (starts from 1) - last_token_ids = torch.ones(batch_size * - (self.num_draft_tokens + 1), - dtype=torch.int32, - device=context_lengths.device) - elif self.is_redrafter_mode: - torch.cuda.nvtx.range_push("last_token_ids_1s") - # update last_token_ids here (buffers already swapped) - last_token_ids = torch.ones(self.host_total_gen_token, - dtype=torch.int32, - device=context_lengths.device) - torch.cuda.nvtx.range_pop() - else: - # For Medusa, last_token_ids should be [bs, seq] and should contain the actual indices (starts from 0) - last_token_ids = torch.arange(self.num_draft_tokens + 1, - dtype=torch.int32, - device=context_lengths.device) - last_token_ids = last_token_ids.expand([batch_size, -1]) - if (use_gpt_attention_plugin - or self.has_rnn_layers) and remove_input_padding: - torch.cuda.nvtx.range_push("last_token_ids_cumsum") - last_token_ids = torch.cumsum(last_token_ids, dim=0).int() - torch.cuda.nvtx.range_pop() - ret = {'last_token_ids': last_token_ids} - - if use_gpt_attention_plugin: - if self.is_redrafter_mode: - torch.cuda.nvtx.range_push("position_ids_update") - # set position_ids - # buffers are swapped but sequence_length is not updated at this point - - if step != 0: - self.buffer['position_ids_base'] += self.buffer[ - 'num_accepted_tokens'] - position_ids = self.buffer['packed_position_ids'].view( - -1)[:self.host_total_gen_token] - if step == 0: - position_ids -= 1 - - torch.cuda.nvtx.range_pop() - else: - position_ids = context_lengths + step - if not remove_input_padding: - position_ids = torch.unsqueeze(position_ids, 1) - - perf_knob_tensor_size = 16 - gen_runtime_perf_knobs = torch.tensor([-1] * perf_knob_tensor_size, - dtype=torch.int64) - if self.multi_block_mode: - gen_runtime_perf_knobs[0] = 1 # multi_block_mode - if self.enable_context_fmha_fp32_acc: - gen_runtime_perf_knobs[1] = 1 # enable_context_fmha_fp32_acc - ret['host_runtime_perf_knobs'] = gen_runtime_perf_knobs - elif self.has_attn_layers: - attention_mask = kwargs.pop('attention_mask') - num_beams = kwargs.pop('num_beams') - attention_mask = torch.cat((attention_mask, - attention_mask.new_ones( - (batch_size * num_beams, 1))), - dim=-1).contiguous() - position_ids = attention_mask.long().cumsum(-1) - 1 - position_ids.masked_fill_(attention_mask == 0, 1) - position_ids = position_ids[:, -1].unsqueeze(-1) - position_ids = position_ids.int() - ret['attention_mask'] = attention_mask - - if self.has_position_embedding and self.has_attn_layers: - ret['position_ids'] = position_ids - if self.is_redrafter_mode: - # buffers are already swapped - # convert spec_decoding_mask to spec_decoding_packed_mask - redrafter_convert_spec_decoding_mask_to_packed_mask( - self, self.buffer['spec_decoding_generation_lengths']) - # NOTE: Generate random tensors using torch - redrafter_prepare_random_tensors(self, batch_size) - torch.cuda.nvtx.range_pop() - - return ret - - def _prepare_cross_attention_mask(self, batch_size, context_lengths, - cross_attention_mask): - cross_attention_mask_for_context = [] - cross_attention_mask_for_gen = [] - max_decoder_input_length = torch.max(context_lengths).item() - for batch_idx in range(batch_size): - decoder_input_length = context_lengths[batch_idx].item() - local_mask_for_context = cross_attention_mask[ - batch_idx][:decoder_input_length, :] - local_mask_for_gen = cross_attention_mask[batch_idx][ - decoder_input_length:, :] - if not self.use_gpt_attention_plugin: - local_mask_for_context = local_mask_for_context.unsqueeze(0) - if not self.remove_input_padding: - local_mask_for_context = torch.nn.functional.pad( - local_mask_for_context, - (0, 0, 0, - (max_decoder_input_length - decoder_input_length)), - "constant", False) - local_mask_for_gen = torch.nn.functional.pad( - local_mask_for_gen, - (0, 0, 0, - (max_decoder_input_length - decoder_input_length)), - "constant", False) - cross_attention_mask_for_context.append(local_mask_for_context) - # add additional dimension for batch size. - cross_attention_mask_for_gen.append(local_mask_for_gen.unsqueeze(0)) - - return torch.concat(cross_attention_mask_for_context), torch.concat( - cross_attention_mask_for_gen) - - def pp_communicate_new_tokens(self, should_stop, cache_indir, - sequence_length): - if self.mapping.is_last_pp_rank(): - for pg in self.mapping.pp_group: - if pg == self.mapping.rank: - continue - should_stop = should_stop.to(self.device) - self.nccl_comm.send(should_stop, pg) - self.nccl_comm.send(cache_indir, pg) - self.nccl_comm.send(sequence_length, pg) - self.nccl_comm.send(self.new_tokens, self.mapping.pp_group[0]) - else: - should_stop = torch.zeros(1, dtype=torch.bool, device=self.device) - self.nccl_comm.recv(should_stop, self.mapping.pp_group[-1]) - self.nccl_comm.recv(cache_indir, self.mapping.pp_group[-1]) - self.nccl_comm.recv(sequence_length, self.mapping.pp_group[-1]) - if self.mapping.is_first_pp_rank(): - self.nccl_comm.recv(self.new_tokens, self.mapping.pp_group[-1]) - return should_stop - - def pp_communicate_final_output_ids(self, final_output_ids, batch_size, - beam_width): - if self.mapping.is_last_pp_rank(): - self.nccl_comm.send(final_output_ids, self.mapping.pp_group[0]) - elif self.mapping.is_first_pp_rank(): - final_output_ids = torch.zeros( - (batch_size, beam_width, self.max_seq_length), - dtype=torch.int32, - device=self.device) - self.nccl_comm.recv(final_output_ids, self.mapping.pp_group[-1]) - return final_output_ids - - def finalize_decoder(self, - context_lengths, - batch_size, - beam_width, - scfg, - in_progress=False): - final_output_ids = None - if self.mapping.is_last_pp_rank(): - # output shape of self.gather_tree: [batch_size, beam_width, output_len] - beam_hyps_args = [ - self.beam_hyps_output_ids_cba, self.beam_hyps_seq_len_cba, - self.beam_hyps_cum_log_probs_cba, - self.beam_hyps_normed_scores_cba, self.beam_hyps_log_probs_cba, - self.beam_hyps_min_normed_scores, self.beam_hyps_num_beams, - self.beam_hyps_is_done - ] - - if scfg.use_beam_hyps and in_progress: - # self.gather_tree modifies these args. - # In streaming mode, this results in incorrect decoding in the following steps. - beam_hyps_args = copy.deepcopy(beam_hyps_args) - - final_output_ids = self.gather_tree( - self.sequence_length_buffer, self.output_ids, self.parent_ids, - self.end_ids, context_lengths, self.cum_log_probs, - self.log_probs, self.log_probs_tiled, *beam_hyps_args, - self.finished, self.length_penalty, batch_size, beam_width, - self.max_seq_length, scfg.use_beam_hyps) - - # Communicate ranks in Pipeline Parallelism - if self.mapping.has_pp(): - final_output_ids = self.pp_communicate_final_output_ids( - final_output_ids, batch_size, beam_width) - - return final_output_ids - - def find_best_medusa_path(self, - batch_size, - input_ids: torch.Tensor, - next_logits, - temp=0): - assert input_ids.shape[-1] == self.num_draft_tokens + 1 - best_path = [0] * batch_size - best_path_len = [1] * batch_size - next_tokens = [None] * batch_size - zero_pad = torch.zeros((batch_size, 1), - dtype=input_ids.dtype, - device=input_ids.device) - input_ids = torch.cat((input_ids, zero_pad), dim=-1) - if temp == 0: - new_tokens_raw = torch.argmax( - next_logits, dim=-1 - ) # TODO: can be done by treating [bs, nT, vocab] as [bs*nT, vocab] and using decoderOp? - new_tokens = torch.cat((new_tokens_raw, zero_pad), dim=-1) - input_paths = [ - input_ids[b, self.medusa_paths] for b in range(batch_size) - ] - new_paths = [ - new_tokens[b, self.medusa_paths] for b in range(batch_size) - ] - for b in range(batch_size): - equality = input_paths[b][:, 1:] == new_paths[b][:, :-1] - paths_correct_len = torch.cumprod(equality.int(), - dim=1).sum(dim=1) - best_path_len[b] = paths_correct_len.max().item() + 1 - if best_path_len[b] > 1: - best_path[b] = torch.argmax(paths_correct_len) - next_tokens[b] = new_paths[b][ - best_path[b]][:best_path_len[b]].clone() - - return best_path, best_path_len, next_tokens - - def filter_medusa_logits(self, batch_size, best_path, best_path_lengths, - medusa_logits): - """ - medusa_logits is of shape [nMH, bs, nMT+1, vocab] - - Returns [nMH, bs, vocab] - """ - filtered_logits = torch.empty( - (self.num_medusa_heads, batch_size, self.vocab_size_padded), - dtype=medusa_logits.dtype, - device=medusa_logits.device) - medusa_logits = medusa_logits.view(self.num_medusa_heads, batch_size, - self.num_draft_tokens + 1, -1) - for b in range(batch_size): - idx = self.medusa_paths[best_path[b], best_path_lengths[b] - 1] - filtered_logits[:, b, ...] = medusa_logits[:, b, idx, ...] - return filtered_logits - - def get_next_medusa_tokens(self, batch_size, next_medusa_logits): - next_medusa_tokens = [ - torch.zeros((batch_size, 1), - dtype=torch.int32, - device=next_medusa_logits.device) - ] # dummy token for now, TODO: update tree_ids and remove this - for i in range(self.num_medusa_heads): - medusa_token = torch.topk(next_medusa_logits[i, :, :], - self.medusa_topks[i], - dim=-1).indices - next_medusa_tokens.append(medusa_token) - next_medusa_tokens = torch.cat(next_medusa_tokens, dim=-1) - return next_medusa_tokens - - def locate_accepted_draft_tokens(self, batch_size, best_path, best_path_len, - draft_paths): - torch.cuda.nvtx.range_push("locate_accepted_draft_tokens") - best_path_len_tensor = best_path_len if isinstance( - best_path_len, torch.Tensor) else torch.tensor( - best_path_len, dtype=torch.int, device='cuda') - accepted_draft_token_counts = torch.maximum( - best_path_len_tensor - 1, - torch.tensor([0], device=best_path_len_tensor.device)) - accepted_draft_token_offsets = torch.zeros(batch_size + 1, - dtype=torch.int32, - device='cuda') - accepted_draft_token_offsets[1:] = torch.cumsum( - accepted_draft_token_counts, dim=0) - accepted_draft_token_offsets_cpu = accepted_draft_token_offsets.to( - 'cpu') - packed_accepted_draft_tokens_indices = torch.empty( - accepted_draft_token_offsets_cpu[batch_size], - dtype=torch.int32, - device='cuda') - for seq_idx in range(batch_size): - cur_draft_paths = draft_paths if self.is_medusa_mode else draft_paths[ - seq_idx] - seq_start = accepted_draft_token_offsets_cpu[seq_idx] - seq_end = accepted_draft_token_offsets_cpu[seq_idx + 1] - seq_accepted_draft_count = seq_end - seq_start - best_path_idx = best_path[seq_idx].cpu() if isinstance( - best_path[seq_idx], torch.Tensor) else best_path[seq_idx] - seq_accepted_token_indices = cur_draft_paths[ - best_path_idx, 1:1 + seq_accepted_draft_count] - packed_accepted_draft_tokens_indices[ - seq_start:seq_end] = seq_accepted_token_indices - 1 - # print("KV offsets & indices", accepted_draft_token_offsets, - # packed_accepted_draft_tokens_indices,) - torch.cuda.nvtx.range_pop() - return accepted_draft_token_offsets, packed_accepted_draft_tokens_indices - - def update_output_ids_by_offset(self, new_generated_ids, offsets): - # output_ids [batch_size, padded_input_length] - # new_generated_ids [batch_size, padded_accepted_length] - # offsets [batch_size] - # FIXME: using fused kernel to update the padded output ids. - batch_size = self.output_ids.shape[0] - for b in range(batch_size): - self.output_ids[b, offsets[b]:( - offsets[b] + self.accept_lengths[b] - )] = new_generated_ids[b][:self.accept_lengths[b]] - return - - def next_medusa_input_ids(self): - # self.new_tokens [batch_size, padded_accepted_length] - # self.accept_lengths [batch_size] - # self.medusa_new_tokens [batch_size, num_draft_tokens] - # FIXME: using fused kernel to generate the new medusa input ids. - batch_size = self.new_tokens.shape[0] - for b in range(batch_size): - self.generation_input_ids[b, 0] = self.new_tokens[ - b, self.accept_lengths[b] - 1] - self.generation_input_ids[b, 1:] = self.medusa_output_tokens[b, :] - - def reorder_kv_cache_for_beam_search( - self, - batch_size: int, - beam_width: int, - max_context_length: int, - step: int, - ): - if self.use_gpt_attention_plugin: - # Do nothing. - return - - # WAR: This degrades the latency performance in beam search - # due to memcpy. Recommend to use gpt attention plugin instead. - assert self.buffer is not None - assert self.parent_ids.shape[:2] == (batch_size, beam_width) - - cache_shape = (batch_size * beam_width, 2, self.get_num_heads_kv(), - max_context_length + step, self.head_size) - - import functools - numel = functools.reduce(lambda x, y: x * y, cache_shape) - - # attention layer num + 1 extra buffer. - num_buffers = self.num_attn_layers + 1 - for i in self.attn_to_general_idx: - # Cyclic buffers, an output becomes the next step's input. - input_idx = (i - (step % num_buffers)) % num_buffers - presents = self.buffer[self.kv_cache_buffer_names[input_idx]] - presents = presents.view(-1)[:numel].view(*cache_shape) - # parent_ids = (batch, beam, max_seq_len) - parent_ids = self.parent_ids[..., - max_context_length + step].view(-1) - - for batch_beam in range(batch_size * beam_width): - batch = batch_beam // beam_width - if parent_ids[batch_beam] != batch_beam % beam_width: - # Update past kv cache to parent beam's cache. - src_bbid = batch * beam_width + parent_ids[batch_beam] - presents[batch_beam, ...] = presents[src_bbid, ...] - - # OPTIMIZE: need to optimize this early-stop workflow. - def early_stop_criteria(self, batch_size, step, should_stop): - for b in range(batch_size): - if self.medusa_should_stop[b]: - self.accept_lengths[b] = 0 - continue - # output sequence length criteria. - prev_total_output_length = self.total_accept_lengths[b] - # end id criteria. - end_id_mask = self.new_tokens[ - b, :self.accept_lengths[b]] == self.end_ids[b] - should_stop_with_end_id = torch.any(end_id_mask) - self.medusa_should_stop[b] = self.medusa_should_stop[b] or ( - prev_total_output_length + self.accept_lengths[b] - >= self.max_new_tokens) or should_stop_with_end_id - # update accept lengths for the current step. - if (prev_total_output_length + self.accept_lengths[b] - >= self.max_new_tokens): - self.accept_lengths[b] = min( - self.max_new_tokens - prev_total_output_length, - self.accept_lengths[b]) - if should_stop_with_end_id: - # get the position of first end_id. - end_id_pos = (end_id_mask).nonzero(as_tuple=True)[0] - self.accept_lengths[b] = min(end_id_pos[0] + 1, - self.accept_lengths[b]) - self.total_accept_lengths[b] += self.accept_lengths[b] - - should_stop[0] = should_stop[0] or (step == self.max_new_tokens - - 1) or torch.all( - self.medusa_should_stop) - return should_stop - - def medusa_decode_and_verify(self, step, batch_size, logits): - medusa_logits = self.buffer['medusa_logits'] - best_path = None - best_path_lengths = None - if step == 0: - # logits buffer is of shape [bs, medusa_tokens+1, vocab] - # but during context phase, we get only [bs, 1, vocab] but contiguous - logits = logits.view(-1)[:batch_size * logits.shape[-1]].view( - batch_size, -1) - next_main_token_logits = logits.to(self.decoder_logits_dtype) - next_main_token = torch.argmax(next_main_token_logits, - dim=-1, - keepdim=True) - self.new_tokens = next_main_token - # NOTE: only one token's medusa logit will be written in. - medusa_logits = medusa_logits.view(self.num_draft_tokens + 1, - -1)[0, ...] - next_medusa_logits = medusa_logits.reshape( - self.num_medusa_heads, batch_size, - -1).to(self.decoder_logits_dtype) - next_medusa_tokens = self.get_next_medusa_tokens( - batch_size, next_medusa_logits) - self.medusa_output_tokens = next_medusa_tokens[:, - self.medusa_tree_ids[ - -self. - num_draft_tokens:]] - self.accept_lengths = torch.ones([batch_size], - dtype=torch.int32, - device=self.device) - else: - next_token_logits = logits.to(self.decoder_logits_dtype) - - best_path, best_path_lengths, next_main_tokens = self.find_best_medusa_path( - batch_size, self.generation_input_ids.view(batch_size, -1), - next_token_logits.view(batch_size, self.num_draft_tokens + 1, - -1)) - self.accept_lengths = torch.tensor(best_path_lengths, - device=self.device) - self.new_tokens = torch.nested.to_padded_tensor( - torch.nested.nested_tensor(next_main_tokens, dtype=torch.int32), - self.end_ids[0]) #FIXME end id padding. - next_medusa_logits = self.filter_medusa_logits( - batch_size, best_path, best_path_lengths, medusa_logits) - next_medusa_tokens = self.get_next_medusa_tokens( - batch_size, next_medusa_logits) - - self.medusa_output_tokens = next_medusa_tokens[:, - self.medusa_tree_ids[ - -self. - num_draft_tokens:]] - return best_path, best_path_lengths - - def process_logits_including_draft(self, step, batch_size, logits, - next_step_buffer): - """ - 1. Process logits to tokens and validate (Medusa) or process outputs (ReDrafter) - 2. Extract early stop criteria here : self.accept_length - 3. Update output ids : needs self.new_tokens and past_sequence_length - 4. Get next input_ids : self.[new_tokens, accept_lengths, medusa_output_tokens] - 5. Update KV cache : self.[sequence_length, num_draft_tokens] - 6. Update sequence_length_buffer and past_kv_length - """ - should_stop = torch.tensor([False], dtype=bool) - if self.is_medusa_mode: - # NOTE: this function call also updates self.[accept_lengths, new_tokens, medusa_output_tokens] - best_path, best_path_lengths = self.medusa_decode_and_verify( - step, batch_size, logits) - last_draft_paths = self.medusa_paths - # print(best_path, self.new_tokens, self.medusa_output_tokens) - last_draft_tokens_len = self.num_draft_tokens if step > 0 else 0 - cur_draft_tokens_len = self.num_draft_tokens - elif self.is_redrafter_mode: - # buffers are swapped at this point - last_draft_tokens = self.buffer['next_draft_tokens'] - new_draft_tokens = self.buffer['draft_tokens'] - last_draft_paths = self.buffer["next_draft_indices"] - last_draft_tokens_len = self.buffer[ - 'next_spec_decoding_generation_lengths'] - 1 if step > 0 else 0 - cur_draft_tokens_len = self.buffer[ - 'spec_decoding_generation_lengths'] - 1 - - best_path, best_path_lengths = process_redrafter_outputs( - self, step, batch_size, last_draft_tokens, new_draft_tokens) - # NOTE: stop criteria - torch.cuda.nvtx.range_push("early_stop_check") - if step == 0: - self.total_accept_lengths = self.accept_lengths.clone() - self.medusa_should_stop = torch.eq(self.new_tokens.reshape(-1), - self.end_ids) - should_stop[0] = torch.equal( - self.new_tokens.reshape(-1), - self.end_ids) or (step == self.max_new_tokens - 1) - else: - should_stop = self.early_stop_criteria(batch_size, step, - should_stop) - torch.cuda.nvtx.range_pop() - # NOTE: self.accept_lengths are the lengths of accepted tokens in the current step - # NOTE: self.sequence_length_buffer = num_past_kv_cache (accepted) + accept_lengths - torch.cuda.nvtx.range_push("update_output_ids") - self.update_output_ids_by_offset( - self.new_tokens, - self.sequence_length_buffer - last_draft_tokens_len) - torch.cuda.nvtx.range_pop() - - if step != self.max_new_tokens - 1 and not should_stop.item(): - if self.is_medusa_mode: - self.next_medusa_input_ids() - if step != 0: - assert best_path is not None and best_path_lengths is not None - accepted_draft_token_offsets, packed_accepted_draft_tokens_indices = self.locate_accepted_draft_tokens( - batch_size, best_path, best_path_lengths, last_draft_paths) - # update the KV cache - torch.cuda.nvtx.range_push("kv_update") - self.kv_cache_updater.update( - accepted_draft_token_offsets, - packed_accepted_draft_tokens_indices, - self.sequence_length_buffer, last_draft_tokens_len) - torch.cuda.nvtx.range_pop() - - self.sequence_length_buffer += self.accept_lengths + cur_draft_tokens_len - last_draft_tokens_len - else: - self.sequence_length_buffer += cur_draft_tokens_len + 1 - - # NOTE: set the accepted tokens for the last step. - if should_stop.item(): - # remove num_draft_tokens for next generation. - # Runtime: denotes kv cache length start positions. - # Output: denotes the length of sequence length (input ids + output ids) - self.sequence_length_buffer += self.accept_lengths - last_draft_tokens_len - - if next_step_buffer is not None: - next_step_buffer['host_past_key_value_lengths'].to_torch().copy_( - self.sequence_length_buffer) - - return should_stop - - def handle_per_step( - self, - *, - cache_indirections: list, - step: int, - batch_size: int, - max_context_length: int, - beam_width: int, - input_ids: torch.Tensor, - hidden_states: torch.Tensor, - scfg: SamplingConfig, - kv_cache_block_offsets: torch.Tensor, - host_kv_cache_block_offsets: torch.Tensor, - cross_kv_cache_block_offsets: torch.Tensor, - host_cross_kv_cache_block_offsets: torch.Tensor, - prompt_embedding_table: torch.Tensor, - tasks: torch.Tensor, - context_lengths: torch.Tensor, - host_context_lengths, - attention_mask: torch.Tensor, - cross_attention_mask_for_context: torch.Tensor, - cross_attention_mask_for_gen: torch.Tensor, - prompt_vocab_size: torch.Tensor, - ite: int, - sequence_limit_lengths: torch.Tensor, - sequence_lengths: torch.Tensor, - next_step_tensors: Dict[str, RuntimeTensor], - stop_words_data, - bad_words_data, - encoder_output: torch.Tensor, - encoder_input_lengths: torch.Tensor, - stopping_criteria: StoppingCriteria, - logits_processor: LogitsProcessor, - output_generation_logits: bool, - **kwargs, - ): - if self.debug_mode: - print( - f"=================================== STEP {step} ==================================" - ) - if step % 2: - context = self.runtime.context_0 - this_src_cache_indirection = cache_indirections[1] - this_tgt_cache_indirection = cache_indirections[0] - next_src_cache_indirection = cache_indirections[0] - else: - context = self.runtime.context_1 - this_src_cache_indirection = cache_indirections[0] - this_tgt_cache_indirection = cache_indirections[1] - next_src_cache_indirection = cache_indirections[1] - - position_ids_raw = kwargs.get('position_ids', None) - skip_cross_attn_blocks = kwargs.get('skip_cross_attn_blocks', None) - language_adapter_routings = kwargs.get('language_adapter_routings', - None) - if step == 0: - model_inputs = self._prepare_context_inputs( - batch_size=batch_size, - context_lengths=context_lengths, - host_context_lengths=host_context_lengths, - use_gpt_attention_plugin=self.use_gpt_attention_plugin, - remove_input_padding=self.remove_input_padding, - max_context_length=max_context_length, - input_ids=input_ids, - pad_id=scfg.pad_id, - eos_id=scfg.end_id) - - if position_ids_raw is None: - # default iota position ids - position_ids = model_inputs.get('position_ids', None) - else: - # user input position ids - if self.remove_input_padding: - position_ids = torch.cat(position_ids_raw, dim=0) - else: - padded_position_ids = torch.nn.utils.rnn.pad_sequence( - position_ids_raw, batch_first=True, padding_value=0) - position_ids = padded_position_ids - last_token_ids = model_inputs.get('last_token_ids') - attention_mask = model_inputs.get('attention_mask', None) - context_runtime_perf_knobs = model_inputs.get( - 'host_runtime_perf_knobs', None) - host_context_progress = torch.tensor([0], dtype=torch.int64) - - if self.paged_kv_cache and self.has_attn_layers: - host_kv_cache_block_offsets = self.pools_kv_cache_manager.get_block_offsets( - beam_width=1) - kv_cache_block_offsets = host_kv_cache_block_offsets.to('cuda') - if self.cross_attention: - host_cross_kv_cache_block_offsets = self.cross_pools_kv_cache_manager.get_block_offsets( - beam_width=1) - cross_kv_cache_block_offsets = host_cross_kv_cache_block_offsets.to( - 'cuda') - - ctx_tensors = self._get_context_shape_buffer( - input_ids, - context_lengths, - host_context_lengths, - position_ids, - last_token_ids, - attention_mask, - cross_attention_mask_for_context, - this_src_cache_indirection, - kv_cache_block_offsets, - host_kv_cache_block_offsets, - cross_kv_cache_block_offsets, - host_cross_kv_cache_block_offsets, - hidden_states, - prompt_embedding_table, - tasks, - prompt_vocab_size, - encoder_output, - encoder_input_lengths, - host_runtime_perf_knobs=context_runtime_perf_knobs, - host_context_progress=host_context_progress, - skip_cross_attn_blocks=skip_cross_attn_blocks, - language_adapter_routings=language_adapter_routings, - ) - - context = self.runtime.ctx_context - self.runtime._set_tensors(context, ctx_tensors) - if self.debug_mode: - self.debug_buffer = { - name: tensor.to_torch() - for name, tensor in ctx_tensors.items() - } - if self.cuda_graph_mode: - # context mode, clean cuda graph instances - self.runtime.cuda_graph_instances = [None for _ in range(2)] - - if self.debug_mode and False: # TODO: after TRT bug is fixed - self.runtime._check_tensors(context) - # dynamic_decoder currently use torch's current stream, so must let TRT enqueue use same stream here - stream = torch.cuda.current_stream().cuda_stream - instance_idx = step % 2 - if self.cuda_graph_mode and self.runtime.cuda_graph_instances[ - instance_idx] is not None: - # launch cuda graph - CUASSERT( - cudart.cudaGraphLaunch( - self.runtime.cuda_graph_instances[instance_idx], stream)) - ok = True - else: - ok = self.runtime._run(context, stream) - - if not ok: - raise RuntimeError(f"Executing TRT engine failed step={step}!") - - # TODO: remove this Windows WAR after https://nvbugs/4460474 is fixed. - if platform.system() == "Windows" or self.debug_mode: - torch.cuda.synchronize() - - context_logits = None - if self.mapping.is_last_pp_rank(): - if step == 0 and self.gather_context_logits: - assert not self.is_medusa_mode and not self.is_redrafter_mode - context_logits = self.buffer['logits'].detach().clone() - # gather last token of context - if self.remove_input_padding: - # reshape self.buffer['logits'] from [bs, max_context_length, vocab] - # to [1, bs * max_context_length, vocab] - # Note that the data are put in the buffer without padding although - # the allocated buffer has padding. - self.buffer['logits'] = self.buffer['logits'].reshape( - [1, -1, self.vocab_size_padded]) - self.buffer['logits'] = torch.index_select( - self.buffer['logits'], 1, - last_token_ids - 1).view(batch_size, - self.vocab_size_padded) - else: - last_token_ids = last_token_ids.reshape(batch_size, 1, 1) - last_token_ids = last_token_ids.expand( - batch_size, 1, self.vocab_size_padded) - 1 - self.buffer['logits'] = torch.gather( - self.buffer['logits'], - dim=1, - index=last_token_ids.to(dtype=torch.int64)).view( - batch_size, self.vocab_size_padded) - - if step == 0 and beam_width > 1: - assert not self.is_medusa_mode and not self.is_redrafter_mode - assert not self.has_rnn_layers - # these tiled tensors are returned by handle_per_step(), so they can relay to the next generation calls - if not self.use_gpt_attention_plugin: - attention_mask = _tile_beam_width(attention_mask, beam_width) - context_lengths = _tile_beam_width(context_lengths, beam_width) - host_context_lengths = _tile_beam_width(host_context_lengths, - beam_width) - if encoder_input_lengths is not None: - encoder_input_lengths = _tile_beam_width( - encoder_input_lengths, beam_width) - - if tasks is not None: - tasks = _tile_beam_width(tasks, beam_width) - - # Move tiling before logit computing of context - if not self.paged_kv_cache: - for key in self.buffer: - # Note: this tiles both self attn cache and cross attn - # cache! both names contain "present_key_value" - if "present_key_value" in key: - if self.use_gpt_attention_plugin: - self.buffer[key] = _tile_beam_width( - self.buffer[key], beam_width) - else: - # In the OOTB path, KV cache should be contiguously - # tiled since TRT engine allocates past_kv cache of - # length context_length, i.e., we need a buffer of - # shape (batch * beam, 2, heads, context_length, head_size). - b, _, h, _, d = self.buffer[key].shape - numel = 2 * b * h * (max_context_length + step) * d - self.buffer[key] = _contiguous_tile_beam_width( - self.buffer[key], numel, beam_width) - - if self.mapping.is_last_pp_rank(): - self.buffer['logits'] = _tile_beam_width( - self.buffer['logits'], beam_width) - - generation_logits = None - if self.mapping.is_last_pp_rank(): - if self.gather_generation_logits or output_generation_logits: - generation_logits = self.buffer['logits'].detach().clone() - - # Initialize sequence_lengths (no paddings) for the generation phase. - if step == 0 and not self.is_medusa_mode and not self.is_redrafter_mode: # Medusa/ReDrafter has its own logic - self.sequence_length_buffer = context_lengths.detach().clone() - - if self.is_redrafter_mode: - # to simplify some processing logic, always swap buffers after execution - exchange_redrafter_buffers(self) - - # NOTE: handle next step. - if not step == self.max_new_tokens - 1: - # Set shape and address for the next step - model_inputs = self._prepare_generation_inputs( - batch_size=batch_size, - context_lengths=context_lengths, - use_gpt_attention_plugin=self.use_gpt_attention_plugin, - remove_input_padding=self.remove_input_padding, - step=step, - num_beams=beam_width, - attention_mask=attention_mask, - ) - - if position_ids_raw is None: - position_ids = model_inputs.get('position_ids', None) - else: - position_ids = torch.cat( - [p[-1:] + step + 1 for p in position_ids_raw], dim=0) - if not self.remove_input_padding: - position_ids = torch.unsqueeze(position_ids, 1) - last_token_ids = model_inputs.get('last_token_ids') - attention_mask = model_inputs.get('attention_mask', None) - gen_runtime_perf_knobs = model_inputs.get('host_runtime_perf_knobs', - None) - host_context_progress = torch.tensor([0], dtype=torch.int64) - - # Prepare for the next step, and always allocate 1 token slot. - if self.paged_kv_cache and self.has_attn_layers: - # Iterate to the next step in KV cache manager. - # Increase number of tokens for all unfinished sequences. - # And allocate new blocks if needed. - # We set this to False for all sequences, since we use only length criterion to stop now - # OPTIMIZE: find a better of adding multiple tokens for paged kv cache. - torch.cuda.nvtx.range_push("paged_kv_alloc") - if self.is_redrafter_mode and self.max_draft_tokens > 0: - add_token_count = (self.max_draft_tokens + - 1) * 2 if step == 0 else torch.max( - self.accept_lengths).item() - assert add_token_count > 0 - for _ in range(add_token_count): - self.pools_kv_cache_manager.step([False] * batch_size) - if self.is_medusa_mode and self.num_draft_tokens > 0: - # Allocate kv cache token slots for next step. - # Make sure there are always > (num_draft_tokens + 1) free token slots. - # Allocate (num_draft_tokens + 1) * 2 for safety as we don't know the current step or next step's accepted lengths. - add_token_count = (self.num_draft_tokens + - 1) * 2 if step == 0 else torch.max( - self.accept_lengths).item() - assert add_token_count > 0 - for _ in range(add_token_count): - self.pools_kv_cache_manager.step([False] * batch_size) - else: - self.pools_kv_cache_manager.step([False] * batch_size) - torch.cuda.nvtx.range_pop() - torch.cuda.nvtx.range_push("paged_kv_post_alloc") - host_kv_cache_block_offsets = self.pools_kv_cache_manager.get_block_offsets( - beam_width) - kv_cache_block_offsets = host_kv_cache_block_offsets.to('cuda') - if self.cross_attention: - host_cross_kv_cache_block_offsets = self.cross_pools_kv_cache_manager.get_block_offsets( - beam_width) - cross_kv_cache_block_offsets = host_cross_kv_cache_block_offsets.to( - 'cuda') - torch.cuda.nvtx.range_pop() - - next_context = self.runtime.context_1 if step % 2 else self.runtime.context_0 - cross_attention_mask_step = None - if cross_attention_mask_for_gen is not None: - # cross_attention_mask_for_gen shape [batch_size, max_output_length, max_encoder_input_length] - decode_step = step - if decode_step == 0: - decode_step += 1 - if self.use_gpt_attention_plugin: - cross_attention_mask_step = cross_attention_mask_for_gen[:, ( - decode_step - 1), :] - else: - cross_attention_mask_step = cross_attention_mask_for_gen[:, ( - decode_step - 1):decode_step, :] - next_step_tensors = self._get_next_step_shape_buffer( - batch_size, - beam_width, - max_context_length, - step, - context_lengths, - host_context_lengths, - position_ids, - last_token_ids, - attention_mask, - cross_attention_mask_step, - next_src_cache_indirection, - kv_cache_block_offsets, - host_kv_cache_block_offsets, - cross_kv_cache_block_offsets, - host_cross_kv_cache_block_offsets, - hidden_states, - prompt_embedding_table, - tasks, - prompt_vocab_size, - encoder_output, - encoder_input_lengths, - host_runtime_perf_knobs=gen_runtime_perf_knobs, - host_context_progress=host_context_progress, - skip_cross_attn_blocks=skip_cross_attn_blocks, - language_adapter_routings=language_adapter_routings) - - # there are some tensors created inside the _get_next_step_shape_buffer, not owned by any object - # needs to pro-long the life time of the tensors inside the next_step_tensors array - # otherwise, it maybe released before the next step actually enqueued - # one way to prolong it is to return the list, and destroy it in next step by assigning new values - torch.cuda.nvtx.range_push("_set_tensors") - self.runtime._set_tensors(next_context, next_step_tensors) - torch.cuda.nvtx.range_pop() - - if logger.level == "verbose": - self.runtime.print_context_info( - next_context, int(next_context == self.runtime.context_1)) - - if self.cuda_graph_mode: - self._capture_cuda_graph_and_instantiate( - next_context, stream, step) - - should_stop = None - logits = None - if self.mapping.is_last_pp_rank(): - logits = self.buffer['logits'] - if self.is_redrafter_mode: - should_stop = self.process_logits_including_draft( - step, batch_size, logits, next_step_tensors) - elif logits is not None: - if self.is_medusa_mode: - should_stop = self.process_logits_including_draft( - step, batch_size, logits, next_step_tensors) - else: - if logits_processor is not None: - final_output_ids = self.finalize_decoder( - context_lengths, - batch_size, - beam_width, - scfg, - in_progress=True) - # keep the shape as same as huggingface stopping_criteria - final_output_ids_ = final_output_ids.reshape( - -1, final_output_ids.size(-1)) - logits = logits_processor(step, final_output_ids_, - logits) - self.buffer['logits'] = logits - # [batch_size x beam_width, vocab_size_padded] -> [batch_size, beam_width, vocab_size_padded] - next_token_logits = logits.reshape( - (batch_size, beam_width, - -1)).to(self.decoder_logits_dtype) - decode_step = step + max_context_length - - stop_words_list_ptrs, stop_words_lens, max_stop_words_len = stop_words_data - bad_words_list_ptrs, bad_words_lens, max_bad_words_len = bad_words_data - - should_stop = self.dynamic_decoder.forward( - next_token_logits, decode_step, max_context_length, - self.max_attention_window_size, self.sink_token_length, - ite, batch_size, self.end_ids, self.embedding_bias_opt, - context_lengths, sequence_limit_lengths, - stop_words_list_ptrs, stop_words_lens, - max_stop_words_len, bad_words_list_ptrs, bad_words_lens, - max_bad_words_len, this_src_cache_indirection, - self.output_ids, self.new_tokens, self.finished, - self.finished, self.sequence_length_buffer, - self.cum_log_probs, self.log_probs, - self.log_probs_tiled, self.parent_ids, - this_tgt_cache_indirection, - self.beam_hyps_output_ids_cba, - self.beam_hyps_seq_len_cba, - self.beam_hyps_cum_log_probs_cba, - self.beam_hyps_normed_scores_cba, - self.beam_hyps_log_probs_cba, - self.beam_hyps_min_normed_scores, - self.beam_hyps_num_beams, self.beam_hyps_is_done, - scfg.use_beam_hyps) - - if not self.use_gpt_attention_plugin: - self.reorder_kv_cache_for_beam_search( - batch_size, beam_width, max_context_length, step) - - if stopping_criteria is not None and not should_stop.item(): - final_output_ids = self.finalize_decoder( - context_lengths, - batch_size, - beam_width, - scfg, - in_progress=True) - # keep the shape as same as huggingface stopping_criteria - final_output_ids_ = final_output_ids.reshape( - -1, final_output_ids.size(-1)) - should_stop[0] = stopping_criteria( - step, final_output_ids_, logits) - - if self.runtime._is_profiling(): - if not context.report_to_profiler(): - logger.warning("Runtime report to profiler failed.") - self.runtime._insert_step_to_profiler(step) - - if self.mapping.has_pp(): - should_stop = self.pp_communicate_new_tokens( - should_stop, this_tgt_cache_indirection, - self.sequence_length_buffer) - - if self.paged_kv_cache and self.has_attn_layers: - if (step >= self.max_new_tokens - 1) or (should_stop is not None - and should_stop.item()): - # Free all blocks in all sequences. - # With in-flight batching and while loop we'll free some sequences, when they are done - self.pools_kv_cache_manager.step([True] * batch_size) - if self.cross_attention: - self.cross_pools_kv_cache_manager.step([True] * batch_size) - - if self.debug_mode: - self.dump_debug_buffers(step) - - if next_step_tensors is not None: - self.debug_buffer = { - name: tensor.to_torch() - for name, tensor in next_step_tensors.items() - } - - return should_stop, next_step_tensors, tasks, context_lengths, host_context_lengths, attention_mask, context_logits, generation_logits, encoder_input_lengths - - def dump_debug_buffers(self, step: int) -> None: - if self.debug_tensors_to_save is not None: - # restricted written tensors according to filter - debug_tensor_names = copy.deepcopy(list(self.debug_buffer.keys())) - for k in debug_tensor_names: - if all([kk not in k for kk in self.debug_tensors_to_save]): - self.debug_buffer.pop(k) - - debug_dir = Path( - f"tllm_debug/PP_{self.mapping.pp_rank}/TP_{self.mapping.tp_rank}/CP_{self.mapping.cp_rank}" - ) - debug_dir.mkdir(parents=True, exist_ok=True) - - for name, t in self.debug_buffer.items(): - # convert tensor name to valid file name - print("Saving: ", name) - fname = name.replace("/", ".") - t = torch_to_numpy(t.float()) - np.save(debug_dir / f"{fname}-step{step}.npy", t) - - txt_format = "%d" if t.dtype in [np.int32, np.int8] else '%.18e' - np.savetxt( - debug_dir / f"{fname}-step{step}.txt", - t.reshape(-1, t.shape[-1]), # savetxt accepts 2 dims only - fmt=txt_format) - - def decode_regular(self, - *, - batch_size: int, - scfg: SamplingConfig, - sequence_lengths: torch.Tensor, - context_lengths: torch.Tensor, - host_context_lengths, - max_context_length: int, - beam_width: int, - cache_indirections: list, - input_ids: torch.Tensor, - hidden_states: torch.Tensor, - prompt_embedding_table: torch.Tensor, - tasks: torch.Tensor, - prompt_vocab_size: torch.Tensor, - ite: int, - sequence_limit_lengths: torch.Tensor, - stop_words_data, - bad_words_data, - output_sequence_lengths: bool = False, - output_generation_logits: bool = False, - return_dict: bool = False, - encoder_output: torch.Tensor = None, - encoder_input_lengths: torch.Tensor = None, - stopping_criteria: StoppingCriteria = None, - logits_processor: LogitsProcessor = None, - cross_attention_mask: List[torch.Tensor] = None, - **kwargs): - kv_cache_block_offsets = None - host_kv_cache_block_offsets = None - cross_kv_cache_block_offsets = None - host_cross_kv_cache_block_offsets = None - attention_mask = None - outputs_context_logits = None - outputs_generation_logits = [] - - def get_outputs_dict(output_ids, num_steps=self.max_new_tokens): - outputs = {} - outputs['output_ids'] = output_ids - if scfg.output_log_probs: - outputs['log_probs'] = self.log_probs - if scfg.output_cum_log_probs: - outputs['cum_log_probs'] = self.cum_log_probs - if output_sequence_lengths: - outputs[ - 'sequence_lengths'] = self.sequence_length_buffer.reshape( - [batch_size, beam_width]) - if self.gather_context_logits: - outputs['context_logits'] = outputs_context_logits - if self.gather_generation_logits or output_generation_logits: - outputs['generation_logits'] = outputs_generation_logits - if self.is_medusa_mode or self.is_redrafter_mode: - outputs['steps_to_finish'] = num_steps - if self.is_medusa_mode: - outputs['medusa_output_tokens'] = self.medusa_output_tokens - outputs['accept_lengths'] = self.accept_lengths - if self.medusa_temperature != 0.0: - outputs['medusa_output_logits'] = self.medusa_output_logits - return outputs - - benchmark_profiler = kwargs.get('benchmark_profiler', None) - generation_phase_step_count = 0 - - if benchmark_profiler is not None and benchmark_profiler.is_recording_perf_profile: - self.runtime._set_profiler() - - def profile_fn(benchmark_profiler_obj, step_count): - if benchmark_profiler_obj is not None: - benchmark_profiler_obj.record_cuda_event('last_token') - benchmark_profiler_obj.record_elapsed_time( - 'first_token', 'last_token', 'generation_time') - benchmark_profiler_obj.add_aux_info('generation_step_count', - step_count) - - # prepare cross attention mask. - cross_attention_mask_for_context = None - cross_attention_mask_for_gen = None - if cross_attention_mask is not None: - cross_attention_mask_for_context, cross_attention_mask_for_gen = self._prepare_cross_attention_mask( - batch_size, context_lengths, cross_attention_mask) - if self.use_gpt_attention_plugin: - # When we use plugin, the data type of cross_attention_mask is bool. - # When we don't use plugin, the data type of cross_attention_mask is int32 - cross_attention_mask_for_context = cross_attention_mask_for_context.to( - torch.bool) - cross_attention_mask_for_gen = cross_attention_mask_for_gen.to( - torch.bool) - - next_step_tensors = None - for step in range(0, self.max_new_tokens): - - should_stop, next_step_tensors, tasks, context_lengths, host_context_lengths, attention_mask, context_logits, generation_logits, encoder_input_lengths = self.handle_per_step( - cache_indirections=cache_indirections, - step=step, - batch_size=batch_size, - max_context_length=max_context_length, - beam_width=beam_width, - input_ids=input_ids, - hidden_states=hidden_states, - scfg=scfg, - kv_cache_block_offsets=kv_cache_block_offsets, - host_kv_cache_block_offsets=host_kv_cache_block_offsets, - cross_kv_cache_block_offsets=cross_kv_cache_block_offsets, - host_cross_kv_cache_block_offsets= - host_cross_kv_cache_block_offsets, - prompt_embedding_table=prompt_embedding_table, - tasks=tasks, - context_lengths=context_lengths, - host_context_lengths=host_context_lengths, - attention_mask=attention_mask, - cross_attention_mask_for_context= - cross_attention_mask_for_context, - cross_attention_mask_for_gen=cross_attention_mask_for_gen, - prompt_vocab_size=prompt_vocab_size, - ite=ite, - sequence_limit_lengths=sequence_limit_lengths, - sequence_lengths=sequence_lengths, - next_step_tensors=next_step_tensors, - stop_words_data=stop_words_data, - bad_words_data=bad_words_data, - encoder_output=encoder_output, - encoder_input_lengths=encoder_input_lengths, - stopping_criteria=stopping_criteria, - logits_processor=logits_processor, - output_generation_logits=output_generation_logits, - **kwargs, - ) - if step == 0: - if benchmark_profiler is not None: - benchmark_profiler.record_cuda_event('first_token') - else: - generation_phase_step_count = generation_phase_step_count + 1 - - if self.mapping.is_last_pp_rank(): - if step == 0 and self.gather_context_logits: - outputs_context_logits = context_logits - if self.gather_generation_logits or output_generation_logits: - outputs_generation_logits.append(generation_logits) - - if should_stop is not None and should_stop.item(): - profile_fn(benchmark_profiler, generation_phase_step_count) - if self.is_medusa_mode or self.is_redrafter_mode: - # just hack away for now - final_output_ids = self.output_ids.clone().unsqueeze(1) - final_output_ids = final_output_ids[:, :, :self. - max_seq_length - - self.max_draft_tokens] - else: - final_output_ids = self.finalize_decoder( - context_lengths, batch_size, beam_width, scfg) - - if self.mapping.is_first_pp_rank(): - if return_dict: - return get_outputs_dict(final_output_ids, step + 1) - else: - return final_output_ids - elif self.mapping.is_last_pp_rank(): - outputs = {} - if self.gather_context_logits: - outputs['context_logits'] = outputs_context_logits - if self.gather_generation_logits or output_generation_logits: - outputs['generation_logits'] = outputs_generation_logits - return outputs - else: - return None - - assert not self.is_medusa_mode and not self.is_redrafter_mode, "the custom decoder doesn't support medusa/redrafter." - - profile_fn(benchmark_profiler, generation_phase_step_count) - - final_output_ids = self.finalize_decoder(context_lengths, batch_size, - beam_width, scfg) - if self.mapping.is_first_pp_rank(): - if return_dict: - return get_outputs_dict(final_output_ids) - else: - return final_output_ids - elif self.mapping.is_last_pp_rank(): - outputs = {} - if self.gather_context_logits: - outputs['context_logits'] = outputs_context_logits - if self.gather_generation_logits or output_generation_logits: - outputs['generation_logits'] = outputs_generation_logits - return outputs - else: - return None - - def decode_stream(self, - *, - batch_size: int, - scfg: SamplingConfig, - sequence_lengths: torch.Tensor, - context_lengths: torch.Tensor, - host_context_lengths, - max_context_length: int, - beam_width: int, - cache_indirections: list, - input_ids: torch.Tensor, - hidden_states: torch.Tensor, - prompt_embedding_table: torch.Tensor, - tasks: torch.Tensor, - prompt_vocab_size: torch.Tensor, - ite: int, - sequence_limit_lengths: torch.Tensor, - stop_words_data, - bad_words_data, - output_sequence_lengths: bool = False, - output_generation_logits: bool = False, - return_dict: bool = False, - encoder_output: torch.Tensor = None, - encoder_input_lengths: torch.Tensor = None, - stopping_criteria: StoppingCriteria = None, - logits_processor: LogitsProcessor = None, - cross_attention_mask: List[torch.Tensor] = None, - **kwargs): - kv_cache_block_offsets = None - host_kv_cache_block_offsets = None - cross_kv_cache_block_offsets = None - host_cross_kv_cache_block_offsets = None - attention_mask = None - outputs_context_logits = None - - def get_outputs_dict(output_ids): - outputs = {} - outputs['output_ids'] = output_ids - if output_sequence_lengths: - outputs[ - 'sequence_lengths'] = self.sequence_length_buffer.reshape( - [batch_size, beam_width]) - if self.gather_context_logits: - outputs['context_logits'] = outputs_context_logits - return outputs - - # prepare cross attention mask. - cross_attention_mask_for_context = None - cross_attention_mask_for_gen = None - if cross_attention_mask is not None: - cross_attention_mask_for_context, cross_attention_mask_for_gen = self._prepare_cross_attention_mask( - batch_size, context_lengths, cross_attention_mask) - - next_step_tensors = None - for step in range(0, self.max_new_tokens): - - should_stop, next_step_tensors, tasks, context_lengths, host_context_lengths, attention_mask, context_logits, generation_logits, encoder_input_lengths = self.handle_per_step( - cache_indirections=cache_indirections, - step=step, - batch_size=batch_size, - max_context_length=max_context_length, - beam_width=beam_width, - input_ids=input_ids, - hidden_states=hidden_states, - scfg=scfg, - kv_cache_block_offsets=kv_cache_block_offsets, - host_kv_cache_block_offsets=host_kv_cache_block_offsets, - cross_kv_cache_block_offsets=cross_kv_cache_block_offsets, - host_cross_kv_cache_block_offsets= - host_cross_kv_cache_block_offsets, - prompt_embedding_table=prompt_embedding_table, - tasks=tasks, - context_lengths=context_lengths, - host_context_lengths=host_context_lengths, - attention_mask=attention_mask, - cross_attention_mask_for_context= - cross_attention_mask_for_context, - cross_attention_mask_for_gen=cross_attention_mask_for_gen, - prompt_vocab_size=prompt_vocab_size, - ite=ite, - sequence_limit_lengths=sequence_limit_lengths, - sequence_lengths=sequence_lengths, - next_step_tensors=next_step_tensors, - stop_words_data=stop_words_data, - bad_words_data=bad_words_data, - encoder_output=encoder_output, - encoder_input_lengths=encoder_input_lengths, - stopping_criteria=stopping_criteria, - logits_processor=logits_processor, - output_generation_logits=output_generation_logits, - ) - if step == 0: - outputs_context_logits = context_logits - if should_stop is not None: - - final_output_ids = self.finalize_decoder(context_lengths, - batch_size, - beam_width, - scfg, - in_progress=True) - - if self.mapping.is_first_pp_rank(): - if return_dict: - yield get_outputs_dict(final_output_ids) - else: - yield final_output_ids - else: - yield None - - if should_stop.item(): - return - - final_output_ids = self.finalize_decoder(context_lengths, batch_size, - beam_width, scfg) - if self.mapping.is_first_pp_rank(): - if return_dict: - yield get_outputs_dict(final_output_ids) - else: - yield final_output_ids - else: - yield None - - def decode_batch(self, - input_ids: Sequence[torch.Tensor], - sampling_config: SamplingConfig, - streaming: bool = False, - **kwargs): - input_ids, context_lengths = _prepare_input_ids(input_ids) - return self.decode(input_ids, - context_lengths, - sampling_config, - streaming=streaming, - **kwargs) - - # As dynamic_decoder uses torch's current stream, we must ensure it runs on the same stream that - # dynamic_decoder was set up with - @cuda_stream_guard - def decode(self, - input_ids: torch.Tensor, - context_lengths: torch.Tensor, - sampling_config: SamplingConfig, - prompt_embedding_table: torch.Tensor = None, - tasks: torch.Tensor = None, - prompt_vocab_size: torch.Tensor = None, - stop_words_list=None, - bad_words_list=None, - streaming: bool = False, - output_sequence_lengths: bool = False, - output_generation_logits: bool = False, - return_dict: bool = False, - encoder_output: torch.Tensor = None, - encoder_input_lengths: torch.Tensor = None, - stopping_criteria: StoppingCriteria = None, - logits_processor: LogitsProcessor = None, - cross_attention_mask: List[torch.Tensor] = None, - **kwargs): - scfg = sampling_config - batch_size = context_lengths.size(0) - beam_width = scfg.num_beams - max_context_length = torch.max(context_lengths).item() - host_context_lengths = context_lengths.cpu() - assert batch_size == self.batch_size, \ - "Given batch size is different from the one used in setup()," \ - "rerun the setup function with the new batch size to avoid buffer overflow." - assert max_context_length <= self.max_context_length, \ - "Given input length is large then the one used in setup()," \ - "rerun the setup function with the new max_context_length to avoid buffer overflow." - assert beam_width == self.beam_width, \ - "Given beam width is different from the one used in setup()," \ - "rerun the setup function with the new beam width to avoid buffer overflow." - assert self.sink_token_length <= torch.min(context_lengths).item(), \ - "Given sink token length is larger than shortest context length," \ - "rerun the setup function with a smaller sink token length." - ite = 0 # index of local batches, will always be 0 if pp_size = 1 - - if self.remove_input_padding and input_ids.dim() == 2: - assert input_ids.shape[ - 0] == 1, "Packed 2D input must have shape [1, ]" - input_ids = input_ids.squeeze(0) - - self.__setup_decoder(input_ids, scfg, host_context_lengths) - if not self.buffer_allocated: - raise RuntimeError('Buffer not allocated, please call setup first!') - - sequence_limit_lengths = torch.full((batch_size, 1), - self.max_seq_length, - dtype=torch.int32, - device=self.device) - - # Sequence_lengths for the dynamic decoder still has the input paddings. - sequence_lengths = torch.full((batch_size * beam_width, 1), - max_context_length, - dtype=torch.int32, - device=self.device) - - cache_indirections = [ - torch.full(( - batch_size, - beam_width, - self.max_attention_window_size, - ), - 0, - dtype=torch.int32, - device=self.device), - torch.full(( - batch_size, - beam_width, - self.max_attention_window_size, - ), - 0, - dtype=torch.int32, - device=self.device) - ] # ping-pong buffers - - hidden_states = None - if self.mapping.has_pp(): - max_num_tokens = max(batch_size * beam_width, - batch_size * self.max_seq_length) - hidden_size = self.hidden_size * self.mapping.tp_size - hidden_states = torch.zeros((1, max_num_tokens, hidden_size)) - - # Init KV cache block manager - if self.paged_kv_cache and self.has_attn_layers: - num_blocks, max_blocks_per_seq = self._get_num_paged_blocks( - self.max_attention_window_size, self.sink_token_length) - - self.buffer[ - f'host_kv_cache_pool_pointers'] = self._memory_pool_allocator.get_kv_cache_pool_pointers( - ) - self.buffer[ - f'host_kv_cache_pool_mapping'] = self._memory_pool_allocator.pool_mapping - - self.pools_kv_cache_manager = PoolsKVCacheManager( - self._memory_pool_allocator.pools_metadata, - max_blocks_per_seq, - num_blocks, - self.tokens_per_block, - self.head_size, - max_attention_window_size=self.max_attention_window_size, - beam_width=beam_width, - sink_token_len=self.sink_token_length) - - if self.cross_attention: - cross_num_blocks, max_cross_blocks_per_seq = self._get_num_paged_blocks( - self.encoder_max_input_length, sink_token_length=0) - self.buffer[ - f'host_cross_kv_cache_pool_pointers'] = self._cross_memory_pool_allocator.get_kv_cache_pool_pointers( - ) - self.buffer[ - f'host_cross_kv_cache_pool_mapping'] = self._cross_memory_pool_allocator.pool_mapping - - self.cross_pools_kv_cache_manager = PoolsKVCacheManager( - self._cross_memory_pool_allocator.pools_metadata, - max_cross_blocks_per_seq, - cross_num_blocks, - self.tokens_per_block, - self.head_size, - max_attention_window_size=self.encoder_max_input_length, - beam_width=beam_width, - sink_token_len=self.sink_token_length) - - # Add sequences to the manager - for bi in range(batch_size): - generation_sequence = GenerationSequence(seq_idx=bi, - batch_idx=bi) - self.pools_kv_cache_manager.add_sequence( - generation_sequence, max_context_length) - if self.cross_attention: - cross_generation_sequence = GenerationSequence(seq_idx=bi, - batch_idx=bi) - self.cross_pools_kv_cache_manager.add_sequence( - cross_generation_sequence, - self.encoder_max_input_length, - always_share_across_beam=True) - # cross attention paged kv cache should always share the context blocks across beams - # due to the fact that we are not adding new key/value cache to cross kv in generation - - if self.is_medusa_mode or self.is_redrafter_mode: - if self.quant_mode.has_kv_cache_quant(): - # Since torch does not support fp8 now, using int8 here. - kv_cache_type = torch.int8 - else: - kv_cache_type = self.dtype if self.paged_kv_cache else self._tensor_dtype( - f'present_key_value_{self.first_layer}') - self.history_max_seq_length = [max_context_length] - self.kv_cache_updater = KVCacheUpdater() - assert not self.cross_attention - assert self.use_gpt_attention_plugin - - if self.paged_kv_cache: - self.kv_cache_updater.init_paged_kv_cache( - self.num_layers, self.get_num_heads_kv(), self.head_size, - kv_cache_type, self.pools_kv_cache_manager, - self.buffer[f'host_kv_cache_pool_pointers']) - else: - past_key_value_list = [ - self.buffer[f'present_key_value_{i}'] - for i in range(self.first_layer, self.last_layer) - ] - self.kv_cache_updater.init_linear_kv_cache( - self.num_layers, self.get_num_heads_kv(), self.head_size, - kv_cache_type, past_key_value_list) - - stop_words_lens = None - stop_words_list_ptrs = None - max_stop_words_len = 0 - if stop_words_list is not None: - stop_words_list = torch.from_numpy(stop_words_list).contiguous().to( - 'cuda') - max_stop_words_len = stop_words_list.shape[2] - stop_words_lens = torch.full((batch_size, ), - max_stop_words_len, - dtype=torch.int32).to('cuda') - stop_words_list_ptrs = torch.zeros((batch_size), dtype=torch.int64) - for bi in range(batch_size): - stop_words_list_ptrs[bi] = stop_words_list.data_ptr( - ) + bi * 2 * max_stop_words_len * stop_words_list.element_size( - ) - stop_words_list_ptrs = stop_words_list_ptrs.to('cuda') - stop_words_data = (stop_words_list_ptrs, stop_words_lens, - max_stop_words_len) - - bad_words_lens = None - bad_words_list_ptrs = None - max_bad_words_len = 0 - if bad_words_list is not None: - bad_words_list = torch.from_numpy(bad_words_list).contiguous().to( - 'cuda') - max_bad_words_len = bad_words_list.shape[2] - bad_words_lens = torch.full((batch_size, ), - max_bad_words_len, - dtype=torch.int32).to('cuda') - bad_words_list_ptrs = torch.zeros((batch_size), dtype=torch.int64) - for bi in range(batch_size): - bad_words_list_ptrs[bi] = bad_words_list.data_ptr( - ) + bi * 2 * max_bad_words_len * bad_words_list.element_size() - bad_words_list_ptrs = bad_words_list_ptrs.to('cuda') - bad_words_data = (bad_words_list_ptrs, bad_words_lens, - max_bad_words_len) - - # start context phase - if streaming: - return self.decode_stream( - batch_size=batch_size, - scfg=scfg, - sequence_lengths=sequence_lengths, - context_lengths=context_lengths, - host_context_lengths=host_context_lengths, - max_context_length=max_context_length, - beam_width=beam_width, - cache_indirections=cache_indirections, - input_ids=input_ids, - hidden_states=hidden_states, - prompt_embedding_table=prompt_embedding_table, - tasks=tasks, - prompt_vocab_size=prompt_vocab_size, - ite=ite, - sequence_limit_lengths=sequence_limit_lengths, - output_generation_logits=output_generation_logits, - stop_words_data=stop_words_data, - bad_words_data=bad_words_data, - output_sequence_lengths=output_sequence_lengths, - return_dict=return_dict, - encoder_output=encoder_output, - encoder_input_lengths=encoder_input_lengths, - stopping_criteria=stopping_criteria, - logits_processor=logits_processor, - cross_attention_mask=cross_attention_mask, - **kwargs, - ) - else: - return self.decode_regular( - batch_size=batch_size, - scfg=scfg, - sequence_lengths=sequence_lengths, - context_lengths=context_lengths, - host_context_lengths=host_context_lengths, - max_context_length=max_context_length, - beam_width=beam_width, - cache_indirections=cache_indirections, - input_ids=input_ids, - hidden_states=hidden_states, - prompt_embedding_table=prompt_embedding_table, - tasks=tasks, - prompt_vocab_size=prompt_vocab_size, - ite=ite, - sequence_limit_lengths=sequence_limit_lengths, - stop_words_data=stop_words_data, - bad_words_data=bad_words_data, - output_sequence_lengths=output_sequence_lengths, - output_generation_logits=output_generation_logits, - return_dict=return_dict, - encoder_output=encoder_output, - encoder_input_lengths=encoder_input_lengths, - stopping_criteria=stopping_criteria, - logits_processor=logits_processor, - cross_attention_mask=cross_attention_mask, - **kwargs, - ) - - -class ChatGLMGenerationSession(GenerationSession): - - def __init__( - self, - model_config: ModelConfig, - engine_buffer, - mapping: Mapping, - debug_mode=False, - debug_tensors_to_save=None, - cuda_graph_mode=False, - stream: torch.cuda.Stream = None, - ): - - super().__init__( - model_config, - engine_buffer, - mapping, - debug_mode, - debug_tensors_to_save, - cuda_graph_mode, - stream, - ) - - self.mask_index_tensor = None - - def _prepare_context_inputs(self, batch_size, context_lengths, - use_gpt_attention_plugin, remove_input_padding, - **kwargs): - - max_context_length = kwargs.pop('max_context_length') - last_token_ids = context_lengths.detach().clone() - - if remove_input_padding: - input_lengths_acc = torch.cumsum(torch.cat( - [torch.IntTensor([0]).cuda(), context_lengths], dim=0), - dim=0) - position_ids = torch.zeros([2, input_lengths_acc[-1]], - dtype=torch.int32) - for i in range(batch_size): - position_ids[0, input_lengths_acc[i]:input_lengths_acc[ - i + 1]] = torch.arange(0, - context_lengths[i], - dtype=torch.int32) - position_ids[0, input_lengths_acc[i + 1] - - 1] = context_lengths[i] - 2 - position_ids[1, input_lengths_acc[i + 1] - 1] = 1 - position_ids = position_ids.int().cuda() - last_token_ids = torch.cumsum(last_token_ids, dim=0).int().cuda() - - # specialization for GLM series models - if kwargs["pad_id"] in [50256, 50259]: - if kwargs["pad_id"] == 50256: # glm_2b / glm_10b - mask_ids = [50260, 50264, 50263] - else: # glm_10b_chinese / glm_large_chinese - mask_ids = [50003, 50008, 50009] - - self.mask_index_tensor = \ - torch.zeros([batch_size], dtype=torch.int32) - position_ids = position_ids.cpu() - for i in range(batch_size): - length = context_lengths[i] - input_ids = kwargs["input_ids"][ - 0:context_lengths[i]] if i == 0 else kwargs[ - "input_ids"][sum(context_lengths[0:i] - ):sum(context_lengths[0:i]) + - length] - mask_index = [ - torch.where(input_ids == id)[0].int() for id in mask_ids - ] - tail_index = torch.Tensor([max_context_length]).int().cuda() - mask_index.append(tail_index) - mask_index = torch.cat(mask_index, dim=0).min() - self.mask_index_tensor[i] = int(mask_index) - position_ids[0][sum(context_lengths[0:i + 1]) - - 1] = int(mask_index) - position_ids = position_ids.cuda() - else: - position_ids = torch.zeros([batch_size, 2, max_context_length], - dtype=torch.int32) - position_ids[:, 0, :] = torch.arange(max_context_length) - - # specialization for GLM series models - if kwargs["pad_id"] in [50256, 50259]: - if kwargs["pad_id"] == 50256: # glm_2b / glm_10b - mask_ids = [50260, 50264, 50263] - else: # glm_10b_chinese / glm_large_chinese - mask_ids = [50003, 50008, 50009] - - self.mask_index_tensor = \ - torch.zeros([batch_size], dtype=torch.int32) - for i in range(batch_size): - length = context_lengths[i] - input_ids = kwargs["input_ids"][i] - mask_index = [ - torch.where(input_ids == id)[0].int() for id in mask_ids - ] - tail_index = torch.Tensor([max_context_length]).int().cuda() - mask_index.append(tail_index) - mask_index = torch.cat(mask_index, dim=0).min() - position_ids[i, 0, length - 1] = int(mask_index) - position_ids[i, 1, length - 1] = 1 - self.mask_index_tensor[i] = int(mask_index) - else: - for i in range(batch_size): - length = context_lengths[i] - position_ids[i, 0, length - 1] = length - 2 - position_ids[i, 1, length - 1] = 1 - - position_ids = position_ids.cuda() - - perf_knob_tensor_size = 16 - context_runtime_perf_knobs = torch.tensor([-1] * perf_knob_tensor_size, - dtype=torch.int64) - - inputs = { - 'position_ids': position_ids, - 'last_token_ids': last_token_ids, - 'host_runtime_perf_knobs': context_runtime_perf_knobs - } - if not use_gpt_attention_plugin: - attention_mask = torch.zeros((batch_size, 1)) - inputs['attention_mask'] = attention_mask - return inputs - - def _prepare_generation_inputs(self, batch_size, context_lengths, - use_gpt_attention_plugin, - remove_input_padding, **kwargs): - - step = kwargs.pop('step') - num_beams = kwargs.pop('num_beams') - last_token_ids = torch.ones_like(context_lengths) - - if remove_input_padding: - - def _tile_beam_width_chatglm(tensor: torch.Tensor, num_beams: int): - new_shape = np.array(tensor.shape) - new_shape[1] = new_shape[1] * num_beams - tile_size = np.ones(new_shape.shape, dtype=np.int32) - tile_size = np.insert(tile_size, 2, num_beams) - new_tensor = torch.unsqueeze(tensor, 2) - new_tensor = new_tensor.tile(tile_size.tolist()) - new_tensor = new_tensor.reshape(new_shape.tolist()) - return new_tensor - - position_ids = torch.zeros([2, batch_size], dtype=torch.int32) - for i in range(batch_size): - position_ids[0, i] = context_lengths[i * num_beams] - 2 - position_ids[1, i] = step + 2 - position_ids = _tile_beam_width_chatglm(position_ids, num_beams) - position_ids = position_ids.int().cuda() - last_token_ids = torch.cumsum(last_token_ids, dim=0).int().cuda() - - if self.mask_index_tensor is not None: # specialization for GLM series models - position_ids = position_ids.cpu() - for i in range(batch_size): - position_ids[0][i] = self.mask_index_tensor[i] - position_ids = position_ids.cuda() - else: - data = [] - if self.mask_index_tensor is not None: # specialization for GLM series models - for i in range(batch_size): - data.append([[self.mask_index_tensor[i]], [step + 2]]) - else: - for i in range(batch_size): - data.append([[context_lengths[i * num_beams] - 2], - [step + 2]]) - position_ids = torch.tensor(data, dtype=torch.int32, device='cuda') - position_ids = _tile_beam_width(position_ids, num_beams) - - perf_knob_tensor_size = 16 - generation_runtime_perf_knobs = torch.tensor([-1] * - perf_knob_tensor_size, - dtype=torch.int64) - - inputs = { - 'position_ids': position_ids, - 'last_token_ids': last_token_ids, - 'host_runtime_perf_knobs': generation_runtime_perf_knobs - } - if not use_gpt_attention_plugin: - attention_mask = torch.zeros((batch_size, 1)) - inputs['attention_mask'] = attention_mask - return inputs - - -class QWenForCausalLMGenerationSession(GenerationSession): - - def __init__( - self, - model_config: ModelConfig, - engine_buffer, - mapping: Mapping, - debug_mode=False, - debug_tensors_to_save=None, - cuda_graph_mode=False, - stream: torch.cuda.Stream = None, - global_max_input_length: int = 2048, - global_max_output_length: int = 4096, - ): - super().__init__(model_config, - engine_buffer, - mapping, - debug_mode, - debug_tensors_to_save=debug_tensors_to_save, - cuda_graph_mode=cuda_graph_mode, - stream=stream) - self.global_max_input_length = global_max_input_length - self.global_max_output_length = global_max_output_length - - def generate( - self, - input_ids: torch.Tensor, - input_lengths: torch.Tensor, - sampling_config: SamplingConfig, - max_new_tokens: int, - runtime_rank: int = 0, - ): - max_input_length = torch.max(input_lengths).item() - max_new_tokens = min(max_new_tokens, - self.global_max_output_length - max_input_length) - # setup batch_size, max_input_length, max_output_len - self.setup(batch_size=input_lengths.size(0), - max_context_length=max_input_length, - max_new_tokens=max_new_tokens) - output_ids = self.decode(input_ids, input_lengths, sampling_config) - with torch.no_grad(): - torch.cuda.synchronize() - if runtime_rank == 0: - outputs = output_ids[:, 0, :] - return outputs diff --git a/tensorrt_llm/runtime/kv_cache_manager.py b/tensorrt_llm/runtime/kv_cache_manager.py deleted file mode 100644 index c2b6c3f9b1a1..000000000000 --- a/tensorrt_llm/runtime/kv_cache_manager.py +++ /dev/null @@ -1,473 +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. -from collections import defaultdict -from typing import List - -import torch - - -class Block(object): - - def __init__(self, block_idx: int): - self.idx = block_idx - self.ref_count = 0 - - def add_link(self): - self.ref_count += 1 - - def remove_link(self): - self.ref_count -= 1 - - def has_link(self) -> bool: - return self.ref_count > 0 - - def is_shared(self) -> bool: - return self.ref_count > 1 - - -class GenerationSequence(object): - - def __init__(self, seq_idx, batch_idx): - self.seq_idx = seq_idx - self.batch_idx = batch_idx - - def get_batch_idx(self) -> int: - """ - Returns idx of sequence in batch - """ - return self.batch_idx - - def get_seq_idx(self) -> int: - """ - Returns sequence idx - """ - return self.seq_idx - - def __eq__(self, another): - return hasattr(another, 'seq_idx') and self.seq_idx == another.seq_idx and \ - hasattr(another, 'batch_idx') and self.batch_idx == another.batch_idx - - def __hash__(self): - return self.seq_idx - - -class BlocksManager(object): - _sizeof = { - torch.float32: 4, - torch.float16: 2, - torch.bfloat16: 2, - torch.int8: 1 - } - - def __init__(self, - *, - num_layers: int, - num_blocks: int, - block_size: int, - max_blocks_per_seq: int = 128, - beam_width: int = 1): - """ - If layers are homogeneous then the expected block pool shape is: [num_blocks, num_layers, 2, block_size] - Otherwise, the expected block pool shape is: [num_blocks, 2, block_size] - """ - - self.max_blocks_per_seq = max_blocks_per_seq - - self.num_layers = num_layers - self.num_blocks = num_blocks - self.block_size = block_size - self.beam_width = beam_width - - self.free_blocks = [] - for bi in range(num_blocks): - self.free_blocks.append(Block(bi)) - - beam_width = self.beam_width - # Here use beam_width instead of self.beam_width to remove cyclic reference between self and - # self.allocated_blocks by preventing capture self, which may cause memory leak. - self.allocated_blocks = defaultdict( - lambda: [[] for _ in range(beam_width)]) - - def has_free_block(self) -> bool: - """ - Returns True if we have at least 1 free block - """ - return len(self.free_blocks) > 0 - - def allocate(self, - owner: GenerationSequence, - share_across_beam: bool = False): - """ - Add block to owner and increase ref count - """ - # Add blocks for whole beam width - block = None - for bi in range(self.beam_width): - if not self.has_free_block(): - raise RuntimeError("Can't allocate new block for KV cache") - - # Use the same block for all seqs in beam if share_across_beam - if block is None or share_across_beam == False: - block = self.free_blocks.pop(0) - # Add one reference to the block - block.add_link() - self.allocated_blocks[owner][bi].append(block) - - def replace_shared_block(self, owner: GenerationSequence, block_idx: int): - """ - Replace the shared block. - Free the shared block, and allocate blocks with share_across_beam=False - """ - if not self.allocated_blocks[owner][0][block_idx].is_shared(): - return - - # Free shared block - for bi in range(self.beam_width): - block = self.allocated_blocks[owner][bi][block_idx] - block.remove_link() - if not block.has_link(): - self.free_blocks.append(block) - - # Allocate new block - for bi in range(self.beam_width): - if not self.has_free_block(): - raise RuntimeError("Can't allocate new block for KV cache") - block = self.free_blocks.pop(0) - block.add_link() - self.allocated_blocks[owner][bi][block_idx] = block - return - - def free(self, owner: GenerationSequence): - """ - Unlink all blocks of given owner. - Moves blocks with ref_count == 0 to free. - Removes owner from allocated blocks. - """ - for bi in range(self.beam_width): - for block in self.allocated_blocks[owner][bi]: - # Move block to free if no one refers to it - block.remove_link() - - # Move block to free if no one refers to it - if not block.has_link(): - self.free_blocks.append(block) - # Remove owner from allocated blocks - self.allocated_blocks.pop(owner) - - def get_number_blocks(self, owner: GenerationSequence) -> int: - """ - Returns number of blocks allocated to the sequence owner - """ - return len(self.allocated_blocks[owner][0]) - - def get_k_or_v_block_offset(self, block_idx, field_idx): - """ - Get offset in memory pool to K or V block. field_idx should be 0 (K) or 1 (V). - """ - return block_idx * self.num_layers * 2 + field_idx - - def get_offset_array(self, beam_width: int) -> torch.Tensor: - """ - Returns array of [batch size, beam_width, 2, max_blocks_per_seq] of offsets - to the allocated blocks in memory pool - """ - assert (beam_width <= self.beam_width) - - def create_nested_list(dims): - """Recursive function to generate nested list.""" - if len(dims) == 1: - return [0 for _ in range(dims[0])] - return [create_nested_list(dims[1:]) for _ in range(dims[0])] - - offset_array = create_nested_list( - (len(self.allocated_blocks), beam_width, 2, - self.max_blocks_per_seq)) - - k_idx = 0 - v_idx = 1 - for owner, beams_blocks in self.allocated_blocks.items(): - for bi in range(beam_width): - for block_linear_idx, block in enumerate(beams_blocks[bi]): - for x_idx in [k_idx, v_idx]: - offset_array[owner.get_batch_idx()][bi][x_idx][ - block_linear_idx] = self.get_k_or_v_block_offset( - block.idx, x_idx) - - self.offset_array = torch.tensor(offset_array, dtype=torch.int32) - return self.offset_array - - def get_continuous_caches(self, memory_pool: torch.Tensor) -> torch.Tensor: - """ - Returns continuous KV caches. - Used only for debug purposes. - """ - assert self.beam_width == 1 - - pool = memory_pool.flatten() - continuous_kv_cache = torch.zeros(len(self.allocated_blocks), - 2, - self.max_blocks_per_seq * - self.block_size, - dtype=pool.dtype, - device="cuda") - k_idx = 0 - v_idx = 1 - for owner, beam_blocks in self.allocated_blocks.items(): - for bi in range(self.beam_width): - for block_linear_idx, block in enumerate(beam_blocks[bi]): - # The batch index. - batch_idx = owner.get_batch_idx() - # The first index in the sequence. - block_offset = block_linear_idx * self.block_size - - for x_idx in [k_idx, v_idx]: - x_start = self.get_k_or_v_block_offset( - block.idx, x_idx) * self.block_size - - continuous_kv_cache[batch_idx][ - x_idx][block_offset:block_offset + - self.block_size] = pool[x_start:x_start + - self.block_size] - - return continuous_kv_cache - - -class KVCacheManager(object): - - def __init__(self, - *, - num_layers: int, - num_blocks: int, - block_size: int, - tokens_per_block: int, - max_blocks_per_seq: int, - max_attention_window_size: int, - sink_token_len: int, - beam_width: int = 1, - use_one_more_block: bool = False): - - self.blocks_manager = BlocksManager( - num_layers=num_layers, - num_blocks=num_blocks, - block_size=block_size, - max_blocks_per_seq=max_blocks_per_seq, - beam_width=beam_width) - - self.tokens_per_block = tokens_per_block - self.max_attention_window_size = max_attention_window_size - self.sink_token_len = sink_token_len - self.beam_width = beam_width - - # The sink tokens are not stored into the same block with other tokens. - # Need to add the bubble after the sink tokens. - if sink_token_len % tokens_per_block == 0: - self.bubble_len = 0 - else: - self.bubble_len = tokens_per_block - sink_token_len % tokens_per_block - - # Token num in the sink blocks - self.sink_block_token_num = self.sink_token_len + self.bubble_len - - # Max token num in the cache - self.max_token_num = self.max_attention_window_size + self.bubble_len - if use_one_more_block: - self.max_token_num += self.tokens_per_block - - self.lens = [] - self.sequences = [] - - def step(self, finished: List[bool]): - """ - Iterate to the next generation step. - Add new blocks where needed and clear finished sequences. - """ - for seq in self.sequences: - batch_idx = seq.get_batch_idx() - # Enable cyclic kv cache when it exceeds the max_token_num - cyclic_token_num = self.max_token_num - self.sink_block_token_num - next_token_idx_in_cache = self.sink_block_token_num + \ - (self.lens[batch_idx] - self.sink_block_token_num) % cyclic_token_num - if not finished[batch_idx] and ( - next_token_idx_in_cache % self.tokens_per_block == 0 or - (next_token_idx_in_cache - self.sink_block_token_num) % - cyclic_token_num == 0): - if self.lens[batch_idx] < self.max_token_num: - self.blocks_manager.allocate(seq) - elif self.beam_width > 1: - # Get next block index - next_block_idx = next_token_idx_in_cache // self.tokens_per_block - # Replace the shared block with the unshared ones - self.blocks_manager.replace_shared_block( - seq, next_block_idx) - - self.lens[batch_idx] += 1 - - # Remove finished sequences - for fi in range(len(finished)): - if finished[fi]: - self.blocks_manager.free(self.sequences[fi]) - self.lens = [l for l, f in zip(self.lens, finished) if not f] - - # Remap sequence ids - new_sequences = [] - batch_idx = 0 - for seq, finish in zip(self.sequences, finished): - if not finish: - seq.batch_idx = batch_idx - new_sequences.append(seq) - batch_idx += 1 - self.sequences = new_sequences - - def add_sequence(self, - sequence: GenerationSequence, - context_len: int, - always_share_across_beam: bool = False): - """ - Add sequence to the manager and allocate minimum amount of blocks for context - """ - seq_len = context_len + self.bubble_len - self.lens.append(seq_len) - self.sequences.append(sequence) - - # Enable cyclic kv cache when inputLength exceeds maxAttentionWindow. - # Note that currently cyclic kv cache doesn't work with shared kv cache of different beams. - enable_cyclic_kv_cache = seq_len >= self.max_token_num - - # Get the final token index in kv cache - final_token_kv_index = self.sink_block_token_num + ( - (seq_len - 1 - self.sink_block_token_num) % - (self.max_token_num - self.sink_block_token_num)) - - # Get block index that with shareAmongBeams=False. - unshared_block_idx = -1 - if (not enable_cyclic_kv_cache or self.beam_width > 1 - or final_token_kv_index % self.tokens_per_block > 0): - unshared_block_idx = final_token_kv_index // self.tokens_per_block + 1 if ( - final_token_kv_index + 1 - ) % self.tokens_per_block == 0 else final_token_kv_index // self.tokens_per_block - - # Get context block num. - # Allocate one more block if there are tokens that can't be shared across beams. - seq_len = min(seq_len, self.max_token_num) - context_blocks = seq_len // self.tokens_per_block - if seq_len % self.tokens_per_block > 0: - context_blocks += 1 - - # Allocate blocks - for i in range(context_blocks): - self.blocks_manager.allocate( - sequence, - share_across_beam=True if always_share_across_beam else - (i != unshared_block_idx)) - - def get_block_offsets(self, beam_width: int) -> torch.Tensor: - """ - Returns array of offsets into memory pools - """ - return self.blocks_manager.get_offset_array(beam_width) - - -class KVCacheUpdater: - - def __init__(self): - self.use_paged_kv_cache = None - self.num_layers = None - self.num_kv_heads = None - self.head_dim = None - self.elt_size = None - self.past_key_value_list = None - self.max_kv_cache_length = None - self.kv_cache_manager = None - self.host_kv_cache_pool_pointers = None - - def init_linear_kv_cache(self, num_layers, num_kv_heads, head_dim, - kv_cache_type, past_key_value_list): - self.use_paged_kv_cache = False - self.num_layers = num_layers - self.num_kv_heads = num_kv_heads - self.head_dim = head_dim - self.past_key_value_list = past_key_value_list - self.elt_size = torch.zeros(1, dtype=kv_cache_type).element_size() - self.max_kv_cache_length = past_key_value_list[0].shape[3] - - def init_paged_kv_cache(self, num_layers, num_kv_heads, head_dim, - kv_cache_type, kv_cache_manager, - host_kv_cache_pool_pointers): - self.use_paged_kv_cache = True - self.num_layers = num_layers - self.num_kv_heads = num_kv_heads - self.head_dim = head_dim - self.kv_cache_manager = kv_cache_manager - self.host_kv_cache_pool_pointers = host_kv_cache_pool_pointers - self.elt_size = torch.zeros(1, dtype=kv_cache_type).element_size() - - def update(self, accepted_draft_token_offsets, - packed_accepted_draft_tokens_indices, sequence_length_buffer, - rewind_tokens): - assert isinstance(rewind_tokens, torch.Tensor) or isinstance( - rewind_tokens, int) - rewind_tokens_tensor = rewind_tokens if isinstance( - rewind_tokens, torch.Tensor) else None - rewind_tokens_count = rewind_tokens if isinstance(rewind_tokens, - int) else 0 - assert self.use_paged_kv_cache is not None - if self.use_paged_kv_cache: - if self.kv_cache_manager.has_single_pool(): - kv_cache_manager = self.kv_cache_manager.get_single_kv_cache_manager( - ) - else: - raise RuntimeError( - "Currently, using KVCacheUpdater with more then single memory pool is not supported" - ) - - host_kv_cache_block_offsets = kv_cache_manager.get_block_offsets(1) - kv_cache_block_offsets = host_kv_cache_block_offsets.to('cuda') - torch.ops.tensorrt_llm.update_kv_cache_draft_token_location( - accepted_draft_token_offsets, - packed_accepted_draft_tokens_indices, - sequence_length_buffer, - True, - self.num_layers, - self.num_kv_heads, - self.head_dim * self.elt_size, - rewind_tokens_count, - kv_cache_manager.max_attention_window_size, - rewind_tokens_tensor, - None, - self.host_kv_cache_pool_pointers, - kv_cache_block_offsets, - kv_cache_manager.blocks_manager.max_blocks_per_seq, - kv_cache_manager.tokens_per_block, - None, - ) - else: - torch.ops.tensorrt_llm.update_kv_cache_draft_token_location( - accepted_draft_token_offsets, - packed_accepted_draft_tokens_indices, - sequence_length_buffer, - False, - self.num_layers, - self.num_kv_heads, - self.head_dim * self.elt_size, - rewind_tokens_count, - self.max_kv_cache_length, - rewind_tokens_tensor, - self.past_key_value_list, - None, - None, - None, - None, - None, - ) diff --git a/tensorrt_llm/runtime/medusa_utils.py b/tensorrt_llm/runtime/medusa_utils.py deleted file mode 100644 index ffd5126c7013..000000000000 --- a/tensorrt_llm/runtime/medusa_utils.py +++ /dev/null @@ -1,182 +0,0 @@ -import copy -from argparse import Namespace -from functools import cmp_to_key -from typing import List - -import numpy as np -import torch - -from tensorrt_llm.logger import logger - - -def path_sorter(a, b): - for i in range(min(len(a), len(b))): - if a[i] != b[i]: - return -1 if a[i] < b[i] else 1 - return 0 # shouldn't reach - - -path_sorting_key = cmp_to_key(path_sorter) - - -def expand_choices_if_needed(medusa_choices: List[List[int]]): - """ - Do a simple check to see if the given choices are path-only or vanilla. - """ - assert len(medusa_choices) > 0 - for c in medusa_choices: - if len(c) > 1: - try: - _ = medusa_choices.index( - [c[0]]) # find the first parent of current path - logger.debug( - "Detected vanilla-style of Medusa choices. No need to expand." - ) - return medusa_choices # if found, just return assuming it is already expanded - except ValueError: - logger.debug( - "Detected path-only style of Medusa choices. Expanding ...") - break - expanded_choices = set() - for c in medusa_choices: - cur = () - for n in c: - cur = (*cur, n) - expanded_choices.add(cur) - expanded_choices = [list(c) for c in expanded_choices] - return expanded_choices - - -def get_packed_mask(num_medusa_tokens, medusa_mask, max_medusa_tokens=None): - max_medusa_tokens = num_medusa_tokens if max_medusa_tokens is None else max_medusa_tokens - num_packed_masks = (max_medusa_tokens + 1 + 32 - 1) // 32 - medusa_packed_mask = torch.zeros((num_medusa_tokens + 1, num_packed_masks), - dtype=torch.int32) - for token_idx in range(num_medusa_tokens + 1): - if token_idx == 0: - medusa_packed_mask[0, 0] = 1 - else: - mask_list = medusa_mask[token_idx - 1, :].tolist() - # insert 1 as there is one extra new token from the original lm head. - mask_list.insert(0, True) - # convert binary bits into 4 int32_t - mask_str_list = [str(int(val)) for val in mask_list] - mask_str_list.reverse() - - for mask_idx in range(num_packed_masks): - if mask_idx * 32 >= len(mask_str_list): - break - mask_32bits_str = ''.join(mask_str_list[-(mask_idx + 1) * 32: - (-mask_idx * 32 - 1)] + - [mask_str_list[(-mask_idx * 32 - 1)]]) - valid_num_bits = len(mask_32bits_str) - first_bit1 = mask_32bits_str[0] == '1' - mask_31bits_str = mask_32bits_str[1:] - mask_31bits = 0 if mask_31bits_str == "" else int( - mask_31bits_str, 2) - if valid_num_bits == 32: - mask_32bits = mask_31bits - first_bit1 * (2**( - valid_num_bits - 1)) - else: - mask_32bits = mask_31bits + first_bit1 * (2**( - valid_num_bits - 1)) - medusa_packed_mask[token_idx, mask_idx] = mask_32bits - return medusa_packed_mask - - -def choices_2_paths(num_medusa_heads, choices): - paths = {} - all_paths = {} - level_counts = [0] * num_medusa_heads - choices.sort(key=len, reverse=True) - for c in choices: - k = ":".join([str(ci) for ci in c]) - if k not in all_paths: - paths[k] = c - for i in range(len(c)): - k = ":".join([str(ci) for ci in c[:i + 1]]) - if k not in all_paths: - all_paths[k] = c[:i + 1] - level_counts[i] += 1 - return list(paths.values()), level_counts, paths, all_paths - - -def get_medusa_topks(num_medusa_heads, paths): - medusa_topks = [0] * num_medusa_heads - for p in paths: - for i, k in enumerate(p): - medusa_topks[i] = max(medusa_topks[i], k + 1) - return medusa_topks - - -def get_medusa_tree(num_medusa_heads, medusa_topks, level_counts, paths): - cum_topks = np.cumsum([0] + medusa_topks) - cum_level_counts = np.cumsum([0] + level_counts) - tree_paths = copy.deepcopy(paths) - medusa_tree_ids = list(np.arange(medusa_topks[0])) - medusa_position_offsets = [0] * medusa_topks[0] - for i in range(1, num_medusa_heads): - last_prefix = "-1" - last = -1 - c = -1 - for pi, p in enumerate(paths): - if i < len(p): - prefix_str = ":".join([str(k) for k in p[:i]]) - if last_prefix != prefix_str or last != p[i]: - # new path - medusa_position_offsets.append(i) - medusa_tree_ids.append(p[i] + cum_topks[i]) - last_prefix = prefix_str - last = p[i] - c += 1 - tree_paths[pi][i] = cum_level_counts[i] + c - return medusa_tree_ids, medusa_position_offsets, tree_paths - - -def get_medusa_mask(medusa_tree_ids, medusa_paths): - medusa_mask = torch.zeros((len(medusa_tree_ids), len(medusa_tree_ids))) - medusa_mask[:, 0] = 1 - for p in medusa_paths: - for i, idx in enumerate(p): - if idx < 0: - continue - for j in range(i + 1): - medusa_mask[idx, p[j]] = 1 - return medusa_mask - - -def _medusa_setup(choices_or_paths, num_medusa_heads=None): - choices = copy.deepcopy(choices_or_paths) - sorted_choices = sorted(choices, key=path_sorting_key) - if num_medusa_heads is None: - num_medusa_heads = max([len(c) for c in sorted_choices]) - paths, level_counts, _, _ = choices_2_paths(num_medusa_heads, - sorted_choices) - paths = sorted(paths, key=path_sorting_key) - medusa_topks = get_medusa_topks(num_medusa_heads, paths) - medusa_tree_ids, medusa_position_offsets, tree_paths = get_medusa_tree( - num_medusa_heads, medusa_topks, level_counts, paths) - - num_medusa_tokens = len(medusa_tree_ids) - # now do the padding before converting to torch.Tensor - medusa_paths = [] - for p in tree_paths: - medusa_paths.append( - torch.tensor([-1] + p + ([-2] * (num_medusa_heads - len(p))))) - medusa_topks = torch.tensor(medusa_topks) - medusa_paths = torch.stack(medusa_paths) + 1 - medusa_tree_ids = torch.tensor([-1] + medusa_tree_ids) + 1 - medusa_position_offsets = torch.tensor([-1] + medusa_position_offsets) + 1 - medusa_mask = get_medusa_mask(medusa_tree_ids, medusa_paths) - medusa_packed_mask = get_packed_mask(num_medusa_tokens, medusa_mask[1:, 1:]) - medusa_spec_decoding_use = torch.tensor([1], device="cpu") - - return Namespace( - medusa_mask=medusa_mask.cuda(), - medusa_packed_mask=medusa_packed_mask.cuda(), - medusa_topks=medusa_topks.cuda(), - medusa_paths=medusa_paths.cuda(), - medusa_tree_ids=medusa_tree_ids.cuda(), - medusa_position_offsets=medusa_position_offsets.cuda(), - medusa_spec_decoding_use=medusa_spec_decoding_use, - ) diff --git a/tensorrt_llm/runtime/memory_pools/memory_pools_allocator.py b/tensorrt_llm/runtime/memory_pools/memory_pools_allocator.py deleted file mode 100644 index 719aca5f8444..000000000000 --- a/tensorrt_llm/runtime/memory_pools/memory_pools_allocator.py +++ /dev/null @@ -1,80 +0,0 @@ -from itertools import chain -from typing import List - -import torch - -import tensorrt_llm -from tensorrt_llm.runtime.memory_pools.pool import Pool - - -class MemoryPoolsAllocator(object): - - def __init__(self, num_blocks, tokens_per_block, head_size): - self._pools_metadata = [] - self._pool_pointers = [] - self._pool_mapping = None - - self._num_blocks = num_blocks - self._tokens_per_block = tokens_per_block - self._head_size = head_size - - def allocate(self, dtype, num_kv_heads_per_layer: List[int], device="cuda"): - if isinstance(dtype, str): - dtype = tensorrt_llm._utils.str_dtype_to_torch(dtype) - - # LayerCachePoolLocator{.indexOfPool, .layerIdxInCachePool}" - layers_mapping = [[-1, -1]] * len(num_kv_heads_per_layer) - unique_nkvh = sorted(set(num_kv_heads_per_layer)) - for index_of_pool, kv_head in enumerate(unique_nkvh): - layers = [ - layer for layer, nkvh in enumerate(num_kv_heads_per_layer) - if nkvh == kv_head - ] - - num_layers = len(layers) - cache_shape = ( - self._num_blocks, - num_layers, - 2, - kv_head, - self._tokens_per_block, - self._head_size, - ) - self._pool_pointers.append( - torch.empty(cache_shape, dtype=dtype, device=device)) - self._pools_metadata.append( - Pool(num_kv_heads=kv_head, num_layers=num_layers)) - - for layer_idx_in_cache_pool, layer in enumerate(layers): - layers_mapping[layer] = [index_of_pool, layer_idx_in_cache_pool] - - assert -1 not in set(chain(*layers_mapping)) - self._pool_mapping = torch.tensor(layers_mapping, dtype=torch.int32) - - def get_kv_cache_pool_pointers(self): - return self._get_primarmy_secondary_pool_pointers() - - def _get_primarmy_secondary_pool_pointers(self): - assert len(self._pool_pointers - ) >= 1, "pool pointers haven't been initiated yet" - data_ptr_pointers = torch.tensor(list( - map(lambda x: x.data_ptr(), self._pool_pointers)), - dtype=torch.int64) - host_kv_cache_pool_pointers = torch.cat( - (data_ptr_pointers.view(-1, 1), - torch.zeros(len(self._pool_pointers), 1, dtype=torch.int64)), - dim=1) - - return host_kv_cache_pool_pointers - - @classmethod - def prepare_num_kv_heads_per_layer(cls, kv_head, num_layers): - return [kv_head] * num_layers - - @property - def pools_metadata(self): - return self._pools_metadata - - @property - def pool_mapping(self): - return self._pool_mapping diff --git a/tensorrt_llm/runtime/memory_pools/pool.py b/tensorrt_llm/runtime/memory_pools/pool.py deleted file mode 100644 index 63308ad0d3f0..000000000000 --- a/tensorrt_llm/runtime/memory_pools/pool.py +++ /dev/null @@ -1,7 +0,0 @@ -from dataclasses import dataclass - - -@dataclass -class Pool(object): - num_kv_heads: int - num_layers: int diff --git a/tensorrt_llm/runtime/memory_pools/pools_kv_cache_manager.py b/tensorrt_llm/runtime/memory_pools/pools_kv_cache_manager.py deleted file mode 100644 index 187117302419..000000000000 --- a/tensorrt_llm/runtime/memory_pools/pools_kv_cache_manager.py +++ /dev/null @@ -1,59 +0,0 @@ -from typing import List - -import torch - -from tensorrt_llm.runtime.kv_cache_manager import (GenerationSequence, - KVCacheManager) -from tensorrt_llm.runtime.memory_pools.pool import Pool - - -class PoolsKVCacheManager(object): - - def __init__(self, pools_metadata: List[Pool], max_blocks_per_seq, - num_blocks, tokens_per_block, head_size, - max_attention_window_size, beam_width, sink_token_len) -> None: - self._num_pools = len(pools_metadata) - self._kv_cache_managers = [] - - for pool in pools_metadata: - block_size = pool.num_kv_heads * tokens_per_block * head_size - self._kv_cache_managers.append( - KVCacheManager( - num_layers=pool.num_layers, - num_blocks=num_blocks, - block_size=block_size, - tokens_per_block=tokens_per_block, - max_blocks_per_seq=max_blocks_per_seq, - max_attention_window_size=max_attention_window_size, - sink_token_len=sink_token_len, - beam_width=beam_width, - )) - - def add_sequence(self, - sequence: GenerationSequence, - context_len: int, - always_share_across_beam: bool = False): - for kv_cache_manager in self._kv_cache_managers: - kv_cache_manager.add_sequence(sequence, context_len, - always_share_across_beam) - - def step(self, finished: List[bool]): - for kv_cache_manager in self._kv_cache_managers: - kv_cache_manager.step(finished) - - def get_block_offsets(self, beam_width: int) -> torch.Tensor: - offsets = [] - for kv_cache_manager in self._kv_cache_managers: - block_offset = kv_cache_manager.get_block_offsets(beam_width) - offsets.append(block_offset) - - return torch.stack(offsets) - - def get_single_kv_cache_manager(self): - assert len(self._kv_cache_managers - ) == 1, f"More then one kv cache manager exists" - - return self._kv_cache_managers[0] - - def has_single_pool(self): - return len(self._kv_cache_managers) == 1 diff --git a/tensorrt_llm/runtime/model_config.py b/tensorrt_llm/runtime/model_config.py new file mode 100644 index 000000000000..31e2d084849b --- /dev/null +++ b/tensorrt_llm/runtime/model_config.py @@ -0,0 +1,105 @@ +# 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. + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, List, Optional + +from .._utils import binding_layer_type_to_str, binding_to_str_dtype +from ..llmapi.kv_cache_type import KVCacheType +from ..quantization import QuantMode + + +@dataclass +class ModelConfig: + max_batch_size: int + max_beam_width: int + vocab_size: int + num_layers: int + num_heads: int + num_kv_heads: int + hidden_size: int + gpt_attention_plugin: bool + gemm_allreduce_plugin: str = None + remove_input_padding: bool = False + model_name: str = "" + kv_cache_type: KVCacheType = KVCacheType.CONTINUOUS + cross_attention: bool = False + head_size: int = None + has_position_embedding: bool = True + has_token_type_embedding: bool = False + tokens_per_block: int = 32 + max_prompt_embedding_table_size: int = 0 + quant_mode: QuantMode = QuantMode(0) + gather_context_logits: bool = False + gather_generation_logits: bool = False + dtype: str = "" + lora_plugin: bool = False + lora_target_modules: List[str] = field(default_factory=list) + trtllm_modules_to_hf_modules: dict = None + skip_cross_kv: bool = False + num_medusa_heads: int = 0 + max_medusa_tokens: int = 0 + paged_state: bool = True + mamba_conv1d_plugin: bool = True + conv_kernel: int = 0 + layer_types: List[str] = field(default_factory=list) + rnn_hidden_size: int = 0 + rnn_head_size: int = 0 + rnn_conv_dim_size: int = 0 + state_size: int = 0 + state_dtype: str = "" + gpu_weights_percent: float = 1.0 + # ReDrafter + redrafter_num_beams: int = 0 + redrafter_draft_len_per_beam: int = 0 + num_kv_heads_per_layer: Optional[List[int]] = None + num_kv_heads_per_cross_attn_layer: Optional[List[int]] = None + skip_cross_attn_blocks: bool = False + # language adapter (typed as Optional[Any]; the concrete config type is + # not imported here) + language_adapter_config: Optional[Any] = None + + @classmethod + def from_model_config_cpp(cls, model_config_cpp) -> "ModelConfig": + """Create a partially initialized ModelConfig instance from a given ModelConfig CPP binding instance. + + Note that each of these classes have fields that don't exist in the other, so the created ModelConfigPython + won't have all of its fields initialized. + """ + return cls( + max_batch_size=model_config_cpp.max_batch_size, + max_beam_width=model_config_cpp.max_beam_width, + vocab_size=model_config_cpp.vocab_size, + num_layers=model_config_cpp.num_layers(), + num_heads=model_config_cpp.num_heads, + num_kv_heads=model_config_cpp.num_kv_heads(0), + hidden_size=model_config_cpp.hidden_size, + remove_input_padding=model_config_cpp.use_packed_input, + kv_cache_type=model_config_cpp.kv_cache_type, + cross_attention=model_config_cpp.use_cross_attention, + head_size=model_config_cpp.head_size, + max_prompt_embedding_table_size=model_config_cpp.max_prompt_embedding_table_size, + quant_mode=QuantMode(model_config_cpp.quant_mode.value), + gather_context_logits=model_config_cpp.compute_context_logits, + gather_generation_logits=model_config_cpp.compute_generation_logits, + gpt_attention_plugin=model_config_cpp.use_gpt_attention_plugin, + dtype=binding_to_str_dtype(model_config_cpp.data_type), + num_kv_heads_per_layer=model_config_cpp.num_kv_heads_per_layer, + tokens_per_block=model_config_cpp.tokens_per_block, + lora_plugin=model_config_cpp.use_lora_plugin, + layer_types=[binding_layer_type_to_str(lt) for lt in model_config_cpp.layer_types], + ) diff --git a/tensorrt_llm/runtime/model_runner.py b/tensorrt_llm/runtime/model_runner.py deleted file mode 100644 index 20c250f940eb..000000000000 --- a/tensorrt_llm/runtime/model_runner.py +++ /dev/null @@ -1,1001 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import copy -import json -import math -from pathlib import Path -from typing import List, Optional, Tuple, Union - -import numpy as np -import tensorrt as trt -import torch - -from .. import profiler -from .._deprecation import emit_engine_arch_deprecation -from .._utils import mpi_comm, mpi_world_size, numpy_to_torch -from ..bindings import MpiComm -from ..bindings.executor import Executor -from ..builder import Engine, EngineConfig, get_engine_version -from ..llmapi.kv_cache_type import KVCacheType -from ..logger import logger -from ..mapping import Mapping -from ..quantization import QuantMode -from .generation import (DISABLE_TORCH_DEVICE_SET, ChatGLMGenerationSession, - GenerationSession, LogitsProcessor, LoraManager, - ModelConfig, QWenForCausalLMGenerationSession, - SamplingConfig, StoppingCriteria, to_word_list_format) - - -def get_engine_name(model: str, dtype: str, tp_size: int, pp_size: int, - rank: int) -> str: - """ - Get the serialized engine file name. - - Args: - model (str): - Model name, e.g., bloom, gpt. - dtype (str): - Data type, e.g., float32, float16, bfloat16, - tp_size (int): - The size of tensor parallel. - pp_size (int): - The size of pipeline parallel. - rank (int): - The rank id. - - Returns: - str: The serialized engine file name. - """ - if pp_size == 1: - return '{}_{}_tp{}_rank{}.engine'.format(model, dtype, tp_size, rank) - return '{}_{}_tp{}_pp{}_rank{}.engine'.format(model, dtype, tp_size, - pp_size, rank) - - -def read_config(config_path: Path) -> Tuple[ModelConfig, dict]: - """ - Read the engine config file and create a ModelConfig instance, return the ModelConfig instance - and other config fields in a dict. - - Args: - config_path (Path): - The path of engine config file. - - Returns: - Tuple[ModelConfig, dict]: A ModelConfig instance and other config fields. - """ - with open(config_path, 'r') as f: - config = json.load(f) - return _builder_to_model_config(config) - - -def _builder_to_model_config(config: dict) -> Tuple[ModelConfig, dict]: - builder_config = config['builder_config'] - model_name = builder_config['name'] - dtype = builder_config['precision'] - tp_size = builder_config['tensor_parallel'] - pp_size = builder_config.get('pipeline_parallel', 1) - kv_cache_type = builder_config.get('kv_cache_type') - if kv_cache_type is not None: - kv_cache_type = KVCacheType(kv_cache_type) - world_size = tp_size * pp_size - assert world_size == mpi_world_size(), \ - f'Engine world size ({tp_size} * {pp_size}) != Runtime world size ({mpi_world_size()})' - - num_heads = builder_config['num_heads'] - assert num_heads % tp_size == 0, \ - f"The number of heads ({num_heads}) is not a multiple of tp_size ({tp_size})" - num_kv_heads = builder_config.get('num_kv_heads', num_heads) - # TODO: multi_query_mode should be removed - multi_query_mode = builder_config.get('multi_query_mode', False) - if multi_query_mode: - logger.warning( - "`multi_query_mode` config is deprecated. Please rebuild the engine." - ) - # num_kv_heads, if exists in config, should override multi_query_mode - if multi_query_mode and ('num_kv_heads' not in builder_config): - num_kv_heads = 1 - num_heads = num_heads // tp_size - num_kv_heads = (num_kv_heads + tp_size - 1) // tp_size - head_size = builder_config.get('head_size', None) - - hidden_size = builder_config['hidden_size'] // tp_size - vocab_size = builder_config['vocab_size'] - num_layers = builder_config['num_layers'] - max_batch_size = builder_config['max_batch_size'] - max_beam_width = builder_config['max_beam_width'] - - cross_attention = builder_config.get('cross_attention', False) - has_position_embedding = builder_config.get('has_position_embedding', True) - has_token_type_embedding = builder_config.get('has_token_type_embedding', - False) - gather_context_logits = builder_config.get('gather_context_logits', False) - gather_generation_logits = builder_config.get('gather_generation_logits', - False) - max_prompt_embedding_table_size = builder_config.get( - 'max_prompt_embedding_table_size', 0) - quant_mode = QuantMode(builder_config.get('quant_mode', 0)) - lora_target_modules = builder_config.get('lora_target_modules') - lora_trtllm_modules_to_hf_modules = builder_config.get( - 'trtllm_modules_to_hf_modules') - max_medusa_token_len = builder_config.get('max_draft_len', 0) - num_medusa_heads = builder_config.get('num_medusa_heads', 0) - - skip_cross_attn_blocks = bool(config['pretrained_config'].get( - 'skip_cross_attn_blocks', False)) - - # ReDrafter - redrafter_num_beams = config['pretrained_config'].get( - 'redrafter_num_beams', 0) - redrafter_draft_len_per_beam = config['pretrained_config'].get( - 'redrafter_draft_len_per_beam', 0) - - plugin_config = config['plugin_config'] - use_gpt_attention_plugin = bool(plugin_config['gpt_attention_plugin']) - gemm_allreduce_plugin = plugin_config['gemm_allreduce_plugin'] - mamba_conv1d_plugin = bool(plugin_config['mamba_conv1d_plugin']) - remove_input_padding = plugin_config['remove_input_padding'] - paged_state = plugin_config['paged_state'] - tokens_per_block = plugin_config['tokens_per_block'] - lora_plugin = plugin_config.get('lora_plugin') - - model_config = ModelConfig( - max_batch_size=max_batch_size, - max_beam_width=max_beam_width, - vocab_size=vocab_size, - num_layers=num_layers, - num_heads=num_heads, - num_kv_heads=num_kv_heads, - hidden_size=hidden_size, - head_size=head_size, - gpt_attention_plugin=use_gpt_attention_plugin, - gemm_allreduce_plugin=gemm_allreduce_plugin, - mamba_conv1d_plugin=mamba_conv1d_plugin, - remove_input_padding=remove_input_padding, - model_name=model_name, - kv_cache_type=kv_cache_type, - paged_state=paged_state, - cross_attention=cross_attention, - has_position_embedding=has_position_embedding, - has_token_type_embedding=has_token_type_embedding, - tokens_per_block=tokens_per_block, - max_prompt_embedding_table_size=max_prompt_embedding_table_size, - quant_mode=quant_mode, - gather_context_logits=gather_context_logits, - gather_generation_logits=gather_generation_logits, - dtype=dtype, - lora_plugin=lora_plugin, - lora_target_modules=lora_target_modules, - trtllm_modules_to_hf_modules=lora_trtllm_modules_to_hf_modules, - num_medusa_heads=num_medusa_heads, - max_medusa_tokens=max_medusa_token_len, - skip_cross_attn_blocks=skip_cross_attn_blocks, - # ReDrafter - redrafter_num_beams=redrafter_num_beams, - redrafter_draft_len_per_beam=redrafter_draft_len_per_beam, - ) - - other_config = { - 'world_size': world_size, - 'tp_size': tp_size, - 'pp_size': pp_size, - 'max_batch_size': builder_config['max_batch_size'], - 'max_input_len': builder_config['max_input_len'], - 'max_output_len': builder_config['max_output_len'], - 'max_beam_width': builder_config['max_beam_width'] - } - return model_config, other_config - - -def _engine_config_to_model_config(engine_config: EngineConfig, - **kwargs) -> ModelConfig: - pretrained_config = engine_config.pretrained_config - build_config = engine_config.build_config - - tp_size = pretrained_config.mapping.tp_size - num_heads = pretrained_config.num_attention_heads // tp_size - num_kv_heads = pretrained_config.num_key_value_heads - num_kv_heads = (num_kv_heads + tp_size - 1) // tp_size - hidden_size = pretrained_config.hidden_size // tp_size - head_size = pretrained_config.head_size - - rnn_config_items = [ - 'conv_kernel', 'layer_types', 'rnn_hidden_size', 'state_size', - 'state_dtype', 'rnn_head_size', 'rnn_conv_dim_size' - ] - rnn_configs_kwargs = {} - for item in rnn_config_items: - if hasattr(pretrained_config, item): - rnn_configs_kwargs[item] = getattr(pretrained_config, item) - - if not hasattr(build_config, 'kv_cache_type'): - logger.Warning( - 'Build config doesn\'t have kv_cache_type, you might need to rebuild your enigne.' - ) - - # TODO(oargov): this is a hack, make it prettier! - if hasattr(pretrained_config, "num_kv_heads_per_layer"): - pp_rank = pretrained_config.mapping.pp_rank - pp_size = pretrained_config.mapping.pp_size - layers_per_pp_rank = pretrained_config.num_hidden_layers // pp_size - first_local_layer = layers_per_pp_rank * pp_rank - first_layer_next_rank = first_local_layer + layers_per_pp_rank - layer_types = getattr(pretrained_config, "layer_types", ["attention"]) - num_attn_layers_lower_ranks = [ - layer_types[layer_idx % len(layer_types)] - for layer_idx in range(first_local_layer) - ].count("attention") - num_local_attn_layers = [ - layer_types[layer_idx % len(layer_types)] - for layer_idx in range(first_local_layer, first_layer_next_rank) - ].count("attention") - num_kv_heads_per_layer = pretrained_config.num_kv_heads_per_layer[ - num_attn_layers_lower_ranks:num_attn_layers_lower_ranks + - num_local_attn_layers] - num_kv_heads_per_layer = [(nheads + tp_size - 1) // tp_size - for nheads in num_kv_heads_per_layer] - - elif hasattr(pretrained_config, "get_layer_num_kv_heads"): - # each layer has a different number of kv heads - attention_layers = [ - layer_idx for layer_idx, layer_type in enumerate( - pretrained_config.layer_types) if layer_type == "attention" - ] if hasattr(pretrained_config, "layer_types") else list( - range(pretrained_config.num_hidden_layers)) - num_kv_heads_per_layer = [ - pretrained_config.get_layer_num_kv_heads(layer_idx) - if layer_idx in attention_layers else 0 - for layer_idx in range(pretrained_config.num_hidden_layers) - ] - else: - num_kv_heads_per_layer = None - - if hasattr(pretrained_config, "num_kv_heads_per_cross_attn_layer"): - num_kv_heads_per_cross_attn_layer = pretrained_config.num_kv_heads_per_cross_attn_layer - else: - num_kv_heads_per_cross_attn_layer = None - - return ModelConfig( - max_batch_size=build_config.max_batch_size, - max_beam_width=build_config.max_beam_width, - vocab_size=pretrained_config.vocab_size, - num_layers=pretrained_config.num_hidden_layers, - num_heads=num_heads, - num_kv_heads=num_kv_heads, - hidden_size=hidden_size, - head_size=head_size, - gpt_attention_plugin=bool( - build_config.plugin_config.gpt_attention_plugin), - gemm_allreduce_plugin=build_config.plugin_config.gemm_allreduce_plugin, - mamba_conv1d_plugin=bool( - build_config.plugin_config.mamba_conv1d_plugin), - remove_input_padding=build_config.plugin_config.remove_input_padding, - paged_state=build_config.plugin_config.paged_state, - tokens_per_block=build_config.plugin_config.tokens_per_block, - quant_mode=pretrained_config.quant_mode, - gather_context_logits=build_config.gather_context_logits, - gather_generation_logits=build_config.gather_generation_logits, - dtype=pretrained_config.dtype, - max_prompt_embedding_table_size=build_config. - max_prompt_embedding_table_size, - lora_plugin=build_config.plugin_config.lora_plugin, - lora_target_modules=build_config.lora_config.lora_target_modules, - trtllm_modules_to_hf_modules=build_config.lora_config. - trtllm_modules_to_hf_modules, - max_medusa_tokens=pretrained_config.max_draft_len if hasattr( - pretrained_config, 'max_draft_len') else 0, - num_medusa_heads=pretrained_config.num_medusa_heads if hasattr( - pretrained_config, 'num_medusa_heads') else 0, - **rnn_configs_kwargs, - num_kv_heads_per_layer=num_kv_heads_per_layer, - num_kv_heads_per_cross_attn_layer=num_kv_heads_per_cross_attn_layer, - redrafter_num_beams=pretrained_config.redrafter_num_beams if hasattr( - pretrained_config, 'redrafter_num_beams') else 0, - redrafter_draft_len_per_beam=pretrained_config. - redrafter_draft_len_per_beam - if hasattr(pretrained_config, 'redrafter_draft_len_per_beam') else 0, - kv_cache_type=getattr(build_config, 'kv_cache_type', - KVCacheType.CONTINUOUS), - cross_attention=getattr(pretrained_config, 'cross_attention', False), - has_position_embedding=getattr(pretrained_config, - 'has_position_embedding', True), - skip_cross_attn_blocks=getattr(pretrained_config, - 'skip_cross_attn_blocks', False), - **kwargs) - - -class ModelRunnerMixin: - - def _check_inputs(self, batch_input_ids: List[torch.Tensor], - sampling_config: SamplingConfig): - batch_size = len(batch_input_ids) - if batch_size > self.max_batch_size: - raise RuntimeError( - f"Input batch size ({batch_size}) exceeds the engine or specified limit ({self.max_batch_size})" - ) - input_lengths = [x.size(0) for x in batch_input_ids] - max_length = max(input_lengths) - if max_length > self.max_input_len: - raise RuntimeError( - f"Maximum input length ({max_length}) exceeds the engine or specified limit ({self.max_input_len})" - ) - if max_length + sampling_config.max_new_tokens > self.max_seq_len: - raise RuntimeError( - f"Maximum input length ({max_length}) + maximum new tokens ({sampling_config.max_new_tokens}) exceeds the engine or specified limit ({self.max_seq_len})" - ) - if sampling_config.num_beams > self.max_beam_width: - raise RuntimeError( - f"Num beams ({sampling_config.num_beams}) exceeds the engine or specified limit ({self.max_beam_width})" - ) - - def _prepare_inputs(self, batch_input_ids: List[torch.Tensor], - pad_id: int) -> Tuple[torch.Tensor]: - # Cast to int32 - batch_input_ids = [x.type(torch.int32) for x in batch_input_ids] - input_lengths = [x.size(0) for x in batch_input_ids] - max_length = max(input_lengths) - - if self.remove_input_padding: - batch_input_ids = torch.concat(batch_input_ids) - else: - # Right padding for trt-llm - paddings = [ - torch.ones(max_length - l, dtype=torch.int32) * pad_id - for l in input_lengths - ] - batch_input_ids = [ - torch.cat([x, pad]) for x, pad in zip(batch_input_ids, paddings) - ] - batch_input_ids = torch.stack(batch_input_ids) - input_lengths = torch.tensor(input_lengths, dtype=torch.int32) - return batch_input_ids, input_lengths - - def _prepare_outputs(self, outputs: Optional[dict], - input_lengths: torch.Tensor) -> dict: - if outputs is not None: - batch_size = input_lengths.size(0) - if 'context_logits' in outputs: - if self.mapping.has_pp(): - # If pp size > 1, the context logits and generation logits are both in last pp - # Last pp rank send context logits and generation logits to rank 0 - if self.mapping.is_last_pp_rank(): - context_logits = outputs['context_logits'] - context_logits_host = context_logits.cpu() - mpi_comm().send(context_logits_host, dest=0) - elif self.mapping.is_first_pp_rank(): - context_logits_host = mpi_comm().recv( - source=self.mapping.prev_pp_rank() - ) # Prev pp rank of rank=0 is the last pp - context_logits = context_logits_host.to( - torch.device('cuda:0')) - outputs['context_logits'] = context_logits - - context_logits = outputs['context_logits'] - - context_logits_output = [] - if self.remove_input_padding: - if isinstance(self.session, Executor) and batch_size > 1: - # The starting position of the context logits buffer of each micro batch is separated - num_batches = self.mapping.pp_size - micro_batch_size = math.ceil(batch_size / - self.mapping.pp_size) - - for i in range(num_batches): - start_idx = i * micro_batch_size - end_idx = min(start_idx + micro_batch_size, - batch_size) - micro_context_logits = context_logits[ - start_idx:end_idx] - micro_input_lengths = input_lengths[ - start_idx:end_idx] - - micro_context_logits = micro_context_logits.flatten( - end_dim=-2) - seg_points = [0] + micro_input_lengths.cumsum( - dim=0).tolist() - context_logits_output += [ - micro_context_logits[s:e] - for s, e in zip(seg_points[:-1], seg_points[1:]) - ] - else: - context_logits = context_logits.flatten(end_dim=-2) - - seg_points = [0] + input_lengths.cumsum(dim=0).tolist() - context_logits_output = [ - context_logits[s:e] - for s, e in zip(seg_points[:-1], seg_points[1:]) - ] - else: - context_logits_output = [ - context_logits[bidx, :input_lengths[bidx]] - for bidx in range(batch_size) - ] - - assert len(context_logits_output) == batch_size - outputs['context_logits'] = context_logits_output - - if 'generation_logits' in outputs: - if self.mapping.has_pp(): - if self.mapping.is_last_pp_rank(): - generation_logits = outputs['generation_logits'] - if isinstance(generation_logits, list): - generation_logits_host = [ - logits.cpu() for logits in generation_logits - ] - else: - generation_logits_host = generation_logits.cpu() - mpi_comm().send(generation_logits_host, dest=0) - elif self.mapping.is_first_pp_rank(): - generation_logits_host = mpi_comm().recv( - source=self.mapping.prev_pp_rank() - ) # Prev pp rank of rank=0 is the last pp - if isinstance(generation_logits_host, list): - generation_logits = [ - logits.to(torch.device('cuda:0')) - for logits in generation_logits_host - ] - else: - generation_logits = generation_logits_host.to( - torch.device('cuda:0')) - outputs['generation_logits'] = generation_logits - - if isinstance(self.session, GenerationSession): - # Convert logits format to be same as GptSession - generation_logits = torch.stack( - outputs['generation_logits'], dim=1) - batch_x_beam, max_gen_len, voc_size = generation_logits.size( - ) - num_beams = batch_x_beam // batch_size - generation_logits = generation_logits.view( - batch_size, num_beams, max_gen_len, voc_size) - outputs['generation_logits'] = generation_logits - - return outputs - - def _prepare_embedding_table(self, prompt_table: Union[str, torch.Tensor]): - if isinstance(prompt_table, str): - prompt_table_data = numpy_to_torch( - np.load(prompt_table)).to(dtype=self.dtype) - else: - assert isinstance( - prompt_table, - torch.Tensor), "Prompt table should be str or torch.Tensor" - prompt_table_data = prompt_table.to(dtype=self.dtype) - torch.cuda.current_stream().synchronize() - - return prompt_table_data - - def _prepare_ptuning(self, prompt_table: Union[str, torch.Tensor], - tasks: str, batch_size: int): - if self.max_prompt_embedding_table_size == 0: - return {} - - if prompt_table is not None: - prompt_table_data = self._prepare_embedding_table(prompt_table) - if len(prompt_table_data.size()) == 3: - _, task_vocab_size, hidden_size = prompt_table_data.size() - elif len(prompt_table_data.size()) == 2: - task_vocab_size, hidden_size = prompt_table_data.size() - task_vocab_size = torch.tensor([task_vocab_size], dtype=torch.int32) - prompt_table_data = prompt_table_data.view(-1, hidden_size) - else: - prompt_table_data = torch.empty( - [1, self.hidden_size * self.mapping.tp_size], dtype=self.dtype) - task_vocab_size = torch.zeros([1], dtype=torch.int32) - if tasks is not None: - tasks = torch.tensor([int(t) for t in tasks.split(',')], - dtype=torch.int32) - assert tasks.size(0) == batch_size, \ - f"Number of supplied tasks ({tasks.size(0)}) must match input batch size ({batch_size})" - else: - tasks = torch.zeros([batch_size], dtype=torch.int32) - - if isinstance(self.session, GenerationSession): - return { - 'prompt_embedding_table': prompt_table_data.cuda(), - 'tasks': tasks.cuda(), - 'prompt_vocab_size': task_vocab_size.cuda() - } - else: - return { - 'embedding_table': prompt_table_data.cuda(), - 'tasks': tasks.cuda(), - 'vocab_size': task_vocab_size.cuda() - } - - -class ModelRunner(ModelRunnerMixin): - """ - An interface class that wraps GenerationSession and provides generation methods. - """ - - def __init__( - self, - session: GenerationSession, - max_batch_size: int, - max_input_len: int, - max_seq_len: int, - max_beam_width: int, - kv_cache_type: KVCacheType, - lora_manager: Optional[LoraManager] = None, - ) -> None: - """ - Create a ModelRunner instance. - You are recommended to use the from_dir method to load the engine and create a ModelRunner instance. - - Args: - session (GenerationSession): - The TensorRT session created from an engine. - max_batch_size (int): - The maximum batch size allowed for the input. - max_input_len (int): - The maximum input length allowed for the input. - max_seq_len (int): - The maximum sequence length (input + new tokens). - max_beam_width (int): - The maximum beam width. - lora_manager (LoraManager): - The LoRA manager to handle LoRA weights. - """ - emit_engine_arch_deprecation("ModelRunner") - self.session = session - self.max_batch_size = max_batch_size - self.max_input_len = max_input_len - self.max_seq_len = max_seq_len - self.max_beam_width = max_beam_width - self.lora_manager = lora_manager - self.kv_cache_type = kv_cache_type - self.enable_context_fmha_fp32_acc = False - self.multi_block_mode = True - - @classmethod - def from_engine( - cls, - engine: Engine, - *, - max_output_len: Optional[int], - lora_dir: Optional[List[str]], - rank: int, - debug_mode: bool, - lora_ckpt_source: str, - medusa_choices: List[List[int]], - stream: torch.cuda.Stream, - gpu_weights_percent: float, - enable_context_fmha_fp32_acc: Optional[bool], - multi_block_mode: Optional[bool], - ) -> 'ModelRunner': - model_config = _engine_config_to_model_config( - engine.config, gpu_weights_percent=gpu_weights_percent) - - if model_config.kv_cache_type == KVCacheType.DISABLED: - assert max_output_len == 1 or max_output_len is None, 'Disabled KV cache is intended for context phase only now.' - - pretrained_config = engine.config.pretrained_config - build_config = engine.config.build_config - max_batch_size = build_config.max_batch_size - max_input_len = build_config.max_input_len - max_seq_len = build_config.max_seq_len - max_beam_width = build_config.max_beam_width - if 'GLM' in pretrained_config.architecture and pretrained_config.chatglm_version in [ - 'glm', 'chatglm' - ]: - session_cls = ChatGLMGenerationSession - else: - session_cls = GenerationSession - engine_buffer = engine.engine - runtime_mapping = pretrained_config.mapping - - if medusa_choices is not None: - assert session_cls == GenerationSession, "Medusa is only supported by GenerationSession" - - assert model_config.max_medusa_tokens > 0, \ - "medusa_chioce is specified but model_config.max_medusa_tokens is 0." - - if MpiComm.size() > runtime_mapping.gpus_per_node: - assert MpiComm.local_size() == runtime_mapping.gpus_per_node - if not DISABLE_TORCH_DEVICE_SET: - torch.cuda.set_device(rank % runtime_mapping.gpus_per_node) - session = session_cls(model_config, - engine_buffer, - runtime_mapping, - debug_mode=debug_mode, - stream=stream) - if session.runtime.engine.streamable_weights_size: - session.runtime._set_weight_streaming(gpu_weights_percent) - - if session.use_lora_plugin: - lora_manager = LoraManager(mapping=runtime_mapping, - model_config=model_config) - if lora_dir is not None: - lora_manager.load_from_ckpt(lora_dir, - model_config=model_config, - ckpt_source=lora_ckpt_source) - else: - lora_manager = None - - runner = cls(session=session, - max_batch_size=max_batch_size, - max_input_len=max_input_len, - max_seq_len=max_seq_len, - max_beam_width=max_beam_width, - kv_cache_type=model_config.kv_cache_type, - lora_manager=lora_manager) - runner.enable_context_fmha_fp32_acc = enable_context_fmha_fp32_acc - runner.multi_block_mode = multi_block_mode - return runner - - @classmethod - def from_dir( - cls, - engine_dir: str, - *, - max_output_len: Optional[int] = None, - lora_dir: Optional[List[str]] = None, - rank: int = 0, - debug_mode: bool = False, - lora_ckpt_source: str = "hf", - medusa_choices: List[List[int]] = None, - stream: torch.cuda.Stream = None, - gpu_weights_percent: float = 1, - enable_context_fmha_fp32_acc: Optional[bool] = None, - multi_block_mode: Optional[bool] = None, - fail_fast_on_attention_window_too_large: bool = False, - ) -> 'ModelRunner': - """ - Create a ModelRunner instance from an engine directory. - - Args: - engine_dir (str): - The directory that contains the serialized engine files and config files. - max_output_len (Optional[int]): - max_output_len, this arg might be available only when loading time, generate will still to check when disable_kv_cache is enabled. - lora_dir (Optional[List[str]]): - The directories that contain LoRA weights. - rank (int): - The runtime rank id. - debug_mode (bool): - Whether or not to turn on the debug mode. - medusa_choices (List[List[int]]): - Medusa choices to use when in Medusa decoding - stream (torch.cuda.Stream): - Stream to use. - multi_block_mode (bool): - Whether to distribute the work across multiple CUDA thread-blocks on the GPU for masked MHA kernel. - fail_fast_on_attention_window_too_large (bool): - Exit with runtime error when attention window is too large to fit even a single sequence in the KV cache. - Note: This parameter is only applicable to C++ runtime (ModelRunnerCpp). - Returns: - ModelRunner: An instance of ModelRunner. - """ - engine_version = get_engine_version(engine_dir) - profiler.start('load tensorrt_llm engine') - # the old engine format - if engine_version is None: - engine_dir = Path(engine_dir) - config_path = engine_dir / "config.json" - model_config, other_config = read_config(config_path) - world_size = other_config.pop('world_size') - tp_size = other_config.pop('tp_size') - pp_size = other_config.pop('pp_size') - max_batch_size = other_config.pop('max_batch_size') - max_input_len = other_config.pop('max_input_len') - max_output_len = other_config.pop('max_output_len') - max_beam_width = other_config.pop('max_beam_width') - runtime_mapping = Mapping(world_size=world_size, - rank=rank, - tp_size=tp_size, - pp_size=pp_size) - - engine_name = get_engine_name(model_config.model_name, - model_config.dtype, tp_size, pp_size, - rank) - serialize_path = engine_dir / engine_name - - with open(serialize_path, 'rb') as f: - engine_buffer = f.read() - - if model_config.model_name in ('chatglm_6b', 'glm_10b'): - session_cls = ChatGLMGenerationSession - elif model_config.model_name == 'qwen': - session_cls = QWenForCausalLMGenerationSession - else: - session_cls = GenerationSession - - if medusa_choices is not None: - assert model_config.max_medusa_tokens > 0, \ - "medusa_choice is specified but model_config.max_medusa_tokens is 0." - - if not DISABLE_TORCH_DEVICE_SET: - torch.cuda.set_device(rank % runtime_mapping.gpus_per_node) - session = session_cls(model_config, - engine_buffer, - runtime_mapping, - debug_mode=debug_mode, - stream=stream) - if session.use_lora_plugin: - lora_manager = LoraManager(mapping=runtime_mapping, - model_config=model_config) - if lora_dir is not None: - lora_manager.load_from_ckpt(lora_dir, - model_config=model_config, - ckpt_source=lora_ckpt_source) - else: - lora_manager = None - - if session.runtime.engine.streamable_weights_size: - session.runtime._set_weight_streaming(gpu_weights_percent) - - profiler.stop('load tensorrt_llm engine') - loading_time = profiler.elapsed_time_in_sec( - "load tensorrt_llm engine") - logger.info(f'Load engine takes: {loading_time} sec') - - runner = cls(session=session, - max_batch_size=max_batch_size, - max_input_len=max_input_len, - max_seq_len=max_input_len + max_output_len, - max_beam_width=max_beam_width, - kv_cache_type=KVCacheType.CONTINUOUS, - lora_manager=lora_manager) - runner.enable_context_fmha_fp32_acc = enable_context_fmha_fp32_acc - runner.multi_block_mode = multi_block_mode - return runner - else: - # the new engine format - engine = Engine.from_dir(engine_dir, rank) - if lora_dir is None: - config_lora_dir = engine.config.build_config.lora_config.lora_dir - if len(config_lora_dir) > 0: - lora_dir = [ - f"{engine_dir}/{dir}" for dir in config_lora_dir - ] - lora_ckpt_source = engine.config.build_config.lora_config.lora_ckpt_source - - runner = ModelRunner.from_engine( - engine=engine, - max_output_len=max_output_len, - lora_dir=lora_dir, - rank=rank, - debug_mode=debug_mode, - lora_ckpt_source=lora_ckpt_source, - medusa_choices=medusa_choices, - stream=stream, - gpu_weights_percent=gpu_weights_percent, - enable_context_fmha_fp32_acc=enable_context_fmha_fp32_acc, - multi_block_mode=multi_block_mode, - ) - profiler.stop('load tensorrt_llm engine') - loading_time = profiler.elapsed_time_in_sec( - "load tensorrt_llm engine") - logger.info(f'Load engine takes: {loading_time} sec') - return runner - - @property - def dtype(self) -> torch.dtype: - return self.session.dtype - - @property - def vocab_size(self) -> int: - return self.session.vocab_size - - @property - def vocab_size_padded(self) -> int: - return self.session.vocab_size_padded - - @property - def hidden_size(self) -> int: - return self.session.hidden_size - - @property - def num_heads(self) -> int: - return self.session.num_heads - - @property - def num_layers(self) -> int: - return self.session.num_layers - - @property - def max_sequence_length(self) -> int: - return self.max_seq_len - - @property - def remove_input_padding(self) -> bool: - return self.session.remove_input_padding - - @property - def use_lora_plugin(self) -> bool: - return self.session.use_lora_plugin - - @property - def max_prompt_embedding_table_size(self) -> int: - return self.session.max_prompt_embedding_table_size - - @property - def mapping(self) -> Mapping: - return self.session.mapping - - @property - def gather_context_logits(self) -> bool: - return self.session.gather_context_logits - - @property - def gather_generation_logits(self) -> bool: - return self.session.gather_generation_logits - - def generate(self, - batch_input_ids: List[torch.Tensor], - position_ids: List[torch.Tensor] = None, - sampling_config: Optional[SamplingConfig] = None, - prompt_table: Optional[Union[str, torch.Tensor]] = None, - prompt_tasks: Optional[str] = None, - lora_uids: Optional[list] = None, - streaming: bool = False, - output_generation_logits: bool = False, - stopping_criteria: Optional[StoppingCriteria] = None, - logits_processor: Optional[LogitsProcessor] = None, - medusa_choices: Optional[List[List[int]]] = None, - encoder_max_input_length: int = None, - encoder_input_features: List[torch.Tensor] = None, - encoder_output_lengths: List[torch.Tensor] = None, - cross_attention_masks: List[torch.Tensor] = None, - **kwargs) -> Union[torch.Tensor, dict]: - """ - Generates sequences of token ids. - The generation-controlling parameters are set in the sampling_config; it will be set to a default one if not passed. - You can override any sampling_config's attributes by passing corresponding parameters. - - Args: - batch_input_ids (List[torch.Tensor]): - A list of input id tensors. Each tensor is of shape (sequence_length, ). - sampling_config (SamplingConfig): - The sampling configuration to be used as base parametrization for the generation call. - The passed **kwargs matching the sampling_config's attributes will override them. - If the sampling_config is not provided, a default will be used. - prompt_table (str or torch.Tensor): - The file path of prompt table (.npy format, exported by nemo_prompt_convert.py) or the prompt table itself. - prompt_tasks (str): - The prompt tuning task ids for the input batch, in format of comma-separated list (e.g., 0,3,1,0). - lora_uids (list): - The uids of LoRA weights for the input batch. Use -1 to disable the LoRA module. - streaming (bool): - Whether or not to use streaming mode for generation. - stopping_criteria (StoppingCriteria): - Custom stopping criteria. - logits_processor (LogitsProcessor): - Custom logits processors. - medusa_choices (List[List[int]]): - Medusa decoding choices. - kwargs (Dict[str, Any]: - Ad hoc parametrization of sampling_config. - The passed **kwargs matching the sampling_config's attributes will override them. - Returns: - torch.Tensor or dict: - If return_dict=False, the method returns generated output_ids. - If return_dict=True, the method returns a dict of output_ids, - sequence_lengths (if sampling_config.output_sequence_lengths=True), - context_logits and generation_logits (if self.gather_context_logits=True - and self.gather_generation_logits=True, respectively). - """ - # Use sampling_config like HF's generation_config - if sampling_config is None: - sampling_config = SamplingConfig(end_id=None, pad_id=None) - else: - sampling_config = copy.deepcopy(sampling_config) - sampling_config.update(**kwargs) - - # To prevent numerical overflow when the temperature is set to 0.0 - # Modify generation.SamplingConfig - if isinstance(sampling_config.temperature, - float) and sampling_config.temperature == 0.0: - logger.warning( - "Convert `temperature=0.0` to `temperature=1.0` and `top_k=1` to prevent overflow." - ) - sampling_config.temperature = 1.0 - sampling_config.top_k = 1 - - self._check_inputs(batch_input_ids, sampling_config) - - if kwargs.get('num_return_sequences', None) is not None: - raise ValueError( - 'num_return_sequences will be ignored since ' - 'num_return_sequences > 1 is not supported on python runtime. ' - 'Please use C++ runtime.') - - batch_size = len(batch_input_ids) - batch_input_ids, input_lengths = self._prepare_inputs( - batch_input_ids, sampling_config.pad_id) - - def maybe_convert_to_words_list_format( - words_list: Optional[Union[list, np.ndarray, torch.Tensor]] - ) -> Optional[np.ndarray]: - if words_list is None or isinstance(words_list, np.ndarray): - return words_list - elif isinstance(words_list, torch.Tensor): - return words_list.numpy() - elif isinstance(words_list, list): - return to_word_list_format(words_list) - else: - raise TypeError( - f"Unexpected words_list type={type(words_list)}. Only list, np.ndarray, and torch.Tensor are supported." - ) - - if cross_attention_masks is not None: - encoder_input_features = torch.concat(encoder_input_features) - encoder_output_lengths = torch.concat(encoder_output_lengths) - - sampling_config.bad_words_list = maybe_convert_to_words_list_format( - sampling_config.bad_words_list) - sampling_config.stop_words_list = maybe_convert_to_words_list_format( - sampling_config.stop_words_list) - - if not self.kv_cache_type and sampling_config.max_new_tokens > 1: - raise RuntimeError( - 'Disabled KV cache is intended for context phase only now.') - - self.session.setup( - batch_size=batch_size, - max_context_length=input_lengths.max().item(), - max_new_tokens=sampling_config.max_new_tokens, - beam_width=sampling_config.num_beams, - max_attention_window_size=sampling_config.max_attention_window_size, - sink_token_length=sampling_config.sink_token_length, - lora_manager=self.lora_manager, - lora_uids=lora_uids, - medusa_choices=medusa_choices, - enable_context_fmha_fp32_acc=self.enable_context_fmha_fp32_acc, - multi_block_mode=self.multi_block_mode, - encoder_max_input_length=encoder_max_input_length, - ) - - batch_input_ids = batch_input_ids.cuda() - input_lengths = input_lengths.cuda() - other_kwargs = self._prepare_ptuning(prompt_table, prompt_tasks, - batch_size) - other_kwargs['skip_cross_attn_blocks'] = kwargs.get( - 'skip_cross_attn_blocks', None) - outputs = self.session.decode( - batch_input_ids, - input_lengths, - sampling_config, - stop_words_list=sampling_config.stop_words_list, - bad_words_list=sampling_config.bad_words_list, - output_sequence_lengths=sampling_config.output_sequence_lengths, - output_generation_logits=output_generation_logits, - return_dict=sampling_config.return_dict, - streaming=streaming, - stopping_criteria=stopping_criteria, - logits_processor=logits_processor, - position_ids=position_ids, - encoder_output=encoder_input_features, - encoder_input_lengths=encoder_output_lengths, - cross_attention_mask=cross_attention_masks, - **other_kwargs) - if sampling_config.return_dict: - if streaming: - outputs = (self._prepare_outputs(curr_outputs, input_lengths) - for curr_outputs in outputs) - else: - outputs = self._prepare_outputs(outputs, input_lengths) - return outputs - - def serialize_engine(self) -> trt.IHostMemory: - """ - Serialize the engine. - - Returns: - bytes: The serialized engine. - """ - return self.session.runtime._serialize_engine() diff --git a/tensorrt_llm/runtime/model_runner_cpp.py b/tensorrt_llm/runtime/model_runner_cpp.py deleted file mode 100644 index 3186a47f697a..000000000000 --- a/tensorrt_llm/runtime/model_runner_cpp.py +++ /dev/null @@ -1,1223 +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. - -import copy -import os -from os.path import join -from pathlib import Path -from typing import Dict, List, Optional, Union - -import torch - -from .. import profiler -from .._deprecation import emit_engine_arch_deprecation -from .._utils import maybe_pin_memory, mpi_broadcast -from ..bindings import DataType, GptJsonConfig, ModelConfig, WorldConfig -from ..bindings import executor as trtllm -from ..bindings.executor import (DecodingMode, ExternalDraftTokensConfig, - OrchestratorConfig, ParallelConfig) -from ..builder import EngineConfig -from ..layers import MropeParams -from ..llmapi.kv_cache_type import KVCacheType -from ..logger import logger -from ..mapping import Mapping -from .generation import LogitsProcessor, LoraManager -from .generation import ModelConfig as ModelConfigPython -from .generation import SamplingConfig, StoppingCriteria -from .model_runner import ModelRunnerMixin, _engine_config_to_model_config - -_bindings_dtype_to_torch_dtype_dict = { - DataType.FLOAT: torch.float, - DataType.HALF: torch.half, - DataType.INT8: torch.int8, - DataType.INT32: torch.int32, - DataType.BOOL: torch.bool, - DataType.UINT8: torch.uint8, - DataType.BF16: torch.bfloat16, - DataType.INT64: torch.int64 -} - -SamplingConfigType = Union[SamplingConfig, trtllm.SamplingConfig] - - -def _world_config_to_mapping(world_config: WorldConfig): - return Mapping(world_size=world_config.size, - rank=world_config.rank, - gpus_per_node=world_config.gpus_per_node, - tp_size=world_config.tensor_parallelism, - pp_size=world_config.pipeline_parallelism, - cp_size=world_config.context_parallelism) - - -class ModelRunnerCpp(ModelRunnerMixin): - """ - An interface class that wraps Executor and provides generation methods. - """ - - def __init__(self, - executor: trtllm.Executor, - max_batch_size: int, - max_input_len: int, - max_seq_len: int, - max_beam_width: int, - model_config: ModelConfig, - world_config: WorldConfig, - use_kv_cache: bool, - lora_manager: Optional[LoraManager] = None) -> None: - emit_engine_arch_deprecation("ModelRunnerCpp") - self.session = executor - self.max_batch_size = max_batch_size - self.max_input_len = max_input_len - self.max_seq_len = max_seq_len - self.max_beam_width = max_beam_width - self.model_config = model_config - self.mapping = _world_config_to_mapping(world_config) - self.world_config = world_config - self.use_kv_cache = use_kv_cache - self.lora_manager = lora_manager - - @classmethod - def from_dir( - cls, - engine_dir: str, - *, - lora_dir: Optional[str] = None, - rank: int = 0, - max_batch_size: Optional[int] = None, - max_input_len: Optional[int] = None, - max_output_len: Optional[int] = None, - max_beam_width: Optional[int] = None, - max_attention_window_size: Optional[list[int]] = None, - sink_token_length: Optional[int] = None, - kv_cache_free_gpu_memory_fraction: Optional[float] = None, - cross_kv_cache_fraction: Optional[float] = None, - medusa_choices: list[list[int]] | None = None, - eagle_choices: list[list[int]] | None = None, - eagle_posterior_threshold: float | None = None, - eagle_use_dynamic_tree: bool = False, - eagle_dynamic_tree_max_top_k: Optional[int] = None, - lookahead_config: list[int] | None = None, - debug_mode: bool = False, - lora_ckpt_source: str = "hf", - use_gpu_direct_storage: bool = False, - gpu_weights_percent: float = 1, - max_tokens_in_paged_kv_cache: int | None = None, - kv_cache_enable_block_reuse: bool = False, - enable_chunked_context: bool = False, - is_enc_dec: bool = False, - multi_block_mode: bool = True, - enable_context_fmha_fp32_acc: Optional[bool] = None, - cuda_graph_mode: Optional[bool] = None, - logits_processor_map: Optional[Dict[str, LogitsProcessor]] = None, - device_ids: List[int] | None = None, - is_orchestrator_mode: bool = False, - use_runtime_defaults: bool = True, - gather_generation_logits: bool = False, - use_variable_beam_width_search: bool = False, - mm_embedding_offloading: bool = False, - fail_fast_on_attention_window_too_large: bool = False, - normalize_log_probs: bool = False, - ) -> 'ModelRunnerCpp': - """ - Create a ModelRunnerCpp instance from an engine directory. - - Args: - engine_dir (str): - The directory that contains the serialized engine files and config files. - lora_dir (str): - The directory that contains LoRA weights. - rank (int): - The runtime rank id. - max_batch_size (int): - The runtime batch size limit. If max_batch_size is not None, it should not - be larger than the engine's max_batch_size; otherwise, the engine's max_batch_size - will be used. - max_input_len (int): - The runtime input length limit. If max_input_len is not None, it should not - be larger than the engine's max_input_len; otherwise, the engine's max_input_len - will be used. - max_output_len (int): - The runtime output length limit. If max_output_len is not None, it should not - be larger than the engine's max_output_len; otherwise, the engine's max_output_len - will be used. - max_beam_width (int): - The runtime beam width limit. If max_beam_width is not None, it should not - be larger than the engine's max_beam_width; otherwise, the engine's max_beam_width - will be used. - max_attention_window_size (List[int]): - The attention window size that controls the sliding window attention / cyclic kv cache behavior. - sink_token_length (int) : - The sink token length, default=0. - kv_cache_free_gpu_memory_fraction (float) : - Free GPU memory fraction that KV cache used. - cross_kv_cache_fraction (float) : - KV Cache fraction reserved for cross attention, should only be used with enc-dec models. - debug_mode (bool): - Whether or not to turn on the debug mode. - medusa_choices (List[List[int]]): - Medusa choices to use when in Medusa decoding. - eagle_choices (List[List[int]]): - Eagle choices to use when in Eagle-1 decoding. - eagle_posterior_threshold float: - Minimum token probability threshold for typical acceptance. - Value different from None enables typical acceptance in Eagle. - eagle_use_dynamic_tree bool: - Whether to use Eagle-2, which is dynamic tree. - eagle_dynamic_tree_max_top_k int: - The maximum number of draft tokens to expand for each node in Eagle-2. - lora_ckpt_source (str): - Source of checkpoint. Should be one of ['hf', 'nemo']. - max_tokens_in_paged_kv_cache (int): - Maximum amount of tokens configured in kv cache. - kv_cache_enable_block_reuse (bool): - Enables block reuse in kv cache. - enable_chunked_context (bool): - Enables chunked context. - is_enc_dec (bool): - Whether the model is encoder-decoder architecture. - multi_block_mode (bool): - Whether to distribute the work across multiple CUDA thread-blocks on the GPU for masked MHA kernel. - enable_context_fmha_fp32_acc (bool): - Enable FMHA runner FP32 accumulation. - cuda_graph_mode (bool): - Whether to use cuda graph for inference. - logits_processor_map (Dict[str, LogitsProcessor]) - A map of logits processor functions indexed by names. A name can be provided later to - the generate() function to specify which logits processor to run. - device_ids (List[int]): - Device indices to run the Executor on. - is_orchestrator_mode (bool): - The mode to run the model-runner, Leader mode by default. - gather_generation_logits (bool): - Enable gathering generation logits. - fail_fast_on_attention_window_too_large (bool): - Whether to fail fast if the attention window(s) are too large to fit even a single sequence in the KVCache. - Returns: - ModelRunnerCpp: An instance of ModelRunnerCpp. - """ - extended_runtime_perf_knob_config = trtllm.ExtendedRuntimePerfKnobConfig( - ) - if multi_block_mode is not None: - extended_runtime_perf_knob_config.multi_block_mode = multi_block_mode - if enable_context_fmha_fp32_acc is not None: - extended_runtime_perf_knob_config.enable_context_fmha_fp32_acc = enable_context_fmha_fp32_acc - if cuda_graph_mode is not None: - extended_runtime_perf_knob_config.cuda_graph_mode = cuda_graph_mode - - model_config = None - is_multimodal = {'vision', 'llm'}.issubset({ - name - for name in os.listdir(engine_dir) - if os.path.isdir(os.path.join(engine_dir, name)) - }) - encoder_path = None - decoder_path = None - - if is_enc_dec: - if is_multimodal: - encoder_path = join(engine_dir, 'vision') - decoder_path = join(engine_dir, 'llm') - else: - encoder_path = join(engine_dir, 'encoder') - decoder_path = join(engine_dir, 'decoder') - - encoder_config_path = Path(encoder_path) / "config.json" - encoder_json_config = GptJsonConfig.parse_file(encoder_config_path) - encoder_model_config = encoder_json_config.model_config - decoder_config_path = Path(decoder_path) / "config.json" - decoder_json_config = GptJsonConfig.parse_file(decoder_config_path) - decoder_model_config = decoder_json_config.model_config - - json_config = decoder_json_config - model_config = decoder_model_config - engine_dir = decoder_path - - if max_input_len is None and not is_multimodal: - max_input_len = encoder_model_config.max_input_len - else: - config_path = Path(engine_dir) / "config.json" - json_config = GptJsonConfig.parse_file(config_path) - model_config = json_config.model_config - - use_kv_cache = KVCacheType.from_cpp( - model_config.kv_cache_type) != KVCacheType.DISABLED - if not model_config.use_cross_attention: - assert cross_kv_cache_fraction is None, "cross_kv_cache_fraction should only be used with enc-dec models." - - if not use_kv_cache: - assert max_output_len == 1 or max_output_len is None, 'Disabled KV cache is intended for context phase only now.' - - # Note: Parallel configuration will be fetched automatically from trtllm.Executor constructor - # by inspecting the json file. These lines serve the purpose of serving vocab_size_padded and - # num_layers properties. - # MPI world size must be 1 in Orchestrator mode - if is_orchestrator_mode: - tp_size = 1 - pp_size = 1 - cp_size = 1 - # Check the count of devices equal to tp_size of engine - # assert len(device_ids) == json_config.tensor_parallelism - else: - tp_size = json_config.tensor_parallelism - pp_size = json_config.pipeline_parallelism - cp_size = json_config.context_parallelism - gpus_per_node = json_config.gpus_per_node - world_config = WorldConfig.mpi(tensor_parallelism=tp_size, - pipeline_parallelism=pp_size, - context_parallelism=cp_size, - gpus_per_node=gpus_per_node) - assert rank == world_config.rank - - engine_config = EngineConfig.from_json_file(f"{engine_dir}/config.json") - if model_config.use_lora_plugin and rank == 0: - mapping = _world_config_to_mapping(world_config) - lora_manager = LoraManager( - mapping=mapping, - model_config=ModelConfigPython.from_model_config_cpp( - model_config)) - if lora_dir is None: - config_lora_dir = engine_config.build_config.lora_config.lora_dir - if len(config_lora_dir) > 0: - lora_dir = [ - f"{engine_dir}/{dir}" for dir in config_lora_dir - ] - lora_ckpt_source = engine_config.build_config.lora_config.lora_ckpt_source - - if lora_dir is not None: - runtime_model_config = _engine_config_to_model_config( - engine_config, gpu_weights_percent=gpu_weights_percent) - # For Executor, only rank 0 can enqueue requests, and should hold all lora weights - lora_manager.load_from_ckpt(lora_dir, - model_config=runtime_model_config, - ckpt_source=lora_ckpt_source) - else: - raise RuntimeError( - f"LoRA weights are unspecified and also unavailable in the engine_dir ({engine_dir})." - ) - - max_lora_rank = engine_config.build_config.lora_config.max_lora_rank - num_lora_modules = engine_config.pretrained_config.num_hidden_layers * \ - len(lora_manager.lora_target_modules + lora_manager.missing_qkv_modules) - num_lora_adapters = min(lora_manager.num_lora_adapters, 8) - peft_cache_config = trtllm.PeftCacheConfig( - num_device_module_layer=max_lora_rank * num_lora_modules * - num_lora_adapters, - num_host_module_layer=max_lora_rank * num_lora_modules * - num_lora_adapters, - ) - else: - lora_manager = None - peft_cache_config = trtllm.PeftCacheConfig() - - if world_config.size > 1: - peft_cache_config = mpi_broadcast(peft_cache_config, 0) - - profiler.start('load tensorrt_llm engine') - - kv_cache_config = trtllm.KvCacheConfig( - free_gpu_memory_fraction=kv_cache_free_gpu_memory_fraction, - max_attention_window=max_attention_window_size, - sink_token_length=sink_token_length, - max_tokens=max_tokens_in_paged_kv_cache, - enable_block_reuse=kv_cache_enable_block_reuse, - cross_kv_cache_fraction=cross_kv_cache_fraction, - runtime_defaults=json_config.runtime_defaults - if use_runtime_defaults else None, - ) - - decoding_config = trtllm.DecodingConfig() - if medusa_choices is not None: - decoding_config.medusa_choices = medusa_choices - if multi_block_mode is not None: - multi_block_mode = False # Medusa doesn't support multi-block mode. - - if eagle_choices is not None or eagle_posterior_threshold is not None or eagle_use_dynamic_tree: - greedy_sampling = eagle_posterior_threshold is None - decoding_config.eagle_config = trtllm.EagleConfig( - eagle_choices, greedy_sampling, eagle_posterior_threshold, - eagle_use_dynamic_tree, eagle_dynamic_tree_max_top_k) - if multi_block_mode is not None: - logger.warning( - f'Multi block mode is not supported for EAGLE. Disabling it.' - ) - multi_block_mode = False # Eagle doesn't support multi-block mode. - - if lookahead_config is not None: - [w, n, g] = lookahead_config - decoding_config.lookahead_decoding_config = trtllm.LookaheadDecodingConfig( - w, n, g) - - if use_variable_beam_width_search: - decoding_config.decoding_mode = DecodingMode.BeamSearch( - ).useVariableBeamWidthSearch(True) - - if max_batch_size is None: - max_batch_size = model_config.max_batch_size - else: - assert max_batch_size <= model_config.max_batch_size - if max_input_len is None: - max_input_len = model_config.max_input_len - # NOTE: remove assertion here for temp fix, - # model_config.max_input_len is not the upper bound of input length. - # If runtime max_input_len is not properly set, - # C++ runtime will throw an error when fetching new requests - if max_output_len is None or is_enc_dec: - max_seq_len = model_config.max_seq_len - else: - max_seq_len = max_input_len + max_output_len - assert max_seq_len <= model_config.max_seq_len - if max_beam_width is None: - max_beam_width = model_config.max_beam_width - else: - assert max_beam_width <= model_config.max_beam_width - - debug_config = None - if debug_mode: - # To debug specific tensors, add tensor names in the following list - # if none provided, all input and output tensors will be dumped - # if not none, it will disable all input/output dump - debug_tensor_names: List[str] = [ - ] # modify this list for specific tensor dump - debug_config = trtllm.DebugConfig( - debug_input_tensors=True, - debug_output_tensors=True, - debug_tensor_names=debug_tensor_names) - - trtllm_config = trtllm.ExecutorConfig( - max_batch_size=max_batch_size, - max_beam_width=max_beam_width, - kv_cache_config=kv_cache_config, - decoding_config=decoding_config, - peft_cache_config=peft_cache_config, - debug_config=debug_config, - use_gpu_direct_storage=use_gpu_direct_storage, - gpu_weights_percent=gpu_weights_percent, - gather_generation_logits=gather_generation_logits, - normalize_log_probs=normalize_log_probs, - ) - trtllm_config.enable_chunked_context = enable_chunked_context - trtllm_config.extended_runtime_perf_knob_config = extended_runtime_perf_knob_config - trtllm_config.mm_embedding_offloading = mm_embedding_offloading - trtllm_config.fail_fast_on_attention_window_too_large = fail_fast_on_attention_window_too_large - if is_orchestrator_mode: - communication_mode = trtllm.CommunicationMode.ORCHESTRATOR - path = str(Path(__file__).parent.parent / 'bin' / 'executorWorker') - orchestrator_config = OrchestratorConfig(True, path) - else: - communication_mode = trtllm.CommunicationMode.LEADER - orchestrator_config = None - - trtllm_config.parallel_config = ParallelConfig( - trtllm.CommunicationType.MPI, - communication_mode, - device_ids=device_ids, - orchestrator_config=orchestrator_config) - - # LogitsPostProcessor in Orchestrator mode is not supported yet. - if not is_orchestrator_mode: - logits_proc_config = trtllm.LogitsPostProcessorConfig() - if logits_processor_map is not None: - logits_proc_config.processor_map = logits_processor_map - trtllm_config.logits_post_processor_config = logits_proc_config - - if is_enc_dec: - executor = trtllm.Executor(encoder_path, decoder_path, - trtllm.ModelType.ENCODER_DECODER, - trtllm_config) - else: - executor = trtllm.Executor(Path(engine_dir), - trtllm.ModelType.DECODER_ONLY, - trtllm_config) - - profiler.stop('load tensorrt_llm engine') - - loading_time = profiler.elapsed_time_in_sec("load tensorrt_llm engine") - logger.info(f'Load engine takes: {loading_time} sec') - - return cls(executor, - max_batch_size=max_batch_size, - max_input_len=max_input_len, - max_seq_len=max_seq_len, - max_beam_width=max_beam_width, - model_config=model_config, - world_config=world_config, - use_kv_cache=use_kv_cache, - lora_manager=lora_manager) - - def _check_inputs(self, batch_input_ids: List[List[int]], - encoder_input_ids: Optional[List[List[int]]], - sampling_config: trtllm.SamplingConfig, max_new_tokens): - batch_size = len(encoder_input_ids) if encoder_input_ids else len( - batch_input_ids) - if batch_size > self.max_batch_size: - raise RuntimeError( - f"Input batch size ({batch_size}) exceeds the engine or specified limit ({self.max_batch_size})" - ) - input_lengths = [ - len(x) for x in encoder_input_ids - ] if encoder_input_ids else [len(x) for x in batch_input_ids] - max_length = max(input_lengths) - if max_length > self.max_input_len: - raise RuntimeError( - f"Maximum input length ({max_length}) exceeds the engine or specified limit ({self.max_input_len})" - ) - if encoder_input_ids: - decoder_max_length = max([len(x) for x in batch_input_ids]) - if decoder_max_length + max_new_tokens > self.max_seq_len: - raise RuntimeError( - f"Decoder prefix tokens ({decoder_max_length}) + maximum new tokens ({max_new_tokens}) exceeds the engine or specified limit ({self.max_seq_len})" - ) - else: - if max_length + max_new_tokens > self.max_seq_len: - raise RuntimeError( - f"Maximum input length ({max_length}) + maximum new tokens ({max_new_tokens}) exceeds the engine or specified limit ({self.max_seq_len})" - ) - if sampling_config.beam_width > self.max_beam_width: - raise RuntimeError( - f"Num beams ({sampling_config.beam_width}) exceeds the engine or specified limit ({self.max_beam_width})" - ) - - @property - def dtype(self) -> torch.dtype: - bindings_dtype = self.model_config.data_type - return _bindings_dtype_to_torch_dtype_dict[bindings_dtype] - - @property - def vocab_size(self) -> int: - return self.model_config.vocab_size - - @property - def vocab_size_padded(self) -> int: - return self.model_config.vocab_size_padded(self.world_config.size) - - @property - def hidden_size(self) -> int: - return self.model_config.hidden_size - - @property - def num_heads(self) -> int: - return self.model_config.num_heads - - @property - def num_layers(self) -> int: - return self.model_config.num_layers( - self.world_config.pipeline_parallelism, - self.world_config.pipeline_parallel_rank, - ) - - @property - def max_sequence_length(self) -> int: - return self.max_seq_len - - @property - def remove_input_padding(self) -> bool: - return self.model_config.use_packed_input - - @property - def max_prompt_embedding_table_size(self) -> int: - return self.model_config.max_prompt_embedding_table_size - - @property - def gather_context_logits(self) -> bool: - return self.model_config.compute_context_logits - - @property - def gather_generation_logits(self) -> bool: - return self.model_config.compute_generation_logits - - def generate( - self, - batch_input_ids: List[torch.Tensor], - *, - position_ids: List[torch.Tensor] = None, - encoder_input_ids: List[torch.Tensor] = None, - encoder_input_features: List[ - torch.Tensor] = None, # TODO: add to doc string - encoder_output_lengths: List[int] = None, - cross_attention_masks: List[ - torch.Tensor] = None, # TODO: add to doc string - mrope_params: Optional[MropeParams] = None, - sampling_config: Optional[SamplingConfig] = None, - lora_uids: Optional[list] = None, - lookahead_config: list[int] | None = None, - streaming: bool = False, - stopping_criteria: Optional[StoppingCriteria] = None, - logits_processor_names: list[str] | None = None, - max_new_tokens: int = 1, - end_id: int | None = None, - pad_id: int | None = None, - bad_words_list: list[list[int]] | None = None, - stop_words_list: list[list[int]] | None = None, - return_dict: bool = False, - output_sequence_lengths: bool = False, - output_generation_logits: bool = False, - output_log_probs: bool = False, - output_cum_log_probs: bool = False, - prompt_table: Optional[Union[str, torch.Tensor]] = None, - prompt_tasks: Optional[str] = None, - input_token_extra_ids: List[List[int]] = None, - return_all_generated_tokens: bool = False, - language_adapter_uids: Optional[List[int]] = None, - mm_embedding_offloading: bool = False, - **kwargs) -> Union[torch.Tensor, dict]: - """ - Generates sequences of token ids. - The generation-controlling parameters are set in the sampling_config; it will be set to a default one if not passed. - You can override any sampling_config's attributes by passing corresponding parameters. - - Args: - batch_input_ids (List[torch.Tensor]): - A list of input id tensors. Each tensor is of shape (sequence_length, ). - position_ids (List[torch.Tensor]): - A list of position id tensors. Each tensor is of shape (sequence_length, ). - encoder_input_ids (List[torch.Tensor]): - A list of encoder input id tensors for encoder-decoder models (optional). Each tensor is of shape (sequence_length, ). - encoder_input_features: (List[torch.Tensor]): - A list of encoder input feature tensors for multimodal encoder-decoder models (optional). Each tensor is of shape (sequence_length, feature_dim). - encoder_output_lengths: (List[int]): - A list of encoder output lengths (optional) if encoder output has different length from encoder input (due to convolution down-sampling, etc.) - sampling_config (SamplingConfig): - The sampling configuration to be used as base parametrization for the generation call. - The passed **kwargs matching the sampling_config's attributes will override them. - If the sampling_config is not provided, a default will be used. - prompt_table (str or torch.Tensor): - The file path of prompt table (.npy format, exported by nemo_prompt_convert.py) or the prompt table itself. - prompt_tasks (str): - The prompt tuning task ids for the input batch, in format of comma-separated list (e.g., 0,3,1,0). - input_token_extra_ids (List[List[int]]): - Input token extra ids for using p-tuning and KV Cache reuse together - lora_uids (list): - The uids of LoRA weights for the input batch. Use -1 to disable the LoRA module. - streaming (bool): - Whether or not to use streaming mode for generation. - stopping_criteria (StoppingCriteria): - Custom stopping criteria. - logits_processor_names (List[str]): - Custom logits processor names. - return_all_generated_tokens (bool): - Whether the full output is returned at each streaming step - kwargs (Dict[str, Any]: - Ad hoc parametrization of sampling_config. - The passed **kwargs matching the sampling_config's attributes will override them. - Returns: - torch.Tensor or dict: - If return_dict=False, the method returns generated output_ids. - If return_dict=True, the method returns a dict of output_ids, - sequence_lengths (if sampling_config.output_sequence_lengths=True), - context_logits and generation_logits (if self.gather_context_logits=True and - self.gather_generation_logits=True, respectively). - """ - # TODO: Check if these can be supported now and support them - if stopping_criteria is not None: - raise RuntimeError( - "Stopping criteria is not supported in C++ session.") - - if not self.use_kv_cache and max_new_tokens > 1: - raise RuntimeError( - 'Disabled KV cache is intended for context phase only now.') - - # If we are in a multi-gpu scenario, only rank 0 continues - if not self.session.can_enqueue_requests(): - return [] - - # Convert tensor input to plain lists - batch_input_ids_list = [a.tolist() for a in batch_input_ids] - encoder_input_ids_list = [a.tolist() for a in encoder_input_ids - ] if encoder_input_ids else None - - if sampling_config is None: - # Convert from old API of SamplingConfig - # Note: Due to a Python3.10 bug one cannot use inspect on it currently - accepted_parameters = [ - "num_beams", - "top_k", - "top_p", - "top_p_min", - "top_p_reset_ids", - "top_p_decay", - "temperature", - "min_tokens", - "beam_search_diversity_rate", - "repetition_penalty", - "presence_penalty", - "frequency_penalty", - "prompt_ignore_length", - "length_penalty", - "early_stopping", - "no_repeat_ngram_size", - "random_seed", - "num_return_sequences", - "min_p", - "beam_width_array", - ] - rename_params = {"num_beams": "beam_width", "random_seed": "seed"} - sampling_params = { - k: v - for k, v in kwargs.items() if k in accepted_parameters - } - for k, v in rename_params.items(): - if k in sampling_params: - sampling_params[v] = sampling_params.pop(k) - if "top_p" in sampling_params and sampling_params["top_p"] == 0.0: - sampling_params["top_p"] = None - - # TODO: improve usage of SamplingConfig. For example, - # construct SamplingConfig for each request, rather than one for the whole batch. - # Here we use beam width array for each request for Variable-Beam-Width-Search. - batch_size = len(batch_input_ids) - use_sampling_config_for_each_request = False - # Just placeholder for non-Variable-Beam-Width-Search - sampling_config_list = [None] * batch_size - if "beam_width_array" in sampling_params and sampling_params[ - "beam_width_array"] is not None and len( - sampling_params["beam_width_array"]) == batch_size: - use_sampling_config_for_each_request = True - sp_copy = copy.deepcopy(sampling_params) - for i in range(batch_size): - bwa = sampling_params["beam_width_array"][i] - sp_copy["beam_width_array"] = bwa - sp_copy["beam_width"] = max(bwa) - sampling_config_list[i] = trtllm.SamplingConfig(**sp_copy) - # Just placeholder for Variable-Beam-Width-Search and for `self._check_inputs` - max_beam_width = max(sc.beam_width - for sc in sampling_config_list) - sampling_params["beam_width"] = max_beam_width - sampling_params["beam_width_array"] = [max_beam_width] * 8 - sampling_config = trtllm.SamplingConfig(**sampling_params) - else: - sampling_config = copy.deepcopy(sampling_config) - - self._check_inputs(batch_input_ids_list, encoder_input_ids_list, - sampling_config, max_new_tokens) - - output_config = trtllm.OutputConfig( - return_context_logits=self.gather_context_logits, - return_generation_logits=self.gather_generation_logits - or output_generation_logits, - return_log_probs=output_log_probs, - ) - - prompt_tuning_configs = self._prepare_ptuning_executor( - batch_input_ids_list, - prompt_table, - prompt_tasks, - input_token_extra_ids, - mm_embedding_offloading=mm_embedding_offloading) - mrope_configs = self._prepare_mrope_executor(batch_input_ids_list, - mrope_params) - - stop_words_list = self._prepare_words_list(stop_words_list, - len(batch_input_ids_list)) - bad_words_list = self._prepare_words_list(bad_words_list, - len(batch_input_ids_list)) - logits_processor_names = self._prepare_names_list( - logits_processor_names, len(batch_input_ids_list)) - - lora_configs = self._prepare_lora_configs(lora_uids, - len(batch_input_ids_list)) - request_lookahead_config = None - if lookahead_config is not None: - [w, n, g] = lookahead_config - request_lookahead_config = trtllm.LookaheadDecodingConfig(w, n, g) - skip_cross_attn_blocks = kwargs.get('skip_cross_attn_blocks', None) - - # Draft-Target-Model speculative decoding - if "draft_tokens_list" in kwargs.keys() and kwargs[ - "draft_tokens_list"] is not None and "draft_logits_list" in kwargs.keys( - ) and kwargs["draft_logits_list"] is not None: - # Use logits to accept - external_draft_tokens_configs = [ - ExternalDraftTokensConfig(draft_tokens, draft_logits) - for draft_tokens, draft_logits in zip( - kwargs["draft_tokens_list"], kwargs["draft_logits_list"]) - ] - is_draft_target_model = True - elif "draft_tokens_list" in kwargs.keys( - ) and kwargs["draft_tokens_list"] is not None: - # Use tokens to accept - external_draft_tokens_configs = [ - ExternalDraftTokensConfig(draft_tokens) - for draft_tokens in kwargs["draft_tokens_list"] - ] - is_draft_target_model = True - else: - external_draft_tokens_configs = [None] * len(batch_input_ids_list) - is_draft_target_model = False - - if language_adapter_uids is None: - language_adapter_uids = [None] * len(batch_input_ids_list) - - requests = [ - trtllm.Request( - input_token_ids=input_ids, - encoder_input_token_ids=encoder_input_ids_list[i] - if encoder_input_ids is not None else None, - encoder_output_length=encoder_output_lengths[i] - if encoder_output_lengths is not None else None, - encoder_input_features=encoder_input_features[i].contiguous() - if encoder_input_features is not None else None, - position_ids=position_ids[i].tolist() - if position_ids is not None else None, - cross_attention_mask=cross_attention_masks[i].contiguous() if - (cross_attention_masks is not None - and cross_attention_masks[i] is not None) else None, - max_tokens=max_new_tokens, - pad_id=pad_id, - end_id=end_id, - stop_words=stop_words, - bad_words=bad_words, - sampling_config=(sampling_config_each_request - if use_sampling_config_for_each_request else - sampling_config), - lookahead_config=request_lookahead_config, - streaming=streaming, - output_config=output_config, - prompt_tuning_config=prompt_tuning_config, - mrope_config=mrope_config, - lora_config=lora_config, - return_all_generated_tokens=return_all_generated_tokens, - logits_post_processor_name=logits_post_processor_name, - external_draft_tokens_config=external_draft_tokens_config, - skip_cross_attn_blocks=skip_cross_attn_blocks, - language_adapter_uid=language_adapter_uid, - ) for i, - (input_ids, stop_words, bad_words, prompt_tuning_config, - mrope_config, lora_config, logits_post_processor_name, - external_draft_tokens_config, language_adapter_uid, - sampling_config_each_request) in enumerate( - zip(batch_input_ids_list, stop_words_list, bad_words_list, - prompt_tuning_configs, mrope_configs, lora_configs, - logits_processor_names, external_draft_tokens_configs, - language_adapter_uids, sampling_config_list)) - ] - - request_ids = self.session.enqueue_requests(requests) - if not streaming: - return self._initialize_and_fill_output( - request_ids=request_ids, - end_id=end_id, - return_dict=return_dict, - output_sequence_lengths=output_sequence_lengths, - output_generation_logits=output_generation_logits, - output_log_probs=output_log_probs, - output_cum_log_probs=output_cum_log_probs, - batch_input_ids=batch_input_ids, - streaming=streaming, - sampling_config=sampling_config, - is_draft_target_model=is_draft_target_model, - ) - else: - return self._stream( - request_ids=request_ids, - end_id=end_id, - return_dict=return_dict, - output_sequence_lengths=output_sequence_lengths, - output_generation_logits=output_generation_logits, - output_log_probs=output_log_probs, - output_cum_log_probs=output_cum_log_probs, - batch_input_ids=batch_input_ids, - batch_input_ids_list=batch_input_ids_list, - streaming=streaming, - return_all_generated_tokens=return_all_generated_tokens, - sampling_config=sampling_config, - is_draft_target_model=is_draft_target_model, - ) - - def _prepare_words_list(self, words_list: List[List[List[int]]], - batch_size: int): - if words_list is None: - return [None] * batch_size - return words_list - - def _prepare_names_list(self, names_list: List[str], batch_size: int): - if names_list is None: - return [None] * batch_size - return names_list - - def _prepare_ptuning_executor(self, batch_input_ids_list, prompt_table, - prompt_tasks, input_token_extra_ids, - mm_embedding_offloading): - if input_token_extra_ids: - assert len(batch_input_ids_list) == len(input_token_extra_ids), \ - f"Batch size of input_token_extra_ids ({len(input_token_extra_ids)}) must be the same as input batch size ({len(batch_input_ids_list)})" - prompt_tuning_configs = len(batch_input_ids_list) * [None] - if prompt_table is not None: - if mm_embedding_offloading: - # CUDA Stream Overlapping Requirements: - # 1. Both memory copy stream and kernel execution stream must be non-default streams - # 2. For host<->device transfers (H2D/D2H), host memory MUST be page-locked (pinned) - # NOTE: pinning is skipped under Confidential Compute - # (see maybe_pin_memory() and prefer_pinned()) - prompt_table_data = maybe_pin_memory( - self._prepare_embedding_table(prompt_table)) - else: - prompt_table_data = self._prepare_embedding_table( - prompt_table).cuda() - if prompt_tasks is not None: - task_indices = [int(t) for t in prompt_tasks.split(',')] - assert len(task_indices) == len(batch_input_ids_list), \ - f"Number of supplied tasks ({len(task_indices)}) must match input batch size ({len(batch_input_ids_list)})" - prompt_tuning_configs = [ - trtllm.PromptTuningConfig( - embedding_table=prompt_table_data[task_indices[i]], - input_token_extra_ids=input_token_extra_ids[i] - if input_token_extra_ids else None) - for i in range(len(batch_input_ids_list)) - ] - else: - prompt_tuning_configs = [ - trtllm.PromptTuningConfig( - embedding_table=prompt_table_data[0], - input_token_extra_ids=input_token_extra_ids[i] - if input_token_extra_ids else None) - for i in range(len(batch_input_ids_list)) - ] - return prompt_tuning_configs - - # TODO: add multimodal input for TRT engine backend - - def _prepare_mrope_executor(self, batch_input_ids_list, mrope: MropeParams): - mrope_configs = len(batch_input_ids_list) * [None] - if mrope != None: - mrope_rotary_cos_sin = mrope.mrope_rotary_cos_sin - assert isinstance( - mrope_rotary_cos_sin, - torch.Tensor), "mrope_rotary_cos_sin should be torch.Tensor" - mrope_rotary_cos_sin_data = mrope_rotary_cos_sin.to( - dtype=torch.float32) - - mrope_position_deltas = mrope.mrope_position_deltas - assert isinstance( - mrope_position_deltas, - torch.Tensor), "mrope_position_deltas should be torch.Tensor" - mrope_position_deltas_data = mrope_position_deltas.to( - dtype=torch.int32) - - mrope_configs = [ - trtllm.MropeConfig( - mrope_rotary_cos_sin=mrope_rotary_cos_sin_data[i], - mrope_position_deltas=mrope_position_deltas_data[i]) - for i in range(len(batch_input_ids_list)) - ] - return mrope_configs - - def _prepare_lora_configs(self, lora_uids, batch_size): - if lora_uids is None: - return [None] * batch_size - assert len(lora_uids) == batch_size - return [ - trtllm.LoraConfig(task_id=int(uid), - weights=self.lora_manager.cpp_lora_weights[uid], - config=self.lora_manager.cpp_lora_config[uid]) - if int(uid) >= 0 else None for uid in lora_uids - ] - - def _get_num_sequences(self, sampling_config: SamplingConfigType): - num_beams = sampling_config.num_beams if isinstance( - sampling_config, SamplingConfig) else sampling_config.beam_width - num_sequences = sampling_config.num_return_sequences or num_beams - assert num_beams == 1 or num_sequences <= num_beams - return num_sequences - - def _initialize_and_fill_output( - self, - *, - request_ids, - end_id, - return_dict, - output_sequence_lengths, - output_generation_logits, - output_log_probs, - output_cum_log_probs, - batch_input_ids, - streaming, - sampling_config: SamplingConfigType, - is_draft_target_model: bool = False, - ): - num_sequences = self._get_num_sequences(sampling_config) - # (batch_size, num_sequences, sequence_len) - output_ids = [[[] for _ in range(num_sequences)] - for _ in range(len(request_ids))] - - all_responses = [] - finished_request_ids = set() - while finished_request_ids != set(request_ids): - responses = self.session.await_responses() - for response in responses: - if response.result.is_final: - finished_request_ids.add(response.request_id) - all_responses.extend(responses) - - return self._fill_output( - responses=all_responses, - output_ids=output_ids, - end_id=end_id, - return_dict=return_dict, - output_sequence_lengths=output_sequence_lengths, - output_generation_logits=output_generation_logits, - output_log_probs=output_log_probs, - output_cum_log_probs=output_cum_log_probs, - batch_input_ids=batch_input_ids, - batch_input_ids_list=[], - streaming=streaming, - request_ids=request_ids, - return_all_generated_tokens=False, - sampling_config=sampling_config, - is_draft_target_model=is_draft_target_model, - ) - - def _stream( - self, - *, - request_ids, - end_id, - return_dict, - output_sequence_lengths, - output_generation_logits, - output_log_probs, - output_cum_log_probs, - batch_input_ids, - batch_input_ids_list, - streaming, - return_all_generated_tokens, - sampling_config: SamplingConfigType, - is_draft_target_model: bool = False, - ): - num_sequences = self._get_num_sequences(sampling_config) - # (batch_size, num_sequences, sequence_len) - output_ids = [[ - copy.deepcopy(batch_input_ids_list[batch_idx]) - for _ in range(num_sequences) - ] for batch_idx in range(len(request_ids))] - - finished_request_ids = set() - while finished_request_ids != set(request_ids): - responses = self.session.await_responses() - for response in responses: - if response.result.is_final: - finished_request_ids.add(response.request_id) - - yield self._fill_output( - responses=responses, - output_ids=output_ids, - end_id=end_id, - return_dict=return_dict, - output_sequence_lengths=output_sequence_lengths, - output_generation_logits=output_generation_logits, - output_log_probs=output_log_probs, - output_cum_log_probs=output_cum_log_probs, - batch_input_ids=batch_input_ids, - batch_input_ids_list=batch_input_ids_list, - streaming=streaming, - request_ids=request_ids, - return_all_generated_tokens=return_all_generated_tokens, - sampling_config=sampling_config, - is_draft_target_model=is_draft_target_model, - ) - - def _fill_output( - self, - *, - responses, - output_ids, - end_id, - return_dict, - output_sequence_lengths, - output_generation_logits, - output_log_probs, - output_cum_log_probs, - batch_input_ids, - batch_input_ids_list, - streaming, - request_ids, - return_all_generated_tokens, - sampling_config: SamplingConfigType, - is_draft_target_model: bool, - ): - cuda_device = torch.device("cuda") - - batch_size = len(batch_input_ids) - num_sequences = len(output_ids[0]) - beam_width = getattr(sampling_config, 'num_beams', - getattr(sampling_config, 'beam_width')) - is_beam_search = beam_width > 1 - - def fill_output_ids(result_token_ids, batch_idx, seq_idx): - # Return shape = (batch_size, num_sequences, seq_len) - if return_all_generated_tokens: - output_ids[batch_idx][seq_idx] = ( - batch_input_ids_list[batch_idx] + result_token_ids) - else: - output_ids[batch_idx][seq_idx] += result_token_ids - - for response in responses: - if response.has_error(): - raise RuntimeError(response.error_msg) - - result = response.result - batch_idx = request_ids.index(response.request_id) - if is_beam_search: - for beam, output_tokens in enumerate(result.output_token_ids): - fill_output_ids(output_tokens, batch_idx, beam) - else: - fill_output_ids(result.output_token_ids[0], batch_idx, - result.sequence_index) - - if output_sequence_lengths: - sequence_lengths = [[len(token_ids) for token_ids in beams] - for beams in output_ids] - - if streaming: - output_ids = copy.deepcopy(output_ids) - - # Pad by end_id tokens (batch, num_sequences, max_seq_len). - for beams in output_ids: - for token_ids in beams: - token_ids += [end_id] * (self.max_seq_len - len(token_ids)) - output_ids = torch.tensor(output_ids, - dtype=torch.int32, - device=cuda_device) - - if return_dict: - outputs = {'output_ids': output_ids} - - input_lengths = torch.tensor([x.size(0) for x in batch_input_ids], - dtype=torch.int32, - device=cuda_device) - - if output_sequence_lengths: - outputs['sequence_lengths'] = torch.tensor(sequence_lengths, - dtype=torch.int32, - device=cuda_device) - - if self.gather_context_logits: - context_logits = None - max_input_len = input_lengths.max() - for response in responses: - result = response.result - logits = result.context_logits - if logits is None: - continue - input_len, vocab_size = logits.shape - if context_logits is None: - context_logits = torch.zeros( - (batch_size, max_input_len, vocab_size), - dtype=logits.dtype, - device=cuda_device) - if result.sequence_index == 0: - batch_idx = request_ids.index(response.request_id) - context_logits[batch_idx, :input_len, :] = logits - assert context_logits is not None - outputs['context_logits'] = context_logits - - if self.gather_generation_logits or output_generation_logits: - gen_logits = None - if is_draft_target_model: - # Put the outputs in a list rather than a tensor since their - # length may vary among requests in a batch - gen_logits = [ - a.result.generation_logits.cuda() for a in responses - if a.result.generation_logits is not None - ] - else: - # The shape of generation logits - # (num_sequences, seq_len, vocab_size) in non-streaming - # (seq_len, num_sequences, vocab_size) in streaming - seq_dim = 0 if streaming else 1 - max_out_len = max( - response.result.generation_logits.size(seq_dim) - for response in responses - if response.result.generation_logits is not None) - vocab_size = responses[0].result.generation_logits.size(-1) - if not streaming: - gen_shape = (num_sequences, max_out_len, vocab_size) - elif streaming and return_all_generated_tokens: - gen_shape = (max_out_len, num_sequences, vocab_size) - else: - # streaming and not return_all_generated_tokens - gen_shape = (1, num_sequences, vocab_size) - logits_dtype = responses[0].result.generation_logits.dtype - gen_logits = torch.zeros((batch_size, *gen_shape), - dtype=logits_dtype, - device=cuda_device) - - for response in responses: - logits = response.result.generation_logits - if logits is None: - continue - seq_len = logits.size(seq_dim) - - batch_idx = request_ids.index(response.request_id) - seq_idx = response.result.sequence_index - if streaming: - if is_beam_search: - # WAR: gen_logits contains all beams, clipping - # the first n beams as a postprocessing. - gen_logits[batch_idx, :seq_len, - ...] = logits[:, :num_sequences, :] - else: - gen_logits[batch_idx, :seq_len, seq_idx, - ...] = logits[:, 0, :] - else: - if is_beam_search: - gen_logits[batch_idx, :, :seq_len, ...] = logits - else: - gen_logits[batch_idx, seq_idx, :seq_len, - ...] = logits[0] - outputs['generation_logits'] = gen_logits - - if output_log_probs: - max_log_probs_len = max( - len(lprobs) for response in responses - for lprobs in response.result.log_probs) - log_probs = torch.zeros( - (batch_size, num_sequences, max_log_probs_len), - dtype=torch.float32) - for response in responses: - batch_idx = request_ids.index(response.request_id) - if is_beam_search: - for beam_idx, lprobs in enumerate( - response.result.log_probs): - log_probs[batch_idx, - beam_idx, :len(lprobs)] = torch.tensor( - lprobs) - else: - seq_idx = response.result.sequence_index - lprobs = response.result.log_probs[0] - log_probs[batch_idx, - seq_idx, :len(lprobs)] = torch.tensor(lprobs) - assert isinstance(log_probs, torch.Tensor) - outputs['log_probs'] = log_probs.to(cuda_device) - - if output_cum_log_probs: - cum_log_probs = torch.zeros((batch_size, num_sequences), - dtype=torch.float32) - for response in responses: - if response.result.cum_log_probs is None: - continue - batch_idx = request_ids.index(response.request_id) - clprobs = torch.tensor(response.result.cum_log_probs) - if is_beam_search: - cum_log_probs[batch_idx, :] = clprobs - else: - seq_idx = response.result.sequence_index - cum_log_probs[batch_idx, seq_idx] = clprobs - outputs['cum_log_probs'] = cum_log_probs.to(cuda_device) - - outputs = self._prepare_outputs(outputs, input_lengths) - else: - outputs = output_ids - - return outputs diff --git a/tensorrt_llm/runtime/multimodal_model_runner.py b/tensorrt_llm/runtime/multimodal_model_runner.py deleted file mode 100644 index 3ab5bb7ed824..000000000000 --- a/tensorrt_llm/runtime/multimodal_model_runner.py +++ /dev/null @@ -1,2744 +0,0 @@ -import json -import os -import sys -from io import BytesIO - -import requests - -# isort: off -import torch -import numpy as np -# isort: on -import math -from typing import Optional, Tuple - -import torch.nn.functional as F - -try: - from cuda.bindings import runtime as cudart -except ImportError: - from cuda import cudart - -from huggingface_hub import hf_hub_download -from PIL import Image, UnidentifiedImageError -from safetensors import safe_open -from torch import nn -from transformers import (AutoConfig, AutoModelForCausalLM, AutoProcessor, - AutoTokenizer) - -from .. import profiler -from .._deprecation import emit_engine_arch_deprecation -from .._utils import (get_hf_rope_theta, maybe_pin_memory, mpi_rank, - prefer_pinned, str_dtype_to_torch, str_dtype_to_trt, - supports_inflight_batching, torch_dtype_to_trt, - trt_dtype_to_torch) -from ..functional import RopeEmbeddingUtils, RotaryScalingType -from ..layers import MropeParams -from ..logger import logger -from .enc_dec_model_runner import EncDecModelRunner -from .model_runner import ModelRunner -from .session import Session, TensorInfo - -try: - import tensorrt_llm.bindings # NOQA - PYTHON_BINDINGS = True -except ImportError: - PYTHON_BINDINGS = False - -if PYTHON_BINDINGS: - from .model_runner_cpp import ModelRunnerCpp - - -class LlavaNextUtils: - # https://github.com/haotian-liu/LLaVA/blob/main/llava/mm_utils.py - - @staticmethod - def select_best_resolution(original_size, possible_resolutions): - """ - Selects the best resolution from a list of possible resolutions based on the original size. - - Args: - original_size (tuple): The original size of the image in the format (width, height). - possible_resolutions (list): A list of possible resolutions in the format [(width1, height1), (width2, height2), ...]. - - Returns: - tuple: The best fit resolution in the format (width, height). - """ - original_width, original_height = original_size - best_fit = None - max_effective_resolution = 0 - min_wasted_resolution = float('inf') - - for width, height in possible_resolutions: - scale = min(width / original_width, height / original_height) - downscaled_width, downscaled_height = int( - original_width * scale), int(original_height * scale) - effective_resolution = min(downscaled_width * downscaled_height, - original_width * original_height) - wasted_resolution = (width * height) - effective_resolution - - if effective_resolution > max_effective_resolution or ( - effective_resolution == max_effective_resolution - and wasted_resolution < min_wasted_resolution): - max_effective_resolution = effective_resolution - min_wasted_resolution = wasted_resolution - best_fit = (width, height) - - return best_fit - - @staticmethod - def get_anyres_image_grid_shape(image_size, - patch_size, - image_grid_pinpoints=None): - """ - Calculate the shape of the image patch grid after the preprocessing for images of any resolution. - - Args: - image_size (tuple): The size of the input image in the format (width, height). - patch_size (int): The size of each image patch. - - Returns: - tuple: The shape of the image patch grid in the format (width, height). - """ - if image_grid_pinpoints is None: - image_grid_pinpoints = [[336, 672], [672, 336], [672, 672], - [1008, 336], [336, 1008]] - width, height = LlavaNextUtils.select_best_resolution( - image_size, image_grid_pinpoints) - return width // patch_size, height // patch_size - - @staticmethod - def unpad_image(tensor, original_size): - """ - Unpads a PyTorch tensor of a padded and resized image. - - Args: - tensor (torch.Tensor): The image tensor, assumed to be in CxHxW format. - original_size (tuple): The original size of the image (width, height). - - Returns: - torch.Tensor: The unpadded image tensor. - """ - original_width, original_height = original_size - current_height, current_width = tensor.shape[1:] - - original_aspect_ratio = original_width / original_height - current_aspect_ratio = current_width / current_height - - if original_aspect_ratio > current_aspect_ratio: - scale_factor = current_width / original_width - new_height = int(original_height * scale_factor) - padding = (current_height - new_height) // 2 - unpadded_tensor = tensor[:, padding:current_height - padding, :] - else: - scale_factor = current_height / original_height - new_width = int(original_width * scale_factor) - padding = (current_width - new_width) // 2 - unpadded_tensor = tensor[:, :, padding:current_width - padding] - - return unpadded_tensor - - @staticmethod - def rearrange_image_features(image_feature, image_newline, image_size): - """ - Combine PyTorch feature grids from image patches. - - Args: - image_feature (torch.Tensor): The feature grids, assumed to be in NxCxHxW format. - image_newline (torch.Tensor): The newline embedding. - image_size (tuple): Size of the original image (width, height). - """ - CLIP_IMAGE_SIZE = 336 - CLIP_PATCH_SIZE = 14 - NUM_PATCHES_PER_SIDE = CLIP_IMAGE_SIZE // CLIP_PATCH_SIZE - if image_feature.shape[0] == 1: - return torch.cat((image_feature, image_newline[None]), dim=0) - - base_image_feature = image_feature[0] - image_feature = image_feature[1:] - height = width = NUM_PATCHES_PER_SIDE - assert height * width == base_image_feature.shape[0] - - num_patch_width, num_patch_height = LlavaNextUtils.get_anyres_image_grid_shape( - image_size, CLIP_IMAGE_SIZE) - image_feature = image_feature.view(num_patch_height, num_patch_width, - height, width, -1) - - image_feature = image_feature.permute(4, 0, 2, 1, 3).contiguous() - image_feature = image_feature.flatten(1, 2).flatten(2, 3) - image_feature = LlavaNextUtils.unpad_image(image_feature, image_size) - image_feature = torch.cat( - (image_feature, image_newline[:, None, None].expand( - *image_feature.shape[:-1], 1)), - dim=-1) - image_feature = image_feature.flatten(1, 2).transpose(0, 1) - image_feature = torch.cat((base_image_feature, image_feature), dim=0) - return image_feature - - -class LlavaOnevisionUtils: - # https://github.com/huggingface/transformers/blob/main/src/transformers/models/llava_onevision/modeling_llava_onevision.py - - @staticmethod - def pack_image_features(image_features, image_sizes, image_newline): - """ - Reshape, unpad and then pack each image_feature into a single image_features tensor containing all visual vectors. - - Args: - image_features (`torch.Tensor` of shape `(num_images, num_patches, image_length, embed_dim)`) - Image feature tensor, each contains all the visual feature of all patches. - image_sizes (`torch.Tensor` of shape `(num_images, 2)`) - Actual image size of each images (H, W). - image_newline (`torch.Tensor` of shape `(embed_dim)`) - New line embedding vector. - Returns: - image_features (`torch.Tensor` of shape `(all_feat_len, embed_dim)`) - """ - - IMAGE_SIZE = 384 - PATCH_SIZE = 14 - MAX_NUM_PATCHES = 9 - - new_image_features = [] - for image_idx, image_feature in enumerate(image_features): - if image_feature.shape[0] > 1: - base_image_feature = image_feature[0] - image_feature = image_feature[1:] - height = width = IMAGE_SIZE // PATCH_SIZE - if height * width != base_image_feature.shape[0]: - raise ValueError( - "The number of patches is not consistent with the image size." - ) - - IMAGE_GRID_PINPOINTS = [[384, 384], [384, 768], [384, 1152], - [384, 1536], [384, 1920], [384, 2304], - [768, 384], [768, 768], [768, 1152], - [768, 1536], [768, 1920], [768, 2304], - [1152, 384], [1152, 768], [1152, 1152], - [1152, 1536], - [1152, 1920], [1152, 2304], [1536, 384], - [1536, 768], [1536, 1152], [1536, 1536], - [1536, 1920], [1536, 2304], [1920, 384], - [1920, 768], [1920, 1152], [1920, 1536], - [1920, 1920], [1920, 2304], [2304, 384], - [2304, 768], [2304, 1152], [2304, 1536], - [2304, 1920], [2304, 2304]] - num_patch_width, num_patch_height = LlavaNextUtils.get_anyres_image_grid_shape( - image_sizes[image_idx][[1, 0]].tolist(), IMAGE_SIZE, - IMAGE_GRID_PINPOINTS) - image_feature = image_feature.view(num_patch_height, - num_patch_width, height, - width, -1) - image_feature = image_feature.permute(4, 0, 2, 1, - 3).contiguous() - image_feature = image_feature.flatten(1, 2).flatten(2, 3) - image_feature = LlavaNextUtils.unpad_image( - image_feature, image_sizes[image_idx][[1, 0]]) - - channels, curr_height, curr_width = image_feature.shape - ratio = math.sqrt(curr_height * curr_width / - (MAX_NUM_PATCHES * height**2)) - if ratio > 1.1: - image_feature = image_feature[None] - image_feature = nn.functional.interpolate( - image_feature, - [int(curr_height // ratio), - int(curr_width // ratio)], - mode="bilinear")[0] - - image_feature = torch.cat( - ( - image_feature, - image_newline[:, None, None].expand( - *image_feature.shape[:-1], 1).to( - image_feature.device, image_feature.dtype), - ), - dim=-1, - ) - image_feature = image_feature.flatten(1, 2).transpose(0, 1) - image_feature = torch.cat((base_image_feature, image_feature), - dim=0) - else: - image_feature = image_feature[0] - if image_newline is not None: - image_feature = torch.cat( - (image_feature, image_newline[None].to(image_feature)), - dim=0) - new_image_features.append(image_feature) - image_features = torch.stack(new_image_features) - return image_features - - @staticmethod - def apply_pooling(image_features): - IMAGE_SIZE = 384 - PATCH_SIZE = 14 - height = width = IMAGE_SIZE // PATCH_SIZE - batch_frames, seq_len, dim = image_features.shape - image_features = image_features.view(batch_frames, height, width, -1) - image_features = image_features.permute(0, 3, 1, 2).contiguous() - - height, width = image_features.shape[2:] - scaled_shape = [math.ceil(height / 2), math.ceil(width / 2)] - image_features = nn.functional.interpolate(image_features, - size=scaled_shape, - mode="bilinear") - - image_features = image_features.permute(0, 2, 3, 1) - image_features = image_features.view(batch_frames, -1, dim) - return image_features - - -class PhiMMUtils: - - @staticmethod - def add_image_newline(image_features, image_newline): - h, w, d = image_features.shape - image_newline = image_newline.expand(h, 1, -1) - image_features_newline = torch.cat([image_features, image_newline], - dim=1).flatten(0, 1) - return image_features_newline - - @staticmethod - def reshape_hd_patches(image_features, h_crop=1, w_crop=1): - n_crops, n_tokens, d = image_features.shape - assert n_crops == h_crop * w_crop - - h = w = int(n_tokens**0.5) - image_features = image_features.reshape( - h_crop, w_crop, h, w, - d).permute(0, 2, 1, 3, 4).reshape(h_crop * h, w_crop * w, d) - return image_features - - @staticmethod - def hd_feature_transform(image_features, - h_crop, - w_crop, - sub_GN, - glb_GN, - patch_mask=None): - glb_image_features = PhiMMUtils.add_image_newline( - PhiMMUtils.reshape_hd_patches(image_features[:1]), sub_GN) - - num_crops = h_crop * w_crop - sub_image_features = PhiMMUtils.reshape_hd_patches( - image_features[1:num_crops + 1], h_crop, w_crop) - if patch_mask is not None: - h, w = patch_mask.shape[1] // 2, patch_mask.shape[2] // 2 - sub_image_mask = (patch_mask[1:num_crops + 1, 0::2, - 0::2].bool().reshape( - h_crop, w_crop, h, - w).permute(0, 2, 1, 3).reshape( - h_crop * h, w_crop * w)) - hh = int(sub_image_mask[:, 0].sum().item()) - ww = int(sub_image_mask[0, :].sum().item()) - sub_image_features = sub_image_features[:hh, :ww] - sub_image_features = PhiMMUtils.add_image_newline( - sub_image_features, sub_GN) - - image_features = torch.cat( - [sub_image_features, - glb_GN.expand(1, -1), glb_image_features]) - return image_features - - @staticmethod - def reshape_audio_chunks(audio_features, chunk_mask=None): - audio_features = audio_features.flatten(0, 1) - if chunk_mask is not None: - # only the last chunk may include paddings - n_tokens = math.ceil(chunk_mask.flatten().sum().item() / 8) - audio_features = audio_features[:n_tokens] - return audio_features - - -class MultimodalModelRunner: - - def __init__(self, args): - emit_engine_arch_deprecation("MultimodalModelRunner") - self.args = args - self.use_trtllm_vision_engine = False - - self.runtime_rank = mpi_rank() - device_id = self.runtime_rank % torch.cuda.device_count() - torch.cuda.set_device(device_id) - self.device = "cuda:%d" % (device_id) - - self.stream = torch.cuda.Stream(torch.cuda.current_device()) - torch.cuda.set_stream(self.stream) - - if self.args.mm_embedding_offloading is None: - self.args.mm_embedding_offloading = self.args.enable_chunked_context - elif self.args.mm_embedding_offloading and not self.args.enable_chunked_context: - logger.warning( - "mm_embedding_offloading requires enable_chunked_context to be True. Setting mm_embedding_offloading to None." - ) - self.args.mm_embedding_offloading = None - - # parse model type from visual engine config - with open(os.path.join(self.visual_engine_dir, "config.json"), - "r") as f: - config = json.load(f) - if 'pretrained_config' in config: - if config['pretrained_config'][ - 'architecture'] == 'LlavaNextForConditionalGeneration': - self.model_type = 'llava_next' - self.vision_precision = config['pretrained_config']['dtype'] - self.use_trtllm_vision_engine = True - else: - logger.error( - "Currently only Llava-NeXT supports TRT-LLM vision engines." - ) - else: - self.model_type = config['builder_config']['model_type'] - self.vision_precision = config['builder_config']['precision'] - if self.model_type == 'pix2struct': - self.vision_precision = 'float16' - self.decoder_llm = not ( - 't5' in self.model_type - or self.model_type in ['nougat', 'pix2struct'] - ) # BLIP2-T5, pix2struct and Nougat are using encoder-decoder models as LLMs - - if self.model_type == 'video-neva': - self.num_frames = config['builder_config'].get('num_frames', None) - if self.model_type == "llava_next": - self.llm_name = AutoConfig.from_pretrained( - self.args.hf_model_dir).text_config._name_or_path - if 'internlm' in self.model_type: - self.args.lora_task_uids = ['0'] * args.batch_size - if self.model_type == "qwen2_vl": - hf_config = AutoConfig.from_pretrained(self.args.hf_model_dir) - self.vision_start_token_id = hf_config.vision_start_token_id - self.vision_end_token_id = hf_config.vision_end_token_id - self.vision_token_id = hf_config.vision_token_id - self.image_token_id = hf_config.image_token_id - self.video_token_id = hf_config.video_token_id - self.spatial_merge_size = hf_config.vision_config.spatial_merge_size - self.max_position_embeddings = hf_config.max_position_embeddings - self.hidden_size = hf_config.hidden_size - self.num_attention_heads = hf_config.num_attention_heads - self.rope_theta = get_hf_rope_theta(hf_config, 10000.0) - if self.model_type == 'llava_onevision': - self.num_frames = self.args.video_num_frames - if self.num_frames is None: - self.num_frames = 8 - assert self.args.video_path is None or self.args.image_path is None - if self.model_type == "pixtral": - hf_config = AutoConfig.from_pretrained(self.args.hf_model_dir) - self.image_size = hf_config.vision_config.image_size - self.patch_size = hf_config.vision_config.patch_size - self.vocab_size = hf_config.text_config.vocab_size - self.image_token_index = hf_config.image_token_index - self.spatial_merge_size = hf_config.spatial_merge_size - - self.audio_input_names = self.audio_output_names = None - if self.model_type == "mllama": - self.vision_input_names = [ - "pixel_values", - "aspect_ratio_ids", - "aspect_ratio_mask", - ] - self.vision_output_names = [ - "encoder_output", - ] - elif self.model_type == "llava_next" and self.use_trtllm_vision_engine: - self.vision_input_names = [ - "pixel_values", - ] - self.vision_output_names = [ - "image_features", - ] - elif self.model_type == "phi-4-multimodal": - self.vision_input_names = ["input", "attention_mask"] - self.audio_input_names = ["input", "attention_mask"] - self.audio_output_names = ["encoder_output"] - self.vision_output_names = ["encoder_output"] - else: - self.vision_input_names = ["input"] - self.vision_output_names = ["encoder_output"] - - self.session = args.session - if self.cpp_e2e: - self.visual_output_shape = config['builder_config'].get( - 'output_shape', None) - if self.decoder_llm: - if not supports_inflight_batching(self.llm_engine_dir): - logger.warning( - "The given engine does not support in-flight batching, both visual engine and LLM fallback to python session" - ) - self.session = 'python' - - if not PYTHON_BINDINGS and 'cpp' in args.session: - logger.warning( - "Python bindings of C++ session is unavailable, both visual engine and LLM fallback to Python session." - ) - self.session = 'python' - - args.debug_mode = False - if args.debug_mode and 'cpp' in args.session: - logger.warning( - "Debug mode is not supported in C++ session for now, both visual engine and LLM fallback to Python session." - ) - self.session = 'python' - - if self.model_type == 'qwen2_vl': - if self.args.session != "cpp_llm_only": - logger.warning( - "Qwen2-vl only support C++ session for now, fallback to C++ session." - ) - self.args.session = "cpp_llm_only" - - if (not (self.model_type in - ('llava', 'vila', 'blip2-opt', 'kosmos-2', 'fuyu', - 'cogvlm', 'neva', "internvl") or 'internlm' - in self.model_type)) and args.session == 'cpp': - logger.warning( - f'C++ end-to-end mode does not support {self.model_type}. Visual engine fallbacks to Python session. See support matrix in README.' - ) - args.session = 'cpp_llm_only' - self.session = args.session - - else: - self.session = 'cpp_llm_only' - - self.init_tokenizer() - self.init_processor() - self.init_image_encoder() - self.init_llm() - - if self.audio_input_names is not None: - with open(os.path.join(self.audio_engine_dir, "config.json"), - "r") as f: - config = json.load(f) - self.audio_precision = config['builder_config']['precision'] - self.init_audio_encoder() - else: - self.audio_encoder_session = self.audio_precision = None - - @property - def cpp_e2e(self): - return self.session == 'cpp' - - @property - def cpp_llm_only(self): - return self.session == 'cpp_llm_only' - - @property - def python_e2e(self): - return self.session == 'python' - - @property - def visual_engine_dir(self): - return os.path.join(self.args.engine_dir, 'vision') - - @property - def audio_engine_dir(self): - return os.path.join(self.args.engine_dir, 'audio') - - @property - def llm_engine_dir(self): - return os.path.join(self.args.engine_dir, 'llm') - - def init_tokenizer(self): - if self.model_type == 'nougat': - from transformers import NougatTokenizerFast - self.tokenizer = NougatTokenizerFast.from_pretrained( - self.args.hf_model_dir) - elif self.model_type == 'neva' or self.model_type == 'video-neva': - from sentencepiece import SentencePieceProcessor - - sp = SentencePieceProcessor( - os.path.join(self.args.hf_model_dir, 'tokenizer.model')) - - class return_obj: - - def __init__(self, input_ids): - self.input_ids = input_ids - - def __getitem__(self, name): - if name in "input_ids": - return self.input_ids - else: - raise AttributeError( - f"'return_obj' has no item '{name}'") - - # sentencepiece does not follow the same interface as HF - class HFTokenizerInterface(): - - def encode(self, x, return_tensors=None, **kwargs): - out = sp.encode(x) - if return_tensors == "pt": - out = torch.tensor(out) - return return_obj(out) - - def __call__(self, x, return_tensors=None, **kwargs): - return self.encode(x, return_tensors, **kwargs) - - def decode(self, x, **kwargs): - return sp.decode(x.tolist()) - - def batch_decode(self, x, **kwargs): - return self.decode(x, **kwargs) - - self.tokenizer = HFTokenizerInterface() - self.tokenizer.eos_token_id = sp.eos_id() - self.tokenizer.bos_token_id = sp.bos_id() - self.tokenizer.pad_token_id = sp.pad_id() - elif self.model_type == 'vila': - self.tokenizer = AutoTokenizer.from_pretrained( - self.args.hf_model_dir + "/llm", - use_fast=False, - use_legacy=False) - else: - use_fast = self.model_type in [ - "phi-3-vision", "phi-4-multimodal", "internvl" - ] - self.tokenizer = AutoTokenizer.from_pretrained( - self.args.hf_model_dir, - use_fast=use_fast, - use_legacy=False, - trust_remote_code=getattr(self.args, "trust_remote_code", - False)) - - self.tokenizer.padding_side = "right" - - def init_processor(self): - from torchvision import transforms - - if 'blip2' in self.model_type: - from transformers import Blip2Processor - self.processor = Blip2Processor.from_pretrained( - self.args.hf_model_dir) - - elif 'nougat' in self.model_type: - from transformers import NougatProcessor - self.processor = NougatProcessor.from_pretrained( - self.args.hf_model_dir) - - elif 'cogvlm' in self.model_type: - image_size = 490 - self.transform = transforms.Compose([ - transforms.Resize( - (image_size, image_size), - interpolation=transforms.InterpolationMode.BICUBIC), - transforms.ToTensor(), - transforms.Normalize((0.48145466, 0.4578275, 0.40821073), - (0.26862954, 0.26130258, 0.27577711)), - transforms.ConvertImageDtype(torch.bfloat16), - ]) - - elif self.model_type in [ - 'phi-3-vision', 'pix2struct', 'llava_next', 'llava', 'fuyu', - 'kosmos-2', 'mllama', 'llava_onevision', 'qwen2_vl', - 'phi-4-multimodal' - ]: - self.processor = AutoProcessor.from_pretrained( - self.args.hf_model_dir, - trust_remote_code=getattr(self.args, "trust_remote_code", - False), - num_crops=16) - - elif 'pixtral' in self.model_type: - self.processor = AutoProcessor.from_pretrained( - self.args.hf_model_dir) - - elif 'internlm' in self.model_type: - image_size = 490 - self.processor = transforms.Compose([ - transforms.Resize( - (image_size, image_size), - interpolation=transforms.InterpolationMode.BICUBIC), - transforms.ToTensor(), - transforms.Normalize((0.48145466, 0.4578275, 0.40821073), - (0.26862954, 0.26130258, 0.27577711)), - ]) - - elif 'internvl' in self.model_type: - from transformers import CLIPImageProcessor - self.processor = CLIPImageProcessor.from_pretrained( - 'OpenGVLab/InternViT-300M-448px' - ) # You can change the InternViT model type according to your InternVL type - - elif self.model_type == "neva": - image_size = 384 - self.transform = transforms.Compose([ - transforms.Resize( - (image_size, image_size), - interpolation=transforms.InterpolationMode.BICUBIC), - transforms.ToTensor(), - transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), - transforms.ConvertImageDtype(torch.float32), - ]) - - elif self.model_type == "video-neva": - pass - - elif self.model_type == "vila": - sys.path.append(self.args.hf_model_dir + "/../VILA") - from llava.mm_utils import process_images - from llava.model import LlavaLlamaConfig # noqa - from transformers import AutoModel - model = AutoModel.from_pretrained( - self.args.hf_model_dir, - device_map='auto', - trust_remote_code=getattr(self.args, "trust_remote_code", - False), - ) - vision_tower = model.get_vision_tower() - vision_tower.image_processor - - def processor(raw_image): - return process_images(raw_image, vision_tower.image_processor, - model.config).to(model.device, - dtype=torch.float16) - - self.processor = processor - - if self.model_type == 'mllama': - from .processor_wrapper import MllamaProcessorWrapper - self.processor = MllamaProcessorWrapper(self.processor, logger) - - def init_image_encoder(self): - - # Phi-4-multimodal uses pytorch engine due to issues with creating TRT engine. - if self.model_type == "phi-4-multimodal": - model = AutoModelForCausalLM.from_pretrained( - self.args.hf_model_dir, - dtype=torch.float16, - trust_remote_code=getattr(self.args, "trust_remote_code", - False), - device_map='cpu') - self.vision_model = model.model.embed_tokens_extend.image_embed.to( - self.device).eval() - self.image_newlines = {} - self.image_newlines['sub_GN'] = self.vision_model.img_projection( - self.vision_model.sub_GN).squeeze() - self.image_newlines['glb_GN'] = self.vision_model.img_projection( - self.vision_model.glb_GN).squeeze() - return - - if self.model_type == "phi-3-vision": - model = AutoModelForCausalLM.from_pretrained( - self.args.hf_model_dir, - dtype=torch.float16, - trust_remote_code=getattr(self.args, "trust_remote_code", - False), - device_map='cpu') - self.vision_model = model.model.vision_embed_tokens.to( - self.device).eval() - - # Test run vision_model.get_img_features to pre-allocate memory for flash attention - image = self.processor(text="<|image_1|>", - images=Image.new('RGB', [10, 10]), - return_tensors="pt")['pixel_values'] - image = image.flatten(0, 1) - image = torch.rand(image.shape, - dtype=str_dtype_to_torch(self.vision_precision), - device=self.device) - self.vision_model.get_img_features(image) - return - - if self.cpp_e2e: - logger.info( - "Using C++ runtime for both visual engine and LLM decoder, skip loading visual engine in Python runtime." - ) - elif self.model_type == "llava_next" and self.use_trtllm_vision_engine: - cudart.cudaSetDevice(self.runtime_rank % torch.cuda.device_count()) - - vision_encoder_path = os.path.join( - self.visual_engine_dir, f"rank{self.runtime_rank}.engine") - logger.info(f'Loading engine from {vision_encoder_path}') - with open(vision_encoder_path, "rb") as f: - engine_buffer = f.read() - logger.info(f'Creating session from engine {vision_encoder_path}') - assert engine_buffer is not None - - self.visual_encoder_session = Session.from_serialized_engine( - engine_buffer) - else: - vision_encoder_path = os.path.join(self.visual_engine_dir, - self.args.visual_engine_name) - logger.info(f'Loading engine from {vision_encoder_path}') - with open(vision_encoder_path, 'rb') as f: - engine_buffer = f.read() - logger.info(f'Creating session from engine {vision_encoder_path}') - self.visual_encoder_session = Session.from_serialized_engine( - engine_buffer) - if self.model_type in ["llava_next", "llava_onevision"]: - self.image_newlines = {} - image_newlines_path = os.path.join(self.visual_engine_dir, - 'image_newlines.safetensors') - with safe_open(image_newlines_path, - framework="pt", - device=self.device) as f: - for k in f.keys(): - self.image_newlines[k] = f.get_tensor(k) - - def init_audio_encoder(self): - assert self.model_type == "phi-4-multimodal" - model = AutoModelForCausalLM.from_pretrained( - self.args.hf_model_dir, - dtype=torch.float16, - trust_remote_code=getattr(self.args, "trust_remote_code", False), - device_map='cpu') - self.audio_model = model.model.embed_tokens_extend.audio_embed.to( - self.device).eval() - - def init_llm(self): - if self.decoder_llm: - cross_kv_cache_fraction = None - if self.model_type == 'mllama': - cross_kv_cache_fraction = self.args.cross_kv_cache_fraction - if self.python_e2e: - logger.info(f'Running LLM with Python runner') - self.model = ModelRunner.from_dir( - self.llm_engine_dir, - rank=tensorrt_llm.mpi_rank(), - debug_mode=False, - stream=self.stream, - enable_context_fmha_fp32_acc=self.args. - enable_context_fmha_fp32_acc, - multi_block_mode=self.args.multi_block_mode, - ) - self.model_config = self.model.session._model_config - elif self.cpp_e2e: - logger.info( - f'Running both visual engine and LLM with Python runner') - self.model = ModelRunnerCpp.from_dir( - self.args.engine_dir, - rank=tensorrt_llm.mpi_rank(), - debug_mode=False, - is_enc_dec=True, # TODO: add a separate model variant here? - enable_context_fmha_fp32_acc=self.args. - enable_context_fmha_fp32_acc) - self.model_config = self.model.model_config - else: - logger.info(f'Running LLM with C++ runner') - self.model = ModelRunnerCpp.from_dir( - self.llm_engine_dir, - rank=tensorrt_llm.mpi_rank(), - debug_mode=False, - enable_chunked_context=self.args.enable_chunked_context, - enable_context_fmha_fp32_acc=self.args. - enable_context_fmha_fp32_acc, - kv_cache_free_gpu_memory_fraction=self.args. - kv_cache_free_gpu_memory_fraction, - cross_kv_cache_fraction=cross_kv_cache_fraction, - multi_block_mode=self.args.multi_block_mode, - mm_embedding_offloading=self.args.mm_embedding_offloading, - ) - self.model_config = self.model.model_config - self.runtime_mapping = self.model.mapping - else: - self.model = EncDecModelRunner.from_engine( - os.path.basename(self.args.hf_model_dir), - self.llm_engine_dir, - skip_encoder=self.model_type in ['nougat', 'pix2struct'], - debug_mode=False, - stream=self.stream, - enable_context_fmha_fp32_acc=self.args. - enable_context_fmha_fp32_acc) - if self.model_type in ['nougat', 'pix2struct']: - self.model_config = self.model.decoder_model_config - self.runtime_mapping = self.model.decoder_runtime_mapping - else: - self.model_config = self.model.encoder_model_config - self.runtime_mapping = self.model.encoder_runtime_mapping - - def video_preprocess(self, video_path): - from decord import VideoReader - if isinstance(video_path, str): - vr = VideoReader(video_path) - num_frames = self.num_frames - if num_frames == -1: - frames = [ - Image.fromarray(frame.asnumpy()[:, :, ::-1]).convert('RGB') - for frame in vr - ] - else: - # equally sliced frames into self.num_frames frames - # if self.num_frames is greater than the number of frames in the video, we will repeat the last frame - num_frames = min(num_frames, len(vr)) - indices = np.linspace(0, len(vr) - 1, num=num_frames, dtype=int) - frames = [ - Image.fromarray( - vr[idx].asnumpy()[:, :, ::-1]).convert('RGB') - for idx in indices - ] - if len(frames) < num_frames: - frames += [frames[-1]] * (num_frames - len(frames)) - else: - frames = self.video_path - - from transformers import CLIPImageProcessor - processor = CLIPImageProcessor.from_pretrained( - "openai/clip-vit-large-patch14", dtype=torch.bfloat16) - frames = processor.preprocess(frames, - return_tensors="pt")['pixel_values'] - # make dtype consistent with vision encoder - media_tensors = frames.to(str_dtype_to_torch( - self.vision_precision)) # [num_frames, 3, H, W] - return media_tensors.unsqueeze(0) #[1, num_frames, 3, H, W] - - def preprocess(self, pre_prompt, post_prompt, image, other_vision_inputs, - other_audio_inputs): - audio = None - # same prompt for single/multiple image(s) - n_prompts_n_images = False - if isinstance(post_prompt, - list) and len(post_prompt) > 1 and image is not None: - if hasattr(image, "pixel_values"): - if len(post_prompt) == image["pixel_values"].shape[0]: - n_prompts_n_images = True - # n prompts and n images - else: - if isinstance( - image, - torch.Tensor) and len(post_prompt) == image.shape[0]: - n_prompts_n_images = True - # n prompts and n images - - if self.model_type == 'kosmos-2': - input_ids = image['input_ids'].clone() - image_mask = image["image_embeds_position_mask"] - image = image['pixel_values'] - input_ids += image_mask * (self.model_config.vocab_size - 4) - input_ids = input_ids.expand(self.args.batch_size, - *input_ids.shape[1:]) - length = input_ids.shape[1] - elif self.model_type == 'phi-3-vision': - input = image - image = input['pixel_values'] - image = image.flatten(0, 1) - elif self.model_type == 'phi-4-multimodal': - input = image - image = input['input_image_embeds'].flatten(0, 1) - other_vision_inputs['attention_mask'] = input[ - 'image_attention_mask'].flatten(0, 1).bool() - - audio = input['input_audio_embeds'] - l, d = audio.shape[1], audio.shape[2] - pad = 4000 - l % 4000 - audio = torch.cat([audio, audio.new_zeros(1, pad, d)], - dim=1).reshape(-1, 4000, d) - audio_mask = audio.new_ones(*audio.shape[:2]) - audio_mask[-1, -pad:] = 0 - other_audio_inputs['attention_mask'] = audio_mask.bool() - elif self.model_type == 'pixtral': - # Hold on to pixel_values and input_ids. - dtype = str_dtype_to_torch(self.vision_precision) - # Shape of pixel values from the processor varies with the raw image. - # So we create a new tensor with a fixed shape as expected by the vision - # encoder and create a corresponding attention mask. - image_size = self.image_size - patch_size = self.patch_size - d_min = torch.finfo(dtype).min - num_patches = (image_size // patch_size) - padded_image = torch.full( - (self.args.batch_size, 3, image_size, image_size), - fill_value=0, - dtype=dtype, - device="cuda") - padded_attention_mask = torch.full( - (self.args.batch_size, num_patches, num_patches), - fill_value=d_min, - dtype=dtype, - device="cuda") - h, w, input_ids = [], [], [] - for img_idx in range(self.args.batch_size): - pixel_values = image["pixel_values"][img_idx] - img_h, img_w = pixel_values.shape[-2:] - padded_image[img_idx, :, :img_h, :img_w] = pixel_values - padded_attention_mask[img_idx, :img_h // patch_size, :img_w // - patch_size] = 0 - input_ids.append(image["input_ids"][img_idx]) - h.append(img_h) - w.append(img_w) - - image = padded_image - other_vision_inputs = { - "attention_mask": padded_attention_mask, - } - elif self.model_type == 'llava_next': - input = image - image = input['pixel_values'] - image = image[0] - image_size = input['image_sizes'][0].cpu() - elif self.model_type == "qwen2_vl": - input = image - image = input['image'] - input_ids = input['input_ids'] - other_vision_inputs['image_grid_thw'].shape[0] - attention_mask = other_vision_inputs['attention_mask_llm'] - other_vision_inputs.pop('attention_mask_llm') - image_grid_thw = other_vision_inputs['image_grid_thw'] - other_vision_inputs.pop('image_grid_thw') - elif self.model_type == 'llava_onevision': - input = image - if self.args.video_path is None: - image = input['pixel_values'] - image = image[0].repeat(self.args.batch_size, 1, 1, 1) - image_size = input['image_sizes'][0] - image_size = image_size.repeat(self.args.batch_size, 1).cpu() - else: - image = input['pixel_values_videos'] - _, _, c, h, w = image.shape - image = image.repeat(self.args.batch_size, 1, 1, 1, 1) - image = image.view(-1, c, h, w) - elif self.model_type == "fuyu": - while len(image["image_patches"]) < self.args.batch_size: - image["image_patches"].append(image["image_patches"][0]) - - profiler.start("Vision encoder") - visual_features, visual_atts, model_runner_input = None, None, None - if image is not None: - model_runner_input = torch.stack( - image['image_patches'], - dim=0) if self.model_type == 'fuyu' else image - - if self.model_type == "phi-3-vision": - visual_features = self.vision_model.get_img_features( - image).reshape(1, image.shape[0], -1, - self.vision_model.image_dim_out) - visual_atts = None - elif self.model_type == "phi-4-multimodal": - visual_features = self.vision_model.get_img_features( - model_runner_input.to( - str_dtype_to_torch(self.vision_precision)), - other_vision_inputs['attention_mask']) - visual_features = self.vision_model.img_projection( - visual_features) - visual_atts = torch.ones(visual_features.size()[:-1], - dtype=torch.long).to( - model_runner_input.device) - else: - if self.cpp_e2e: - # If using E2E C++ runtime, visual_features will not be computed here in Python runtime. - # Instead, it only contains a shape read from the engine config, and is used for generating - # decoder prompt later - logger.info( - 'Skip running visual engine, get visual output shape from engine config.' - ) - model_runner_input = model_runner_input.to( - str_dtype_to_torch(self.vision_precision)) - batch_size = model_runner_input.shape[0] - output_shape = list(self.visual_output_shape) - output_shape[0] = batch_size - if self.model_type == 'fuyu': - output_shape[1] = model_runner_input.shape[ - 2] # fuyu's output patch number is not fixed, same as input patch number - visual_features = TensorInfo( - 'encoder_output', - str_dtype_to_trt(self.vision_precision), - tuple(output_shape)) - atts_shape = visual_features.shape[:-1] - visual_atts = TensorInfo('image_atts', None, - tuple(atts_shape)) - model_runner_input = torch.vsplit( - model_runner_input, model_runner_input.shape[0]) - else: - visual_features, visual_atts = self.get_visual_features( - model_runner_input, other_vision_inputs) - model_runner_input = None - profiler.stop("Vision encoder") - - profiler.start("Audio encoder") - audio_features = None - if audio is not None: - audio_features = self.get_audio_features(audio, other_audio_inputs) - profiler.stop("Audio encoder") - - if self.model_type == 'fuyu': - input_ids = image['input_ids'].to(torch.int32) - image_patches_indices = image['image_patches_indices'].to( - torch.int32) - - input_ids = input_ids.expand(self.args.batch_size, - *input_ids.shape[1:]) - image_patches_indices = image_patches_indices.expand( - self.args.batch_size, *image_patches_indices.shape[1:]) - - input_ids = self.ptuning_setup_fuyu(input_ids, - image_patches_indices) - input_ids = torch.stack(input_ids, dim=0).to('cpu') - length = input_ids.shape[1] - if not self.cpp_e2e: # TODO: bs > 1 for C++ E2E Fuyu - visual_features = visual_features.repeat( - self.args.batch_size, 1, 1) - elif self.model_type == 'qwen2_vl': - length = input_ids.shape[1] - input_lengths = torch.IntTensor([length] * self.args.batch_size).to( - torch.int32) - input_ids, ptuning_args, mrope_args = self.setup_fake_prompts_qwen2vl( - visual_features, input_ids, image_grid_thw, attention_mask, - input_lengths) - return input_ids, input_lengths, ptuning_args, visual_features, mrope_args - - elif self.model_type == 'kosmos-2': - visual_features = visual_features.squeeze( - ) if visual_features is not None else None - elif self.model_type == 'vila': - if n_prompts_n_images: - input_ids = self.tokenizer_image_token(self.args.batch_size, - pre_prompt[0], - post_prompt, - self.tokenizer) - else: - input_ids = self.tokenizer_image_token(self.args.batch_size, - pre_prompt[0], - post_prompt[0], - self.tokenizer) - batch_split_prompts = self.split_prompt_by_images(input_ids) - if not n_prompts_n_images: - first_batch_split_prompts = batch_split_prompts[0] - # compute prompt length + visual length - length = sum( - [ids.shape[1] for ids in first_batch_split_prompts]) - if self.args.batch_size == 1 and len(image) > 1: - # mode 1: multiple image as a whole, flatten visual dims - length += visual_atts.shape[0] * visual_atts.shape[1] - else: - length += visual_atts.shape[1] - input_lengths = torch.IntTensor( - [length] * self.args.batch_size).to(torch.int32) - input_ids, ptuning_args = self.setup_fake_prompts_vila( - self.args.batch_size, visual_features, - first_batch_split_prompts, input_lengths) - else: - # mode 2: multiple different prompts corresponding to multiple images (1-1 correspondence) - length = [ - sum([ids.shape[1] for ids in batch_split_prompt]) - for batch_split_prompt in batch_split_prompts - ] - length = [l + visual_atts.shape[1] for l in length] - input_lengths = torch.IntTensor(length).to(torch.int32) - input_ids, ptuning_args = self.setup_fake_prompts_vila( - self.args.batch_size, visual_features, batch_split_prompts, - input_lengths) - return input_ids, input_lengths, ptuning_args, visual_features, model_runner_input - elif self.model_type == 'phi-3-vision': - image_sizes = input["image_sizes"] - profiler.start("Feature transform") - visual_features = self.vision_model.hd_feature_transform( - visual_features, image_sizes) - profiler.stop("Feature transform") - input_ids = input["input_ids"].clone() - input_ids = input_ids.expand(self.args.batch_size, - *input_ids.shape[1:]) - num_img_tokens = [visual_features.shape[0]] - input_ids = self.ptuning_setup_phi3(visual_features=visual_features, - audio_features=None, - input_ids=input_ids, - num_img_tokens=num_img_tokens, - num_aud_tokens=None) - visual_features = visual_features.unsqueeze(0).repeat( - self.args.batch_size, 1, 1) - length = input_ids.shape[1] - elif self.model_type == 'phi-4-multimodal': - h, w = input["image_sizes"][0] - image_attention_mask = input.get("image_attention_mask") - if image_attention_mask is not None: - image_attention_mask = image_attention_mask[0].bool() - patch_size = 336 if self.model_type == 'phi-3-vision' else 448 - profiler.start("Feature transform") - visual_features = PhiMMUtils.hd_feature_transform( - visual_features, - h // patch_size, - w // patch_size, - self.image_newlines["sub_GN"], - self.image_newlines["glb_GN"], - patch_mask=image_attention_mask) - profiler.stop("Feature transform") - input_ids = input["input_ids"].clone() - input_ids = input_ids.expand(self.args.batch_size, - *input_ids.shape[1:]) - num_img_tokens = [visual_features.shape[0]] - if audio_features is not None: - dim = audio_features.shape[-1] // 2 - audio_features = audio_features[..., -dim:] - audio_features = PhiMMUtils.reshape_audio_chunks( - audio_features, other_audio_inputs["attention_mask"]) - num_aud_tokens = [audio_features.shape[0]] - else: - num_aud_tokens = None - input_ids = self.ptuning_setup_phi3(visual_features=visual_features, - audio_features=audio_features, - input_ids=input_ids, - num_img_tokens=num_img_tokens, - num_aud_tokens=num_aud_tokens) - visual_features = visual_features.unsqueeze(0).repeat( - self.args.batch_size, 1, 1) - if audio_features is not None: - audio_features = audio_features.unsqueeze(0).repeat( - self.args.batch_size, 1, 1) - length = input_ids.shape[1] - - elif self.model_type == 'pixtral': - relevant_patch_size = self.patch_size * self.spatial_merge_size - output_img_size = self.image_size // relevant_patch_size - # Note: max_h * max_w shall serve as the `tokens_per_task` in ptuning prompt table. - max_h = max(h) // relevant_patch_size - max_w = max(w) // relevant_patch_size - visual_embed_dim = visual_features.shape[-1] - relevant_visual_features = torch.zeros(self.args.batch_size, - max_h * max_w, - visual_embed_dim) - for img_idx in range(self.args.batch_size): - complete_features = visual_features[img_idx] - complete_features = complete_features.reshape( - output_img_size, output_img_size, visual_embed_dim) - relevant_h = h[img_idx] // relevant_patch_size - relevant_w = w[img_idx] // relevant_patch_size - flattened_features = complete_features[:relevant_h, : - relevant_w, :].flatten( - 0, 1) - relevant_visual_features[img_idx, :relevant_h * - relevant_w, :] = flattened_features - visual_features = relevant_visual_features - input_ids = self.ptuning_setup_pixtral(input_ids=input_ids) - # Note: length is not used for pixtral model downstream. Setting it to a list - # of length of input_ids causes errors downstream. So, supplying a placeholder. - length = input_ids[0].shape[0] - - elif self.model_type == 'llava_next': - visual_features = LlavaNextUtils.rearrange_image_features( - visual_features, self.image_newlines["image_newline"], - image_size) - input_ids = self.ptuning_setup_llava_next(visual_features, - pre_prompt, post_prompt) - length = input_ids.shape[1] - elif self.model_type == 'mllama': - pre_input_ids = self.tokenizer(pre_prompt, - return_tensors="pt", - padding=True).input_ids - if n_prompts_n_images: - length = [pre_input_ids.shape[1]] * self.args.batch_size - else: - length = pre_input_ids.shape[1] - post_input_ids = None - elif self.model_type == 'llava_onevision': - if self.args.video_path is None: - visual_features = torch.split(visual_features, - visual_features.shape[0] // - self.args.batch_size, - dim=0) - visual_features = LlavaOnevisionUtils.pack_image_features( - visual_features, - image_size, - image_newline=self.image_newlines["image_newline"], - ) - else: - visual_features = LlavaOnevisionUtils.apply_pooling( - visual_features) - visual_features = visual_features.reshape( - self.args.batch_size, - self.num_frames * visual_features.shape[1], -1) - image_newline = self.image_newlines["image_newline"][ - None, None, :].repeat(self.args.batch_size, 1, - 1).to(visual_features.device) - visual_features = torch.cat((visual_features, image_newline), - dim=1) - - pre_input_ids = self.tokenizer(pre_prompt, - return_tensors="pt", - padding=True).input_ids - post_input_ids = self.tokenizer(post_prompt, - return_tensors="pt", - padding=True).input_ids - length = pre_input_ids.shape[1] + visual_features.shape[ - 1] + post_input_ids.shape[1] - else: - pre_input_ids = self.tokenizer(pre_prompt, - return_tensors="pt", - padding=True).input_ids - if post_prompt[0] is not None: - post_input_encoded = self.tokenizer(post_prompt, - return_tensors="pt", - padding=True) - post_input_ids = post_input_encoded.input_ids - if n_prompts_n_images and 'neva' not in self.model_type: - post_input_attention_mask = post_input_encoded.attention_mask - post_input_ids = [ - input_id[mask.bool()] for input_id, mask in zip( - post_input_ids, post_input_attention_mask) - ] - - if self.model_type == 'video-neva': - length = pre_input_ids.shape[1] + post_input_ids.shape[ - 1] + visual_atts.shape[2] * visual_atts.shape[1] - elif self.model_type == 'internvl': - length = pre_input_ids.shape[1] + post_input_ids.shape[ - 1] + visual_atts.shape[0] * visual_atts.shape[1] - else: - if n_prompts_n_images: - length = [ - pre_input_ids.shape[1] + visual_atts.shape[1] + - post_input_id.shape[0] - for post_input_id in post_input_ids - ] - else: - length = pre_input_ids.shape[1] + post_input_ids.shape[ - 1] + visual_atts.shape[1] - else: - post_input_ids = None - assert pre_input_ids.shape[0] == visual_atts.shape[0] - if visual_atts.shape[0] == 1: - length = pre_input_ids.shape[1] + visual_atts.shape[1] - else: - length = [ - pre_input_ids.shape[1] + visual_atts.shape[1] - for _ in range(visual_atts.shape[0]) - ] - - if n_prompts_n_images: - if isinstance(length, int): length = [length] - assert isinstance(length, list) - input_lengths = torch.IntTensor(length).to(torch.int32) - else: - assert isinstance(length, int) - input_lengths = torch.IntTensor([length] * self.args.batch_size).to( - torch.int32) - - if self.model_type in [ - 'fuyu', 'kosmos-2', 'phi-3-vision', 'llava_next', 'pixtral' - ]: - return input_ids, input_lengths, [ - visual_features - ], visual_features, model_runner_input - if self.model_type == 'phi-4-multimodal': - multimodal_features = torch.cat([visual_features, audio_features], - dim=1) - return input_ids, input_lengths, [ - multimodal_features - ], multimodal_features, model_runner_input - - input_ids, ptuning_args = self.setup_fake_prompts( - visual_features, pre_input_ids, post_input_ids, input_lengths) - - return input_ids, input_lengths, ptuning_args, visual_features, model_runner_input - - @staticmethod - def tokenizer_image_token(batch_size, - pre_prompt, - post_prompt, - tokenizer, - image_token_index=-200): - if isinstance(post_prompt, list): - prompts = [pre_prompt + item for item in post_prompt] - else: - prompts = [pre_prompt + post_prompt] - - def insert_separator(X, sep): - return [ - ele for sublist in zip(X, [sep] * len(X)) for ele in sublist - ][:-1] - - result = [] - for prompt in prompts: - prompt_chunks = [ - tokenizer(chunk).input_ids for chunk in prompt.split("") - ] - input_ids = [] - offset = 0 - if (len(prompt_chunks) > 0 and len(prompt_chunks[0]) > 0 - and prompt_chunks[0][0] == tokenizer.bos_token_id): - offset = 1 - input_ids.append(prompt_chunks[0][0]) - - for x in insert_separator(prompt_chunks, - [image_token_index] * (offset + 1)): - input_ids.extend(x[offset:]) - - input_ids = torch.tensor(input_ids, dtype=torch.long) - input_ids[input_ids == image_token_index] = 0 - result.append(input_ids) - - if not isinstance(post_prompt, list): - result = result[0].unsqueeze(0).expand(batch_size, -1) - return result - - def split_prompt_by_images(self, tensor): - batch_splits = [] - for batch in tensor: - # Find indices where value is zero () - zero_indices = (batch == 0).nonzero(as_tuple=False).squeeze(0) - # Add starting point for slicing - start_idx = 0 - splits = [] - for idx in zero_indices: - if start_idx != idx: # Ensure not slicing zero-length tensors - splits.append(batch[start_idx:idx].unsqueeze(0)) - start_idx = idx + 1 # Move start index past the zero - if start_idx < len( - batch): # Handle last segment if it's not zero-ending - splits.append(batch[start_idx:].unsqueeze(0)) - # Remove empty tensors resulting from consecutive zeros - splits = [split for split in splits if split.numel() > 0] - batch_splits.append(splits) - - return batch_splits - - def prepare_position_ids_for_cogvlm(self, input_ids): - batch_size = len(input_ids) - position_ids = torch.arange(input_ids.shape[1]) - position_ids[2:1227] = 2 - position_ids[1227:] = torch.arange(3, input_ids.shape[1] + 1 - 1225) - - position_ids = position_ids.to(torch.int32).to('cuda') - input_position_ids = [] - for i in range(batch_size): - input_position_ids.append(position_ids) - - return input_position_ids - - def generate(self, - pre_prompt, - post_prompt, - image, - decoder_input_ids, - max_new_tokens, - other_vision_inputs={}, - other_audio_inputs={}, - other_decoder_inputs={}): - profiler.start("Generate") - profiler.start("Preprocess") - if 'qwen2_vl' in self.model_type: - input_ids, input_lengths, ptuning_args, visual_features, mrope_args = self.preprocess( - pre_prompt, post_prompt, image, other_vision_inputs, - other_audio_inputs) - mrope_params = MropeParams( - mrope_rotary_cos_sin=mrope_args[0], - mrope_position_deltas=mrope_args[1], - ) - else: - input_ids, input_lengths, ptuning_args, visual_features, model_runner_input = self.preprocess( - pre_prompt, post_prompt, image, other_vision_inputs, - other_audio_inputs) - profiler.stop("Preprocess") - - # use prompt tuning to pass multimodal features - # model.generate() expects the following params (see layers/embedding.py): - # args[0]: prompt embedding table, [batch_size, multimodal_len, hidden_size], later flattened to [batch_size * multimodal_len, hidden_size] - # args[1]: prompt task ids, [batch_size]. in multimodal case, arange(batch_size), i.e. in VILA batching mode 2, each image is treated separately in the batch instead of concated together (although the prompt embedding table has to be concated) - # args[2]: prompt task vocab size, [1]. assuming all table has the same length, which in multimodal case equals to multimodal_len - profiler.start("LLM") - if self.decoder_llm and self.model_type != "mllama": - end_id = self.tokenizer.eos_token_id - if 'opt' in self.model_type and 'blip2' in self.model_type: - # For BLIP2-OPT, model outputs a "\n" at the end. - # we avoid it by using newline as the end token - end_id = self.tokenizer.encode("\n", - add_special_tokens=False)[0] - - if self.model_type == 'cogvlm': - input_position_ids = self.prepare_position_ids_for_cogvlm( - input_ids) - - prompt_tasks = None - prompt_table = None - if not self.cpp_e2e: - batch_size = len(input_ids) - prompt_tasks = ",".join( - np.arange(batch_size, dtype=np.int32).astype(str)) - prompt_table = torch.stack([ptuning_args[0]]) - prompt_table = prompt_table.view(batch_size, -1, - prompt_table.shape[-1]) - - output_ids = self.model.generate( - input_ids, - input_position_ids=input_position_ids - if self.model_type == 'cogvlm' else None, - mrope_params=mrope_params - if self.model_type == 'qwen2_vl' else None, - encoder_input_features=model_runner_input - if self.cpp_e2e else None, - sampling_config=None, - prompt_table=prompt_table, - prompt_tasks=prompt_tasks, - max_new_tokens=max_new_tokens, - end_id=end_id, - pad_id=self.tokenizer.pad_token_id - if self.tokenizer.pad_token_id is not None else - self.tokenizer.all_special_ids[0], - top_k=self.args.top_k, - top_p=self.args.top_p, - temperature=self.args.temperature, - repetition_penalty=self.args.repetition_penalty, - num_beams=self.args.num_beams, - lora_uids=self.args.lora_task_uids, - output_sequence_lengths=False, - return_dict=False, - mm_embedding_offloading=self.args.mm_embedding_offloading) - elif self.model_type == "mllama": - # When image is passed: - # the shape of visual_features is [bs, 1, 4, 1025, hidden_size] - # the shape of cross_attention_mask is [bs, decode_input_len, 1, 4] - # When image is None, create dummy visual_features and cross_attention_mask - if visual_features is None: - visual_features = torch.zeros([ - self.args.batch_size, 1, 4, 1, - self.model_config.hidden_size * self.runtime_mapping.tp_size - ], - dtype=self.model.dtype, - device=self.device) - dummy_cross_attention_mask = torch.zeros( - [self.args.batch_size, input_ids.shape[1], 1, 4], - dtype=bool, - device=self.device) - skip_cross_attn_blocks = torch.ones([1], - dtype=torch.bool, - device='cpu') - else: - skip_cross_attn_blocks = torch.zeros([1], - dtype=torch.bool, - device='cpu') - - visual_features = visual_features.to(self.model.dtype).chunk( - self.args.batch_size, dim=0) - encoder_input_features = [] - cross_attention_masks = [] - encoder_output_lengths = [] - for batch_idx in range(self.args.batch_size): - visual_feature = visual_features[batch_idx] - num_vision_tokens = visual_feature.shape[3] - visual_feature = visual_feature.reshape( - [-1, visual_feature.shape[-1]]) - encoder_max_input_length = visual_feature.shape[0] - encoder_input_lengths = torch.IntTensor( - [encoder_max_input_length]).to(visual_feature.device) - - # prepare cross_attention_mask of context phase - if 'cross_attention_mask' in other_decoder_inputs: - cross_attention_mask = other_decoder_inputs[ - 'cross_attention_mask'][batch_idx] - else: - cross_attention_mask = dummy_cross_attention_mask[batch_idx] - text_total_length, *_ = cross_attention_mask.shape - cross_attention_mask = cross_attention_mask.repeat_interleave( - num_vision_tokens, dim=2) - cross_attention_mask = cross_attention_mask.view( - text_total_length, -1) - cross_attention_mask = cross_attention_mask.unsqueeze(1) - cross_attention_mask = cross_attention_mask.to( - visual_feature.device).to(torch.bool).reshape( - [-1, cross_attention_mask.shape[-1]]) - - # prepare cross_attention_mask for generation phase and concat them - tmp_mask = [cross_attention_mask] + [ - cross_attention_mask[-1:, :] for _ in range(max_new_tokens) - ] - cross_attention_mask = torch.concat(tmp_mask) - - encoder_input_features.append(visual_feature) - cross_attention_masks.append(cross_attention_mask) - encoder_output_lengths.append(encoder_input_lengths) - - outputs = self.model.generate( - batch_input_ids=input_ids, - encoder_input_ids=None, - encoder_input_features=encoder_input_features, - encoder_output_lengths=encoder_output_lengths, - cross_attention_masks=cross_attention_masks, - max_new_tokens=max_new_tokens, - # max_attention_window_size=args.max_attention_window_size, - # sink_token_length=args.sink_token_length, - end_id=self.tokenizer.eos_token_id, - pad_id=self.tokenizer.pad_token_id, - temperature=self.args.temperature, - top_k=self.args.top_k, - top_p=self.args.top_p, - num_beams=self.args.num_beams, - # length_penalty=args.length_penalty, - # early_stopping=args.early_stopping, - # beam_width_array=args.beam_width_array, - repetition_penalty=self.args.repetition_penalty, - # presence_penalty=args.presence_penalty, - # frequency_penalty=args.frequency_penalty, - # stop_words_list=stop_words_list, - # bad_words_list=bad_words_list, - # output_cum_log_probs=(args.output_cum_log_probs_npy != None), - # output_log_probs=(args.output_log_probs_npy != None), - # random_seed=args.random_seed, - # lora_uids=args.lora_task_uids, - # prompt_table=args.prompt_table_path, - # prompt_tasks=args.prompt_tasks, - # streaming=args.streaming, - output_sequence_lengths=True, - # no_repeat_ngram_size=self.args.no_repeat_ngram_size, - return_dict=True, - # medusa_choices=args.medusa_choices, - # return_all_generated_tokens=args.return_all_generated_tokens, - # input_token_extra_ids=input_token_extra_ids, - encoder_max_input_length=encoder_max_input_length, - skip_cross_attn_blocks=skip_cross_attn_blocks, - ) - if mpi_rank() == 0: - output_ids = outputs["output_ids"] - else: - if self.model_type in ['nougat', 'pix2struct']: - # Trim encoder input_ids to match visual features shape - ids_shape = (self.args.batch_size, visual_features.shape[1]) - if self.model_type == 'nougat': - input_ids = torch.zeros(ids_shape, dtype=torch.int32) - elif self.model_type == 'pix2struct': - input_ids = torch.ones(ids_shape, dtype=torch.int32) - - output_ids = self.model.generate( - input_ids, - decoder_input_ids, - max_new_tokens, - num_beams=self.args.num_beams, - bos_token_id=self.tokenizer.bos_token_id, - pad_token_id=self.tokenizer.pad_token_id, - eos_token_id=self.tokenizer.eos_token_id, - debug_mode=False, - prompt_embedding_table=ptuning_args[0], - prompt_tasks=ptuning_args[1], - prompt_vocab_size=ptuning_args[2]) - - # Reset input_lengths to match decoder_input_ids - input_lengths = torch.ones(input_lengths.shape, - dtype=input_lengths.dtype) - profiler.stop("LLM") - - if mpi_rank() == 0: - # Extract a list of tensors of shape beam_width x output_ids. - profiler.start("Tokenizer decode") - output_beams_list = [ - self.tokenizer.batch_decode( - output_ids[batch_idx, :, input_lengths[batch_idx]:], - skip_special_tokens=True) for batch_idx in range( - min(self.args.batch_size, input_lengths.shape[0])) - ] - - stripped_text = [[ - output_beams_list[batch_idx][beam_idx].strip() - for beam_idx in range(self.args.num_beams) - ] for batch_idx in range( - min(self.args.batch_size, input_lengths.shape[0]))] - profiler.stop("Tokenizer decode") - profiler.stop("Generate") - return stripped_text - else: - profiler.stop("Generate") - return None - - def get_visual_features(self, image, other_vision_inputs): - visual_features = { - self.vision_input_names[0]: - image.to(str_dtype_to_torch(self.vision_precision)), - } - if self.model_type == "qwen2_vl": - other_vision_inputs['attention_mask'] = other_vision_inputs[ - 'attention_mask'].to(str_dtype_to_torch(self.vision_precision)) - for key, tensor in other_vision_inputs.items(): - visual_features.update({key: tensor}) - - tensor_info = [ - TensorInfo(self.vision_input_names[0], - str_dtype_to_trt(self.vision_precision), image.shape), - ] - for key, tensor in other_vision_inputs.items(): - tensor_info.append( - TensorInfo(key, torch_dtype_to_trt(tensor.dtype), tensor.shape)) - - visual_output_info = self.visual_encoder_session.infer_shapes( - tensor_info) - self.visual_encoder_session.set_shapes(visual_features) - visual_outputs = { - t.name: - torch.empty(tuple(t.shape), - dtype=trt_dtype_to_torch(t.dtype), - device=image.device) - for t in visual_output_info - } - - ok = self.visual_encoder_session.run(visual_features, visual_outputs, - self.stream.cuda_stream) - assert ok, "Runtime execution failed for vision encoder session" - self.stream.synchronize() - - image_embeds = visual_outputs[self.vision_output_names[0]] - - if self.args.mm_embedding_offloading: - # CUDA Stream Overlapping Requirements: - # 1. Both memory copy stream and kernel execution stream must be non-default streams - # 2. For host<->device transfers (H2D/D2H), host memory MUST be page-locked (pinned) - # NOTE: pinning is skipped under Confidential Compute - # (see maybe_pin_memory() and prefer_pinned()) - pinned_embeds = torch.empty_like(image_embeds, - device='cpu', - pin_memory=prefer_pinned()) - pinned_embeds.copy_(image_embeds, non_blocking=True) - image_embeds = pinned_embeds - - image_atts = torch.ones(image_embeds.size()[:-1], - dtype=torch.long).to(image.device) - - return image_embeds, image_atts - - def get_audio_features(self, audio, other_audio_inputs): - tmp_features, _ = self.audio_model.encoder( - audio.to(str_dtype_to_torch(self.audio_precision)), - other_audio_inputs['attention_mask']) - speech_out = self.audio_model.audio_projection['speech'](tmp_features) - vision_out = self.audio_model.audio_projection['vision'](tmp_features) - return torch.cat((speech_out, vision_out), dim=-1) - - def setup_fake_prompts_vila(self, batch_size, visual_features, - split_input_ids, input_lengths): - # visual_features (num_images, feature_len, token_embed) - # Assemble fake prompts which points to image embedding actually - fake_prompt_counter = self.model_config.vocab_size - if batch_size == 1: - # only check for multi-image inference (mode 1) - assert len(visual_features) <= len( - split_input_ids - ), "Unexpected number of visual features. Please check # in prompt and the #image files." - - input_ids = [] - if batch_size == 1: - # mode 1: multiple image as a whole, concat all prompts together, 
...
-            input_ids = [split_input_ids[0]]
-            for idx in range(
-                    len(visual_features)
-            ):  # TODO:alternatively make TensorInfo iterable if this breaks others
-                fake_prompt_id = torch.arange(
-                    fake_prompt_counter,
-                    fake_prompt_counter + visual_features.shape[1])
-                fake_prompt_counter += visual_features.shape[1]
-                fake_prompt_id = fake_prompt_id.unsqueeze(0)
-                input_ids.append(fake_prompt_id)
-                # in case no inter or post prompt
-                if len(split_input_ids) > idx + 1:
-                    input_ids.append(split_input_ids[idx + 1])
-            input_ids = torch.cat(input_ids, dim=1).contiguous().to(torch.int32)
-            input_ids = input_ids.reshape(batch_size, -1)
-
-        elif batch_size > 1:
-            # mode 2: each image have individual prompt, 

-            for idx in range(len(visual_features)):
-                input_ids.append(split_input_ids[idx][0])
-                fake_prompt_id = torch.arange(
-                    fake_prompt_counter,
-                    fake_prompt_counter + visual_features.shape[1])
-                fake_prompt_id = fake_prompt_id.unsqueeze(0)
-                input_ids.append(fake_prompt_id)
-                if len(split_input_ids[idx]) > 1:
-                    input_ids.append(split_input_ids[idx][1])
-            result = []
-            for i in range(0, len(input_ids), 3):
-                # Concatenate every 3 items (
, , )
-                concatenated = torch.cat(input_ids[i:i + 3],
-                                         dim=1).to(torch.int32).squeeze(0)
-                result.append(concatenated)
-            input_ids = result
-
-        if (self.decoder_llm
-                or self.runtime_mapping.is_first_pp_rank()) and isinstance(
-                    visual_features, torch.Tensor):
-            ptuning_args = self.ptuning_setup(visual_features, input_ids,
-                                              input_lengths)
-        else:
-            ptuning_args = [None, None, None]
-
-        return input_ids, ptuning_args
-
-    def setup_fake_prompts(self, visual_features, pre_input_ids, post_input_ids,
-                           input_lengths):
-        # Assemble fake prompts which points to image embedding actually
-        if hasattr(self, 'num_frames') and (visual_features.shape[1]
-                                            == self.num_frames):
-            visual_features = visual_features.view(visual_features.shape[0], -1,
-                                                   visual_features.shape[-1])
-
-        if visual_features is not None:
-            if self.python_e2e:
-                # Non-IFB Mode(used in python session): All requests in a batch have their prompt_table concatenated in
-                # a shape of (bs*vision_embedding_len, vision_hidden). So only one fake_prompt_id is needed for the
-                # entire batch, with values from 0 to bs * vision_embedding_len-1.
-                fake_prompt_id = torch.arange(
-                    self.model_config.vocab_size, self.model_config.vocab_size +
-                    visual_features.shape[0] * visual_features.shape[1])
-                fake_prompt_id = fake_prompt_id.reshape(
-                    visual_features.shape[0], visual_features.shape[1])
-            else:
-                # IFB Mode(used in c++ session): Each request's prompt_table is independent and requires a fake_prompt_id
-                # for each request, with values ranging from 0 to vision_embedding_len-1.
-                fake_prompt_id = torch.arange(
-                    self.model_config.vocab_size,
-                    self.model_config.vocab_size + visual_features.shape[1])
-                fake_prompt_id = fake_prompt_id.repeat(visual_features.shape[0],
-                                                       1)
-
-        if 'internvl' in self.model_type:
-            fake_prompt_id = fake_prompt_id.reshape(1, -1)
-
-        if 'cogvlm' in self.model_type:
-            input_ids = torch.cat(
-                [pre_input_ids[:, 0:1], fake_prompt_id, pre_input_ids[:, 1:]],
-                dim=1).contiguous().to(torch.int32)
-        elif self.model_type == 'mllama':
-            input_ids = pre_input_ids.contiguous().to(torch.int32)
-        else:
-            if post_input_ids is not None:
-                if isinstance(post_input_ids, list):
-                    pre_input_fake_prompt_ids = [
-                        pre_input_ids[:len(fake_prompt_id)], fake_prompt_id
-                    ]
-                    pre_input_fake_prompt_ids = torch.cat(
-                        pre_input_fake_prompt_ids,
-                        dim=1).contiguous().to(torch.int32)
-                    input_ids = [
-                        torch.cat((pre_input_fake_prompt_id,
-                                   post_input_id)).contiguous().to(torch.int32)
-                        for pre_input_fake_prompt_id, post_input_id in zip(
-                            pre_input_fake_prompt_ids, post_input_ids)
-                    ]
-                else:
-                    input_ids = [pre_input_ids, fake_prompt_id, post_input_ids]
-                    input_ids = torch.cat(input_ids,
-                                          dim=1).contiguous().to(torch.int32)
-            else:
-                input_ids = [fake_prompt_id, pre_input_ids]
-                input_ids = torch.cat(input_ids,
-                                      dim=1).contiguous().to(torch.int32)
-
-        if (self.decoder_llm or self.runtime_mapping.is_first_pp_rank()
-            ) and self.model_type != "mllama" and isinstance(
-                visual_features, torch.Tensor):
-            ptuning_args = self.ptuning_setup(visual_features, input_ids,
-                                              input_lengths)
-        else:
-            ptuning_args = [None, None, None]
-
-        return input_ids, ptuning_args
-
-    def get_rope_index(
-        self,
-        input_ids: torch.IntTensor,
-        image_grid_thw: Optional[torch.LongTensor] = None,
-        video_grid_thw: Optional[torch.LongTensor] = None,
-        attention_mask: Optional[torch.Tensor] = None,
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
-        """
-        Calculate the 3D rope index based on image and video's temporal, height and width in LLM.
-
-        Explanation:
-            Each embedding sequence contains vision embedding and text embedding or just contains text embedding.
-
-            For pure text embedding sequence, the rotary position embedding has no difference with modern LLMs.
-            Examples:
-                input_ids: [T T T T T], here T is for text.
-                temporal position_ids: [0, 1, 2, 3, 4]
-                height position_ids: [0, 1, 2, 3, 4]
-                width position_ids: [0, 1, 2, 3, 4]
-
-            For vision and text embedding sequence, we calculate 3D rotary position embedding for vision part
-            and 1D rotary position embedding for text part.
-            Examples:
-                Assume we have a video input with 3 temporal patches, 2 height patches and 2 width patches.
-                input_ids: [V V V V V V V V V V V V T T T T T], here V is for vision.
-                vision temporal position_ids: [0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2]
-                vision height position_ids: [0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1]
-                vision width position_ids: [0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1]
-                text temporal position_ids: [3, 4, 5, 6, 7]
-                text height position_ids: [3, 4, 5, 6, 7]
-                text width position_ids: [3, 4, 5, 6, 7]
-                Here we calculate the text start position_ids as the max vision position_ids plus 1.
-
-        Args:
-            input_ids (`torch.IntTensor` of shape `(batch_size, sequence_length)`):
-                Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide
-                it.
-            image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
-                The temporal, height and width of feature shape of each image in LLM.
-            video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*):
-                The temporal, height and width of feature shape of each video in LLM.
-            attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
-                Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
-
-                - 1 for tokens that are **not masked**,
-                - 0 for tokens that are **masked**.
-
-        Returns:
-            position_ids (`torch.IntTensor` of shape `(3, batch_size, sequence_length)`)
-            mrope_position_deltas (`torch.Tensor` of shape `(batch_size)`)
-        """
-        spatial_merge_size = self.spatial_merge_size
-        image_token_id = self.image_token_id
-        video_token_id = self.video_token_id
-        vision_start_token_id = self.vision_start_token_id
-        mrope_position_deltas = []
-        if image_grid_thw is not None or video_grid_thw is not None:
-            total_input_ids = input_ids
-            position_ids = torch.ones(3,
-                                      input_ids.shape[0],
-                                      input_ids.shape[1],
-                                      dtype=input_ids.dtype,
-                                      device=input_ids.device)
-            image_index, video_index = 0, 0
-            for i, input_ids in enumerate(total_input_ids):
-                if attention_mask is not None:
-                    input_ids = input_ids[attention_mask[i] == 1]
-                image_nums, video_nums = 0, 0
-                vision_start_indices = torch.argwhere(
-                    input_ids == vision_start_token_id).squeeze(1)
-                vision_tokens = input_ids[vision_start_indices + 1]
-                image_nums = (vision_tokens == image_token_id).sum()
-                video_nums = (vision_tokens == video_token_id).sum()
-                input_tokens = input_ids.tolist()
-                llm_pos_ids_list: list = []
-                st = 0
-                remain_images, remain_videos = image_nums, video_nums
-                for _ in range(image_nums + video_nums):
-                    if image_token_id in input_tokens and remain_images > 0:
-                        ed_image = input_tokens.index(image_token_id, st)
-                    else:
-                        ed_image = len(input_tokens) + 1
-                    if video_token_id in input_tokens and remain_videos > 0:
-                        ed_video = input_tokens.index(video_token_id, st)
-                    else:
-                        ed_video = len(input_tokens) + 1
-                    if ed_image < ed_video:
-                        t, h, w = (
-                            image_grid_thw[image_index][0],
-                            image_grid_thw[image_index][1],
-                            image_grid_thw[image_index][2],
-                        )
-                        image_index += 1
-                        remain_images -= 1
-                        ed = ed_image
-                    else:
-                        t, h, w = (
-                            video_grid_thw[video_index][0],
-                            video_grid_thw[video_index][1],
-                            video_grid_thw[video_index][2],
-                        )
-                        video_index += 1
-                        remain_videos -= 1
-                        ed = ed_video
-                    llm_grid_t, llm_grid_h, llm_grid_w = (
-                        t.item(),
-                        h.item() // spatial_merge_size,
-                        w.item() // spatial_merge_size,
-                    )
-                    text_len = ed - st
-
-                    st_idx = llm_pos_ids_list[-1].max() + 1 if len(
-                        llm_pos_ids_list) > 0 else 0
-                    llm_pos_ids_list.append(
-                        torch.arange(text_len).view(1, -1).expand(3, -1) +
-                        st_idx)
-
-                    t_index = torch.arange(llm_grid_t).view(-1, 1).expand(
-                        -1, llm_grid_h * llm_grid_w).flatten()
-                    h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(
-                        llm_grid_t, -1, llm_grid_w).flatten()
-                    w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(
-                        llm_grid_t, llm_grid_h, -1).flatten()
-                    llm_pos_ids_list.append(
-                        torch.stack([t_index, h_index, w_index]) + text_len +
-                        st_idx)
-                    st = ed + llm_grid_t * llm_grid_h * llm_grid_w
-
-                if st < len(input_tokens):
-                    st_idx = llm_pos_ids_list[-1].max() + 1 if len(
-                        llm_pos_ids_list) > 0 else 0
-                    text_len = len(input_tokens) - st
-                    llm_pos_ids_list.append(
-                        torch.arange(text_len).view(1, -1).expand(3, -1) +
-                        st_idx)
-
-                llm_positions = torch.cat(llm_pos_ids_list,
-                                          dim=1).reshape(3, -1)
-                position_ids[..., i, attention_mask[i] == 1] = llm_positions.to(
-                    position_ids.device)
-                mrope_position_deltas.append(llm_positions.max() + 1 -
-                                             len(total_input_ids[i]))
-            mrope_position_deltas = torch.tensor(
-                mrope_position_deltas, device=input_ids.device).unsqueeze(1)
-            return position_ids, mrope_position_deltas
-        else:
-            if attention_mask is not None:
-                position_ids = attention_mask.long().cumsum(-1) - 1
-                position_ids.masked_fill_(attention_mask == 0, 1)
-                position_ids = position_ids.unsqueeze(0).expand(3, -1, -1).to(
-                    input_ids.device)
-                max_position_ids = position_ids.max(0, keepdim=False)[0].max(
-                    -1, keepdim=True)[0]
-                mrope_position_deltas = max_position_ids + 1 - attention_mask.shape[
-                    -1]
-            else:
-                position_ids = (torch.arange(input_ids.shape[1],
-                                             device=input_ids.device).view(
-                                                 1, 1, -1).expand(
-                                                     3, input_ids.shape[0], -1))
-                mrope_position_deltas = torch.zeros(
-                    [input_ids.shape[0], 1],
-                    device=input_ids.device,
-                    dtype=input_ids.dtype,
-                )
-
-            return position_ids, mrope_position_deltas
-
-    def setup_fake_prompts_qwen2vl(self, visual_features, input_ids,
-                                   vision_grid_thws, attention_mask,
-                                   input_lengths):
-
-        visual_features = torch.unsqueeze(visual_features, 0)
-
-        # Get the rope index
-        # From HF's preprocess code
-        mrope_position_ids, mrope_position_deltas = self.get_rope_index(
-            input_ids,
-            image_grid_thw=vision_grid_thws,
-            video_grid_thw=None,
-            attention_mask=attention_mask,
-        )
-
-        # This is where we convert input_ids of image features into fake_prompt_ids mapping for TRT-LLM engine.
-        masks = (input_ids == self.image_token_id) | (
-            input_ids == self.vision_token_id) | (input_ids
-                                                  == self.video_token_id)
-        cumulative_counts = masks.cumsum(dim=1)
-        values = (self.model_config.vocab_size - 1) + cumulative_counts
-        input_ids[masks] = values[masks]
-
-        if self.decoder_llm or self.runtime_mapping.is_first_pp_rank():
-            ptuning_args = self.ptuning_setup(visual_features, input_ids,
-                                              input_lengths)
-        else:
-            ptuning_args = [None, None, None]
-
-        # This does not have dependency on input.
-        # Switch to attributes to use across iterations.
-        if not hasattr(self, 'rotary_cos_sin'):
-            inv_freq, rotary_cos_sin = RopeEmbeddingUtils.create_sinusoidal_positions_for_attention_plugin(
-                num_pos=self.max_position_embeddings,
-                dim=int(self.hidden_size / self.num_attention_heads),
-                theta=float(self.rope_theta),
-                scale_type=RotaryScalingType.mrope)
-            self.rotary_cos_sin = torch.from_numpy(rotary_cos_sin).to(
-                visual_features.device)
-            self.rotary_cos_sin = self.rotary_cos_sin.reshape(
-                self.max_position_embeddings,
-                int(self.hidden_size / self.num_attention_heads / 2), 2)
-            self.cos_ori = self.rotary_cos_sin[:, :, 0]
-            self.sin_ori = self.rotary_cos_sin[:, :, 1]
-
-        mrope_position_ids = mrope_position_ids.transpose(1, 0)
-        mrope_position_ids_padding = torch.zeros(
-            mrope_position_ids.shape[:-1] + (self.max_position_embeddings, ),
-            dtype=torch.int32,
-            device=visual_features.device)
-        mrope_position_ids_padding[:, :, :mrope_position_ids.
-                                   shape[-1]] = mrope_position_ids
-        cos = self.cos_ori[mrope_position_ids_padding]
-        sin = self.sin_ori[mrope_position_ids_padding]
-
-        mrope_section = [16, 24, 24]
-        cos = torch.cat([
-            m[:, i % 3] for i, m in enumerate(cos.split(mrope_section, dim=-1))
-        ],
-                        dim=-1).unsqueeze(-1)
-        sin = torch.cat([
-            m[:, i % 3] for i, m in enumerate(sin.split(mrope_section, dim=-1))
-        ],
-                        dim=-1).unsqueeze(-1)
-        concat_cos_sin = torch.concatenate((cos, sin), axis=-1)
-        concat_cos_sin = concat_cos_sin.reshape(concat_cos_sin.shape[0], -1)
-
-        mrope_args = [concat_cos_sin, mrope_position_deltas]
-        return input_ids, ptuning_args, mrope_args
-
-    def ptuning_setup_fuyu(self, input_ids, image_patches_indices):
-        res_input_ids = []
-        for cur_input_ids, cur_image_patches_indices in zip(
-                input_ids, image_patches_indices):
-            # Truncate input_ids to the length of image_patches_indices
-            cur_image_patches_indices = cur_image_patches_indices[:len(
-                cur_input_ids)]
-            # Get ids of the image_patches
-            non_zero_mask = cur_image_patches_indices != -1
-            # Replace input_ids with image_patches_indices values (where the patches are placed)
-            cur_input_ids = cur_input_ids.masked_scatter(
-                non_zero_mask,
-                cur_image_patches_indices[non_zero_mask] +
-                self.model_config.vocab_size,
-            )
-            res_input_ids.append(cur_input_ids)
-        return res_input_ids
-
-    def ptuning_setup_pixtral(self, input_ids):
-        # input_ids obtained from processor has token_ids for text as well as image tokens
-        # where each image token is represented by the same image_token_index.
-        image_token_index = self.image_token_index
-        vocab_size = self.vocab_size
-        # Replace all image tokens with a unique token_id > text_vacab_size.
-        # This shall be used to lookup the prompt table.
-        for img_idx in range(self.args.batch_size):
-            # Note: We reset replacer to text_vocab_size for each sample. This is as opposed to doing `replacer = vocab_size + img_idx * tokens_per_task`.
-            # That part of the look-up manipulation is done by the `task_ids` input to PromptEmbedding forward.
-            replacer = vocab_size
-            for token_idx in range(len(input_ids[img_idx])):
-                if input_ids[img_idx][token_idx] == image_token_index:
-                    input_ids[img_idx][token_idx] = replacer
-                    replacer += 1
-        return input_ids
-
-    def ptuning_setup_llava_next(self, visual_features, pre_prompt,
-                                 post_prompt):
-        input_ids = []
-        fake_prompt_ids = list(
-            range(self.model_config.vocab_size,
-                  self.model_config.vocab_size + visual_features.shape[0]))
-        input_ids = self.tokenizer.encode(
-            pre_prompt[0]) + fake_prompt_ids + self.tokenizer.encode(
-                post_prompt[0])[self.tokenizer.add_bos_token:]
-        input_ids = [input_ids] * len(pre_prompt)
-        input_ids = torch.tensor(input_ids)
-        return input_ids
-
-    def ptuning_setup_phi3(self, visual_features, audio_features, input_ids,
-                           num_img_tokens, num_aud_tokens):
-        fake_prompt_id = torch.arange(
-            self.model_config.vocab_size,
-            self.model_config.vocab_size + visual_features.shape[0])
-        if self.model_type == "phi-3-vision":
-            MAX_INPUT_ID = int(1e9)
-            positions = torch.nonzero(
-                (input_ids < 0) & (input_ids > -MAX_INPUT_ID), as_tuple=False)
-        elif self.model_type == "phi-4-multimodal":
-            IMAGE_TOKEN_ID = 200010
-            positions = torch.nonzero(input_ids == IMAGE_TOKEN_ID,
-                                      as_tuple=False)
-        idx = 0
-        for _, cnt in enumerate(num_img_tokens):
-            input_ids[positions[idx, 0], positions[idx, 1]:positions[idx, 1] +
-                      cnt] = fake_prompt_id[idx:idx + cnt]
-            idx += cnt
-
-        if self.model_type == "phi-4-multimodal" and audio_features is not None:
-            prompt_id_offset = self.model_config.vocab_size + visual_features.shape[
-                0]
-            fake_prompt_id = torch.arange(
-                prompt_id_offset, prompt_id_offset + audio_features.shape[0])
-            AUDIO_TOKEN_ID = 200011
-            positions = torch.nonzero(input_ids == AUDIO_TOKEN_ID,
-                                      as_tuple=False)
-            idx = 0
-            for _, cnt in enumerate(num_aud_tokens):
-                input_ids[positions[idx, 0], positions[idx,
-                                                       1]:positions[idx, 1] +
-                          cnt] = fake_prompt_id[idx:idx + cnt]
-                idx += cnt
-        return input_ids
-
-    def ptuning_setup(self, prompt_table, input_ids, input_lengths):
-        hidden_size = self.model_config.hidden_size * self.runtime_mapping.tp_size
-        if prompt_table is not None:
-            task_vocab_size = torch.tensor(
-                [prompt_table.shape[1]],
-                dtype=torch.int32,
-            ).cuda()
-            prompt_table = prompt_table.view(
-                (prompt_table.shape[0] * prompt_table.shape[1],
-                 prompt_table.shape[2]))
-
-            assert prompt_table.shape[
-                1] == hidden_size, "Prompt table dimensions do not match hidden size"
-
-            if hasattr(self.model_config, 'dtype'):
-                prompt_table = prompt_table.cuda().to(
-                    dtype=str_dtype_to_torch(self.model_config.dtype))
-            else:
-                if self.args.mm_embedding_offloading:
-                    # CUDA Stream Overlapping Requirements:
-                    # 1. Both memory copy stream and kernel execution stream must be non-default streams
-                    # 2. For host<->device transfers (H2D/D2H), host memory MUST be page-locked (pinned)
-                    # NOTE: pinning is skipped under Confidential Compute
-                    # (see maybe_pin_memory() and prefer_pinned())
-                    prompt_table = maybe_pin_memory(prompt_table).to(
-                        dtype=self.model.dtype)
-                else:
-                    prompt_table = prompt_table.cuda().to(
-                        dtype=self.model.dtype)
-        else:
-            prompt_table = torch.empty([1, hidden_size]).cuda()
-            task_vocab_size = torch.zeros([1]).cuda()
-
-        remove_input_padding = self.model_config.remove_input_padding if hasattr(
-            self.model_config,
-            'remove_input_padding') else self.model_config.use_packed_input
-        if remove_input_padding:
-            tasks = torch.zeros([torch.sum(input_lengths)],
-                                dtype=torch.int32).cuda()
-            if self.decoder_llm: tasks = tasks.unsqueeze(0)
-        else:
-            if not isinstance(input_ids, list):
-                tasks = torch.zeros(input_ids.shape, dtype=torch.int32).cuda()
-            else:
-                max_length = max(input_id.size(-1) for input_id in input_ids)
-                tasks = torch.zeros((len(input_ids), max_length),
-                                    dtype=torch.int32).cuda()
-
-        return [prompt_table, tasks, task_vocab_size]
-
-    def load_test_data(self, image_path=None, video_path=None):
-
-        def load_images(image_paths):
-            if isinstance(image_paths, str):
-                image_paths = [image_paths]
-            images = []
-            for image_path in image_paths:
-                if image_path.startswith("http") or image_path.startswith(
-                        "https"):
-                    logger.info(f"downloading image from url {image_path}")
-                    try:
-                        response = requests.get(image_path, timeout=5)
-                        response.raise_for_status()
-                        if 'image' not in response.headers.get(
-                                'Content-Type', ''):
-                            raise Exception(
-                                f"URL does not point to an image: {image_path}."
-                            )
-                        image = Image.open(BytesIO(
-                            response.content)).convert("RGB")
-                    except (UnidentifiedImageError, IOError):
-                        raise Exception(
-                            f"Cannot identify image file at URL: {image_path}.")
-                    except Exception as e:
-                        raise Exception(
-                            f"Failed to download image from url {image_path}: {e}"
-                        )
-                else:
-                    image = Image.open(image_path).convert("RGB")
-                images.append(image)
-            return images if len(images) > 1 else images[0]
-
-        if "vila" in self.model_type:
-            if image_path is None:
-                img_urls = [
-                    'https://github.com/NVlabs/VILA/blob/6b941da19e31ddfdfaa60160908ccf0978d96615/demo_images/av.png?raw=true',
-                    'https://storage.googleapis.com/sfr-vision-language-research/LAVIS/assets/merlion.png'
-                ] * 4
-
-                img_urls = img_urls[:self.args.batch_size]
-                self.args.image_path = ",".join(img_urls)
-                images = load_images(img_urls)
-            else:
-                if isinstance(image_path, str):
-                    image_path = image_path.split(self.args.path_sep)
-                images = load_images(image_path)
-        elif "pixtral" in self.model_type:
-            if image_path is None:
-                image_urls = [
-                    "https://storage.googleapis.com/sfr-vision-language-research/LAVIS/assets/merlion.png",
-                    "https://www.ilankelman.org/stopsigns/australia.jpg",
-                    "https://huggingface.co/microsoft/kosmos-2-patch14-224/resolve/main/snowman.png",
-                    "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg",
-                ]
-                while len(image_urls) < self.args.batch_size:
-                    image_urls *= 2
-                image_urls = image_urls[:self.args.batch_size]
-                self.args.image_path = ",".join(image_urls)
-                images = load_images(image_urls)
-            else:
-                if isinstance(image_path, str):
-                    image_path = image_path.split(self.args.path_sep)
-                images = load_images(image_path)
-            images = [images] if not isinstance(images, list) else images
-        elif "nougat" in self.model_type:
-            filepath = hf_hub_download(
-                repo_id="hf-internal-testing/fixtures_docvqa",
-                filename="nougat_paper.png",
-                revision="ec57bf8c8b1653a209c13f6e9ee66b12df0fc2db",
-                repo_type="dataset")
-            images = Image.open(filepath)
-        elif "fuyu" in self.model_type:
-            filepath = hf_hub_download(repo_id="adept/fuyu-8b",
-                                       filename="skateboard.png",
-                                       repo_type='model')
-            images = Image.open(filepath)
-        elif "kosmos" in self.model_type:
-            if image_path is None:
-                image_path = 'https://huggingface.co/microsoft/kosmos-2-patch14-224/resolve/main/snowman.png'
-            images = load_images(image_path)
-        elif "pix2struct" in self.model_type:
-            if image_path is None:
-                image_path = 'https://raw.githubusercontent.com/vis-nlp/ChartQA/main/ChartQA%20Dataset/val/png/multi_col_40963.png'
-            images = load_images(image_path)
-        elif "video-neva" in self.model_type:
-            images = video_path
-        elif "internlm" in self.model_type:
-            img_url = "https://huggingface.co/internlm/internlm-xcomposer2-vl-7b/resolve/main/image1.webp"
-            images = load_images(img_url)
-        elif "internvl" in self.model_type:
-            if image_path is None:
-                img_url = 'https://huggingface.co/OpenGVLab/InternVL2-4B/resolve/main/examples/image1.jpg'
-                images = Image.open(
-                    requests.get(img_url, stream=True,
-                                 timeout=5).raw).convert('RGB')
-            else:
-                images = Image.open(image_path).convert('RGB')
-        elif "qwen2_vl" in self.model_type:
-            images = []
-            if self.args.image_path is None:
-                img_url = 'https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg'
-                image = Image.open(
-                    requests.get(img_url, stream=True,
-                                 timeout=5).raw).convert('RGB')
-                image = image.resize((504, 504))
-                images.append(image)
-            else:
-                images = []
-                for image_path in self.args.image_path:
-                    image = Image.open(image_path).convert('RGB')
-                    image = image.resize((504, 504))
-                    images.append(image)
-        elif "llava_onevision" in self.model_type and self.args.video_path is not None:
-            if self.args.video_path == 'llava-onevision-accuracy':
-                self.args.video_path = hf_hub_download(
-                    repo_id="raushan-testing-hf/videos-test",
-                    filename="sample_demo_1.mp4",
-                    repo_type="dataset")
-            import av
-            with av.open(self.args.video_path) as container:
-                total_frames = container.streams.video[0].frames
-                assert total_frames >= self.num_frames
-                indices = np.arange(0, total_frames,
-                                    total_frames / self.num_frames).astype(int)
-                frames = []
-                container.seek(0)
-                start_index = indices[0]
-                end_index = indices[-1]
-                for i, frame in enumerate(container.decode(video=0)):
-                    if i > end_index:
-                        break
-                    if i >= start_index and i in indices:
-                        frames.append(frame)
-                images = np.stack(
-                    [x.to_ndarray(format="rgb24") for x in frames])
-            images = torch.tensor(images)
-        else:
-            if self.model_type != 'mllama':
-                if image_path is None:
-                    if self.model_type == "llava":
-                        image_path = [
-                            'https://storage.googleapis.com/sfr-vision-language-research/LAVIS/assets/merlion.png'
-                        ] * 8
-                        image_path = image_path[:self.args.batch_size]
-                        self.args.image_path = ",".join(image_path)
-                    else:
-                        image_path = 'https://storage.googleapis.com/sfr-vision-language-research/LAVIS/assets/merlion.png'
-                else:
-                    if isinstance(image_path, str):
-                        image_path = image_path.split(self.args.path_sep)
-
-            images = load_images(image_path) if image_path is not None else None
-        return images
-
-    def load_test_audio(self, audio_path):
-        if self.model_type != "phi-4-multimodal":
-            return None
-
-        assert audio_path is not None
-        import soundfile
-        audio = soundfile.read(audio_path)
-        return audio
-
-    def setup_inputs(self, input_text, raw_image, raw_audio=None):
-        from ..tools.multimodal_builder import compute_rotary_pos_emb
-        other_vision_inputs = {}
-        other_audio_inputs = {}
-        other_decoder_inputs = {}
-        if self.model_type not in ['qwen2_vl', 'vila', 'llava']:
-            input_text = input_text[0] if isinstance(input_text,
-                                                     list) else input_text
-
-        if 'blip2' in self.model_type:
-            image = self.processor(raw_image, input_text,
-                                   return_tensors="pt")['pixel_values']
-            if input_text is None:
-                input_text = "Question: which city is this? Answer:"
-            pre_prompt = input_text
-            post_prompt = None
-        elif 'internlm' in self.model_type:
-            #Feed the raw image into vis_processor, to get processed image
-            image = self.processor(raw_image).unsqueeze(0).cuda()
-            if input_text is None:
-                input_text = "Please describe this image in detail."
-            pre_prompt = ''
-            meta_instruction = 'You are an AI assistant whose name is InternLM-XComposer (浦语·灵笔).\n'
-            '- InternLM-XComposer (浦语·灵笔) is a multi-modality conversational language model that is developed by Shanghai AI Laboratory (上海人工智能实验室). It is designed to be helpful, honest, and harmless.\n'
-            '- InternLM-XComposer (浦语·灵笔) can understand and communicate fluently in the language chosen by the user such as English and 中文.\n'
-            '- InternLM-XComposer (浦语·灵笔) is capable of comprehending and articulating responses effectively based on the provided image.',
-            pre_prompt += f"""[UNUSED_TOKEN_146]system\n{meta_instruction}[UNUSED_TOKEN_145]\n"""
-            pre_prompt += f"""[UNUSED_TOKEN_146]user\n"""
-            post_prompt = f"""{input_text}[UNUSED_TOKEN_145]\n[UNUSED_TOKEN_146]assistant\n"""
-        elif 'qwen2_vl' in self.model_type:
-            from qwen_vl_utils import process_vision_info
-            from transformers.models.qwen2_vl.modeling_qwen2_vl import \
-                VisionRotaryEmbedding
-            hf_config = AutoConfig.from_pretrained(self.args.hf_model_dir)
-            if input_text is None:
-                input_text = ["Question: Describe this image. Answer:"
-                              ] * self.args.batch_size
-            messages = [[{
-                "role":
-                "user",
-                "content": [
-                    {
-                        "type": "image",
-                        "image": raw_image[idx],
-                    },
-                    {
-                        "type": "text",
-                        "text": input_text[idx],
-                    },
-                ],
-            }] for idx in range(self.args.batch_size)]
-
-            texts = [
-                self.processor.apply_chat_template(msg,
-                                                   tokenize=False,
-                                                   add_generation_prompt=True)
-                for msg in messages
-            ]
-            image_inputs, video_inputs = process_vision_info(messages)
-            inputs = self.processor(
-                text=texts,
-                images=image_inputs,
-                videos=video_inputs,
-                padding=True,
-                return_tensors="pt",
-            )
-            inputs = inputs.to(self.device)
-            image = inputs['pixel_values']
-            image_grid_thw = inputs['image_grid_thw']
-            input_ids = inputs['input_ids']
-            attention_mask = inputs['attention_mask']
-            cu_seqlens = torch.repeat_interleave(
-                image_grid_thw[:, 1] * image_grid_thw[:, 2],
-                image_grid_thw[:, 0]).cumsum(dim=0, dtype=torch.int32)
-            cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0)
-
-            seq_length = image.shape[0]
-            # Create block indices using bucketing
-            block_indices = torch.bucketize(torch.arange(seq_length,
-                                                         device=image.device),
-                                            cu_seqlens,
-                                            right=True) - 1
-
-            # Generate block diagonal mask using matrix expansion
-            attention_mask_vit = torch.where(
-                block_indices.view(-1, 1) == block_indices.view(1, -1),
-                torch.zeros((), device=image.device, dtype=image.dtype),
-                torch.full((),
-                           torch.finfo(torch.float16).min,
-                           device=image.device,
-                           dtype=image.dtype)).unsqueeze(0)
-
-            decoder_input_ids = None
-            post_prompt = None
-            pre_prompt = None
-            images_qwenvl = {
-                "image": image,
-                "input_ids": input_ids,
-            }
-            rotary_pos_emb = compute_rotary_pos_emb(
-                image_grid_thw, hf_config, VisionRotaryEmbedding).to("cuda")
-            other_vision_inputs['attention_mask_llm'] = attention_mask
-            other_vision_inputs['image_grid_thw'] = image_grid_thw
-            other_vision_inputs['attention_mask'] = attention_mask_vit
-            other_vision_inputs['rotary_pos_emb'] = rotary_pos_emb
-            return input_text, pre_prompt, post_prompt, images_qwenvl, decoder_input_ids, other_vision_inputs, other_audio_inputs, other_decoder_inputs
-        elif 'nougat' in self.model_type:
-            image = self.processor(raw_image,
-                                   return_tensors="pt")['pixel_values']
-            # Nougat doesn't need text prompt (mBART use single token to start generation), just leave a dummy one here
-            if input_text is None:
-                input_text = "Question: which city is this? Answer:"
-            pre_prompt = input_text
-            post_prompt = None
-
-        elif 'cogvlm' in self.model_type:
-            image = self.transform(raw_image).unsqueeze(0)
-            if input_text is None:
-                input_text = " [INST] which city is this? [/INST] "
-            pre_prompt = input_text
-            post_prompt = None
-
-        elif 'phi-3-vision' in self.model_type:
-            pre_prompt = "<|user|>\n<|image_1|>\n"
-            if input_text is None:
-                input_text = "Which city is this?"
-            post_prompt = input_text + "<|end|>\n<|assistant|>\n"
-            prompt = pre_prompt + post_prompt
-            image = self.processor(text=prompt,
-                                   images=raw_image,
-                                   return_tensors="pt")
-        elif 'phi-4-multimodal' in self.model_type:
-            pre_prompt = "<|user|><|image_1|><|audio_1|>"
-            post_prompt = "<|end|><|assistant|>"
-            prompt = pre_prompt + post_prompt
-            image = self.processor(text=prompt,
-                                   images=[raw_image],
-                                   audios=[raw_audio],
-                                   return_tensors="pt")
-
-        elif 'pixtral' in self.model_type:
-            # Send image and text prompt to processor.
-            pre_prompt = "[INST][IMG]"
-            if input_text is None:
-                input_text = "What is in the image?"
-            post_prompt = "[/INST]"
-            prompt = pre_prompt + input_text + post_prompt
-            dtype = str_dtype_to_torch(self.vision_precision)
-            image = {'pixel_values': [], 'input_ids': []}
-            for img_idx in range(self.args.batch_size):
-                image_info = self.processor(text=prompt,
-                                            images=[raw_image[img_idx]],
-                                            return_tensors="pt").to(dtype)
-                image['pixel_values'].append(image_info['pixel_values'].to(
-                    self.device))
-                image['input_ids'].append(image_info['input_ids'][0].to(
-                    self.device))
-
-        elif 'internvl' in self.model_type:
-            pre_prompt = "<|system|>\n你是由上海人工智能实验室联合商汤科技开发的书生多模态大模型，英文名叫InternVL, 是一个有用无害的人工智能助手。<|end|><|user|>\n\n"
-            if input_text is None:
-                input_text = "Please describe the image shortly."
-            post_prompt = input_text + "<|end|><|assistant|>\n"
-            prompt = pre_prompt + post_prompt
-            image = self.processor(images=raw_image,
-                                   return_tensors='pt').pixel_values
-
-        elif self.model_type == "pix2struct":
-            if input_text is None:
-                input_text = ""
-            inputs = self.processor(
-                images=raw_image,
-                text=input_text,
-                return_tensors="pt",
-            )
-            image = inputs['flattened_patches']
-            image = image.expand(self.args.batch_size, -1, -1).contiguous()
-            pre_prompt = ""
-            post_prompt = None
-
-        elif self.model_type == "neva":
-            image = self.transform(raw_image).unsqueeze(0)
-
-            if input_text is None:
-                input_text = "Hi! What is in this image?"
-
-            pre_prompt = "System\n\nUser\n"
-            post_prompt = f"\n{input_text}\nAssistant\n"
-
-        elif self.model_type == "video-neva":
-
-            image = self.video_preprocess(
-                raw_image)  # shape (1, num_frames, 3, H, W)
-
-            if input_text is None:
-                input_text = "Hi! What is in this video?"
-
-            # SteerLM prompt template
-            pre_prompt = """System\nA chat between a curious user and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user's questions.\n\nUser"""
-            post_prompt = f"\n{input_text}\nAssistant\nquality:4,toxicity:0,humor:0,creativity:0,helpfulness:4,correctness:4,coherence:4,complexity:4,verbosity:4\n" ""
-
-        elif self.model_type == "llava_next":
-            if self.llm_name == "mistralai/Mistral-7B-Instruct-v0.2":
-                pre_prompt = "[INST] "
-                if input_text is None:
-                    input_text = "Question: which city is this? Answer:"
-                post_prompt = f"\n{input_text} [/INST]"
-                prompt = pre_prompt + post_prompt
-
-            elif self.llm_name == "NousResearch/Nous-Hermes-2-Yi-34B":
-                pre_prompt = "<|im_start|>system\nAnswer the questions.<|im_end|><|im_start|>user\n"
-                if input_text is None:
-                    input_text = "Question: which city is this? Answer:"
-                post_prompt = f"\n{input_text}<|im_end|><|im_start|>assistant\n"
-                prompt = pre_prompt + post_prompt
-
-            else:
-                raise Exception(
-                    f"Prompt template for {self.llm_name} for not included currently"
-                )
-
-            image = self.processor(text=prompt,
-                                   images=raw_image,
-                                   return_tensors="pt")
-
-        elif self.model_type == 'vila':
-            if input_text is None:
-                input_text = "\n Please elaborate what you see in the images?"
-            if '8b' in self.args.hf_model_dir.lower():
-                pre_prompt = "<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n\nYou are a helpful language and vision assistant. You are able to understand the visual content that the user provides, and assist the user with a variety of tasks using natural language.<|eot_id|><|start_header_id|>user<|end_header_id|>\n\n"
-                post_prompt = "<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
-            elif '40b' in self.args.hf_model_dir.lower():
-                pre_prompt = "<|im_start|>system\nAnswer the questions.<|im_end|><|im_start|>user\n"
-                post_prompt = "<|im_end|><|im_start|>assistant\n"
-            else:
-                pre_prompt = "A chat between a curious user and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user's questions. USER: "
-                post_prompt = " ASSISTANT:"
-            if isinstance(input_text, list):
-                post_prompt = [input + post_prompt for input in input_text]
-            else:
-                post_prompt = input_text + post_prompt
-            if not isinstance(raw_image, list):
-                raw_image = [raw_image]
-            image = self.processor(raw_image)
-
-        elif self.model_type in ['llava', 'fuyu', 'kosmos-2']:
-            if self.model_type == "llava":
-                pre_prompt = "A chat between a curious user and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user's questions. USER: "
-                post_prompt = " ASSISTANT:"
-                if input_text is None:
-                    input_text = "\n Which city is this? Answer:"
-                if isinstance(input_text, list):
-                    post_prompt = [input + post_prompt for input in input_text]
-                else:
-                    post_prompt = input_text + post_prompt
-            elif self.model_type == 'fuyu':
-                pre_prompt = "Describe this image:"
-                post_prompt = None
-                if input_text is None:
-                    input_text = "Answer the following VQAv2 question based on the image: How many people are in the image?\n"
-            elif self.model_type == "kosmos-2":
-                pre_prompt = ""
-                post_prompt = None
-                if input_text is None:
-                    input_text = "An image of"
-
-            if self.model_type not in ['fuyu', 'kosmos-2']:
-                post_prompt = " ASSISTANT:"
-                if isinstance(input_text, list):
-                    post_prompt = [input + post_prompt for input in input_text]
-                else:
-                    post_prompt = input_text + post_prompt
-            else:
-                post_prompt = None
-            if self.model_type in ['fuyu', 'kosmos-2']:
-                image = self.processor(text=input_text,
-                                       images=raw_image,
-                                       return_tensors='pt')
-            else:
-                if isinstance(input_text, list):
-                    image = self.processor(text=input_text[0],
-                                           images=raw_image,
-                                           padding=True,
-                                           return_tensors="pt")['pixel_values']
-                else:
-                    image = self.processor(text=input_text,
-                                           images=raw_image,
-                                           return_tensors="pt")['pixel_values']
-
-        elif self.model_type in ['mllama']:
-            if raw_image is not None:
-                input_text = self.processor.apply_chat_template(
-                    images=raw_image, text=input_text)
-                inputs = self.processor(images=raw_image,
-                                        text=input_text,
-                                        return_tensors="pt")
-                other_vision_inputs = {
-                    "aspect_ratio_ids":
-                    inputs["aspect_ratio_ids"].to(self.device).expand(
-                        self.args.batch_size, -1).contiguous(),
-                    "aspect_ratio_mask":
-                    inputs["aspect_ratio_mask"].to(self.device).expand(
-                        self.args.batch_size, -1, -1).contiguous(),
-                }
-                other_decoder_inputs = {
-                    "cross_attention_mask":
-                    inputs["cross_attention_mask"].to(self.device).expand(
-                        self.args.batch_size, -1, -1, -1).contiguous(),
-                }
-                pre_prompt = input_text
-                post_prompt = None
-                image = inputs["pixel_values"]
-            else:
-                pre_prompt = input_text
-                post_prompt = None
-                image = None
-                logger.warning(
-                    "image_path is None. Will not pass image as input, skipping the vision encoder."
-                )
-                image = None
-        elif self.model_type in ['llava_onevision']:
-            pre_prompt = "<|im_start|>user " + "